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