Actual source code: ex4.c

  1: static char help[] = "Two-level system for Landau Damping using Vlasov-Poisson equations\n";

  3: /*
  4:   Moment Equations:

  6:     Discretize the moment equations using finite elements and project the moments into the finite element space. Use DMSWARM_REMAP_PFAK, which guarantees that the FE approximation is weakly equivalent to the true moment. The first moment, number density, is given by

  8:       \int dx \phi_i n_f = \int dx \phi_i n_p
  9:       \int dx \phi_i n_f = \int dx \phi_i \int dv f
 10:       \int dx \phi_i n_f = \int dx \phi_i \int dv \sum_p w_p \delta(x - x_p) \delta(v - v_p)
 11:       \int dx \phi_i n_f = \int dx \phi_i \sum_p w_p \delta(x - x_p)
 12:                    M n_F = M_p w_p

 14:     where

 16:       (M_p){ip} = \phi_i(x_p)

 18:     which is just a scaled version of the charge density. The second moment, momentum density, is given by

 20:       \int dx \phi_i p_f = m \int dx \phi_i \int dv v f
 21:       \int dx \phi_i p_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) v_p
 22:                    M p_F = M_p v_p w_p

 24:     And finally the third moment, pressure, is given by

 26:       \int dx \phi_i pr_f = m \int dx \phi_i \int dv (v - u)^2 f
 27:       \int dx \phi_i pr_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) (v_p - u)^2
 28:                    M pr_F = M_p (v_p - u)^2 w_p
 29:                           = M_p (v_p - p_F(x_p) / m n_F(x_p))^2 w_p
 30:                           = M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p

 32:     Here we need all FEM basis functions \phi_i that see that particle p.

 34:   To run the code with particles sinusoidally perturbed in x space use the test "pp_poisson_bsi_1d_4" or "pp_poisson_bsi_2d_4"
 35:   According to Lukas, good damping results come at ~16k particles per cell

 37:   Swarm CellDMs
 38:   =============
 39:   Name: "space"
 40:   Fields: DMSwarmPICField_coor, "velocity"
 41:   Coordinates: DMSwarmPICField_coor

 43:   Name: "velocity"
 44:   Fields: "w_q"
 45:   Coordinates: "velocity"

 47:   Name: "moments"
 48:   Fields: "w_q"
 49:   Coordinates: DMSwarmPICField_coor

 51:   Name: "moment fields"
 52:   Fields: "velocity"
 53:   Coordinates: DMSwarmPICField_coor

 55:   To visualize the maximum electric field use

 57:     -efield_monitor

 59:   To monitor velocity moments of the distribution use

 61:     -ptof_pc_type lu -moments_monitor

 63:   To monitor the particle positions in phase space use

 65:     -positions_monitor

 67:   To monitor the charge density, E field, and potential use

 69:     -poisson_monitor

 71:   To monitor the remapping field use

 73:     -remap_uf_view draw

 75:   To visualize the swarm distribution use

 77:     -ts_monitor_hg_swarm

 79:   To visualize the particles, we can use

 81:     -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
 82: */
 83: #include <petsctao.h>
 84: #include <petscts.h>
 85: #include <petscdmplex.h>
 86: #include <petscdmswarm.h>
 87: #include <petscfe.h>
 88: #include <petscds.h>
 89: #include <petscbag.h>
 90: #include <petscdraw.h>
 91: #include <petsc/private/petscfeimpl.h>
 92: #include <petsc/private/dmswarmimpl.h>
 93: #include "petscdm.h"
 94: #include "petscdmlabel.h"

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

 99: const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
100: typedef enum {
101:   EM_PRIMAL,
102:   EM_MIXED,
103:   EM_COULOMB,
104:   EM_NONE
105: } EMType;

107: typedef enum {
108:   V0,
109:   X0,
110:   T0,
111:   M0,
112:   Q0,
113:   PHI0,
114:   POISSON,
115:   VLASOV,
116:   SIGMA,
117:   NUM_CONSTANTS
118: } ConstantType;

120: typedef enum {
121:   E_MONITOR_NONE,
122:   E_MONITOR_FULL,
123:   E_MONITOR_QUIET
124: } EMonitorType;
125: const char *const EMonitorTypes[] = {"NONE", "FULL", "QUIET", "EMonitorType", "E_MONITOR_", NULL};

127: typedef struct {
128:   PetscScalar v0; /* Velocity scale, often the thermal velocity */
129:   PetscScalar t0; /* Time scale */
130:   PetscScalar x0; /* Space scale */
131:   PetscScalar m0; /* Mass scale */
132:   PetscScalar q0; /* Charge scale */
133:   PetscScalar kb;
134:   PetscScalar epsi0;
135:   PetscScalar phi0;          /* Potential scale */
136:   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
137:   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
138:   PetscReal   sigma;         /* Nondimensional charge per length in x */
139: } Parameter;

141: typedef struct {
142:   PetscInt         s;    // Starting sample (we ignore some in the beginning)
143:   PetscInt         e;    // Ending sample
144:   PetscInt         per;  // Period of fitting
145:   const PetscReal *t;    // Time for each sample
146:   const PetscReal *Emax; // Emax for each sample
147: } EmaxCtx;

149: typedef struct {
150:   PetscBag     bag;                  // Problem parameters
151:   PetscBool    error;                // Flag for printing the error
152:   PetscInt     remapFreq;            // Number of timesteps between remapping
153:   EMonitorType efield_monitor;       // Flag to show electric field monitor
154:   PetscBool    moment_monitor;       // Flag to show distribution moment monitor
155:   PetscBool    moment_field_monitor; // Flag to show moment field monitor
156:   PetscBool    positions_monitor;    // Flag to show particle positins at each time step
157:   PetscBool    poisson_monitor;      // Flag to display charge, E field, and potential at each solve
158:   PetscBool    initial_monitor;      // Flag to monitor the initial conditions
159:   PetscInt     velocity_monitor;     // Cell to monitor the velocity distribution for
160:   PetscBool    perturbed_weights;    // Uniformly sample x,v space with gaussian weights
161:   PetscInt     ostep;                // Print the energy at each ostep time steps
162:   PetscInt     numParticles;
163:   PetscReal    timeScale;              /* Nondimensionalizing time scale */
164:   PetscReal    charges[2];             /* The charges of each species */
165:   PetscReal    masses[2];              /* The masses of each species */
166:   PetscReal    thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
167:   PetscReal    cosine_coefficients[2]; /*(alpha, k)*/
168:   PetscReal    totalWeight;
169:   PetscReal    stepSize;
170:   PetscInt     steps;
171:   PetscReal    initVel;
172:   EMType       em;           // Type of electrostatic model
173:   SNES         snes;         // EM solver
174:   DM           dmMom;        // The DM for moment fields
175:   DM           dmN;          // The DM for number density fields
176:   IS           isN;          // The IS mapping dmN into dmMom
177:   Mat          MN;           // The finite element mass matrix for number density
178:   DM           dmP;          // The DM for momentum density fields
179:   IS           isP;          // The IS mapping dmP into dmMom
180:   Mat          MP;           // The finite element mass matrix for momentum density
181:   DM           dmE;          // The DM for energy density (pressure) fields
182:   IS           isE;          // The IS mapping dmE into dmMom
183:   Mat          ME;           // The finite element mass matrix for energy density (pressure)
184:   DM           dmPot;        // The DM for potential
185:   Mat          fftPot;       // Fourier Transform operator for the potential
186:   Vec          fftX, fftY;   //   FFT vectors with phases added (complex parts)
187:   IS           fftReal;      //   The indices for real parts
188:   IS           isPot;        // The IS for potential, or NULL in primal
189:   Mat          M;            // The finite element mass matrix for potential
190:   PetscFEGeom *fegeom;       // Geometric information for the DM cells
191:   PetscDrawHG  drawhgic_x;   // Histogram of the particle weight in each X cell
192:   PetscDrawHG  drawhgic_v;   // Histogram of the particle weight in each X cell
193:   PetscDrawHG  drawhgcell_v; // Histogram of the particle weight in a given cell
194:   PetscBool    validE;       // Flag to indicate E-field in swarm is valid
195:   PetscReal    drawlgEmin;   // The minimum lg(E) to plot
196:   PetscDrawLG  drawlgE;      // Logarithm of maximum electric field
197:   PetscDrawSP  drawspE;      // Electric field at particle positions
198:   PetscDrawSP  drawspX;      // Particle positions
199:   PetscViewer  viewerRho;    // Charge density viewer
200:   PetscViewer  viewerRhoHat; // Charge density Fourier Transform viewer
201:   PetscViewer  viewerPhi;    // Potential viewer
202:   PetscViewer  viewerN;      // Number density viewer
203:   PetscViewer  viewerP;      // Momentum density viewer
204:   PetscViewer  viewerE;      // Energy density (pressure) viewer
205:   PetscViewer  viewerNRes;   // Number density residual viewer
206:   PetscViewer  viewerPRes;   // Momentum density residual viewer
207:   PetscViewer  viewerERes;   // Energy density (pressure) residual viewer
208:   PetscDrawLG  drawlgMomRes; // Residuals for the moment equations
209:   DM           swarm;        // The particle swarm
210:   PetscRandom  random;       // Used for particle perturbations
211:   PetscBool    twostream;    // Flag for activating 2-stream setup
212:   PetscBool    checkweights; // Check weight normalization
213:   PetscInt     checkVRes;    // Flag to check/output velocity residuals for nightly tests
214:   PetscBool    checkLandau;  // Check the Landau damping result
215:   EmaxCtx      emaxCtx;      // Information for fit to decay profile
216:   PetscReal    gamma;        // The damping rate for Landau damping
217:   PetscReal    omega;        // The perturbed oscillation frequency for Landau damping

219:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
220: } AppCtx;

222: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
223: {
224:   PetscFunctionBeginUser;
225:   PetscInt d                      = 2;
226:   PetscInt maxSpecies             = 2;
227:   options->error                  = PETSC_FALSE;
228:   options->remapFreq              = 1;
229:   options->efield_monitor         = E_MONITOR_NONE;
230:   options->moment_monitor         = PETSC_FALSE;
231:   options->moment_field_monitor   = PETSC_FALSE;
232:   options->initial_monitor        = PETSC_FALSE;
233:   options->perturbed_weights      = PETSC_FALSE;
234:   options->poisson_monitor        = PETSC_FALSE;
235:   options->positions_monitor      = PETSC_FALSE;
236:   options->velocity_monitor       = -1;
237:   options->ostep                  = 100;
238:   options->timeScale              = 2.0e-14;
239:   options->charges[0]             = -1.0;
240:   options->charges[1]             = 1.0;
241:   options->masses[0]              = 1.0;
242:   options->masses[1]              = 1000.0;
243:   options->thermal_energy[0]      = 1.0;
244:   options->thermal_energy[1]      = 1.0;
245:   options->cosine_coefficients[0] = 0.01;
246:   options->cosine_coefficients[1] = 0.5;
247:   options->initVel                = 1;
248:   options->totalWeight            = 1.0;
249:   options->drawhgic_x             = NULL;
250:   options->drawhgic_v             = NULL;
251:   options->drawhgcell_v           = NULL;
252:   options->validE                 = PETSC_FALSE;
253:   options->drawlgEmin             = -6;
254:   options->drawlgE                = NULL;
255:   options->drawspE                = NULL;
256:   options->drawspX                = NULL;
257:   options->viewerRho              = NULL;
258:   options->viewerRhoHat           = NULL;
259:   options->viewerPhi              = NULL;
260:   options->viewerN                = NULL;
261:   options->viewerP                = NULL;
262:   options->viewerE                = NULL;
263:   options->viewerNRes             = NULL;
264:   options->viewerPRes             = NULL;
265:   options->viewerERes             = NULL;
266:   options->drawlgMomRes           = NULL;
267:   options->em                     = EM_COULOMB;
268:   options->snes                   = NULL;
269:   options->dmMom                  = NULL;
270:   options->dmN                    = NULL;
271:   options->MN                     = NULL;
272:   options->dmP                    = NULL;
273:   options->MP                     = NULL;
274:   options->dmE                    = NULL;
275:   options->ME                     = NULL;
276:   options->dmPot                  = NULL;
277:   options->fftPot                 = NULL;
278:   options->fftX                   = NULL;
279:   options->fftY                   = NULL;
280:   options->fftReal                = NULL;
281:   options->isPot                  = NULL;
282:   options->M                      = NULL;
283:   options->numParticles           = 32768;
284:   options->twostream              = PETSC_FALSE;
285:   options->checkweights           = PETSC_FALSE;
286:   options->checkVRes              = 0;
287:   options->checkLandau            = PETSC_FALSE;
288:   options->emaxCtx.s              = 50;
289:   options->emaxCtx.per            = 100;

291:   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
292:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", __FILE__, options->error, &options->error, NULL));
293:   PetscCall(PetscOptionsInt("-remap_freq", "Number", __FILE__, options->remapFreq, &options->remapFreq, NULL));
294:   PetscCall(PetscOptionsEnum("-efield_monitor", "Flag to record and plot log(max E) over time", __FILE__, EMonitorTypes, (PetscEnum)options->efield_monitor, (PetscEnum *)&options->efield_monitor, NULL));
295:   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", __FILE__, options->drawlgEmin, &options->drawlgEmin, NULL));
296:   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", __FILE__, options->moment_monitor, &options->moment_monitor, NULL));
297:   PetscCall(PetscOptionsBool("-moment_field_monitor", "Flag to show moment fields", __FILE__, options->moment_field_monitor, &options->moment_field_monitor, NULL));
298:   PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", __FILE__, options->initial_monitor, &options->initial_monitor, NULL));
299:   PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", __FILE__, options->positions_monitor, &options->positions_monitor, NULL));
300:   PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", __FILE__, options->poisson_monitor, &options->poisson_monitor, NULL));
301:   PetscCall(PetscOptionsInt("-velocity_monitor", "Cell to show velocity histograms", __FILE__, options->velocity_monitor, &options->velocity_monitor, NULL));
302:   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", __FILE__, options->twostream, &options->twostream, NULL));
303:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", __FILE__, options->perturbed_weights, &options->perturbed_weights, NULL));
304:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", __FILE__, options->checkweights, &options->checkweights, NULL));
305:   PetscCall(PetscOptionsBool("-check_landau", "Check the decay from Landau damping", __FILE__, options->checkLandau, &options->checkLandau, NULL));
306:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", __FILE__, options->ostep, &options->ostep, NULL));
307:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", __FILE__, options->timeScale, &options->timeScale, NULL));
308:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", __FILE__, options->checkVRes, &options->checkVRes, NULL));
309:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", __FILE__, options->initVel, &options->initVel, NULL));
310:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", __FILE__, options->totalWeight, &options->totalWeight, NULL));
311:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", __FILE__, options->cosine_coefficients, &d, NULL));
312:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", __FILE__, options->charges, &maxSpecies, NULL));
313:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", __FILE__, EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
314:   PetscCall(PetscOptionsInt("-emax_start_step", "First time step to use for Emax fits", __FILE__, options->emaxCtx.s, &options->emaxCtx.s, NULL));
315:   PetscCall(PetscOptionsInt("-emax_solve_step", "Number of time steps between Emax fits", __FILE__, options->emaxCtx.per, &options->emaxCtx.per, NULL));
316:   PetscOptionsEnd();

318:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
319:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
320:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
321:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
322:   PetscFunctionReturn(PETSC_SUCCESS);
323: }

