Actual source code: ex3.c

  1: static char help[] = "Landau Damping test using Vlasov-Poisson equations\n";

  3: /*
  4:   This example solves the Vlasov-Poisson system for Landau damping (6X + 6V).
  5:   The system is solved using a Particle-In-Cell (PIC) method with DMSwarm for particles and DMPlex/PetscFE for the field.
  6:   This particular example uses the velocity mesh from DMPlexLandauCreateVelocitySpace for 3D velocity space.

  8:   Options:
  9:   -particle_monitor [prefix] : Output particle data (x, v, w, E) to binary files at each output step.
 10:                                Optional prefix for filenames (default: "particles").

 12: */
 13: #include <petscts.h>
 14: #include <petscdmplex.h>
 15: #include <petscdmswarm.h>
 16: #include <petscfe.h>
 17: #include <petscds.h>
 18: #include <petscbag.h>
 19: #include <petscdraw.h>
 20: #include <petscviewer.h>
 21: #include <petsclandau.h>
 22: #include <petscdmcomposite.h>
 23: #include <petsc/private/dmpleximpl.h>
 24: #include <petsc/private/petscfeimpl.h>
 25: #include <petsc/private/dmswarmimpl.h>
 26: #include "petscdm.h"
 27: #include "petscdmlabel.h"

 29: PETSC_EXTERN PetscErrorCode stream(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
 30: PETSC_EXTERN PetscErrorCode line(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);

 32: typedef enum {
 33:   V0,
 34:   X0,
 35:   T0,
 36:   M0,
 37:   Q0,
 38:   PHI0,
 39:   POISSON,
 40:   VLASOV,
 41:   SIGMA,
 42:   NUM_CONSTANTS
 43: } ConstantType;
 44: typedef struct {
 45:   PetscScalar v0; /* Velocity scale, often the thermal velocity */
 46:   PetscScalar t0; /* Time scale */
 47:   PetscScalar x0; /* Space scale */
 48:   PetscScalar m0; /* Mass scale */
 49:   PetscScalar q0; /* Charge scale */
 50:   PetscScalar kb;
 51:   PetscScalar epsi0;
 52:   PetscScalar phi0;          /* Potential scale */
 53:   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
 54:   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
 55:   PetscReal   sigma;         /* Nondimensional charge per length in x */
 56: } Parameter;

 58: typedef struct {
 59:   PetscBag      bag;            // Problem parameters
 60:   PetscBool     error;          // Flag for printing the error
 61:   PetscBool     efield_monitor; // Flag to show electric field monitor
 62:   PetscBool     moment_monitor; // Flag to show distribution moment monitor
 63:   char          particle_monitor_prefix[PETSC_MAX_PATH_LEN];
 64:   PetscBool     particle_monitor; // Flag to output particle data
 65:   PetscInt      ostep;            // Print the energy at each ostep time steps
 66:   PetscInt      numParticles;
 67:   PetscReal     timeScale;              /* Nondimensionalizing time scale */
 68:   PetscReal     charges[2];             /* The charges of each species */
 69:   PetscReal     masses[2];              /* The masses of each species */
 70:   PetscReal     thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
 71:   PetscReal     cosine_coefficients[2]; /*(alpha, k)*/
 72:   PetscReal     totalWeight;
 73:   PetscReal     stepSize;
 74:   PetscInt      steps;
 75:   PetscReal     initVel;
 76:   SNES          snes;       // EM solver
 77:   DM            dmPot;      // The DM for potential
 78:   Mat           M;          // The finite element mass matrix for potential
 79:   PetscFEGeom  *fegeom;     // Geometric information for the DM cells
 80:   PetscBool     validE;     // Flag to indicate E-field in swarm is valid
 81:   PetscReal     drawlgEmin; // The minimum lg(E) to plot
 82:   PetscDrawLG   drawlgE;    // Logarithm of maximum electric field
 83:   DM            swarm;
 84:   PetscBool     checkweights;
 85:   PetscInt      checkVRes; // Flag to check/output velocity residuals for nightly tests
 86:   DM            landau_pack;
 87:   PetscBool     use_landau_velocity_space;
 88:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
 89: } AppCtx;

 91: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
 92: {
 93:   PetscFunctionBeginUser;
 94:   PetscInt d                         = 2;
 95:   PetscInt maxSpecies                = 2;
 96:   options->error                     = PETSC_FALSE;
 97:   options->efield_monitor            = PETSC_FALSE;
 98:   options->moment_monitor            = PETSC_FALSE;
 99:   options->particle_monitor          = PETSC_FALSE;
100:   options->ostep                     = 100;
101:   options->timeScale                 = 2.0e-14;
102:   options->charges[0]                = -1.0;
103:   options->charges[1]                = 1.0;
104:   options->masses[0]                 = 1.0;
105:   options->masses[1]                 = 1000.0;
106:   options->thermal_energy[0]         = 1.0;
107:   options->thermal_energy[1]         = 1.0;
108:   options->cosine_coefficients[0]    = 0.01;
109:   options->cosine_coefficients[1]    = 0.5;
110:   options->initVel                   = 1;
111:   options->totalWeight               = 1.0;
112:   options->validE                    = PETSC_FALSE;
113:   options->drawlgEmin                = -6;
114:   options->drawlgE                   = NULL;
115:   options->snes                      = NULL;
116:   options->dmPot                     = NULL;
117:   options->M                         = NULL;
118:   options->numParticles              = 32768;
119:   options->checkweights              = PETSC_FALSE;
120:   options->checkVRes                 = 0;
121:   options->landau_pack               = NULL;
122:   options->use_landau_velocity_space = PETSC_FALSE;

124:   PetscOptionsBegin(comm, "", "Landau Damping options", "DMSWARM");
125:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex3.c", options->error, &options->error, NULL));
126:   PetscCall(PetscOptionsBool("-efield_monitor", "Flag to plot log(max E) over time", "ex3.c", options->efield_monitor, &options->efield_monitor, NULL));
127:   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", "ex3.c", options->drawlgEmin, &options->drawlgEmin, NULL));
128:   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", "ex3.c", options->moment_monitor, &options->moment_monitor, NULL));
129:   PetscCall(PetscOptionsString("-particle_monitor", "Prefix for particle data files", "ex3.c", options->particle_monitor_prefix, options->particle_monitor_prefix, sizeof(options->particle_monitor_prefix), &options->particle_monitor));
130:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex3.c", options->checkweights, &options->checkweights, NULL));
131:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex3.c", options->checkVRes, &options->checkVRes, NULL));
132:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex3.c", options->ostep, &options->ostep, NULL));
133:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex3.c", options->timeScale, &options->timeScale, NULL));
134:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex3.c", options->initVel, &options->initVel, NULL));
135:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex3.c", options->totalWeight, &options->totalWeight, NULL));
136:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex3.c", options->cosine_coefficients, &d, NULL));
137:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex3.c", options->charges, &maxSpecies, NULL));
138:   PetscCall(PetscOptionsBool("-use_landau_velocity_space", "Use Landau velocity space", "ex3.c", options->use_landau_velocity_space, &options->use_landau_velocity_space, NULL));
139:   PetscOptionsEnd();

141:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
142:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
143:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
144:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
145:   PetscFunctionReturn(PETSC_SUCCESS);
146: }

148: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
149: {
150:   MPI_Comm comm;

152:   PetscFunctionBeginUser;
153:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
154:   if (user->efield_monitor) {
155:     PetscDraw     draw;
156:     PetscDrawAxis axis;

158:     PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
159:     PetscCall(PetscDrawSetSave(draw, "ex3_Efield"));
160:     PetscCall(PetscDrawSetFromOptions(draw));
161:     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
162:     PetscCall(PetscDrawDestroy(&draw));
163:     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
164:     PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
165:     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
166:     PetscCall(PetscDrawLGSetUseMarkers(user->drawlgE, PETSC_FALSE));
167:   }
168:   PetscFunctionReturn(PETSC_SUCCESS);
169: }

171: static PetscErrorCode DestroyContext(AppCtx *user)
172: {
173:   PetscFunctionBeginUser;
174:   if (user->landau_pack) PetscCall(DMPlexLandauDestroyVelocitySpace(&user->landau_pack));
175:   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
176:   PetscCall(PetscBagDestroy(&user->bag));
177:   PetscFunctionReturn(PETSC_SUCCESS);
178: }

180: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
181: {
182:   const PetscScalar *w;
183:   PetscInt           Np;

185:   PetscFunctionBeginUser;
186:   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
187:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
188:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
189:   for (PetscInt p = 0; p < Np; ++p) PetscCheck(w[p] >= 0.0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " has negative weight %g", p, w[p]);
190:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
191:   PetscFunctionReturn(PETSC_SUCCESS);
192: }

194: static void f0_Dirichlet(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 f0[])
195: {
196:   for (PetscInt d = 0; d < dim; ++d) f0[0] += 0.5 * PetscSqr(u_x[d]);
197: }

199: static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
200: {
201:   PetscDS        ds;
202:   const PetscInt field = 0;
203:   PetscInt       Nf;
204:   void          *ctx;

206:   PetscFunctionBegin;
207:   PetscCall(DMGetApplicationContext(dm, &ctx));
208:   PetscCall(DMGetDS(dm, &ds));
209:   PetscCall(PetscDSGetNumFields(ds, &Nf));
210:   PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
211:   PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
212:   PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
213:   PetscFunctionReturn(PETSC_SUCCESS);
214: }

216: static void f0_grad_phi2(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 f0[])
217: {
218:   f0[0] = 0.;
219:   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d]); // + d * dim  cause segfault
220: }

