xref: /petsc/src/ts/tutorials/hamiltonian/ex4.c (revision 340b794c6bb4bf4f05a91c896fd500c569804341)
1 static char help[] = "Two-level system for Landau Damping using Vlasov-Poisson equations\n";
2 
3 /*
4   Moment Equations:
5 
6     We will discretize the moment equations using finite elements, and we will project the moments into the finite element space We will use the PFAK method, which guarantees that our FE approximation is weakly equivalent to the true moment. The first moment, number density, is given by
7 
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
13 
14     where
15 
16       (M_p){ip} = \phi_i(x_p)
17 
18     which is just a scaled version of the charge density. The second moment, momentum density, is given by
19 
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
23 
24     And finally the third moment, pressure, is given by
25 
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
31 
32     Here we need all FEM basis functions \phi_i that see that particle p.
33 
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
36 
37   Swarm CellDMs
38   =============
39   Name: "space"
40   Fields: DMSwarmPICField_coor, "velocity"
41   Coordinates: DMSwarmPICField_coor
42 
43   Name: "velocity"
44   Fields: "w_q"
45   Coordinates: "velocity"
46 
47   Name: "moments"
48   Fields: "w_q"
49   Coordinates: DMSwarmPICField_coor
50 
51   Name: "moment fields"
52   Fields: "velocity"
53   Coordinates: DMSwarmPICField_coor
54 
55   To visualize the maximum electric field use
56 
57     -efield_monitor
58 
59   To monitor velocity moments of the distribution use
60 
61     -ptof_pc_type lu -moments_monitor
62 
63   To monitor the particle positions in phase space use
64 
65     -positions_monitor
66 
67   To monitor the charge density, E field, and potential use
68 
69     -poisson_monitor
70 
71   To monitor the remapping field use
72 
73     -remap_uf_view draw
74 
75   To visualize the swarm distribution use
76 
77     -ts_monitor_hg_swarm
78 
79   To visualize the particles, we can use
80 
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> /* For interpolation */
92 #include <petsc/private/dmswarmimpl.h> /* For swarm debugging */
93 #include "petscdm.h"
94 #include "petscdmlabel.h"
95 
96 PETSC_EXTERN PetscErrorCode stream(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
97 PETSC_EXTERN PetscErrorCode line(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
98 
99 const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
100 typedef enum {
101   EMPRIMAL,
102   EM_MIXED,
103   EM_COULOMB,
104   EMNONE
105 } EMType;
106 
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;
119 
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};
126 
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;
140 
141 typedef struct {
142   PetscBag     bag;                  // Problem parameters
143   PetscBool    error;                // Flag for printing the error
144   PetscInt     remapFreq;            // Number of timesteps between remapping
145   EMonitorType efield_monitor;       // Flag to show electric field monitor
146   PetscBool    moment_monitor;       // Flag to show distribution moment monitor
147   PetscBool    moment_field_monitor; // Flag to show moment field monitor
148   PetscBool    positions_monitor;    // Flag to show particle positins at each time step
149   PetscBool    poisson_monitor;      // Flag to display charge, E field, and potential at each solve
150   PetscBool    initial_monitor;      // Flag to monitor the initial conditions
151   PetscInt     velocity_monitor;     // Cell to monitor the velocity distribution for
152   PetscBool    perturbed_weights;    // Uniformly sample x,v space with gaussian weights
153   PetscInt     ostep;                // Print the energy at each ostep time steps
154   PetscInt     numParticles;
155   PetscReal    timeScale;              /* Nondimensionalizing time scale */
156   PetscReal    charges[2];             /* The charges of each species */
157   PetscReal    masses[2];              /* The masses of each species */
158   PetscReal    thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
159   PetscReal    cosine_coefficients[2]; /*(alpha, k)*/
160   PetscReal    totalWeight;
161   PetscReal    stepSize;
162   PetscInt     steps;
163   PetscReal    initVel;
164   EMType       em;           // Type of electrostatic model
165   SNES         snes;         // EM solver
166   DM           dmMom;        // The DM for moment fields
167   DM           dmN;          // The DM for number density fields
168   IS           isN;          // The IS mapping dmN into dmMom
169   Mat          MN;           // The finite element mass matrix for number density
170   DM           dmP;          // The DM for momentum density fields
171   IS           isP;          // The IS mapping dmP into dmMom
172   Mat          MP;           // The finite element mass matrix for momentum density
173   DM           dmE;          // The DM for energy density (pressure) fields
174   IS           isE;          // The IS mapping dmE into dmMom
175   Mat          ME;           // The finite element mass matrix for energy density (pressure)
176   DM           dmPot;        // The DM for potential
177   Mat          fftPot;       // Fourier Transform operator for the potential
178   Vec          fftX, fftY;   //   FFT vectors with phases added (complex parts)
179   IS           fftReal;      //   The indices for real parts
180   IS           isPot;        // The IS for potential, or NULL in primal
181   Mat          M;            // The finite element mass matrix for potential
182   PetscFEGeom *fegeom;       // Geometric information for the DM cells
183   PetscDrawHG  drawhgic_x;   // Histogram of the particle weight in each X cell
184   PetscDrawHG  drawhgic_v;   // Histogram of the particle weight in each X cell
185   PetscDrawHG  drawhgcell_v; // Histogram of the particle weight in a given cell
186   PetscBool    validE;       // Flag to indicate E-field in swarm is valid
187   PetscReal    drawlgEmin;   // The minimum lg(E) to plot
188   PetscDrawLG  drawlgE;      // Logarithm of maximum electric field
189   PetscDrawSP  drawspE;      // Electric field at particle positions
190   PetscDrawSP  drawspX;      // Particle positions
191   PetscViewer  viewerRho;    // Charge density viewer
192   PetscViewer  viewerRhoHat; // Charge density Fourier Transform viewer
193   PetscViewer  viewerPhi;    // Potential viewer
194   PetscViewer  viewerN;      // Number density viewer
195   PetscViewer  viewerP;      // Momentum density viewer
196   PetscViewer  viewerE;      // Energy density (pressure) viewer
197   DM           swarm;        // The particle swarm
198   PetscRandom  random;       // Used for particle perturbations
199   PetscBool    twostream;    // Flag for activating 2-stream setup
200   PetscBool    checkweights; // Check weight normalization
201   PetscInt     checkVRes;    // Flag to check/output velocity residuals for nightly tests
202   PetscBool    checkLandau;  // Check the Landau damping result
203   PetscReal    gamma;        // The damping rate for Landau damping
204   PetscReal    omega;        // The perturbed oscillation frequency for Landau damping
205 
206   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
207 } AppCtx;
208 
209 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
210 {
211   PetscFunctionBeginUser;
212   PetscInt d                      = 2;
213   PetscInt maxSpecies             = 2;
214   options->error                  = PETSC_FALSE;
215   options->remapFreq              = 1;
216   options->efield_monitor         = E_MONITOR_NONE;
217   options->moment_monitor         = PETSC_FALSE;
218   options->moment_field_monitor   = PETSC_FALSE;
219   options->initial_monitor        = PETSC_FALSE;
220   options->perturbed_weights      = PETSC_FALSE;
221   options->poisson_monitor        = PETSC_FALSE;
222   options->positions_monitor      = PETSC_FALSE;
223   options->velocity_monitor       = -1;
224   options->ostep                  = 100;
225   options->timeScale              = 2.0e-14;
226   options->charges[0]             = -1.0;
227   options->charges[1]             = 1.0;
228   options->masses[0]              = 1.0;
229   options->masses[1]              = 1000.0;
230   options->thermal_energy[0]      = 1.0;
231   options->thermal_energy[1]      = 1.0;
232   options->cosine_coefficients[0] = 0.01;
233   options->cosine_coefficients[1] = 0.5;
234   options->initVel                = 1;
235   options->totalWeight            = 1.0;
236   options->drawhgic_x             = NULL;
237   options->drawhgic_v             = NULL;
238   options->drawhgcell_v           = NULL;
239   options->drawlgEmin             = -6;
240   options->drawlgE                = NULL;
241   options->drawspE                = NULL;
242   options->drawspX                = NULL;
243   options->viewerRho              = NULL;
244   options->viewerRhoHat           = NULL;
245   options->viewerPhi              = NULL;
246   options->viewerN                = NULL;
247   options->viewerP                = NULL;
248   options->viewerE                = NULL;
249   options->em                     = EM_COULOMB;
250   options->snes                   = NULL;
251   options->dmMom                  = NULL;
252   options->dmN                    = NULL;
253   options->MN                     = NULL;
254   options->dmP                    = NULL;
255   options->MP                     = NULL;
256   options->dmE                    = NULL;
257   options->ME                     = NULL;
258   options->dmPot                  = NULL;
259   options->fftPot                 = NULL;
260   options->fftX                   = NULL;
261   options->fftY                   = NULL;
262   options->fftReal                = NULL;
263   options->isPot                  = NULL;
264   options->M                      = NULL;
265   options->numParticles           = 32768;
266   options->twostream              = PETSC_FALSE;
267   options->checkweights           = PETSC_FALSE;
268   options->checkVRes              = 0;
269   options->checkLandau            = PETSC_FALSE;
270 
271   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
272   PetscCall(PetscOptionsBool("-error", "Flag to print the error", __FILE__, options->error, &options->error, NULL));
273   PetscCall(PetscOptionsInt("-remap_freq", "Number", __FILE__, options->remapFreq, &options->remapFreq, NULL));
274   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));
275   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", __FILE__, options->drawlgEmin, &options->drawlgEmin, NULL));
276   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", __FILE__, options->moment_monitor, &options->moment_monitor, NULL));
277   PetscCall(PetscOptionsBool("-moment_field_monitor", "Flag to show moment fields", __FILE__, options->moment_field_monitor, &options->moment_field_monitor, NULL));
278   PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", __FILE__, options->initial_monitor, &options->initial_monitor, NULL));
279   PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", __FILE__, options->positions_monitor, &options->positions_monitor, NULL));
280   PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", __FILE__, options->poisson_monitor, &options->poisson_monitor, NULL));
281   PetscCall(PetscOptionsInt("-velocity_monitor", "Cell to show velocity histograms", __FILE__, options->velocity_monitor, &options->velocity_monitor, NULL));
282   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", __FILE__, options->twostream, &options->twostream, NULL));
283   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", __FILE__, options->perturbed_weights, &options->perturbed_weights, NULL));
284   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", __FILE__, options->checkweights, &options->checkweights, NULL));
285   PetscCall(PetscOptionsBool("-check_landau", "Check the decay from Landau damping", __FILE__, options->checkLandau, &options->checkLandau, NULL));
286   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", __FILE__, options->ostep, &options->ostep, NULL));
287   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", __FILE__, options->timeScale, &options->timeScale, NULL));
288   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", __FILE__, options->checkVRes, &options->checkVRes, NULL));
289   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", __FILE__, options->initVel, &options->initVel, NULL));
290   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", __FILE__, options->totalWeight, &options->totalWeight, NULL));
291   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", __FILE__, options->cosine_coefficients, &d, NULL));
292   PetscCall(PetscOptionsRealArray("-charges", "Species charges", __FILE__, options->charges, &maxSpecies, NULL));
293   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", __FILE__, EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
294   PetscOptionsEnd();
295 
296   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
297   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
298   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
299   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
300   PetscFunctionReturn(PETSC_SUCCESS);
301 }
302 
303 static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
304 {
305   MPI_Comm comm;
306 
307   PetscFunctionBeginUser;
308   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
309   if (user->efield_monitor) {
310     PetscDraw     draw;
311     PetscDrawAxis axis;
312 
313     if (user->efield_monitor == E_MONITOR_FULL) {
314       PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
315       PetscCall(PetscDrawSetSave(draw, "ex2_Efield"));
316       PetscCall(PetscDrawSetFromOptions(draw));
317     } else {
318       PetscCall(PetscDrawOpenNull(comm, &draw));
319     }
320     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
321     PetscCall(PetscDrawDestroy(&draw));
322     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
323     PetscCall(PetscDrawAxisSetLabels(axis, "Max Electric Field", "time", "E_max"));
324     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
325   }
326 
327   if (user->initial_monitor) {
328     PetscDraw     drawic_x, drawic_v;
329     PetscDrawAxis axis1, axis2;
330     PetscReal     dmboxlower[2], dmboxupper[2];
331     PetscInt      dim, cStart, cEnd;
332 
333     PetscCall(DMGetDimension(sw, &dim));
334     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
335     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
336 
337     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &drawic_x));
338     PetscCall(PetscDrawSetSave(drawic_x, "ex2_ic_x"));
339     PetscCall(PetscDrawSetFromOptions(drawic_x));
340     PetscCall(PetscDrawHGCreate(drawic_x, (int)dim, &user->drawhgic_x));
341     PetscCall(PetscDrawHGCalcStats(user->drawhgic_x, PETSC_TRUE));
342     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
343     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart)));
344     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "weight"));
345     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 0));
346     PetscCall(PetscDrawDestroy(&drawic_x));
347 
348     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &drawic_v));
349     PetscCall(PetscDrawSetSave(drawic_v, "ex9_ic_v"));
350     PetscCall(PetscDrawSetFromOptions(drawic_v));
351     PetscCall(PetscDrawHGCreate(drawic_v, (int)dim, &user->drawhgic_v));
352     PetscCall(PetscDrawHGCalcStats(user->drawhgic_v, PETSC_TRUE));
353     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
354     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 21));
355     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "weight"));
356     PetscCall(PetscDrawAxisSetLimits(axis2, -6, 6, 0, 0));
357     PetscCall(PetscDrawDestroy(&drawic_v));
358   }
359 
360   if (user->velocity_monitor >= 0) {
361     DM            vdm;
362     DMSwarmCellDM celldm;
363     PetscDraw     drawcell_v;
364     PetscDrawAxis axis;
365     PetscReal     dmboxlower[2], dmboxupper[2];
366     PetscInt      dim;
367     char          title[PETSC_MAX_PATH_LEN];
368 
369     PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
370     PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
371     PetscCall(DMGetDimension(vdm, &dim));
372     PetscCall(DMGetBoundingBox(vdm, dmboxlower, dmboxupper));
373 
374     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Cell %" PetscInt_FMT ": Velocity Distribution", user->velocity_monitor));
375     PetscCall(PetscDrawCreate(comm, NULL, title, 400, 300, 400, 300, &drawcell_v));
376     PetscCall(PetscDrawSetSave(drawcell_v, "ex2_cell_v"));
377     PetscCall(PetscDrawSetFromOptions(drawcell_v));
378     PetscCall(PetscDrawHGCreate(drawcell_v, (int)dim, &user->drawhgcell_v));
379     PetscCall(PetscDrawHGCalcStats(user->drawhgcell_v, PETSC_TRUE));
380     PetscCall(PetscDrawHGGetAxis(user->drawhgcell_v, &axis));
381     PetscCall(PetscDrawHGSetNumberBins(user->drawhgcell_v, 21));
382     PetscCall(PetscDrawAxisSetLabels(axis, "V_x Distribution", "V", "weight"));
383     PetscCall(PetscDrawAxisSetLimits(axis, dmboxlower[0], dmboxupper[0], 0, 0));
384     PetscCall(PetscDrawDestroy(&drawcell_v));
385   }
386 
387   if (user->positions_monitor) {
388     PetscDraw     draw;
389     PetscDrawAxis axis;
390 
391     PetscCall(PetscDrawCreate(comm, NULL, "Particle Position", 0, 0, 400, 300, &draw));
392     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
393     PetscCall(PetscDrawSetFromOptions(draw));
394     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspX));
395     PetscCall(PetscDrawDestroy(&draw));
396     PetscCall(PetscDrawSPSetDimension(user->drawspX, 1));
397     PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis));
398     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
399     PetscCall(PetscDrawSPReset(user->drawspX));
400   }
401   if (user->poisson_monitor) {
402     Vec           rho, rhohat, phi;
403     PetscDraw     draw;
404     PetscDrawAxis axis;
405 
406     PetscCall(PetscDrawCreate(comm, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
407     PetscCall(PetscDrawSetFromOptions(draw));
408     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
409     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspE));
410     PetscCall(PetscDrawDestroy(&draw));
411     PetscCall(PetscDrawSPSetDimension(user->drawspE, 1));
412     PetscCall(PetscDrawSPGetAxis(user->drawspE, &axis));
413     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
414     PetscCall(PetscDrawSPReset(user->drawspE));
415 
416     PetscCall(PetscViewerDrawOpen(comm, NULL, "Charge Density", 0, 0, 400, 300, &user->viewerRho));
417     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRho, "rho_"));
418     PetscCall(PetscViewerDrawGetDraw(user->viewerRho, 0, &draw));
419     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
420     PetscCall(PetscViewerSetFromOptions(user->viewerRho));
421     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
422     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
423     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
424 
425     PetscInt dim, N;
426 
427     PetscCall(DMGetDimension(user->dmPot, &dim));
428     if (dim == 1) {
429       PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
430       PetscCall(VecGetSize(rhohat, &N));
431       PetscCall(MatCreateFFT(comm, dim, &N, MATFFTW, &user->fftPot));
432       PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
433       PetscCall(MatCreateVecs(user->fftPot, &user->fftX, &user->fftY));
434       PetscCall(ISCreateStride(PETSC_COMM_SELF, N, 0, 1, &user->fftReal));
435     }
436 
437     PetscCall(PetscViewerDrawOpen(comm, NULL, "rhohat: Charge Density FT", 0, 0, 400, 300, &user->viewerRhoHat));
438     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRhoHat, "rhohat_"));
439     PetscCall(PetscViewerDrawGetDraw(user->viewerRhoHat, 0, &draw));
440     PetscCall(PetscDrawSetSave(draw, "ex9_rho_ft"));
441     PetscCall(PetscViewerSetFromOptions(user->viewerRhoHat));
442     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
443     PetscCall(PetscObjectSetName((PetscObject)rhohat, "charge_density_ft"));
444     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
445 
446     PetscCall(PetscViewerDrawOpen(comm, NULL, "Potential", 400, 0, 400, 300, &user->viewerPhi));
447     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPhi, "phi_"));
448     PetscCall(PetscViewerDrawGetDraw(user->viewerPhi, 0, &draw));
449     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
450     PetscCall(PetscViewerSetFromOptions(user->viewerPhi));
451     PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
452     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
453     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
454   }
455   if (user->moment_field_monitor) {
456     Vec       n, p, e;
457     PetscDraw draw;
458 
459     PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density", 0, 800, 400, 300, &user->viewerN));
460     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerN, "n_"));
461     PetscCall(PetscViewerDrawGetDraw(user->viewerN, 0, &draw));
462     PetscCall(PetscDrawSetSave(draw, "ex9_n_spatial"));
463     PetscCall(PetscViewerSetFromOptions(user->viewerN));
464     PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n));
465     PetscCall(PetscObjectSetName((PetscObject)n, "Number Density"));
466     PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n));
467 
468     PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density", 400, 800, 400, 300, &user->viewerP));
469     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerP, "p_"));
470     PetscCall(PetscViewerDrawGetDraw(user->viewerP, 0, &draw));
471     PetscCall(PetscDrawSetSave(draw, "ex9_p_spatial"));
472     PetscCall(PetscViewerSetFromOptions(user->viewerP));
473     PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p));
474     PetscCall(PetscObjectSetName((PetscObject)p, "Momentum Density"));
475     PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p));
476 
477     PetscCall(PetscViewerDrawOpen(comm, NULL, "Emergy Density (Pressure)", 800, 800, 400, 300, &user->viewerE));
478     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerE, "e_"));
479     PetscCall(PetscViewerDrawGetDraw(user->viewerE, 0, &draw));
480     PetscCall(PetscDrawSetSave(draw, "ex9_e_spatial"));
481     PetscCall(PetscViewerSetFromOptions(user->viewerE));
482     PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e));
483     PetscCall(PetscObjectSetName((PetscObject)e, "Energy Density (Pressure)"));
484     PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e));
485   }
486   PetscFunctionReturn(PETSC_SUCCESS);
487 }
488 
489 static PetscErrorCode DestroyContext(AppCtx *user)
490 {
491   PetscFunctionBeginUser;
492   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
493   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
494   PetscCall(PetscDrawHGDestroy(&user->drawhgcell_v));
495 
496   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
497   PetscCall(PetscDrawSPDestroy(&user->drawspE));
498   PetscCall(PetscDrawSPDestroy(&user->drawspX));
499   PetscCall(PetscViewerDestroy(&user->viewerRho));
500   PetscCall(PetscViewerDestroy(&user->viewerRhoHat));
501   PetscCall(MatDestroy(&user->fftPot));
502   PetscCall(VecDestroy(&user->fftX));
503   PetscCall(VecDestroy(&user->fftY));
504   PetscCall(ISDestroy(&user->fftReal));
505   PetscCall(PetscViewerDestroy(&user->viewerPhi));
506   PetscCall(PetscViewerDestroy(&user->viewerN));
507   PetscCall(PetscViewerDestroy(&user->viewerP));
508   PetscCall(PetscViewerDestroy(&user->viewerE));
509 
510   PetscCall(PetscBagDestroy(&user->bag));
511   PetscFunctionReturn(PETSC_SUCCESS);
512 }
513 
514 static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
515 {
516   const PetscScalar *w;
517   PetscInt           Np;
518 
519   PetscFunctionBeginUser;
520   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
521   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
522   PetscCall(DMSwarmGetLocalSize(sw, &Np));
523   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]);
524   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
525   PetscFunctionReturn(PETSC_SUCCESS);
526 }
527 
528 static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *user)
529 {
530   DMSwarmCellDM celldm;
531   DM            vdm;
532   Vec           u[1];
533   const char   *fields[1] = {"w_q"};
534 
535   PetscFunctionBegin;
536   PetscCall(DMSwarmSetCellDMActive(sw, "velocity"));
537   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
538   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
539   PetscCall(DMGetGlobalVector(vdm, &u[0]));
540   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
541   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
542   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
543   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
544   PetscFunctionReturn(PETSC_SUCCESS);
545 }
546 
547 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[])
548 {
549   f0[0] = 0.;
550   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d * dim + d]);
551 }
552 
553 typedef struct {
554   PetscInt         s;    // Starting sample (we ignore some in the beginning)
555   PetscInt         e;    // Ending sample
556   const PetscReal *t;    // Time for each sample
557   const PetscReal *Emax; // Emax for each sample
558 } EmaxCtx;
559 
560 // Our model is E_max(t) = C e^{-gamma t} |cos(omega t - phi)|
561 static PetscErrorCode ComputeEmaxResidual(Tao tao, Vec x, Vec res, void *user)
562 {
563   EmaxCtx           *ctx = (EmaxCtx *)user;
564   const PetscScalar *a;
565   PetscScalar       *F;
566   PetscReal          C, gamma, omega, phi;
567 
568   PetscFunctionBegin;
569   PetscCall(VecGetArrayRead(x, &a));
570   PetscCall(VecGetArray(res, &F));
571   C     = PetscRealPart(a[0]);
572   gamma = PetscRealPart(a[1]);
573   omega = PetscRealPart(a[2]);
574   phi   = PetscRealPart(a[3]);
575   PetscCall(VecRestoreArrayRead(x, &a));
576   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));
577   PetscCall(VecRestoreArray(res, &F));
578   PetscFunctionReturn(PETSC_SUCCESS);
579 }
580 
581 // The Jacobian of the residual J = dr(x)/dx
582 static PetscErrorCode ComputeEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *user)
583 {
584   EmaxCtx           *ctx = (EmaxCtx *)user;
585   const PetscScalar *a;
586   PetscScalar       *jac;
587   PetscReal          C, gamma, omega, phi;
588   const PetscInt     n = ctx->e - ctx->s;
589 
590   PetscFunctionBegin;
591   PetscCall(VecGetArrayRead(x, &a));
592   C     = PetscRealPart(a[0]);
593   gamma = PetscRealPart(a[1]);
594   omega = PetscRealPart(a[2]);
595   phi   = PetscRealPart(a[3]);
596   PetscCall(VecRestoreArrayRead(x, &a));
597   PetscCall(MatDenseGetArray(J, &jac));
598   for (PetscInt i = 0; i < n; ++i) {
599     const PetscInt k = i + ctx->s;
600 
601     jac[i * 4 + 0] = -PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
602     jac[i * 4 + 1] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
603     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);
604     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);
605   }
606   PetscCall(MatDenseRestoreArray(J, &jac));
607   PetscFunctionReturn(PETSC_SUCCESS);
608 }
609 
610 // Our model is log_10 E_max(t) = log_10 C - gamma t log_10 e + log_10 |cos(omega t - phi)|
611 static PetscErrorCode ComputeLogEmaxResidual(Tao tao, Vec x, Vec res, void *user)
612 {
613   EmaxCtx           *ctx = (EmaxCtx *)user;
614   const PetscScalar *a;
615   PetscScalar       *F;
616   PetscReal          C, gamma, omega, phi;
617 
618   PetscFunctionBegin;
619   PetscCall(VecGetArrayRead(x, &a));
620   PetscCall(VecGetArray(res, &F));
621   C     = PetscRealPart(a[0]);
622   gamma = PetscRealPart(a[1]);
623   omega = PetscRealPart(a[2]);
624   phi   = PetscRealPart(a[3]);
625   PetscCall(VecRestoreArrayRead(x, &a));
626   for (PetscInt i = ctx->s; i < ctx->e; ++i) {
627     if (C < 0) {
628       F[i - ctx->s] = 1e10;
629       continue;
630     }
631     F[i - ctx->s] = ctx->Emax[i] - (PetscLog10Real(C) - gamma * ctx->t[i] * PetscLog10Real(PETSC_E) + PetscLog10Real(PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi))));
632   }
633   PetscCall(VecRestoreArray(res, &F));
634   PetscFunctionReturn(PETSC_SUCCESS);
635 }
636 
637 // The Jacobian of the residual J = dr(x)/dx
638 static PetscErrorCode ComputeLogEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *user)
639 {
640   EmaxCtx           *ctx = (EmaxCtx *)user;
641   const PetscScalar *a;
642   PetscScalar       *jac;
643   PetscReal          C, omega, phi;
644   const PetscInt     n = ctx->e - ctx->s;
645 
646   PetscFunctionBegin;
647   PetscCall(VecGetArrayRead(x, &a));
648   C     = PetscRealPart(a[0]);
649   omega = PetscRealPart(a[2]);
650   phi   = PetscRealPart(a[3]);
651   PetscCall(VecRestoreArrayRead(x, &a));
652   PetscCall(MatDenseGetArray(J, &jac));
653   for (PetscInt i = 0; i < n; ++i) {
654     const PetscInt k = i + ctx->s;
655 
656     jac[0 * n + i] = -1. / (PetscLog10Real(PETSC_E) * C);
657     jac[1 * n + i] = ctx->t[k] * PetscLog10Real(PETSC_E);
658     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)));
659     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)));
660   }
661   PetscCall(MatDenseRestoreArray(J, &jac));
662   PetscCall(MatViewFromOptions(J, NULL, "-emax_jac_view"));
663   PetscFunctionReturn(PETSC_SUCCESS);
664 }
665 
666 static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
667 {
668   AppCtx     *user = (AppCtx *)ctx;
669   DM          sw;
670   PetscScalar intESq;
671   PetscReal  *E, *x, *weight;
672   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
673   PetscReal   pmoments[4]; /* \int f, \int v f, \int v^2 f */
674   PetscInt   *species, dim, Np, gNp;
675   MPI_Comm    comm;
676   PetscMPIInt rank;
677 
678   PetscFunctionBeginUser;
679   if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
680   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
681   PetscCallMPI(MPI_Comm_rank(comm, &rank));
682   PetscCall(TSGetDM(ts, &sw));
683   PetscCall(DMGetDimension(sw, &dim));
684   PetscCall(DMSwarmGetLocalSize(sw, &Np));
685   PetscCall(DMSwarmGetSize(sw, &gNp));
686   PetscCall(DMSwarmSortGetAccess(sw));
687   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
688   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
689   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
690   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
691 
692   for (PetscInt p = 0; p < Np; ++p) {
693     for (PetscInt d = 0; d < 1; ++d) {
694       PetscReal temp = PetscAbsReal(E[p * dim + d]);
695       if (temp > Emax) Emax = temp;
696     }
697     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
698     sum += E[p * dim];
699     chargesum += user->charges[0] * weight[p];
700   }
701   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
702   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
703   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;
704 
705   PetscDS ds;
706   Vec     phi;
707 
708   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
709   PetscCall(DMGetDS(user->dmPot, &ds));
710   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
711   PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user));
712   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
713 
714   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
715   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
716   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
717   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
718   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
719   if (user->efield_monitor == E_MONITOR_FULL) {
720     PetscDraw draw;
721 
722     PetscCall(PetscDrawLGDraw(user->drawlgE));
723     PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
724     PetscCall(PetscDrawSave(draw));
725 
726     PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
727     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));
728     PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
729   }
730 
731   // Compute decay rate and frequency
732   EmaxCtx ectx;
733 
734   ectx.s = 50;
735   PetscCall(PetscDrawLGGetData(user->drawlgE, NULL, &ectx.e, &ectx.t, &ectx.Emax));
736   if (!rank && !(ectx.e % 100)) {
737     Tao          tao;
738     Mat          J;
739     Vec          x, r;
740     PetscScalar *a;
741     PetscBool    fitLog = PETSC_TRUE, debug = PETSC_FALSE;
742 
743     PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
744     PetscCall(TaoSetOptionsPrefix(tao, "emax_"));
745     PetscCall(VecCreateSeq(PETSC_COMM_SELF, 4, &x));
746     PetscCall(TaoSetSolution(tao, x));
747     PetscCall(VecCreateSeq(PETSC_COMM_SELF, ectx.e - ectx.s, &r));
748     if (fitLog) PetscCall(TaoSetResidualRoutine(tao, r, ComputeLogEmaxResidual, &ectx));
749     else PetscCall(TaoSetResidualRoutine(tao, r, ComputeEmaxResidual, &ectx));
750     PetscCall(VecDestroy(&r));
751     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, ectx.e - ectx.s, 4, NULL, &J));
752     if (fitLog) PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeLogEmaxJacobian, &ectx));
753     else PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeEmaxJacobian, &ectx));
754     PetscCall(MatDestroy(&J));
755     PetscCall(TaoSetFromOptions(tao));
756     PetscCall(VecGetArray(x, &a));
757     a[0] = 0.02;
758     a[1] = 0.15;
759     a[2] = 1.4;
760     a[3] = 0.45;
761     PetscCall(VecRestoreArray(x, &a));
762     PetscCall(TaoSolve(tao));
763     if (debug) {
764       PetscCall(PetscPrintf(PETSC_COMM_SELF, "t = ["));
765       for (PetscInt i = 0; i < ectx.e; ++i) {
766         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
767         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ectx.t[i]));
768       }
769       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
770       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax = ["));
771       for (PetscInt i = 0; i < ectx.e; ++i) {
772         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
773         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ectx.Emax[i]));
774       }
775       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
776     }
777     PetscDraw     draw;
778     PetscDrawAxis axis;
779     char          title[PETSC_MAX_PATH_LEN];
780 
781     PetscCall(VecGetArray(x, &a));
782     user->gamma = a[1];
783     user->omega = a[2];
784     if (user->efield_monitor == E_MONITOR_FULL) {
785       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax Fit: gamma %g omega %g C %g phi %g\n", a[1], a[2], a[0], a[3]));
786       PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
787       PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Max Electric Field gamma %.4g omega %.4g", a[1], a[2]));
788       PetscCall(PetscDrawSetTitle(draw, title));
789       PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
790       PetscCall(PetscDrawAxisSetLabels(axis, title, "time", "E_max"));
791     }
792     PetscCall(VecRestoreArray(x, &a));
793     PetscCall(VecDestroy(&x));
794     PetscCall(TaoDestroy(&tao));
795   }
796   PetscFunctionReturn(PETSC_SUCCESS);
797 }
798 
799 static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
800 {
801   AppCtx   *user = (AppCtx *)ctx;
802   DM        sw;
803   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */
804 
805   PetscFunctionBeginUser;
806   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
807   PetscCall(TSGetDM(ts, &sw));
808 
809   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
810   PetscCall(computeVelocityFEMMoments(sw, fmoments, user));
811 
812   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]));
813   PetscFunctionReturn(PETSC_SUCCESS);
814 }
815 
816 static void f0_pres(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[])
817 {
818   const PetscReal v = a[0];
819   const PetscReal n = u[0];
820   const PetscReal p = u[1];
821 
822   f1[0] = PetscSqr(v - p / n);
823 }
824 
825 static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
826 {
827   u[0] = 0.0;
828   return PETSC_SUCCESS;
829 }
830 
831 /*
832     M_p w_p
833       - Make M_p with "moments"
834       - Get w_p from Swarm
835     M_p v_p w_p
836       - Get v_p from Swarm
837       - pointwise multiply v_p and w_p
838     M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p
839       - ProjectField(sw, {n, p} U, {v_p} A, tmp_p)
840       - pointwise multiply tmp_p and w_p
841 
842 */
843 static PetscErrorCode ComputeMomentFields(TS ts)
844 {
845   AppCtx   *user;
846   DM        sw;
847   KSP       ksp;
848   Mat       M_p;
849   Vec       f, v, vf, vuf;
850   Vec       m, n, nRhs, p, pRhs, e, eRhs;
851   PetscReal dt;
852   PetscInt  Nts;
853 
854   PetscFunctionBegin;
855   PetscCall(TSGetStepNumber(ts, &Nts));
856   PetscCall(TSGetTimeStep(ts, &dt));
857   PetscCall(TSGetDM(ts, &sw));
858   PetscCall(DMGetApplicationContext(sw, (void *)&user));
859   PetscCall(DMSwarmSetCellDMActive(sw, "moment fields"));
860   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
861   // TODO Will have to create different M_p in higher dimensions for velocity
862   PetscCall(DMCreateMassMatrix(sw, user->dmN, &M_p));
863   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
864   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &v));
865   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
866 
867   PetscCall(DMGetGlobalVector(user->dmN, &nRhs));
868   PetscCall(PetscObjectSetName((PetscObject)nRhs, "Weak number density"));
869   PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n));
870   PetscCall(DMGetGlobalVector(user->dmP, &pRhs));
871   PetscCall(PetscObjectSetName((PetscObject)pRhs, "Weak momentum density"));
872   PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p));
873   PetscCall(DMGetGlobalVector(user->dmE, &eRhs));
874   PetscCall(PetscObjectSetName((PetscObject)eRhs, "Weak energy density (pressure)"));
875   PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e));
876 
877   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
878   PetscCall(MatViewFromOptions(user->MN, NULL, "-mn_view"));
879   PetscCall(MatViewFromOptions(user->MP, NULL, "-mp_view"));
880   PetscCall(MatViewFromOptions(user->ME, NULL, "-me_view"));
881   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
882 
883   PetscCall(VecDuplicate(f, &vf));
884   PetscCall(VecDuplicate(f, &vuf));
885 
886   PetscCall(MatMultTranspose(M_p, f, nRhs));
887 
888   PetscCall(VecPointwiseMult(vf, f, v));
889   PetscCall(MatMultTranspose(M_p, vf, pRhs));
890   PetscCall(VecDestroy(&vf));
891   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &v));
892   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
893 
894   PetscPointFn *funcs[1] = {f0_pres};
895 
896   PetscCall(DMGetGlobalVector(user->dmMom, &m));
897   PetscCall(VecISCopy(m, user->isN, SCATTER_FORWARD, n));
898   PetscCall(VecISCopy(m, user->isP, SCATTER_FORWARD, p));
899   PetscCall(DMProjectField(sw, 0., m, funcs, INSERT_VALUES, vuf));
900   PetscCall(DMRestoreGlobalVector(user->dmMom, &m));
901   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
902   PetscCall(VecPointwiseMult(vuf, f, vuf));
903   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
904   PetscCall(MatMultTranspose(M_p, vuf, eRhs));
905   PetscCall(VecDestroy(&vuf));
906 
907   PetscCall(MatDestroy(&M_p));
908 
909   PetscCall(KSPCreate(PetscObjectComm((PetscObject)sw), &ksp));
910   PetscCall(KSPSetOptionsPrefix(ksp, "mom_proj_"));
911   PetscCall(KSPSetOperators(ksp, user->MN, user->MN));
912   PetscCall(KSPSetFromOptions(ksp));
913   PetscCall(KSPSolve(ksp, nRhs, n));
914   PetscCall(KSPSetOperators(ksp, user->MP, user->MP));
915   PetscCall(KSPSetFromOptions(ksp));
916   PetscCall(KSPSolve(ksp, pRhs, p));
917   PetscCall(KSPSetOperators(ksp, user->ME, user->ME));
918   PetscCall(KSPSetFromOptions(ksp));
919   PetscCall(KSPSolve(ksp, eRhs, e));
920   PetscCall(KSPDestroy(&ksp));
921 
922   // Check moment residual
923   // TODO Fix global2local here
924   PetscSimplePointFn *exacts[3] = {zero, zero, zero};
925   Vec                 mold, m_t, locF, phi;
926   PetscReal           res[3];
927 
928   PetscCall(DMGetNamedGlobalVector(user->dmMom, "mold", &mold));
929   if (!Nts) PetscCall(VecSet(mold, 0.));
930 
931   PetscCall(DMGetGlobalVector(user->dmMom, &m));
932   PetscCall(DMGetGlobalVector(user->dmMom, &m_t));
933   PetscCall(DMGetGlobalVector(user->dmMom, &locF));
934   PetscCall(VecISCopy(m, user->isN, SCATTER_FORWARD, n));
935   PetscCall(VecISCopy(m, user->isP, SCATTER_FORWARD, p));
936   PetscCall(VecISCopy(m, user->isE, SCATTER_FORWARD, e));
937   PetscCall(VecAXPBYPCZ(m_t, -1.0, 1.0, 0.0, mold, m));
938   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
939   PetscCall(DMSetAuxiliaryVec(user->dmMom, NULL, 0, 0, phi));
940   PetscCall(DMPlexTSComputeIFunctionFEM(user->dmMom, PETSC_MIN_REAL, m, m_t, locF, user));
941   PetscCall(DMSetAuxiliaryVec(user->dmMom, NULL, 0, 0, NULL));
942   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
943   PetscCall(VecViewFromOptions(locF, NULL, "-moment_residual_view"));
944   PetscCall(VecNorm(locF, NORM_2, &res[0]));
945   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Moment Residual: %g\n", (double)res[0]));
946   PetscCall(DMComputeL2FieldDiff(user->dmMom, 0., exacts, NULL, locF, res));
947   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Moment Residuals: %g %g %g\n", (double)res[0], (double)res[1], (double)res[2]));
948   PetscCall(VecCopy(m, mold));
949   PetscCall(DMRestoreGlobalVector(user->dmMom, &locF));
950   PetscCall(DMRestoreGlobalVector(user->dmMom, &m));
951   PetscCall(DMRestoreNamedGlobalVector(user->dmMom, "mold", &mold));
952 
953   PetscCall(DMRestoreGlobalVector(user->dmN, &nRhs));
954   PetscCall(DMRestoreGlobalVector(user->dmP, &pRhs));
955   PetscCall(DMRestoreGlobalVector(user->dmE, &eRhs));
956   PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n));
957   PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p));
958   PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e));
959   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
960   PetscFunctionReturn(PETSC_SUCCESS);
961 }
962 
963 static PetscErrorCode MonitorMomentFields(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
964 {
965   AppCtx *user = (AppCtx *)ctx;
966   Vec     n, p, e;
967 
968   PetscFunctionBeginUser;
969   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
970   PetscCall(ComputeMomentFields(ts));
971 
972   PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n));
973   PetscCall(VecView(n, user->viewerN));
974   PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n));
975 
976   PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p));
977   PetscCall(VecView(p, user->viewerP));
978   PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p));
979 
980   PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e));
981   PetscCall(VecView(e, user->viewerE));
982   PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e));
983   PetscFunctionReturn(PETSC_SUCCESS);
984 }
985 
986 PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
987 {
988   AppCtx    *user = (AppCtx *)ctx;
989   DM         sw;
990   PetscDraw  drawic_x, drawic_v;
991   PetscReal *weight, *pos, *vel;
992   PetscInt   dim, Np;
993 
994   PetscFunctionBegin;
995   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
996   if (step == 0) {
997     PetscCall(TSGetDM(ts, &sw));
998     PetscCall(DMGetDimension(sw, &dim));
999     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1000 
1001     PetscCall(PetscDrawHGReset(user->drawhgic_x));
1002     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &drawic_x));
1003     PetscCall(PetscDrawClear(drawic_x));
1004     PetscCall(PetscDrawFlush(drawic_x));
1005 
1006     PetscCall(PetscDrawHGReset(user->drawhgic_v));
1007     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &drawic_v));
1008     PetscCall(PetscDrawClear(drawic_v));
1009     PetscCall(PetscDrawFlush(drawic_v));
1010 
1011     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
1012     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1013     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1014     for (PetscInt p = 0; p < Np; ++p) {
1015       PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_x, pos[p * dim], weight[p]));
1016       PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_v, vel[p * dim], weight[p]));
1017     }
1018     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
1019     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1020     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1021 
1022     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
1023     PetscCall(PetscDrawHGSave(user->drawhgic_x));
1024     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
1025     PetscCall(PetscDrawHGSave(user->drawhgic_v));
1026   }
1027   PetscFunctionReturn(PETSC_SUCCESS);
1028 }
1029 
1030 // Right now, make the complete velocity histogram
1031 PetscErrorCode MonitorVelocity(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
1032 {
1033   AppCtx      *user = (AppCtx *)ctx;
1034   DM           sw, dm;
1035   Vec          ks;
1036   PetscProbFn *cdf;
1037   PetscDraw    drawcell_v;
1038   PetscScalar *ksa;
1039   PetscReal   *weight, *vel;
1040   PetscInt    *pidx;
1041   PetscInt     dim, Npc, cStart, cEnd, cell = user->velocity_monitor;
1042 
1043   PetscFunctionBegin;
1044   PetscCall(TSGetDM(ts, &sw));
1045   PetscCall(DMGetDimension(sw, &dim));
1046 
1047   PetscCall(DMSwarmGetCellDM(sw, &dm));
1048   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1049   PetscCall(VecCreate(PetscObjectComm((PetscObject)dm), &ks));
1050   PetscCall(PetscObjectSetName((PetscObject)ks, "KS Statistic by Cell"));
1051   PetscCall(VecSetSizes(ks, cEnd - cStart, PETSC_DETERMINE));
1052   PetscCall(VecSetFromOptions(ks));
1053   switch (dim) {
1054   case 1:
1055     //cdf = PetscCDFMaxwellBoltzmann1D;
1056     cdf = PetscCDFGaussian1D;
1057     break;
1058   case 2:
1059     cdf = PetscCDFMaxwellBoltzmann2D;
1060     break;
1061   case 3:
1062     cdf = PetscCDFMaxwellBoltzmann3D;
1063     break;
1064   default:
1065     SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported", dim);
1066   }
1067 
1068   PetscCall(PetscDrawHGReset(user->drawhgcell_v));
1069   PetscCall(PetscDrawHGGetDraw(user->drawhgcell_v, &drawcell_v));
1070   PetscCall(PetscDrawClear(drawcell_v));
1071   PetscCall(PetscDrawFlush(drawcell_v));
1072 
1073   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1074   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1075   PetscCall(DMSwarmSortGetAccess(sw));
1076   PetscCall(VecGetArrayWrite(ks, &ksa));
1077   for (PetscInt c = cStart; c < cEnd; ++c) {
1078     Vec          cellv, cellw;
1079     PetscScalar *cella, *cellaw;
1080     PetscReal    totWgt = 0.;
1081 
1082     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1083     PetscCall(VecCreate(PETSC_COMM_SELF, &cellv));
1084     PetscCall(VecSetBlockSize(cellv, dim));
1085     PetscCall(VecSetSizes(cellv, Npc * dim, Npc));
1086     PetscCall(VecSetFromOptions(cellv));
1087     PetscCall(VecCreate(PETSC_COMM_SELF, &cellw));
1088     PetscCall(VecSetSizes(cellw, Npc, Npc));
1089     PetscCall(VecSetFromOptions(cellw));
1090     PetscCall(VecGetArrayWrite(cellv, &cella));
1091     PetscCall(VecGetArrayWrite(cellw, &cellaw));
1092     for (PetscInt q = 0; q < Npc; ++q) {
1093       const PetscInt p = pidx[q];
1094       if (c == cell) PetscCall(PetscDrawHGAddWeightedValue(user->drawhgcell_v, vel[p * dim], weight[p]));
1095       for (PetscInt d = 0; d < dim; ++d) cella[q * dim + d] = vel[p * dim + d];
1096       cellaw[q] = weight[p];
1097       totWgt += weight[p];
1098     }
1099     PetscCall(VecRestoreArrayWrite(cellv, &cella));
1100     PetscCall(VecRestoreArrayWrite(cellw, &cellaw));
1101     PetscCall(VecScale(cellw, 1. / totWgt));
1102     PetscCall(PetscProbComputeKSStatisticWeighted(cellv, cellw, cdf, &ksa[c - cStart]));
1103     PetscCall(VecDestroy(&cellv));
1104     PetscCall(VecDestroy(&cellw));
1105     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1106   }
1107   PetscCall(VecRestoreArrayWrite(ks, &ksa));
1108   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1109   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1110   PetscCall(DMSwarmSortRestoreAccess(sw));
1111 
1112   PetscReal minalpha, maxalpha;
1113   PetscInt  mincell, maxcell;
1114 
1115   PetscCall(VecFilter(ks, PETSC_SMALL));
1116   PetscCall(VecMin(ks, &mincell, &minalpha));
1117   PetscCall(VecMax(ks, &maxcell, &maxalpha));
1118   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));
1119   PetscCall(VecViewFromOptions(ks, NULL, "-ks_view"));
1120   PetscCall(VecDestroy(&ks));
1121 
1122   PetscCall(PetscDrawHGDraw(user->drawhgcell_v));
1123   PetscCall(PetscDrawHGSave(user->drawhgcell_v));
1124   PetscFunctionReturn(PETSC_SUCCESS);
1125 }
1126 
1127 static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
1128 {
1129   AppCtx         *user = (AppCtx *)ctx;
1130   DM              dm, sw;
1131   PetscScalar    *x, *v, *weight;
1132   PetscReal       lower[3], upper[3], speed;
1133   const PetscInt *s;
1134   PetscInt        dim, cStart, cEnd, c;
1135 
1136   PetscFunctionBeginUser;
1137   if (step > 0 && step % user->ostep == 0) {
1138     PetscCall(TSGetDM(ts, &sw));
1139     PetscCall(DMSwarmGetCellDM(sw, &dm));
1140     PetscCall(DMGetDimension(dm, &dim));
1141     PetscCall(DMGetBoundingBox(dm, lower, upper));
1142     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1143     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1144     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1145     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1146     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
1147     PetscCall(DMSwarmSortGetAccess(sw));
1148     PetscCall(PetscDrawSPReset(user->drawspX));
1149     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], lower[1], upper[1]));
1150     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], -12, 12));
1151     for (c = 0; c < cEnd - cStart; ++c) {
1152       PetscInt *pidx, Npc, q;
1153       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1154       for (q = 0; q < Npc; ++q) {
1155         const PetscInt p = pidx[q];
1156         if (s[p] == 0) {
1157           speed = 0.;
1158           for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]);
1159           speed = PetscSqrtReal(speed);
1160           if (dim == 1) {
1161             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &v[p * dim], &speed));
1162           } else {
1163             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
1164           }
1165         } else if (s[p] == 1) {
1166           PetscCall(PetscDrawSPAddPoint(user->drawspX, &x[p * dim], &v[p * dim]));
1167         }
1168       }
1169       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1170     }
1171     PetscCall(PetscDrawSPDraw(user->drawspX, PETSC_TRUE));
1172     PetscDraw draw;
1173     PetscCall(PetscDrawSPGetDraw(user->drawspX, &draw));
1174     PetscCall(PetscDrawSave(draw));
1175     PetscCall(DMSwarmSortRestoreAccess(sw));
1176     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1177     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1178     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1179     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
1180   }
1181   PetscFunctionReturn(PETSC_SUCCESS);
1182 }
1183 
1184 static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
1185 {
1186   AppCtx *user = (AppCtx *)ctx;
1187   DM      dm, sw;
1188 
1189   PetscFunctionBeginUser;
1190   if (step > 0 && step % user->ostep == 0) {
1191     PetscCall(TSGetDM(ts, &sw));
1192     PetscCall(DMSwarmGetCellDM(sw, &dm));
1193 
1194     if (user->validE) {
1195       PetscScalar *x, *E, *weight;
1196       PetscReal    lower[3], upper[3], xval;
1197       PetscDraw    draw;
1198       PetscInt     dim, cStart, cEnd;
1199 
1200       PetscCall(DMGetDimension(dm, &dim));
1201       PetscCall(DMGetBoundingBox(dm, lower, upper));
1202       PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1203 
1204       PetscCall(PetscDrawSPReset(user->drawspE));
1205       PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1206       PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1207       PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1208 
1209       PetscCall(DMSwarmSortGetAccess(sw));
1210       for (PetscInt c = 0; c < cEnd - cStart; ++c) {
1211         PetscReal Eavg = 0.0;
1212         PetscInt *pidx, Npc;
1213 
1214         PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1215         for (PetscInt q = 0; q < Npc; ++q) {
1216           const PetscInt p = pidx[q];
1217           Eavg += E[p * dim];
1218         }
1219         Eavg /= Npc;
1220         xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
1221         PetscCall(PetscDrawSPAddPoint(user->drawspE, &xval, &Eavg));
1222         PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1223       }
1224       PetscCall(PetscDrawSPDraw(user->drawspE, PETSC_TRUE));
1225       PetscCall(PetscDrawSPGetDraw(user->drawspE, &draw));
1226       PetscCall(PetscDrawSave(draw));
1227       PetscCall(DMSwarmSortRestoreAccess(sw));
1228       PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1229       PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1230       PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1231     }
1232 
1233     Vec rho, rhohat, phi;
1234 
1235     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
1236     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
1237     PetscCall(VecView(rho, user->viewerRho));
1238     PetscCall(VecISCopy(user->fftX, user->fftReal, SCATTER_FORWARD, rho));
1239     PetscCall(MatMult(user->fftPot, user->fftX, user->fftY));
1240     PetscCall(VecFilter(user->fftY, PETSC_SMALL));
1241     PetscCall(VecViewFromOptions(user->fftX, NULL, "-real_view"));
1242     PetscCall(VecViewFromOptions(user->fftY, NULL, "-fft_view"));
1243     PetscCall(VecISCopy(user->fftY, user->fftReal, SCATTER_REVERSE, rhohat));
1244     PetscCall(VecSetValue(rhohat, 0, 0., INSERT_VALUES)); // Remove large DC component
1245     PetscCall(VecView(rhohat, user->viewerRhoHat));
1246     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
1247     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
1248 
1249     PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1250     PetscCall(VecView(phi, user->viewerPhi));
1251     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
1252   }
1253   PetscFunctionReturn(PETSC_SUCCESS);
1254 }
1255 
1256 static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
1257 {
1258   PetscBag   bag;
1259   Parameter *p;
1260 
1261   PetscFunctionBeginUser;
1262   /* setup PETSc parameter bag */
1263   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
1264   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
1265   bag = ctx->bag;
1266   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
1267   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
1268   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
1269   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
1270   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
1271   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
1272   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
1273   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));
1274 
1275   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
1276   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
1277   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
1278   PetscCall(PetscBagSetFromOptions(bag));
1279   {
1280     PetscViewer       viewer;
1281     PetscViewerFormat format;
1282     PetscBool         flg;
1283 
1284     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
1285     if (flg) {
1286       PetscCall(PetscViewerPushFormat(viewer, format));
1287       PetscCall(PetscBagView(bag, viewer));
1288       PetscCall(PetscViewerFlush(viewer));
1289       PetscCall(PetscViewerPopFormat(viewer));
1290       PetscCall(PetscViewerDestroy(&viewer));
1291     }
1292   }
1293   PetscFunctionReturn(PETSC_SUCCESS);
1294 }
1295 
1296 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
1297 {
1298   PetscFunctionBeginUser;
1299   PetscCall(DMCreate(comm, dm));
1300   PetscCall(DMSetType(*dm, DMPLEX));
1301   PetscCall(DMSetFromOptions(*dm));
1302   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
1303   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
1304 
1305   // Cache the mesh geometry
1306   DMField         coordField;
1307   IS              cellIS;
1308   PetscQuadrature quad;
1309   PetscReal      *wt, *pt;
1310   PetscInt        cdim, cStart, cEnd;
1311 
1312   PetscCall(DMGetCoordinateField(*dm, &coordField));
1313   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
1314   PetscCall(DMGetCoordinateDim(*dm, &cdim));
1315   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
1316   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
1317   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
1318   PetscCall(PetscMalloc1(1, &wt));
1319   PetscCall(PetscMalloc1(2, &pt));
1320   wt[0] = 1.;
1321   pt[0] = -1.;
1322   pt[1] = -1.;
1323   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
1324   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom));
1325   PetscCall(PetscQuadratureDestroy(&quad));
1326   PetscCall(ISDestroy(&cellIS));
1327   PetscFunctionReturn(PETSC_SUCCESS);
1328 }
1329 
1330 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[])
1331 {
1332   f0[0] = -constants[SIGMA];
1333 }
1334 
1335 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[])
1336 {
1337   PetscInt d;
1338   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
1339 }
1340 
1341 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[])
1342 {
1343   PetscInt d;
1344   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
1345 }
1346 
1347 /*
1348    /  I   -grad\ / q \ = /0\
1349    \-div    0  / \phi/   \f/
1350 */
1351 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[])
1352 {
1353   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
1354 }
1355 
1356 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[])
1357 {
1358   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
1359 }
1360 
1361 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[])
1362 {
1363   f0[0] += constants[SIGMA];
1364   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
1365 }
1366 
1367 /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
1368 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[])
1369 {
1370   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
1371 }
1372 
1373 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[])
1374 {
1375   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
1376 }
1377 
1378 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[])
1379 {
1380   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
1381 }
1382 
1383 static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
1384 {
1385   PetscFE   fephi, feq;
1386   PetscDS   ds;
1387   PetscBool simplex;
1388   PetscInt  dim;
1389 
1390   PetscFunctionBeginUser;
1391   PetscCall(DMGetDimension(dm, &dim));
1392   PetscCall(DMPlexIsSimplex(dm, &simplex));
1393   if (user->em == EM_MIXED) {
1394     DMLabel        label;
1395     const PetscInt id = 1;
1396 
1397     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
1398     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
1399     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
1400     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1401     PetscCall(PetscFECopyQuadrature(feq, fephi));
1402     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
1403     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
1404     PetscCall(DMCreateDS(dm));
1405     PetscCall(PetscFEDestroy(&fephi));
1406     PetscCall(PetscFEDestroy(&feq));
1407 
1408     PetscCall(DMGetLabel(dm, "marker", &label));
1409     PetscCall(DMGetDS(dm, &ds));
1410 
1411     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
1412     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
1413     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
1414     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
1415     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));
1416 
1417     PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "wall", label, 1, &id, 0, 0, NULL, (void (*)(void))zero, NULL, NULL, NULL));
1418 
1419   } else {
1420     MatNullSpace nullsp;
1421     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
1422     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1423     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
1424     PetscCall(DMCreateDS(dm));
1425     PetscCall(DMGetDS(dm, &ds));
1426     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
1427     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
1428     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
1429     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
1430     PetscCall(MatNullSpaceDestroy(&nullsp));
1431     PetscCall(PetscFEDestroy(&fephi));
1432   }
1433   PetscFunctionReturn(PETSC_SUCCESS);
1434 }
1435 
1436 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
1437 {
1438   SNES         snes;
1439   Mat          J;
1440   MatNullSpace nullSpace;
1441 
1442   PetscFunctionBeginUser;
1443   PetscCall(CreateFEM(dm, user));
1444   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
1445   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
1446   PetscCall(SNESSetDM(snes, dm));
1447   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
1448   PetscCall(SNESSetFromOptions(snes));
1449 
1450   PetscCall(DMCreateMatrix(dm, &J));
1451   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
1452   PetscCall(MatSetNullSpace(J, nullSpace));
1453   PetscCall(MatNullSpaceDestroy(&nullSpace));
1454   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
1455   PetscCall(MatDestroy(&J));
1456   if (user->em == EM_MIXED) {
1457     const PetscInt potential = 1;
1458 
1459     PetscCall(DMCreateSubDM(dm, 1, &potential, &user->isPot, &user->dmPot));
1460   } else {
1461     user->dmPot = dm;
1462     PetscCall(PetscObjectReference((PetscObject)user->dmPot));
1463   }
1464   PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M));
1465   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
1466   user->snes = snes;
1467   PetscFunctionReturn(PETSC_SUCCESS);
1468 }
1469 
1470 // Conservation of mass (m = 1.0)
1471 // n_t + 1/ m p_x = 0
1472 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[])
1473 {
1474   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_t[uOff[0]] + u_x[uOff_x[1] + d * dim + d];
1475 }
1476 
1477 // Conservation of momentum (m = 1, e = 1)
1478 // p_t + (u p)_x = -pr_x + e n E
1479 // p_t + (div u) p + u . grad p = -pr_x + e n E
1480 // p_t + (div p) p / n - (p . grad n) p / n^2 + p / n . grad p = -pr_x + e n E
1481 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[])
1482 {
1483   const PetscScalar n = u[uOff[0]];
1484 
1485   for (PetscInt d = 0; d < dim; ++d) {
1486     PetscReal divp = 0.;
1487 
1488     f0[d] += u_t[uOff[1] + d];
1489     for (PetscInt e = 0; e < dim; ++e) {
1490       f0[d] += u[uOff[1] + e] * u_x[uOff_x[1] + d * dim + e] / n;                    // p / n . grad p
1491       f0[d] -= (u[uOff[1] + e] * u_x[uOff_x[0] + e]) * u[uOff[1] + d] / PetscSqr(n); // -(p . grad n) p / n^2
1492       divp += u_x[uOff_x[1] + e * dim + e];
1493     }
1494     f0[d] += divp * u[uOff[1] + d] / n; // (div p) p / n
1495     f0[d] += u_x[uOff_x[2] + d];        // pr_x
1496     f0[d] -= n * a[d];                  // -e n E
1497   }
1498 }
1499 
1500 // Conservation of energy
1501 // pr_t + (u pr)_x = -3 pr u_x - q_x
1502 // pr_t + (div u) pr + u . grad pr = -3 pr (div u) - q_x
1503 // pr_t + 4 (div u) pr + u . grad pr = -q_x
1504 // pr_t + 4 div p pr / n - 4 (p . grad n) pr / n^2 + p . grad pr / n = -q_x
1505 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[])
1506 {
1507   const PetscScalar n    = u[uOff[0]];
1508   const PetscScalar pr   = u[uOff[2]];
1509   PetscReal         divp = 0.;
1510 
1511   f0[0] += u_t[uOff[2]];
1512   for (PetscInt d = 0; d < dim; ++d) {
1513     f0[0] += u[uOff[1] + d] * u_x[uOff_x[2] + d] / n;                     // p . grad pr / n
1514     f0[0] -= 4. * u[uOff[1] + d] * u_x[uOff_x[0] + d] * pr / PetscSqr(n); // -4 (p . grad n) pr / n^2
1515     divp += u_x[uOff_x[1] + d * dim + d];
1516   }
1517   f0[0] += 4. * divp * pr / n; // 4 div p pr / n
1518 }
1519 
1520 static PetscErrorCode SetupMomentProblem(DM dm, AppCtx *ctx)
1521 {
1522   PetscDS ds;
1523 
1524   PetscFunctionBegin;
1525   PetscCall(DMGetDS(dm, &ds));
1526   PetscCall(PetscDSSetResidual(ds, 0, f0_mass, NULL));
1527   PetscCall(PetscDSSetResidual(ds, 1, f0_momentum, NULL));
1528   PetscCall(PetscDSSetResidual(ds, 2, f0_energy, NULL));
1529   //PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_mass_uu, NULL, NULL, NULL));
1530   PetscFunctionReturn(PETSC_SUCCESS);
1531 }
1532 
1533 static PetscErrorCode CreateMomentFields(DM odm, AppCtx *user)
1534 {
1535   DM             dm;
1536   PetscFE        fe;
1537   DMPolytopeType ct;
1538   PetscInt       dim, cStart;
1539 
1540   PetscFunctionBeginUser;
1541   PetscCall(DMClone(odm, &dm));
1542   PetscCall(DMGetDimension(dm, &dim));
1543   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, NULL));
1544   PetscCall(DMPlexGetCellType(dm, cStart, &ct));
1545   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
1546   PetscCall(PetscObjectSetName((PetscObject)fe, "number density"));
1547   PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fe));
1548   PetscCall(PetscFEDestroy(&fe));
1549   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, dim, ct, NULL, PETSC_DETERMINE, &fe));
1550   PetscCall(PetscObjectSetName((PetscObject)fe, "momentum density"));
1551   PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fe));
1552   PetscCall(PetscFEDestroy(&fe));
1553   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
1554   PetscCall(PetscObjectSetName((PetscObject)fe, "energy density"));
1555   PetscCall(DMSetField(dm, 2, NULL, (PetscObject)fe));
1556   PetscCall(PetscFEDestroy(&fe));
1557   PetscCall(DMCreateDS(dm));
1558   PetscCall(SetupMomentProblem(dm, user));
1559   user->dmMom = dm;
1560   PetscInt field;
1561 
1562   field = 0;
1563   PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isN, &user->dmN));
1564   PetscCall(DMCreateMassMatrix(user->dmN, user->dmN, &user->MN));
1565   field = 1;
1566   PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isP, &user->dmP));
1567   PetscCall(DMCreateMassMatrix(user->dmP, user->dmP, &user->MP));
1568   field = 2;
1569   PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isE, &user->dmE));
1570   PetscCall(DMCreateMassMatrix(user->dmE, user->dmE, &user->ME));
1571   PetscFunctionReturn(PETSC_SUCCESS);
1572 }
1573 
1574 PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1575 {
1576   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1577   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
1578   return PETSC_SUCCESS;
1579 }
1580 PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1581 {
1582   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1583   return PETSC_SUCCESS;
1584 }
1585 
1586 PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1587 {
1588   const PetscReal alpha = scale ? scale[0] : 0.0;
1589   const PetscReal k     = scale ? scale[1] : 1.;
1590   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
1591   return PETSC_SUCCESS;
1592 }
1593 
1594 PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1595 {
1596   const PetscReal alpha = scale ? scale[0] : 0.;
1597   const PetscReal k     = scale ? scale[0] : 1.;
1598   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
1599   return PETSC_SUCCESS;
1600 }
1601 
1602 static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
1603 {
1604   PetscFE        fe;
1605   DMPolytopeType ct;
1606   PetscInt       dim, cStart;
1607   const char    *prefix = "v";
1608 
1609   PetscFunctionBegin;
1610   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
1611   PetscCall(DMSetType(*vdm, DMPLEX));
1612   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
1613   PetscCall(DMSetFromOptions(*vdm));
1614   PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
1615   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));
1616 
1617   PetscCall(DMGetDimension(*vdm, &dim));
1618   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
1619   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
1620   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
1621   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
1622   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
1623   PetscCall(DMCreateDS(*vdm));
1624   PetscCall(PetscFEDestroy(&fe));
1625   PetscFunctionReturn(PETSC_SUCCESS);
1626 }
1627 
1628 /*
1629   InitializeParticles_Centroid - Initialize a regular grid of particles.
1630 
1631   Input Parameters:
1632 + sw      - The `DMSWARM`
1633 - force1D - Treat the spatial domain as 1D
1634 
1635   Notes:
1636   This functions sets the species, cellid, spatial coordinate, and velocity fields for all particles.
1637 
1638   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
1639   and velocity meshes.
1640 */
1641 static PetscErrorCode InitializeParticles_Centroid(DM sw)
1642 {
1643   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
1644   DMSwarmCellDM celldm;
1645   DM            xdm, vdm;
1646   PetscReal     vmin[3], vmax[3];
1647   PetscReal    *x, *v;
1648   PetscInt     *species, *cellid;
1649   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
1650   PetscBool     flg;
1651   MPI_Comm      comm;
1652   const char   *cellidname;
1653 
1654   PetscFunctionBegin;
1655   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
1656 
1657   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
1658   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1659   PetscCall(PetscOptionsInt("-dm_swarMNum_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
1660   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
1661   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
1662   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
1663   PetscOptionsEnd();
1664   debug = swarm->printCoords;
1665 
1666   PetscCall(DMGetDimension(sw, &dim));
1667   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1668   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
1669 
1670   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1671   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1672   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
1673 
1674   // One particle per centroid on the tensor product grid
1675   Npc = (vcEnd - vcStart) * Ns;
1676   Np  = (xcEnd - xcStart) * Npc;
1677   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
1678   if (debug) {
1679     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
1680     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
1681     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
1682     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
1683     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
1684     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
1685   }
1686 
1687   // Set species and cellid
1688   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
1689   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
1690   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1691   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
1692   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
1693     for (PetscInt s = 0; s < Ns; ++s) {
1694       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
1695         species[p] = s;
1696         cellid[p]  = c;
1697       }
1698     }
1699   }
1700   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1701   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));
1702 
1703   // Set particle coordinates
1704   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1705   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1706   PetscCall(DMSwarmSortGetAccess(sw));
1707   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
1708   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
1709   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
1710     const PetscInt cell = c + xcStart;
1711     PetscInt      *pidx, Npc;
1712     PetscReal      centroid[3], volume;
1713 
1714     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1715     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
1716     for (PetscInt s = 0; s < Ns; ++s) {
1717       for (PetscInt q = 0; q < Npc / Ns; ++q) {
1718         const PetscInt p = pidx[q * Ns + s];
1719 
1720         for (PetscInt d = 0; d < dim; ++d) {
1721           x[p * dim + d] = centroid[d];
1722           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
1723         }
1724         if (debug > 1) {
1725           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
1726           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
1727           for (PetscInt d = 0; d < dim; ++d) {
1728             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1729             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
1730           }
1731           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
1732           for (PetscInt d = 0; d < dim; ++d) {
1733             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1734             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
1735           }
1736           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
1737         }
1738       }
1739     }
1740     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1741   }
1742   PetscCall(DMSwarmSortRestoreAccess(sw));
1743   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1744   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1745   PetscFunctionReturn(PETSC_SUCCESS);
1746 }
1747 
1748 /*
1749   InitializeWeights - Compute weight for each local particle
1750 
1751   Input Parameters:
1752 + sw          - The `DMSwarm`
1753 . totalWeight - The sum of all particle weights
1754 . func        - The PDF for the particle spatial distribution
1755 - param       - The PDF parameters
1756 
1757   Notes:
1758   The PDF for velocity is assumed to be a Gaussian
1759 
1760   The particle weights are returned in the `w_q` field of `sw`.
1761 */
1762 static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
1763 {
1764   DM               xdm, vdm;
1765   DMSwarmCellDM    celldm;
1766   PetscScalar     *weight;
1767   PetscQuadrature  xquad;
1768   const PetscReal *xq, *xwq;
1769   const PetscInt   order = 5;
1770   PetscReal        xi0[3];
1771   PetscReal        xwtot = 0., pwtot = 0.;
1772   PetscInt         xNq;
1773   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
1774   MPI_Comm         comm;
1775   PetscMPIInt      rank;
1776 
1777   PetscFunctionBegin;
1778   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
1779   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1780   PetscCall(DMGetDimension(sw, &dim));
1781   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1782   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1783   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
1784   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1785   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1786   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
1787 
1788   // Setup Quadrature for spatial and velocity weight calculations
1789   PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad));
1790   PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq));
1791   for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0;
1792 
1793   // Integrate the density function to get the weights of particles in each cell
1794   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
1795   PetscCall(DMSwarmSortGetAccess(sw));
1796   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1797   for (PetscInt c = xcStart; c < xcEnd; ++c) {
1798     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
1799     PetscInt          *pidx, Npc;
1800     PetscInt           xNc;
1801     const PetscScalar *xarray;
1802     PetscScalar       *xcoords = NULL;
1803     PetscBool          xisDG;
1804 
1805     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1806     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1807     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);
1808     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
1809     for (PetscInt q = 0; q < xNq; ++q) {
1810       // Transform quadrature points from ref space to real space
1811       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
1812       // Get probability density at quad point
1813       //   No need to scale xqr since PDF will be periodic
1814       PetscCall((*func)(xqr, param, &xden));
1815       xw += xden * (xwq[q] * xdetJ);
1816     }
1817     xwtot += xw;
1818     if (debug) {
1819       IS              globalOrdering;
1820       const PetscInt *ordering;
1821 
1822       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
1823       PetscCall(ISGetIndices(globalOrdering, &ordering));
1824       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));
1825       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
1826     }
1827     // Set weights to be Gaussian in velocity cells
1828     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
1829       const PetscInt     p  = pidx[vc * Ns + 0];
1830       PetscReal          vw = 0.;
1831       PetscInt           vNc;
1832       const PetscScalar *varray;
1833       PetscScalar       *vcoords = NULL;
1834       PetscBool          visDG;
1835 
1836       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1837       // TODO: Fix 2 stream Ask Joe
1838       //   Two stream function from 1/2pi v^2 e^(-v^2/2)
1839       //   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.)));
1840       vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.)));
1841 
1842       weight[p] = totalWeight * vw * xw;
1843       pwtot += weight[p];
1844       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 1: %g, %g, %g", p, xw, vw, totalWeight);
1845       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1846       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
1847     }
1848     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1849     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1850   }
1851   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1852   PetscCall(DMSwarmSortRestoreAccess(sw));
1853   PetscCall(PetscQuadratureDestroy(&xquad));
1854 
1855   if (debug) {
1856     PetscReal wtot[2] = {pwtot, xwtot}, gwtot[2];
1857 
1858     PetscCall(PetscSynchronizedFlush(comm, NULL));
1859     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1860     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
1861   }
1862   PetscFunctionReturn(PETSC_SUCCESS);
1863 }
1864 
1865 static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
1866 {
1867   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
1868   PetscInt  dim;
1869 
1870   PetscFunctionBegin;
1871   PetscCall(DMGetDimension(sw, &dim));
1872   PetscCall(InitializeParticles_Centroid(sw));
1873   PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
1874   PetscFunctionReturn(PETSC_SUCCESS);
1875 }
1876 
1877 static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
1878 {
1879   DM         dm;
1880   PetscInt  *species;
1881   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
1882   PetscInt   Np, dim;
1883 
1884   PetscFunctionBegin;
1885   PetscCall(DMSwarmGetCellDM(sw, &dm));
1886   PetscCall(DMGetDimension(sw, &dim));
1887   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1888   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
1889   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1890   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1891   for (PetscInt p = 0; p < Np; ++p) {
1892     totalWeight += weight[p];
1893     totalCharge += user->charges[species[p]] * weight[p];
1894   }
1895   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1896   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1897   {
1898     Parameter *param;
1899     PetscReal  Area;
1900 
1901     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1902     switch (dim) {
1903     case 1:
1904       Area = (gmax[0] - gmin[0]);
1905       break;
1906     case 2:
1907       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1908       break;
1909     case 3:
1910       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1911       break;
1912     default:
1913       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1914     }
1915     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1916     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1917     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)user->charges[0], (double)global_charge, (double)Area));
1918     param->sigma = PetscAbsReal(global_charge / (Area));
1919 
1920     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
1921     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,
1922                           (double)param->vlasovNumber));
1923   }
1924   /* Setup Constants */
1925   {
1926     PetscDS    ds;
1927     Parameter *param;
1928     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1929     PetscScalar constants[NUM_CONSTANTS];
1930     constants[SIGMA]   = param->sigma;
1931     constants[V0]      = param->v0;
1932     constants[T0]      = param->t0;
1933     constants[X0]      = param->x0;
1934     constants[M0]      = param->m0;
1935     constants[Q0]      = param->q0;
1936     constants[PHI0]    = param->phi0;
1937     constants[POISSON] = param->poissonNumber;
1938     constants[VLASOV]  = param->vlasovNumber;
1939     PetscCall(DMGetDS(dm, &ds));
1940     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
1941   }
1942   PetscFunctionReturn(PETSC_SUCCESS);
1943 }
1944 
1945 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1946 {
1947   DMSwarmCellDM celldm;
1948   DM            vdm;
1949   PetscReal     v0[2] = {1., 0.};
1950   PetscInt      dim;
1951 
1952   PetscFunctionBeginUser;
1953   PetscCall(DMGetDimension(dm, &dim));
1954   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
1955   PetscCall(DMSetType(*sw, DMSWARM));
1956   PetscCall(DMSetDimension(*sw, dim));
1957   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
1958   PetscCall(DMSetApplicationContext(*sw, user));
1959 
1960   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
1961   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
1962   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
1963   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
1964 
1965   const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
1966 
1967   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
1968   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1969   PetscCall(DMSwarmCellDMDestroy(&celldm));
1970 
1971   const char *vfieldnames[2] = {"w_q"};
1972 
1973   PetscCall(CreateVelocityDM(*sw, &vdm));
1974   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
1975   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1976   PetscCall(DMSwarmCellDMDestroy(&celldm));
1977   PetscCall(DMDestroy(&vdm));
1978 
1979   DM mdm;
1980 
1981   PetscCall(DMClone(dm, &mdm));
1982   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
1983   PetscCall(DMCopyDisc(dm, mdm));
1984   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
1985   PetscCall(DMDestroy(&mdm));
1986   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1987   PetscCall(DMSwarmCellDMDestroy(&celldm));
1988 
1989   DM mfdm;
1990 
1991   PetscCall(DMClone(dm, &mfdm));
1992   PetscCall(PetscObjectSetName((PetscObject)mfdm, "moment fields"));
1993   PetscCall(DMCopyDisc(dm, mfdm));
1994   PetscCall(DMSwarmCellDMCreate(mfdm, 1, &fieldnames[1], 1, fieldnames, &celldm));
1995   PetscCall(DMDestroy(&mfdm));
1996   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1997   PetscCall(DMSwarmCellDMDestroy(&celldm));
1998 
1999   PetscCall(DMSetFromOptions(*sw));
2000   PetscCall(DMSetUp(*sw));
2001 
2002   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
2003   user->swarm = *sw;
2004   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
2005   if (user->perturbed_weights) {
2006     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
2007   } else {
2008     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
2009     PetscCall(DMSwarmInitializeCoordinates(*sw));
2010     PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
2011   }
2012   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
2013   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
2014   PetscFunctionReturn(PETSC_SUCCESS);
2015 }
2016 
2017 static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
2018 {
2019   AppCtx     *user;
2020   PetscReal  *coords;
2021   PetscInt   *species, dim, Np, Ns;
2022   PetscMPIInt size;
2023 
2024   PetscFunctionBegin;
2025   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
2026   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
2027   PetscCall(DMGetDimension(sw, &dim));
2028   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2029   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
2030   PetscCall(DMGetApplicationContext(sw, (void *)&user));
2031 
2032   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2033   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
2034   for (PetscInt p = 0; p < Np; ++p) {
2035     PetscReal *pcoord = &coords[p * dim];
2036     PetscReal  pE[3]  = {0., 0., 0.};
2037 
2038     /* Calculate field at particle p due to particle q */
2039     for (PetscInt q = 0; q < Np; ++q) {
2040       PetscReal *qcoord = &coords[q * dim];
2041       PetscReal  rpq[3], r, r3, q_q;
2042 
2043       if (p == q) continue;
2044       q_q = user->charges[species[q]] * 1.;
2045       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
2046       r = DMPlex_NormD_Internal(dim, rpq);
2047       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
2048       r3 = PetscPowRealInt(r, 3);
2049       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
2050     }
2051     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
2052   }
2053   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
2054   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2055   PetscFunctionReturn(PETSC_SUCCESS);
2056 }
2057 
2058 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
2059 {
2060   DM         dm;
2061   AppCtx    *user;
2062   PetscDS    ds;
2063   PetscFE    fe;
2064   KSP        ksp;
2065   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
2066   Vec        rho;         // Charge density, M^{-1} rhoRhs
2067   Vec        phi, locPhi; // Potential
2068   Vec        f;           // Particle weights
2069   PetscReal *coords;
2070   PetscInt   dim, cStart, cEnd, Np;
2071 
2072   PetscFunctionBegin;
2073   PetscCall(DMGetApplicationContext(sw, (void *)&user));
2074   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
2075   PetscCall(DMGetDimension(sw, &dim));
2076   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2077 
2078   PetscCall(SNESGetDM(snes, &dm));
2079   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
2080   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
2081   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
2082   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
2083   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
2084 
2085   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
2086   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
2087   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
2088 
2089   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
2090   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
2091 
2092   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
2093   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
2094   PetscCall(KSPSetOperators(ksp, user->M, user->M));
2095   PetscCall(KSPSetFromOptions(ksp));
2096   PetscCall(KSPSolve(ksp, rhoRhs, rho));
2097 
2098   PetscCall(VecScale(rhoRhs, -1.0));
2099 
2100   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
2101   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
2102   PetscCall(KSPDestroy(&ksp));
2103 
2104   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
2105   PetscCall(VecSet(phi, 0.0));
2106   PetscCall(SNESSolve(snes, rhoRhs, phi));
2107   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
2108   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
2109 
2110   PetscCall(DMGetLocalVector(dm, &locPhi));
2111   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
2112   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
2113   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
2114   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));
2115 
2116   PetscCall(DMGetDS(dm, &ds));
2117   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
2118   PetscCall(DMSwarmSortGetAccess(sw));
2119   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
2120   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2121 
2122   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
2123   PetscTabulation tab;
2124   PetscReal      *pcoord, *refcoord;
2125   PetscFEGeom    *chunkgeom = NULL;
2126   PetscInt        maxNcp    = 0;
2127 
2128   for (PetscInt c = cStart; c < cEnd; ++c) {
2129     PetscInt Ncp;
2130 
2131     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
2132     maxNcp = PetscMax(maxNcp, Ncp);
2133   }
2134   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2135   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2136   // This can raise an FP_INEXACT in the dgemm inside
2137   PetscCall(PetscFPTrapPush(PETSC_FP_TRAP_OFF));
2138   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
2139   PetscCall(PetscFPTrapPop());
2140   for (PetscInt c = cStart; c < cEnd; ++c) {
2141     PetscScalar *clPhi = NULL;
2142     PetscInt    *points;
2143     PetscInt     Ncp;
2144 
2145     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
2146     for (PetscInt cp = 0; cp < Ncp; ++cp)
2147       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
2148     {
2149       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
2150       for (PetscInt i = 0; i < Ncp; ++i) {
2151         const PetscReal x0[3] = {-1., -1., -1.};
2152         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
2153       }
2154     }
2155     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
2156     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2157     for (PetscInt cp = 0; cp < Ncp; ++cp) {
2158       const PetscReal *basisDer = tab->T[1];
2159       const PetscInt   p        = points[cp];
2160 
2161       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
2162       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
2163       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
2164     }
2165     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2166     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2167   }
2168   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2169   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2170   PetscCall(PetscTabulationDestroy(&tab));
2171   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2172   PetscCall(DMSwarmSortRestoreAccess(sw));
2173   PetscCall(DMRestoreLocalVector(dm, &locPhi));
2174   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
2175   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
2176   PetscFunctionReturn(PETSC_SUCCESS);
2177 }
2178 
2179 static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, Mat M_p, PetscReal E[])
2180 {
2181   DM         dm;
2182   AppCtx    *user;
2183   PetscDS    ds;
2184   PetscFE    fe;
2185   KSP        ksp;
2186   Vec        rhoRhs, rhoRhsFull;   // Weak charge density, \int phi_i rho, and embedding in mixed problem
2187   Vec        rho;                  // Charge density, M^{-1} rhoRhs
2188   Vec        phi, locPhi, phiFull; // Potential and embedding in mixed problem
2189   Vec        f;                    // Particle weights
2190   PetscReal *coords;
2191   PetscInt   dim, cStart, cEnd, Np;
2192 
2193   PetscFunctionBegin;
2194   PetscCall(DMGetApplicationContext(sw, &user));
2195   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
2196   PetscCall(DMGetDimension(sw, &dim));
2197   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2198 
2199   PetscCall(SNESGetDM(snes, &dm));
2200   PetscCall(DMGetGlobalVector(user->dmPot, &rhoRhs));
2201   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
2202   PetscCall(DMGetGlobalVector(dm, &rhoRhsFull));
2203   PetscCall(PetscObjectSetName((PetscObject)rhoRhsFull, "Weak charge density"));
2204   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
2205   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
2206   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
2207 
2208   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
2209   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
2210   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
2211 
2212   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
2213   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
2214 
2215   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
2216   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
2217   PetscCall(KSPSetOperators(ksp, user->M, user->M));
2218   PetscCall(KSPSetFromOptions(ksp));
2219   PetscCall(KSPSolve(ksp, rhoRhs, rho));
2220 
2221   PetscCall(VecISCopy(rhoRhsFull, user->isPot, SCATTER_FORWARD, rhoRhs));
2222   //PetscCall(VecScale(rhoRhsFull, -1.0));
2223 
2224   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
2225   PetscCall(VecViewFromOptions(rhoRhsFull, NULL, "-rho_full_view"));
2226   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
2227   PetscCall(DMRestoreGlobalVector(user->dmPot, &rhoRhs));
2228   PetscCall(KSPDestroy(&ksp));
2229 
2230   PetscCall(DMGetGlobalVector(dm, &phiFull));
2231   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
2232   PetscCall(VecSet(phiFull, 0.0));
2233   PetscCall(SNESSolve(snes, rhoRhsFull, phiFull));
2234   PetscCall(DMRestoreGlobalVector(dm, &rhoRhsFull));
2235   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
2236 
2237   PetscCall(VecISCopy(phiFull, user->isPot, SCATTER_REVERSE, phi));
2238   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
2239 
2240   PetscCall(DMGetLocalVector(dm, &locPhi));
2241   PetscCall(DMGlobalToLocalBegin(dm, phiFull, INSERT_VALUES, locPhi));
2242   PetscCall(DMGlobalToLocalEnd(dm, phiFull, INSERT_VALUES, locPhi));
2243   PetscCall(DMRestoreGlobalVector(dm, &phiFull));
2244   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));
2245 
2246   PetscCall(DMGetDS(dm, &ds));
2247   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
2248   PetscCall(DMSwarmSortGetAccess(sw));
2249   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
2250   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2251 
2252   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
2253   PetscTabulation tab;
2254   PetscReal      *pcoord, *refcoord;
2255   PetscFEGeom    *chunkgeom = NULL;
2256   PetscInt        maxNcp    = 0;
2257 
2258   for (PetscInt c = cStart; c < cEnd; ++c) {
2259     PetscInt Ncp;
2260 
2261     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
2262     maxNcp = PetscMax(maxNcp, Ncp);
2263   }
2264   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2265   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2266   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
2267   for (PetscInt c = cStart; c < cEnd; ++c) {
2268     PetscScalar *clPhi = NULL;
2269     PetscInt    *points;
2270     PetscInt     Ncp;
2271 
2272     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
2273     for (PetscInt cp = 0; cp < Ncp; ++cp)
2274       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
2275     {
2276       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
2277       for (PetscInt i = 0; i < Ncp; ++i) {
2278         const PetscReal x0[3] = {-1., -1., -1.};
2279         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
2280       }
2281     }
2282     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
2283     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2284     for (PetscInt cp = 0; cp < Ncp; ++cp) {
2285       const PetscInt p = points[cp];
2286 
2287       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
2288       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
2289       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
2290       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
2291     }
2292     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
2293     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2294   }
2295   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
2296   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
2297   PetscCall(PetscTabulationDestroy(&tab));
2298   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2299   PetscCall(DMSwarmSortRestoreAccess(sw));
2300   PetscCall(DMRestoreLocalVector(dm, &locPhi));
2301   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
2302   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
2303   PetscFunctionReturn(PETSC_SUCCESS);
2304 }
2305 
2306 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
2307 {
2308   AppCtx    *user;
2309   Mat        M_p;
2310   PetscReal *E;
2311   PetscInt   dim, Np;
2312 
2313   PetscFunctionBegin;
2314   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
2315   PetscValidHeaderSpecific(sw, DM_CLASSID, 2);
2316   PetscCall(DMGetDimension(sw, &dim));
2317   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2318   PetscCall(DMGetApplicationContext(sw, &user));
2319 
2320   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
2321   // TODO: Could share sort context with space cellDM
2322   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
2323   PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p));
2324   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
2325 
2326   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2327   PetscCall(PetscArrayzero(E, Np * dim));
2328   user->validE = PETSC_TRUE;
2329 
2330   switch (user->em) {
2331   case EM_COULOMB:
2332     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
2333     break;
2334   case EMPRIMAL:
2335     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
2336     break;
2337   case EM_MIXED:
2338     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, M_p, E));
2339     break;
2340   case EMNONE:
2341     break;
2342   default:
2343     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[user->em]);
2344   }
2345   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2346   PetscCall(MatDestroy(&M_p));
2347   PetscFunctionReturn(PETSC_SUCCESS);
2348 }
2349 
2350 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
2351 {
2352   DM                 sw;
2353   SNES               snes = ((AppCtx *)ctx)->snes;
2354   const PetscScalar *u;
2355   PetscScalar       *g;
2356   PetscReal         *E, m_p = 1., q_p = -1.;
2357   PetscInt           dim, d, Np, p;
2358 
2359   PetscFunctionBeginUser;
2360   PetscCall(TSGetDM(ts, &sw));
2361   PetscCall(ComputeFieldAtParticles(snes, sw));
2362 
2363   PetscCall(DMGetDimension(sw, &dim));
2364   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2365   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2366   PetscCall(VecGetArrayRead(U, &u));
2367   PetscCall(VecGetArray(G, &g));
2368   Np /= 2 * dim;
2369   for (p = 0; p < Np; ++p) {
2370     for (d = 0; d < dim; ++d) {
2371       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
2372       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
2373     }
2374   }
2375   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2376   PetscCall(VecRestoreArrayRead(U, &u));
2377   PetscCall(VecRestoreArray(G, &g));
2378   PetscFunctionReturn(PETSC_SUCCESS);
2379 }
2380 
2381 /* J_{ij} = dF_i/dx_j
2382    J_p = (  0   1)
2383          (-w^2  0)
2384    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
2385         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
2386 */
2387 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
2388 {
2389   DM               sw;
2390   const PetscReal *coords, *vel;
2391   PetscInt         dim, d, Np, p, rStart;
2392 
2393   PetscFunctionBeginUser;
2394   PetscCall(TSGetDM(ts, &sw));
2395   PetscCall(DMGetDimension(sw, &dim));
2396   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2397   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
2398   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2399   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2400   Np /= 2 * dim;
2401   for (p = 0; p < Np; ++p) {
2402     const PetscReal x0      = coords[p * dim + 0];
2403     const PetscReal vy0     = vel[p * dim + 1];
2404     const PetscReal omega   = vy0 / x0;
2405     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};
2406 
2407     for (d = 0; d < dim; ++d) {
2408       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
2409       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
2410     }
2411   }
2412   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2413   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2414   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
2415   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
2416   PetscFunctionReturn(PETSC_SUCCESS);
2417 }
2418 
2419 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
2420 {
2421   AppCtx            *user = (AppCtx *)ctx;
2422   DM                 sw;
2423   const PetscScalar *v;
2424   PetscScalar       *xres;
2425   PetscInt           Np, p, d, dim;
2426 
2427   PetscFunctionBeginUser;
2428   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
2429   PetscCall(TSGetDM(ts, &sw));
2430   PetscCall(DMGetDimension(sw, &dim));
2431   PetscCall(VecGetLocalSize(Xres, &Np));
2432   PetscCall(VecGetArrayRead(V, &v));
2433   PetscCall(VecGetArray(Xres, &xres));
2434   Np /= dim;
2435   for (p = 0; p < Np; ++p) {
2436     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
2437   }
2438   PetscCall(VecRestoreArrayRead(V, &v));
2439   PetscCall(VecRestoreArray(Xres, &xres));
2440   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
2441   PetscFunctionReturn(PETSC_SUCCESS);
2442 }
2443 
2444 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
2445 {
2446   DM                 sw;
2447   AppCtx            *user = (AppCtx *)ctx;
2448   SNES               snes = ((AppCtx *)ctx)->snes;
2449   const PetscScalar *x;
2450   PetscScalar       *vres;
2451   PetscReal         *E, m_p, q_p;
2452   PetscInt           Np, p, dim, d;
2453   Parameter         *param;
2454 
2455   PetscFunctionBeginUser;
2456   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
2457   PetscCall(TSGetDM(ts, &sw));
2458   PetscCall(ComputeFieldAtParticles(snes, sw));
2459 
2460   PetscCall(DMGetDimension(sw, &dim));
2461   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2462   PetscCall(PetscBagGetData(user->bag, (void **)&param));
2463   m_p = user->masses[0] * param->m0;
2464   q_p = user->charges[0] * param->q0;
2465   PetscCall(VecGetLocalSize(Vres, &Np));
2466   PetscCall(VecGetArrayRead(X, &x));
2467   PetscCall(VecGetArray(Vres, &vres));
2468   Np /= dim;
2469   for (p = 0; p < Np; ++p) {
2470     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
2471   }
2472   PetscCall(VecRestoreArrayRead(X, &x));
2473   /*
2474     Synchronized, ordered output for parallel/sequential test cases.
2475     In the 1D (on the 2D mesh) case, every y component should be zero.
2476   */
2477   if (user->checkVRes) {
2478     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
2479     PetscInt  step;
2480 
2481     PetscCall(TSGetStepNumber(ts, &step));
2482     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
2483     for (PetscInt p = 0; p < Np; ++p) {
2484       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
2485       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]));
2486     }
2487     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
2488   }
2489   PetscCall(VecRestoreArray(Vres, &vres));
2490   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2491   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
2492   PetscFunctionReturn(PETSC_SUCCESS);
2493 }
2494 
2495 static PetscErrorCode CreateSolution(TS ts)
2496 {
2497   DM       sw;
2498   Vec      u;
2499   PetscInt dim, Np;
2500 
2501   PetscFunctionBegin;
2502   PetscCall(TSGetDM(ts, &sw));
2503   PetscCall(DMGetDimension(sw, &dim));
2504   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2505   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
2506   PetscCall(VecSetBlockSize(u, dim));
2507   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
2508   PetscCall(VecSetUp(u));
2509   PetscCall(TSSetSolution(ts, u));
2510   PetscCall(VecDestroy(&u));
2511   PetscFunctionReturn(PETSC_SUCCESS);
2512 }
2513 
2514 static PetscErrorCode SetProblem(TS ts)
2515 {
2516   AppCtx *user;
2517   DM      sw;
2518 
2519   PetscFunctionBegin;
2520   PetscCall(TSGetDM(ts, &sw));
2521   PetscCall(DMGetApplicationContext(sw, (void **)&user));
2522   // Define unified system for (X, V)
2523   {
2524     Mat      J;
2525     PetscInt dim, Np;
2526 
2527     PetscCall(DMGetDimension(sw, &dim));
2528     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2529     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
2530     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
2531     PetscCall(MatSetBlockSize(J, 2 * dim));
2532     PetscCall(MatSetFromOptions(J));
2533     PetscCall(MatSetUp(J));
2534     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
2535     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
2536     PetscCall(MatDestroy(&J));
2537   }
2538   /* Define split system for X and V */
2539   {
2540     Vec             u;
2541     IS              isx, isv, istmp;
2542     const PetscInt *idx;
2543     PetscInt        dim, Np, rstart;
2544 
2545     PetscCall(TSGetSolution(ts, &u));
2546     PetscCall(DMGetDimension(sw, &dim));
2547     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2548     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
2549     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
2550     PetscCall(ISGetIndices(istmp, &idx));
2551     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
2552     PetscCall(ISRestoreIndices(istmp, &idx));
2553     PetscCall(ISDestroy(&istmp));
2554     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
2555     PetscCall(ISGetIndices(istmp, &idx));
2556     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
2557     PetscCall(ISRestoreIndices(istmp, &idx));
2558     PetscCall(ISDestroy(&istmp));
2559     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
2560     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
2561     PetscCall(ISDestroy(&isx));
2562     PetscCall(ISDestroy(&isv));
2563     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
2564     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
2565   }
2566   PetscFunctionReturn(PETSC_SUCCESS);
2567 }
2568 
2569 static PetscErrorCode DMSwarmTSRedistribute(TS ts)
2570 {
2571   DM        sw;
2572   Vec       u;
2573   PetscReal t, maxt, dt;
2574   PetscInt  n, maxn;
2575 
2576   PetscFunctionBegin;
2577   PetscCall(TSGetDM(ts, &sw));
2578   PetscCall(TSGetTime(ts, &t));
2579   PetscCall(TSGetMaxTime(ts, &maxt));
2580   PetscCall(TSGetTimeStep(ts, &dt));
2581   PetscCall(TSGetStepNumber(ts, &n));
2582   PetscCall(TSGetMaxSteps(ts, &maxn));
2583 
2584   PetscCall(TSReset(ts));
2585   PetscCall(TSSetDM(ts, sw));
2586   PetscCall(TSSetFromOptions(ts));
2587   PetscCall(TSSetTime(ts, t));
2588   PetscCall(TSSetMaxTime(ts, maxt));
2589   PetscCall(TSSetTimeStep(ts, dt));
2590   PetscCall(TSSetStepNumber(ts, n));
2591   PetscCall(TSSetMaxSteps(ts, maxn));
2592 
2593   PetscCall(CreateSolution(ts));
2594   PetscCall(SetProblem(ts));
2595   PetscCall(TSGetSolution(ts, &u));
2596   PetscFunctionReturn(PETSC_SUCCESS);
2597 }
2598 
2599 PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
2600 {
2601   DM        sw, cdm;
2602   PetscInt  Np;
2603   PetscReal low[2], high[2];
2604   AppCtx   *user = (AppCtx *)ctx;
2605 
2606   sw = user->swarm;
2607   PetscCall(DMSwarmGetCellDM(sw, &cdm));
2608   // Get the bounding box so we can equally space the particles
2609   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
2610   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2611   // shift it by h/2 so nothing is initialized directly on a boundary
2612   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
2613   x[1] = 0.;
2614   return PETSC_SUCCESS;
2615 }
2616 
2617 /*
2618   InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.
2619 
2620   Input Parameters:
2621 + ts         - The TS
2622 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem
2623 
2624   Output Parameters:
2625 . u - The initialized solution vector
2626 
2627   Level: advanced
2628 
2629 .seealso: InitializeSolve()
2630 */
2631 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
2632 {
2633   DM       sw;
2634   Vec      u, gc, gv;
2635   IS       isx, isv;
2636   PetscInt dim;
2637   AppCtx  *user;
2638 
2639   PetscFunctionBeginUser;
2640   PetscCall(TSGetDM(ts, &sw));
2641   PetscCall(DMGetApplicationContext(sw, &user));
2642   PetscCall(DMGetDimension(sw, &dim));
2643   if (useInitial) {
2644     PetscReal v0[2] = {1., 0.};
2645     if (user->perturbed_weights) {
2646       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
2647     } else {
2648       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
2649       PetscCall(DMSwarmInitializeCoordinates(sw));
2650       PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
2651     }
2652     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2653     PetscCall(DMSwarmTSRedistribute(ts));
2654   }
2655   PetscCall(DMSetUp(sw));
2656   PetscCall(TSGetSolution(ts, &u));
2657   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2658   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2659   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2660   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2661   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
2662   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
2663   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2664   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2665   PetscFunctionReturn(PETSC_SUCCESS);
2666 }
2667 
2668 static PetscErrorCode InitializeSolve(TS ts, Vec u)
2669 {
2670   PetscFunctionBegin;
2671   PetscCall(TSSetSolution(ts, u));
2672   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
2673   PetscFunctionReturn(PETSC_SUCCESS);
2674 }
2675 
2676 static PetscErrorCode MigrateParticles(TS ts)
2677 {
2678   DM               sw, cdm;
2679   const PetscReal *L;
2680   AppCtx          *ctx;
2681 
2682   PetscFunctionBeginUser;
2683   PetscCall(TSGetDM(ts, &sw));
2684   PetscCall(DMGetApplicationContext(sw, &ctx));
2685   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
2686   {
2687     Vec        u, gc, gv, position, momentum;
2688     IS         isx, isv;
2689     PetscReal *pos, *mom;
2690 
2691     PetscCall(TSGetSolution(ts, &u));
2692     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2693     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2694     PetscCall(VecGetSubVector(u, isx, &position));
2695     PetscCall(VecGetSubVector(u, isv, &momentum));
2696     PetscCall(VecGetArray(position, &pos));
2697     PetscCall(VecGetArray(momentum, &mom));
2698     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2699     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2700     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
2701     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
2702 
2703     PetscCall(DMSwarmGetCellDM(sw, &cdm));
2704     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
2705     PetscCheck(L, PetscObjectComm((PetscObject)cdm), PETSC_ERR_ARG_WRONG, "Mesh must be periodic");
2706     if ((L[0] || L[1]) >= 0.) {
2707       PetscReal *x, *v, upper[3], lower[3];
2708       PetscInt   Np, dim;
2709 
2710       PetscCall(DMSwarmGetLocalSize(sw, &Np));
2711       PetscCall(DMGetDimension(cdm, &dim));
2712       PetscCall(DMGetBoundingBox(cdm, lower, upper));
2713       PetscCall(VecGetArray(gc, &x));
2714       PetscCall(VecGetArray(gv, &v));
2715       for (PetscInt p = 0; p < Np; ++p) {
2716         for (PetscInt d = 0; d < dim; ++d) {
2717           if (pos[p * dim + d] < lower[d]) {
2718             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
2719           } else if (pos[p * dim + d] > upper[d]) {
2720             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
2721           } else {
2722             x[p * dim + d] = pos[p * dim + d];
2723           }
2724           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]);
2725           v[p * dim + d] = mom[p * dim + d];
2726         }
2727       }
2728       PetscCall(VecRestoreArray(gc, &x));
2729       PetscCall(VecRestoreArray(gv, &v));
2730     }
2731     PetscCall(VecRestoreArray(position, &pos));
2732     PetscCall(VecRestoreArray(momentum, &mom));
2733     PetscCall(VecRestoreSubVector(u, isx, &position));
2734     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2735     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2736     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2737   }
2738   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2739   PetscInt step;
2740 
2741   PetscCall(TSGetStepNumber(ts, &step));
2742   if (!(step % ctx->remapFreq)) {
2743     // Monitor electric field before we destroy it
2744     PetscReal ptime;
2745     PetscInt  step;
2746 
2747     PetscCall(TSGetStepNumber(ts, &step));
2748     PetscCall(TSGetTime(ts, &ptime));
2749     if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx));
2750     if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx));
2751     PetscCall(DMSwarmRemap(sw));
2752     ctx->validE = PETSC_FALSE;
2753   }
2754   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
2755   PetscCall(DMSwarmTSRedistribute(ts));
2756   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2757   PetscFunctionReturn(PETSC_SUCCESS);
2758 }
2759 
2760 int main(int argc, char **argv)
2761 {
2762   DM        dm, sw;
2763   TS        ts;
2764   Vec       u;
2765   PetscReal dt;
2766   PetscInt  maxn;
2767   AppCtx    user;
2768 
2769   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2770   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2771   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2772   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2773   PetscCall(CreatePoisson(dm, &user));
2774   PetscCall(CreateMomentFields(dm, &user));
2775   PetscCall(CreateSwarm(dm, &user, &sw));
2776   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2777   PetscCall(InitializeConstants(sw, &user));
2778   PetscCall(DMSetApplicationContext(sw, &user));
2779 
2780   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2781   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2782   PetscCall(TSSetDM(ts, sw));
2783   PetscCall(TSSetMaxTime(ts, 0.1));
2784   PetscCall(TSSetTimeStep(ts, 0.00001));
2785   PetscCall(TSSetMaxSteps(ts, 100));
2786   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
2787 
2788   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2789   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
2790   if (user.moment_field_monitor) PetscCall(TSMonitorSet(ts, MonitorMomentFields, &user, NULL));
2791   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2792   if (user.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2793   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));
2794   if (user.velocity_monitor >= 0) PetscCall(TSMonitorSet(ts, MonitorVelocity, &user, NULL));
2795 
2796   PetscCall(TSSetFromOptions(ts));
2797   PetscCall(TSGetTimeStep(ts, &dt));
2798   PetscCall(TSGetMaxSteps(ts, &maxn));
2799   user.steps    = maxn;
2800   user.stepSize = dt;
2801   PetscCall(SetupContext(dm, sw, &user));
2802   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2803   PetscCall(TSSetPostStep(ts, MigrateParticles));
2804   PetscCall(CreateSolution(ts));
2805   PetscCall(TSGetSolution(ts, &u));
2806   PetscCall(TSComputeInitialCondition(ts, u));
2807   PetscCall(CheckNonNegativeWeights(sw, &user));
2808   PetscCall(TSSolve(ts, NULL));
2809 
2810   if (user.checkLandau) {
2811     // We should get a lookup table based on charge density and \hat k
2812     const PetscReal gammaEx = -0.15336;
2813     const PetscReal omegaEx = 1.4156;
2814     const PetscReal tol     = 1e-2;
2815 
2816     PetscCheck(PetscAbsReal((user.gamma - gammaEx) / gammaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau gamma %g != %g", user.gamma, gammaEx);
2817     PetscCheck(PetscAbsReal((user.omega - omegaEx) / omegaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau omega %g != %g", user.omega, omegaEx);
2818   }
2819 
2820   PetscCall(SNESDestroy(&user.snes));
2821   PetscCall(DMDestroy(&user.dmN));
2822   PetscCall(ISDestroy(&user.isN));
2823   PetscCall(MatDestroy(&user.MN));
2824   PetscCall(DMDestroy(&user.dmP));
2825   PetscCall(ISDestroy(&user.isP));
2826   PetscCall(MatDestroy(&user.MP));
2827   PetscCall(DMDestroy(&user.dmE));
2828   PetscCall(ISDestroy(&user.isE));
2829   PetscCall(MatDestroy(&user.ME));
2830   PetscCall(DMDestroy(&user.dmMom));
2831   PetscCall(DMDestroy(&user.dmPot));
2832   PetscCall(ISDestroy(&user.isPot));
2833   PetscCall(MatDestroy(&user.M));
2834   PetscCall(PetscFEGeomDestroy(&user.fegeom));
2835   PetscCall(TSDestroy(&ts));
2836   PetscCall(DMDestroy(&sw));
2837   PetscCall(DMDestroy(&dm));
2838   PetscCall(DestroyContext(&user));
2839   PetscCall(PetscFinalize());
2840   return 0;
2841 }
2842 
2843 /*TEST
2844 
2845   build:
2846     requires: !complex double
2847 
2848   # This tests that we can compute the correct decay rate and frequency
2849   #   For gold runs, use -dm_plex_box_faces 160 -vdm_plex_box_faces 450 -remap_dm_plex_box_faces 80,150
2850   test:
2851     suffix: 0
2852     args: -cosine_coefficients 0.01 -charges -1. -perturbed_weights -total_weight 1. \
2853           -dm_plex_dim 1 -dm_plex_box_faces 80 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 \
2854             -dm_plex_box_bd periodic -dm_plex_hash_location \
2855           -vdm_plex_dim 1 -vdm_plex_box_faces 220 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 \
2856             -vpetscspace_degree 2 -vdm_plex_hash_location \
2857           -remap_freq 100 -dm_swarm_remap_type pfak -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 \
2858             -remap_dm_plex_box_faces 40,110 -remap_dm_plex_box_bd periodic,none \
2859             -remap_dm_plex_box_lower 0.,-6. -remap_dm_plex_box_upper 12.5664,6. \
2860             -remap_petscspace_degree 1 -remap_dm_plex_hash_location \
2861             -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu \
2862           -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged \
2863             -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu \
2864           -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_dt 0.03 -ts_max_steps 500 \
2865             -ts_max_time 100 \
2866           -emax_tao_type brgn -emax_tao_max_it 100 -emax_tao_brgn_regularization_type l2pure \
2867             -emax_tao_brgn_regularizer_weight 1e-5 -tao_brgn_subsolver_tao_bnk_ksp_rtol 1e-12 \
2868           -output_step 1 -efield_monitor quiet
2869 
2870 TEST*/
2871