325: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *ctx)
326: {
327:   MPI_Comm comm;

329:   PetscFunctionBeginUser;
330:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
331:   if (ctx->efield_monitor) {
332:     PetscDraw     draw;
333:     PetscDrawAxis axis;

335:     if (ctx->efield_monitor == E_MONITOR_FULL) {
336:       PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 0, 400, 300, &draw));
337:       PetscCall(PetscDrawSetSave(draw, "ex2_Efield"));
338:       PetscCall(PetscDrawSetFromOptions(draw));
339:     } else {
340:       PetscCall(PetscDrawOpenNull(comm, &draw));
341:     }
342:     PetscCall(PetscDrawLGCreate(draw, 1, &ctx->drawlgE));
343:     PetscCall(PetscDrawDestroy(&draw));
344:     PetscCall(PetscDrawLGGetAxis(ctx->drawlgE, &axis));
345:     PetscCall(PetscDrawAxisSetLabels(axis, "Max Electric Field", "time", "E_max"));
346:     PetscCall(PetscDrawLGSetLimits(ctx->drawlgE, 0., ctx->steps * ctx->stepSize, ctx->drawlgEmin, 0.));
347:   }

349:   if (ctx->initial_monitor) {
350:     PetscDraw     drawic_x, drawic_v;
351:     PetscDrawAxis axis1, axis2;
352:     PetscReal     dmboxlower[2], dmboxupper[2];
353:     PetscInt      dim, cStart, cEnd;

355:     PetscCall(DMGetDimension(sw, &dim));
356:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
357:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

359:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &drawic_x));
360:     PetscCall(PetscDrawSetSave(drawic_x, "ex2_ic_x"));
361:     PetscCall(PetscDrawSetFromOptions(drawic_x));
362:     PetscCall(PetscDrawHGCreate(drawic_x, (int)dim, &ctx->drawhgic_x));
363:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgic_x, PETSC_TRUE));
364:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgic_x, &axis1));
365:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgic_x, (int)(cEnd - cStart)));
366:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "weight"));
367:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 0));
368:     PetscCall(PetscDrawDestroy(&drawic_x));

370:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &drawic_v));
371:     PetscCall(PetscDrawSetSave(drawic_v, "ex9_ic_v"));
372:     PetscCall(PetscDrawSetFromOptions(drawic_v));
373:     PetscCall(PetscDrawHGCreate(drawic_v, (int)dim, &ctx->drawhgic_v));
374:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgic_v, PETSC_TRUE));
375:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgic_v, &axis2));
376:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgic_v, 21));
377:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "weight"));
378:     PetscCall(PetscDrawAxisSetLimits(axis2, -6, 6, 0, 0));
379:     PetscCall(PetscDrawDestroy(&drawic_v));
380:   }

382:   if (ctx->velocity_monitor >= 0) {
383:     DM            vdm;
384:     DMSwarmCellDM celldm;
385:     PetscDraw     drawcell_v;
386:     PetscDrawAxis axis;
387:     PetscReal     dmboxlower[2], dmboxupper[2];
388:     PetscInt      dim;
389:     char          title[PETSC_MAX_PATH_LEN];

391:     PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
392:     PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
393:     PetscCall(DMGetDimension(vdm, &dim));
394:     PetscCall(DMGetBoundingBox(vdm, dmboxlower, dmboxupper));

396:     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Cell %" PetscInt_FMT ": Velocity Distribution", ctx->velocity_monitor));
397:     PetscCall(PetscDrawCreate(comm, NULL, title, 400, 300, 400, 300, &drawcell_v));
398:     PetscCall(PetscDrawSetSave(drawcell_v, "ex2_cell_v"));
399:     PetscCall(PetscDrawSetFromOptions(drawcell_v));
400:     PetscCall(PetscDrawHGCreate(drawcell_v, (int)dim, &ctx->drawhgcell_v));
401:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgcell_v, PETSC_TRUE));
402:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgcell_v, &axis));
403:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgcell_v, 21));
404:     PetscCall(PetscDrawAxisSetLabels(axis, "V_x Distribution", "V", "weight"));
405:     PetscCall(PetscDrawAxisSetLimits(axis, dmboxlower[0], dmboxupper[0], 0, 0));
406:     PetscCall(PetscDrawDestroy(&drawcell_v));
407:   }

409:   if (ctx->positions_monitor) {
410:     PetscDraw     draw;
411:     PetscDrawAxis axis;

413:     PetscCall(PetscDrawCreate(comm, NULL, "Particle Position", 0, 0, 400, 300, &draw));
414:     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
415:     PetscCall(PetscDrawSetFromOptions(draw));
416:     PetscCall(PetscDrawSPCreate(draw, 10, &ctx->drawspX));
417:     PetscCall(PetscDrawDestroy(&draw));
418:     PetscCall(PetscDrawSPSetDimension(ctx->drawspX, 1));
419:     PetscCall(PetscDrawSPGetAxis(ctx->drawspX, &axis));
420:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
421:     PetscCall(PetscDrawSPReset(ctx->drawspX));
422:   }
423:   if (ctx->poisson_monitor) {
424:     Vec           rho, rhohat, phi;
425:     PetscDraw     draw;
426:     PetscDrawAxis axis;

428:     PetscCall(PetscDrawCreate(comm, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
429:     PetscCall(PetscDrawSetFromOptions(draw));
430:     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
431:     PetscCall(PetscDrawSPCreate(draw, 10, &ctx->drawspE));
432:     PetscCall(PetscDrawDestroy(&draw));
433:     PetscCall(PetscDrawSPSetDimension(ctx->drawspE, 1));
434:     PetscCall(PetscDrawSPGetAxis(ctx->drawspE, &axis));
435:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
436:     PetscCall(PetscDrawSPReset(ctx->drawspE));

438:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Charge Density", 0, 0, 400, 300, &ctx->viewerRho));
439:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerRho, "rho_"));
440:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerRho, 0, &draw));
441:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
442:     PetscCall(PetscViewerSetFromOptions(ctx->viewerRho));
443:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
444:     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
445:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));

447:     PetscInt dim, N;

449:     PetscCall(DMGetDimension(ctx->dmPot, &dim));
450:     if (dim == 1) {
451:       PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
452:       PetscCall(VecGetSize(rhohat, &N));
453:       PetscCall(MatCreateFFT(comm, dim, &N, MATFFTW, &ctx->fftPot));
454:       PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
455:       PetscCall(MatCreateVecs(ctx->fftPot, &ctx->fftX, &ctx->fftY));
456:       PetscCall(ISCreateStride(PETSC_COMM_SELF, N, 0, 1, &ctx->fftReal));
457:     }

459:     PetscCall(PetscViewerDrawOpen(comm, NULL, "rhohat: Charge Density FT", 0, 0, 400, 300, &ctx->viewerRhoHat));
460:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerRhoHat, "rhohat_"));
461:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerRhoHat, 0, &draw));
462:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_ft"));
463:     PetscCall(PetscViewerSetFromOptions(ctx->viewerRhoHat));
464:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
465:     PetscCall(PetscObjectSetName((PetscObject)rhohat, "charge_density_ft"));
466:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));

468:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Potential", 400, 0, 400, 300, &ctx->viewerPhi));
469:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerPhi, "phi_"));
470:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerPhi, 0, &draw));
471:     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
472:     PetscCall(PetscViewerSetFromOptions(ctx->viewerPhi));
473:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
474:     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
475:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
476:   }
477:   if (ctx->moment_field_monitor) {
478:     Vec       n, p, e;
479:     Vec       nres, pres, eres;
480:     PetscDraw draw;

482:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density", 400, 0, 400, 300, &ctx->viewerN));
483:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerN, "n_"));
484:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerN, 0, &draw));
485:     PetscCall(PetscDrawSetSave(draw, "ex4_n_spatial"));
486:     PetscCall(PetscViewerSetFromOptions(ctx->viewerN));
487:     PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
488:     PetscCall(PetscObjectSetName((PetscObject)n, "Number Density"));
489:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));

491:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density", 800, 0, 400, 300, &ctx->viewerP));
492:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerP, "p_"));
493:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerP, 0, &draw));
494:     PetscCall(PetscDrawSetSave(draw, "ex4_p_spatial"));
495:     PetscCall(PetscViewerSetFromOptions(ctx->viewerP));
496:     PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
497:     PetscCall(PetscObjectSetName((PetscObject)p, "Momentum Density"));
498:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));

500:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Emergy Density (Pressure)", 1200, 0, 400, 300, &ctx->viewerE));
501:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerE, "e_"));
502:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerE, 0, &draw));
503:     PetscCall(PetscDrawSetSave(draw, "ex4_e_spatial"));
504:     PetscCall(PetscViewerSetFromOptions(ctx->viewerE));
505:     PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));
506:     PetscCall(PetscObjectSetName((PetscObject)e, "Energy Density (Pressure)"));
507:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));

509:     PetscDrawAxis axis;

511:     PetscCall(PetscDrawCreate(comm, NULL, "Moment Residual", 0, 320, 400, 300, &draw));
512:     PetscCall(PetscDrawSetSave(draw, "ex4_moment_res"));
513:     PetscCall(PetscDrawSetFromOptions(draw));
514:     PetscCall(PetscDrawLGCreate(draw, 3, &ctx->drawlgMomRes));
515:     PetscCall(PetscDrawDestroy(&draw));
516:     PetscCall(PetscDrawLGGetAxis(ctx->drawlgMomRes, &axis));
517:     PetscCall(PetscDrawAxisSetLabels(axis, "Moment Residial", "time", "Residual Norm"));
518:     PetscCall(PetscDrawLGSetLimits(ctx->drawlgMomRes, 0., ctx->steps * ctx->stepSize, -8, 0));

520:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density Residual", 400, 300, 400, 300, &ctx->viewerNRes));
521:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerNRes, "nres_"));
522:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerNRes, 0, &draw));
523:     PetscCall(PetscDrawSetSave(draw, "ex4_nres_spatial"));
524:     PetscCall(PetscViewerSetFromOptions(ctx->viewerNRes));
525:     PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
526:     PetscCall(PetscObjectSetName((PetscObject)nres, "Number Density Residual"));
527:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));

529:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density Residual", 800, 300, 400, 300, &ctx->viewerPRes));
530:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerPRes, "pres_"));
531:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerPRes, 0, &draw));
532:     PetscCall(PetscDrawSetSave(draw, "ex4_pres_spatial"));
533:     PetscCall(PetscViewerSetFromOptions(ctx->viewerPRes));
534:     PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
535:     PetscCall(PetscObjectSetName((PetscObject)pres, "Momentum Density Residual"));
536:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));

538:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Energy Density Residual", 1200, 300, 400, 300, &ctx->viewerERes));
539:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerERes, "eres_"));
540:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerERes, 0, &draw));
541:     PetscCall(PetscDrawSetSave(draw, "ex4_eres_spatial"));
542:     PetscCall(PetscViewerSetFromOptions(ctx->viewerERes));
543:     PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
544:     PetscCall(PetscObjectSetName((PetscObject)eres, "Energy Density Residual"));
545:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
546:   }
547:   PetscFunctionReturn(PETSC_SUCCESS);
548: }

550: static PetscErrorCode DestroyContext(AppCtx *ctx)
551: {
552:   PetscFunctionBeginUser;
553:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgic_x));
554:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgic_v));
555:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgcell_v));

557:   PetscCall(PetscDrawLGDestroy(&ctx->drawlgE));
558:   PetscCall(PetscDrawSPDestroy(&ctx->drawspE));
559:   PetscCall(PetscDrawSPDestroy(&ctx->drawspX));
560:   PetscCall(PetscViewerDestroy(&ctx->viewerRho));
561:   PetscCall(PetscViewerDestroy(&ctx->viewerRhoHat));
562:   PetscCall(MatDestroy(&ctx->fftPot));
563:   PetscCall(VecDestroy(&ctx->fftX));
564:   PetscCall(VecDestroy(&ctx->fftY));
565:   PetscCall(ISDestroy(&ctx->fftReal));
566:   PetscCall(PetscViewerDestroy(&ctx->viewerPhi));
567:   PetscCall(PetscViewerDestroy(&ctx->viewerN));
568:   PetscCall(PetscViewerDestroy(&ctx->viewerP));
569:   PetscCall(PetscViewerDestroy(&ctx->viewerE));
570:   PetscCall(PetscViewerDestroy(&ctx->viewerNRes));
571:   PetscCall(PetscViewerDestroy(&ctx->viewerPRes));
572:   PetscCall(PetscViewerDestroy(&ctx->viewerERes));
573:   PetscCall(PetscDrawLGDestroy(&ctx->drawlgMomRes));

575:   PetscCall(PetscBagDestroy(&ctx->bag));
576:   PetscFunctionReturn(PETSC_SUCCESS);
577: }

579: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *ctx)
580: {
581:   const PetscScalar *w;
582:   PetscInt           Np;

584:   PetscFunctionBeginUser;
585:   if (!ctx->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
586:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
587:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
588:   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]);
589:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
590:   PetscFunctionReturn(PETSC_SUCCESS);
591: }

593: 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[])
594: {
595:   for (PetscInt d = 0; d < dim; ++d) f0[0] += 0.5 * PetscSqr(u_x[d]);
596: }

598: static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
599: {
600:   PetscDS        ds;
601:   const PetscInt field = 0;
602:   PetscInt       Nf;
603:   void          *ctx;

605:   PetscFunctionBegin;
606:   PetscCall(DMGetApplicationContext(dm, &ctx));
607:   PetscCall(DMGetDS(dm, &ds));
608:   PetscCall(PetscDSGetNumFields(ds, &Nf));
609:   PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
610:   PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
611:   PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
612:   PetscFunctionReturn(PETSC_SUCCESS);
613: }

615: static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *ctx)
616: {
617:   DMSwarmCellDM celldm;
618:   DM            vdm;
619:   Vec           u[1];
620:   const char   *fields[1] = {"w_q"};

622:   PetscFunctionBegin;
623:   PetscCall(DMSwarmSetCellDMActive(sw, "velocity"));
624:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
625:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
626:   PetscCall(DMGetGlobalVector(vdm, &u[0]));
627:   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
628:   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
629:   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
630:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
631:   PetscFunctionReturn(PETSC_SUCCESS);
632: }

634: 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[])
635: {
636:   f0[0] = 0.;
637:   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d * dim + d]);
638: }

640: // Our model is E_max(t) = C e^{-gamma t} |cos(omega t - phi)|
641: static PetscErrorCode ComputeEmaxResidual(Tao tao, Vec x, Vec res, void *Ctx)
642: {
643:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
644:   const PetscScalar *a;
645:   PetscScalar       *F;
646:   PetscReal          C, gamma, omega, phi;

648:   PetscFunctionBegin;
649:   PetscCall(VecGetArrayRead(x, &a));
650:   PetscCall(VecGetArray(res, &F));
651:   C     = PetscRealPart(a[0]);
652:   gamma = PetscRealPart(a[1]);
653:   omega = PetscRealPart(a[2]);
654:   phi   = PetscRealPart(a[3]);
655:   PetscCall(VecRestoreArrayRead(x, &a));
656:   for (PetscInt i = ctx->s; i < ctx->e; ++i) F[i - ctx->s] = PetscPowReal(10., ctx->Emax[i]) - C * PetscExpReal(-gamma * ctx->t[i]) * PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi));
657:   PetscCall(VecRestoreArray(res, &F));
658:   PetscFunctionReturn(PETSC_SUCCESS);
659: }