222: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
223: {
224:   AppCtx     *user = (AppCtx *)ctx;
225:   DM          sw;
226:   PetscScalar intESq;
227:   PetscReal  *E, *x, *weight;
228:   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
229:   PetscReal   pmoments[16]; /* \int f, \int v f, \int v^2 f */
230:   PetscInt   *species, dim, Np, gNp;
231:   MPI_Comm    comm;

233:   PetscFunctionBeginUser;
234:   if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
235:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
236:   PetscCall(TSGetDM(ts, &sw));
237:   PetscCall(DMGetDimension(sw, &dim));
238:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
239:   PetscCall(DMSwarmGetSize(sw, &gNp));
240:   PetscCall(DMSwarmSortGetAccess(sw));
241:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
242:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
243:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
244:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

246:   for (PetscInt p = 0; p < Np; ++p) {
247:     PetscReal Emag = 0.;
248:     for (PetscInt d = 0; d < dim; ++d) {
249:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
250:       if (temp > Emax) Emax = temp;
251:       Emag += PetscSqr(E[p * dim + d]);
252:     }
253:     Enorm += PetscSqrtReal(Emag);
254:     sum += E[p * dim];
255:     chargesum += user->charges[0] * weight[p];
256:   }
257:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
258:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Enorm, 1, MPIU_REAL, MPIU_SUM, comm));
259:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &sum, 1, MPIU_REAL, MPIU_SUM, comm));
260:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &chargesum, 1, MPIU_REAL, MPIU_SUM, comm));
261:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
262:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;
263:   if (lgEmax < user->drawlgEmin) lgEmax = user->drawlgEmin;

265:   PetscDS ds;
266:   Vec     phi;

268:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
269:   PetscCall(DMGetDS(user->dmPot, &ds));
270:   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
271:   PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user));
272:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

274:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
275:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
276:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
277:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
278:   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
279:   PetscCall(PetscDrawLGDraw(user->drawlgE));
280:   PetscDraw draw;
281:   PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
282:   PetscCall(PetscDrawSave(draw));

284:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
285:   PetscCall(PetscPrintf(comm, "E: %f\t%+e\t%e\t%f\t%20.15e\t%f\t%f\t%f\t%20.15e\t%20.15e\t%20.15e\t%" PetscInt_FMT "\t(%" PetscInt_FMT ")\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[1 + dim], (double)PetscSqrtReal(intESq), gNp, step));
286:   PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
287:   PetscFunctionReturn(PETSC_SUCCESS);
288: }

290: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
291: {
292:   DM        sw;
293:   PetscReal pmoments[16]; /* \int f, \int v f, \int v^2 f */

295:   PetscFunctionBeginUser;
296:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
297:   PetscCall(TSGetDM(ts, &sw));

299:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));

301:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3]));
302:   PetscFunctionReturn(PETSC_SUCCESS);
303: }

305: static PetscErrorCode MonitorParticles(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
306: {
307:   AppCtx            *user = (AppCtx *)ctx;
308:   DM                 sw;
309:   PetscInt           Np, dim;
310:   const PetscReal   *x, *v, *E;
311:   const PetscScalar *w;
312:   PetscViewer        viewer;
313:   char               filename[64];
314:   MPI_Comm           comm;

316:   PetscFunctionBeginUser;
317:   if (step < 0 || step % user->ostep != 0) PetscFunctionReturn(PETSC_SUCCESS);
318:   PetscCall(TSGetDM(ts, &sw));
319:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
320:   PetscCall(DMGetDimension(sw, &dim));
321:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

323:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
324:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
325:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
326:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));

328:   if (user->particle_monitor_prefix[0]) {
329:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "%s_step_%d.bin", user->particle_monitor_prefix, (int)step));
330:   } else {
331:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "particles_step_%d.bin", (int)step));
332:   }
333:   PetscCall(PetscViewerBinaryOpen(comm, filename, FILE_MODE_WRITE, &viewer));

335:   {
336:     Vec vx, vv, vw, vE;
337:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, x, &vx));
338:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, v, &vv));
339:     PetscCall(VecCreateMPIWithArray(comm, 1, Np, PETSC_DECIDE, w, &vw));
340:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, E, &vE));

342:     PetscCall(VecView(vx, viewer));
343:     PetscCall(VecView(vv, viewer));
344:     PetscCall(VecView(vw, viewer));
345:     PetscCall(VecView(vE, viewer));

347:     PetscCall(VecDestroy(&vx));
348:     PetscCall(VecDestroy(&vv));
349:     PetscCall(VecDestroy(&vw));
350:     PetscCall(VecDestroy(&vE));
351:   }
352:   PetscCall(PetscViewerDestroy(&viewer));

354:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
355:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
356:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
357:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
358:   PetscFunctionReturn(PETSC_SUCCESS);
359: }

361: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
362: {
363:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (4. * PETSC_PI * 4. * PETSC_PI);
364:   return PETSC_SUCCESS;
365: }
366: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
367: {
368:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
369:   return PETSC_SUCCESS;
370: }

372: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
373: {
374:   const PetscReal alpha = scale ? scale[0] : 0.0;
375:   const PetscReal k     = scale ? scale[1] : 1.;
376:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
377:   return PETSC_SUCCESS;
378: }

380: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
381: {
382:   const PetscReal alpha = scale ? scale[0] : 0.;
383:   const PetscReal k     = scale ? scale[1] : 1.;
384:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
385:   return PETSC_SUCCESS;
386: }

