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