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