388: PetscErrorCode PetscPDFCosine3D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
389: {
390:   const PetscReal alpha = scale ? scale[0] : 0.;
391:   const PetscReal k     = scale ? scale[1] : 1.;
392:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
393:   return PETSC_SUCCESS;
394: }

396: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
397: {
398:   PetscFE        fe;
399:   DMPolytopeType ct;
400:   PetscInt       dim, cStart;
401:   const char    *prefix = "v";
402:   AppCtx        *user;

404:   PetscFunctionBegin;
405:   PetscCall(DMGetDimension(sw, &dim));
406:   PetscCall(DMGetApplicationContext(sw, &user));
407:   if (dim == 3 && user->use_landau_velocity_space) {
408:     LandauCtx *ctx;
409:     Vec        X;
410:     Mat        J;

412:     PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, prefix, &X, &J, &user->landau_pack));
413:     PetscCall(DMGetApplicationContext(user->landau_pack, &ctx));
414:     *vdm = ctx->plex[0];
415:     PetscCall(PetscObjectReference((PetscObject)*vdm));
416:     PetscCall(VecDestroy(&X));
417:     PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
418:   } else {
419:     PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
420:     PetscCall(DMSetType(*vdm, DMPLEX));
421:     PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
422:     PetscCall(DMSetFromOptions(*vdm));
423:     PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
424:   }
425:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

427:   PetscCall(DMGetDimension(*vdm, &dim));
428:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
429:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
430:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
431:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
432:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
433:   PetscCall(DMCreateDS(*vdm));
434:   PetscCall(PetscFEDestroy(&fe));
435:   PetscFunctionReturn(PETSC_SUCCESS);
436: }

438: /*
439:   InitializeParticles_Centroid - Initialize a regular grid of particles.

441:   Input Parameters:
442: + sw      - The `DMSWARM`
443: - force1D - Treat the spatial domain as 1D

445:   Notes:
446:   This functions sets the species, cellid, spatial coordinate, and velocity fields for all particles.

448:   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
449:   and velocity meshes.
450: */
451: static PetscErrorCode InitializeParticles_Centroid(DM sw)
452: {
453:   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
454:   DMSwarmCellDM celldm;
455:   DM            xdm, vdm;
456:   PetscReal     vmin[3], vmax[3];
457:   PetscReal    *x, *v;
458:   PetscInt     *species, *cellid;
459:   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
460:   PetscBool     flg;
461:   MPI_Comm      comm;
462:   const char   *cellidname;

464:   PetscFunctionBegin;
465:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

467:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
468:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
469:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
470:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
471:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
472:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
473:   PetscOptionsEnd();
474:   debug = swarm->printCoords;

476:   PetscCall(DMGetDimension(sw, &dim));
477:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
478:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

480:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
481:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
482:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

484:   // One particle per centroid on the tensor product grid
485:   Npc = (vcEnd - vcStart) * Ns;
486:   Np  = (xcEnd - xcStart) * Npc;
487:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
488:   if (debug) {
489:     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
490:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
491:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
492:     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
493:     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
494:     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
495:   }

497:   // Set species and cellid
498:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
499:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
500:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
501:   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
502:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
503:     for (PetscInt s = 0; s < Ns; ++s) {
504:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
505:         species[p] = s;
506:         cellid[p]  = c;
507:       }
508:     }
509:   }
510:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
511:   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));

513:   // Set particle coordinates
514:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
515:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
516:   PetscCall(DMSwarmSortGetAccess(sw));
517:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
518:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
519:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
520:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
521:     const PetscInt cell = c + xcStart;
522:     PetscInt      *pidx, Npc;
523:     PetscReal      centroid[3], volume;

525:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
526:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
527:     for (PetscInt s = 0; s < Ns; ++s) {
528:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
529:         const PetscInt p  = pidx[q * Ns + s];
530:         const PetscInt vc = vcStart + q;
531:         PetscReal      vcentroid[3], vvolume;

533:         PetscCall(DMPlexComputeCellGeometryFVM(vdm, vc, &vvolume, vcentroid, NULL));
534:         for (PetscInt d = 0; d < dim; ++d) {
535:           x[p * dim + d] = centroid[d];
536:           v[p * dim + d] = vcentroid[d];
537:         }

539:         if (debug > 1) {
540:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
541:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
542:           for (PetscInt d = 0; d < dim; ++d) {
543:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
544:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
545:           }
546:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
547:           for (PetscInt d = 0; d < dim; ++d) {
548:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
549:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
550:           }
551:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
552:         }
553:       }
554:     }
555:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
556:   }
557:   PetscCall(DMSwarmSortRestoreAccess(sw));
558:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
559:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
560:   PetscFunctionReturn(PETSC_SUCCESS);
561: }

563: /*
564:   InitializeWeights - Compute weight for each local particle

566:   Input Parameters:
567: + sw          - The `DMSwarm`
568: . totalWeight - The sum of all particle weights
569: . func        - The PDF for the particle spatial distribution
570: - param       - The PDF parameters

572:   Notes:
573:   The PDF for velocity is assumed to be a Gaussian

575:   The particle weights are returned in the `w_q` field of `sw`.
576: */
577: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
578: {
579:   DM               xdm, vdm;
580:   DMSwarmCellDM    celldm;
581:   PetscScalar     *weight;
582:   PetscQuadrature  xquad;
583:   const PetscReal *xq, *xwq;
584:   const PetscInt   order = 5;
585:   PetscReal        xi0[3];
586:   PetscReal        xwtot = 0., pwtot = 0.;
587:   PetscInt         xNq;
588:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
589:   MPI_Comm         comm;
590:   PetscMPIInt      rank;

592:   PetscFunctionBegin;
593:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
594:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
595:   PetscCall(DMGetDimension(sw, &dim));
596:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
597:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
598:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
599:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
600:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
601:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

603:   // Setup Quadrature for spatial and velocity weight calculations
604:   PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad));
605:   PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq));
606:   for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0;

608:   // Integrate the density function to get the weights of particles in each cell
609:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
610:   PetscCall(DMSwarmSortGetAccess(sw));
611:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
612:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
613:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
614:     PetscInt          *pidx, Npc;
615:     PetscInt           xNc;
616:     const PetscScalar *xarray;
617:     PetscScalar       *xcoords = NULL;
618:     PetscBool          xisDG;

620:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
621:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
622:     PetscCheck(Npc == (vcEnd - vcStart) * Ns, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Number of particles %" PetscInt_FMT " in cell (rank %d) != %" PetscInt_FMT " number of velocity vertices", Npc, rank, (vcEnd - vcStart) * Ns);
623:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
624:     for (PetscInt q = 0; q < xNq; ++q) {
625:       // Transform quadrature points from ref space to real space
626:       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
627:       // Get probability density at quad point
628:       //   No need to scale xqr since PDF will be periodic
629:       PetscCall((*func)(xqr, param, &xden));
630:       xw += xden * (xwq[q] * xdetJ);
631:     }
632:     xwtot += xw;
633:     if (debug) {
634:       IS              globalOrdering;
635:       const PetscInt *ordering;

637:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
638:       PetscCall(ISGetIndices(globalOrdering, &ordering));
639:       PetscCall(PetscSynchronizedPrintf(comm, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(xcoords[0]), (double)PetscRealPart(xcoords[0 + dim]), (double)xw));
640:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
641:     }
642:     // Set weights to be Gaussian in velocity cells
643:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
644:       const PetscInt     p  = pidx[vc * Ns + 0];
645:       PetscReal          vw = 0.;
646:       PetscInt           vNc;
647:       const PetscScalar *varray;
648:       PetscScalar       *vcoords = NULL;
649:       PetscBool          visDG;

651:       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
652:       PetscCheck(vNc > 0 && vNc % dim == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Velocity cell %" PetscInt_FMT " has invalid coordinates (vNc=%" PetscInt_FMT ", dim=%" PetscInt_FMT ")", vc, vNc, dim);
653:       {
654:         PetscReal vmin[3] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
655:         PetscReal vmax[3] = {-PETSC_MAX_REAL, -PETSC_MAX_REAL, -PETSC_MAX_REAL};
656:         PetscInt  numVert = vNc / dim;
657:         for (PetscInt i = 0; i < numVert; ++i) {
658:           for (PetscInt d = 0; d < dim; ++d) {
659:             PetscReal coord = PetscRealPart(vcoords[i * dim + d]);
660:             vmin[d]         = PetscMin(vmin[d], coord);
661:             vmax[d]         = PetscMax(vmax[d], coord);
662:           }
663:         }
664:         vw = 1.0;
665:         for (PetscInt d = 0; d < dim; ++d) vw *= 0.5 * (PetscErfReal(vmax[d] / PetscSqrtReal(2.)) - PetscErfReal(vmin[d] / PetscSqrtReal(2.)));
666:         PetscCheck(PetscIsNormalReal(vw), PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " velocity weight is not normal: vw=%g, vmin=(%g,%g,%g), vmax=(%g,%g,%g)", p, vw, vmin[0], vmin[1], vmin[2], vmax[0], vmax[1], vmax[2]);
667:       }

669:       weight[p] = totalWeight * vw * xw;
670:       pwtot += weight[p];
671:       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 10: weight=%g, xw=%g, vw=%g, totalWeight=%g", p, weight[p], xw, vw, totalWeight);
672:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
673:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
674:     }
675:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
676:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
677:   }
678:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
679:   PetscCall(DMSwarmSortRestoreAccess(sw));
680:   PetscCall(PetscQuadratureDestroy(&xquad));

682:   if (debug) {
683:     PetscReal wtot[2] = {pwtot, xwtot}, gwtot[2];

685:     PetscCall(PetscSynchronizedFlush(comm, NULL));
686:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
687:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
688:   }
689:   PetscFunctionReturn(PETSC_SUCCESS);
690: }

692: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
693: {
694:   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
695:   PetscInt  dim;

697:   PetscFunctionBegin;
698:   PetscCall(DMGetDimension(sw, &dim));
699:   PetscCall(InitializeParticles_Centroid(sw));
700:   PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : (dim == 2 ? PetscPDFCosine2D : PetscPDFCosine3D), scale));
701:   PetscFunctionReturn(PETSC_SUCCESS);
702: }

704: static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
705: {
706:   DM         dm;
707:   PetscInt  *species;
708:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
709:   PetscInt   Np, dim;

711:   PetscFunctionBegin;
712:   PetscCall(DMSwarmGetCellDM(sw, &dm));
713:   PetscCall(DMGetDimension(sw, &dim));
714:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
715:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
716:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
717:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
718:   for (PetscInt p = 0; p < Np; ++p) {
719:     totalWeight += weight[p];
720:     totalCharge += user->charges[species[p]] * weight[p];
721:   }
722:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
723:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
724:   {
725:     Parameter *param;
726:     PetscReal  Area;

728:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
729:     switch (dim) {
730:     case 1:
731:       Area = (gmax[0] - gmin[0]);
732:       break;
733:     case 2:
734:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
735:       break;
736:     case 3:
737:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
738:       break;
739:     default:
740:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
741:     }
742:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
743:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
744:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)user->charges[0], (double)global_charge, (double)Area));
745:     param->sigma = PetscAbsReal(global_charge / (Area));

747:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
748:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "(x0,v0,t0,m0,q0,phi0): (%e, %e, %e, %e, %e, %e) - (P, V) = (%e, %e)\n", (double)param->x0, (double)param->v0, (double)param->t0, (double)param->m0, (double)param->q0, (double)param->phi0, (double)param->poissonNumber,
749:                           (double)param->vlasovNumber));
750:   }
751:   /* Setup Constants */
752:   {
753:     PetscDS    ds;
754:     Parameter *param;
755:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
756:     PetscScalar constants[NUM_CONSTANTS];
757:     constants[SIGMA]   = param->sigma;
758:     constants[V0]      = param->v0;
759:     constants[T0]      = param->t0;
760:     constants[X0]      = param->x0;
761:     constants[M0]      = param->m0;
762:     constants[Q0]      = param->q0;
763:     constants[PHI0]    = param->phi0;
764:     constants[POISSON] = param->poissonNumber;
765:     constants[VLASOV]  = param->vlasovNumber;
766:     PetscCall(DMGetDS(dm, &ds));
767:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
768:   }
769:   PetscFunctionReturn(PETSC_SUCCESS);
770: }

772: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
773: {
774:   PetscBag   bag;
775:   Parameter *p;

777:   PetscFunctionBeginUser;
778:   /* setup PETSc parameter bag */
779:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
780:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
781:   bag = ctx->bag;
782:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
783:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
784:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
785:   PetscCall(PetscBagRegisterScalar(bag, &p->phi0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
786:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
787:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
788:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
789:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

791:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
792:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
793:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
794:   PetscCall(PetscBagSetFromOptions(bag));
795:   {
796:     PetscViewer       viewer;
797:     PetscViewerFormat format;
798:     PetscBool         flg;

800:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
801:     if (flg) {
802:       PetscCall(PetscViewerPushFormat(viewer, format));
803:       PetscCall(PetscBagView(bag, viewer));
804:       PetscCall(PetscViewerFlush(viewer));
805:       PetscCall(PetscViewerPopFormat(viewer));
806:       PetscCall(PetscViewerDestroy(&viewer));
807:     }
808:   }
809:   PetscFunctionReturn(PETSC_SUCCESS);
810: }

812: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
813: {
814:   const char *prefix = "x";

816:   PetscFunctionBeginUser;
817:   PetscCall(DMCreate(comm, dm));
818:   PetscCall(DMSetType(*dm, DMPLEX));
819:   PetscCall(DMPlexSetOptionsPrefix(*dm, prefix));
820:   PetscCall(DMSetFromOptions(*dm));
821:   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
822:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

824:   // Cache the mesh geometry
825:   DMField         coordField;
826:   IS              cellIS;
827:   PetscQuadrature quad;
828:   PetscReal      *wt, *pt;
829:   PetscInt        cdim, cStart, cEnd;

831:   PetscCall(DMGetCoordinateField(*dm, &coordField));
832:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
833:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
834:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
835:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
836:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
837:   PetscCall(PetscMalloc1(1, &wt));
838:   PetscCall(PetscMalloc1(cdim, &pt));
839:   wt[0] = 1.;
840:   for (PetscInt d = 0; d < cdim; ++d) pt[d] = -1.;
841:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
842:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom));
843:   PetscCall(PetscQuadratureDestroy(&quad));
844:   PetscCall(ISDestroy(&cellIS));
845:   PetscFunctionReturn(PETSC_SUCCESS);
846: }

848: static void ion_f0(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 f0[])
849: {
850:   f0[0] = -constants[SIGMA];
851: }

853: 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[])
854: {
855:   PetscInt d;
856:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
857: }

859: static void laplacian_g3(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[])
860: {
861:   PetscInt d;
862:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
863: }