661: // The Jacobian of the residual J = dr(x)/dx
662: static PetscErrorCode ComputeEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *Ctx)
663: {
664:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
665:   const PetscScalar *a;
666:   PetscScalar       *jac;
667:   PetscReal          C, gamma, omega, phi;
668:   const PetscInt     n = ctx->e - ctx->s;

670:   PetscFunctionBegin;
671:   PetscCall(VecGetArrayRead(x, &a));
672:   C     = PetscRealPart(a[0]);
673:   gamma = PetscRealPart(a[1]);
674:   omega = PetscRealPart(a[2]);
675:   phi   = PetscRealPart(a[3]);
676:   PetscCall(VecRestoreArrayRead(x, &a));
677:   PetscCall(MatDenseGetArray(J, &jac));
678:   for (PetscInt i = 0; i < n; ++i) {
679:     const PetscInt k = i + ctx->s;

681:     jac[i * 4 + 0] = -PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
682:     jac[i * 4 + 1] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
683:     jac[i * 4 + 2] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi);
684:     jac[i * 4 + 3] = -C * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi);
685:   }
686:   PetscCall(MatDenseRestoreArray(J, &jac));
687:   PetscFunctionReturn(PETSC_SUCCESS);
688: }

690: // Our model is log_10 E_max(t) = log_10 C - gamma t log_10 e + log_10 |cos(omega t - phi)|
691: static PetscErrorCode ComputeLogEmaxResidual(Tao tao, Vec x, Vec res, void *Ctx)
692: {
693:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
694:   const PetscScalar *a;
695:   PetscScalar       *F;
696:   PetscReal          C, gamma, omega, phi;

698:   PetscFunctionBegin;
699:   PetscCall(VecGetArrayRead(x, &a));
700:   PetscCall(VecGetArray(res, &F));
701:   C     = PetscRealPart(a[0]);
702:   gamma = PetscRealPart(a[1]);
703:   omega = PetscRealPart(a[2]);
704:   phi   = PetscRealPart(a[3]);
705:   PetscCall(VecRestoreArrayRead(x, &a));
706:   for (PetscInt i = ctx->s; i < ctx->e; ++i) {
707:     if (C < 0) {
708:       F[i - ctx->s] = 1e10;
709:       continue;
710:     }
711:     F[i - ctx->s] = ctx->Emax[i] - (PetscLog10Real(C) - gamma * ctx->t[i] * PetscLog10Real(PETSC_E) + PetscLog10Real(PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi))));
712:   }
713:   PetscCall(VecRestoreArray(res, &F));
714:   PetscFunctionReturn(PETSC_SUCCESS);
715: }

717: // The Jacobian of the residual J = dr(x)/dx
718: static PetscErrorCode ComputeLogEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *Ctx)
719: {
720:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
721:   const PetscScalar *a;
722:   PetscScalar       *jac;
723:   PetscReal          C, omega, phi;
724:   const PetscInt     n = ctx->e - ctx->s;

726:   PetscFunctionBegin;
727:   PetscCall(VecGetArrayRead(x, &a));
728:   C     = PetscRealPart(a[0]);
729:   omega = PetscRealPart(a[2]);
730:   phi   = PetscRealPart(a[3]);
731:   PetscCall(VecRestoreArrayRead(x, &a));
732:   PetscCall(MatDenseGetArray(J, &jac));
733:   for (PetscInt i = 0; i < n; ++i) {
734:     const PetscInt k = i + ctx->s;

736:     jac[0 * n + i] = -1. / (PetscLog10Real(PETSC_E) * C);
737:     jac[1 * n + i] = ctx->t[k] * PetscLog10Real(PETSC_E);
738:     jac[2 * n + i] = (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * ctx->t[k] * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)));
739:     jac[3 * n + i] = -(PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)));
740:   }
741:   PetscCall(MatDenseRestoreArray(J, &jac));
742:   PetscCall(MatViewFromOptions(J, NULL, "-emax_jac_view"));
743:   PetscFunctionReturn(PETSC_SUCCESS);
744: }

746: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
747: {
748:   AppCtx     *ctx = (AppCtx *)Ctx;
749:   DM          sw;
750:   PetscScalar intESq;
751:   PetscReal  *E, *x, *weight;
752:   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
753:   PetscReal   pmoments[4]; /* \int f, \int v f, \int v^2 f */
754:   PetscInt   *species, dim, Np, gNp;
755:   MPI_Comm    comm;
756:   PetscMPIInt rank;

758:   PetscFunctionBeginUser;
759:   if (step < 0 || !ctx->validE) PetscFunctionReturn(PETSC_SUCCESS);
760:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
761:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
762:   PetscCall(TSGetDM(ts, &sw));
763:   PetscCall(DMGetDimension(sw, &dim));
764:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
765:   PetscCall(DMSwarmGetSize(sw, &gNp));
766:   PetscCall(DMSwarmSortGetAccess(sw));
767:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
768:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
769:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
770:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

772:   for (PetscInt p = 0; p < Np; ++p) {
773:     for (PetscInt d = 0; d < 1; ++d) {
774:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
775:       if (temp > Emax) Emax = temp;
776:     }
777:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
778:     sum += E[p * dim];
779:     chargesum += ctx->charges[0] * weight[p];
780:   }
781:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
782:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
783:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : ctx->drawlgEmin;

785:   PetscDS ds;
786:   Vec     phi;

788:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
789:   PetscCall(DMGetDS(ctx->dmPot, &ds));
790:   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
791:   PetscCall(DMPlexComputeIntegralFEM(ctx->dmPot, phi, &intESq, ctx));
792:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

794:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
795:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
796:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
797:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
798:   PetscCall(PetscDrawLGAddPoint(ctx->drawlgE, &t, &lgEmax));
799:   if (ctx->efield_monitor == E_MONITOR_FULL) {
800:     PetscDraw draw;

802:     PetscCall(PetscDrawLGDraw(ctx->drawlgE));
803:     PetscCall(PetscDrawLGGetDraw(ctx->drawlgE, &draw));
804:     PetscCall(PetscDrawSave(draw));

806:     PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
807:     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));
808:     PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
809:   }

811:   // Compute decay rate and frequency
812:   PetscCall(PetscDrawLGGetData(ctx->drawlgE, NULL, &ctx->emaxCtx.e, &ctx->emaxCtx.t, &ctx->emaxCtx.Emax));
813:   if (!rank && !(ctx->emaxCtx.e % ctx->emaxCtx.per)) {
814:     Tao          tao;
815:     Mat          J;
816:     Vec          x, r;
817:     PetscScalar *a;
818:     PetscBool    fitLog = PETSC_TRUE, debug = PETSC_FALSE;

820:     PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
821:     PetscCall(TaoSetOptionsPrefix(tao, "emax_"));
822:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, 4, &x));
823:     PetscCall(TaoSetSolution(tao, x));
824:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, ctx->emaxCtx.e - ctx->emaxCtx.s, &r));
825:     if (fitLog) PetscCall(TaoSetResidualRoutine(tao, r, ComputeLogEmaxResidual, &ctx->emaxCtx));
826:     else PetscCall(TaoSetResidualRoutine(tao, r, ComputeEmaxResidual, &ctx->emaxCtx));
827:     PetscCall(VecDestroy(&r));
828:     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, ctx->emaxCtx.e - ctx->emaxCtx.s, 4, NULL, &J));
829:     if (fitLog) PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeLogEmaxJacobian, &ctx->emaxCtx));
830:     else PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeEmaxJacobian, &ctx->emaxCtx));
831:     PetscCall(MatDestroy(&J));
832:     PetscCall(TaoSetFromOptions(tao));
833:     PetscCall(VecGetArray(x, &a));
834:     a[0] = 0.02;
835:     a[1] = 0.15;
836:     a[2] = 1.4;
837:     a[3] = 0.45;
838:     PetscCall(VecRestoreArray(x, &a));
839:     PetscCall(TaoSolve(tao));
840:     if (debug) {
841:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "t = ["));
842:       for (PetscInt i = 0; i < ctx->emaxCtx.e; ++i) {
843:         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
844:         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ctx->emaxCtx.t[i]));
845:       }
846:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
847:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax = ["));
848:       for (PetscInt i = 0; i < ctx->emaxCtx.e; ++i) {
849:         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
850:         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ctx->emaxCtx.Emax[i]));
851:       }
852:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
853:     }
854:     PetscDraw     draw;
855:     PetscDrawAxis axis;
856:     char          title[PETSC_MAX_PATH_LEN];

858:     PetscCall(VecGetArray(x, &a));
859:     ctx->gamma = a[1];
860:     ctx->omega = a[2];
861:     if (ctx->efield_monitor == E_MONITOR_FULL) {
862:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax Fit: gamma %g omega %g C %g phi %g\n", a[1], a[2], a[0], a[3]));
863:       PetscCall(PetscDrawLGGetDraw(ctx->drawlgE, &draw));
864:       PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Max Electric Field gamma %.4g omega %.4g", a[1], a[2]));
865:       PetscCall(PetscDrawSetTitle(draw, title));
866:       PetscCall(PetscDrawLGGetAxis(ctx->drawlgE, &axis));
867:       PetscCall(PetscDrawAxisSetLabels(axis, title, "time", "E_max"));
868:     }
869:     PetscCall(VecRestoreArray(x, &a));
870:     PetscCall(VecDestroy(&x));
871:     PetscCall(TaoDestroy(&tao));
872:   }
873:   PetscFunctionReturn(PETSC_SUCCESS);
874: }

876: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
877: {
878:   AppCtx   *ctx = (AppCtx *)Ctx;
879:   DM        sw;
880:   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */

882:   PetscFunctionBeginUser;
883:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
884:   PetscCall(TSGetDM(ts, &sw));

886:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
887:   PetscCall(computeVelocityFEMMoments(sw, fmoments, ctx));

889:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3], (double)fmoments[0], (double)fmoments[1], (double)fmoments[2]));
890:   PetscFunctionReturn(PETSC_SUCCESS);
891: }

893: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, PetscCtx ctx)
894: {
895:   u[0] = 0.0;
896:   return PETSC_SUCCESS;
897: }

899: /*
900:     M_p w_p
901:       - Make M_p with "moments"
902:       - Get w_p from Swarm
903:     M_p v_p w_p
904:       - Get v_p from Swarm
905:       - pointwise multiply v_p and w_p
906:     M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p
907:       - ProjectField(sw, {n, p} U, {v_p} A, tmp_p)
908:       - pointwise multiply tmp_p and w_p

910:   Projection works fpr swarms
911:     Fields are FE from the CellDM, and aux fields are the swarm fields
912: */
913: static PetscErrorCode ComputeMomentFields(TS ts)
914: {
915:   AppCtx   *ctx;
916:   DM        sw;
917:   KSP       ksp;
918:   Mat       M_p, D_p;
919:   Vec       f, v, E, tmpMom;
920:   Vec       m, mold, mfluxold, mres, n, nrhs, nflux, nres, p, prhs, pflux, pres, e, erhs, eflux, eres;
921:   PetscReal dt, t;
922:   PetscInt  Nts;

924:   PetscFunctionBegin;
925:   PetscCall(TSGetStepNumber(ts, &Nts));
926:   PetscCall(TSGetTimeStep(ts, &dt));
927:   PetscCall(TSGetTime(ts, &t));
928:   PetscCall(TSGetDM(ts, &sw));
929:   PetscCall(DMGetApplicationContext(sw, &ctx));
930:   PetscCall(DMSwarmSetCellDMActive(sw, "moment fields"));
931:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
932:   // TODO In higher dimensions, we will have to create different M_p and D_p for each field
933:   PetscCall(DMCreateMassMatrix(sw, ctx->dmN, &M_p));
934:   PetscCall(DMCreateGradientMatrix(sw, ctx->dmN, &D_p));
935:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
936:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &v));
937:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "E_field", &E));
938:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

940:   PetscCall(MatViewFromOptions(ctx->MN, NULL, "-mn_view"));
941:   PetscCall(MatViewFromOptions(ctx->MP, NULL, "-mp_view"));
942:   PetscCall(MatViewFromOptions(ctx->ME, NULL, "-me_view"));
943:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

945:   PetscCall(DMGetGlobalVector(ctx->dmN, &nrhs));
946:   PetscCall(DMGetGlobalVector(ctx->dmN, &nflux));
947:   PetscCall(PetscObjectSetName((PetscObject)nrhs, "Weak number density"));
948:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
949:   PetscCall(DMGetGlobalVector(ctx->dmP, &prhs));
950:   PetscCall(DMGetGlobalVector(ctx->dmP, &pflux));
951:   PetscCall(PetscObjectSetName((PetscObject)prhs, "Weak momentum density"));
952:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
953:   PetscCall(DMGetGlobalVector(ctx->dmE, &erhs));
954:   PetscCall(DMGetGlobalVector(ctx->dmE, &eflux));
955:   PetscCall(PetscObjectSetName((PetscObject)erhs, "Weak energy density (pressure)"));
956:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));

958:   // Compute moments and fluxes
959:   PetscCall(VecDuplicate(f, &tmpMom));

961:   PetscCall(MatMultTranspose(M_p, f, nrhs));

963:   PetscCall(VecPointwiseMult(tmpMom, f, v));
964:   PetscCall(MatMultTranspose(M_p, tmpMom, prhs));
965:   PetscCall(MatMultTranspose(D_p, tmpMom, nflux));

967:   PetscCall(VecPointwiseMult(tmpMom, tmpMom, v));
968:   PetscCall(MatMultTranspose(M_p, tmpMom, erhs));
969:   PetscCall(MatMultTranspose(D_p, tmpMom, pflux));

971:   PetscCall(VecPointwiseMult(tmpMom, tmpMom, v));
972:   PetscCall(MatMultTranspose(D_p, tmpMom, eflux));

974:   PetscCall(VecPointwiseMult(tmpMom, f, E));
975:   PetscCall(MatMultTransposeAdd(M_p, tmpMom, pflux, pflux));

977:   PetscCall(VecPointwiseMult(tmpMom, v, E));
978:   PetscCall(VecScale(tmpMom, 2.));
979:   PetscCall(MatMultTransposeAdd(M_p, tmpMom, eflux, eflux));

981:   PetscCall(VecDestroy(&tmpMom));
982:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &v));
983:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
984:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "E_field", &E));

986:   PetscCall(MatDestroy(&M_p));
987:   PetscCall(MatDestroy(&D_p));

989:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)sw), &ksp));
990:   PetscCall(KSPSetOptionsPrefix(ksp, "mom_proj_"));
991:   PetscCall(KSPSetOperators(ksp, ctx->MN, ctx->MN));
992:   PetscCall(KSPSetFromOptions(ksp));
993:   PetscCall(KSPSolve(ksp, nrhs, n));
994:   PetscCall(KSPSetOperators(ksp, ctx->MP, ctx->MP));
995:   PetscCall(KSPSetFromOptions(ksp));
996:   PetscCall(KSPSolve(ksp, prhs, p));
997:   PetscCall(KSPSetOperators(ksp, ctx->ME, ctx->ME));
998:   PetscCall(KSPSetFromOptions(ksp));
999:   PetscCall(KSPSolve(ksp, erhs, e));
1000:   PetscCall(KSPDestroy(&ksp));
1001:   PetscCall(DMRestoreGlobalVector(ctx->dmN, &nrhs));
1002:   PetscCall(DMRestoreGlobalVector(ctx->dmP, &prhs));
1003:   PetscCall(DMRestoreGlobalVector(ctx->dmE, &erhs));

1005:   // Check moment residual
1006:   // TODO Fix global2local here
1007:   PetscReal res[3], logres[3];

1009:   PetscCall(DMGetGlobalVector(ctx->dmMom, &m));
1010:   PetscCall(VecISCopy(m, ctx->isN, SCATTER_FORWARD, n));
1011:   PetscCall(VecISCopy(m, ctx->isP, SCATTER_FORWARD, p));
1012:   PetscCall(VecISCopy(m, ctx->isE, SCATTER_FORWARD, e));
1013:   PetscCall(DMGetNamedGlobalVector(ctx->dmMom, "mold", &mold));
1014:   PetscCall(DMGetNamedGlobalVector(ctx->dmMom, "mfluxold", &mfluxold));
1015:   if (!Nts) goto end;

