Actual source code: ex2.c

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

  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

  7:   To visualize the efield use

  9:     -monitor_efield

 11:   To monitor moments of the distribution use

 13:     -ptof_pc_type lu -monitor_moments

 15:   To visualize the swarm distribution use

 17:     -ts_monitor_hg_swarm

 19:   To visualize the particles, we can use

 21:     -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500

 23: For a Landau Damping verification run, we use

 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

 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>
 43: #include <petsc/private/petscfeimpl.h>
 44: #include <petsc/private/dmswarmimpl.h>
 45: #include "petscdm.h"
 46: #include "petscdmlabel.h"

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

 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;

 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;

 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    drawic_x;
110:   PetscDraw    drawic_v;
111:   PetscDraw    drawic_w;
112:   PetscDrawHG  drawhgic_x;
113:   PetscDrawHG  drawhgic_v;
114:   PetscDrawHG  drawhgic_w;
115:   PetscReal    drawlgEmin;        // The minimum lg(E) to plot
116:   PetscDrawLG  drawlgE;           // Logarithm of maximum electric field
117:   PetscDrawSP  drawspE;           // Electric field at particle positions
118:   PetscDrawSP  drawspX;           // Particle positions
119:   PetscViewer  viewerRho;         // Charge density viewer
120:   PetscViewer  viewerPhi;         // Potential viewer
121:   PetscBool    monitor_positions; /* Flag to show particle positins at each time step */
122:   DM           swarm;
123:   PetscRandom  random;
124:   PetscBool    twostream;
125:   PetscBool    checkweights;
126:   PetscInt     checkVRes; /* Flag to check/output velocity residuals for nightly tests */

128:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
129: } AppCtx;

131: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
132: {
133:   PetscFunctionBeginUser;
134:   PetscInt d                      = 2;
135:   PetscInt maxSpecies             = 2;
136:   options->error                  = PETSC_FALSE;
137:   options->efield_monitor         = PETSC_FALSE;
138:   options->moment_monitor         = PETSC_FALSE;
139:   options->initial_monitor        = PETSC_FALSE;
140:   options->fake_1D                = PETSC_FALSE;
141:   options->perturbed_weights      = PETSC_FALSE;
142:   options->poisson_monitor        = PETSC_FALSE;
143:   options->ostep                  = 100;
144:   options->timeScale              = 2.0e-14;
145:   options->charges[0]             = -1.0;
146:   options->charges[1]             = 1.0;
147:   options->masses[0]              = 1.0;
148:   options->masses[1]              = 1000.0;
149:   options->thermal_energy[0]      = 1.0;
150:   options->thermal_energy[1]      = 1.0;
151:   options->cosine_coefficients[0] = 0.01;
152:   options->cosine_coefficients[1] = 0.5;
153:   options->initVel                = 1;
154:   options->totalWeight            = 1.0;
155:   options->drawic_x               = NULL;
156:   options->drawic_v               = NULL;
157:   options->drawic_w               = NULL;
158:   options->drawhgic_x             = NULL;
159:   options->drawhgic_v             = NULL;
160:   options->drawhgic_w             = NULL;
161:   options->drawlgEmin             = -6;
162:   options->drawlgE                = NULL;
163:   options->drawspE                = NULL;
164:   options->drawspX                = NULL;
165:   options->viewerRho              = NULL;
166:   options->viewerPhi              = NULL;
167:   options->em                     = EM_COULOMB;
168:   options->numParticles           = 32768;
169:   options->monitor_positions      = PETSC_FALSE;
170:   options->twostream              = PETSC_FALSE;
171:   options->checkweights           = PETSC_FALSE;
172:   options->checkVRes              = 0;

174:   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
175:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex2.c", options->error, &options->error, NULL));
176:   PetscCall(PetscOptionsBool("-monitor_efield", "Flag to plot log(max E) over time", "ex2.c", options->efield_monitor, &options->efield_monitor, NULL));
177:   PetscCall(PetscOptionsReal("-monitor_efield_min", "Minimum E field to plot", "ex2.c", options->drawlgEmin, &options->drawlgEmin, NULL));
178:   PetscCall(PetscOptionsBool("-monitor_moments", "Flag to show moments table", "ex2.c", options->moment_monitor, &options->moment_monitor, NULL));
179:   PetscCall(PetscOptionsBool("-monitor_ics", "Flag to show initial condition histograms", "ex2.c", options->initial_monitor, &options->initial_monitor, NULL));
180:   PetscCall(PetscOptionsBool("-monitor_positions", "The flag to show particle positions", "ex2.c", options->monitor_positions, &options->monitor_positions, NULL));
181:   PetscCall(PetscOptionsBool("-monitor_poisson", "The flag to show charges, Efield and potential solve", "ex2.c", options->poisson_monitor, &options->poisson_monitor, NULL));
182:   PetscCall(PetscOptionsBool("-fake_1D", "Flag to run a 1D simulation (but really in 2D)", "ex2.c", options->fake_1D, &options->fake_1D, NULL));
183:   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", "ex2.c", options->twostream, &options->twostream, NULL));
184:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", "ex2.c", options->perturbed_weights, &options->perturbed_weights, NULL));
185:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex2.c", options->checkweights, &options->checkweights, NULL));
186:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex2.c", options->checkVRes, &options->checkVRes, NULL));
187:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex2.c", options->ostep, &options->ostep, NULL));
188:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex2.c", options->timeScale, &options->timeScale, NULL));
189:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex2.c", options->initVel, &options->initVel, NULL));
190:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex2.c", options->totalWeight, &options->totalWeight, NULL));
191:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex2.c", options->cosine_coefficients, &d, NULL));
192:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex2.c", options->charges, &maxSpecies, NULL));
193:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex2.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
194:   PetscOptionsEnd();

196:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
197:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
198:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
199:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
200:   PetscFunctionReturn(PETSC_SUCCESS);
201: }

203: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
204: {
205:   PetscFunctionBeginUser;
206:   if (user->efield_monitor) {
207:     PetscDraw     draw;
208:     PetscDrawAxis axis;

210:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
211:     PetscCall(PetscDrawSetSave(draw, "ex9_Efield"));
212:     PetscCall(PetscDrawSetFromOptions(draw));
213:     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
214:     PetscCall(PetscDrawDestroy(&draw));
215:     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
216:     PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
217:     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
218:   }
219:   if (user->initial_monitor) {
220:     PetscDrawAxis axis1, axis2, axis3;
221:     PetscReal     dmboxlower[2], dmboxupper[2];
222:     PetscInt      dim, cStart, cEnd;
223:     PetscCall(DMGetDimension(sw, &dim));
224:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
225:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

227:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &user->drawic_x));
228:     PetscCall(PetscDrawSetSave(user->drawic_x, "ex9_ic_x"));
229:     PetscCall(PetscDrawSetFromOptions(user->drawic_x));
230:     PetscCall(PetscDrawHGCreate(user->drawic_x, (int)dim, &user->drawhgic_x));
231:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
232:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart)));
233:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "counts"));
234:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 1500));

236:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &user->drawic_v));
237:     PetscCall(PetscDrawSetSave(user->drawic_v, "ex9_ic_v"));
238:     PetscCall(PetscDrawSetFromOptions(user->drawic_v));
239:     PetscCall(PetscDrawHGCreate(user->drawic_v, (int)dim, &user->drawhgic_v));
240:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
241:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 1000));
242:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "counts"));
243:     PetscCall(PetscDrawAxisSetLimits(axis2, -1, 1, 0, 1500));

245:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_w", 800, 300, 400, 300, &user->drawic_w));
246:     PetscCall(PetscDrawSetSave(user->drawic_w, "ex9_ic_w"));
247:     PetscCall(PetscDrawSetFromOptions(user->drawic_w));
248:     PetscCall(PetscDrawHGCreate(user->drawic_w, (int)dim, &user->drawhgic_w));
249:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_w, &axis3));
250:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_w, 10));
251:     PetscCall(PetscDrawAxisSetLabels(axis3, "Initial W Distribution", "weight", "counts"));
252:     PetscCall(PetscDrawAxisSetLimits(axis3, 0, 0.01, 0, 5000));
253:   }
254:   if (user->monitor_positions) {
255:     PetscDraw     draw;
256:     PetscDrawAxis axis;

258:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Particle Position", 0, 0, 400, 300, &draw));
259:     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
260:     PetscCall(PetscDrawSetFromOptions(draw));
261:     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspX));
262:     PetscCall(PetscDrawDestroy(&draw));
263:     PetscCall(PetscDrawSPSetDimension(user->drawspX, 1));
264:     PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis));
265:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
266:     PetscCall(PetscDrawSPReset(user->drawspX));
267:   }
268:   if (user->poisson_monitor) {
269:     Vec           rho, phi;
270:     PetscDraw     draw;
271:     PetscDrawAxis axis;

273:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
274:     PetscCall(PetscDrawSetFromOptions(draw));
275:     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
276:     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspE));
277:     PetscCall(PetscDrawDestroy(&draw));
278:     PetscCall(PetscDrawSPSetDimension(user->drawspE, 1));
279:     PetscCall(PetscDrawSPGetAxis(user->drawspE, &axis));
280:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
281:     PetscCall(PetscDrawSPReset(user->drawspE));