865: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
866: {
867:   PetscFE      fephi;
868:   PetscDS      ds;
869:   PetscBool    simplex;
870:   PetscInt     dim;
871:   MatNullSpace nullsp;

873:   PetscFunctionBeginUser;
874:   PetscCall(DMGetDimension(dm, &dim));
875:   PetscCall(DMPlexIsSimplex(dm, &simplex));

877:   PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
878:   PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
879:   PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
880:   PetscCall(DMCreateDS(dm));
881:   PetscCall(DMGetDS(dm, &ds));
882:   PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
883:   PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
884:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
885:   PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
886:   PetscCall(MatNullSpaceDestroy(&nullsp));
887:   PetscCall(PetscFEDestroy(&fephi));
888:   PetscFunctionReturn(PETSC_SUCCESS);
889: }

891: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
892: {
893:   SNES         snes;
894:   Mat          J;
895:   MatNullSpace nullSpace;

897:   PetscFunctionBeginUser;
898:   PetscCall(CreateFEM(dm, user));
899:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
900:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
901:   PetscCall(SNESSetDM(snes, dm));
902:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
903:   PetscCall(SNESSetFromOptions(snes));

905:   PetscCall(DMCreateMatrix(dm, &J));
906:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
907:   PetscCall(MatSetNullSpace(J, nullSpace));
908:   PetscCall(MatNullSpaceDestroy(&nullSpace));
909:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
910:   PetscCall(MatDestroy(&J));

912:   user->dmPot = dm;
913:   PetscCall(PetscObjectReference((PetscObject)user->dmPot));

915:   PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M));
916:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
917:   user->snes = snes;
918:   PetscFunctionReturn(PETSC_SUCCESS);
919: }

921: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
922: {
923:   DMSwarmCellDM celldm;
924:   PetscInt      dim;

926:   PetscFunctionBeginUser;
927:   PetscCall(DMGetDimension(dm, &dim));
928:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
929:   PetscCall(DMSetType(*sw, DMSWARM));
930:   PetscCall(DMSetDimension(*sw, dim));
931:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
932:   PetscCall(DMSetApplicationContext(*sw, user));

934:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
935:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
936:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
937:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));

939:   const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};

941:   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
942:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
943:   PetscCall(DMSwarmCellDMDestroy(&celldm));

945:   const char *vfieldnames[2] = {"w_q"};
946:   DM          vdm;

948:   PetscCall(CreateVelocityDM(*sw, &vdm));
949:   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
950:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
951:   PetscCall(DMSwarmCellDMDestroy(&celldm));
952:   PetscCall(DMDestroy(&vdm));

954:   DM mdm;

956:   PetscCall(DMClone(dm, &mdm));
957:   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
958:   PetscCall(DMCopyDisc(dm, mdm));
959:   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
960:   PetscCall(DMDestroy(&mdm));
961:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
962:   PetscCall(DMSwarmCellDMDestroy(&celldm));

964:   PetscCall(DMSetFromOptions(*sw));
965:   PetscCall(DMSetUp(*sw));

967:   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
968:   user->swarm = *sw;
969:   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
970:   PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
971:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
972:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
973:   PetscFunctionReturn(PETSC_SUCCESS);
974: }

976: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
977: {
978:   DM         dm;
979:   AppCtx    *user;
980:   PetscDS    ds;
981:   PetscFE    fe;
982:   KSP        ksp;
983:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
984:   Vec        rho;         // Charge density, M^{-1} rhoRhs
985:   Vec        phi, locPhi; // Potential
986:   Vec        f;           // Particle weights
987:   PetscReal *coords;
988:   PetscInt   dim, cStart, cEnd, Np;

990:   PetscFunctionBegin;
991:   PetscCall(DMGetApplicationContext(sw, (void *)&user));
992:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
993:   PetscCall(DMGetDimension(sw, &dim));
994:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

996:   PetscCall(SNESGetDM(snes, &dm));
997:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
998:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
999:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
1000:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1001:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

1003:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1004:   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1005:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

1007:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1008:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1010:   // Low-pass filter rhoRhs
1011:   PetscInt window = 0;
1012:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-rho_average", &window, NULL));
1013:   if (window) {
1014:     PetscScalar *a;
1015:     PetscInt     n;
1016:     PetscReal    width = 2. * window + 1.;

1018:     // This will only work for P_1
1019:     //   This works because integration against a test function is linear
1020:     //   Do a real integral against weight function for higher order
1021:     PetscCall(VecGetLocalSize(rhoRhs, &n));
1022:     PetscCall(VecGetArrayWrite(rhoRhs, &a));
1023:     for (PetscInt i = 0; i < n; ++i) {
1024:       PetscScalar avg = a[i];
1025:       for (PetscInt j = 1; j <= window; ++j) avg += a[(i - j + n) % n] + a[(i + j) % n];
1026:       a[i] = avg / width;
1027:       //a[i] = (a[(i - 1 + n) % n] + a[i] + a[(i + 1) % n]) / 3.;
1028:     }
1029:     PetscCall(VecRestoreArrayWrite(rhoRhs, &a));
1030:   }

1032:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1033:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1034:   PetscCall(KSPSetOperators(ksp, user->M, user->M));
1035:   PetscCall(KSPSetFromOptions(ksp));
1036:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

1038:   PetscCall(VecScale(rhoRhs, -1.0));

1040:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1041:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
1042:   PetscCall(KSPDestroy(&ksp));

1044:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1045:   PetscCall(VecSet(phi, 0.0));
1046:   PetscCall(SNESSolve(snes, rhoRhs, phi));
1047:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
1048:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1050:   PetscCall(DMGetLocalVector(dm, &locPhi));
1051:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1052:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1053:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
1054:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

1056:   PetscCall(DMGetDS(dm, &ds));
1057:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1058:   PetscCall(DMSwarmSortGetAccess(sw));
1059:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1060:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1062:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1063:   PetscTabulation tab;
1064:   PetscReal      *pcoord, *refcoord;
1065:   PetscFEGeom    *chunkgeom = NULL;
1066:   PetscInt        maxNcp    = 0;

1068:   for (PetscInt c = cStart; c < cEnd; ++c) {
1069:     PetscInt Ncp;

1071:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
1072:     maxNcp = PetscMax(maxNcp, Ncp);
1073:   }
1074:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1075:   PetscCall(PetscArrayzero(refcoord, maxNcp * dim));
1076:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1077:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
1078:   for (PetscInt c = cStart; c < cEnd; ++c) {
1079:     PetscScalar *clPhi = NULL;
1080:     PetscInt    *points;
1081:     PetscInt     Ncp;

1083:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1084:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1085:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1086:     {
1087:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1088:       for (PetscInt i = 0; i < Ncp; ++i) {
1089:         const PetscReal x0[3] = {-1., -1., -1.};
1090:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1091:       }
1092:     }
1093:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
1094:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1095:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1096:       const PetscReal *basisDer = tab->T[1];
1097:       const PetscInt   p        = points[cp];

1099:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1100:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
1101:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
1102:     }
1103:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1104:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1105:   }
1106:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1107:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1108:   PetscCall(PetscTabulationDestroy(&tab));
1109:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1110:   PetscCall(DMSwarmSortRestoreAccess(sw));
1111:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1112:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1113:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1114:   PetscFunctionReturn(PETSC_SUCCESS);
1115: }

1117: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
1118: {
1119:   AppCtx    *user;
1120:   Mat        M_p;
1121:   PetscReal *E;
1122:   PetscInt   dim, Np;

1124:   PetscFunctionBegin;
1127:   PetscCall(DMGetDimension(sw, &dim));
1128:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1129:   PetscCall(DMGetApplicationContext(sw, &user));

1131:   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
1132:   // TODO: Could share sort context with space cellDM
1133:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
1134:   PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p));
1135:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));

1137:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1138:   PetscCall(PetscArrayzero(E, Np * dim));
1139:   user->validE = PETSC_TRUE;

