1 static char help[] = "Two stream instability from Birdsal and Langdon with DMSwarm and TS basic symplectic integrators\n"; 2 3 #include <petscdmplex.h> 4 #include <petscfe.h> 5 #include <petscdmswarm.h> 6 #include <petscds.h> 7 #include <petscksp.h> 8 #include <petsc/private/petscfeimpl.h> 9 #include <petsc/private/tsimpl.h> 10 #include <petscts.h> 11 #include <petscmath.h> 12 13 typedef struct { 14 PetscInt dim; /* The topological mesh dimension */ 15 PetscBool simplex; /* Flag for simplices or tensor cells */ 16 PetscBool bdm; /* Flag for mixed form poisson */ 17 PetscBool monitor; /* Flag for use of the TS monitor */ 18 PetscBool uniform; /* Flag to uniformly space particles in x */ 19 char meshFilename[PETSC_MAX_PATH_LEN]; /* Name of the mesh filename if any */ 20 PetscReal sigma; /* Linear charge per box length */ 21 PetscReal timeScale; /* Nondimensionalizing time scaling */ 22 PetscInt particlesPerCell; /* The number of partices per cell */ 23 PetscReal particleRelDx; /* Relative particle position perturbation compared to average cell diameter h */ 24 PetscInt k; /* Mode number for test function */ 25 PetscReal momentTol; /* Tolerance for checking moment conservation */ 26 SNES snes; /* SNES object */ 27 PetscInt steps; /* TS iterations */ 28 PetscReal stepSize; /* Time stepper step size */ 29 PetscErrorCode (*func)(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar *, void *); 30 } AppCtx; 31 32 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options) { 33 PetscFunctionBeginUser; 34 options->dim = 2; 35 options->simplex = PETSC_TRUE; 36 options->monitor = PETSC_TRUE; 37 options->particlesPerCell = 1; 38 options->k = 1; 39 options->particleRelDx = 1.e-20; 40 options->momentTol = 100. * PETSC_MACHINE_EPSILON; 41 options->sigma = 1.; 42 options->timeScale = 1.0e-6; 43 options->uniform = PETSC_FALSE; 44 options->steps = 1; 45 options->stepSize = 0.01; 46 options->bdm = PETSC_FALSE; 47 48 PetscOptionsBegin(comm, "", "Two Stream options", "DMPLEX"); 49 PetscCall(PetscStrcpy(options->meshFilename, "")); 50 PetscCall(PetscOptionsInt("-dim", "The topological mesh dimension", "ex2.c", options->dim, &options->dim, NULL)); 51 PetscCall(PetscOptionsInt("-steps", "TS steps to take", "ex2.c", options->steps, &options->steps, NULL)); 52 PetscCall(PetscOptionsBool("-monitor", "To use the TS monitor or not", "ex2.c", options->monitor, &options->monitor, NULL)); 53 PetscCall(PetscOptionsBool("-simplex", "The flag for simplices or tensor cells", "ex2.c", options->simplex, &options->simplex, NULL)); 54 PetscCall(PetscOptionsBool("-uniform", "Uniform particle spacing", "ex2.c", options->uniform, &options->uniform, NULL)); 55 PetscCall(PetscOptionsBool("-bdm", "Use H1 instead of C0", "ex2.c", options->bdm, &options->bdm, NULL)); 56 PetscCall(PetscOptionsString("-mesh", "Name of the mesh filename if any", "ex2.c", options->meshFilename, options->meshFilename, PETSC_MAX_PATH_LEN, NULL)); 57 PetscCall(PetscOptionsInt("-k", "Mode number of test", "ex5.c", options->k, &options->k, NULL)); 58 PetscCall(PetscOptionsInt("-particlesPerCell", "Number of particles per cell", "ex2.c", options->particlesPerCell, &options->particlesPerCell, NULL)); 59 PetscCall(PetscOptionsReal("-sigma", "parameter", "<1>", options->sigma, &options->sigma, PETSC_NULL)); 60 PetscCall(PetscOptionsReal("-stepSize", "parameter", "<1e-2>", options->stepSize, &options->stepSize, PETSC_NULL)); 61 PetscCall(PetscOptionsReal("-timeScale", "parameter", "<1>", options->timeScale, &options->timeScale, PETSC_NULL)); 62 PetscCall(PetscOptionsReal("-particle_perturbation", "Relative perturbation of particles (0,1)", "ex2.c", options->particleRelDx, &options->particleRelDx, NULL)); 63 PetscOptionsEnd(); 64 PetscFunctionReturn(0); 65 } 66 67 static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user) { 68 PetscFunctionBeginUser; 69 PetscCall(DMCreate(comm, dm)); 70 PetscCall(DMSetType(*dm, DMPLEX)); 71 PetscCall(DMSetFromOptions(*dm)); 72 PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view")); 73 PetscFunctionReturn(0); 74 } 75 76 static void laplacian_f1(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[]) { 77 PetscInt d; 78 for (d = 0; d < dim; ++d) f1[d] = u_x[d]; 79 } 80 81 static void laplacian(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g3[]) { 82 PetscInt d; 83 for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0; 84 } 85 86 static PetscErrorCode CreateFEM(DM dm, AppCtx *user) { 87 PetscFE fe; 88 PetscDS ds; 89 DMPolytopeType ct; 90 PetscBool simplex; 91 PetscInt dim, cStart; 92 93 PetscFunctionBeginUser; 94 PetscCall(DMGetDimension(dm, &dim)); 95 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, NULL)); 96 PetscCall(DMPlexGetCellType(dm, cStart, &ct)); 97 simplex = DMPolytopeTypeGetNumVertices(ct) == DMPolytopeTypeGetDim(ct) + 1 ? PETSC_TRUE : PETSC_FALSE; 98 PetscCall(PetscFECreateDefault(PetscObjectComm((PetscObject)dm), dim, 1, simplex, NULL, -1, &fe)); 99 PetscCall(PetscObjectSetName((PetscObject)fe, "potential")); 100 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fe)); 101 PetscCall(DMCreateDS(dm)); 102 PetscCall(PetscFEDestroy(&fe)); 103 PetscCall(DMGetDS(dm, &ds)); 104 PetscCall(PetscDSSetResidual(ds, 0, NULL, laplacian_f1)); 105 PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian)); 106 PetscFunctionReturn(0); 107 } 108 109 /* 110 Initialize particle coordinates uniformly and with opposing velocities 111 */ 112 static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user) { 113 PetscRandom rnd, rndp; 114 PetscReal interval = user->particleRelDx; 115 PetscScalar value, *vals; 116 PetscReal *centroid, *coords, *xi0, *v0, *J, *invJ, detJ, *initialConditions, normalized_vel; 117 PetscInt *cellid, cStart; 118 PetscInt Ncell, Np = user->particlesPerCell, p, c, dim, d; 119 120 PetscFunctionBeginUser; 121 PetscCall(DMGetDimension(dm, &dim)); 122 PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw)); 123 PetscCall(DMSetType(*sw, DMSWARM)); 124 PetscCall(DMSetDimension(*sw, dim)); 125 PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd)); 126 PetscCall(PetscRandomSetInterval(rnd, 0.0, 1.0)); 127 PetscCall(PetscRandomSetFromOptions(rnd)); 128 PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rndp)); 129 PetscCall(PetscRandomSetInterval(rndp, -interval, interval)); 130 PetscCall(PetscRandomSetFromOptions(rndp)); 131 PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC)); 132 PetscCall(DMSwarmSetCellDM(*sw, dm)); 133 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR)); 134 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", dim, PETSC_REAL)); 135 PetscCall(DMSwarmFinalizeFieldRegister(*sw)); 136 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &Ncell)); 137 PetscCall(DMSwarmSetLocalSizes(*sw, Ncell * Np, 0)); 138 PetscCall(DMSetFromOptions(*sw)); 139 PetscCall(DMSwarmGetField(*sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 140 PetscCall(DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid)); 141 PetscCall(DMSwarmGetField(*sw, "w_q", NULL, NULL, (void **)&vals)); 142 PetscCall(DMSwarmGetField(*sw, "kinematics", NULL, NULL, (void **)&initialConditions)); 143 PetscCall(PetscMalloc5(dim, ¢roid, dim, &xi0, dim, &v0, dim * dim, &J, dim * dim, &invJ)); 144 for (c = cStart; c < Ncell; c++) { 145 if (Np == 1) { 146 PetscCall(DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL)); 147 cellid[c] = c; 148 for (d = 0; d < dim; ++d) coords[c * dim + d] = centroid[d]; 149 } else { 150 for (d = 0; d < dim; ++d) xi0[d] = -1.0; 151 PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ)); /* affine */ 152 for (p = 0; p < Np; ++p) { 153 const PetscInt n = c * Np + p; 154 PetscReal refcoords[3], spacing; 155 156 cellid[n] = c; 157 if (user->uniform) { 158 spacing = 2. / Np; 159 PetscCall(PetscRandomGetValue(rnd, &value)); 160 for (d = 0; d < dim; ++d) refcoords[d] = d == 0 ? -1. + spacing / 2. + p * spacing + value / 100. : 0.; 161 } else { 162 for (d = 0; d < dim; ++d) { 163 PetscCall(PetscRandomGetValue(rnd, &value)); 164 refcoords[d] = d == 0 ? PetscRealPart(value) : 0.; 165 } 166 } 167 CoordinatesRefToReal(dim, dim, xi0, v0, J, refcoords, &coords[n * dim]); 168 /* constant particle weights */ 169 for (d = 0; d < dim; ++d) vals[n] = user->sigma / Np; 170 } 171 } 172 } 173 PetscCall(PetscFree5(centroid, xi0, v0, J, invJ)); 174 normalized_vel = 1.; 175 for (c = 0; c < Ncell; ++c) { 176 for (p = 0; p < Np; ++p) { 177 if (p % 2 == 0) { 178 for (d = 0; d < dim; ++d) initialConditions[(c * Np + p) * dim + d] = d == 0 ? normalized_vel : 0.; 179 } else { 180 for (d = 0; d < dim; ++d) initialConditions[(c * Np + p) * dim + d] = d == 0 ? -(normalized_vel) : 0.; 181 } 182 } 183 } 184 PetscCall(DMSwarmRestoreField(*sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 185 PetscCall(DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid)); 186 PetscCall(DMSwarmRestoreField(*sw, "w_q", NULL, NULL, (void **)&vals)); 187 PetscCall(DMSwarmRestoreField(*sw, "kinematics", NULL, NULL, (void **)&initialConditions)); 188 PetscCall(PetscRandomDestroy(&rnd)); 189 PetscCall(PetscRandomDestroy(&rndp)); 190 PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles")); 191 PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view")); 192 PetscCall(DMLocalizeCoordinates(*sw)); 193 PetscFunctionReturn(0); 194 } 195 196 /* Solve for particle position updates */ 197 static PetscErrorCode RHSFunction1(TS ts, PetscReal t, Vec V, Vec Posres, void *ctx) { 198 const PetscScalar *v; 199 PetscScalar *posres; 200 PetscInt Np, p, dim, d; 201 DM dm; 202 203 PetscFunctionBeginUser; 204 PetscCall(VecGetLocalSize(Posres, &Np)); 205 PetscCall(VecGetArray(Posres, &posres)); 206 PetscCall(VecGetArrayRead(V, &v)); 207 PetscCall(TSGetDM(ts, &dm)); 208 PetscCall(DMGetDimension(dm, &dim)); 209 Np /= dim; 210 for (p = 0; p < Np; ++p) { 211 for (d = 0; d < dim; ++d) posres[p * dim + d] = v[p * dim + d]; 212 } 213 PetscCall(VecRestoreArrayRead(V, &v)); 214 PetscCall(VecRestoreArray(Posres, &posres)); 215 PetscFunctionReturn(0); 216 } 217 218 /* 219 Solve for the gradient of the electric field and apply force to particles. 220 */ 221 static PetscErrorCode RHSFunction2(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx) { 222 AppCtx *user = (AppCtx *)ctx; 223 DM dm, plex; 224 PetscDS prob; 225 PetscFE fe; 226 Mat M_p; 227 Vec phi, locPhi, rho, f; 228 const PetscScalar *x; 229 PetscScalar *vres; 230 PetscReal *coords, phi_0; 231 PetscInt dim, d, cStart, cEnd, cell, cdim; 232 PetscReal m_e = 9.11e-31, q_e = 1.60e-19, epsi_0 = 8.85e-12; 233 234 PetscFunctionBeginUser; 235 PetscObjectSetName((PetscObject)X, "rhsf2 position"); 236 VecViewFromOptions(X, NULL, "-rhsf2_x_view"); 237 PetscCall(VecGetArrayRead(X, &x)); 238 PetscCall(VecGetArray(Vres, &vres)); 239 PetscCall(TSGetDM(ts, &dm)); 240 PetscCall(DMGetDimension(dm, &dim)); 241 PetscCall(SNESGetDM(user->snes, &plex)); 242 PetscCall(DMGetCoordinateDim(plex, &cdim)); 243 PetscCall(DMGetDS(plex, &prob)); 244 PetscCall(PetscDSGetDiscretization(prob, 0, (PetscObject *)&fe)); 245 PetscCall(DMGetGlobalVector(plex, &phi)); 246 PetscCall(DMGetLocalVector(plex, &locPhi)); 247 PetscCall(DMCreateMassMatrix(dm, plex, &M_p)); 248 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view")); 249 PetscCall(DMGetGlobalVector(plex, &rho)); 250 PetscCall(DMSwarmCreateGlobalVectorFromField(dm, "w_q", &f)); 251 PetscCall(PetscObjectSetName((PetscObject)f, "weights vector")); 252 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 253 PetscCall(MatMultTranspose(M_p, f, rho)); 254 PetscCall(DMSwarmDestroyGlobalVectorFromField(dm, "w_q", &f)); 255 PetscCall(PetscObjectSetName((PetscObject)rho, "rho")); 256 PetscCall(VecViewFromOptions(rho, NULL, "-poisson_rho_view")); 257 /* Take nullspace out of rhs */ 258 { 259 PetscScalar sum; 260 PetscInt n; 261 phi_0 = (user->sigma * user->sigma * user->sigma) * (user->timeScale * user->timeScale) / (m_e * q_e * epsi_0); 262 263 PetscCall(VecGetSize(rho, &n)); 264 PetscCall(VecSum(rho, &sum)); 265 PetscCall(VecShift(rho, -sum / n)); 266 267 PetscCall(VecSum(rho, &sum)); 268 PetscCheck(PetscAbsScalar(sum) <= 1.0e-10, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Charge should have no DC component %g", (double)PetscAbsScalar(sum)); 269 PetscCall(VecScale(rho, phi_0)); 270 } 271 PetscCall(VecSet(phi, 0.0)); 272 PetscCall(SNESSolve(user->snes, rho, phi)); 273 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view")); 274 PetscCall(DMRestoreGlobalVector(plex, &rho)); 275 PetscCall(MatDestroy(&M_p)); 276 PetscCall(DMGlobalToLocalBegin(plex, phi, INSERT_VALUES, locPhi)); 277 PetscCall(DMGlobalToLocalEnd(plex, phi, INSERT_VALUES, locPhi)); 278 PetscCall(DMSwarmSortGetAccess(dm)); 279 PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 280 PetscCall(DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd)); 281 for (cell = cStart; cell < cEnd; ++cell) { 282 PetscTabulation tab; 283 PetscReal v[3], J[9], invJ[9], detJ; 284 PetscScalar *ph = PETSC_NULL; 285 PetscReal *pcoord = PETSC_NULL; 286 PetscReal *refcoord = PETSC_NULL; 287 PetscInt *points = PETSC_NULL, Ncp, cp; 288 PetscScalar gradPhi[3]; 289 290 PetscCall(DMPlexComputeCellGeometryFEM(plex, cell, NULL, v, J, invJ, &detJ)); 291 PetscCall(DMSwarmSortGetPointsPerCell(dm, cell, &Ncp, &points)); 292 PetscCall(DMGetWorkArray(dm, Ncp * cdim, MPIU_REAL, &pcoord)); 293 PetscCall(DMGetWorkArray(dm, Ncp * cdim, MPIU_REAL, &refcoord)); 294 for (cp = 0; cp < Ncp; ++cp) { 295 for (d = 0; d < cdim; ++d) pcoord[cp * cdim + d] = coords[points[cp] * cdim + d]; 296 } 297 PetscCall(DMPlexCoordinatesToReference(plex, cell, Ncp, pcoord, refcoord)); 298 PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab)); 299 PetscCall(DMPlexVecGetClosure(plex, NULL, locPhi, cell, NULL, &ph)); 300 for (cp = 0; cp < Ncp; ++cp) { 301 const PetscInt p = points[cp]; 302 gradPhi[0] = 0.0; 303 gradPhi[1] = 0.0; 304 gradPhi[2] = 0.0; 305 const PetscReal *basisDer = tab->T[1]; 306 307 PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, ph, cdim, invJ, NULL, cp, gradPhi)); 308 for (d = 0; d < cdim; ++d) vres[p * cdim + d] = d == 0 ? gradPhi[d] : 0.; 309 } 310 PetscCall(DMPlexVecRestoreClosure(plex, NULL, locPhi, cell, NULL, &ph)); 311 PetscCall(PetscTabulationDestroy(&tab)); 312 PetscCall(DMRestoreWorkArray(dm, Ncp * cdim, MPIU_REAL, &pcoord)); 313 PetscCall(DMRestoreWorkArray(dm, Ncp * cdim, MPIU_REAL, &refcoord)); 314 PetscCall(PetscFree(points)); 315 } 316 PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 317 PetscCall(DMSwarmSortRestoreAccess(dm)); 318 PetscCall(DMRestoreLocalVector(plex, &locPhi)); 319 PetscCall(DMRestoreGlobalVector(plex, &phi)); 320 PetscCall(VecRestoreArray(Vres, &vres)); 321 PetscCall(VecRestoreArrayRead(X, &x)); 322 PetscCall(VecViewFromOptions(Vres, NULL, "-vel_res_view")); 323 PetscFunctionReturn(0); 324 } 325 326 int main(int argc, char **argv) { 327 PetscInt i, par; 328 PetscInt locSize, p, d, dim, Np, step, *idx1, *idx2; 329 TS ts; 330 DM dm, sw; 331 AppCtx user; 332 MPI_Comm comm; 333 Vec coorVec, kinVec, probVec, solution, position, momentum; 334 const PetscScalar *coorArr, *kinArr; 335 PetscReal ftime = 10., *probArr, *probVecArr; 336 IS is1, is2; 337 PetscReal *coor, *kin, *pos, *mom; 338 339 PetscFunctionBeginUser; 340 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 341 comm = PETSC_COMM_WORLD; 342 PetscCall(ProcessOptions(comm, &user)); 343 /* Create dm and particles */ 344 PetscCall(CreateMesh(comm, &dm, &user)); 345 PetscCall(CreateFEM(dm, &user)); 346 PetscCall(CreateParticles(dm, &sw, &user)); 347 PetscCall(SNESCreate(comm, &user.snes)); 348 PetscCall(SNESSetDM(user.snes, dm)); 349 PetscCall(DMPlexSetSNESLocalFEM(dm, &user, &user, &user)); 350 PetscCall(SNESSetFromOptions(user.snes)); 351 { 352 Mat J; 353 MatNullSpace nullSpace; 354 355 PetscCall(DMCreateMatrix(dm, &J)); 356 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace)); 357 PetscCall(MatSetNullSpace(J, nullSpace)); 358 PetscCall(MatNullSpaceDestroy(&nullSpace)); 359 PetscCall(SNESSetJacobian(user.snes, J, J, NULL, NULL)); 360 PetscCall(MatDestroy(&J)); 361 } 362 /* Place TSSolve in a loop to handle resetting the TS at every manual call of TSStep() */ 363 PetscCall(TSCreate(comm, &ts)); 364 PetscCall(TSSetMaxTime(ts, ftime)); 365 PetscCall(TSSetTimeStep(ts, user.stepSize)); 366 PetscCall(TSSetMaxSteps(ts, 100000)); 367 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 368 for (step = 0; step < user.steps; ++step) { 369 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &kinVec)); 370 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &coorVec)); 371 PetscCall(VecViewFromOptions(kinVec, NULL, "-ic_vec_view")); 372 PetscCall(DMGetDimension(sw, &dim)); 373 PetscCall(VecGetLocalSize(kinVec, &locSize)); 374 PetscCall(PetscMalloc1(locSize, &idx1)); 375 PetscCall(PetscMalloc1(locSize, &idx2)); 376 PetscCall(PetscMalloc1(2 * locSize, &probArr)); 377 Np = locSize / dim; 378 PetscCall(VecGetArrayRead(kinVec, &kinArr)); 379 PetscCall(VecGetArrayRead(coorVec, &coorArr)); 380 for (p = 0; p < Np; ++p) { 381 for (d = 0; d < dim; ++d) { 382 probArr[p * 2 * dim + d] = coorArr[p * dim + d]; 383 probArr[(p * 2 + 1) * dim + d] = kinArr[p * dim + d]; 384 } 385 } 386 PetscCall(VecRestoreArrayRead(kinVec, &kinArr)); 387 PetscCall(VecRestoreArrayRead(coorVec, &coorArr)); 388 /* Allocate for IS Strides that will contain x, y and vx, vy */ 389 for (p = 0; p < Np; ++p) { 390 for (d = 0; d < dim; ++d) { 391 idx1[p * dim + d] = (p * 2 + 0) * dim + d; 392 idx2[p * dim + d] = (p * 2 + 1) * dim + d; 393 } 394 } 395 396 PetscCall(ISCreateGeneral(comm, locSize, idx1, PETSC_OWN_POINTER, &is1)); 397 PetscCall(ISCreateGeneral(comm, locSize, idx2, PETSC_OWN_POINTER, &is2)); 398 /* DM needs to be set before splits so it propogates to sub TSs */ 399 PetscCall(TSSetDM(ts, sw)); 400 PetscCall(TSSetType(ts, TSBASICSYMPLECTIC)); 401 PetscCall(TSRHSSplitSetIS(ts, "position", is1)); 402 PetscCall(TSRHSSplitSetIS(ts, "momentum", is2)); 403 PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunction1, &user)); 404 PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunction2, &user)); 405 PetscCall(TSSetTime(ts, step * user.stepSize)); 406 if (step == 0) PetscCall(TSSetFromOptions(ts)); 407 /* Compose vector from array for TS solve with all kinematic variables */ 408 PetscCall(VecCreate(comm, &probVec)); 409 PetscCall(VecSetBlockSize(probVec, 1)); 410 PetscCall(VecSetSizes(probVec, PETSC_DECIDE, 2 * locSize)); 411 PetscCall(VecSetUp(probVec)); 412 PetscCall(VecGetArray(probVec, &probVecArr)); 413 for (i = 0; i < 2 * locSize; ++i) probVecArr[i] = probArr[i]; 414 PetscCall(VecRestoreArray(probVec, &probVecArr)); 415 PetscCall(TSSetSolution(ts, probVec)); 416 PetscCall(PetscFree(probArr)); 417 PetscCall(VecViewFromOptions(kinVec, NULL, "-ic_view")); 418 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &kinVec)); 419 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &coorVec)); 420 PetscCall(TSMonitor(ts, step, ts->ptime, ts->vec_sol)); 421 if (!ts->steprollback) PetscCall(TSPreStep(ts)); 422 PetscCall(TSStep(ts)); 423 if (ts->steprollback) PetscCall(TSPostEvaluate(ts)); 424 if (!ts->steprollback) { 425 TSPostStep(ts); 426 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coor)); 427 PetscCall(DMSwarmGetField(sw, "kinematics", NULL, NULL, (void **)&kin)); 428 PetscCall(TSGetSolution(ts, &solution)); 429 PetscCall(VecGetSubVector(solution, is1, &position)); 430 PetscCall(VecGetSubVector(solution, is2, &momentum)); 431 PetscCall(VecGetArray(position, &pos)); 432 PetscCall(VecGetArray(momentum, &mom)); 433 for (par = 0; par < Np; ++par) { 434 for (d = 0; d < dim; ++d) { 435 if (pos[par * dim + d] < 0.) { 436 coor[par * dim + d] = pos[par * dim + d] + 2. * PETSC_PI; 437 } else if (pos[par * dim + d] > 2. * PETSC_PI) { 438 coor[par * dim + d] = pos[par * dim + d] - 2. * PETSC_PI; 439 } else { 440 coor[par * dim + d] = pos[par * dim + d]; 441 } 442 kin[par * dim + d] = mom[par * dim + d]; 443 } 444 } 445 PetscCall(VecRestoreArray(position, &pos)); 446 PetscCall(VecRestoreArray(momentum, &mom)); 447 PetscCall(VecRestoreSubVector(solution, is1, &position)); 448 PetscCall(VecRestoreSubVector(solution, is2, &momentum)); 449 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coor)); 450 PetscCall(DMSwarmRestoreField(sw, "kinematics", NULL, NULL, (void **)&kin)); 451 } 452 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE)); 453 PetscCall(DMLocalizeCoordinates(sw)); 454 PetscCall(TSReset(ts)); 455 PetscCall(VecDestroy(&probVec)); 456 PetscCall(ISDestroy(&is1)); 457 PetscCall(ISDestroy(&is2)); 458 } 459 PetscCall(SNESDestroy(&user.snes)); 460 PetscCall(TSDestroy(&ts)); 461 PetscCall(DMDestroy(&sw)); 462 PetscCall(DMDestroy(&dm)); 463 PetscCall(PetscFinalize()); 464 return 0; 465 } 466 467 /*TEST 468 469 build: 470 requires: triangle !single !complex 471 test: 472 suffix: bsi1q3 473 args: -particlesPerCell 200\ 474 -petscspace_degree 2\ 475 -petscfe_default_quadrature_order 3\ 476 -ts_basicsymplectic_type 1\ 477 -pc_type svd\ 478 -uniform\ 479 -sigma 1.0e-8\ 480 -timeScale 2.0e-14\ 481 -stepSize 1.0e-2\ 482 -ts_monitor_sp_swarm\ 483 -steps 10\ 484 -dm_view\ 485 -dm_plex_simplex 0 -dm_plex_dim 2\ 486 -dm_plex_box_lower 0,-1\ 487 -dm_plex_box_upper 6.283185307179586,1\ 488 -dm_plex_box_bd periodic,none\ 489 -dm_plex_box_faces 4,1 490 test: 491 suffix: bsi2q3 492 args: -particlesPerCell 200\ 493 -petscspace_degree 2\ 494 -petscfe_default_quadrature_order 3\ 495 -ts_basicsymplectic_type 2\ 496 -pc_type svd\ 497 -uniform\ 498 -sigma 1.0e-8\ 499 -timeScale 2.0e-14\ 500 -stepSize 1.0e-2\ 501 -ts_monitor_sp_swarm\ 502 -steps 10\ 503 -dm_view\ 504 -dm_plex_simplex 0 -dm_plex_dim 2\ 505 -dm_plex_box_lower 0,-1\ 506 -dm_plex_box_upper 6.283185307179586,1\ 507 -dm_plex_box_bd periodic,none\ 508 -dm_plex_box_faces 4,1 509 TEST*/ 510