283:     PetscCall(PetscViewerDrawOpen(PETSC_COMM_WORLD, NULL, "Charge Density", 0, 0, 400, 300, &user->viewerRho));
284:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRho, "rho_"));
285:     PetscCall(PetscViewerDrawGetDraw(user->viewerRho, 0, &draw));
286:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
287:     PetscCall(PetscViewerSetFromOptions(user->viewerRho));
288:     PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
289:     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
290:     PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));

292:     PetscCall(PetscViewerDrawOpen(PETSC_COMM_WORLD, NULL, "Potential", 400, 0, 400, 300, &user->viewerPhi));
293:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPhi, "phi_"));
294:     PetscCall(PetscViewerDrawGetDraw(user->viewerPhi, 0, &draw));
295:     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
296:     PetscCall(PetscViewerSetFromOptions(user->viewerPhi));
297:     PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
298:     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
299:     PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
300:   }
301:   PetscFunctionReturn(PETSC_SUCCESS);
302: }

304: static PetscErrorCode DestroyContext(AppCtx *user)
305: {
306:   PetscFunctionBeginUser;
307:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
308:   PetscCall(PetscDrawDestroy(&user->drawic_x));
309:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
310:   PetscCall(PetscDrawDestroy(&user->drawic_v));
311:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_w));
312:   PetscCall(PetscDrawDestroy(&user->drawic_w));

314:   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
315:   PetscCall(PetscDrawSPDestroy(&user->drawspE));
316:   PetscCall(PetscDrawSPDestroy(&user->drawspX));
317:   PetscCall(PetscViewerDestroy(&user->viewerRho));
318:   PetscCall(PetscViewerDestroy(&user->viewerPhi));

320:   PetscCall(PetscBagDestroy(&user->bag));
321:   PetscFunctionReturn(PETSC_SUCCESS);
322: }

324: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
325: {
326:   const PetscScalar *w;
327:   PetscInt           Np;

329:   PetscFunctionBeginUser;
330:   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
331:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
332:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
333:   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]);
334:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
335:   PetscFunctionReturn(PETSC_SUCCESS);
336: }

338: static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[3], AppCtx *user)
339: {
340:   DM          vdm;
341:   Vec         u[1];
342:   const char *fields[1] = {"w_q"};

344:   PetscFunctionBegin;
345:   PetscCall(PetscObjectQuery((PetscObject)sw, "__vdm__", (PetscObject *)&vdm));
346:   PetscCall(DMGetGlobalVector(vdm, &u[0]));
347:   PetscCall(DMSwarmPushCellDM(sw, vdm, 1, fields, "velocity"));
348:   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
349:   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
350:   PetscCall(DMSwarmPopCellDM(sw));
351:   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
352:   PetscFunctionReturn(PETSC_SUCCESS);
353: }

355: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
356: {
357:   AppCtx    *user = (AppCtx *)ctx;
358:   DM         sw;
359:   PetscReal *E, *x, *weight;
360:   PetscReal  Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
361:   PetscReal  pmoments[4]; /* \int f, \int v f, \int v^2 f */
362:   PetscInt  *species, dim, Np;

364:   PetscFunctionBeginUser;
365:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
366:   PetscCall(TSGetDM(ts, &sw));
367:   PetscCall(DMGetDimension(sw, &dim));
368:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
369:   PetscCall(DMSwarmSortGetAccess(sw));
370:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
371:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
372:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
373:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

375:   for (PetscInt p = 0; p < Np; ++p) {
376:     for (PetscInt d = 0; d < 1; ++d) {
377:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
378:       if (temp > Emax) Emax = temp;
379:     }
380:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
381:     sum += E[p * dim];
382:     chargesum += user->charges[0] * weight[p];
383:   }
384:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
385:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;

387:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
388:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
389:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
390:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
391:   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
392:   PetscCall(PetscDrawLGDraw(user->drawlgE));
393:   PetscDraw draw;
394:   PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
395:   PetscCall(PetscDrawSave(draw));

397:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
398:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%+e\t%e\t%f\t%f\t%f\t%f\t%f\t%f\t%f\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[1 + dim]));
399:   PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
400:   PetscFunctionReturn(PETSC_SUCCESS);
401: }

403: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
404: {
405:   AppCtx   *user = (AppCtx *)ctx;
406:   DM        sw;
407:   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */

409:   PetscFunctionBeginUser;
410:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
411:   PetscCall(TSGetDM(ts, &sw));

413:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
414:   PetscCall(computeVelocityFEMMoments(sw, fmoments, user));

416:   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]));
417:   PetscFunctionReturn(PETSC_SUCCESS);
418: }

420: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
421: {
422:   AppCtx            *user = (AppCtx *)ctx;
423:   DM                 dm, sw;
424:   const PetscScalar *u;
425:   PetscReal         *weight, *pos, *vel;
426:   PetscInt           dim, p, Np, cStart, cEnd;

428:   PetscFunctionBegin;
429:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
430:   PetscCall(TSGetDM(ts, &sw));
431:   PetscCall(DMSwarmGetCellDM(sw, &dm));
432:   PetscCall(DMGetDimension(sw, &dim));
433:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
434:   PetscCall(DMSwarmSortGetAccess(sw));
435:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

437:   if (step == 0) {
438:     PetscCall(PetscDrawHGReset(user->drawhgic_x));
439:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &user->drawic_x));
440:     PetscCall(PetscDrawClear(user->drawic_x));
441:     PetscCall(PetscDrawFlush(user->drawic_x));

443:     PetscCall(PetscDrawHGReset(user->drawhgic_v));
444:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &user->drawic_v));
445:     PetscCall(PetscDrawClear(user->drawic_v));
446:     PetscCall(PetscDrawFlush(user->drawic_v));

448:     PetscCall(PetscDrawHGReset(user->drawhgic_w));
449:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_w, &user->drawic_w));
450:     PetscCall(PetscDrawClear(user->drawic_w));
451:     PetscCall(PetscDrawFlush(user->drawic_w));

453:     PetscCall(VecGetArrayRead(U, &u));
454:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
455:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
456:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));

458:     PetscCall(VecGetLocalSize(U, &Np));
459:     Np /= dim * 2;
460:     for (p = 0; p < Np; ++p) {
461:       PetscCall(PetscDrawHGAddValue(user->drawhgic_x, pos[p * dim]));
462:       PetscCall(PetscDrawHGAddValue(user->drawhgic_v, vel[p * dim]));
463:       PetscCall(PetscDrawHGAddValue(user->drawhgic_w, weight[p]));
464:     }

466:     PetscCall(VecRestoreArrayRead(U, &u));
467:     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
468:     PetscCall(PetscDrawHGSave(user->drawhgic_x));

470:     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
471:     PetscCall(PetscDrawHGSave(user->drawhgic_v));

473:     PetscCall(PetscDrawHGDraw(user->drawhgic_w));
474:     PetscCall(PetscDrawHGSave(user->drawhgic_w));

476:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
477:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
478:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
479:   }
480:   PetscFunctionReturn(PETSC_SUCCESS);
481: }

483: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
484: {
485:   AppCtx         *user = (AppCtx *)ctx;
486:   DM              dm, sw;
487:   PetscScalar    *x, *v, *weight;
488:   PetscReal       lower[3], upper[3], speed;
489:   const PetscInt *s;
490:   PetscInt        dim, cStart, cEnd, c;

492:   PetscFunctionBeginUser;
493:   if (step > 0 && step % user->ostep == 0) {
494:     PetscCall(TSGetDM(ts, &sw));
495:     PetscCall(DMSwarmGetCellDM(sw, &dm));
496:     PetscCall(DMGetDimension(dm, &dim));
497:     PetscCall(DMGetBoundingBox(dm, lower, upper));
498:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
499:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
500:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
501:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
502:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
503:     PetscCall(DMSwarmSortGetAccess(sw));
504:     PetscCall(PetscDrawSPReset(user->drawspX));
505:     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], lower[1], upper[1]));
506:     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], -12, 12));
507:     for (c = 0; c < cEnd - cStart; ++c) {
508:       PetscInt *pidx, Npc, q;
509:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
510:       for (q = 0; q < Npc; ++q) {
511:         const PetscInt p = pidx[q];
512:         if (s[p] == 0) {
513:           speed = PetscSqrtReal(PetscSqr(v[p * dim]) + PetscSqr(v[p * dim + 1]));
514:           if (dim == 1 || user->fake_1D) {
515:             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &v[p * dim], &speed));
516:           } else {
517:             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
518:           }
519:         } else if (s[p] == 1) {
520:           PetscCall(PetscDrawSPAddPoint(user->drawspX, &x[p * dim], &v[p * dim]));
521:         }
522:       }
523:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
524:     }
525:     PetscCall(PetscDrawSPDraw(user->drawspX, PETSC_TRUE));
526:     PetscDraw draw;
527:     PetscCall(PetscDrawSPGetDraw(user->drawspX, &draw));
528:     PetscCall(PetscDrawSave(draw));
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: }

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;
543:   PetscReal    lower[3], upper[3], xval;
544:   PetscInt     dim, cStart, cEnd, c;

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));