1017:   // e = \Tr{\tau}
1018:   // M_p w^{k+1} - M_p w^k - \Delta t D_p (w^k \vb{v}^k) = 0
1019:   // M_p \vb{p}^{k+1} - M_p \vb{p}^k - \Delta t D_p \tau - e \Delta t M_p \left( n \vb{E} \right) = 0
1020:   // M_p e^{k+1} - M_p e^k - \Delta t D_p \vb{Q} - 2 e \Delta t M_p \left( \vb{p} \cdot \vb{E} \right) = 0
1021:   PetscCall(DMGetGlobalVector(ctx->dmMom, &mres));
1022:   PetscCall(VecCopy(mfluxold, mres));
1023:   PetscCall(VecAXPBYPCZ(mres, 1. / dt, -1. / dt, -1., m, mold));

1025:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
1026:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
1027:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
1028:   PetscCall(VecISCopy(mres, ctx->isN, SCATTER_REVERSE, nres));
1029:   PetscCall(VecISCopy(mres, ctx->isP, SCATTER_REVERSE, pres));
1030:   PetscCall(VecISCopy(mres, ctx->isE, SCATTER_REVERSE, eres));
1031:   PetscCall(VecNorm(nres, NORM_2, &res[0]));
1032:   PetscCall(VecNorm(pres, NORM_2, &res[1]));
1033:   PetscCall(VecNorm(eres, NORM_2, &res[2]));
1034:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Mass Residual: %g\n", (double)res[0]));
1035:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Momentum Residual: %g\n", (double)res[1]));
1036:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Energy Residual: %g\n", (double)res[2]));
1037:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));
1038:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));
1039:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
1040:   PetscCall(DMRestoreGlobalVector(ctx->dmMom, &mres));

1042:   for (PetscInt i = 0; i < 3; ++i) logres[i] = PetscLog10Real(res[i]);
1043:   PetscCall(PetscDrawLGAddCommonPoint(ctx->drawlgMomRes, t, logres));
1044:   PetscCall(PetscDrawLGDraw(ctx->drawlgMomRes));
1045:   {
1046:     PetscDraw draw;

1048:     PetscCall(PetscDrawLGGetDraw(ctx->drawlgMomRes, &draw));
1049:     PetscCall(PetscDrawSave(draw));
1050:   }

1052: end:
1053:   PetscCall(VecCopy(m, mold));
1054:   PetscCall(DMRestoreGlobalVector(ctx->dmMom, &m));
1055:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmMom, "mold", &mold));
1056:   PetscCall(VecISCopy(mfluxold, ctx->isN, SCATTER_FORWARD, nflux));
1057:   PetscCall(VecISCopy(mfluxold, ctx->isP, SCATTER_FORWARD, pflux));
1058:   PetscCall(VecISCopy(mfluxold, ctx->isE, SCATTER_FORWARD, eflux));
1059:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmMom, "mfluxold", &mfluxold));

1061:   PetscCall(DMRestoreGlobalVector(ctx->dmN, &nflux));
1062:   PetscCall(DMRestoreGlobalVector(ctx->dmP, &pflux));
1063:   PetscCall(DMRestoreGlobalVector(ctx->dmE, &eflux));
1064:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));
1065:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));
1066:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));
1067:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
1068:   PetscFunctionReturn(PETSC_SUCCESS);
1069: }

1071: static PetscErrorCode MonitorMomentFields(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
1072: {
1073:   AppCtx *ctx = (AppCtx *)Ctx;
1074:   Vec     n, p, e;
1075:   Vec     nres, pres, eres;

1077:   PetscFunctionBeginUser;
1078:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
1079:   PetscCall(ComputeMomentFields(ts));

1081:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
1082:   PetscCall(VecView(n, ctx->viewerN));
1083:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));

1085:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
1086:   PetscCall(VecView(p, ctx->viewerP));
1087:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));

1089:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));
1090:   PetscCall(VecView(e, ctx->viewerE));
1091:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));

1093:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
1094:   PetscCall(VecView(nres, ctx->viewerNRes));
1095:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));

1097:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
1098:   PetscCall(VecView(pres, ctx->viewerPRes));
1099:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));

1101:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
1102:   PetscCall(VecView(eres, ctx->viewerERes));
1103:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
1104:   PetscFunctionReturn(PETSC_SUCCESS);
1105: }

1107: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
1108: {
1109:   AppCtx    *ctx = (AppCtx *)Ctx;
1110:   DM         sw;
1111:   PetscDraw  drawic_x, drawic_v;
1112:   PetscReal *weight, *pos, *vel;
1113:   PetscInt   dim, Np;

1115:   PetscFunctionBegin;
1116:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
1117:   if (step == 0) {
1118:     PetscCall(TSGetDM(ts, &sw));
1119:     PetscCall(DMGetDimension(sw, &dim));
1120:     PetscCall(DMSwarmGetLocalSize(sw, &Np));

1122:     PetscCall(PetscDrawHGReset(ctx->drawhgic_x));
1123:     PetscCall(PetscDrawHGGetDraw(ctx->drawhgic_x, &drawic_x));
1124:     PetscCall(PetscDrawClear(drawic_x));
1125:     PetscCall(PetscDrawFlush(drawic_x));

1127:     PetscCall(PetscDrawHGReset(ctx->drawhgic_v));
1128:     PetscCall(PetscDrawHGGetDraw(ctx->drawhgic_v, &drawic_v));
1129:     PetscCall(PetscDrawClear(drawic_v));
1130:     PetscCall(PetscDrawFlush(drawic_v));

1132:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
1133:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1134:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1135:     for (PetscInt p = 0; p < Np; ++p) {
1136:       PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgic_x, pos[p * dim], weight[p]));
1137:       PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgic_v, vel[p * dim], weight[p]));
1138:     }
1139:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
1140:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1141:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));

1143:     PetscCall(PetscDrawHGDraw(ctx->drawhgic_x));
1144:     PetscCall(PetscDrawHGSave(ctx->drawhgic_x));
1145:     PetscCall(PetscDrawHGDraw(ctx->drawhgic_v));
1146:     PetscCall(PetscDrawHGSave(ctx->drawhgic_v));
1147:   }
1148:   PetscFunctionReturn(PETSC_SUCCESS);
1149: }

1151: // Right now, make the complete velocity histogram
1152: PetscErrorCode MonitorVelocity(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
1153: {
1154:   AppCtx      *ctx = (AppCtx *)Ctx;
1155:   DM           sw, dm;
1156:   Vec          ks;
1157:   PetscProbFn *cdf;
1158:   PetscDraw    drawcell_v;
1159:   PetscScalar *ksa;
1160:   PetscReal   *weight, *vel;
1161:   PetscInt    *pidx;
1162:   PetscInt     dim, Npc, cStart, cEnd, cell = ctx->velocity_monitor;

1164:   PetscFunctionBegin;
1165:   PetscCall(TSGetDM(ts, &sw));
1166:   PetscCall(DMGetDimension(sw, &dim));

1168:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1169:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1170:   PetscCall(VecCreate(PetscObjectComm((PetscObject)dm), &ks));
1171:   PetscCall(PetscObjectSetName((PetscObject)ks, "KS Statistic by Cell"));
1172:   PetscCall(VecSetSizes(ks, cEnd - cStart, PETSC_DETERMINE));
1173:   PetscCall(VecSetFromOptions(ks));
1174:   switch (dim) {
1175:   case 1:
1176:     //cdf = PetscCDFMaxwellBoltzmann1D;
1177:     cdf = PetscCDFGaussian1D;
1178:     break;
1179:   case 2:
1180:     cdf = PetscCDFMaxwellBoltzmann2D;
1181:     break;
1182:   case 3:
1183:     cdf = PetscCDFMaxwellBoltzmann3D;
1184:     break;
1185:   default:
1186:     SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported", dim);
1187:   }

1189:   PetscCall(PetscDrawHGReset(ctx->drawhgcell_v));
1190:   PetscCall(PetscDrawHGGetDraw(ctx->drawhgcell_v, &drawcell_v));
1191:   PetscCall(PetscDrawClear(drawcell_v));
1192:   PetscCall(PetscDrawFlush(drawcell_v));

1194:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1195:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1196:   PetscCall(DMSwarmSortGetAccess(sw));
1197:   PetscCall(VecGetArrayWrite(ks, &ksa));
1198:   for (PetscInt c = cStart; c < cEnd; ++c) {
1199:     Vec          cellv, cellw;
1200:     PetscScalar *cella, *cellaw;
1201:     PetscReal    totWgt = 0.;

1203:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1204:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellv));
1205:     PetscCall(VecSetBlockSize(cellv, dim));
1206:     PetscCall(VecSetSizes(cellv, Npc * dim, Npc));
1207:     PetscCall(VecSetFromOptions(cellv));
1208:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellw));
1209:     PetscCall(VecSetSizes(cellw, Npc, Npc));
1210:     PetscCall(VecSetFromOptions(cellw));
1211:     PetscCall(VecGetArrayWrite(cellv, &cella));
1212:     PetscCall(VecGetArrayWrite(cellw, &cellaw));
1213:     for (PetscInt q = 0; q < Npc; ++q) {
1214:       const PetscInt p = pidx[q];
1215:       if (c == cell) PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgcell_v, vel[p * dim], weight[p]));
1216:       for (PetscInt d = 0; d < dim; ++d) cella[q * dim + d] = vel[p * dim + d];
1217:       cellaw[q] = weight[p];
1218:       totWgt += weight[p];
1219:     }
1220:     PetscCall(VecRestoreArrayWrite(cellv, &cella));
1221:     PetscCall(VecRestoreArrayWrite(cellw, &cellaw));
1222:     PetscCall(VecScale(cellw, 1. / totWgt));
1223:     PetscCall(PetscProbComputeKSStatisticWeighted(cellv, cellw, cdf, &ksa[c - cStart]));
1224:     PetscCall(VecDestroy(&cellv));
1225:     PetscCall(VecDestroy(&cellw));
1226:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1227:   }
1228:   PetscCall(VecRestoreArrayWrite(ks, &ksa));
1229:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1230:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1231:   PetscCall(DMSwarmSortRestoreAccess(sw));

1233:   PetscReal minalpha, maxalpha;
1234:   PetscInt  mincell, maxcell;

1236:   PetscCall(VecFilter(ks, PETSC_SMALL));
1237:   PetscCall(VecMin(ks, &mincell, &minalpha));
1238:   PetscCall(VecMax(ks, &maxcell, &maxalpha));
1239:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)dm), "Step %" PetscInt_FMT ": Min/Max KS statistic %g/%g in cell %" PetscInt_FMT "/%" PetscInt_FMT "\n", step, minalpha, maxalpha, mincell, maxcell));
1240:   PetscCall(VecViewFromOptions(ks, NULL, "-ks_view"));
1241:   PetscCall(VecDestroy(&ks));

1243:   PetscCall(PetscDrawHGDraw(ctx->drawhgcell_v));
1244:   PetscCall(PetscDrawHGSave(ctx->drawhgcell_v));
1245:   PetscFunctionReturn(PETSC_SUCCESS);
1246: }

1248: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
1249: {
1250:   AppCtx         *ctx = (AppCtx *)Ctx;
1251:   DM              dm, sw;
1252:   PetscDrawAxis   axis;
1253:   char            title[1024];
1254:   PetscScalar    *x, *v, *weight;
1255:   PetscReal       lower[3], upper[3], speed;
1256:   const PetscInt *s;
1257:   PetscInt        dim, cStart, cEnd, c;

1259:   PetscFunctionBeginUser;
1260:   if (step > 0 && step % ctx->ostep == 0) {
1261:     PetscCall(TSGetDM(ts, &sw));
1262:     PetscCall(DMSwarmGetCellDM(sw, &dm));
1263:     PetscCall(DMGetDimension(dm, &dim));
1264:     PetscCall(DMGetBoundingBox(dm, lower, upper));
1265:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1266:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1267:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1268:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1269:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
1270:     PetscCall(DMSwarmSortGetAccess(sw));
1271:     PetscCall(PetscDrawSPReset(ctx->drawspX));
1272:     PetscCall(PetscDrawSPGetAxis(ctx->drawspX, &axis));
1273:     PetscCall(PetscSNPrintf(title, 1024, "Step %" PetscInt_FMT " Time: %g", step, (double)t));
1274:     PetscCall(PetscDrawAxisSetLabels(axis, title, "x", "v"));
1275:     PetscCall(PetscDrawSPSetLimits(ctx->drawspX, lower[0], upper[0], lower[1], upper[1]));
1276:     PetscCall(PetscDrawSPSetLimits(ctx->drawspX, lower[0], upper[0], -12, 12));
1277:     for (c = 0; c < cEnd - cStart; ++c) {
1278:       PetscInt *pidx, Npc, q;
1279:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1280:       for (q = 0; q < Npc; ++q) {
1281:         const PetscInt p = pidx[q];
1282:         if (s[p] == 0) {
1283:           speed = 0.;
1284:           for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]);
1285:           speed = PetscSqrtReal(speed);
1286:           if (dim == 1) {
1287:             PetscCall(PetscDrawSPAddPointColorized(ctx->drawspX, &x[p * dim], &v[p * dim], &speed));
1288:           } else {
1289:             PetscCall(PetscDrawSPAddPointColorized(ctx->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
1290:           }
1291:         } else if (s[p] == 1) {
1292:           PetscCall(PetscDrawSPAddPoint(ctx->drawspX, &x[p * dim], &v[p * dim]));
1293:         }
1294:       }
1295:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1296:     }
1297:     PetscCall(PetscDrawSPDraw(ctx->drawspX, PETSC_TRUE));
1298:     PetscDraw draw;
1299:     PetscCall(PetscDrawSPGetDraw(ctx->drawspX, &draw));
1300:     PetscCall(PetscDrawSave(draw));
1301:     PetscCall(DMSwarmSortRestoreAccess(sw));
1302:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1303:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1304:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1305:     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
1306:   }
1307:   PetscFunctionReturn(PETSC_SUCCESS);
1308: }

1310: static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
1311: {
1312:   AppCtx *ctx = (AppCtx *)Ctx;
1313:   DM      dm, sw;

1315:   PetscFunctionBeginUser;
1316:   if (step > 0 && step % ctx->ostep == 0) {
1317:     PetscCall(TSGetDM(ts, &sw));
1318:     PetscCall(DMSwarmGetCellDM(sw, &dm));

1320:     if (ctx->validE) {
1321:       PetscScalar *x, *E, *weight;
1322:       PetscReal    lower[3], upper[3], xval;
1323:       PetscDraw    draw;
1324:       PetscInt     dim, cStart, cEnd;

1326:       PetscCall(DMGetDimension(dm, &dim));
1327:       PetscCall(DMGetBoundingBox(dm, lower, upper));
1328:       PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

1330:       PetscCall(PetscDrawSPReset(ctx->drawspE));
1331:       PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1332:       PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1333:       PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

1335:       PetscCall(DMSwarmSortGetAccess(sw));
1336:       for (PetscInt c = 0; c < cEnd - cStart; ++c) {
1337:         PetscReal Eavg = 0.0;
1338:         PetscInt *pidx, Npc;

1340:         PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1341:         for (PetscInt q = 0; q < Npc; ++q) {
1342:           const PetscInt p = pidx[q];
1343:           Eavg += E[p * dim];
1344:         }
1345:         Eavg /= Npc;
1346:         xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
1347:         PetscCall(PetscDrawSPAddPoint(ctx->drawspE, &xval, &Eavg));
1348:         PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1349:       }
1350:       PetscCall(PetscDrawSPDraw(ctx->drawspE, PETSC_TRUE));
1351:       PetscCall(PetscDrawSPGetDraw(ctx->drawspE, &draw));
1352:       PetscCall(PetscDrawSave(draw));
1353:       PetscCall(DMSwarmSortRestoreAccess(sw));
1354:       PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1355:       PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1356:       PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1357:     }

1359:     Vec rho, rhohat, phi;

1361:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
1362:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
1363:     PetscCall(VecView(rho, ctx->viewerRho));
1364:     PetscCall(VecISCopy(ctx->fftX, ctx->fftReal, SCATTER_FORWARD, rho));
1365:     PetscCall(MatMult(ctx->fftPot, ctx->fftX, ctx->fftY));
1366:     PetscCall(VecFilter(ctx->fftY, PETSC_SMALL));
1367:     PetscCall(VecViewFromOptions(ctx->fftX, NULL, "-real_view"));
1368:     PetscCall(VecViewFromOptions(ctx->fftY, NULL, "-fft_view"));
1369:     PetscCall(VecISCopy(ctx->fftY, ctx->fftReal, SCATTER_REVERSE, rhohat));
1370:     PetscCall(VecSetValue(rhohat, 0, 0., INSERT_VALUES)); // Remove large DC component
1371:     PetscCall(VecView(rhohat, ctx->viewerRhoHat));
1372:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
1373:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));

