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