554:     PetscCall(PetscDrawSPReset(user->drawspE));
555:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
556:     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
557:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

559:     PetscCall(DMSwarmSortGetAccess(sw));
560:     for (c = 0; c < cEnd - cStart; ++c) {
561:       PetscReal Eavg = 0.0;
562:       PetscInt *pidx, Npc, q;
563:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
564:       for (q = 0; q < Npc; ++q) {
565:         const PetscInt p = pidx[q];
566:         Eavg += E[p * dim];
567:       }
568:       Eavg /= Npc;
569:       xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
570:       PetscCall(PetscDrawSPAddPoint(user->drawspE, &xval, &Eavg));
571:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
572:     }
573:     PetscCall(PetscDrawSPDraw(user->drawspE, PETSC_TRUE));
574:     PetscDraw draw;
575:     PetscCall(PetscDrawSPGetDraw(user->drawspE, &draw));
576:     PetscCall(PetscDrawSave(draw));
577:     PetscCall(DMSwarmSortRestoreAccess(sw));
578:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
579:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
580:     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));

582:     Vec rho, phi;

584:     PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
585:     PetscCall(VecView(rho, user->viewerRho));
586:     PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));

588:     PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
589:     PetscCall(VecView(phi, user->viewerPhi));
590:     PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
591:   }
592:   PetscFunctionReturn(PETSC_SUCCESS);
593: }

595: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
596: {
597:   PetscBag   bag;
598:   Parameter *p;

600:   PetscFunctionBeginUser;
601:   /* setup PETSc parameter bag */
602:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
603:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
604:   bag = ctx->bag;
605:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
606:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
607:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
608:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
609:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
610:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
611:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
612:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

614:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
615:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
616:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
617:   PetscCall(PetscBagSetFromOptions(bag));
618:   {
619:     PetscViewer       viewer;
620:     PetscViewerFormat format;
621:     PetscBool         flg;

623:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
624:     if (flg) {
625:       PetscCall(PetscViewerPushFormat(viewer, format));
626:       PetscCall(PetscBagView(bag, viewer));
627:       PetscCall(PetscViewerFlush(viewer));
628:       PetscCall(PetscViewerPopFormat(viewer));
629:       PetscCall(PetscViewerDestroy(&viewer));
630:     }
631:   }
632:   PetscFunctionReturn(PETSC_SUCCESS);
633: }

635: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
636: {
637:   PetscFunctionBeginUser;
638:   PetscCall(DMCreate(comm, dm));
639:   PetscCall(DMSetType(*dm, DMPLEX));
640:   PetscCall(DMSetFromOptions(*dm));
641:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

643:   // Cache the mesh geometry
644:   DMField         coordField;
645:   IS              cellIS;
646:   PetscQuadrature quad;
647:   PetscReal      *wt, *pt;
648:   PetscInt        cdim, cStart, cEnd;

650:   PetscCall(DMGetCoordinateField(*dm, &coordField));
651:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
652:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
653:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
654:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
655:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
656:   PetscCall(PetscMalloc1(1, &wt));
657:   PetscCall(PetscMalloc1(2, &pt));
658:   wt[0] = 1.;
659:   pt[0] = -1.;
660:   pt[1] = -1.;
661:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
662:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FALSE, &user->fegeom));
663:   PetscCall(PetscQuadratureDestroy(&quad));
664:   PetscCall(ISDestroy(&cellIS));
665:   PetscFunctionReturn(PETSC_SUCCESS);
666: }

668: 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[])
669: {
670:   f0[0] = -constants[SIGMA];
671: }

673: 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[])
674: {
675:   PetscInt d;
676:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
677: }

679: 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[])
680: {
681:   PetscInt d;
682:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
683: }

685: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
686: {
687:   *u = 0.0;
688:   return PETSC_SUCCESS;
689: }

691: /*
692:    /  I   -grad\ / q \ = /0\
693:    \-div    0  / \phi/   \f/
694: */
695: 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[])
696: {
697:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
698: }

700: 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[])
701: {
702:   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
703: }

705: 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[])
706: {
707:   f0[0] += constants[SIGMA];
708:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
709: }

711: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
712: 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[])
713: {
714:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
715: }

717: 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[])
718: {
719:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
720: }

722: 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[])
723: {
724:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
725: }

727: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
728: {
729:   PetscFE   fephi, feq;
730:   PetscDS   ds;
731:   PetscBool simplex;
732:   PetscInt  dim;

734:   PetscFunctionBeginUser;
735:   PetscCall(DMGetDimension(dm, &dim));
736:   PetscCall(DMPlexIsSimplex(dm, &simplex));
737:   if (user->em == EM_MIXED) {
738:     DMLabel        label;
739:     const PetscInt id = 1;

741:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
742:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
743:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
744:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
745:     PetscCall(PetscFECopyQuadrature(feq, fephi));
746:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
747:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
748:     PetscCall(DMCreateDS(dm));
749:     PetscCall(PetscFEDestroy(&fephi));
750:     PetscCall(PetscFEDestroy(&feq));

752:     PetscCall(DMGetLabel(dm, "marker", &label));
753:     PetscCall(DMGetDS(dm, &ds));

755:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
756:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
757:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
758:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
759:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

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

763:   } else {
764:     MatNullSpace nullsp;
765:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
766:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
767:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
768:     PetscCall(DMCreateDS(dm));
769:     PetscCall(DMGetDS(dm, &ds));
770:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
771:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
772:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
773:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
774:     PetscCall(MatNullSpaceDestroy(&nullsp));
775:     PetscCall(PetscFEDestroy(&fephi));
776:   }
777:   PetscFunctionReturn(PETSC_SUCCESS);
778: }

780: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
781: {
782:   SNES         snes;
783:   Mat          J;
784:   MatNullSpace nullSpace;

786:   PetscFunctionBeginUser;
787:   PetscCall(CreateFEM(dm, user));
788:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
789:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
790:   PetscCall(SNESSetDM(snes, dm));
791:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
792:   PetscCall(SNESSetFromOptions(snes));

794:   PetscCall(DMCreateMatrix(dm, &J));
795:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
796:   PetscCall(MatSetNullSpace(J, nullSpace));
797:   PetscCall(MatNullSpaceDestroy(&nullSpace));
798:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
799:   PetscCall(MatDestroy(&J));
800:   PetscCall(DMCreateMassMatrix(dm, dm, &user->M));
801:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
802:   user->snes = snes;
803:   PetscFunctionReturn(PETSC_SUCCESS);
804: }

806: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
807: {
808:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
809:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
810:   return PETSC_SUCCESS;
811: }
812: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
813: {
814:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
815:   return PETSC_SUCCESS;
816: }

818: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
819: {
820:   const PetscReal alpha = scale ? scale[0] : 0.0;
821:   const PetscReal k     = scale ? scale[1] : 1.;
822:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
823:   return PETSC_SUCCESS;
824: }

826: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
827: {
828:   const PetscReal alpha = scale ? scale[0] : 0.;
829:   const PetscReal k     = scale ? scale[0] : 1.;
830:   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
831:   return PETSC_SUCCESS;
832: }

834: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
835: {
836:   PetscFE        fe;
837:   DMPolytopeType ct;
838:   PetscInt       dim, cStart;
839:   const char    *prefix = "v";

841:   PetscFunctionBegin;
842:   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
843:   PetscCall(DMSetType(*vdm, DMPLEX));
844:   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
845:   PetscCall(DMSetFromOptions(*vdm));
846:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

848:   PetscCall(DMGetDimension(*vdm, &dim));
849:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
850:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
851:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
852:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
853:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
854:   PetscCall(DMCreateDS(*vdm));
855:   PetscCall(PetscFEDestroy(&fe));
856:   PetscFunctionReturn(PETSC_SUCCESS);
857: }