1375:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
1376:     PetscCall(VecView(phi, ctx->viewerPhi));
1377:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
1378:   }
1379:   PetscFunctionReturn(PETSC_SUCCESS);
1380: }

1382: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
1383: {
1384:   PetscBag   bag;
1385:   Parameter *p;

1387:   PetscFunctionBeginUser;
1388:   /* setup PETSc parameter bag */
1389:   PetscCall(PetscBagGetData(ctx->bag, &p));
1390:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
1391:   bag = ctx->bag;
1392:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
1393:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
1394:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
1395:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
1396:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
1397:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
1398:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
1399:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

1401:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
1402:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
1403:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
1404:   PetscCall(PetscBagSetFromOptions(bag));
1405:   {
1406:     PetscViewer       viewer;
1407:     PetscViewerFormat format;
1408:     PetscBool         flg;

1410:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
1411:     if (flg) {
1412:       PetscCall(PetscViewerPushFormat(viewer, format));
1413:       PetscCall(PetscBagView(bag, viewer));
1414:       PetscCall(PetscViewerFlush(viewer));
1415:       PetscCall(PetscViewerPopFormat(viewer));
1416:       PetscCall(PetscViewerDestroy(&viewer));
1417:     }
1418:   }
1419:   PetscFunctionReturn(PETSC_SUCCESS);
1420: }

1422: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *ctx, DM *dm)
1423: {
1424:   DMField         coordField;
1425:   IS              cellIS;
1426:   PetscQuadrature quad;
1427:   PetscReal      *wt, *pt;
1428:   PetscInt        cdim, cStart, cEnd;

1430:   PetscFunctionBeginUser;
1431:   PetscCall(DMCreate(comm, dm));
1432:   PetscCall(DMSetType(*dm, DMPLEX));
1433:   PetscCall(DMSetFromOptions(*dm));
1434:   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
1435:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

1437:   // Cache the mesh geometry
1438:   PetscCall(DMGetCoordinateField(*dm, &coordField));
1439:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
1440:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
1441:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
1442:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
1443:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
1444:   PetscCall(PetscMalloc1(1, &wt));
1445:   PetscCall(PetscMalloc1(2, &pt));
1446:   wt[0] = 1.;
1447:   pt[0] = -1.;
1448:   pt[1] = -1.;
1449:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
1450:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &ctx->fegeom));
1451:   PetscCall(PetscQuadratureDestroy(&quad));
1452:   PetscCall(ISDestroy(&cellIS));
1453:   PetscFunctionReturn(PETSC_SUCCESS);
1454: }

1456: 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[])
1457: {
1458:   f0[0] = -constants[SIGMA];
1459: }

1461: 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[])
1462: {
1463:   PetscInt d;
1464:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
1465: }

1467: 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[])
1468: {
1469:   PetscInt d;
1470:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
1471: }

1473: /*
1474:    /  I   -grad\ / q \ = /0\
1475:    \-div    0  / \phi/   \f/
1476: */
1477: static void f0_q(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[])
1478: {
1479:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
1480: }

1482: static void f1_q(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[])
1483: {
1484:   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
1485: }

1487: static void f0_phi_backgroundCharge(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[])
1488: {
1489:   f0[0] += constants[SIGMA];
1490:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
1491: }

1493: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
1494: static void g0_qq(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 g0[])
1495: {
1496:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
1497: }

1499: static void g2_qphi(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 g2[])
1500: {
1501:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
1502: }

1504: static void g1_phiq(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 g1[])
1505: {
1506:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
1507: }

1509: static PetscErrorCode CreateFEM(DM dm, AppCtx *ctx)
1510: {
1511:   PetscFE   fephi, feq;
1512:   PetscDS   ds;
1513:   PetscBool simplex;
1514:   PetscInt  dim;

1516:   PetscFunctionBeginUser;
1517:   PetscCall(DMGetDimension(dm, &dim));
1518:   PetscCall(DMPlexIsSimplex(dm, &simplex));
1519:   if (ctx->em == EM_MIXED) {
1520:     DMLabel        label;
1521:     const PetscInt id = 1;

1523:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
1524:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
1525:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
1526:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1527:     PetscCall(PetscFECopyQuadrature(feq, fephi));
1528:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
1529:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
1530:     PetscCall(DMCreateDS(dm));
1531:     PetscCall(PetscFEDestroy(&fephi));
1532:     PetscCall(PetscFEDestroy(&feq));

1534:     PetscCall(DMGetLabel(dm, "marker", &label));
1535:     PetscCall(DMGetDS(dm, &ds));

1537:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
1538:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
1539:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
1540:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
1541:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

1543:     PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "wall", label, 1, &id, 0, 0, NULL, (PetscVoidFn *)zero, NULL, NULL, NULL));

1545:   } else {
1546:     MatNullSpace nullsp;
1547:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
1548:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1549:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
1550:     PetscCall(DMCreateDS(dm));
1551:     PetscCall(DMGetDS(dm, &ds));
1552:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
1553:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
1554:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
1555:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
1556:     PetscCall(MatNullSpaceDestroy(&nullsp));
1557:     PetscCall(PetscFEDestroy(&fephi));
1558:   }
1559:   PetscFunctionReturn(PETSC_SUCCESS);
1560: }

1562: static PetscErrorCode CreatePoisson(DM dm, AppCtx *ctx)
1563: {
1564:   SNES         snes;
1565:   Mat          J;
1566:   MatNullSpace nullSpace;

1568:   PetscFunctionBeginUser;
1569:   PetscCall(CreateFEM(dm, ctx));
1570:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
1571:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
1572:   PetscCall(SNESSetDM(snes, dm));
1573:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, ctx));
1574:   PetscCall(SNESSetFromOptions(snes));

1576:   PetscCall(DMCreateMatrix(dm, &J));
1577:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
1578:   PetscCall(MatSetNullSpace(J, nullSpace));
1579:   PetscCall(MatNullSpaceDestroy(&nullSpace));
1580:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
1581:   PetscCall(MatDestroy(&J));
1582:   if (ctx->em == EM_MIXED) {
1583:     const PetscInt potential = 1;

1585:     PetscCall(DMCreateSubDM(dm, 1, &potential, &ctx->isPot, &ctx->dmPot));
1586:   } else {
1587:     ctx->dmPot = dm;
1588:     PetscCall(PetscObjectReference((PetscObject)ctx->dmPot));
1589:   }
1590:   PetscCall(DMCreateMassMatrix(ctx->dmPot, ctx->dmPot, &ctx->M));
1591:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
1592:   ctx->snes = snes;
1593:   PetscFunctionReturn(PETSC_SUCCESS);
1594: }

1596: // Conservation of mass (m = 1.0)
1597: // n_t + 1/ m p_x = 0
1598: static void f0_mass(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[])
1599: {
1600:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_t[uOff[0]] + u_x[uOff_x[1] + d * dim + d];
1601: }

1603: // Conservation of momentum (m = 1, e = 1)
1604: // p_t + (u p)_x = -pr_x + e n E
1605: // p_t + (div u) p + u . grad p = -pr_x + e n E
1606: // p_t + (div p) p / n - (p . grad n) p / n^2 + p / n . grad p = -pr_x + e n E
1607: static void f0_momentum(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[])
1608: {
1609:   const PetscScalar n = u[uOff[0]];

1611:   for (PetscInt d = 0; d < dim; ++d) {
1612:     PetscReal divp = 0.;

1614:     f0[d] += u_t[uOff[1] + d];
1615:     for (PetscInt e = 0; e < dim; ++e) {
1616:       f0[d] += u[uOff[1] + e] * u_x[uOff_x[1] + d * dim + e] / n;                    // p / n . grad p
1617:       f0[d] -= (u[uOff[1] + e] * u_x[uOff_x[0] + e]) * u[uOff[1] + d] / PetscSqr(n); // -(p . grad n) p / n^2
1618:       divp += u_x[uOff_x[1] + e * dim + e];
1619:     }
1620:     f0[d] += divp * u[uOff[1] + d] / n; // (div p) p / n
1621:     f0[d] += u_x[uOff_x[2] + d];        // pr_x
1622:     f0[d] -= n * a[d];                  // -e n E
1623:   }
1624: }

1626: // Conservation of energy
1627: // pr_t + (u pr)_x = -3 pr u_x - q_x
1628: // pr_t + (div u) pr + u . grad pr = -3 pr (div u) - q_x
1629: // pr_t + 4 (div u) pr + u . grad pr = -q_x
1630: // pr_t + 4 div p pr / n - 4 (p . grad n) pr / n^2 + p . grad pr / n = -q_x
1631: static void f0_energy(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[])
1632: {
1633:   const PetscScalar n    = u[uOff[0]];
1634:   const PetscScalar pr   = u[uOff[2]];
1635:   PetscReal         divp = 0.;

1637:   f0[0] += u_t[uOff[2]];
1638:   for (PetscInt d = 0; d < dim; ++d) {
1639:     f0[0] += u[uOff[1] + d] * u_x[uOff_x[2] + d] / n;                     // p . grad pr / n
1640:     f0[0] -= 4. * u[uOff[1] + d] * u_x[uOff_x[0] + d] * pr / PetscSqr(n); // -4 (p . grad n) pr / n^2
1641:     divp += u_x[uOff_x[1] + d * dim + d];
1642:   }
1643:   f0[0] += 4. * divp * pr / n; // 4 div p pr / n
1644: }

1646: static PetscErrorCode SetupMomentProblem(DM dm, AppCtx *ctx)
1647: {
1648:   PetscDS ds;

1650:   PetscFunctionBegin;
1651:   PetscCall(DMGetDS(dm, &ds));
1652:   PetscCall(PetscDSSetResidual(ds, 0, f0_mass, NULL));
1653:   PetscCall(PetscDSSetResidual(ds, 1, f0_momentum, NULL));
1654:   PetscCall(PetscDSSetResidual(ds, 2, f0_energy, NULL));
1655:   //PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_mass_uu, NULL, NULL, NULL));
1656:   PetscFunctionReturn(PETSC_SUCCESS);
1657: }

1659: static PetscErrorCode CreateMomentFields(DM odm, AppCtx *ctx)
1660: {
1661:   DM             dm;
1662:   PetscFE        fe;
1663:   DMPolytopeType ct;
1664:   PetscInt       dim, cStart;

1666:   PetscFunctionBeginUser;
1667:   PetscCall(DMClone(odm, &dm));
1668:   PetscCall(DMGetDimension(dm, &dim));
1669:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, NULL));
1670:   PetscCall(DMPlexGetCellType(dm, cStart, &ct));
1671:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
1672:   PetscCall(PetscObjectSetName((PetscObject)fe, "number density"));
1673:   PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fe));
1674:   PetscCall(PetscFEDestroy(&fe));
1675:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, dim, ct, NULL, PETSC_DETERMINE, &fe));
1676:   PetscCall(PetscObjectSetName((PetscObject)fe, "momentum density"));
1677:   PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fe));
1678:   PetscCall(PetscFEDestroy(&fe));
1679:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
1680:   PetscCall(PetscObjectSetName((PetscObject)fe, "energy density"));
1681:   PetscCall(DMSetField(dm, 2, NULL, (PetscObject)fe));
1682:   PetscCall(PetscFEDestroy(&fe));
1683:   PetscCall(DMCreateDS(dm));
1684:   PetscCall(SetupMomentProblem(dm, ctx));
1685:   ctx->dmMom = dm;
1686:   PetscInt field;

1688:   field = 0;
1689:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isN, &ctx->dmN));
1690:   PetscCall(DMCreateMassMatrix(ctx->dmN, ctx->dmN, &ctx->MN));
1691:   field = 1;
1692:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isP, &ctx->dmP));
1693:   PetscCall(DMCreateMassMatrix(ctx->dmP, ctx->dmP, &ctx->MP));
1694:   field = 2;
1695:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isE, &ctx->dmE));
1696:   PetscCall(DMCreateMassMatrix(ctx->dmE, ctx->dmE, &ctx->ME));
1697:   PetscFunctionReturn(PETSC_SUCCESS);
1698: }

1700: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1701: {
1702:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1703:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
1704:   return PETSC_SUCCESS;
1705: }
1706: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1707: {
1708:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1709:   return PETSC_SUCCESS;
1710: }

1712: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1713: {
1714:   const PetscReal alpha = scale ? scale[0] : 0.0;
1715:   const PetscReal k     = scale ? scale[1] : 1.;
1716:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
1717:   return PETSC_SUCCESS;
1718: }

1720: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1721: {
1722:   const PetscReal alpha = scale ? scale[0] : 0.;
1723:   const PetscReal k     = scale ? scale[0] : 1.;
1724:   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
1725:   return PETSC_SUCCESS;
1726: }

1728: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
1729: {
1730:   PetscFE        fe;
1731:   DMPolytopeType ct;
1732:   PetscInt       dim, cStart;
1733:   const char    *prefix = "v";

1735:   PetscFunctionBegin;
1736:   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
1737:   PetscCall(DMSetType(*vdm, DMPLEX));
1738:   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
1739:   PetscCall(DMSetFromOptions(*vdm));
1740:   PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
1741:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

1743:   PetscCall(DMGetDimension(*vdm, &dim));
1744:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
1745:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
1746:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
1747:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
1748:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
1749:   PetscCall(DMCreateDS(*vdm));
1750:   PetscCall(PetscFEDestroy(&fe));
1751:   PetscFunctionReturn(PETSC_SUCCESS);
1752: }

1754: /*
1755:   InitializeParticles_Centroid - Initialize a regular grid of particles.

1757:   Input Parameters:
1758: + sw      - The `DMSWARM`
1759: - force1D - Treat the spatial domain as 1D

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

1764:   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
1765:   and velocity meshes.
1766: */
1767: static PetscErrorCode InitializeParticles_Centroid(DM sw)
1768: {
1769:   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
1770:   DMSwarmCellDM celldm;
1771:   DM            xdm, vdm;
1772:   PetscReal     vmin[3], vmax[3];
1773:   PetscReal    *x, *v;
1774:   PetscInt     *species, *cellid;
1775:   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
1776:   PetscBool     flg;
1777:   MPI_Comm      comm;
1778:   const char   *cellidname;

1780:   PetscFunctionBegin;
1781:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

1783:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
1784:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1785:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
1786:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
1787:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
1788:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
1789:   PetscOptionsEnd();
1790:   debug = swarm->printCoords;

1792:   PetscCall(DMGetDimension(sw, &dim));
1793:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1794:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

1796:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1797:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1798:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

1800:   // One particle per centroid on the tensor product grid
1801:   Npc = (vcEnd - vcStart) * Ns;
1802:   Np  = (xcEnd - xcStart) * Npc;
1803:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
1804:   if (debug) {
1805:     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
1806:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
1807:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
1808:     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
1809:     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
1810:     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
1811:   }

1813:   // Set species and cellid
1814:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
1815:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
1816:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1817:   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
1818:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
1819:     for (PetscInt s = 0; s < Ns; ++s) {
1820:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
1821:         species[p] = s;
1822:         cellid[p]  = c;
1823:       }
1824:     }
1825:   }
1826:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1827:   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));

