xref: /petsc/src/ts/tutorials/hamiltonian/ex2.c (revision fcb342c8126fa5278aa1d9f32bf8eea14700142e)
1 static char help[] = "Landau Damping/Two Stream instability test using Vlasov-Poisson equations\n";
2 
3 /*
4   TODO:
5   - Cache mesh geometry
6   - Move electrostatic solver to MG+SVD
7 
8   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"
9   According to Lukas, good damping results come at ~16k particles per cell
10 
11   To visualize the maximum electric field use
12 
13     -efield_monitor
14 
15   To monitor velocity moments of the distribution use
16 
17     -ptof_pc_type lu -moments_monitor
18 
19   To monitor the particle positions in phase space use
20 
21     -positions_monitor
22 
23   To monitor the charge density, E field, and potential use
24 
25     -poisson_monitor
26 
27   To monitor the remapping field use
28 
29     -remap_uf_view draw
30 
31   To visualize the swarm distribution use
32 
33     -ts_monitor_hg_swarm
34 
35   To visualize the particles, we can use
36 
37     -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
38 
39 For a Landau Damping verification run, we use
40 
41   # Physics
42   -cosine_coefficients 0.01 -dm_swarm_num_species 1 -charges -1. -perturbed_weights -total_weight 1.
43   # Spatial Mesh
44   -dm_plex_dim 1 -dm_plex_box_faces 40 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic  -dm_plex_hash_location
45   # Velocity Mesh
46   -vdm_plex_dim 1 -vdm_plex_box_faces 40 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 -vpetscspace_degree 2 -vdm_plex_hash_location
47   # Remap Space
48   -dm_swarm_remap_type pfak -remap_freq 100
49   -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 -remap_dm_plex_box_faces 20,20 -remap_dm_plex_box_bd periodic,none -remap_dm_plex_box_lower 0.,-6.
50     -remap_dm_plex_box_upper 12.5664,6. -remap_petscspace_degree 1 -remap_dm_plex_hash_location
51   # Remap Solve
52   -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu
53   # EM Solve
54   -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu
55   # Timestepping
56   -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_dt 0.03 -ts_max_steps 1500 -ts_max_time 100
57   # Monitoring
58   -output_step 1 -check_vel_res -efield_monitor -poisson_monitor -positions_monitor -dm_swarm_print_coords 0 -remap_uf_view draw
59     -ftop_ksp_lsqr_monitor -ftop_ksp_converged_reason
60 
61 */
62 #include <petscts.h>
63 #include <petscdmplex.h>
64 #include <petscdmswarm.h>
65 #include <petscfe.h>
66 #include <petscds.h>
67 #include <petscbag.h>
68 #include <petscdraw.h>
69 #include <petsc/private/dmpleximpl.h>  /* For norm and dot */
70 #include <petsc/private/petscfeimpl.h> /* For interpolation */
71 #include <petsc/private/dmswarmimpl.h> /* For swarm debugging */
72 #include "petscdm.h"
73 #include "petscdmlabel.h"
74 
75 PETSC_EXTERN PetscErrorCode stream(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
76 PETSC_EXTERN PetscErrorCode line(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
77 
78 const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
79 typedef enum {
80   EM_PRIMAL,
81   EM_MIXED,
82   EM_COULOMB,
83   EM_NONE
84 } EMType;
85 
86 typedef enum {
87   V0,
88   X0,
89   T0,
90   M0,
91   Q0,
92   PHI0,
93   POISSON,
94   VLASOV,
95   SIGMA,
96   NUM_CONSTANTS
97 } ConstantType;
98 typedef struct {
99   PetscScalar v0; /* Velocity scale, often the thermal velocity */
100   PetscScalar t0; /* Time scale */
101   PetscScalar x0; /* Space scale */
102   PetscScalar m0; /* Mass scale */
103   PetscScalar q0; /* Charge scale */
104   PetscScalar kb;
105   PetscScalar epsi0;
106   PetscScalar phi0;          /* Potential scale */
107   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
108   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
109   PetscReal   sigma;         /* Nondimensional charge per length in x */
110 } Parameter;
111 
112 typedef struct {
113   PetscBag     bag;               // Problem parameters
114   PetscBool    error;             // Flag for printing the error
115   PetscInt     remapFreq;         // Number of timesteps between remapping
116   PetscBool    efield_monitor;    // Flag to show electric field monitor
117   PetscBool    moment_monitor;    // Flag to show distribution moment monitor
118   PetscBool    positions_monitor; // Flag to show particle positins at each time step
119   PetscBool    poisson_monitor;   // Flag to display charge, E field, and potential at each solve
120   PetscBool    initial_monitor;   // Flag to monitor the initial conditions
121   PetscBool    perturbed_weights; // Uniformly sample x,v space with gaussian weights
122   PetscInt     ostep;             // Print the energy at each ostep time steps
123   PetscInt     numParticles;
124   PetscReal    timeScale;              /* Nondimensionalizing time scale */
125   PetscReal    charges[2];             /* The charges of each species */
126   PetscReal    masses[2];              /* The masses of each species */
127   PetscReal    thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
128   PetscReal    cosine_coefficients[2]; /*(alpha, k)*/
129   PetscReal    totalWeight;
130   PetscReal    stepSize;
131   PetscInt     steps;
132   PetscReal    initVel;
133   EMType       em;     // Type of electrostatic model
134   SNES         snes;   // EM solver
135   Mat          M;      // The finite element mass matrix
136   PetscFEGeom *fegeom; // Geometric information for the DM cells
137   PetscDraw    drawic_x;
138   PetscDraw    drawic_v;
139   PetscDraw    drawic_w;
140   PetscDrawHG  drawhgic_x;
141   PetscDrawHG  drawhgic_v;
142   PetscDrawHG  drawhgic_w;
143   PetscBool    validE;     // Flag to indicate E-field in swarm is valid
144   PetscReal    drawlgEmin; // The minimum lg(E) to plot
145   PetscDrawLG  drawlgE;    // Logarithm of maximum electric field
146   PetscDrawSP  drawspE;    // Electric field at particle positions
147   PetscDrawSP  drawspX;    // Particle positions
148   PetscViewer  viewerRho;  // Charge density viewer
149   PetscViewer  viewerPhi;  // Potential viewer
150   DM           swarm;
151   PetscRandom  random;
152   PetscBool    twostream;
153   PetscBool    checkweights;
154   PetscInt     checkVRes; /* Flag to check/output velocity residuals for nightly tests */
155 
156   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
157 } AppCtx;
158 
159 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
160 {
161   PetscFunctionBeginUser;
162   PetscInt d                      = 2;
163   PetscInt maxSpecies             = 2;
164   options->error                  = PETSC_FALSE;
165   options->remapFreq              = 1;
166   options->efield_monitor         = PETSC_FALSE;
167   options->moment_monitor         = PETSC_FALSE;
168   options->initial_monitor        = PETSC_FALSE;
169   options->perturbed_weights      = PETSC_FALSE;
170   options->poisson_monitor        = PETSC_FALSE;
171   options->positions_monitor      = PETSC_FALSE;
172   options->ostep                  = 100;
173   options->timeScale              = 2.0e-14;
174   options->charges[0]             = -1.0;
175   options->charges[1]             = 1.0;
176   options->masses[0]              = 1.0;
177   options->masses[1]              = 1000.0;
178   options->thermal_energy[0]      = 1.0;
179   options->thermal_energy[1]      = 1.0;
180   options->cosine_coefficients[0] = 0.01;
181   options->cosine_coefficients[1] = 0.5;
182   options->initVel                = 1;
183   options->totalWeight            = 1.0;
184   options->drawic_x               = NULL;
185   options->drawic_v               = NULL;
186   options->drawic_w               = NULL;
187   options->drawhgic_x             = NULL;
188   options->drawhgic_v             = NULL;
189   options->drawhgic_w             = NULL;
190   options->drawlgEmin             = -6;
191   options->drawlgE                = NULL;
192   options->drawspE                = NULL;
193   options->drawspX                = NULL;
194   options->viewerRho              = NULL;
195   options->viewerPhi              = NULL;
196   options->em                     = EM_COULOMB;
197   options->numParticles           = 32768;
198   options->twostream              = PETSC_FALSE;
199   options->checkweights           = PETSC_FALSE;
200   options->checkVRes              = 0;
201 
202   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
203   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex2.c", options->error, &options->error, NULL));
204   PetscCall(PetscOptionsInt("-remap_freq", "Number", "ex2.c", options->remapFreq, &options->remapFreq, NULL));
205   PetscCall(PetscOptionsBool("-efield_monitor", "Flag to plot log(max E) over time", "ex2.c", options->efield_monitor, &options->efield_monitor, NULL));
206   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", "ex2.c", options->drawlgEmin, &options->drawlgEmin, NULL));
207   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", "ex2.c", options->moment_monitor, &options->moment_monitor, NULL));
208   PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", "ex2.c", options->initial_monitor, &options->initial_monitor, NULL));
209   PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", "ex2.c", options->positions_monitor, &options->positions_monitor, NULL));
210   PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", "ex2.c", options->poisson_monitor, &options->poisson_monitor, NULL));
211   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", "ex2.c", options->twostream, &options->twostream, NULL));
212   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", "ex2.c", options->perturbed_weights, &options->perturbed_weights, NULL));
213   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex2.c", options->checkweights, &options->checkweights, NULL));
214   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex2.c", options->checkVRes, &options->checkVRes, NULL));
215   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex2.c", options->ostep, &options->ostep, NULL));
216   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex2.c", options->timeScale, &options->timeScale, NULL));
217   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex2.c", options->initVel, &options->initVel, NULL));
218   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex2.c", options->totalWeight, &options->totalWeight, NULL));
219   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex2.c", options->cosine_coefficients, &d, NULL));
220   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex2.c", options->charges, &maxSpecies, NULL));
221   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex2.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
222   PetscOptionsEnd();
223 
224   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
225   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
226   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
227   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
228   PetscFunctionReturn(PETSC_SUCCESS);
229 }
230 
231 static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
232 {
233   PetscFunctionBeginUser;
234   if (user->efield_monitor) {
235     PetscDraw     draw;
236     PetscDrawAxis axis;
237 
238     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
239     PetscCall(PetscDrawSetSave(draw, "ex9_Efield"));
240     PetscCall(PetscDrawSetFromOptions(draw));
241     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
242     PetscCall(PetscDrawDestroy(&draw));
243     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
244     PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
245     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
246   }
247   if (user->initial_monitor) {
248     PetscDrawAxis axis1, axis2, axis3;
249     PetscReal     dmboxlower[2], dmboxupper[2];
250     PetscInt      dim, cStart, cEnd;
251     PetscCall(DMGetDimension(sw, &dim));
252     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
253     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
254 
255     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &user->drawic_x));
256     PetscCall(PetscDrawSetSave(user->drawic_x, "ex9_ic_x"));
257     PetscCall(PetscDrawSetFromOptions(user->drawic_x));
258     PetscCall(PetscDrawHGCreate(user->drawic_x, (int)dim, &user->drawhgic_x));
259     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
260     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart)));
261     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "counts"));
262     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 1500));
263 
264     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &user->drawic_v));
265     PetscCall(PetscDrawSetSave(user->drawic_v, "ex9_ic_v"));
266     PetscCall(PetscDrawSetFromOptions(user->drawic_v));
267     PetscCall(PetscDrawHGCreate(user->drawic_v, (int)dim, &user->drawhgic_v));
268     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
269     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 1000));
270     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "counts"));
271     PetscCall(PetscDrawAxisSetLimits(axis2, -1, 1, 0, 1500));
272 
273     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_w", 800, 300, 400, 300, &user->drawic_w));
274     PetscCall(PetscDrawSetSave(user->drawic_w, "ex9_ic_w"));
275     PetscCall(PetscDrawSetFromOptions(user->drawic_w));
276     PetscCall(PetscDrawHGCreate(user->drawic_w, (int)dim, &user->drawhgic_w));
277     PetscCall(PetscDrawHGGetAxis(user->drawhgic_w, &axis3));
278     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_w, 10));
279     PetscCall(PetscDrawAxisSetLabels(axis3, "Initial W Distribution", "weight", "counts"));
280     PetscCall(PetscDrawAxisSetLimits(axis3, 0, 0.01, 0, 5000));
281   }
282   if (user->positions_monitor) {
283     PetscDraw     draw;
284     PetscDrawAxis axis;
285 
286     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Particle Position", 0, 0, 400, 300, &draw));
287     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
288     PetscCall(PetscDrawSetFromOptions(draw));
289     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspX));
290     PetscCall(PetscDrawDestroy(&draw));
291     PetscCall(PetscDrawSPSetDimension(user->drawspX, 1));
292     PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis));
293     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
294     PetscCall(PetscDrawSPReset(user->drawspX));
295   }
296   if (user->poisson_monitor) {
297     Vec           rho, phi;
298     PetscDraw     draw;
299     PetscDrawAxis axis;
300 
301     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
302     PetscCall(PetscDrawSetFromOptions(draw));
303     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
304     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspE));
305     PetscCall(PetscDrawDestroy(&draw));
306     PetscCall(PetscDrawSPSetDimension(user->drawspE, 1));
307     PetscCall(PetscDrawSPGetAxis(user->drawspE, &axis));
308     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
309     PetscCall(PetscDrawSPReset(user->drawspE));
310 
311     PetscCall(PetscViewerDrawOpen(PETSC_COMM_WORLD, NULL, "Charge Density", 0, 0, 400, 300, &user->viewerRho));
312     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRho, "rho_"));
313     PetscCall(PetscViewerDrawGetDraw(user->viewerRho, 0, &draw));
314     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
315     PetscCall(PetscViewerSetFromOptions(user->viewerRho));
316     PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
317     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
318     PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));
319 
320     PetscCall(PetscViewerDrawOpen(PETSC_COMM_WORLD, NULL, "Potential", 400, 0, 400, 300, &user->viewerPhi));
321     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPhi, "phi_"));
322     PetscCall(PetscViewerDrawGetDraw(user->viewerPhi, 0, &draw));
323     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
324     PetscCall(PetscViewerSetFromOptions(user->viewerPhi));
325     PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
326     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
327     PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
328   }
329   PetscFunctionReturn(PETSC_SUCCESS);
330 }
331 
332 static PetscErrorCode DestroyContext(AppCtx *user)
333 {
334   PetscFunctionBeginUser;
335   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
336   PetscCall(PetscDrawDestroy(&user->drawic_x));
337   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
338   PetscCall(PetscDrawDestroy(&user->drawic_v));
339   PetscCall(PetscDrawHGDestroy(&user->drawhgic_w));
340   PetscCall(PetscDrawDestroy(&user->drawic_w));
341 
342   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
343   PetscCall(PetscDrawSPDestroy(&user->drawspE));
344   PetscCall(PetscDrawSPDestroy(&user->drawspX));
345   PetscCall(PetscViewerDestroy(&user->viewerRho));
346   PetscCall(PetscViewerDestroy(&user->viewerPhi));
347 
348   PetscCall(PetscBagDestroy(&user->bag));
349   PetscFunctionReturn(PETSC_SUCCESS);
350 }
351 
352 static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
353 {
354   const PetscScalar *w;
355   PetscInt           Np;
356 
357   PetscFunctionBeginUser;
358   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
359   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
360   PetscCall(DMSwarmGetLocalSize(sw, &Np));
361   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]);
362   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
363   PetscFunctionReturn(PETSC_SUCCESS);
364 }
365 
366 static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *user)
367 {
368   DMSwarmCellDM celldm;
369   DM            vdm;
370   Vec           u[1];
371   const char   *fields[1] = {"w_q"};
372 
373   PetscFunctionBegin;
374   PetscCall(DMSwarmSetCellDMActive(sw, "velocity"));
375   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
376   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
377   PetscCall(DMGetGlobalVector(vdm, &u[0]));
378   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
379   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
380   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
381   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
382   PetscFunctionReturn(PETSC_SUCCESS);
383 }
384 
385 static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
386 {
387   AppCtx    *user = (AppCtx *)ctx;
388   DM         sw;
389   PetscReal *E, *x, *weight;
390   PetscReal  Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
391   PetscReal  pmoments[4]; /* \int f, \int v f, \int v^2 f */
392   PetscInt  *species, dim, Np;
393 
394   PetscFunctionBeginUser;
395   if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
396   PetscCall(TSGetDM(ts, &sw));
397   PetscCall(DMGetDimension(sw, &dim));
398   PetscCall(DMSwarmGetLocalSize(sw, &Np));
399   PetscCall(DMSwarmSortGetAccess(sw));
400   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
401   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
402   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
403   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
404 
405   for (PetscInt p = 0; p < Np; ++p) {
406     for (PetscInt d = 0; d < 1; ++d) {
407       PetscReal temp = PetscAbsReal(E[p * dim + d]);
408       if (temp > Emax) Emax = temp;
409     }
410     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
411     sum += E[p * dim];
412     chargesum += user->charges[0] * weight[p];
413   }
414   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
415   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;
416 
417   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
418   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
419   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
420   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
421   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
422   PetscCall(PetscDrawLGDraw(user->drawlgE));
423   PetscDraw draw;
424   PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
425   PetscCall(PetscDrawSave(draw));
426 
427   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
428   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%+e\t%e\t%f\t%f\t%f\t%f\t%f\t%f\t%f\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]));
429   PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
430   PetscFunctionReturn(PETSC_SUCCESS);
431 }
432 
433 static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
434 {
435   AppCtx   *user = (AppCtx *)ctx;
436   DM        sw;
437   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */
438 
439   PetscFunctionBeginUser;
440   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
441   PetscCall(TSGetDM(ts, &sw));
442 
443   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
444   PetscCall(computeVelocityFEMMoments(sw, fmoments, user));
445 
446   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]));
447   PetscFunctionReturn(PETSC_SUCCESS);
448 }
449 
450 PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
451 {
452   AppCtx            *user = (AppCtx *)ctx;
453   DM                 dm, sw;
454   const PetscScalar *u;
455   PetscReal         *weight, *pos, *vel;
456   PetscInt           dim, p, Np, cStart, cEnd;
457 
458   PetscFunctionBegin;
459   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
460   PetscCall(TSGetDM(ts, &sw));
461   PetscCall(DMSwarmGetCellDM(sw, &dm));
462   PetscCall(DMGetDimension(sw, &dim));
463   PetscCall(DMSwarmGetLocalSize(sw, &Np));
464   PetscCall(DMSwarmSortGetAccess(sw));
465   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
466 
467   if (step == 0) {
468     PetscCall(PetscDrawHGReset(user->drawhgic_x));
469     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &user->drawic_x));
470     PetscCall(PetscDrawClear(user->drawic_x));
471     PetscCall(PetscDrawFlush(user->drawic_x));
472 
473     PetscCall(PetscDrawHGReset(user->drawhgic_v));
474     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &user->drawic_v));
475     PetscCall(PetscDrawClear(user->drawic_v));
476     PetscCall(PetscDrawFlush(user->drawic_v));
477 
478     PetscCall(PetscDrawHGReset(user->drawhgic_w));
479     PetscCall(PetscDrawHGGetDraw(user->drawhgic_w, &user->drawic_w));
480     PetscCall(PetscDrawClear(user->drawic_w));
481     PetscCall(PetscDrawFlush(user->drawic_w));
482 
483     PetscCall(VecGetArrayRead(U, &u));
484     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
485     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
486     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
487 
488     PetscCall(VecGetLocalSize(U, &Np));
489     Np /= dim * 2;
490     for (p = 0; p < Np; ++p) {
491       PetscCall(PetscDrawHGAddValue(user->drawhgic_x, pos[p * dim]));
492       PetscCall(PetscDrawHGAddValue(user->drawhgic_v, vel[p * dim]));
493       PetscCall(PetscDrawHGAddValue(user->drawhgic_w, weight[p]));
494     }
495 
496     PetscCall(VecRestoreArrayRead(U, &u));
497     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
498     PetscCall(PetscDrawHGSave(user->drawhgic_x));
499 
500     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
501     PetscCall(PetscDrawHGSave(user->drawhgic_v));
502 
503     PetscCall(PetscDrawHGDraw(user->drawhgic_w));
504     PetscCall(PetscDrawHGSave(user->drawhgic_w));
505 
506     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
507     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
508     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
509   }
510   PetscFunctionReturn(PETSC_SUCCESS);
511 }
512 
513 static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
514 {
515   AppCtx         *user = (AppCtx *)ctx;
516   DM              dm, sw;
517   PetscScalar    *x, *v, *weight;
518   PetscReal       lower[3], upper[3], speed;
519   const PetscInt *s;
520   PetscInt        dim, cStart, cEnd, c;
521 
522   PetscFunctionBeginUser;
523   if (step > 0 && step % user->ostep == 0) {
524     PetscCall(TSGetDM(ts, &sw));
525     PetscCall(DMSwarmGetCellDM(sw, &dm));
526     PetscCall(DMGetDimension(dm, &dim));
527     PetscCall(DMGetBoundingBox(dm, lower, upper));
528     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
529     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
530     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
531     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
532     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
533     PetscCall(DMSwarmSortGetAccess(sw));
534     PetscCall(PetscDrawSPReset(user->drawspX));
535     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], lower[1], upper[1]));
536     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], -12, 12));
537     for (c = 0; c < cEnd - cStart; ++c) {
538       PetscInt *pidx, Npc, q;
539       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
540       for (q = 0; q < Npc; ++q) {
541         const PetscInt p = pidx[q];
542         if (s[p] == 0) {
543           speed = 0.;
544           for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]);
545           speed = PetscSqrtReal(speed);
546           if (dim == 1) {
547             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &v[p * dim], &speed));
548           } else {
549             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
550           }
551         } else if (s[p] == 1) {
552           PetscCall(PetscDrawSPAddPoint(user->drawspX, &x[p * dim], &v[p * dim]));
553         }
554       }
555       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
556     }
557     PetscCall(PetscDrawSPDraw(user->drawspX, PETSC_TRUE));
558     PetscDraw draw;
559     PetscCall(PetscDrawSPGetDraw(user->drawspX, &draw));
560     PetscCall(PetscDrawSave(draw));
561     PetscCall(DMSwarmSortRestoreAccess(sw));
562     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
563     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
564     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
565     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
566   }
567   PetscFunctionReturn(PETSC_SUCCESS);
568 }
569 
570 static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
571 {
572   AppCtx *user = (AppCtx *)ctx;
573   DM      dm, sw;
574 
575   PetscFunctionBeginUser;
576   if (step > 0 && step % user->ostep == 0) {
577     PetscCall(TSGetDM(ts, &sw));
578     PetscCall(DMSwarmGetCellDM(sw, &dm));
579 
580     if (user->validE) {
581       PetscScalar *x, *E, *weight;
582       PetscReal    lower[3], upper[3], xval;
583       PetscDraw    draw;
584       PetscInt     dim, cStart, cEnd;
585 
586       PetscCall(DMGetDimension(dm, &dim));
587       PetscCall(DMGetBoundingBox(dm, lower, upper));
588       PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
589 
590       PetscCall(PetscDrawSPReset(user->drawspE));
591       PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
592       PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
593       PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
594 
595       PetscCall(DMSwarmSortGetAccess(sw));
596       for (PetscInt c = 0; c < cEnd - cStart; ++c) {
597         PetscReal Eavg = 0.0;
598         PetscInt *pidx, Npc;
599 
600         PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
601         for (PetscInt q = 0; q < Npc; ++q) {
602           const PetscInt p = pidx[q];
603           Eavg += E[p * dim];
604         }
605         Eavg /= Npc;
606         xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
607         PetscCall(PetscDrawSPAddPoint(user->drawspE, &xval, &Eavg));
608         PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
609       }
610       PetscCall(PetscDrawSPDraw(user->drawspE, PETSC_TRUE));
611       PetscCall(PetscDrawSPGetDraw(user->drawspE, &draw));
612       PetscCall(PetscDrawSave(draw));
613       PetscCall(DMSwarmSortRestoreAccess(sw));
614       PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
615       PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
616       PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
617     }
618 
619     Vec rho, phi;
620 
621     PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
622     PetscCall(VecView(rho, user->viewerRho));
623     PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));
624 
625     PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
626     PetscCall(VecView(phi, user->viewerPhi));
627     PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
628   }
629   PetscFunctionReturn(PETSC_SUCCESS);
630 }
631 
632 static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
633 {
634   PetscBag   bag;
635   Parameter *p;
636 
637   PetscFunctionBeginUser;
638   /* setup PETSc parameter bag */
639   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
640   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
641   bag = ctx->bag;
642   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
643   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
644   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
645   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
646   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
647   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
648   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
649   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));
650 
651   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
652   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
653   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
654   PetscCall(PetscBagSetFromOptions(bag));
655   {
656     PetscViewer       viewer;
657     PetscViewerFormat format;
658     PetscBool         flg;
659 
660     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
661     if (flg) {
662       PetscCall(PetscViewerPushFormat(viewer, format));
663       PetscCall(PetscBagView(bag, viewer));
664       PetscCall(PetscViewerFlush(viewer));
665       PetscCall(PetscViewerPopFormat(viewer));
666       PetscCall(PetscViewerDestroy(&viewer));
667     }
668   }
669   PetscFunctionReturn(PETSC_SUCCESS);
670 }
671 
672 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
673 {
674   PetscFunctionBeginUser;
675   PetscCall(DMCreate(comm, dm));
676   PetscCall(DMSetType(*dm, DMPLEX));
677   PetscCall(DMSetFromOptions(*dm));
678   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
679   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
680 
681   // Cache the mesh geometry
682   DMField         coordField;
683   IS              cellIS;
684   PetscQuadrature quad;
685   PetscReal      *wt, *pt;
686   PetscInt        cdim, cStart, cEnd;
687 
688   PetscCall(DMGetCoordinateField(*dm, &coordField));
689   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
690   PetscCall(DMGetCoordinateDim(*dm, &cdim));
691   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
692   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
693   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
694   PetscCall(PetscMalloc1(1, &wt));
695   PetscCall(PetscMalloc1(2, &pt));
696   wt[0] = 1.;
697   pt[0] = -1.;
698   pt[1] = -1.;
699   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
700   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FALSE, &user->fegeom));
701   PetscCall(PetscQuadratureDestroy(&quad));
702   PetscCall(ISDestroy(&cellIS));
703   PetscFunctionReturn(PETSC_SUCCESS);
704 }
705 
706 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[])
707 {
708   f0[0] = -constants[SIGMA];
709 }
710 
711 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[])
712 {
713   PetscInt d;
714   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
715 }
716 
717 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[])
718 {
719   PetscInt d;
720   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
721 }
722 
723 static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
724 {
725   *u = 0.0;
726   return PETSC_SUCCESS;
727 }
728 
729 /*
730    /  I   -grad\ / q \ = /0\
731    \-div    0  / \phi/   \f/
732 */
733 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[])
734 {
735   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
736 }
737 
738 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[])
739 {
740   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
741 }
742 
743 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[])
744 {
745   f0[0] += constants[SIGMA];
746   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
747 }
748 
749 /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
750 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[])
751 {
752   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
753 }
754 
755 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[])
756 {
757   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
758 }
759 
760 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[])
761 {
762   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
763 }
764 
765 static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
766 {
767   PetscFE   fephi, feq;
768   PetscDS   ds;
769   PetscBool simplex;
770   PetscInt  dim;
771 
772   PetscFunctionBeginUser;
773   PetscCall(DMGetDimension(dm, &dim));
774   PetscCall(DMPlexIsSimplex(dm, &simplex));
775   if (user->em == EM_MIXED) {
776     DMLabel        label;
777     const PetscInt id = 1;
778 
779     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
780     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
781     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
782     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
783     PetscCall(PetscFECopyQuadrature(feq, fephi));
784     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
785     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
786     PetscCall(DMCreateDS(dm));
787     PetscCall(PetscFEDestroy(&fephi));
788     PetscCall(PetscFEDestroy(&feq));
789 
790     PetscCall(DMGetLabel(dm, "marker", &label));
791     PetscCall(DMGetDS(dm, &ds));
792 
793     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
794     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
795     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
796     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
797     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));
798 
799     PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "wall", label, 1, &id, 0, 0, NULL, (void (*)(void))zero, NULL, NULL, NULL));
800 
801   } else {
802     MatNullSpace nullsp;
803     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
804     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
805     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
806     PetscCall(DMCreateDS(dm));
807     PetscCall(DMGetDS(dm, &ds));
808     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
809     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
810     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
811     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
812     PetscCall(MatNullSpaceDestroy(&nullsp));
813     PetscCall(PetscFEDestroy(&fephi));
814   }
815   PetscFunctionReturn(PETSC_SUCCESS);
816 }
817 
818 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
819 {
820   SNES         snes;
821   Mat          J;
822   MatNullSpace nullSpace;
823 
824   PetscFunctionBeginUser;
825   PetscCall(CreateFEM(dm, user));
826   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
827   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
828   PetscCall(SNESSetDM(snes, dm));
829   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
830   PetscCall(SNESSetFromOptions(snes));
831 
832   PetscCall(DMCreateMatrix(dm, &J));
833   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
834   PetscCall(MatSetNullSpace(J, nullSpace));
835   PetscCall(MatNullSpaceDestroy(&nullSpace));
836   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
837   PetscCall(MatDestroy(&J));
838   PetscCall(DMCreateMassMatrix(dm, dm, &user->M));
839   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
840   user->snes = snes;
841   PetscFunctionReturn(PETSC_SUCCESS);
842 }
843 
844 PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
845 {
846   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
847   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
848   return PETSC_SUCCESS;
849 }
850 PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
851 {
852   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
853   return PETSC_SUCCESS;
854 }
855 
856 PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
857 {
858   const PetscReal alpha = scale ? scale[0] : 0.0;
859   const PetscReal k     = scale ? scale[1] : 1.;
860   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
861   return PETSC_SUCCESS;
862 }
863 
864 PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
865 {
866   const PetscReal alpha = scale ? scale[0] : 0.;
867   const PetscReal k     = scale ? scale[0] : 1.;
868   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
869   return PETSC_SUCCESS;
870 }
871 
872 static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
873 {
874   PetscFE        fe;
875   DMPolytopeType ct;
876   PetscInt       dim, cStart;
877   const char    *prefix = "v";
878 
879   PetscFunctionBegin;
880   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
881   PetscCall(DMSetType(*vdm, DMPLEX));
882   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
883   PetscCall(DMSetFromOptions(*vdm));
884   PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
885   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));
886 
887   PetscCall(DMGetDimension(*vdm, &dim));
888   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
889   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
890   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
891   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
892   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
893   PetscCall(DMCreateDS(*vdm));
894   PetscCall(PetscFEDestroy(&fe));
895   PetscFunctionReturn(PETSC_SUCCESS);
896 }
897 
898 /*
899   InitializeParticles_Centroid - Initialize a regular grid of particles.
900 
901   Input Parameters:
902 + sw      - The `DMSWARM`
903 - force1D - Treat the spatial domain as 1D
904 
905   Notes:
906   This functions sets the species, cellid, spatial coordinate, and velocity fields for all particles.
907 
908   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
909   and velocity meshes.
910 */
911 static PetscErrorCode InitializeParticles_Centroid(DM sw)
912 {
913   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
914   DMSwarmCellDM celldm;
915   DM            xdm, vdm;
916   PetscReal     vmin[3], vmax[3];
917   PetscReal    *x, *v;
918   PetscInt     *species, *cellid;
919   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
920   PetscBool     flg;
921   MPI_Comm      comm;
922   const char   *cellidname;
923 
924   PetscFunctionBegin;
925   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
926 
927   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
928   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
929   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
930   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
931   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
932   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
933   PetscOptionsEnd();
934   debug = swarm->printCoords;
935 
936   PetscCall(DMGetDimension(sw, &dim));
937   PetscCall(DMSwarmGetCellDM(sw, &xdm));
938   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
939 
940   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
941   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
942   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
943 
944   // One particle per centroid on the tensor product grid
945   Npc = (vcEnd - vcStart) * Ns;
946   Np  = (xcEnd - xcStart) * Npc;
947   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
948   if (debug) {
949     PetscInt gNp;
950     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
951     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
952   }
953 
954   // Set species and cellid
955   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
956   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
957   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
958   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
959   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
960     for (PetscInt s = 0; s < Ns; ++s) {
961       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
962         species[p] = s;
963         cellid[p]  = c;
964       }
965     }
966   }
967   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
968   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));
969 
970   // Set particle coordinates
971   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
972   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
973   PetscCall(DMSwarmSortGetAccess(sw));
974   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
975   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
976   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
977     const PetscInt cell = c + xcStart;
978     PetscInt      *pidx, Npc;
979     PetscReal      centroid[3], volume;
980 
981     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
982     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
983     for (PetscInt s = 0; s < Ns; ++s) {
984       for (PetscInt q = 0; q < Npc / Ns; ++q) {
985         const PetscInt p = pidx[q * Ns + s];
986 
987         for (PetscInt d = 0; d < dim; ++d) {
988           x[p * dim + d] = centroid[d];
989           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
990         }
991         if (debug > 1) {
992           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
993           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
994           for (PetscInt d = 0; d < dim; ++d) {
995             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
996             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
997           }
998           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
999           for (PetscInt d = 0; d < dim; ++d) {
1000             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1001             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
1002           }
1003           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
1004         }
1005       }
1006     }
1007     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1008   }
1009   PetscCall(DMSwarmSortRestoreAccess(sw));
1010   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1011   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1012   PetscFunctionReturn(PETSC_SUCCESS);
1013 }
1014 
1015 /*
1016   InitializeWeights - Compute weight for each local particle
1017 
1018   Input Parameters:
1019 + sw          - The `DMSwarm`
1020 . totalWeight - The sum of all particle weights
1021 . func        - The PDF for the particle spatial distribution
1022 - param       - The PDF parameters
1023 
1024   Notes:
1025   The PDF for velocity is assumed to be a Gaussian
1026 
1027   The particle weights are returned in the `w_q` field of `sw`.
1028 */
1029 static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFunc func, const PetscReal param[])
1030 {
1031   DM               xdm, vdm;
1032   DMSwarmCellDM    celldm;
1033   PetscScalar     *weight;
1034   PetscQuadrature  xquad;
1035   const PetscReal *xq, *xwq;
1036   const PetscInt   order = 5;
1037   PetscReal        xi0[3];
1038   PetscReal        xwtot = 0., pwtot = 0.;
1039   PetscInt         xNq;
1040   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
1041   MPI_Comm         comm;
1042   PetscMPIInt      rank;
1043 
1044   PetscFunctionBegin;
1045   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
1046   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1047   PetscCall(DMGetDimension(sw, &dim));
1048   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1049   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1050   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
1051   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1052   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1053   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
1054 
1055   // Setup Quadrature for spatial and velocity weight calculations
1056   PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad));
1057   PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq));
1058   for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0;
1059 
1060   // Integrate the density function to get the weights of particles in each cell
1061   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
1062   PetscCall(DMSwarmSortGetAccess(sw));
1063   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1064   for (PetscInt c = xcStart; c < xcEnd; ++c) {
1065     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
1066     PetscInt          *pidx, Npc;
1067     PetscInt           xNc;
1068     const PetscScalar *xarray;
1069     PetscScalar       *xcoords = NULL;
1070     PetscBool          xisDG;
1071 
1072     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1073     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1074     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);
1075     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
1076     for (PetscInt q = 0; q < xNq; ++q) {
1077       // Transform quadrature points from ref space to real space
1078       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
1079       // Get probability density at quad point
1080       //   No need to scale xqr since PDF will be periodic
1081       PetscCall((*func)(xqr, param, &xden));
1082       xw += xden * (xwq[q] * xdetJ);
1083     }
1084     xwtot += xw;
1085     if (debug) {
1086       IS              globalOrdering;
1087       const PetscInt *ordering;
1088 
1089       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
1090       PetscCall(ISGetIndices(globalOrdering, &ordering));
1091       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));
1092       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
1093     }
1094     // Set weights to be Gaussian in velocity cells
1095     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
1096       const PetscInt     p  = pidx[vc * Ns + 0];
1097       PetscReal          vw = 0.;
1098       PetscInt           vNc;
1099       const PetscScalar *varray;
1100       PetscScalar       *vcoords = NULL;
1101       PetscBool          visDG;
1102 
1103       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1104       // TODO: Fix 2 stream Ask Joe
1105       //   Two stream function from 1/2pi v^2 e^(-v^2/2)
1106       //   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.)));
1107       vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.)));
1108 
1109       weight[p] = totalWeight * vw * xw;
1110       pwtot += weight[p];
1111       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeeded 1: %g, %g, %g", p, xw, vw, totalWeight);
1112       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1113       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
1114     }
1115     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1116     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1117   }
1118   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1119   PetscCall(DMSwarmSortRestoreAccess(sw));
1120   PetscCall(PetscQuadratureDestroy(&xquad));
1121 
1122   if (debug) {
1123     PetscReal wtot[2] = {pwtot, xwtot}, gwtot[2];
1124 
1125     PetscCall(PetscSynchronizedFlush(comm, NULL));
1126     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1127     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
1128   }
1129   PetscFunctionReturn(PETSC_SUCCESS);
1130 }
1131 
1132 static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
1133 {
1134   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
1135   PetscInt  dim;
1136 
1137   PetscFunctionBegin;
1138   PetscCall(DMGetDimension(sw, &dim));
1139   PetscCall(InitializeParticles_Centroid(sw));
1140   PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
1141   PetscFunctionReturn(PETSC_SUCCESS);
1142 }
1143 
1144 static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
1145 {
1146   DM         dm;
1147   PetscInt  *species;
1148   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
1149   PetscInt   Np, dim;
1150 
1151   PetscFunctionBegin;
1152   PetscCall(DMSwarmGetCellDM(sw, &dm));
1153   PetscCall(DMGetDimension(sw, &dim));
1154   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1155   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
1156   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1157   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1158   for (PetscInt p = 0; p < Np; ++p) {
1159     totalWeight += weight[p];
1160     totalCharge += user->charges[species[p]] * weight[p];
1161   }
1162   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1163   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1164   {
1165     Parameter *param;
1166     PetscReal  Area;
1167 
1168     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1169     switch (dim) {
1170     case 1:
1171       Area = (gmax[0] - gmin[0]);
1172       break;
1173     case 2:
1174       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1175       break;
1176     case 3:
1177       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1178       break;
1179     default:
1180       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1181     }
1182     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1183     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1184     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));
1185     param->sigma = PetscAbsReal(global_charge / (Area));
1186 
1187     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
1188     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,
1189                           (double)param->vlasovNumber));
1190   }
1191   /* Setup Constants */
1192   {
1193     PetscDS    ds;
1194     Parameter *param;
1195     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1196     PetscScalar constants[NUM_CONSTANTS];
1197     constants[SIGMA]   = param->sigma;
1198     constants[V0]      = param->v0;
1199     constants[T0]      = param->t0;
1200     constants[X0]      = param->x0;
1201     constants[M0]      = param->m0;
1202     constants[Q0]      = param->q0;
1203     constants[PHI0]    = param->phi0;
1204     constants[POISSON] = param->poissonNumber;
1205     constants[VLASOV]  = param->vlasovNumber;
1206     PetscCall(DMGetDS(dm, &ds));
1207     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
1208   }
1209   PetscFunctionReturn(PETSC_SUCCESS);
1210 }
1211 
1212 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1213 {
1214   DMSwarmCellDM celldm;
1215   DM            vdm;
1216   PetscReal     v0[2] = {1., 0.};
1217   PetscInt      dim;
1218 
1219   PetscFunctionBeginUser;
1220   PetscCall(DMGetDimension(dm, &dim));
1221   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
1222   PetscCall(DMSetType(*sw, DMSWARM));
1223   PetscCall(DMSetDimension(*sw, dim));
1224   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
1225   PetscCall(DMSetApplicationContext(*sw, user));
1226 
1227   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
1228   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
1229   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
1230   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
1231 
1232   const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
1233 
1234   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
1235   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1236   PetscCall(DMSwarmCellDMDestroy(&celldm));
1237 
1238   const char *vfieldnames[1] = {"w_q"};
1239 
1240   PetscCall(CreateVelocityDM(*sw, &vdm));
1241   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
1242   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1243   PetscCall(DMSwarmCellDMDestroy(&celldm));
1244   PetscCall(DMDestroy(&vdm));
1245 
1246   PetscCall(DMSetFromOptions(*sw));
1247   PetscCall(DMSetUp(*sw));
1248 
1249   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
1250   user->swarm = *sw;
1251   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
1252   if (user->perturbed_weights) {
1253     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
1254   } else {
1255     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
1256     PetscCall(DMSwarmInitializeCoordinates(*sw));
1257     PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
1258   }
1259   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
1260   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
1261   PetscFunctionReturn(PETSC_SUCCESS);
1262 }
1263 
1264 static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
1265 {
1266   AppCtx     *user;
1267   PetscReal  *coords;
1268   PetscInt   *species, dim, Np, Ns;
1269   PetscMPIInt size;
1270 
1271   PetscFunctionBegin;
1272   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
1273   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
1274   PetscCall(DMGetDimension(sw, &dim));
1275   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1276   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1277   PetscCall(DMGetApplicationContext(sw, (void *)&user));
1278 
1279   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1280   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1281   for (PetscInt p = 0; p < Np; ++p) {
1282     PetscReal *pcoord = &coords[p * dim];
1283     PetscReal  pE[3]  = {0., 0., 0.};
1284 
1285     /* Calculate field at particle p due to particle q */
1286     for (PetscInt q = 0; q < Np; ++q) {
1287       PetscReal *qcoord = &coords[q * dim];
1288       PetscReal  rpq[3], r, r3, q_q;
1289 
1290       if (p == q) continue;
1291       q_q = user->charges[species[q]] * 1.;
1292       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1293       r = DMPlex_NormD_Internal(dim, rpq);
1294       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
1295       r3 = PetscPowRealInt(r, 3);
1296       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
1297     }
1298     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
1299   }
1300   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1301   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1302   PetscFunctionReturn(PETSC_SUCCESS);
1303 }
1304 
1305 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[])
1306 {
1307   DM         dm;
1308   AppCtx    *user;
1309   PetscDS    ds;
1310   PetscFE    fe;
1311   Mat        M_p;
1312   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
1313   Vec        rho;         // Charge density, M^{-1} rhoRhs
1314   Vec        phi, locPhi; // Potential
1315   Vec        f;           // Particle weights
1316   PetscReal *coords;
1317   PetscInt   dim, cStart, cEnd, Np;
1318 
1319   PetscFunctionBegin;
1320   PetscCall(DMGetApplicationContext(sw, (void *)&user));
1321   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1322   PetscCall(DMGetDimension(sw, &dim));
1323   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1324 
1325   KSP          ksp;
1326   const char **oldFields;
1327   PetscInt     Nf;
1328   const char **tmp;
1329 
1330   PetscCall(SNESGetDM(snes, &dm));
1331   PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp));
1332   PetscCall(PetscMalloc1(Nf, &oldFields));
1333   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f]));
1334   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1335   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
1336   PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields));
1337   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f]));
1338   PetscCall(PetscFree(oldFields));
1339 
1340   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
1341   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
1342   PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
1343   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1344   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1345 
1346   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1347   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1348   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1349 
1350   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1351   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1352 
1353   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1354   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1355   PetscCall(KSPSetOperators(ksp, user->M, user->M));
1356   PetscCall(KSPSetFromOptions(ksp));
1357   PetscCall(KSPSolve(ksp, rhoRhs, rho));
1358 
1359   PetscCall(VecScale(rhoRhs, -1.0));
1360 
1361   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1362   PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));
1363   PetscCall(KSPDestroy(&ksp));
1364   PetscCall(MatDestroy(&M_p));
1365 
1366   PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
1367   PetscCall(VecSet(phi, 0.0));
1368   PetscCall(SNESSolve(snes, rhoRhs, phi));
1369   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
1370   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1371 
1372   PetscCall(DMGetLocalVector(dm, &locPhi));
1373   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1374   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1375   PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
1376   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));
1377 
1378   PetscCall(DMGetDS(dm, &ds));
1379   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1380   PetscCall(DMSwarmSortGetAccess(sw));
1381   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1382   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1383 
1384   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1385   PetscTabulation tab;
1386   PetscReal      *pcoord, *refcoord;
1387   PetscFEGeom    *chunkgeom = NULL;
1388   PetscInt        maxNcp    = 0;
1389 
1390   for (PetscInt c = cStart; c < cEnd; ++c) {
1391     PetscInt Ncp;
1392 
1393     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
1394     maxNcp = PetscMax(maxNcp, Ncp);
1395   }
1396   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1397   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1398   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
1399   for (PetscInt c = cStart; c < cEnd; ++c) {
1400     PetscScalar *clPhi = NULL;
1401     PetscInt    *points;
1402     PetscInt     Ncp;
1403 
1404     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1405     for (PetscInt cp = 0; cp < Ncp; ++cp)
1406       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1407     {
1408       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1409       for (PetscInt i = 0; i < Ncp; ++i) {
1410         const PetscReal x0[3] = {-1., -1., -1.};
1411         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1412       }
1413     }
1414     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
1415     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1416     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1417       const PetscReal *basisDer = tab->T[1];
1418       const PetscInt   p        = points[cp];
1419 
1420       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1421       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
1422       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
1423     }
1424     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1425     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1426   }
1427   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1428   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1429   PetscCall(PetscTabulationDestroy(&tab));
1430   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1431   PetscCall(DMSwarmSortRestoreAccess(sw));
1432   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1433   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1434   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1435   PetscFunctionReturn(PETSC_SUCCESS);
1436 }
1437 
1438 static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
1439 {
1440   AppCtx         *user;
1441   DM              dm, potential_dm;
1442   KSP             ksp;
1443   IS              potential_IS;
1444   PetscDS         ds;
1445   PetscFE         fe;
1446   Mat             M_p, M;
1447   Vec             phi, locPhi, rho, f, temp_rho, rho0;
1448   PetscQuadrature q;
1449   PetscReal      *coords;
1450   PetscInt        dim, cStart, cEnd, Np, pot_field = 1;
1451   const char    **oldFields;
1452   PetscInt        Nf;
1453   const char    **tmp;
1454 
1455   PetscFunctionBegin;
1456   PetscCall(DMGetApplicationContext(sw, &user));
1457   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1458   PetscCall(DMGetDimension(sw, &dim));
1459   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1460 
1461   /* Create the charges rho */
1462   PetscCall(SNESGetDM(snes, &dm));
1463   PetscCall(DMGetGlobalVector(dm, &rho));
1464   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
1465 
1466   PetscCall(DMCreateSubDM(dm, 1, &pot_field, &potential_IS, &potential_dm));
1467 
1468   PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp));
1469   PetscCall(PetscMalloc1(Nf, &oldFields));
1470   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f]));
1471   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1472   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
1473   PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields));
1474   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f]));
1475   PetscCall(PetscFree(oldFields));
1476 
1477   PetscCall(DMCreateMassMatrix(potential_dm, potential_dm, &M));
1478   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1479   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1480   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
1481   PetscCall(PetscObjectSetName((PetscObject)temp_rho, "Mf"));
1482   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1483   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1484   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1485   PetscCall(MatMultTranspose(M_p, f, temp_rho));
1486   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1487   PetscCall(DMGetGlobalVector(potential_dm, &rho0));
1488   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Mixed Compute"));
1489 
1490   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1491   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1492   PetscCall(KSPSetOperators(ksp, M, M));
1493   PetscCall(KSPSetFromOptions(ksp));
1494   PetscCall(KSPSolve(ksp, temp_rho, rho0));
1495   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1496 
1497   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
1498   PetscCall(VecScale(rho, 0.25));
1499   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1500   PetscCall(VecViewFromOptions(temp_rho, NULL, "-temprho_view"));
1501   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1502   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
1503   PetscCall(DMRestoreGlobalVector(potential_dm, &rho0));
1504 
1505   PetscCall(MatDestroy(&M_p));
1506   PetscCall(MatDestroy(&M));
1507   PetscCall(KSPDestroy(&ksp));
1508   PetscCall(DMDestroy(&potential_dm));
1509   PetscCall(ISDestroy(&potential_IS));
1510 
1511   PetscCall(DMGetGlobalVector(dm, &phi));
1512   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1513   PetscCall(VecSet(phi, 0.0));
1514   PetscCall(SNESSolve(snes, rho, phi));
1515   PetscCall(DMRestoreGlobalVector(dm, &rho));
1516 
1517   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1518 
1519   PetscCall(DMGetLocalVector(dm, &locPhi));
1520   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1521   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1522   PetscCall(DMRestoreGlobalVector(dm, &phi));
1523   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));
1524 
1525   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1526   PetscCall(DMGetDS(dm, &ds));
1527   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1528   PetscCall(DMSwarmSortGetAccess(sw));
1529   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1530   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1531   PetscCall(PetscFEGetQuadrature(fe, &q));
1532   PetscFEGeom *chunkgeom = NULL;
1533   for (PetscInt c = cStart; c < cEnd; ++c) {
1534     PetscTabulation tab;
1535     PetscScalar    *clPhi = NULL;
1536     PetscReal      *pcoord, *refcoord;
1537     PetscInt       *points;
1538     PetscInt        Ncp;
1539 
1540     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1541     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1542     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1543     for (PetscInt cp = 0; cp < Ncp; ++cp)
1544       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1545     {
1546       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1547       for (PetscInt i = 0; i < Ncp; ++i) {
1548         // Apply the inverse affine transformation for each point
1549         const PetscReal x0[3] = {-1., -1., -1.};
1550         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1551       }
1552     }
1553     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1554     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1555 
1556     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1557       const PetscInt p = points[cp];
1558 
1559       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1560       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
1561       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
1562       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -2.0;
1563     }
1564     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1565     PetscCall(PetscTabulationDestroy(&tab));
1566     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1567     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1568     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1569   }
1570   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1571   PetscCall(DMSwarmSortRestoreAccess(sw));
1572   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1573   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1574   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1575   PetscFunctionReturn(PETSC_SUCCESS);
1576 }
1577 
1578 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
1579 {
1580   AppCtx  *ctx;
1581   PetscInt dim, Np;
1582 
1583   PetscFunctionBegin;
1584   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
1585   PetscValidHeaderSpecific(sw, DM_CLASSID, 2);
1586   PetscAssertPointer(E, 3);
1587   PetscCall(DMGetDimension(sw, &dim));
1588   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1589   PetscCall(DMGetApplicationContext(sw, &ctx));
1590   PetscCall(PetscArrayzero(E, Np * dim));
1591   ctx->validE = PETSC_TRUE;
1592 
1593   switch (ctx->em) {
1594   case EM_PRIMAL:
1595     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
1596     break;
1597   case EM_COULOMB:
1598     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1599     break;
1600   case EM_MIXED:
1601     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
1602     break;
1603   case EM_NONE:
1604     break;
1605   default:
1606     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
1607   }
1608   PetscFunctionReturn(PETSC_SUCCESS);
1609 }
1610 
1611 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1612 {
1613   DM                 sw;
1614   SNES               snes = ((AppCtx *)ctx)->snes;
1615   const PetscScalar *u;
1616   PetscScalar       *g;
1617   PetscReal         *E, m_p = 1., q_p = -1.;
1618   PetscInt           dim, d, Np, p;
1619 
1620   PetscFunctionBeginUser;
1621   PetscCall(TSGetDM(ts, &sw));
1622   PetscCall(DMGetDimension(sw, &dim));
1623   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1624   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1625   PetscCall(VecGetArrayRead(U, &u));
1626   PetscCall(VecGetArray(G, &g));
1627 
1628   PetscCall(ComputeFieldAtParticles(snes, sw, E));
1629 
1630   Np /= 2 * dim;
1631   for (p = 0; p < Np; ++p) {
1632     for (d = 0; d < dim; ++d) {
1633       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1634       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1635     }
1636   }
1637   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1638   PetscCall(VecRestoreArrayRead(U, &u));
1639   PetscCall(VecRestoreArray(G, &g));
1640   PetscFunctionReturn(PETSC_SUCCESS);
1641 }
1642 
1643 /* J_{ij} = dF_i/dx_j
1644    J_p = (  0   1)
1645          (-w^2  0)
1646    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1647         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1648 */
1649 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1650 {
1651   DM               sw;
1652   const PetscReal *coords, *vel;
1653   PetscInt         dim, d, Np, p, rStart;
1654 
1655   PetscFunctionBeginUser;
1656   PetscCall(TSGetDM(ts, &sw));
1657   PetscCall(DMGetDimension(sw, &dim));
1658   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1659   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1660   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1661   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1662   Np /= 2 * dim;
1663   for (p = 0; p < Np; ++p) {
1664     const PetscReal x0      = coords[p * dim + 0];
1665     const PetscReal vy0     = vel[p * dim + 1];
1666     const PetscReal omega   = vy0 / x0;
1667     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};
1668 
1669     for (d = 0; d < dim; ++d) {
1670       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1671       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1672     }
1673   }
1674   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1675   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1676   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1677   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1678   PetscFunctionReturn(PETSC_SUCCESS);
1679 }
1680 
1681 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1682 {
1683   AppCtx            *user = (AppCtx *)ctx;
1684   DM                 sw;
1685   const PetscScalar *v;
1686   PetscScalar       *xres;
1687   PetscInt           Np, p, d, dim;
1688 
1689   PetscFunctionBeginUser;
1690   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
1691   PetscCall(TSGetDM(ts, &sw));
1692   PetscCall(DMGetDimension(sw, &dim));
1693   PetscCall(VecGetLocalSize(Xres, &Np));
1694   PetscCall(VecGetArrayRead(V, &v));
1695   PetscCall(VecGetArray(Xres, &xres));
1696   Np /= dim;
1697   for (p = 0; p < Np; ++p) {
1698     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
1699   }
1700   PetscCall(VecRestoreArrayRead(V, &v));
1701   PetscCall(VecRestoreArray(Xres, &xres));
1702   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
1703   PetscFunctionReturn(PETSC_SUCCESS);
1704 }
1705 
1706 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1707 {
1708   DM                 sw;
1709   AppCtx            *user = (AppCtx *)ctx;
1710   SNES               snes = ((AppCtx *)ctx)->snes;
1711   const PetscScalar *x;
1712   PetscReal         *E, m_p, q_p;
1713   PetscScalar       *vres;
1714   PetscInt           Np, p, dim, d;
1715   Parameter         *param;
1716 
1717   PetscFunctionBeginUser;
1718   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
1719   PetscCall(TSGetDM(ts, &sw));
1720   PetscCall(DMGetDimension(sw, &dim));
1721   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1722   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1723   m_p = user->masses[0] * param->m0;
1724   q_p = user->charges[0] * param->q0;
1725   PetscCall(VecGetLocalSize(Vres, &Np));
1726   PetscCall(VecGetArrayRead(X, &x));
1727   PetscCall(VecGetArray(Vres, &vres));
1728   PetscCall(ComputeFieldAtParticles(snes, sw, E));
1729 
1730   Np /= dim;
1731   for (p = 0; p < Np; ++p) {
1732     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1733   }
1734   PetscCall(VecRestoreArrayRead(X, &x));
1735   /*
1736     Synchronized, ordered output for parallel/sequential test cases.
1737     In the 1D (on the 2D mesh) case, every y component should be zero.
1738   */
1739   if (user->checkVRes) {
1740     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1741     PetscInt  step;
1742 
1743     PetscCall(TSGetStepNumber(ts, &step));
1744     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
1745     for (PetscInt p = 0; p < Np; ++p) {
1746       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
1747       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]));
1748     }
1749     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1750   }
1751   PetscCall(VecRestoreArray(Vres, &vres));
1752   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1753   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
1754   PetscFunctionReturn(PETSC_SUCCESS);
1755 }
1756 
1757 static PetscErrorCode CreateSolution(TS ts)
1758 {
1759   DM       sw;
1760   Vec      u;
1761   PetscInt dim, Np;
1762 
1763   PetscFunctionBegin;
1764   PetscCall(TSGetDM(ts, &sw));
1765   PetscCall(DMGetDimension(sw, &dim));
1766   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1767   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1768   PetscCall(VecSetBlockSize(u, dim));
1769   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1770   PetscCall(VecSetUp(u));
1771   PetscCall(TSSetSolution(ts, u));
1772   PetscCall(VecDestroy(&u));
1773   PetscFunctionReturn(PETSC_SUCCESS);
1774 }
1775 
1776 static PetscErrorCode SetProblem(TS ts)
1777 {
1778   AppCtx *user;
1779   DM      sw;
1780 
1781   PetscFunctionBegin;
1782   PetscCall(TSGetDM(ts, &sw));
1783   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1784   // Define unified system for (X, V)
1785   {
1786     Mat      J;
1787     PetscInt dim, Np;
1788 
1789     PetscCall(DMGetDimension(sw, &dim));
1790     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1791     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1792     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1793     PetscCall(MatSetBlockSize(J, 2 * dim));
1794     PetscCall(MatSetFromOptions(J));
1795     PetscCall(MatSetUp(J));
1796     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1797     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1798     PetscCall(MatDestroy(&J));
1799   }
1800   /* Define split system for X and V */
1801   {
1802     Vec             u;
1803     IS              isx, isv, istmp;
1804     const PetscInt *idx;
1805     PetscInt        dim, Np, rstart;
1806 
1807     PetscCall(TSGetSolution(ts, &u));
1808     PetscCall(DMGetDimension(sw, &dim));
1809     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1810     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1811     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1812     PetscCall(ISGetIndices(istmp, &idx));
1813     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1814     PetscCall(ISRestoreIndices(istmp, &idx));
1815     PetscCall(ISDestroy(&istmp));
1816     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1817     PetscCall(ISGetIndices(istmp, &idx));
1818     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1819     PetscCall(ISRestoreIndices(istmp, &idx));
1820     PetscCall(ISDestroy(&istmp));
1821     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1822     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1823     PetscCall(ISDestroy(&isx));
1824     PetscCall(ISDestroy(&isv));
1825     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1826     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1827   }
1828   PetscFunctionReturn(PETSC_SUCCESS);
1829 }
1830 
1831 static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1832 {
1833   DM        sw;
1834   Vec       u;
1835   PetscReal t, maxt, dt;
1836   PetscInt  n, maxn;
1837 
1838   PetscFunctionBegin;
1839   PetscCall(TSGetDM(ts, &sw));
1840   PetscCall(TSGetTime(ts, &t));
1841   PetscCall(TSGetMaxTime(ts, &maxt));
1842   PetscCall(TSGetTimeStep(ts, &dt));
1843   PetscCall(TSGetStepNumber(ts, &n));
1844   PetscCall(TSGetMaxSteps(ts, &maxn));
1845 
1846   PetscCall(TSReset(ts));
1847   PetscCall(TSSetDM(ts, sw));
1848   PetscCall(TSSetFromOptions(ts));
1849   PetscCall(TSSetTime(ts, t));
1850   PetscCall(TSSetMaxTime(ts, maxt));
1851   PetscCall(TSSetTimeStep(ts, dt));
1852   PetscCall(TSSetStepNumber(ts, n));
1853   PetscCall(TSSetMaxSteps(ts, maxn));
1854 
1855   PetscCall(CreateSolution(ts));
1856   PetscCall(SetProblem(ts));
1857   PetscCall(TSGetSolution(ts, &u));
1858   PetscFunctionReturn(PETSC_SUCCESS);
1859 }
1860 
1861 PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
1862 {
1863   DM        sw, cdm;
1864   PetscInt  Np;
1865   PetscReal low[2], high[2];
1866   AppCtx   *user = (AppCtx *)ctx;
1867 
1868   sw = user->swarm;
1869   PetscCall(DMSwarmGetCellDM(sw, &cdm));
1870   // Get the bounding box so we can equally space the particles
1871   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
1872   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1873   // shift it by h/2 so nothing is initialized directly on a boundary
1874   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
1875   x[1] = 0.;
1876   return PETSC_SUCCESS;
1877 }
1878 
1879 /*
1880   InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.
1881 
1882   Input Parameters:
1883 + ts         - The TS
1884 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem
1885 
1886   Output Parameters:
1887 . u - The initialized solution vector
1888 
1889   Level: advanced
1890 
1891 .seealso: InitializeSolve()
1892 */
1893 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1894 {
1895   DM       sw;
1896   Vec      u, gc, gv;
1897   IS       isx, isv;
1898   PetscInt dim;
1899   AppCtx  *user;
1900 
1901   PetscFunctionBeginUser;
1902   PetscCall(TSGetDM(ts, &sw));
1903   PetscCall(DMGetApplicationContext(sw, &user));
1904   PetscCall(DMGetDimension(sw, &dim));
1905   if (useInitial) {
1906     PetscReal v0[2] = {1., 0.};
1907     if (user->perturbed_weights) {
1908       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1909     } else {
1910       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
1911       PetscCall(DMSwarmInitializeCoordinates(sw));
1912       PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
1913     }
1914     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1915     PetscCall(DMSwarmTSRedistribute(ts));
1916   }
1917   PetscCall(DMSetUp(sw));
1918   PetscCall(TSGetSolution(ts, &u));
1919   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1920   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1921   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1922   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1923   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1924   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1925   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1926   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1927   PetscFunctionReturn(PETSC_SUCCESS);
1928 }
1929 
1930 static PetscErrorCode InitializeSolve(TS ts, Vec u)
1931 {
1932   PetscFunctionBegin;
1933   PetscCall(TSSetSolution(ts, u));
1934   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1935   PetscFunctionReturn(PETSC_SUCCESS);
1936 }
1937 
1938 static PetscErrorCode MigrateParticles(TS ts)
1939 {
1940   DM               sw, cdm;
1941   const PetscReal *L;
1942   AppCtx          *ctx;
1943 
1944   PetscFunctionBeginUser;
1945   PetscCall(TSGetDM(ts, &sw));
1946   PetscCall(DMGetApplicationContext(sw, &ctx));
1947   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1948   {
1949     Vec        u, gc, gv, position, momentum;
1950     IS         isx, isv;
1951     PetscReal *pos, *mom;
1952 
1953     PetscCall(TSGetSolution(ts, &u));
1954     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1955     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1956     PetscCall(VecGetSubVector(u, isx, &position));
1957     PetscCall(VecGetSubVector(u, isv, &momentum));
1958     PetscCall(VecGetArray(position, &pos));
1959     PetscCall(VecGetArray(momentum, &mom));
1960     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1961     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1962     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1963     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
1964 
1965     PetscCall(DMSwarmGetCellDM(sw, &cdm));
1966     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
1967     if ((L[0] || L[1]) >= 0.) {
1968       PetscReal *x, *v, upper[3], lower[3];
1969       PetscInt   Np, dim;
1970 
1971       PetscCall(DMSwarmGetLocalSize(sw, &Np));
1972       PetscCall(DMGetDimension(cdm, &dim));
1973       PetscCall(DMGetBoundingBox(cdm, lower, upper));
1974       PetscCall(VecGetArray(gc, &x));
1975       PetscCall(VecGetArray(gv, &v));
1976       for (PetscInt p = 0; p < Np; ++p) {
1977         for (PetscInt d = 0; d < dim; ++d) {
1978           if (pos[p * dim + d] < lower[d]) {
1979             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
1980           } else if (pos[p * dim + d] > upper[d]) {
1981             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
1982           } else {
1983             x[p * dim + d] = pos[p * dim + d];
1984           }
1985           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]);
1986           v[p * dim + d] = mom[p * dim + d];
1987         }
1988       }
1989       PetscCall(VecRestoreArray(gc, &x));
1990       PetscCall(VecRestoreArray(gv, &v));
1991     }
1992     PetscCall(VecRestoreArray(position, &pos));
1993     PetscCall(VecRestoreArray(momentum, &mom));
1994     PetscCall(VecRestoreSubVector(u, isx, &position));
1995     PetscCall(VecRestoreSubVector(u, isv, &momentum));
1996     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1997     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1998   }
1999   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2000   PetscInt step;
2001 
2002   PetscCall(TSGetStepNumber(ts, &step));
2003   if (!(step % ctx->remapFreq)) {
2004     // Monitor electric field before we destroy it
2005     PetscReal ptime;
2006     PetscInt  step;
2007 
2008     PetscCall(TSGetStepNumber(ts, &step));
2009     PetscCall(TSGetTime(ts, &ptime));
2010     if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx));
2011     if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx));
2012     PetscCall(DMSwarmRemap(sw));
2013     ctx->validE = PETSC_FALSE;
2014   }
2015   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
2016   PetscCall(DMSwarmTSRedistribute(ts));
2017   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2018   PetscFunctionReturn(PETSC_SUCCESS);
2019 }
2020 
2021 int main(int argc, char **argv)
2022 {
2023   DM        dm, sw;
2024   TS        ts;
2025   Vec       u;
2026   PetscReal dt;
2027   PetscInt  maxn;
2028   AppCtx    user;
2029 
2030   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2031   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2032   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2033   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2034   PetscCall(CreatePoisson(dm, &user));
2035   PetscCall(CreateSwarm(dm, &user, &sw));
2036   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2037   PetscCall(InitializeConstants(sw, &user));
2038   PetscCall(DMSetApplicationContext(sw, &user));
2039 
2040   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2041   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2042   PetscCall(TSSetDM(ts, sw));
2043   PetscCall(TSSetMaxTime(ts, 0.1));
2044   PetscCall(TSSetTimeStep(ts, 0.00001));
2045   PetscCall(TSSetMaxSteps(ts, 100));
2046   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
2047 
2048   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2049   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
2050   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2051   if (user.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2052   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));
2053 
2054   PetscCall(TSSetFromOptions(ts));
2055   PetscCall(TSGetTimeStep(ts, &dt));
2056   PetscCall(TSGetMaxSteps(ts, &maxn));
2057   user.steps    = maxn;
2058   user.stepSize = dt;
2059   PetscCall(SetupContext(dm, sw, &user));
2060   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2061   PetscCall(TSSetPostStep(ts, MigrateParticles));
2062   PetscCall(CreateSolution(ts));
2063   PetscCall(TSGetSolution(ts, &u));
2064   PetscCall(TSComputeInitialCondition(ts, u));
2065   PetscCall(CheckNonNegativeWeights(sw, &user));
2066   PetscCall(TSSolve(ts, NULL));
2067 
2068   PetscCall(SNESDestroy(&user.snes));
2069   PetscCall(MatDestroy(&user.M));
2070   PetscCall(PetscFEGeomDestroy(&user.fegeom));
2071   PetscCall(TSDestroy(&ts));
2072   PetscCall(DMDestroy(&sw));
2073   PetscCall(DMDestroy(&dm));
2074   PetscCall(DestroyContext(&user));
2075   PetscCall(PetscFinalize());
2076   return 0;
2077 }
2078 
2079 /*TEST
2080 
2081    build:
2082     requires: !complex double
2083 
2084   # This tests that we can put particles in a box and compute the Coulomb force
2085   # Recommend -draw_size 500,500
2086    testset:
2087      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2088      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 20 \
2089              -dm_plex_box_lower 0 -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2090            -vdm_plex_simplex 0 \
2091            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2092            -sigma 1.0e-8 -timeScale 2.0e-14 \
2093            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2094            -output_step 50 -check_vel_res
2095      test:
2096        suffix: none_1d
2097        requires:
2098        args: -em_type none -error
2099      test:
2100        suffix: coulomb_1d
2101        args: -em_type coulomb
2102 
2103    # for viewers
2104    #-ts_monitor_sp_swarm_phase -ts_monitor_sp_swarm -em_snes_monitor -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0
2105    testset:
2106      nsize: {{1 2}}
2107      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2108      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 36 \
2109              -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2110            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2111              -vdm_plex_box_lower -3 -vdm_plex_box_upper 3 \
2112            -dm_swarm_num_species 1 -twostream -charges -1.,1. -sigma 1.0e-8 \
2113              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2114            -ts_type basicsymplectic -ts_basicsymplectic_type 2 \
2115              -ts_dt 0.01 -ts_max_time 5 -ts_max_steps 10 \
2116            -em_snes_atol 1.e-15 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2117            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1
2118      test:
2119        suffix: two_stream_c0
2120        args: -em_type primal -petscfe_default_quadrature_order 2 -petscspace_degree 2 -em_pc_type svd
2121      test:
2122        suffix: two_stream_rt
2123        requires: superlu_dist
2124        args: -em_type mixed \
2125                -potential_petscspace_degree 0 \
2126                -potential_petscdualspace_lagrange_use_moments \
2127                -potential_petscdualspace_lagrange_moment_order 2 \
2128                -field_petscspace_degree 1 -field_petscfe_default_quadrature_order 1 \
2129              -em_snes_error_if_not_converged \
2130              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2131              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2132                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2133                -em_fieldsplit_field_pc_type lu \
2134                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2135                -em_fieldsplit_potential_pc_type svd
2136 
2137    # For an eyeball check, we use
2138    # -ts_max_steps 1000 -dm_plex_box_faces 10,1 -vdm_plex_box_faces 2000 -efield_monitor
2139    # For verification, we use
2140    # -ts_max_steps 1000 -dm_plex_box_faces 100,1 -vdm_plex_box_faces 8000 -dm_swarm_num_particles 800000 -monitor_efield
2141    # -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
2142    testset:
2143      nsize: {{1 2}}
2144      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2145      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 10 \
2146              -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2147            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2148              -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 \
2149              -vpetscspace_degree 2 -vdm_plex_hash_location \
2150            -dm_swarm_num_species 1 -charges -1.,1. \
2151              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2152            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2153              -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 1 \
2154            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2155            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1
2156 
2157      test:
2158        suffix: uniform_equilibrium_1d
2159        args: -cosine_coefficients 0.0,0.5 -em_type primal -petscspace_degree 1 -em_pc_type svd
2160      test:
2161        suffix: uniform_equilibrium_1d_real
2162        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2163                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2164              -cosine_coefficients 0.0 -em_type primal -petscspace_degree 1 -em_pc_type svd
2165      test:
2166        suffix: uniform_primal_1d
2167        args: -em_type primal -petscspace_degree 1 -em_pc_type svd
2168      test:
2169        suffix: uniform_primal_1d_real
2170        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2171                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2172              -cosine_coefficients 0.01 -em_type primal -petscspace_degree 1 -em_pc_type svd
2173      # NOT WORKING -ftop_pc_type bjacobi -ftop_sub_pc_type lu -ftop_sub_pc_factor_shift_type nonzero
2174      test:
2175        suffix: uniform_primal_1d_real_pfak
2176        nsize: 1
2177        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2178                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2179              -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 -remap_dm_plex_box_faces 10,10 -remap_dm_plex_box_bd periodic,none \
2180                -remap_dm_plex_box_lower 0.,-10. -remap_dm_plex_box_upper 12.5664,10. \
2181                -remap_petscspace_degree 2 -remap_dm_plex_hash_location \
2182              -remap_freq 1 -dm_swarm_remap_type pfak \
2183                -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 \
2184                -ptof_pc_type lu \
2185              -cosine_coefficients 0.01 -em_type primal -petscspace_degree 1 -em_pc_type svd -em_proj_pc_type lu
2186      test:
2187        requires: superlu_dist
2188        suffix: uniform_mixed_1d
2189        args: -em_type mixed \
2190                -potential_petscspace_degree 0 \
2191                -potential_petscdualspace_lagrange_use_moments \
2192                -potential_petscdualspace_lagrange_moment_order 2 \
2193                -field_petscspace_degree 1 \
2194              -em_snes_error_if_not_converged \
2195              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2196              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2197                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2198                -em_fieldsplit_field_pc_type lu \
2199                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2200                -em_fieldsplit_potential_pc_type svd
2201 
2202 TEST*/
2203