859: static PetscErrorCode InitializeParticles_Centroid(DM sw, PetscBool force1D)
860: {
861:   DM_Swarm  *swarm = (DM_Swarm *)sw->data;
862:   DM         xdm, vdm;
863:   PetscReal  vmin[3], vmax[3];
864:   PetscReal *x, *v;
865:   PetscInt  *species, *cellid;
866:   PetscInt   dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
867:   PetscBool  flg;
868:   MPI_Comm   comm;

870:   PetscFunctionBegin;
871:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

873:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
874:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
875:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
876:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
877:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
878:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
879:   PetscOptionsEnd();
880:   debug = swarm->printCoords;

882:   PetscCall(DMGetDimension(sw, &dim));
883:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
884:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

886:   PetscCall(PetscObjectQuery((PetscObject)sw, "__vdm__", (PetscObject *)&vdm));
887:   if (!vdm) {
888:     PetscCall(CreateVelocityDM(sw, &vdm));
889:     PetscCall(PetscObjectCompose((PetscObject)sw, "__vdm__", (PetscObject)vdm));
890:     PetscCall(DMDestroy(&vdm));
891:     PetscCall(PetscObjectQuery((PetscObject)sw, "__vdm__", (PetscObject *)&vdm));
892:   }
893:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

895:   // One particle per centroid on the tensor product grid
896:   Npc = (vcEnd - vcStart) * Ns;
897:   Np  = (xcEnd - xcStart) * Npc;
898:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
899:   if (debug) {
900:     PetscInt gNp;
901:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
902:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
903:   }

905:   // Set species and cellid
906:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
907:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
908:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
909:     for (PetscInt s = 0; s < Ns; ++s) {
910:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
911:         species[p] = s;
912:         cellid[p]  = c;
913:       }
914:     }
915:   }
916:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
917:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));

919:   // Set particle coordinates
920:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
921:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
922:   PetscCall(DMSwarmSortGetAccess(sw));
923:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
924:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
925:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
926:     const PetscInt cell = c + xcStart;
927:     PetscInt      *pidx, Npc;
928:     PetscReal      centroid[3], volume;

930:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
931:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
932:     for (PetscInt s = 0; s < Ns; ++s) {
933:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
934:         const PetscInt p = pidx[q * Ns + s];

936:         for (PetscInt d = 0; d < dim; ++d) {
937:           x[p * dim + d] = centroid[d];
938:           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
939:           if (force1D && d > 0) v[p * dim + d] = 0.;
940:         }
941:       }
942:     }
943:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
944:   }
945:   PetscCall(DMSwarmSortRestoreAccess(sw));
946:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
947:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
948:   PetscFunctionReturn(PETSC_SUCCESS);
949: }

951: /*
952:   InitializeWeights - Compute weight for each local particle

954:   Input Parameters:
955: + sw          - The `DMSwarm`
956: . totalWeight - The sum of all particle weights
957: . force1D     - Flag to treat the problem as 1D
958: . func        - The PDF for the particle spatial distribution
959: - param       - The PDF parameters

961:   Notes:
962:   The PDF for velocity is assumed to be a Gaussian

964:   The particle weights are returned in the `w_q` field of `sw`.
965: */
966: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscBool force1D, PetscProbFunc func, const PetscReal param[])
967: {
968:   DM               xdm, vdm;
969:   PetscScalar     *weight;
970:   PetscQuadrature  xquad;
971:   const PetscReal *xq, *xwq;
972:   const PetscInt   order = 5;
973:   PetscReal       *xqd   = NULL, xi0[3];
974:   PetscReal        xwtot = 0., pwtot = 0.;
975:   PetscInt         xNq;
976:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
977:   MPI_Comm         comm;
978:   PetscMPIInt      rank;

980:   PetscFunctionBegin;
981:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
982:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
983:   PetscCall(DMGetDimension(sw, &dim));
984:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
985:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
986:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
987:   PetscCall(PetscObjectQuery((PetscObject)sw, "__vdm__", (PetscObject *)&vdm));
988:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

990:   // Setup Quadrature for spatial and velocity weight calculations
991:   if (force1D) PetscCall(PetscDTGaussTensorQuadrature(1, 1, order, -1.0, 1.0, &xquad));
992:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad));
993:   PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq));
994:   if (force1D) {
995:     PetscCall(PetscCalloc1(xNq * dim, &xqd));
996:     for (PetscInt q = 0; q < xNq; ++q) xqd[q * dim] = xq[q];
997:   }
998:   for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0;

1000:   // Integrate the density function to get the weights of particles in each cell
1001:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
1002:   PetscCall(DMSwarmSortGetAccess(sw));
1003:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1004:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
1005:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
1006:     PetscInt          *pidx, Npc;
1007:     PetscInt           xNc;
1008:     const PetscScalar *xarray;
1009:     PetscScalar       *xcoords = NULL;
1010:     PetscBool          xisDG;

1012:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1013:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1014:     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);
1015:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
1016:     for (PetscInt q = 0; q < xNq; ++q) {
1017:       // Transform quadrature points from ref space to real space
1018:       if (force1D) CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xqd[q * dim], xqr);
1019:       else CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
1020:       // Get probability density at quad point
1021:       //   No need to scale xqr since PDF will be periodic
1022:       PetscCall((*func)(xqr, param, &xden));
1023:       if (force1D) xdetJ = xJ[0]; // Only want x contribution
1024:       xw += xden * (xwq[q] * xdetJ);
1025:     }
1026:     xwtot += xw;
1027:     if (debug) {
1028:       IS              globalOrdering;
1029:       const PetscInt *ordering;

1031:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
1032:       PetscCall(ISGetIndices(globalOrdering, &ordering));
1033:       PetscCall(PetscSynchronizedPrintf(comm, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(xcoords[0]), (double)PetscRealPart(xcoords[0 + dim]), (double)xw));
1034:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
1035:     }
1036:     // Set weights to be Gaussian in velocity cells
1037:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
1038:       const PetscInt     p  = pidx[vc * Ns + 0];
1039:       PetscReal          vw = 0.;
1040:       PetscInt           vNc;
1041:       const PetscScalar *varray;
1042:       PetscScalar       *vcoords = NULL;
1043:       PetscBool          visDG;

1045:       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1046:       // TODO: Fix 2 stream Ask Joe
1047:       //   Two stream function from 1/2pi v^2 e^(-v^2/2)
1048:       //   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.)));
1049:       vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.)));

1051:       weight[p] = totalWeight * vw * xw;
1052:       pwtot += weight[p];
1053:       PetscCheck(weight[p] <= 1., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeeded 1: %g, %g, %g", p, xw, vw, totalWeight);
1054:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1055:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
1056:     }
1057:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1058:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1059:   }
1060:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1061:   PetscCall(DMSwarmSortRestoreAccess(sw));
1062:   PetscCall(PetscQuadratureDestroy(&xquad));
1063:   if (force1D) PetscCall(PetscFree(xqd));

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

1068:     PetscCall(PetscSynchronizedFlush(comm, NULL));
1069:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1070:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
1071:   }
1072:   PetscFunctionReturn(PETSC_SUCCESS);
1073: }

1075: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
1076: {
1077:   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
1078:   PetscInt  dim;

1080:   PetscFunctionBegin;
1081:   PetscCall(DMGetDimension(sw, &dim));
1082:   PetscCall(InitializeParticles_Centroid(sw, user->fake_1D));
1083:   PetscCall(InitializeWeights(sw, user->totalWeight, user->fake_1D, dim == 1 || user->fake_1D ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
1084:   PetscFunctionReturn(PETSC_SUCCESS);
1085: }

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;

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;

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));

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: }

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;

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));

1178:   for (PetscInt p = 0; p < Np; ++p) {
1179:     PetscReal a[3] = {0., 0., 0.}, vel[3] = {0., 0., 0.};

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: }

1195: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1196: {
1197:   PetscReal v0[2] = {1., 0.};
1198:   PetscInt  dim;

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(DMSwarmFinalizeFieldRegister(*sw));
1214:   PetscCall(DMSetApplicationContext(*sw, user));
1215:   PetscCall(DMSetFromOptions(*sw));
1216:   user->swarm = *sw;
1217:   if (user->perturbed_weights) {
1218:     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
1219:   } else {
1220:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
1221:     PetscCall(DMSwarmInitializeCoordinates(*sw));
1222:     if (user->fake_1D) {
1223:       PetscCall(InitializeVelocities_Fake1D(*sw, user));
1224:     } else {
1225:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
1226:     }
1227:   }
1228:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
1229:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
1230:   {
1231:     Vec gc, gc0, gv, gv0;

1233:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1234:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1235:     PetscCall(VecCopy(gc, gc0));
1236:     PetscCall(VecViewFromOptions(gc, NULL, "-ic_x_view"));
1237:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1238:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1239:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv));
1240:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0));
1241:     PetscCall(VecCopy(gv, gv0));
1242:     PetscCall(VecViewFromOptions(gv, NULL, "-ic_v_view"));
1243:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv));
1244:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0));
1245:   }
1246:   {
1247:     const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};

1249:     PetscCall(DMSwarmVectorDefineFields(*sw, 2, fieldnames));
1250:   }
1251:   PetscFunctionReturn(PETSC_SUCCESS);
1252: }