1829:   // Set particle coordinates
1830:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1831:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1832:   PetscCall(DMSwarmSortGetAccess(sw));
1833:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
1834:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
1835:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
1836:     const PetscInt cell = c + xcStart;
1837:     PetscInt      *pidx, Npc;
1838:     PetscReal      centroid[3], volume;

1840:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1841:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
1842:     for (PetscInt s = 0; s < Ns; ++s) {
1843:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
1844:         const PetscInt p = pidx[q * Ns + s];

1846:         for (PetscInt d = 0; d < dim; ++d) {
1847:           x[p * dim + d] = centroid[d];
1848:           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
1849:         }
1850:         if (debug > 1) {
1851:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
1852:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
1853:           for (PetscInt d = 0; d < dim; ++d) {
1854:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1855:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
1856:           }
1857:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
1858:           for (PetscInt d = 0; d < dim; ++d) {
1859:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1860:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
1861:           }
1862:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
1863:         }
1864:       }
1865:     }
1866:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1867:   }
1868:   PetscCall(DMSwarmSortRestoreAccess(sw));
1869:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1870:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1871:   PetscFunctionReturn(PETSC_SUCCESS);
1872: }

1874: /*
1875:   InitializeWeights - Compute weight for each local particle

1877:   Input Parameters:
1878: + sw          - The `DMSwarm`
1879: . totalWeight - The sum of all particle weights
1880: . func        - The PDF for the particle spatial distribution
1881: - param       - The PDF parameters

1883:   Notes:
1884:   The PDF for velocity is assumed to be a Gaussian

1886:   The particle weights are returned in the `w_q` field of `sw`.
1887: */
1888: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
1889: {
1890:   DM               xdm, vdm;
1891:   DMSwarmCellDM    celldm;
1892:   PetscScalar     *weight;
1893:   PetscQuadrature  xquad;
1894:   const PetscReal *xq, *xwq;
1895:   const PetscInt   order = 5;
1896:   PetscReal        xi0[3];
1897:   PetscReal        xwtot = 0., pwtot = 0.;
1898:   PetscInt         xNq;
1899:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
1900:   MPI_Comm         comm;
1901:   PetscMPIInt      rank;

1903:   PetscFunctionBegin;
1904:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
1905:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1906:   PetscCall(DMGetDimension(sw, &dim));
1907:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1908:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1909:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
1910:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1911:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1912:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

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

1919:   // Integrate the density function to get the weights of particles in each cell
1920:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
1921:   PetscCall(DMSwarmSortGetAccess(sw));
1922:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1923:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
1924:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
1925:     PetscInt          *pidx, Npc;
1926:     PetscInt           xNc;
1927:     const PetscScalar *xarray;
1928:     PetscScalar       *xcoords = NULL;
1929:     PetscBool          xisDG;

1931:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1932:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1933:     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);
1934:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
1935:     for (PetscInt q = 0; q < xNq; ++q) {
1936:       // Transform quadrature points from ref space to real space
1937:       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
1938:       // Get probability density at quad point
1939:       //   No need to scale xqr since PDF will be periodic
1940:       PetscCall((*func)(xqr, param, &xden));
1941:       xw += xden * (xwq[q] * xdetJ);
1942:     }
1943:     xwtot += xw;
1944:     if (debug) {
1945:       IS              globalOrdering;
1946:       const PetscInt *ordering;

1948:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
1949:       PetscCall(ISGetIndices(globalOrdering, &ordering));
1950:       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));
1951:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
1952:     }
1953:     // Set weights to be Gaussian in velocity cells
1954:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
1955:       const PetscInt     p  = pidx[vc * Ns + 0];
1956:       PetscReal          vw = 0.;
1957:       PetscInt           vNc;
1958:       const PetscScalar *varray;
1959:       PetscScalar       *vcoords = NULL;
1960:       PetscBool          visDG;

1962:       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1963:       // TODO: Fix 2 stream Ask Joe
1964:       //   Two stream function from 1/2pi v^2 e^(-v^2/2)
1965:       //   vw = 1. / (PetscSqrtReal(2 * PETSC_PI)) * (((coords_v[0] * PetscExpReal(-PetscSqr(coords_v[0]) / 2.)) - (coords_v[1] * PetscExpReal(-PetscSqr(coords_v[1]) / 2.)))) - 0.5 * PetscErfReal(coords_v[0] / PetscSqrtReal(2.)) + 0.5 * (PetscErfReal(coords_v[1] / PetscSqrtReal(2.)));
1966:       vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.)));

1968:       weight[p] = totalWeight * vw * xw;
1969:       pwtot += weight[p];
1970:       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 1: %g, %g, %g", p, xw, vw, totalWeight);
1971:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1972:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
1973:     }
1974:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1975:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1976:   }
1977:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1978:   PetscCall(DMSwarmSortRestoreAccess(sw));
1979:   PetscCall(PetscQuadratureDestroy(&xquad));

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

1984:     PetscCall(PetscSynchronizedFlush(comm, NULL));
1985:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1986:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
1987:   }
1988:   PetscFunctionReturn(PETSC_SUCCESS);
1989: }

1991: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *ctx)
1992: {
1993:   PetscReal scale[2] = {ctx->cosine_coefficients[0], ctx->cosine_coefficients[1]};
1994:   PetscInt  dim;

1996:   PetscFunctionBegin;
1997:   PetscCall(DMGetDimension(sw, &dim));
1998:   PetscCall(InitializeParticles_Centroid(sw));
1999:   PetscCall(InitializeWeights(sw, ctx->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
2000:   PetscFunctionReturn(PETSC_SUCCESS);
2001: }

2003: static PetscErrorCode InitializeConstants(DM sw, AppCtx *ctx)
2004: {
2005:   DM         dm;
2006:   PetscInt  *species;
2007:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
2008:   PetscInt   Np, dim;

2010:   PetscFunctionBegin;
2011:   PetscCall(DMSwarmGetCellDM(sw, &dm));
2012:   PetscCall(DMGetDimension(sw, &dim));
2013:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2014:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
2015:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
2016:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
2017:   for (PetscInt p = 0; p < Np; ++p) {
2018:     totalWeight += weight[p];
2019:     totalCharge += ctx->charges[species[p]] * weight[p];
2020:   }
2021:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
2022:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
2023:   {
2024:     Parameter *param;
2025:     PetscReal  Area;

2027:     PetscCall(PetscBagGetData(ctx->bag, &param));
2028:     switch (dim) {
2029:     case 1:
2030:       Area = (gmax[0] - gmin[0]);
2031:       break;
2032:     case 2:
2033:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
2034:       break;
2035:     case 3:
2036:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
2037:       break;
2038:     default:
2039:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
2040:     }
2041:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
2042:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
2043:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, ctx->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)ctx->charges[0], (double)global_charge, (double)Area));
2044:     param->sigma = PetscAbsReal(global_charge / (Area));

2046:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
2047:     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,
2048:                           (double)param->vlasovNumber));
2049:   }
2050:   /* Setup Constants */
2051:   {
2052:     PetscDS    ds;
2053:     Parameter *param;
2054:     PetscCall(PetscBagGetData(ctx->bag, &param));
2055:     PetscScalar constants[NUM_CONSTANTS];
2056:     constants[SIGMA]   = param->sigma;
2057:     constants[V0]      = param->v0;
2058:     constants[T0]      = param->t0;
2059:     constants[X0]      = param->x0;
2060:     constants[M0]      = param->m0;
2061:     constants[Q0]      = param->q0;
2062:     constants[PHI0]    = param->phi0;
2063:     constants[POISSON] = param->poissonNumber;
2064:     constants[VLASOV]  = param->vlasovNumber;
2065:     PetscCall(DMGetDS(dm, &ds));
2066:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
2067:   }
2068:   PetscFunctionReturn(PETSC_SUCCESS);
2069: }

2071: static PetscErrorCode CreateSwarm(DM dm, AppCtx *ctx, DM *sw)
2072: {
2073:   DMSwarmCellDM celldm;
2074:   DM            vdm;
2075:   PetscReal     v0[2] = {1., 0.};
2076:   PetscInt      dim;

2078:   PetscFunctionBeginUser;
2079:   PetscCall(DMGetDimension(dm, &dim));
2080:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
2081:   PetscCall(DMSetType(*sw, DMSWARM));
2082:   PetscCall(DMSetDimension(*sw, dim));
2083:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
2084:   PetscCall(DMSetApplicationContext(*sw, ctx));

2086:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
2087:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
2088:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
2089:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));

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

2093:   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
2094:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
2095:   PetscCall(DMSwarmCellDMDestroy(&celldm));

2097:   const char *vfieldnames[2] = {"w_q"};

2099:   PetscCall(CreateVelocityDM(*sw, &vdm));
2100:   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
2101:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
2102:   PetscCall(DMSwarmCellDMDestroy(&celldm));
2103:   PetscCall(DMDestroy(&vdm));

2105:   DM mdm;

2107:   PetscCall(DMClone(dm, &mdm));
2108:   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
2109:   PetscCall(DMCopyDisc(dm, mdm));
2110:   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
2111:   PetscCall(DMDestroy(&mdm));
2112:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
2113:   PetscCall(DMSwarmCellDMDestroy(&celldm));

2115:   DM mfdm;

2117:   PetscCall(DMClone(dm, &mfdm));
2118:   PetscCall(PetscObjectSetName((PetscObject)mfdm, "moment fields"));
2119:   PetscCall(DMCopyDisc(dm, mfdm));
2120:   PetscCall(DMSwarmCellDMCreate(mfdm, 1, &fieldnames[1], 1, fieldnames, &celldm));
2121:   PetscCall(DMDestroy(&mfdm));
2122:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
2123:   PetscCall(DMSwarmCellDMDestroy(&celldm));

2125:   PetscCall(DMSetFromOptions(*sw));
2126:   PetscCall(DMSetUp(*sw));

2128:   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
2129:   ctx->swarm = *sw;
2130:   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
2131:   if (ctx->perturbed_weights) {
2132:     PetscCall(InitializeParticles_PerturbedWeights(*sw, ctx));
2133:   } else {
2134:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
2135:     PetscCall(DMSwarmInitializeCoordinates(*sw));
2136:     PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
2137:   }
2138:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
2139:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
2140:   PetscFunctionReturn(PETSC_SUCCESS);
2141: }

2143: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
2144: {
2145:   AppCtx     *ctx;
2146:   PetscReal  *coords;
2147:   PetscInt   *species, dim, Np, Ns;
2148:   PetscMPIInt size;

2150:   PetscFunctionBegin;
2151:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
2152:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
2153:   PetscCall(DMGetDimension(sw, &dim));
2154:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2155:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
2156:   PetscCall(DMGetApplicationContext(sw, &ctx));

2158:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2159:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
2160:   for (PetscInt p = 0; p < Np; ++p) {
2161:     PetscReal *pcoord = &coords[p * dim];
2162:     PetscReal  pE[3]  = {0., 0., 0.};

2164:     /* Calculate field at particle p due to particle q */
2165:     for (PetscInt q = 0; q < Np; ++q) {
2166:       PetscReal *qcoord = &coords[q * dim];
2167:       PetscReal  rpq[3], r, r3, q_q;

2169:       if (p == q) continue;
2170:       q_q = ctx->charges[species[q]] * 1.;
2171:       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
2172:       r = DMPlex_NormD_Internal(dim, rpq);
2173:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
2174:       r3 = PetscPowRealInt(r, 3);
2175:       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
2176:     }
2177:     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
2178:   }
2179:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
2180:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2181:   PetscFunctionReturn(PETSC_SUCCESS);
2182: }

2184: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
2185: {
2186:   DM         dm;
2187:   AppCtx    *ctx;
2188:   PetscDS    ds;
2189:   PetscFE    fe;
2190:   KSP        ksp;
2191:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
2192:   Vec        rho;         // Charge density, M^{-1} rhoRhs
2193:   Vec        phi, locPhi; // Potential
2194:   Vec        f;           // Particle weights
2195:   PetscReal *coords;
2196:   PetscInt   dim, cStart, cEnd, Np;

2198:   PetscFunctionBegin;
2199:   PetscCall(DMGetApplicationContext(sw, &ctx));
2200:   PetscCall(PetscLogEventBegin(ctx->ESolveEvent, snes, sw, 0, 0));
2201:   PetscCall(DMGetDimension(sw, &dim));
2202:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

2204:   PetscCall(SNESGetDM(snes, &dm));
2205:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
2206:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
2207:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
2208:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
2209:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

2211:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
2212:   PetscCall(MatViewFromOptions(ctx->M, NULL, "-m_view"));
2213:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

2215:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
2216:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

2218:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
2219:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
2220:   PetscCall(KSPSetOperators(ksp, ctx->M, ctx->M));
2221:   PetscCall(KSPSetFromOptions(ksp));
2222:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

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

2226:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
2227:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
2228:   PetscCall(KSPDestroy(&ksp));

2230:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
2231:   PetscCall(VecSet(phi, 0.0));
2232:   PetscCall(SNESSolve(snes, rhoRhs, phi));
2233:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
2234:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

2236:   PetscCall(DMGetLocalVector(dm, &locPhi));
2237:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
2238:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
2239:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
2240:   PetscCall(PetscLogEventEnd(ctx->ESolveEvent, snes, sw, 0, 0));

2242:   PetscCall(DMGetDS(dm, &ds));
2243:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
2244:   PetscCall(DMSwarmSortGetAccess(sw));
2245:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
2246:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

2248:   PetscCall(PetscLogEventBegin(ctx->ETabEvent, snes, sw, 0, 0));
2249:   PetscTabulation tab;
2250:   PetscReal      *pcoord, *refcoord;
2251:   PetscFEGeom    *chunkgeom = NULL;
2252:   PetscInt        maxNcp    = 0;

2254:   for (PetscInt c = cStart; c < cEnd; ++c) {
2255:     PetscInt Ncp;

2257:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
2258:     maxNcp = PetscMax(maxNcp, Ncp);
2259:   }
2260:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2261:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2262:   // This can raise an FP_INEXACT in the dgemm inside
2263:   PetscCall(PetscFPTrapPush(PETSC_FP_TRAP_OFF));
2264:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
2265:   PetscCall(PetscFPTrapPop());
2266:   for (PetscInt c = cStart; c < cEnd; ++c) {
2267:     PetscScalar *clPhi = NULL;
2268:     PetscInt    *points;
2269:     PetscInt     Ncp;

2271:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
2272:     for (PetscInt cp = 0; cp < Ncp; ++cp)
2273:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
2274:     {
2275:       PetscCall(PetscFEGeomGetChunk(ctx->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
2276:       for (PetscInt i = 0; i < Ncp; ++i) {
2277:         const PetscReal x0[3] = {-1., -1., -1.};
2278:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
2279:       }
2280:     }
2281:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
2282:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2283:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
2284:       const PetscReal *basisDer = tab->T[1];
2285:       const PetscInt   p        = points[cp];

2287:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
2288:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
2289:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
2290:     }
2291:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2292:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2293:   }
2294:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2295:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2296:   PetscCall(PetscTabulationDestroy(&tab));
2297:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2298:   PetscCall(DMSwarmSortRestoreAccess(sw));
2299:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
2300:   PetscCall(PetscFEGeomRestoreChunk(ctx->fegeom, 0, 1, &chunkgeom));
2301:   PetscCall(PetscLogEventEnd(ctx->ETabEvent, snes, sw, 0, 0));
2302:   PetscFunctionReturn(PETSC_SUCCESS);
2303: }

2305: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, Mat M_p, PetscReal E[])
2306: {
2307:   DM         dm;
2308:   AppCtx    *ctx;
2309:   PetscDS    ds;
2310:   PetscFE    fe;
2311:   KSP        ksp;
2312:   Vec        rhoRhs, rhoRhsFull;   // Weak charge density, \int phi_i rho, and embedding in mixed problem
2313:   Vec        rho;                  // Charge density, M^{-1} rhoRhs
2314:   Vec        phi, locPhi, phiFull; // Potential and embedding in mixed problem
2315:   Vec        f;                    // Particle weights
2316:   PetscReal *coords;
2317:   PetscInt   dim, cStart, cEnd, Np;

2319:   PetscFunctionBegin;
2320:   PetscCall(DMGetApplicationContext(sw, &ctx));
2321:   PetscCall(PetscLogEventBegin(ctx->ESolveEvent, snes, sw, 0, 0));
2322:   PetscCall(DMGetDimension(sw, &dim));
2323:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

2325:   PetscCall(SNESGetDM(snes, &dm));
2326:   PetscCall(DMGetGlobalVector(ctx->dmPot, &rhoRhs));
2327:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
2328:   PetscCall(DMGetGlobalVector(dm, &rhoRhsFull));
2329:   PetscCall(PetscObjectSetName((PetscObject)rhoRhsFull, "Weak charge density"));
2330:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
2331:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
2332:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

2334:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
2335:   PetscCall(MatViewFromOptions(ctx->M, NULL, "-m_view"));
2336:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

2338:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
2339:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

2341:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
2342:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
2343:   PetscCall(KSPSetOperators(ksp, ctx->M, ctx->M));
2344:   PetscCall(KSPSetFromOptions(ksp));
2345:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

2347:   PetscCall(VecISCopy(rhoRhsFull, ctx->isPot, SCATTER_FORWARD, rhoRhs));
2348:   //PetscCall(VecScale(rhoRhsFull, -1.0));

2350:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
2351:   PetscCall(VecViewFromOptions(rhoRhsFull, NULL, "-rho_full_view"));
2352:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
2353:   PetscCall(DMRestoreGlobalVector(ctx->dmPot, &rhoRhs));
2354:   PetscCall(KSPDestroy(&ksp));

2356:   PetscCall(DMGetGlobalVector(dm, &phiFull));
2357:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
2358:   PetscCall(VecSet(phiFull, 0.0));
2359:   PetscCall(SNESSolve(snes, rhoRhsFull, phiFull));
2360:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhsFull));
2361:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