1141:   PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
1142:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1143:   PetscCall(MatDestroy(&M_p));
1144:   PetscFunctionReturn(PETSC_SUCCESS);
1145: }

1147: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1148: {
1149:   DM                 sw;
1150:   SNES               snes = ((AppCtx *)ctx)->snes;
1151:   const PetscScalar *u;
1152:   PetscScalar       *g;
1153:   PetscReal         *E, m_p = 1., q_p = -1.;
1154:   PetscInt           dim, d, Np, p;

1156:   PetscFunctionBeginUser;
1157:   PetscCall(TSGetDM(ts, &sw));
1158:   PetscCall(ComputeFieldAtParticles(snes, sw));

1160:   PetscCall(DMGetDimension(sw, &dim));
1161:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1162:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1163:   PetscCall(VecGetArrayRead(U, &u));
1164:   PetscCall(VecGetArray(G, &g));
1165:   Np /= 2 * dim;
1166:   for (p = 0; p < Np; ++p) {
1167:     for (d = 0; d < dim; ++d) {
1168:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1169:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1170:     }
1171:   }
1172:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1173:   PetscCall(VecRestoreArrayRead(U, &u));
1174:   PetscCall(VecRestoreArray(G, &g));
1175:   PetscFunctionReturn(PETSC_SUCCESS);
1176: }

1178: /* J_{ij} = dF_i/dx_j
1179:    J_p = (  0   1)
1180:          (-w^2  0)
1181:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1182:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1183: */
1184: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1185: {
1186:   DM               sw;
1187:   const PetscReal *coords, *vel;
1188:   PetscInt         dim, d, Np, p, rStart;

1190:   PetscFunctionBeginUser;
1191:   PetscCall(TSGetDM(ts, &sw));
1192:   PetscCall(DMGetDimension(sw, &dim));
1193:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1194:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1195:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1196:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1197:   Np /= 2 * dim;
1198:   for (p = 0; p < Np; ++p) {
1199:     // TODO This is not right because dv/dx has the electric field in it
1200:     PetscScalar vals[4] = {0., 1., -1, 0.};

1202:     for (d = 0; d < dim; ++d) {
1203:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1204:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1205:     }
1206:   }
1207:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1208:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1209:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1210:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1211:   PetscFunctionReturn(PETSC_SUCCESS);
1212: }

1214: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1215: {
1216:   AppCtx            *user = (AppCtx *)ctx;
1217:   DM                 sw;
1218:   const PetscScalar *v;
1219:   PetscScalar       *xres;
1220:   PetscInt           Np, p, d, dim;

1222:   PetscFunctionBeginUser;
1223:   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
1224:   PetscCall(TSGetDM(ts, &sw));
1225:   PetscCall(DMGetDimension(sw, &dim));
1226:   PetscCall(VecGetLocalSize(Xres, &Np));
1227:   PetscCall(VecGetArrayRead(V, &v));
1228:   PetscCall(VecGetArray(Xres, &xres));
1229:   Np /= dim;
1230:   for (p = 0; p < Np; ++p) {
1231:     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
1232:   }
1233:   PetscCall(VecRestoreArrayRead(V, &v));
1234:   PetscCall(VecRestoreArray(Xres, &xres));
1235:   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
1236:   PetscFunctionReturn(PETSC_SUCCESS);
1237: }

1239: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1240: {
1241:   DM                 sw;
1242:   AppCtx            *user = (AppCtx *)ctx;
1243:   SNES               snes = ((AppCtx *)ctx)->snes;
1244:   const PetscScalar *x;
1245:   PetscScalar       *vres;
1246:   PetscReal         *E, m_p, q_p;
1247:   PetscInt           Np, p, dim, d;
1248:   Parameter         *param;

1250:   PetscFunctionBeginUser;
1251:   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
1252:   PetscCall(TSGetDM(ts, &sw));
1253:   PetscCall(ComputeFieldAtParticles(snes, sw));

1255:   PetscCall(DMGetDimension(sw, &dim));
1256:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1257:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1258:   m_p = user->masses[0] * param->m0;
1259:   q_p = user->charges[0] * param->q0;
1260:   PetscCall(VecGetLocalSize(Vres, &Np));
1261:   PetscCall(VecGetArrayRead(X, &x));
1262:   PetscCall(VecGetArray(Vres, &vres));
1263:   Np /= dim;
1264:   for (p = 0; p < Np; ++p) {
1265:     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1266:   }
1267:   PetscCall(VecRestoreArrayRead(X, &x));
1268:   /*
1269:     Synchronized, ordered output for parallel/sequential test cases.
1270:     In the 1D (on the 2D mesh) case, every y component should be zero.
1271:   */
1272:   if (user->checkVRes) {
1273:     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1274:     PetscInt  step;

1276:     PetscCall(TSGetStepNumber(ts, &step));
1277:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
1278:     for (PetscInt p = 0; p < Np; ++p) {
1279:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
1280:       PetscCheck(PetscAbsScalar(vres[p * dim + 1]) < PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Y velocity should be 0., not %g", (double)PetscRealPart(vres[p * dim + 1]));
1281:     }
1282:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1283:   }
1284:   PetscCall(VecRestoreArray(Vres, &vres));
1285:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1286:   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
1287:   PetscFunctionReturn(PETSC_SUCCESS);
1288: }

1290: /* Discrete Gradients Formulation: S, F, gradF (G) */
1291: PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
1292: {
1293:   PetscScalar vals[4] = {0., 1., -1., 0.};
1294:   DM          sw;
1295:   PetscInt    dim, d, Np, p, rStart;

1297:   PetscFunctionBeginUser;
1298:   PetscCall(TSGetDM(ts, &sw));
1299:   PetscCall(DMGetDimension(sw, &dim));
1300:   PetscCall(VecGetLocalSize(U, &Np));
1301:   PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
1302:   Np /= 2 * dim;
1303:   for (p = 0; p < Np; ++p) {
1304:     for (d = 0; d < dim; ++d) {
1305:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1306:       PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
1307:     }
1308:   }
1309:   PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
1310:   PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
1311:   PetscFunctionReturn(PETSC_SUCCESS);
1312: }

1314: PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
1315: {
1316:   AppCtx            *user = (AppCtx *)ctx;
1317:   DM                 sw;
1318:   Vec                phi;
1319:   const PetscScalar *u;
1320:   PetscInt           dim, Np, cStart, cEnd;
1321:   PetscReal         *vel, *coords, m_p = 1.;

1323:   PetscFunctionBeginUser;
1324:   PetscCall(TSGetDM(ts, &sw));
1325:   PetscCall(DMGetDimension(sw, &dim));
1326:   PetscCall(DMPlexGetHeightStratum(user->dmPot, 0, &cStart, &cEnd));

1328:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1329:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
1330:   PetscCall(computeFieldEnergy(user->dmPot, phi, F));
1331:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

1333:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1334:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1335:   PetscCall(DMSwarmSortGetAccess(sw));
1336:   PetscCall(VecGetArrayRead(U, &u));
1337:   PetscCall(VecGetLocalSize(U, &Np));
1338:   Np /= 2 * dim;
1339:   for (PetscInt c = cStart; c < cEnd; ++c) {
1340:     PetscInt *points;
1341:     PetscInt  Ncp;

1343:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1344:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1345:       const PetscInt  p  = points[cp];
1346:       const PetscReal v2 = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);

1348:       *F += 0.5 * m_p * v2;
1349:     }
1350:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1351:   }
1352:   PetscCall(VecRestoreArrayRead(U, &u));
1353:   PetscCall(DMSwarmSortRestoreAccess(sw));
1354:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1355:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1356:   PetscFunctionReturn(PETSC_SUCCESS);
1357: }