1254: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
1255: {
1256:   AppCtx     *user;
1257:   PetscReal  *coords;
1258:   PetscInt   *species, dim, Np, Ns;
1259:   PetscMPIInt size;

1261:   PetscFunctionBegin;
1262:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
1263:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
1264:   PetscCall(DMGetDimension(sw, &dim));
1265:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1266:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1267:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1269:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1270:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1271:   for (PetscInt p = 0; p < Np; ++p) {
1272:     PetscReal *pcoord = &coords[p * dim];
1273:     PetscReal  pE[3]  = {0., 0., 0.};

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

1280:       if (p == q) continue;
1281:       q_q = user->charges[species[q]] * 1.;
1282:       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1283:       r = DMPlex_NormD_Internal(dim, rpq);
1284:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
1285:       r3 = PetscPowRealInt(r, 3);
1286:       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
1287:     }
1288:     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
1289:   }
1290:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1291:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1292:   PetscFunctionReturn(PETSC_SUCCESS);
1293: }

1295: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[])
1296: {
1297:   DM         dm;
1298:   AppCtx    *user;
1299:   PetscDS    ds;
1300:   PetscFE    fe;
1301:   Mat        M_p;
1302:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
1303:   Vec        rho;         // Charge density, M^{-1} rhoRhs
1304:   Vec        phi, locPhi; // Potential
1305:   Vec        f;           // Particle weights
1306:   PetscReal *coords;
1307:   PetscInt   dim, cStart, cEnd, Np;

1309:   PetscFunctionBegin;
1310:   PetscCall(DMGetApplicationContext(sw, (void *)&user));
1311:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1312:   PetscCall(DMGetDimension(sw, &dim));
1313:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

1315:   KSP          ksp;
1316:   const char **oldFields;
1317:   PetscInt     Nf;
1318:   const char **tmp;

1320:   PetscCall(SNESGetDM(snes, &dm));
1321:   PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp));
1322:   PetscCall(PetscMalloc1(Nf, &oldFields));
1323:   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f]));
1324:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1325:   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
1326:   PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields));
1327:   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f]));
1328:   PetscCall(PetscFree(oldFields));

1330:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
1331:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
1332:   PetscCall(DMGetNamedGlobalVector(dm, "rho", &rho));
1333:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1334:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

1336:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1337:   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1338:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

1340:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1341:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1343:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1344:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1345:   PetscCall(KSPSetOperators(ksp, user->M, user->M));
1346:   PetscCall(KSPSetFromOptions(ksp));
1347:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

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

1351:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1352:   PetscCall(DMRestoreNamedGlobalVector(dm, "rho", &rho));
1353:   PetscCall(KSPDestroy(&ksp));
1354:   PetscCall(MatDestroy(&M_p));

1356:   PetscCall(DMGetNamedGlobalVector(dm, "phi", &phi));
1357:   PetscCall(VecSet(phi, 0.0));
1358:   PetscCall(SNESSolve(snes, rhoRhs, phi));
1359:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
1360:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1362:   PetscCall(DMGetLocalVector(dm, &locPhi));
1363:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1364:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1365:   PetscCall(DMRestoreNamedGlobalVector(dm, "phi", &phi));
1366:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

1368:   PetscCall(DMGetDS(dm, &ds));
1369:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1370:   PetscCall(DMSwarmSortGetAccess(sw));
1371:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1372:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1374:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1375:   PetscFEGeom *chunkgeom = NULL;
1376:   for (PetscInt c = cStart; c < cEnd; ++c) {
1377:     PetscTabulation tab;
1378:     PetscScalar    *clPhi = NULL;
1379:     PetscReal      *pcoord, *refcoord;
1380:     PetscInt       *points;
1381:     PetscInt        Ncp;

1383:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1384:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1385:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1386:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1387:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1388:     {
1389:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1390:       for (PetscInt i = 0; i < Ncp; ++i) {
1391:         const PetscReal x0[3] = {-1., -1., -1.};
1392:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1393:       }
1394:     }
1395:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1396:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1397:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1398:       const PetscReal *basisDer = tab->T[1];
1399:       const PetscInt   p        = points[cp];

1401:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1402:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
1403:       for (PetscInt d = 0; d < dim; ++d) {
1404:         E[p * dim + d] *= -1.0;
1405:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1406:       }
1407:     }
1408:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1409:     PetscCall(PetscTabulationDestroy(&tab));
1410:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1411:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1412:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1413:   }
1414:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1415:   PetscCall(DMSwarmSortRestoreAccess(sw));
1416:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1417:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1418:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1419:   PetscFunctionReturn(PETSC_SUCCESS);
1420: }

1422: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
1423: {
1424:   AppCtx         *user;
1425:   DM              dm, potential_dm;
1426:   KSP             ksp;
1427:   IS              potential_IS;
1428:   PetscDS         ds;
1429:   PetscFE         fe;
1430:   Mat             M_p, M;
1431:   Vec             phi, locPhi, rho, f, temp_rho, rho0;
1432:   PetscQuadrature q;
1433:   PetscReal      *coords;
1434:   PetscInt        dim, cStart, cEnd, Np, pot_field = 1;
1435:   const char    **oldFields;
1436:   PetscInt        Nf;
1437:   const char    **tmp;

1439:   PetscFunctionBegin;
1440:   PetscCall(DMGetApplicationContext(sw, &user));
1441:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1442:   PetscCall(DMGetDimension(sw, &dim));
1443:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

1445:   /* Create the charges rho */
1446:   PetscCall(SNESGetDM(snes, &dm));
1447:   PetscCall(DMGetGlobalVector(dm, &rho));
1448:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));

1450:   PetscCall(DMCreateSubDM(dm, 1, &pot_field, &potential_IS, &potential_dm));

1452:   PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp));
1453:   PetscCall(PetscMalloc1(Nf, &oldFields));
1454:   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f]));
1455:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1456:   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
1457:   PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields));
1458:   for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f]));
1459:   PetscCall(PetscFree(oldFields));

1461:   PetscCall(DMCreateMassMatrix(potential_dm, potential_dm, &M));
1462:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1463:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1464:   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
1465:   PetscCall(PetscObjectSetName((PetscObject)temp_rho, "Mf"));
1466:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1467:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1468:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1469:   PetscCall(MatMultTranspose(M_p, f, temp_rho));
1470:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1471:   PetscCall(DMGetGlobalVector(potential_dm, &rho0));
1472:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Mixed Compute"));

1474:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1475:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1476:   PetscCall(KSPSetOperators(ksp, M, M));
1477:   PetscCall(KSPSetFromOptions(ksp));
1478:   PetscCall(KSPSolve(ksp, temp_rho, rho0));
1479:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1481:   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
1482:   PetscCall(VecScale(rho, 0.25));
1483:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1484:   PetscCall(VecViewFromOptions(temp_rho, NULL, "-temprho_view"));
1485:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1486:   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
1487:   PetscCall(DMRestoreGlobalVector(potential_dm, &rho0));

1489:   PetscCall(MatDestroy(&M_p));
1490:   PetscCall(MatDestroy(&M));
1491:   PetscCall(KSPDestroy(&ksp));
1492:   PetscCall(DMDestroy(&potential_dm));
1493:   PetscCall(ISDestroy(&potential_IS));

1495:   PetscCall(DMGetGlobalVector(dm, &phi));
1496:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1497:   PetscCall(VecSet(phi, 0.0));
1498:   PetscCall(SNESSolve(snes, rho, phi));
1499:   PetscCall(DMRestoreGlobalVector(dm, &rho));

1501:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1503:   PetscCall(DMGetLocalVector(dm, &locPhi));
1504:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1505:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1506:   PetscCall(DMRestoreGlobalVector(dm, &phi));
1507:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

1509:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1510:   PetscCall(DMGetDS(dm, &ds));
1511:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1512:   PetscCall(DMSwarmSortGetAccess(sw));
1513:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1514:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1515:   PetscCall(PetscFEGetQuadrature(fe, &q));
1516:   PetscFEGeom *chunkgeom = NULL;
1517:   for (PetscInt c = cStart; c < cEnd; ++c) {
1518:     PetscTabulation tab;
1519:     PetscScalar    *clPhi = NULL;
1520:     PetscReal      *pcoord, *refcoord;
1521:     PetscInt       *points;
1522:     PetscInt        Ncp;

1524:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1525:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1526:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1527:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1528:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1529:     {
1530:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1531:       for (PetscInt i = 0; i < Ncp; ++i) {
1532:         // Apply the inverse affine transformation for each point
1533:         const PetscReal x0[3] = {-1., -1., -1.};
1534:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1535:       }
1536:     }
1537:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1538:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));

1540:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1541:       const PetscInt p = points[cp];

1543:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1544:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
1545:       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
1546:       for (PetscInt d = 0; d < dim; ++d) {
1547:         E[p * dim + d] *= -2.0;
1548:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1549:       }
1550:     }
1551:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1552:     PetscCall(PetscTabulationDestroy(&tab));
1553:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1554:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1555:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1556:   }
1557:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1558:   PetscCall(DMSwarmSortRestoreAccess(sw));
1559:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1560:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1561:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1562:   PetscFunctionReturn(PETSC_SUCCESS);
1563: }