2363:   PetscCall(VecISCopy(phiFull, ctx->isPot, SCATTER_REVERSE, phi));
2364:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

2366:   PetscCall(DMGetLocalVector(dm, &locPhi));
2367:   PetscCall(DMGlobalToLocalBegin(dm, phiFull, INSERT_VALUES, locPhi));
2368:   PetscCall(DMGlobalToLocalEnd(dm, phiFull, INSERT_VALUES, locPhi));
2369:   PetscCall(DMRestoreGlobalVector(dm, &phiFull));
2370:   PetscCall(PetscLogEventEnd(ctx->ESolveEvent, snes, sw, 0, 0));

2372:   PetscCall(DMGetDS(dm, &ds));
2373:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
2374:   PetscCall(DMSwarmSortGetAccess(sw));
2375:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
2376:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

2378:   PetscCall(PetscLogEventBegin(ctx->ETabEvent, snes, sw, 0, 0));
2379:   PetscTabulation tab;
2380:   PetscReal      *pcoord, *refcoord;
2381:   PetscFEGeom    *chunkgeom = NULL;
2382:   PetscInt        maxNcp    = 0;

2384:   for (PetscInt c = cStart; c < cEnd; ++c) {
2385:     PetscInt Ncp;

2387:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
2388:     maxNcp = PetscMax(maxNcp, Ncp);
2389:   }
2390:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2391:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2392:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
2393:   for (PetscInt c = cStart; c < cEnd; ++c) {
2394:     PetscScalar *clPhi = NULL;
2395:     PetscInt    *points;
2396:     PetscInt     Ncp;

2398:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
2399:     for (PetscInt cp = 0; cp < Ncp; ++cp)
2400:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
2401:     {
2402:       PetscCall(PetscFEGeomGetChunk(ctx->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
2403:       for (PetscInt i = 0; i < Ncp; ++i) {
2404:         const PetscReal x0[3] = {-1., -1., -1.};
2405:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
2406:       }
2407:     }
2408:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
2409:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2410:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
2411:       const PetscInt p = points[cp];

2413:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
2414:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
2415:       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
2416:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
2417:     }
2418:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2419:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2420:   }
2421:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2422:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2423:   PetscCall(PetscTabulationDestroy(&tab));
2424:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2425:   PetscCall(DMSwarmSortRestoreAccess(sw));
2426:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
2427:   PetscCall(PetscFEGeomRestoreChunk(ctx->fegeom, 0, 1, &chunkgeom));
2428:   PetscCall(PetscLogEventEnd(ctx->ETabEvent, snes, sw, 0, 0));
2429:   PetscFunctionReturn(PETSC_SUCCESS);
2430: }

2432: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
2433: {
2434:   AppCtx    *ctx;
2435:   Mat        M_p;
2436:   PetscReal *E;
2437:   PetscInt   dim, Np;

2439:   PetscFunctionBegin;
2442:   PetscCall(DMGetDimension(sw, &dim));
2443:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2444:   PetscCall(DMGetApplicationContext(sw, &ctx));

2446:   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
2447:   // TODO: Could share sort context with space cellDM
2448:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
2449:   PetscCall(DMCreateMassMatrix(sw, ctx->dmPot, &M_p));
2450:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));

2452:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2453:   PetscCall(PetscArrayzero(E, Np * dim));
2454:   ctx->validE = PETSC_TRUE;

2456:   switch (ctx->em) {
2457:   case EM_COULOMB:
2458:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
2459:     break;
2460:   case EM_PRIMAL:
2461:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
2462:     break;
2463:   case EM_MIXED:
2464:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, M_p, E));
2465:     break;
2466:   case EM_NONE:
2467:     break;
2468:   default:
2469:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
2470:   }
2471:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2472:   PetscCall(MatDestroy(&M_p));
2473:   PetscFunctionReturn(PETSC_SUCCESS);
2474: }

2476: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, PetscCtx ctx)
2477: {
2478:   DM                 sw;
2479:   SNES               snes = ((AppCtx *)ctx)->snes;
2480:   const PetscScalar *u;
2481:   PetscScalar       *g;
2482:   PetscReal         *E, m_p = 1., q_p = -1.;
2483:   PetscInt           dim, d, Np, p;

2485:   PetscFunctionBeginUser;
2486:   PetscCall(TSGetDM(ts, &sw));
2487:   PetscCall(ComputeFieldAtParticles(snes, sw));

2489:   PetscCall(DMGetDimension(sw, &dim));
2490:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2491:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2492:   PetscCall(VecGetArrayRead(U, &u));
2493:   PetscCall(VecGetArray(G, &g));
2494:   Np /= 2 * dim;
2495:   for (p = 0; p < Np; ++p) {
2496:     for (d = 0; d < dim; ++d) {
2497:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
2498:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
2499:     }
2500:   }
2501:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2502:   PetscCall(VecRestoreArrayRead(U, &u));
2503:   PetscCall(VecRestoreArray(G, &g));
2504:   PetscFunctionReturn(PETSC_SUCCESS);
2505: }

2507: /* J_{ij} = dF_i/dx_j
2508:    J_p = (  0   1)
2509:          (-w^2  0)
2510:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
2511:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
2512: */
2513: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, PetscCtx ctx)
2514: {
2515:   DM               sw;
2516:   const PetscReal *coords, *vel;
2517:   PetscInt         dim, d, Np, p, rStart;

2519:   PetscFunctionBeginUser;
2520:   PetscCall(TSGetDM(ts, &sw));
2521:   PetscCall(DMGetDimension(sw, &dim));
2522:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2523:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
2524:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2525:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2526:   Np /= 2 * dim;
2527:   for (p = 0; p < Np; ++p) {
2528:     // TODO This is not right because dv/dx has the electric field in it
2529:     PetscScalar vals[4] = {0., 1., -1., 0.};

2531:     for (d = 0; d < dim; ++d) {
2532:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
2533:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
2534:     }
2535:   }
2536:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2537:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2538:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
2539:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
2540:   PetscFunctionReturn(PETSC_SUCCESS);
2541: }

2543: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *Ctx)
2544: {
2545:   AppCtx            *ctx = (AppCtx *)Ctx;
2546:   DM                 sw;
2547:   const PetscScalar *v;
2548:   PetscScalar       *xres;
2549:   PetscInt           Np, p, d, dim;

2551:   PetscFunctionBeginUser;
2552:   PetscCall(PetscLogEventBegin(ctx->RhsXEvent, ts, 0, 0, 0));
2553:   PetscCall(TSGetDM(ts, &sw));
2554:   PetscCall(DMGetDimension(sw, &dim));
2555:   PetscCall(VecGetLocalSize(Xres, &Np));
2556:   PetscCall(VecGetArrayRead(V, &v));
2557:   PetscCall(VecGetArray(Xres, &xres));
2558:   Np /= dim;
2559:   for (p = 0; p < Np; ++p) {
2560:     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
2561:   }
2562:   PetscCall(VecRestoreArrayRead(V, &v));
2563:   PetscCall(VecRestoreArray(Xres, &xres));
2564:   PetscCall(PetscLogEventEnd(ctx->RhsXEvent, ts, 0, 0, 0));
2565:   PetscFunctionReturn(PETSC_SUCCESS);
2566: }

2568: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *Ctx)
2569: {
2570:   DM                 sw;
2571:   AppCtx            *ctx  = (AppCtx *)Ctx;
2572:   SNES               snes = ((AppCtx *)ctx)->snes;
2573:   const PetscScalar *x;
2574:   PetscScalar       *vres;
2575:   PetscReal         *E, m_p, q_p;
2576:   PetscInt           Np, p, dim, d;
2577:   Parameter         *param;

2579:   PetscFunctionBeginUser;
2580:   PetscCall(PetscLogEventBegin(ctx->RhsVEvent, ts, 0, 0, 0));
2581:   PetscCall(TSGetDM(ts, &sw));
2582:   PetscCall(ComputeFieldAtParticles(snes, sw));

2584:   PetscCall(DMGetDimension(sw, &dim));
2585:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2586:   PetscCall(PetscBagGetData(ctx->bag, &param));
2587:   m_p = ctx->masses[0] * param->m0;
2588:   q_p = ctx->charges[0] * param->q0;
2589:   PetscCall(VecGetLocalSize(Vres, &Np));
2590:   PetscCall(VecGetArrayRead(X, &x));
2591:   PetscCall(VecGetArray(Vres, &vres));
2592:   Np /= dim;
2593:   for (p = 0; p < Np; ++p) {
2594:     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
2595:   }
2596:   PetscCall(VecRestoreArrayRead(X, &x));
2597:   /*
2598:     Synchronized, ordered output for parallel/sequential test cases.
2599:     In the 1D (on the 2D mesh) case, every y component should be zero.
2600:   */
2601:   if (ctx->checkVRes) {
2602:     PetscBool pr = ctx->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
2603:     PetscInt  step;

2605:     PetscCall(TSGetStepNumber(ts, &step));
2606:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
2607:     for (PetscInt p = 0; p < Np; ++p) {
2608:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
2609:       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]));
2610:     }
2611:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
2612:   }
2613:   PetscCall(VecRestoreArray(Vres, &vres));
2614:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2615:   PetscCall(PetscLogEventEnd(ctx->RhsVEvent, ts, 0, 0, 0));
2616:   PetscFunctionReturn(PETSC_SUCCESS);
2617: }

2619: /* Discrete Gradients Formulation: S, F, gradF (G) */
2620: PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, PetscCtx ctx)
2621: {
2622:   PetscScalar vals[4] = {0., 1., -1., 0.};
2623:   DM          sw;
2624:   PetscInt    dim, d, Np, p, rStart;

2626:   PetscFunctionBeginUser;
2627:   PetscCall(TSGetDM(ts, &sw));
2628:   PetscCall(DMGetDimension(sw, &dim));
2629:   PetscCall(VecGetLocalSize(U, &Np));
2630:   PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
2631:   Np /= 2 * dim;
2632:   for (p = 0; p < Np; ++p) {
2633:     for (d = 0; d < dim; ++d) {
2634:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
2635:       PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
2636:     }
2637:   }
2638:   PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
2639:   PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
2640:   PetscFunctionReturn(PETSC_SUCCESS);
2641: }

2643: PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *Ctx)
2644: {
2645:   AppCtx            *ctx = (AppCtx *)Ctx;
2646:   DM                 sw;
2647:   Vec                phi;
2648:   const PetscScalar *u;
2649:   PetscInt           dim, Np, cStart, cEnd;
2650:   PetscReal         *vel, *coords, m_p = 1.;

2652:   PetscFunctionBeginUser;
2653:   PetscCall(TSGetDM(ts, &sw));
2654:   PetscCall(DMGetDimension(sw, &dim));
2655:   PetscCall(DMPlexGetHeightStratum(ctx->dmPot, 0, &cStart, &cEnd));

2657:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
2658:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
2659:   PetscCall(computeFieldEnergy(ctx->dmPot, phi, F));
2660:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

2662:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2663:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2664:   PetscCall(DMSwarmSortGetAccess(sw));
2665:   PetscCall(VecGetArrayRead(U, &u));
2666:   PetscCall(VecGetLocalSize(U, &Np));
2667:   Np /= 2 * dim;
2668:   for (PetscInt c = cStart; c < cEnd; ++c) {
2669:     PetscInt *points;
2670:     PetscInt  Ncp;

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

2677:       *F += 0.5 * m_p * v2;
2678:     }
2679:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2680:   }
2681:   PetscCall(VecRestoreArrayRead(U, &u));
2682:   PetscCall(DMSwarmSortRestoreAccess(sw));
2683:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2684:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2685:   PetscFunctionReturn(PETSC_SUCCESS);
2686: }

2688: /* dF/dx = q E   dF/dv = v */
2689: PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, PetscCtx ctx)
2690: {
2691:   DM                 sw;
2692:   SNES               snes = ((AppCtx *)ctx)->snes;
2693:   const PetscReal   *coords, *vel, *E;
2694:   const PetscScalar *u;
2695:   PetscScalar       *g;
2696:   PetscReal          m_p = 1., q_p = -1.;
2697:   PetscInt           dim, d, Np, p;

2699:   PetscFunctionBeginUser;
2700:   PetscCall(TSGetDM(ts, &sw));
2701:   PetscCall(DMGetDimension(sw, &dim));
2702:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2703:   PetscCall(VecGetArrayRead(U, &u));
2704:   PetscCall(VecGetArray(G, &g));

2706:   PetscLogEvent COMPUTEFIELD;
2707:   PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
2708:   PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
2709:   PetscCall(ComputeFieldAtParticles(snes, sw));
2710:   PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
2711:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2712:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2713:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2714:   for (p = 0; p < Np; ++p) {
2715:     for (d = 0; d < dim; ++d) {
2716:       g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
2717:       g[(p * 2 + 1) * dim + d] = m_p * u[(p * 2 + 1) * dim + d];
2718:     }
2719:   }
2720:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2721:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2722:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2723:   PetscCall(VecRestoreArrayRead(U, &u));
2724:   PetscCall(VecRestoreArray(G, &g));
2725:   PetscFunctionReturn(PETSC_SUCCESS);
2726: }