1359: /* dF/dx = q E   dF/dv = v */
1360: PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1361: {
1362:   DM                 sw;
1363:   SNES               snes = ((AppCtx *)ctx)->snes;
1364:   const PetscReal   *coords, *vel, *E;
1365:   const PetscScalar *u;
1366:   PetscScalar       *g;
1367:   PetscReal          m_p = 1., q_p = -1.;
1368:   PetscInt           dim, d, Np, p;

1370:   PetscFunctionBeginUser;
1371:   PetscCall(TSGetDM(ts, &sw));
1372:   PetscCall(DMGetDimension(sw, &dim));
1373:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1374:   PetscCall(VecGetArrayRead(U, &u));
1375:   PetscCall(VecGetArray(G, &g));

1377:   PetscLogEvent COMPUTEFIELD;
1378:   PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
1379:   PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
1380:   PetscCall(ComputeFieldAtParticles(snes, sw));
1381:   PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
1382:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1383:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1384:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1385:   for (p = 0; p < Np; ++p) {
1386:     for (d = 0; d < dim; ++d) {
1387:       g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
1388:       g[(p * 2 + 1) * dim + d] = u[(p * 2 + 1) * dim + d];
1389:     }
1390:   }
1391:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1392:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1393:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1394:   PetscCall(VecRestoreArrayRead(U, &u));
1395:   PetscCall(VecRestoreArray(G, &g));
1396:   PetscFunctionReturn(PETSC_SUCCESS);
1397: }

1399: static PetscErrorCode CreateSolution(TS ts)
1400: {
1401:   DM       sw;
1402:   Vec      u;
1403:   PetscInt dim, Np;

1405:   PetscFunctionBegin;
1406:   PetscCall(TSGetDM(ts, &sw));
1407:   PetscCall(DMGetDimension(sw, &dim));
1408:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1409:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1410:   PetscCall(VecSetBlockSize(u, dim));
1411:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1412:   PetscCall(VecSetUp(u));
1413:   PetscCall(TSSetSolution(ts, u));
1414:   PetscCall(VecDestroy(&u));
1415:   PetscFunctionReturn(PETSC_SUCCESS);
1416: }

1418: static PetscErrorCode SetProblem(TS ts)
1419: {
1420:   AppCtx *user;
1421:   DM      sw;

1423:   PetscFunctionBegin;
1424:   PetscCall(TSGetDM(ts, &sw));
1425:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1426:   // Define unified system for (X, V)
1427:   {
1428:     Mat      J;
1429:     PetscInt dim, Np;

1431:     PetscCall(DMGetDimension(sw, &dim));
1432:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1433:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1434:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1435:     PetscCall(MatSetBlockSize(J, 2 * dim));
1436:     PetscCall(MatSetFromOptions(J));
1437:     PetscCall(MatSetUp(J));
1438:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1439:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1440:     PetscCall(MatDestroy(&J));
1441:   }
1442:   /* Define split system for X and V */
1443:   {
1444:     Vec             u;
1445:     IS              isx, isv, istmp;
1446:     const PetscInt *idx;
1447:     PetscInt        dim, Np, rstart;

1449:     PetscCall(TSGetSolution(ts, &u));
1450:     PetscCall(DMGetDimension(sw, &dim));
1451:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1452:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1453:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1454:     PetscCall(ISGetIndices(istmp, &idx));
1455:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1456:     PetscCall(ISRestoreIndices(istmp, &idx));
1457:     PetscCall(ISDestroy(&istmp));
1458:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1459:     PetscCall(ISGetIndices(istmp, &idx));
1460:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1461:     PetscCall(ISRestoreIndices(istmp, &idx));
1462:     PetscCall(ISDestroy(&istmp));
1463:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1464:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1465:     PetscCall(ISDestroy(&isx));
1466:     PetscCall(ISDestroy(&isv));
1467:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1468:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1469:   }
1470:   // Define symplectic formulation U_t = S . G, where G = grad F
1471:   {
1472:     PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, user));
1473:   }
1474:   PetscFunctionReturn(PETSC_SUCCESS);
1475: }

1477: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1478: {
1479:   DM        sw;
1480:   Vec       u;
1481:   PetscReal t, maxt, dt;
1482:   PetscInt  n, maxn;

1484:   PetscFunctionBegin;
1485:   PetscCall(TSGetDM(ts, &sw));
1486:   PetscCall(TSGetTime(ts, &t));
1487:   PetscCall(TSGetMaxTime(ts, &maxt));
1488:   PetscCall(TSGetTimeStep(ts, &dt));
1489:   PetscCall(TSGetStepNumber(ts, &n));
1490:   PetscCall(TSGetMaxSteps(ts, &maxn));

1492:   PetscCall(TSReset(ts));
1493:   PetscCall(TSSetDM(ts, sw));
1494:   PetscCall(TSSetFromOptions(ts));
1495:   PetscCall(TSSetTime(ts, t));
1496:   PetscCall(TSSetMaxTime(ts, maxt));
1497:   PetscCall(TSSetTimeStep(ts, dt));
1498:   PetscCall(TSSetStepNumber(ts, n));
1499:   PetscCall(TSSetMaxSteps(ts, maxn));

1501:   PetscCall(CreateSolution(ts));
1502:   PetscCall(SetProblem(ts));
1503:   PetscCall(TSGetSolution(ts, &u));
1504:   PetscFunctionReturn(PETSC_SUCCESS);
1505: }

1507: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
1508: {
1509:   DM        sw, cdm;
1510:   PetscInt  Np;
1511:   PetscReal low[2], high[2];
1512:   AppCtx   *user = (AppCtx *)ctx;

1514:   sw = user->swarm;
1515:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
1516:   // Get the bounding box so we can equally space the particles
1517:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
1518:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1519:   // shift it by h/2 so nothing is initialized directly on a boundary
1520:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
1521:   x[1] = 0.;
1522:   return PETSC_SUCCESS;
1523: }

1525: /*
1526:   InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.

1528:   Input Parameters:
1529: + ts         - The TS
1530: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

1532:   Output Parameters:
1533: . u - The initialized solution vector

1535:   Level: advanced

1537: .seealso: InitializeSolve()
1538: */
1539: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1540: {
1541:   DM       sw;
1542:   Vec      u, gc, gv;
1543:   IS       isx, isv;
1544:   PetscInt dim;
1545:   AppCtx  *user;

1547:   PetscFunctionBeginUser;
1548:   PetscCall(TSGetDM(ts, &sw));
1549:   PetscCall(DMGetApplicationContext(sw, &user));
1550:   PetscCall(DMGetDimension(sw, &dim));
1551:   if (useInitial) {
1552:     PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1553:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1554:     PetscCall(DMSwarmTSRedistribute(ts));
1555:   }
1556:   PetscCall(DMSetUp(sw));
1557:   PetscCall(TSGetSolution(ts, &u));
1558:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1559:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1560:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1561:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1562:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1563:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1564:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1565:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1566:   PetscFunctionReturn(PETSC_SUCCESS);
1567: }

1569: static PetscErrorCode InitializeSolve(TS ts, Vec u)
1570: {
1571:   PetscFunctionBegin;
1572:   PetscCall(TSSetSolution(ts, u));
1573:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1574:   PetscFunctionReturn(PETSC_SUCCESS);
1575: }