1565: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
1566: {
1567:   AppCtx  *ctx;
1568:   PetscInt dim, Np;

1570:   PetscFunctionBegin;
1573:   PetscAssertPointer(E, 3);
1574:   PetscCall(DMGetDimension(sw, &dim));
1575:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1576:   PetscCall(DMGetApplicationContext(sw, &ctx));
1577:   PetscCall(PetscArrayzero(E, Np * dim));

1579:   switch (ctx->em) {
1580:   case EM_PRIMAL:
1581:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
1582:     break;
1583:   case EM_COULOMB:
1584:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1585:     break;
1586:   case EM_MIXED:
1587:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
1588:     break;
1589:   case EM_NONE:
1590:     break;
1591:   default:
1592:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
1593:   }
1594:   PetscFunctionReturn(PETSC_SUCCESS);
1595: }

1597: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1598: {
1599:   DM                 sw;
1600:   SNES               snes = ((AppCtx *)ctx)->snes;
1601:   const PetscReal   *coords, *vel;
1602:   const PetscScalar *u;
1603:   PetscScalar       *g;
1604:   PetscReal         *E, m_p = 1., q_p = -1.;
1605:   PetscInt           dim, d, Np, p;

1607:   PetscFunctionBeginUser;
1608:   PetscCall(TSGetDM(ts, &sw));
1609:   PetscCall(DMGetDimension(sw, &dim));
1610:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1611:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1612:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1613:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1614:   PetscCall(VecGetArrayRead(U, &u));
1615:   PetscCall(VecGetArray(G, &g));

1617:   PetscCall(ComputeFieldAtParticles(snes, sw, E));

1619:   Np /= 2 * dim;
1620:   for (p = 0; p < Np; ++p) {
1621:     for (d = 0; d < dim; ++d) {
1622:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1623:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1624:     }
1625:   }
1626:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1627:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1628:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1629:   PetscCall(VecRestoreArrayRead(U, &u));
1630:   PetscCall(VecRestoreArray(G, &g));
1631:   PetscFunctionReturn(PETSC_SUCCESS);
1632: }

1634: /* J_{ij} = dF_i/dx_j
1635:    J_p = (  0   1)
1636:          (-w^2  0)
1637:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1638:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1639: */
1640: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1641: {
1642:   DM               sw;
1643:   const PetscReal *coords, *vel;
1644:   PetscInt         dim, d, Np, p, rStart;

1646:   PetscFunctionBeginUser;
1647:   PetscCall(TSGetDM(ts, &sw));
1648:   PetscCall(DMGetDimension(sw, &dim));
1649:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1650:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1651:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1652:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1653:   Np /= 2 * dim;
1654:   for (p = 0; p < Np; ++p) {
1655:     const PetscReal x0      = coords[p * dim + 0];
1656:     const PetscReal vy0     = vel[p * dim + 1];
1657:     const PetscReal omega   = vy0 / x0;
1658:     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};

1660:     for (d = 0; d < dim; ++d) {
1661:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1662:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1663:     }
1664:   }
1665:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1666:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1667:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1668:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1669:   PetscFunctionReturn(PETSC_SUCCESS);
1670: }

1672: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1673: {
1674:   AppCtx            *user = (AppCtx *)ctx;
1675:   DM                 sw;
1676:   const PetscScalar *v;
1677:   PetscScalar       *xres;
1678:   PetscInt           Np, p, d, dim;

1680:   PetscFunctionBeginUser;
1681:   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
1682:   PetscCall(TSGetDM(ts, &sw));
1683:   PetscCall(DMGetDimension(sw, &dim));
1684:   PetscCall(VecGetLocalSize(Xres, &Np));
1685:   PetscCall(VecGetArrayRead(V, &v));
1686:   PetscCall(VecGetArray(Xres, &xres));
1687:   Np /= dim;
1688:   for (p = 0; p < Np; ++p) {
1689:     for (d = 0; d < dim; ++d) {
1690:       xres[p * dim + d] = v[p * dim + d];
1691:       if (user->fake_1D && d > 0) xres[p * dim + d] = 0;
1692:     }
1693:   }
1694:   PetscCall(VecRestoreArrayRead(V, &v));
1695:   PetscCall(VecRestoreArray(Xres, &xres));
1696:   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
1697:   PetscFunctionReturn(PETSC_SUCCESS);
1698: }