2728: static PetscErrorCode CreateSolution(TS ts)
2729: {
2730:   DM       sw;
2731:   Vec      u;
2732:   PetscInt dim, Np;

2734:   PetscFunctionBegin;
2735:   PetscCall(TSGetDM(ts, &sw));
2736:   PetscCall(DMGetDimension(sw, &dim));
2737:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2738:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
2739:   PetscCall(VecSetBlockSize(u, dim));
2740:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
2741:   PetscCall(VecSetUp(u));
2742:   PetscCall(TSSetSolution(ts, u));
2743:   PetscCall(VecDestroy(&u));
2744:   PetscFunctionReturn(PETSC_SUCCESS);
2745: }

2747: static PetscErrorCode SetProblem(TS ts)
2748: {
2749:   AppCtx *ctx;
2750:   DM      sw;

2752:   PetscFunctionBegin;
2753:   PetscCall(TSGetDM(ts, &sw));
2754:   PetscCall(DMGetApplicationContext(sw, &ctx));
2755:   // Define unified system for (X, V)
2756:   {
2757:     Mat      J;
2758:     PetscInt dim, Np;

2760:     PetscCall(DMGetDimension(sw, &dim));
2761:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2762:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
2763:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
2764:     PetscCall(MatSetBlockSize(J, 2 * dim));
2765:     PetscCall(MatSetFromOptions(J));
2766:     PetscCall(MatSetUp(J));
2767:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, ctx));
2768:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, ctx));
2769:     PetscCall(MatDestroy(&J));
2770:   }
2771:   /* Define split system for X and V */
2772:   {
2773:     Vec             u;
2774:     IS              isx, isv, istmp;
2775:     const PetscInt *idx;
2776:     PetscInt        dim, Np, rstart;

2778:     PetscCall(TSGetSolution(ts, &u));
2779:     PetscCall(DMGetDimension(sw, &dim));
2780:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2781:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
2782:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
2783:     PetscCall(ISGetIndices(istmp, &idx));
2784:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
2785:     PetscCall(ISRestoreIndices(istmp, &idx));
2786:     PetscCall(ISDestroy(&istmp));
2787:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
2788:     PetscCall(ISGetIndices(istmp, &idx));
2789:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
2790:     PetscCall(ISRestoreIndices(istmp, &idx));
2791:     PetscCall(ISDestroy(&istmp));
2792:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
2793:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
2794:     PetscCall(ISDestroy(&isx));
2795:     PetscCall(ISDestroy(&isv));
2796:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, ctx));
2797:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, ctx));
2798:   }
2799:   // Define symplectic formulation U_t = S . G, where G = grad F
2800:   {
2801:     PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, ctx));
2802:   }
2803:   PetscFunctionReturn(PETSC_SUCCESS);
2804: }

2806: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
2807: {
2808:   DM        sw;
2809:   Vec       u;
2810:   PetscReal t, maxt, dt;
2811:   PetscInt  n, maxn;

2813:   PetscFunctionBegin;
2814:   PetscCall(TSGetDM(ts, &sw));
2815:   PetscCall(TSGetTime(ts, &t));
2816:   PetscCall(TSGetMaxTime(ts, &maxt));
2817:   PetscCall(TSGetTimeStep(ts, &dt));
2818:   PetscCall(TSGetStepNumber(ts, &n));
2819:   PetscCall(TSGetMaxSteps(ts, &maxn));

2821:   PetscCall(TSReset(ts));
2822:   PetscCall(TSSetDM(ts, sw));
2823:   PetscCall(TSSetFromOptions(ts));
2824:   PetscCall(TSSetTime(ts, t));
2825:   PetscCall(TSSetMaxTime(ts, maxt));
2826:   PetscCall(TSSetTimeStep(ts, dt));
2827:   PetscCall(TSSetStepNumber(ts, n));
2828:   PetscCall(TSSetMaxSteps(ts, maxn));

2830:   PetscCall(CreateSolution(ts));
2831:   PetscCall(SetProblem(ts));
2832:   PetscCall(TSGetSolution(ts, &u));
2833:   PetscFunctionReturn(PETSC_SUCCESS);
2834: }

2836: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *Ctx)
2837: {
2838:   DM        sw, cdm;
2839:   PetscInt  Np;
2840:   PetscReal low[2], high[2];
2841:   AppCtx   *ctx = (AppCtx *)Ctx;

2843:   sw = ctx->swarm;
2844:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
2845:   // Get the bounding box so we can equally space the particles
2846:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
2847:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2848:   // shift it by h/2 so nothing is initialized directly on a boundary
2849:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
2850:   x[1] = 0.;
2851:   return PETSC_SUCCESS;
2852: }

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

2857:   Input Parameters:
2858: + ts         - The TS
2859: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

2861:   Output Parameters:
2862: . u - The initialized solution vector

2864:   Level: advanced

2866: .seealso: InitializeSolve()
2867: */
2868: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
2869: {
2870:   DM       sw;
2871:   Vec      u, gc, gv;
2872:   IS       isx, isv;
2873:   PetscInt dim;
2874:   AppCtx  *ctx;

2876:   PetscFunctionBeginUser;
2877:   PetscCall(TSGetDM(ts, &sw));
2878:   PetscCall(DMGetApplicationContext(sw, &ctx));
2879:   PetscCall(DMGetDimension(sw, &dim));
2880:   if (useInitial) {
2881:     PetscReal v0[2] = {1., 0.};
2882:     if (ctx->perturbed_weights) {
2883:       PetscCall(InitializeParticles_PerturbedWeights(sw, ctx));
2884:     } else {
2885:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
2886:       PetscCall(DMSwarmInitializeCoordinates(sw));
2887:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
2888:     }
2889:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2890:     PetscCall(DMSwarmTSRedistribute(ts));
2891:   }
2892:   PetscCall(DMSetUp(sw));
2893:   PetscCall(TSGetSolution(ts, &u));
2894:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2895:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2896:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2897:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2898:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
2899:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
2900:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2901:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2902:   PetscFunctionReturn(PETSC_SUCCESS);
2903: }

2905: static PetscErrorCode InitializeSolve(TS ts, Vec u)
2906: {
2907:   PetscFunctionBegin;
2908:   PetscCall(TSSetSolution(ts, u));
2909:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
2910:   PetscFunctionReturn(PETSC_SUCCESS);
2911: }

2913: static PetscErrorCode MigrateParticles(TS ts)
2914: {
2915:   DM               sw, cdm;
2916:   const PetscReal *L;
2917:   AppCtx          *ctx;

2919:   PetscFunctionBeginUser;
2920:   PetscCall(TSGetDM(ts, &sw));
2921:   PetscCall(DMGetApplicationContext(sw, &ctx));
2922:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
2923:   {
2924:     Vec        u, gc, gv, position, momentum;
2925:     IS         isx, isv;
2926:     PetscReal *pos, *mom;

2928:     PetscCall(TSGetSolution(ts, &u));
2929:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2930:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2931:     PetscCall(VecGetSubVector(u, isx, &position));
2932:     PetscCall(VecGetSubVector(u, isv, &momentum));
2933:     PetscCall(VecGetArray(position, &pos));
2934:     PetscCall(VecGetArray(momentum, &mom));
2935:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2936:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2937:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
2938:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

2940:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
2941:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
2942:     PetscCheck(L, PetscObjectComm((PetscObject)cdm), PETSC_ERR_ARG_WRONG, "Mesh must be periodic");
2943:     if ((L[0] || L[1]) >= 0.) {
2944:       PetscReal *x, *v, upper[3], lower[3];
2945:       PetscInt   Np, dim;

2947:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
2948:       PetscCall(DMGetDimension(cdm, &dim));
2949:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
2950:       PetscCall(VecGetArray(gc, &x));
2951:       PetscCall(VecGetArray(gv, &v));
2952:       for (PetscInt p = 0; p < Np; ++p) {
2953:         for (PetscInt d = 0; d < dim; ++d) {
2954:           if (pos[p * dim + d] < lower[d]) {
2955:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
2956:           } else if (pos[p * dim + d] > upper[d]) {
2957:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
2958:           } else {
2959:             x[p * dim + d] = pos[p * dim + d];
2960:           }
2961:           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]);
2962:           v[p * dim + d] = mom[p * dim + d];
2963:         }
2964:       }
2965:       PetscCall(VecRestoreArray(gc, &x));
2966:       PetscCall(VecRestoreArray(gv, &v));
2967:     }
2968:     PetscCall(VecRestoreArray(position, &pos));
2969:     PetscCall(VecRestoreArray(momentum, &mom));
2970:     PetscCall(VecRestoreSubVector(u, isx, &position));
2971:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2972:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2973:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2974:   }
2975:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2976:   PetscInt step;

2978:   PetscCall(TSGetStepNumber(ts, &step));
2979:   if (!(step % ctx->remapFreq)) {
2980:     // Monitor electric field before we destroy it
2981:     PetscReal ptime;
2982:     PetscInt  step;

2984:     PetscCall(TSGetStepNumber(ts, &step));
2985:     PetscCall(TSGetTime(ts, &ptime));
2986:     if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx));
2987:     if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx));
2988:     PetscCall(DMSwarmRemap(sw));
2989:     ctx->validE = PETSC_FALSE;
2990:   }
2991:   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
2992:   PetscCall(DMSwarmTSRedistribute(ts));
2993:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2994:   PetscFunctionReturn(PETSC_SUCCESS);
2995: }

2997: int main(int argc, char **argv)
2998: {
2999:   DM        dm, sw;
3000:   TS        ts;
3001:   Vec       u;
3002:   PetscReal dt;
3003:   PetscInt  maxn;
3004:   AppCtx    ctx;

3006:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
3007:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &ctx));
3008:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &ctx.bag));
3009:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &ctx, &dm));
3010:   PetscCall(CreatePoisson(dm, &ctx));
3011:   PetscCall(CreateMomentFields(dm, &ctx));
3012:   PetscCall(CreateSwarm(dm, &ctx, &sw));
3013:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &ctx));
3014:   PetscCall(InitializeConstants(sw, &ctx));
3015:   PetscCall(DMSetApplicationContext(sw, &ctx));

3017:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
3018:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
3019:   PetscCall(TSSetDM(ts, sw));
3020:   PetscCall(TSSetMaxTime(ts, 0.1));
3021:   PetscCall(TSSetTimeStep(ts, 0.00001));
3022:   PetscCall(TSSetMaxSteps(ts, 100));
3023:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

3025:   if (ctx.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &ctx, NULL));
3026:   if (ctx.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &ctx, NULL));
3027:   if (ctx.moment_field_monitor) PetscCall(TSMonitorSet(ts, MonitorMomentFields, &ctx, NULL));
3028:   if (ctx.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &ctx, NULL));
3029:   if (ctx.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &ctx, NULL));
3030:   if (ctx.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &ctx, NULL));
3031:   if (ctx.velocity_monitor >= 0) PetscCall(TSMonitorSet(ts, MonitorVelocity, &ctx, NULL));

3033:   PetscCall(TSSetFromOptions(ts));
3034:   PetscCall(TSGetTimeStep(ts, &dt));
3035:   PetscCall(TSGetMaxSteps(ts, &maxn));
3036:   ctx.steps    = maxn;
3037:   ctx.stepSize = dt;
3038:   PetscCall(SetupContext(dm, sw, &ctx));
3039:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
3040:   PetscCall(TSSetPostStep(ts, MigrateParticles));
3041:   PetscCall(CreateSolution(ts));
3042:   PetscCall(TSGetSolution(ts, &u));
3043:   PetscCall(TSComputeInitialCondition(ts, u));
3044:   PetscCall(CheckNonNegativeWeights(sw, &ctx));
3045:   PetscCall(TSSolve(ts, NULL));

3047:   if (ctx.checkLandau) {
3048:     // We should get a lookup table based on charge density and \hat k
3049:     const PetscReal gammaEx = -0.15336;
3050:     const PetscReal omegaEx = 1.4156;
3051:     const PetscReal tol     = 1e-2;

3053:     PetscCheck(PetscAbsReal((ctx.gamma - gammaEx) / gammaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau gamma %g != %g", ctx.gamma, gammaEx);
3054:     PetscCheck(PetscAbsReal((ctx.omega - omegaEx) / omegaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau omega %g != %g", ctx.omega, omegaEx);
3055:   }

3057:   PetscCall(SNESDestroy(&ctx.snes));
3058:   PetscCall(DMDestroy(&ctx.dmN));
3059:   PetscCall(ISDestroy(&ctx.isN));
3060:   PetscCall(MatDestroy(&ctx.MN));
3061:   PetscCall(DMDestroy(&ctx.dmP));
3062:   PetscCall(ISDestroy(&ctx.isP));
3063:   PetscCall(MatDestroy(&ctx.MP));
3064:   PetscCall(DMDestroy(&ctx.dmE));
3065:   PetscCall(ISDestroy(&ctx.isE));
3066:   PetscCall(MatDestroy(&ctx.ME));
3067:   PetscCall(DMDestroy(&ctx.dmMom));
3068:   PetscCall(DMDestroy(&ctx.dmPot));
3069:   PetscCall(ISDestroy(&ctx.isPot));
3070:   PetscCall(MatDestroy(&ctx.M));
3071:   PetscCall(PetscFEGeomDestroy(&ctx.fegeom));
3072:   PetscCall(TSDestroy(&ts));
3073:   PetscCall(DMDestroy(&sw));
3074:   PetscCall(DMDestroy(&dm));
3075:   PetscCall(DestroyContext(&ctx));
3076:   PetscCall(PetscFinalize());
3077:   return 0;
3078: }

3080: /*TEST

3082:   build:
3083:     requires: !complex double

3085:   # This tests that we can compute the correct decay rate and frequency
3086:   #   For gold runs, use -dm_plex_box_faces 160 -vdm_plex_box_faces 450 -remap_dm_plex_box_faces 80,150 -ts_max_steps 1000
3087:   #                      -remap_freq 100 -emax_start_step 50 -emax_solve_step 100
3088:   testset:
3089:     args: -cosine_coefficients 0.01 -charges -1. -perturbed_weights -total_weight 1. \
3090:           -dm_plex_dim 1 -dm_plex_box_faces 80 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 \
3091:             -dm_plex_box_bd periodic -dm_plex_hash_location \
3092:           -vdm_plex_dim 1 -vdm_plex_box_faces 220 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 \
3093:             -vpetscspace_degree 2 -vdm_plex_hash_location \
3094:           -remap_freq 1 -dm_swarm_remap_type pfak -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 \
3095:             -remap_dm_plex_box_faces 40,110 -remap_dm_plex_box_bd periodic,none \
3096:             -remap_dm_plex_box_lower 0.,-6. -remap_dm_plex_box_upper 12.5664,6. \
3097:             -remap_petscspace_degree 1 -remap_dm_plex_hash_location \
3098:             -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu \
3099:           -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged \
3100:             -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu \
3101:           -ts_time_step 0.03 -ts_max_steps 2 -ts_max_time 100 \
3102:           -emax_tao_type brgn -emax_tao_max_it 100 -emax_tao_brgn_regularization_type l2pure \
3103:             -emax_tao_brgn_regularizer_weight 1e-5 -emax_tao_brgn_subsolver_tao_bnk_ksp_rtol 1e-12 \
3104:             -emax_start_step 1 -emax_solve_step 1 \
3105:           -output_step 1 -efield_monitor quiet

3107:     test:
3108:       suffix: landau_damping_1d_bs
3109:       args: -ts_type basicsymplectic -ts_basicsymplectic_type 1

3111:     test:
3112:       suffix: landau_damping_1d_dg
3113:       args: -ts_type discgrad -ts_discgrad_type average -snes_type qn

3115: TEST*/