1577: static PetscErrorCode MigrateParticles(TS ts)
1578: {
1579:   DM               sw, cdm;
1580:   const PetscReal *L;
1581:   AppCtx          *ctx;

1583:   PetscFunctionBeginUser;
1584:   PetscCall(TSGetDM(ts, &sw));
1585:   PetscCall(DMGetApplicationContext(sw, &ctx));
1586:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1587:   {
1588:     Vec        u, gc, gv, position, momentum;
1589:     IS         isx, isv;
1590:     PetscReal *pos, *mom;

1592:     PetscCall(TSGetSolution(ts, &u));
1593:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1594:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1595:     PetscCall(VecGetSubVector(u, isx, &position));
1596:     PetscCall(VecGetSubVector(u, isv, &momentum));
1597:     PetscCall(VecGetArray(position, &pos));
1598:     PetscCall(VecGetArray(momentum, &mom));
1599:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1600:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1601:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1602:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

1604:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
1605:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
1606:     if ((L[0] || L[1]) >= 0.) {
1607:       PetscReal *x, *v, upper[3], lower[3];
1608:       PetscInt   Np, dim;

1610:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
1611:       PetscCall(DMGetDimension(cdm, &dim));
1612:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
1613:       PetscCall(VecGetArray(gc, &x));
1614:       PetscCall(VecGetArray(gv, &v));
1615:       for (PetscInt p = 0; p < Np; ++p) {
1616:         for (PetscInt d = 0; d < dim; ++d) {
1617:           if (pos[p * dim + d] < lower[d]) {
1618:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
1619:           } else if (pos[p * dim + d] > upper[d]) {
1620:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
1621:           } else {
1622:             x[p * dim + d] = pos[p * dim + d];
1623:           }
1624:           PetscCheck(x[p * dim + d] >= lower[d] && x[p * dim + d] <= upper[d], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "p: %" PetscInt_FMT "x[%" PetscInt_FMT "] %g", p, d, (double)x[p * dim + d]);
1625:           v[p * dim + d] = mom[p * dim + d];
1626:         }
1627:       }
1628:       PetscCall(VecRestoreArray(gc, &x));
1629:       PetscCall(VecRestoreArray(gv, &v));
1630:     }
1631:     PetscCall(VecRestoreArray(position, &pos));
1632:     PetscCall(VecRestoreArray(momentum, &mom));
1633:     PetscCall(VecRestoreSubVector(u, isx, &position));
1634:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
1635:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1636:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1637:   }
1638:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1639:   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
1640:   PetscCall(DMSwarmTSRedistribute(ts));
1641:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
1642:   {
1643:     const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
1644:     PetscCall(DMSwarmVectorDefineFields(sw, 2, fieldnames));
1645:   }
1646:   PetscFunctionReturn(PETSC_SUCCESS);
1647: }

1649: int main(int argc, char **argv)
1650: {
1651:   DM        dm, sw;
1652:   TS        ts;
1653:   Vec       u;
1654:   PetscReal dt;
1655:   PetscInt  maxn;
1656:   AppCtx    user;

1658:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
1659:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
1660:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
1661:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
1662:   PetscCall(CreatePoisson(dm, &user));
1663:   PetscCall(CreateSwarm(dm, &user, &sw));
1664:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
1665:   PetscCall(InitializeConstants(sw, &user));
1666:   PetscCall(DMSetApplicationContext(sw, &user));

1668:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
1669:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
1670:   PetscCall(TSSetDM(ts, sw));
1671:   PetscCall(TSSetMaxTime(ts, 0.1));
1672:   PetscCall(TSSetTimeStep(ts, 0.00001));
1673:   PetscCall(TSSetMaxSteps(ts, 100));
1674:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

1676:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
1677:   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
1678:   if (user.particle_monitor) PetscCall(TSMonitorSet(ts, MonitorParticles, &user, NULL));

1680:   PetscCall(TSSetFromOptions(ts));
1681:   PetscCall(TSGetTimeStep(ts, &dt));
1682:   PetscCall(TSGetMaxSteps(ts, &maxn));
1683:   user.steps    = maxn;
1684:   user.stepSize = dt;
1685:   PetscCall(SetupContext(dm, sw, &user));
1686:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
1687:   PetscCall(TSSetPostStep(ts, MigrateParticles));
1688:   PetscCall(CreateSolution(ts));
1689:   PetscCall(TSGetSolution(ts, &u));
1690:   PetscCall(TSComputeInitialCondition(ts, u));
1691:   PetscCall(CheckNonNegativeWeights(sw, &user));
1692:   PetscCall(TSSolve(ts, NULL));

1694:   PetscCall(SNESDestroy(&user.snes));
1695:   PetscCall(DMDestroy(&user.dmPot));
1696:   PetscCall(MatDestroy(&user.M));
1697:   PetscCall(PetscFEGeomDestroy(&user.fegeom));
1698:   PetscCall(TSDestroy(&ts));
1699:   PetscCall(DMDestroy(&sw));
1700:   PetscCall(DMDestroy(&dm));
1701:   PetscCall(DestroyContext(&user));
1702:   PetscCall(PetscFinalize());
1703:   return 0;
1704: }

1706: /*TEST

1708:    build:
1709:     requires: !complex double

1711:    testset:
1712:      nsize: 2
1713:      args: -cosine_coefficients 0.01 -charges -1. -total_weight 1. -xdm_plex_hash_location -vpetscspace_degree 2 -petscspace_degree 1 -em_snes_atol 1.e-12 \
1714:      -em_snes_error_if_not_converged -em_ksp_error_if_not_converged -em_ksp_type cg -em_pc_type gamg -em_mg_coarse_ksp_type preonly -em_mg_coarse_pc_type svd -em_proj_ksp_type cg \
1715:      -em_proj_pc_type gamg -em_proj_mg_coarse_ksp_type preonly -em_proj_mg_coarse_pc_type svd -ts_time_step 0.03 -xdm_plex_simplex 0 \
1716:      -ts_max_time 100 -output_step 1 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -check_weights -ts_max_steps 60

1718:      test:
1719:        suffix: landau_damping_1d
1720:        args: -xdm_plex_dim 1 -xdm_plex_box_faces 60 -xdm_plex_box_lower 0. -xdm_plex_box_upper 12.5664 -xdm_plex_box_bd periodic -vdm_plex_dim 1 -vdm_plex_box_faces 60 \
1721:        -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none

1723:      test:
1724:        suffix: landau_damping_2d
1725:        args: -xdm_plex_dim 2 -xdm_plex_box_bd periodic,periodic -vdm_plex_dim 2 -xdm_plex_box_lower 0.,-.5 -vdm_plex_box_lower -6,-6 -vdm_plex_box_upper 6,6 -xdm_plex_box_faces 6,3 \
1726:        -xdm_plex_box_upper 12.5664,.5 -vdm_plex_box_faces 15,15 -vdm_plex_box_bd none,none -vdm_plex_hash_location -vdm_plex_simplex 0

1728:      test:
1729:        suffix: landau_damping_3d
1730:        args: -xdm_plex_dim 3 -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_plex_dim 3 -vdm_plex_box_faces 4,4,4 \
1731:        -vdm_plex_box_lower -6,-6,-6 -vdm_plex_box_upper 6,6,6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none,none,none

1733:      test:
1734:        requires: !defined(PETSC_USE_DMLANDAU_2D)
1735:        suffix: sphere_3d
1736:        nsize: 1
1737:        args: -use_landau_velocity_space -xdm_plex_dim 3 -vdm_landau_thermal_temps 1 -vdm_landau_device_type cpu -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 \
1738:        -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_landau_verbose 2 -vdm_landau_sphere -vdm_landau_map_sphere \
1739:        -vdm_landau_domain_radius 6,6,6 -vdm_landau_sphere_inner_radius_90degree_scale .35 -vdm_refine 1

1741: TEST*/