1700: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1701: {
1702:   DM                 sw;
1703:   AppCtx            *user = (AppCtx *)ctx;
1704:   SNES               snes = ((AppCtx *)ctx)->snes;
1705:   const PetscScalar *x;
1706:   const PetscReal   *coords, *vel;
1707:   PetscReal         *E, m_p, q_p;
1708:   PetscScalar       *vres;
1709:   PetscInt           Np, p, dim, d;
1710:   Parameter         *param;

1712:   PetscFunctionBeginUser;
1713:   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
1714:   PetscCall(TSGetDM(ts, &sw));
1715:   PetscCall(DMGetDimension(sw, &dim));
1716:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1717:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1718:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1719:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1720:   m_p = user->masses[0] * param->m0;
1721:   q_p = user->charges[0] * param->q0;
1722:   PetscCall(VecGetLocalSize(Vres, &Np));
1723:   PetscCall(VecGetArrayRead(X, &x));
1724:   PetscCall(VecGetArray(Vres, &vres));
1725:   PetscCall(ComputeFieldAtParticles(snes, sw, E));

1727:   Np /= dim;
1728:   for (p = 0; p < Np; ++p) {
1729:     for (d = 0; d < dim; ++d) {
1730:       vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1731:       if (user->fake_1D && d > 0) vres[p * dim + d] = 0.;
1732:     }
1733:   }
1734:   PetscCall(VecRestoreArrayRead(X, &x));
1735:   /*
1736:     Synchronized, ordered output for parallel/sequential test cases.
1737:     In the 1D (on the 2D mesh) case, every y component should be zero.
1738:   */
1739:   if (user->checkVRes) {
1740:     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1741:     PetscInt  step;

1743:     PetscCall(TSGetStepNumber(ts, &step));
1744:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
1745:     for (PetscInt p = 0; p < Np; ++p) {
1746:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
1747:       PetscCheck(PetscAbsScalar(vres[p * dim + 1]) < PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Y velocity should be 0., not %g", (double)PetscRealPart(vres[p * dim + 1]));
1748:     }
1749:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1750:   }
1751:   PetscCall(VecRestoreArray(Vres, &vres));
1752:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1753:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1754:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1755:   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
1756:   PetscFunctionReturn(PETSC_SUCCESS);
1757: }

1759: static PetscErrorCode CreateSolution(TS ts)
1760: {
1761:   DM       sw;
1762:   Vec      u;
1763:   PetscInt dim, Np;

1765:   PetscFunctionBegin;
1766:   PetscCall(TSGetDM(ts, &sw));
1767:   PetscCall(DMGetDimension(sw, &dim));
1768:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1769:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1770:   PetscCall(VecSetBlockSize(u, dim));
1771:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1772:   PetscCall(VecSetUp(u));
1773:   PetscCall(TSSetSolution(ts, u));
1774:   PetscCall(VecDestroy(&u));
1775:   PetscFunctionReturn(PETSC_SUCCESS);
1776: }

1778: static PetscErrorCode SetProblem(TS ts)
1779: {
1780:   AppCtx *user;
1781:   DM      sw;

1783:   PetscFunctionBegin;
1784:   PetscCall(TSGetDM(ts, &sw));
1785:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1786:   // Define unified system for (X, V)
1787:   {
1788:     Mat      J;
1789:     PetscInt dim, Np;

1791:     PetscCall(DMGetDimension(sw, &dim));
1792:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1793:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1794:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1795:     PetscCall(MatSetBlockSize(J, 2 * dim));
1796:     PetscCall(MatSetFromOptions(J));
1797:     PetscCall(MatSetUp(J));
1798:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1799:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1800:     PetscCall(MatDestroy(&J));
1801:   }
1802:   /* Define split system for X and V */
1803:   {
1804:     Vec             u;
1805:     IS              isx, isv, istmp;
1806:     const PetscInt *idx;
1807:     PetscInt        dim, Np, rstart;

1809:     PetscCall(TSGetSolution(ts, &u));
1810:     PetscCall(DMGetDimension(sw, &dim));
1811:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1812:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1813:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1814:     PetscCall(ISGetIndices(istmp, &idx));
1815:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1816:     PetscCall(ISRestoreIndices(istmp, &idx));
1817:     PetscCall(ISDestroy(&istmp));
1818:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1819:     PetscCall(ISGetIndices(istmp, &idx));
1820:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1821:     PetscCall(ISRestoreIndices(istmp, &idx));
1822:     PetscCall(ISDestroy(&istmp));
1823:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1824:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1825:     PetscCall(ISDestroy(&isx));
1826:     PetscCall(ISDestroy(&isv));
1827:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1828:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1829:   }
1830:   PetscFunctionReturn(PETSC_SUCCESS);
1831: }

1833: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1834: {
1835:   DM        sw;
1836:   Vec       u;
1837:   PetscReal t, maxt, dt;
1838:   PetscInt  n, maxn;

1840:   PetscFunctionBegin;
1841:   PetscCall(TSGetDM(ts, &sw));
1842:   PetscCall(TSGetTime(ts, &t));
1843:   PetscCall(TSGetMaxTime(ts, &maxt));
1844:   PetscCall(TSGetTimeStep(ts, &dt));
1845:   PetscCall(TSGetStepNumber(ts, &n));
1846:   PetscCall(TSGetMaxSteps(ts, &maxn));

1848:   PetscCall(TSReset(ts));
1849:   PetscCall(TSSetDM(ts, sw));
1850:   PetscCall(TSSetFromOptions(ts));
1851:   PetscCall(TSSetTime(ts, t));
1852:   PetscCall(TSSetMaxTime(ts, maxt));
1853:   PetscCall(TSSetTimeStep(ts, dt));
1854:   PetscCall(TSSetStepNumber(ts, n));
1855:   PetscCall(TSSetMaxSteps(ts, maxn));

1857:   PetscCall(CreateSolution(ts));
1858:   PetscCall(SetProblem(ts));
1859:   PetscCall(TSGetSolution(ts, &u));
1860:   PetscFunctionReturn(PETSC_SUCCESS);
1861: }

1863: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
1864: {
1865:   DM        sw, cdm;
1866:   PetscInt  Np;
1867:   PetscReal low[2], high[2];
1868:   AppCtx   *user = (AppCtx *)ctx;

1870:   sw = user->swarm;
1871:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
1872:   // Get the bounding box so we can equally space the particles
1873:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
1874:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1875:   // shift it by h/2 so nothing is initialized directly on a boundary
1876:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
1877:   x[1] = 0.;
1878:   return PETSC_SUCCESS;
1879: }

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

1884:   Input Parameters:
1885: + ts         - The TS
1886: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

1888:   Output Parameters:
1889: . u - The initialized solution vector

1891:   Level: advanced

1893: .seealso: InitializeSolve()
1894: */
1895: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1896: {
1897:   DM       sw;
1898:   Vec      u, gc, gv, gc0, gv0;
1899:   IS       isx, isv;
1900:   PetscInt dim;
1901:   AppCtx  *user;

1903:   PetscFunctionBeginUser;
1904:   PetscCall(TSGetDM(ts, &sw));
1905:   PetscCall(DMGetApplicationContext(sw, &user));
1906:   PetscCall(DMGetDimension(sw, &dim));
1907:   if (useInitial) {
1908:     PetscReal v0[2] = {1., 0.};
1909:     if (user->perturbed_weights) {
1910:       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1911:     } else {
1912:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
1913:       PetscCall(DMSwarmInitializeCoordinates(sw));
1914:       if (user->fake_1D) {
1915:         PetscCall(InitializeVelocities_Fake1D(sw, user));
1916:       } else {
1917:         PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
1918:       }
1919:     }
1920:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1921:     PetscCall(DMSwarmTSRedistribute(ts));
1922:   }
1923:   PetscCall(TSGetSolution(ts, &u));
1924:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1925:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1926:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1927:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
1928:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1929:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
1930:   if (useInitial) {
1931:     PetscCall(VecCopy(gc, gc0));
1932:     PetscCall(VecCopy(gv, gv0));
1933:   }
1934:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1935:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1936:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1937:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
1938:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1939:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
1940:   PetscFunctionReturn(PETSC_SUCCESS);
1941: }

1943: static PetscErrorCode InitializeSolve(TS ts, Vec u)
1944: {
1945:   PetscFunctionBegin;
1946:   PetscCall(TSSetSolution(ts, u));
1947:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1948:   PetscFunctionReturn(PETSC_SUCCESS);
1949: }

1951: static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
1952: {
1953:   MPI_Comm           comm;
1954:   DM                 sw;
1955:   AppCtx            *user;
1956:   const PetscScalar *u;
1957:   const PetscReal   *coords, *vel;
1958:   PetscScalar       *e;
1959:   PetscReal          t;
1960:   PetscInt           dim, Np, p;

1962:   PetscFunctionBeginUser;
1963:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
1964:   PetscCall(TSGetDM(ts, &sw));
1965:   PetscCall(DMGetApplicationContext(sw, &user));
1966:   PetscCall(DMGetDimension(sw, &dim));
1967:   PetscCall(TSGetSolveTime(ts, &t));
1968:   PetscCall(VecGetArray(E, &e));
1969:   PetscCall(VecGetArrayRead(U, &u));
1970:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1971:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1972:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1973:   Np /= 2 * dim;
1974:   for (p = 0; p < Np; ++p) {
1975:     /* TODO generalize initial conditions and project into plane instead of assuming x-y */
1976:     const PetscReal    r0    = DMPlex_NormD_Internal(dim, &coords[p * dim]);
1977:     const PetscReal    th0   = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]);
1978:     const PetscReal    v0    = DMPlex_NormD_Internal(dim, &vel[p * dim]);
1979:     const PetscReal    omega = v0 / r0;
1980:     const PetscReal    ct    = PetscCosReal(omega * t + th0);
1981:     const PetscReal    st    = PetscSinReal(omega * t + th0);
1982:     const PetscScalar *x     = &u[(p * 2 + 0) * dim];
1983:     const PetscScalar *v     = &u[(p * 2 + 1) * dim];
1984:     const PetscReal    xe[3] = {r0 * ct, r0 * st, 0.0};
1985:     const PetscReal    ve[3] = {-v0 * st, v0 * ct, 0.0};
1986:     PetscInt           d;

1988:     for (d = 0; d < dim; ++d) {
1989:       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
1990:       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
1991:     }
1992:     if (user->error) {
1993:       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
1994:       const PetscReal exen = 0.5 * PetscSqr(v0);
1995:       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)));
1996:     }
1997:   }
1998:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1999:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
2000:   PetscCall(VecRestoreArrayRead(U, &u));
2001:   PetscCall(VecRestoreArray(E, &e));
2002:   PetscFunctionReturn(PETSC_SUCCESS);
2003: }

2005: static PetscErrorCode MigrateParticles(TS ts)
2006: {
2007:   DM               sw, cdm;
2008:   const PetscReal *L;

2010:   PetscFunctionBeginUser;
2011:   PetscCall(TSGetDM(ts, &sw));
2012:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
2013:   {
2014:     Vec        u, gc, gv, position, momentum;
2015:     IS         isx, isv;
2016:     PetscReal *pos, *mom;

2018:     PetscCall(TSGetSolution(ts, &u));
2019:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2020:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2021:     PetscCall(VecGetSubVector(u, isx, &position));
2022:     PetscCall(VecGetSubVector(u, isv, &momentum));
2023:     PetscCall(VecGetArray(position, &pos));
2024:     PetscCall(VecGetArray(momentum, &mom));
2025:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2026:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2027:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
2028:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

2030:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
2031:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
2032:     if ((L[0] || L[1]) >= 0.) {
2033:       PetscReal *x, *v, upper[3], lower[3];
2034:       PetscInt   Np, dim;

2036:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
2037:       PetscCall(DMGetDimension(cdm, &dim));
2038:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
2039:       PetscCall(VecGetArray(gc, &x));
2040:       PetscCall(VecGetArray(gv, &v));
2041:       for (PetscInt p = 0; p < Np; ++p) {
2042:         for (PetscInt d = 0; d < dim; ++d) {
2043:           if (pos[p * dim + d] < lower[d]) {
2044:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
2045:           } else if (pos[p * dim + d] > upper[d]) {
2046:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
2047:           } else {
2048:             x[p * dim + d] = pos[p * dim + d];
2049:           }
2050:           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]);
2051:           v[p * dim + d] = mom[p * dim + d];
2052:         }
2053:       }
2054:       PetscCall(VecRestoreArray(gc, &x));
2055:       PetscCall(VecRestoreArray(gv, &v));
2056:     }
2057:     PetscCall(VecRestoreArray(position, &pos));
2058:     PetscCall(VecRestoreArray(momentum, &mom));
2059:     PetscCall(VecRestoreSubVector(u, isx, &position));
2060:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2061:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2062:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2063:   }
2064:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2065:   PetscCall(DMSwarmTSRedistribute(ts));
2066:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2067:   PetscFunctionReturn(PETSC_SUCCESS);
2068: }

2070: int main(int argc, char **argv)
2071: {
2072:   DM        dm, sw;
2073:   TS        ts;
2074:   Vec       u;
2075:   PetscReal dt;
2076:   PetscInt  maxn;
2077:   AppCtx    user;

2079:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2080:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2081:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2082:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2083:   PetscCall(CreatePoisson(dm, &user));
2084:   PetscCall(CreateSwarm(dm, &user, &sw));
2085:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2086:   PetscCall(InitializeConstants(sw, &user));
2087:   PetscCall(DMSetApplicationContext(sw, &user));

2089:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2090:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2091:   PetscCall(TSSetDM(ts, sw));
2092:   PetscCall(TSSetMaxTime(ts, 0.1));
2093:   PetscCall(TSSetTimeStep(ts, 0.00001));
2094:   PetscCall(TSSetMaxSteps(ts, 100));
2095:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

2097:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2098:   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
2099:   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2100:   if (user.monitor_positions) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2101:   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));

2103:   PetscCall(TSSetFromOptions(ts));
2104:   PetscCall(TSGetTimeStep(ts, &dt));
2105:   PetscCall(TSGetMaxSteps(ts, &maxn));
2106:   user.steps    = maxn;
2107:   user.stepSize = dt;
2108:   PetscCall(SetupContext(dm, sw, &user));
2109:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2110:   PetscCall(TSSetComputeExactError(ts, ComputeError));
2111:   PetscCall(TSSetPostStep(ts, MigrateParticles));
2112:   PetscCall(CreateSolution(ts));
2113:   PetscCall(TSGetSolution(ts, &u));
2114:   PetscCall(TSComputeInitialCondition(ts, u));
2115:   PetscCall(CheckNonNegativeWeights(sw, &user));
2116:   PetscCall(TSSolve(ts, NULL));

2118:   PetscCall(SNESDestroy(&user.snes));
2119:   PetscCall(MatDestroy(&user.M));
2120:   PetscCall(PetscFEGeomDestroy(&user.fegeom));
2121:   PetscCall(TSDestroy(&ts));
2122:   PetscCall(DMDestroy(&sw));
2123:   PetscCall(DMDestroy(&dm));
2124:   PetscCall(DestroyContext(&user));
2125:   PetscCall(PetscFinalize());
2126:   return 0;
2127: }

2129: /*TEST

2131:    build:
2132:     requires: !complex double

2134:   # This tests that we can put particles in a box and compute the Coulomb force
2135:   # Recommend -draw_size 500,500
2136:    testset:
2137:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2138:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 20,1 \
2139:              -dm_plex_box_lower 0,-1 -dm_plex_box_upper 12.5664,1 \
2140:              -dm_plex_box_bd periodic,none \
2141:            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2142:            -sigma 1.0e-8 -timeScale 2.0e-14 \
2143:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2144:              -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 0 \
2145:            -output_step 50 -check_vel_res
2146:      test:
2147:        suffix: none_1d
2148:        args: -em_type none -error
2149:      test:
2150:        suffix: coulomb_1d
2151:        args: -em_type coulomb

2153:    # for viewers
2154:    #-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
2155:    testset:
2156:      nsize: {{1 2}}
2157:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2158:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 36,1 \
2159:              -dm_plex_box_lower 0.,-0.5 -dm_plex_box_upper 12.5664,0.5 \
2160:              -dm_plex_box_bd periodic,none \
2161:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2162:              -vdm_plex_box_lower -3 -vdm_plex_box_upper 3 \
2163:            -dm_swarm_num_species 1 -twostream -charges -1.,1. -sigma 1.0e-8 \
2164:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2165:            -ts_type basicsymplectic -ts_basicsymplectic_type 2 \
2166:              -ts_dt 0.01 -ts_max_time 5 -ts_max_steps 10 \
2167:            -em_snes_atol 1.e-15 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2168:            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1
2169:      test:
2170:        suffix: two_stream_c0
2171:        args: -em_type primal -petscfe_default_quadrature_order 2 -petscspace_degree 2 -em_pc_type svd
2172:      test:
2173:        suffix: two_stream_rt
2174:        requires: superlu_dist
2175:        args: -em_type mixed \
2176:                -potential_petscspace_degree 0 \
2177:                -potential_petscdualspace_lagrange_use_moments \
2178:                -potential_petscdualspace_lagrange_moment_order 2 \
2179:                -field_petscspace_degree 2 -field_petscfe_default_quadrature_order 1 \
2180:                -field_petscspace_type sum \
2181:                  -field_petscspace_variables 2 \
2182:                  -field_petscspace_components 2 \
2183:                  -field_petscspace_sum_spaces 2 \
2184:                  -field_petscspace_sum_concatenate true \
2185:                  -field_sumcomp_0_petscspace_variables 2 \
2186:                  -field_sumcomp_0_petscspace_type tensor \
2187:                  -field_sumcomp_0_petscspace_tensor_spaces 2 \
2188:                  -field_sumcomp_0_petscspace_tensor_uniform false \
2189:                  -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
2190:                  -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
2191:                  -field_sumcomp_1_petscspace_variables 2 \
2192:                  -field_sumcomp_1_petscspace_type tensor \
2193:                  -field_sumcomp_1_petscspace_tensor_spaces 2 \
2194:                  -field_sumcomp_1_petscspace_tensor_uniform false \
2195:                  -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
2196:                  -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
2197:                -field_petscdualspace_form_degree -1 \
2198:                -field_petscdualspace_order 1 \
2199:                -field_petscdualspace_lagrange_trimmed true \
2200:              -em_snes_error_if_not_converged \
2201:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2202:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2203:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2204:                -em_fieldsplit_field_pc_type lu \
2205:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2206:                -em_fieldsplit_potential_pc_type svd

2208:    # For an eyeball check, we use
2209:    # -ts_max_steps 1000 -dm_plex_box_faces 10,1 -vdm_plex_box_faces 2000 -monitor_efield
2210:    # For verification, we use
2211:    # -ts_max_steps 1000 -dm_plex_box_faces 100,1 -vdm_plex_box_faces 8000 -dm_swarm_num_particles 800000 -monitor_efield
2212:    # -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
2213:    testset:
2214:      nsize: {{1 2}}
2215:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2216:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 10,1 \
2217:              -dm_plex_box_lower 0.,-0.5 -dm_plex_box_upper 12.5664,0.5 \
2218:              -dm_plex_box_bd periodic,none \
2219:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2220:              -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 \
2221:              -vpetscspace_degree 2 -vdm_plex_hash_location \
2222:            -dm_swarm_num_species 1 -charges -1.,1. \
2223:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2224:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2225:              -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 1 \
2226:            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2227:            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1

2229:      test:
2230:        suffix: uniform_equilibrium_1d
2231:        args: -cosine_coefficients 0.0,0.5 -em_type primal -petscspace_degree 1 -em_pc_type svd
2232:      test:
2233:        suffix: uniform_equilibrium_1d_real
2234:        args: -dm_plex_dim 1 -dm_plex_simplex 1 -fake_1D 0 -dm_plex_box_faces 10 \
2235:                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2236:              -cosine_coefficients 0.0 -em_type primal -petscspace_degree 1 -em_pc_type svd
2237:      test:
2238:        suffix: uniform_primal_1d
2239:        args: -em_type primal -petscspace_degree 1 -em_pc_type svd
2240:      test:
2241:        suffix: uniform_primal_1d_real
2242:        args: -dm_plex_dim 1 -dm_plex_simplex 1 -fake_1D 0 -dm_plex_box_faces 10 \
2243:                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2244:              -cosine_coefficients 0.01 -em_type primal -petscspace_degree 1 -em_pc_type svd
2245:      test:
2246:        requires: superlu_dist
2247:        suffix: uniform_mixed_1d
2248:        args: -em_type mixed \
2249:                -potential_petscspace_degree 0 \
2250:                -potential_petscdualspace_lagrange_use_moments \
2251:                -potential_petscdualspace_lagrange_moment_order 2 \
2252:                -field_petscspace_degree 2 -field_petscfe_default_quadrature_order 1 \
2253:                -field_petscspace_type sum \
2254:                  -field_petscspace_variables 2 \
2255:                  -field_petscspace_components 2 \
2256:                  -field_petscspace_sum_spaces 2 \
2257:                  -field_petscspace_sum_concatenate true \
2258:                  -field_sumcomp_0_petscspace_variables 2 \
2259:                  -field_sumcomp_0_petscspace_type tensor \
2260:                  -field_sumcomp_0_petscspace_tensor_spaces 2 \
2261:                  -field_sumcomp_0_petscspace_tensor_uniform false \
2262:                  -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
2263:                  -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
2264:                  -field_sumcomp_1_petscspace_variables 2 \
2265:                  -field_sumcomp_1_petscspace_type tensor \
2266:                  -field_sumcomp_1_petscspace_tensor_spaces 2 \
2267:                  -field_sumcomp_1_petscspace_tensor_uniform false \
2268:                  -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
2269:                  -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
2270:                -field_petscdualspace_form_degree -1 \
2271:                -field_petscdualspace_order 1 \
2272:                -field_petscdualspace_lagrange_trimmed true \
2273:              -em_snes_error_if_not_converged \
2274:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2275:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2276:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2277:                -em_fieldsplit_field_pc_type lu \
2278:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2279:                -em_fieldsplit_potential_pc_type svd

2281: TEST*/