Actual source code: ex2.c

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

  3: /*
  4:   TODO:
  5:   - Cache mesh geometry
  6:   - Move electrostatic solver to MG+SVD

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

 11:   To visualize the maximum electric field use

 13:     -efield_monitor

 15:   To monitor velocity moments of the distribution use

 17:     -ptof_pc_type lu -moments_monitor

 19:   To monitor the particle positions in phase space use

 21:     -positions_monitor

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

 25:     -poisson_monitor

 27:   To monitor the remapping field use

 29:     -remap_uf_view draw

 31:   To visualize the swarm distribution use

 33:     -ts_monitor_hg_swarm

 35:   To visualize the particles, we can use

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

 39: For a Landau Damping verification run, we use

 41:   # Physics
 42:   -cosine_coefficients 0.01 -dm_swarm_num_species 1 -charges -1. -perturbed_weights -total_weight 1.
 43:   # Spatial Mesh
 44:   -dm_plex_dim 1 -dm_plex_box_faces 40 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic  -dm_plex_hash_location
 45:   # Velocity Mesh
 46:   -vdm_plex_dim 1 -vdm_plex_box_faces 40 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 -vpetscspace_degree 2 -vdm_plex_hash_location
 47:   # Remap Space
 48:   -dm_swarm_remap_type pfak -remap_freq 100
 49:   -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 -remap_dm_plex_box_faces 20,20 -remap_dm_plex_box_bd periodic,none -remap_dm_plex_box_lower 0.,-6.
 50:     -remap_dm_plex_box_upper 12.5664,6. -remap_petscspace_degree 1 -remap_dm_plex_hash_location
 51:   # Remap Solve
 52:   -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu
 53:   # EM Solve
 54:   -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu
 55:   # Timestepping
 56:   -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_dt 0.03 -ts_max_steps 1500 -ts_max_time 100
 57:   # Monitoring
 58:   -output_step 1 -check_vel_res -efield_monitor -poisson_monitor -positions_monitor -dm_swarm_print_coords 0 -remap_uf_view draw
 59:     -ftop_ksp_lsqr_monitor -ftop_ksp_converged_reason

 61: */
 62: #include <petscts.h>
 63: #include <petscdmplex.h>
 64: #include <petscdmswarm.h>
 65: #include <petscfe.h>
 66: #include <petscds.h>
 67: #include <petscbag.h>
 68: #include <petscdraw.h>
 69: #include <petsc/private/dmpleximpl.h>
 70: #include <petsc/private/petscfeimpl.h>
 71: #include <petsc/private/dmswarmimpl.h>
 72: #include "petscdm.h"
 73: #include "petscdmlabel.h"

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

 78: const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
 79: typedef enum {
 80:   EM_PRIMAL,
 81:   EM_MIXED,
 82:   EM_COULOMB,
 83:   EM_NONE
 84: } EMType;

 86: typedef enum {
 87:   V0,
 88:   X0,
 89:   T0,
 90:   M0,
 91:   Q0,
 92:   PHI0,
 93:   POISSON,
 94:   VLASOV,
 95:   SIGMA,
 96:   NUM_CONSTANTS
 97: } ConstantType;
 98: typedef struct {
 99:   PetscScalar v0; /* Velocity scale, often the thermal velocity */
100:   PetscScalar t0; /* Time scale */
101:   PetscScalar x0; /* Space scale */
102:   PetscScalar m0; /* Mass scale */
103:   PetscScalar q0; /* Charge scale */
104:   PetscScalar kb;
105:   PetscScalar epsi0;
106:   PetscScalar phi0;          /* Potential scale */
107:   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
108:   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
109:   PetscReal   sigma;         /* Nondimensional charge per length in x */
110: } Parameter;

112: typedef struct {
113:   PetscBag     bag;               // Problem parameters
114:   PetscBool    error;             // Flag for printing the error
115:   PetscInt     remapFreq;         // Number of timesteps between remapping
116:   PetscBool    efield_monitor;    // Flag to show electric field monitor
117:   PetscBool    moment_monitor;    // Flag to show distribution moment monitor
118:   PetscBool    positions_monitor; // Flag to show particle positins at each time step
119:   PetscBool    poisson_monitor;   // Flag to display charge, E field, and potential at each solve
120:   PetscBool    initial_monitor;   // Flag to monitor the initial conditions
121:   PetscInt     velocity_monitor;  // Cell to monitor the velocity distribution for
122:   PetscBool    perturbed_weights; // Uniformly sample x,v space with gaussian weights
123:   PetscInt     ostep;             // Print the energy at each ostep time steps
124:   PetscInt     numParticles;
125:   PetscReal    timeScale;              /* Nondimensionalizing time scale */
126:   PetscReal    charges[2];             /* The charges of each species */
127:   PetscReal    masses[2];              /* The masses of each species */
128:   PetscReal    thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
129:   PetscReal    cosine_coefficients[2]; /*(alpha, k)*/
130:   PetscReal    totalWeight;
131:   PetscReal    stepSize;
132:   PetscInt     steps;
133:   PetscReal    initVel;
134:   EMType       em;           // Type of electrostatic model
135:   SNES         snes;         // EM solver
136:   DM           dmPot;        // The DM for potential
137:   Mat          fftPot;       // Fourier Transform operator for the potential
138:   Vec          fftX, fftY;   //   FFT vectors with phases added (complex parts)
139:   IS           fftReal;      //   The indices for real parts
140:   IS           isPot;        // The IS for potential, or NULL in primal
141:   Mat          M;            // The finite element mass matrix for potential
142:   PetscFEGeom *fegeom;       // Geometric information for the DM cells
143:   PetscDrawHG  drawhgic_x;   // Histogram of the particle weight in each X cell
144:   PetscDrawHG  drawhgic_v;   // Histogram of the particle weight in each X cell
145:   PetscDrawHG  drawhgcell_v; // Histogram of the particle weight in a given cell
146:   PetscBool    validE;       // Flag to indicate E-field in swarm is valid
147:   PetscReal    drawlgEmin;   // The minimum lg(E) to plot
148:   PetscDrawLG  drawlgE;      // Logarithm of maximum electric field
149:   PetscDrawSP  drawspE;      // Electric field at particle positions
150:   PetscDrawSP  drawspX;      // Particle positions
151:   PetscViewer  viewerRho;    // Charge density viewer
152:   PetscViewer  viewerRhoHat; // Charge density Fourier Transform viewer
153:   PetscViewer  viewerPhi;    // Potential viewer
154:   DM           swarm;
155:   PetscRandom  random;
156:   PetscBool    twostream;
157:   PetscBool    checkweights;
158:   PetscInt     checkVRes; // Flag to check/output velocity residuals for nightly tests

160:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
161: } AppCtx;

163: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
164: {
165:   PetscFunctionBeginUser;
166:   PetscInt d                      = 2;
167:   PetscInt maxSpecies             = 2;
168:   options->error                  = PETSC_FALSE;
169:   options->remapFreq              = 1;
170:   options->efield_monitor         = PETSC_FALSE;
171:   options->moment_monitor         = PETSC_FALSE;
172:   options->initial_monitor        = PETSC_FALSE;
173:   options->perturbed_weights      = PETSC_FALSE;
174:   options->poisson_monitor        = PETSC_FALSE;
175:   options->positions_monitor      = PETSC_FALSE;
176:   options->velocity_monitor       = -1;
177:   options->ostep                  = 100;
178:   options->timeScale              = 2.0e-14;
179:   options->charges[0]             = -1.0;
180:   options->charges[1]             = 1.0;
181:   options->masses[0]              = 1.0;
182:   options->masses[1]              = 1000.0;
183:   options->thermal_energy[0]      = 1.0;
184:   options->thermal_energy[1]      = 1.0;
185:   options->cosine_coefficients[0] = 0.01;
186:   options->cosine_coefficients[1] = 0.5;
187:   options->initVel                = 1;
188:   options->totalWeight            = 1.0;
189:   options->drawhgic_x             = NULL;
190:   options->drawhgic_v             = NULL;
191:   options->drawhgcell_v           = NULL;
192:   options->validE                 = PETSC_FALSE;
193:   options->drawlgEmin             = -6;
194:   options->drawlgE                = NULL;
195:   options->drawspE                = NULL;
196:   options->drawspX                = NULL;
197:   options->viewerRho              = NULL;
198:   options->viewerRhoHat           = NULL;
199:   options->viewerPhi              = NULL;
200:   options->em                     = EM_COULOMB;
201:   options->snes                   = NULL;
202:   options->dmPot                  = NULL;
203:   options->fftPot                 = NULL;
204:   options->fftX                   = NULL;
205:   options->fftY                   = NULL;
206:   options->fftReal                = NULL;
207:   options->isPot                  = NULL;
208:   options->M                      = NULL;
209:   options->numParticles           = 32768;
210:   options->twostream              = PETSC_FALSE;
211:   options->checkweights           = PETSC_FALSE;
212:   options->checkVRes              = 0;

214:   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
215:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex2.c", options->error, &options->error, NULL));
216:   PetscCall(PetscOptionsInt("-remap_freq", "Number", "ex2.c", options->remapFreq, &options->remapFreq, NULL));
217:   PetscCall(PetscOptionsBool("-efield_monitor", "Flag to plot log(max E) over time", "ex2.c", options->efield_monitor, &options->efield_monitor, NULL));
218:   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", "ex2.c", options->drawlgEmin, &options->drawlgEmin, NULL));
219:   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", "ex2.c", options->moment_monitor, &options->moment_monitor, NULL));
220:   PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", "ex2.c", options->initial_monitor, &options->initial_monitor, NULL));
221:   PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", "ex2.c", options->positions_monitor, &options->positions_monitor, NULL));
222:   PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", "ex2.c", options->poisson_monitor, &options->poisson_monitor, NULL));
223:   PetscCall(PetscOptionsInt("-velocity_monitor", "Cell to show velocity histograms", "ex2.c", options->velocity_monitor, &options->velocity_monitor, NULL));
224:   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", "ex2.c", options->twostream, &options->twostream, NULL));
225:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", "ex2.c", options->perturbed_weights, &options->perturbed_weights, NULL));
226:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex2.c", options->checkweights, &options->checkweights, NULL));
227:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex2.c", options->checkVRes, &options->checkVRes, NULL));
228:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex2.c", options->ostep, &options->ostep, NULL));
229:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex2.c", options->timeScale, &options->timeScale, NULL));
230:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex2.c", options->initVel, &options->initVel, NULL));
231:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex2.c", options->totalWeight, &options->totalWeight, NULL));
232:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex2.c", options->cosine_coefficients, &d, NULL));
233:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex2.c", options->charges, &maxSpecies, NULL));
234:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex2.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
235:   PetscOptionsEnd();

237:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
238:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
239:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
240:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
241:   PetscFunctionReturn(PETSC_SUCCESS);
242: }

244: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
245: {
246:   MPI_Comm comm;

248:   PetscFunctionBeginUser;
249:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
250:   if (user->efield_monitor) {
251:     PetscDraw     draw;
252:     PetscDrawAxis axis;

254:     PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
255:     PetscCall(PetscDrawSetSave(draw, "ex2_Efield"));
256:     PetscCall(PetscDrawSetFromOptions(draw));
257:     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
258:     PetscCall(PetscDrawDestroy(&draw));
259:     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
260:     PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
261:     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
262:   }

264:   if (user->initial_monitor) {
265:     PetscDraw     drawic_x, drawic_v;
266:     PetscDrawAxis axis1, axis2;
267:     PetscReal     dmboxlower[2], dmboxupper[2];
268:     PetscInt      dim, cStart, cEnd;

270:     PetscCall(DMGetDimension(sw, &dim));
271:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
272:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

274:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &drawic_x));
275:     PetscCall(PetscDrawSetSave(drawic_x, "ex2_ic_x"));
276:     PetscCall(PetscDrawSetFromOptions(drawic_x));
277:     PetscCall(PetscDrawHGCreate(drawic_x, (int)dim, &user->drawhgic_x));
278:     PetscCall(PetscDrawHGCalcStats(user->drawhgic_x, PETSC_TRUE));
279:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
280:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart)));
281:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "weight"));
282:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 0));
283:     PetscCall(PetscDrawDestroy(&drawic_x));

285:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &drawic_v));
286:     PetscCall(PetscDrawSetSave(drawic_v, "ex9_ic_v"));
287:     PetscCall(PetscDrawSetFromOptions(drawic_v));
288:     PetscCall(PetscDrawHGCreate(drawic_v, (int)dim, &user->drawhgic_v));
289:     PetscCall(PetscDrawHGCalcStats(user->drawhgic_v, PETSC_TRUE));
290:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
291:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 21));
292:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "weight"));
293:     PetscCall(PetscDrawAxisSetLimits(axis2, -6, 6, 0, 0));
294:     PetscCall(PetscDrawDestroy(&drawic_v));
295:   }

297:   if (user->velocity_monitor >= 0) {
298:     DM            vdm;
299:     DMSwarmCellDM celldm;
300:     PetscDraw     drawcell_v;
301:     PetscDrawAxis axis;
302:     PetscReal     dmboxlower[2], dmboxupper[2];
303:     PetscInt      dim;
304:     char          title[PETSC_MAX_PATH_LEN];

306:     PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
307:     PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
308:     PetscCall(DMGetDimension(vdm, &dim));
309:     PetscCall(DMGetBoundingBox(vdm, dmboxlower, dmboxupper));

311:     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Cell %" PetscInt_FMT ": Velocity Distribution", user->velocity_monitor));
312:     PetscCall(PetscDrawCreate(comm, NULL, title, 400, 300, 400, 300, &drawcell_v));
313:     PetscCall(PetscDrawSetSave(drawcell_v, "ex2_cell_v"));
314:     PetscCall(PetscDrawSetFromOptions(drawcell_v));
315:     PetscCall(PetscDrawHGCreate(drawcell_v, (int)dim, &user->drawhgcell_v));
316:     PetscCall(PetscDrawHGCalcStats(user->drawhgcell_v, PETSC_TRUE));
317:     PetscCall(PetscDrawHGGetAxis(user->drawhgcell_v, &axis));
318:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgcell_v, 21));
319:     PetscCall(PetscDrawAxisSetLabels(axis, "V_x Distribution", "V", "weight"));
320:     PetscCall(PetscDrawAxisSetLimits(axis, dmboxlower[0], dmboxupper[0], 0, 0));
321:     PetscCall(PetscDrawDestroy(&drawcell_v));
322:   }

324:   if (user->positions_monitor) {
325:     PetscDraw     draw;
326:     PetscDrawAxis axis;

328:     PetscCall(PetscDrawCreate(comm, NULL, "Particle Position", 0, 0, 400, 300, &draw));
329:     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
330:     PetscCall(PetscDrawSetFromOptions(draw));
331:     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspX));
332:     PetscCall(PetscDrawDestroy(&draw));
333:     PetscCall(PetscDrawSPSetDimension(user->drawspX, 1));
334:     PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis));
335:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
336:     PetscCall(PetscDrawSPReset(user->drawspX));
337:   }
338:   if (user->poisson_monitor) {
339:     Vec           rho, rhohat, phi;
340:     PetscDraw     draw;
341:     PetscDrawAxis axis;

343:     PetscCall(PetscDrawCreate(comm, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
344:     PetscCall(PetscDrawSetFromOptions(draw));
345:     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
346:     PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspE));
347:     PetscCall(PetscDrawDestroy(&draw));
348:     PetscCall(PetscDrawSPSetDimension(user->drawspE, 1));
349:     PetscCall(PetscDrawSPGetAxis(user->drawspE, &axis));
350:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
351:     PetscCall(PetscDrawSPReset(user->drawspE));

353:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Charge Density", 0, 0, 400, 300, &user->viewerRho));
354:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRho, "rho_"));
355:     PetscCall(PetscViewerDrawGetDraw(user->viewerRho, 0, &draw));
356:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
357:     PetscCall(PetscViewerSetFromOptions(user->viewerRho));
358:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
359:     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
360:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));

362:     PetscInt dim, N;

364:     PetscCall(DMGetDimension(user->dmPot, &dim));
365:     if (dim == 1) {
366:       PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
367:       PetscCall(VecGetSize(rhohat, &N));
368:       PetscCall(MatCreateFFT(comm, dim, &N, MATFFTW, &user->fftPot));
369:       PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
370:       PetscCall(MatCreateVecs(user->fftPot, &user->fftX, &user->fftY));
371:       PetscCall(ISCreateStride(PETSC_COMM_SELF, N, 0, 1, &user->fftReal));
372:     }

374:     PetscCall(PetscViewerDrawOpen(comm, NULL, "rhohat: Charge Density FT", 0, 0, 400, 300, &user->viewerRhoHat));
375:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRhoHat, "rhohat_"));
376:     PetscCall(PetscViewerDrawGetDraw(user->viewerRhoHat, 0, &draw));
377:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_ft"));
378:     PetscCall(PetscViewerSetFromOptions(user->viewerRhoHat));
379:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
380:     PetscCall(PetscObjectSetName((PetscObject)rhohat, "charge_density_ft"));
381:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));

383:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Potential", 400, 0, 400, 300, &user->viewerPhi));
384:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPhi, "phi_"));
385:     PetscCall(PetscViewerDrawGetDraw(user->viewerPhi, 0, &draw));
386:     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
387:     PetscCall(PetscViewerSetFromOptions(user->viewerPhi));
388:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
389:     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
390:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
391:   }
392:   PetscFunctionReturn(PETSC_SUCCESS);
393: }

395: static PetscErrorCode DestroyContext(AppCtx *user)
396: {
397:   PetscFunctionBeginUser;
398:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
399:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
400:   PetscCall(PetscDrawHGDestroy(&user->drawhgcell_v));

402:   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
403:   PetscCall(PetscDrawSPDestroy(&user->drawspE));
404:   PetscCall(PetscDrawSPDestroy(&user->drawspX));
405:   PetscCall(PetscViewerDestroy(&user->viewerRho));
406:   PetscCall(PetscViewerDestroy(&user->viewerRhoHat));
407:   PetscCall(MatDestroy(&user->fftPot));
408:   PetscCall(VecDestroy(&user->fftX));
409:   PetscCall(VecDestroy(&user->fftY));
410:   PetscCall(ISDestroy(&user->fftReal));
411:   PetscCall(PetscViewerDestroy(&user->viewerPhi));

413:   PetscCall(PetscBagDestroy(&user->bag));
414:   PetscFunctionReturn(PETSC_SUCCESS);
415: }

417: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
418: {
419:   const PetscScalar *w;
420:   PetscInt           Np;

422:   PetscFunctionBeginUser;
423:   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
424:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
425:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
426:   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]);
427:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
428:   PetscFunctionReturn(PETSC_SUCCESS);
429: }

431: static void f0_Dirichlet(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
432: {
433:   for (PetscInt d = 0; d < dim; ++d) f0[0] += 0.5 * PetscSqr(u_x[d]);
434: }

436: static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
437: {
438:   PetscDS        ds;
439:   const PetscInt field = 0;
440:   PetscInt       Nf;
441:   void          *ctx;

443:   PetscFunctionBegin;
444:   PetscCall(DMGetApplicationContext(dm, &ctx));
445:   PetscCall(DMGetDS(dm, &ds));
446:   PetscCall(PetscDSGetNumFields(ds, &Nf));
447:   PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
448:   PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
449:   PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
450:   PetscFunctionReturn(PETSC_SUCCESS);
451: }

453: static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *user)
454: {
455:   DMSwarmCellDM celldm;
456:   DM            vdm;
457:   Vec           u[1];
458:   const char   *fields[1] = {"w_q"};

460:   PetscFunctionBegin;
461:   PetscCall(DMSwarmSetCellDMActive(sw, "velocity"));
462:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
463:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
464: #if 0
465:   PetscReal  *v, pvmin[3] = {0., 0., 0.}, pvmax[3] = {0., 0., 0.}, vmin[3], vmax[3], fact = 1.;
466:   PetscInt    dim, Np;

468:   PetscCall(PetscObjectQuery((PetscObject)sw, "__vdm__", (PetscObject *)&vdm));
469:   // Check for particles outside the velocity grid
470:   PetscCall(DMGetDimension(sw, &dim));
471:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
472:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
473:   for (PetscInt p = 0; p < Np; ++p) {
474:     for (PetscInt d = 0; d < dim; ++d) {
475:       pvmin[d] = PetscMin(pvmax[d], v[p * dim + d]);
476:       pvmax[d] = PetscMax(pvmax[d], v[p * dim + d]);
477:     }
478:   }
479:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
480:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
481:   // To avoid particle loss, we enlarge the velocity grid if necessary
482:   for (PetscInt d = 0; d < dim; ++d) {
483:     fact = PetscMax(fact, pvmax[d] / vmax[d]);
484:     fact = PetscMax(fact, pvmin[d] / vmin[d]);
485:   }
486:   if (fact > 1.) {
487:     Vec coordinates, coordinatesLocal;

489:     fact *= 1.1;
490:     PetscCall(PetscPrintf(PETSC_COMM_SELF, "Expanding velocity grid by %g\n", fact));
491:     PetscCall(DMGetCoordinatesLocal(vdm, &coordinatesLocal));
492:     PetscCall(DMGetCoordinates(vdm, &coordinates));
493:     PetscCall(VecScale(coordinatesLocal, fact));
494:     PetscCall(VecScale(coordinates, fact));
495:     PetscCall(PetscGridHashDestroy(&((DM_Plex *)vdm->data)->lbox));
496:   }
497: #endif
498:   PetscCall(DMGetGlobalVector(vdm, &u[0]));
499:   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
500:   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
501:   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
502:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
503:   PetscFunctionReturn(PETSC_SUCCESS);
504: }

506: static void f0_grad_phi2(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
507: {
508:   f0[0] = 0.;
509:   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d * dim + d]);
510: }

512: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
513: {
514:   AppCtx     *user = (AppCtx *)ctx;
515:   DM          sw;
516:   PetscScalar intESq;
517:   PetscReal  *E, *x, *weight;
518:   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
519:   PetscReal   pmoments[4]; /* \int f, \int v f, \int v^2 f */
520:   PetscInt   *species, dim, Np, gNp;
521:   MPI_Comm    comm;

523:   PetscFunctionBeginUser;
524:   if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
525:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
526:   PetscCall(TSGetDM(ts, &sw));
527:   PetscCall(DMGetDimension(sw, &dim));
528:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
529:   PetscCall(DMSwarmGetSize(sw, &gNp));
530:   PetscCall(DMSwarmSortGetAccess(sw));
531:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
532:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
533:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
534:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

536:   for (PetscInt p = 0; p < Np; ++p) {
537:     for (PetscInt d = 0; d < 1; ++d) {
538:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
539:       if (temp > Emax) Emax = temp;
540:     }
541:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
542:     sum += E[p * dim];
543:     chargesum += user->charges[0] * weight[p];
544:   }
545:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
546:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
547:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;

549:   PetscDS ds;
550:   Vec     phi;

552:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
553:   PetscCall(DMGetDS(user->dmPot, &ds));
554:   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
555:   PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user));
556:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

558:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
559:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
560:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
561:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
562:   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
563:   PetscCall(PetscDrawLGDraw(user->drawlgE));
564:   PetscDraw draw;
565:   PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
566:   PetscCall(PetscDrawSave(draw));

568:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
569:   PetscCall(PetscPrintf(comm, "E: %f\t%+e\t%e\t%f\t%20.15e\t%f\t%f\t%f\t%20.15e\t%20.15e\t%20.15e\t%" PetscInt_FMT "\t(%" PetscInt_FMT ")\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[1 + dim], (double)PetscSqrtReal(intESq), gNp, step));
570:   PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
571:   PetscFunctionReturn(PETSC_SUCCESS);
572: }

574: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
575: {
576:   AppCtx   *user = (AppCtx *)ctx;
577:   DM        sw;
578:   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */

580:   PetscFunctionBeginUser;
581:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
582:   PetscCall(TSGetDM(ts, &sw));

584:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
585:   PetscCall(computeVelocityFEMMoments(sw, fmoments, user));

587:   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]));
588:   PetscFunctionReturn(PETSC_SUCCESS);
589: }

591: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
592: {
593:   AppCtx    *user = (AppCtx *)ctx;
594:   DM         sw;
595:   PetscDraw  drawic_x, drawic_v;
596:   PetscReal *weight, *pos, *vel;
597:   PetscInt   dim, Np;

599:   PetscFunctionBegin;
600:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
601:   if (step == 0) {
602:     PetscCall(TSGetDM(ts, &sw));
603:     PetscCall(DMGetDimension(sw, &dim));
604:     PetscCall(DMSwarmGetLocalSize(sw, &Np));

606:     PetscCall(PetscDrawHGReset(user->drawhgic_x));
607:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &drawic_x));
608:     PetscCall(PetscDrawClear(drawic_x));
609:     PetscCall(PetscDrawFlush(drawic_x));

611:     PetscCall(PetscDrawHGReset(user->drawhgic_v));
612:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &drawic_v));
613:     PetscCall(PetscDrawClear(drawic_v));
614:     PetscCall(PetscDrawFlush(drawic_v));

616:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
617:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
618:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
619:     for (PetscInt p = 0; p < Np; ++p) {
620:       PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_x, pos[p * dim], weight[p]));
621:       PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_v, vel[p * dim], weight[p]));
622:     }
623:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
624:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
625:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));

627:     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
628:     PetscCall(PetscDrawHGSave(user->drawhgic_x));
629:     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
630:     PetscCall(PetscDrawHGSave(user->drawhgic_v));
631:   }
632:   PetscFunctionReturn(PETSC_SUCCESS);
633: }

635: // Right now, make the complete velocity histogram
636: PetscErrorCode MonitorVelocity(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
637: {
638:   AppCtx      *user = (AppCtx *)ctx;
639:   DM           sw, dm;
640:   Vec          ks;
641:   PetscProbFn *cdf;
642:   PetscDraw    drawcell_v;
643:   PetscScalar *ksa;
644:   PetscReal   *weight, *vel;
645:   PetscInt    *pidx;
646:   PetscInt     dim, Npc, cStart, cEnd, cell = user->velocity_monitor;

648:   PetscFunctionBegin;
649:   PetscCall(TSGetDM(ts, &sw));
650:   PetscCall(DMGetDimension(sw, &dim));

652:   PetscCall(DMSwarmGetCellDM(sw, &dm));
653:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
654:   PetscCall(VecCreate(PetscObjectComm((PetscObject)dm), &ks));
655:   PetscCall(PetscObjectSetName((PetscObject)ks, "KS Statistic by Cell"));
656:   PetscCall(VecSetSizes(ks, cEnd - cStart, PETSC_DETERMINE));
657:   PetscCall(VecSetFromOptions(ks));
658:   switch (dim) {
659:   case 1:
660:     //cdf = PetscCDFMaxwellBoltzmann1D;
661:     cdf = PetscCDFGaussian1D;
662:     break;
663:   case 2:
664:     cdf = PetscCDFMaxwellBoltzmann2D;
665:     break;
666:   case 3:
667:     cdf = PetscCDFMaxwellBoltzmann3D;
668:     break;
669:   default:
670:     SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported", dim);
671:   }

673:   PetscCall(PetscDrawHGReset(user->drawhgcell_v));
674:   PetscCall(PetscDrawHGGetDraw(user->drawhgcell_v, &drawcell_v));
675:   PetscCall(PetscDrawClear(drawcell_v));
676:   PetscCall(PetscDrawFlush(drawcell_v));

678:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
679:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
680:   PetscCall(DMSwarmSortGetAccess(sw));
681:   PetscCall(VecGetArrayWrite(ks, &ksa));
682:   for (PetscInt c = cStart; c < cEnd; ++c) {
683:     Vec          cellv, cellw;
684:     PetscScalar *cella, *cellaw;
685:     PetscReal    totWgt = 0.;

687:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
688:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellv));
689:     PetscCall(VecSetBlockSize(cellv, dim));
690:     PetscCall(VecSetSizes(cellv, Npc * dim, Npc));
691:     PetscCall(VecSetFromOptions(cellv));
692:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellw));
693:     PetscCall(VecSetSizes(cellw, Npc, Npc));
694:     PetscCall(VecSetFromOptions(cellw));
695:     PetscCall(VecGetArrayWrite(cellv, &cella));
696:     PetscCall(VecGetArrayWrite(cellw, &cellaw));
697:     for (PetscInt q = 0; q < Npc; ++q) {
698:       const PetscInt p = pidx[q];
699:       if (c == cell) PetscCall(PetscDrawHGAddWeightedValue(user->drawhgcell_v, vel[p * dim], weight[p]));
700:       for (PetscInt d = 0; d < dim; ++d) cella[q * dim + d] = vel[p * dim + d];
701:       cellaw[q] = weight[p];
702:       totWgt += weight[p];
703:     }
704:     PetscCall(VecRestoreArrayWrite(cellv, &cella));
705:     PetscCall(VecRestoreArrayWrite(cellw, &cellaw));
706:     PetscCall(VecScale(cellw, 1. / totWgt));
707:     PetscCall(PetscProbComputeKSStatisticWeighted(cellv, cellw, cdf, &ksa[c - cStart]));
708:     PetscCall(VecDestroy(&cellv));
709:     PetscCall(VecDestroy(&cellw));
710:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
711:   }
712:   PetscCall(VecRestoreArrayWrite(ks, &ksa));
713:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
714:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
715:   PetscCall(DMSwarmSortRestoreAccess(sw));

717:   PetscReal minalpha, maxalpha;
718:   PetscInt  mincell, maxcell;

720:   PetscCall(VecFilter(ks, PETSC_SMALL));
721:   PetscCall(VecMin(ks, &mincell, &minalpha));
722:   PetscCall(VecMax(ks, &maxcell, &maxalpha));
723:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)dm), "Step %" PetscInt_FMT ": Min/Max KS statistic %g/%g in cell %" PetscInt_FMT "/%" PetscInt_FMT "\n", step, minalpha, maxalpha, mincell, maxcell));
724:   PetscCall(VecViewFromOptions(ks, NULL, "-ks_view"));
725:   PetscCall(VecDestroy(&ks));

727:   PetscCall(PetscDrawHGDraw(user->drawhgcell_v));
728:   PetscCall(PetscDrawHGSave(user->drawhgcell_v));
729:   PetscFunctionReturn(PETSC_SUCCESS);
730: }

732: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
733: {
734:   AppCtx         *user = (AppCtx *)ctx;
735:   DM              dm, sw;
736:   PetscDrawAxis   axis;
737:   char            title[1024];
738:   PetscScalar    *x, *v, *weight;
739:   PetscReal       lower[3], upper[3], speed;
740:   const PetscInt *s;
741:   PetscInt        dim, cStart, cEnd, c;

743:   PetscFunctionBeginUser;
744:   if (step > 0 && step % user->ostep == 0) {
745:     PetscCall(TSGetDM(ts, &sw));
746:     PetscCall(DMSwarmGetCellDM(sw, &dm));
747:     PetscCall(DMGetDimension(dm, &dim));
748:     PetscCall(DMGetBoundingBox(dm, lower, upper));
749:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
750:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
751:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
752:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
753:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
754:     PetscCall(DMSwarmSortGetAccess(sw));
755:     PetscCall(PetscDrawSPReset(user->drawspX));
756:     PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis));
757:     PetscCall(PetscSNPrintf(title, 1024, "Step %" PetscInt_FMT " Time: %g", step, (double)t));
758:     PetscCall(PetscDrawAxisSetLabels(axis, title, "x", "v"));
759:     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], lower[1], upper[1]));
760:     PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], -12, 12));
761:     for (c = 0; c < cEnd - cStart; ++c) {
762:       PetscInt *pidx, Npc, q;
763:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
764:       for (q = 0; q < Npc; ++q) {
765:         const PetscInt p = pidx[q];
766:         if (s[p] == 0) {
767:           speed = 0.;
768:           for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]);
769:           speed = PetscSqrtReal(speed);
770:           if (dim == 1) {
771:             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &v[p * dim], &speed));
772:           } else {
773:             PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
774:           }
775:         } else if (s[p] == 1) {
776:           PetscCall(PetscDrawSPAddPoint(user->drawspX, &x[p * dim], &v[p * dim]));
777:         }
778:       }
779:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
780:     }
781:     PetscCall(PetscDrawSPDraw(user->drawspX, PETSC_TRUE));
782:     PetscDraw draw;
783:     PetscCall(PetscDrawSPGetDraw(user->drawspX, &draw));
784:     PetscCall(PetscDrawSave(draw));
785:     PetscCall(DMSwarmSortRestoreAccess(sw));
786:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
787:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
788:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
789:     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
790:   }
791:   PetscFunctionReturn(PETSC_SUCCESS);
792: }

794: static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
795: {
796:   AppCtx *user = (AppCtx *)ctx;
797:   DM      dm, sw;

799:   PetscFunctionBeginUser;
800:   if (step > 0 && step % user->ostep == 0) {
801:     PetscCall(TSGetDM(ts, &sw));
802:     PetscCall(DMSwarmGetCellDM(sw, &dm));

804:     if (user->validE) {
805:       PetscScalar *x, *E, *weight;
806:       PetscReal    lower[3], upper[3], xval;
807:       PetscDraw    draw;
808:       PetscInt     dim, cStart, cEnd;

810:       PetscCall(DMGetDimension(dm, &dim));
811:       PetscCall(DMGetBoundingBox(dm, lower, upper));
812:       PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

814:       PetscCall(PetscDrawSPReset(user->drawspE));
815:       PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
816:       PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
817:       PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

819:       PetscCall(DMSwarmSortGetAccess(sw));
820:       for (PetscInt c = 0; c < cEnd - cStart; ++c) {
821:         PetscReal Eavg = 0.0;
822:         PetscInt *pidx, Npc;

824:         PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
825:         for (PetscInt q = 0; q < Npc; ++q) {
826:           const PetscInt p = pidx[q];
827:           Eavg += E[p * dim];
828:         }
829:         Eavg /= Npc;
830:         xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
831:         PetscCall(PetscDrawSPAddPoint(user->drawspE, &xval, &Eavg));
832:         PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
833:       }
834:       PetscCall(PetscDrawSPDraw(user->drawspE, PETSC_TRUE));
835:       PetscCall(PetscDrawSPGetDraw(user->drawspE, &draw));
836:       PetscCall(PetscDrawSave(draw));
837:       PetscCall(DMSwarmSortRestoreAccess(sw));
838:       PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
839:       PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
840:       PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
841:     }

843:     Vec rho, rhohat, phi;

845:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
846:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat));
847:     PetscCall(VecView(rho, user->viewerRho));
848:     PetscCall(VecISCopy(user->fftX, user->fftReal, SCATTER_FORWARD, rho));
849:     PetscCall(MatMult(user->fftPot, user->fftX, user->fftY));
850:     PetscCall(VecFilter(user->fftY, PETSC_SMALL));
851:     PetscCall(VecViewFromOptions(user->fftX, NULL, "-real_view"));
852:     PetscCall(VecViewFromOptions(user->fftY, NULL, "-fft_view"));
853:     PetscCall(VecISCopy(user->fftY, user->fftReal, SCATTER_REVERSE, rhohat));
854:     PetscCall(VecSetValue(rhohat, 0, 0., INSERT_VALUES)); // Remove large DC component
855:     PetscCall(VecView(rhohat, user->viewerRhoHat));
856:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
857:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat));

859:     PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
860:     PetscCall(VecView(phi, user->viewerPhi));
861:     PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
862:   }
863:   PetscFunctionReturn(PETSC_SUCCESS);
864: }

866: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
867: {
868:   PetscBag   bag;
869:   Parameter *p;

871:   PetscFunctionBeginUser;
872:   /* setup PETSc parameter bag */
873:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
874:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
875:   bag = ctx->bag;
876:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
877:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
878:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
879:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
880:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
881:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
882:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
883:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

885:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
886:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
887:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
888:   PetscCall(PetscBagSetFromOptions(bag));
889:   {
890:     PetscViewer       viewer;
891:     PetscViewerFormat format;
892:     PetscBool         flg;

894:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
895:     if (flg) {
896:       PetscCall(PetscViewerPushFormat(viewer, format));
897:       PetscCall(PetscBagView(bag, viewer));
898:       PetscCall(PetscViewerFlush(viewer));
899:       PetscCall(PetscViewerPopFormat(viewer));
900:       PetscCall(PetscViewerDestroy(&viewer));
901:     }
902:   }
903:   PetscFunctionReturn(PETSC_SUCCESS);
904: }

906: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
907: {
908:   PetscFunctionBeginUser;
909:   PetscCall(DMCreate(comm, dm));
910:   PetscCall(DMSetType(*dm, DMPLEX));
911:   PetscCall(DMSetFromOptions(*dm));
912:   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
913:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

915:   // Cache the mesh geometry
916:   DMField         coordField;
917:   IS              cellIS;
918:   PetscQuadrature quad;
919:   PetscReal      *wt, *pt;
920:   PetscInt        cdim, cStart, cEnd;

922:   PetscCall(DMGetCoordinateField(*dm, &coordField));
923:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
924:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
925:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
926:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
927:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
928:   PetscCall(PetscMalloc1(1, &wt));
929:   PetscCall(PetscMalloc1(2, &pt));
930:   wt[0] = 1.;
931:   pt[0] = -1.;
932:   pt[1] = -1.;
933:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
934:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom));
935:   PetscCall(PetscQuadratureDestroy(&quad));
936:   PetscCall(ISDestroy(&cellIS));
937:   PetscFunctionReturn(PETSC_SUCCESS);
938: }

940: 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[])
941: {
942:   f0[0] = -constants[SIGMA];
943: }

945: 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[])
946: {
947:   PetscInt d;
948:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
949: }

951: 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[])
952: {
953:   PetscInt d;
954:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
955: }

957: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
958: {
959:   *u = 0.0;
960:   return PETSC_SUCCESS;
961: }

963: /*
964:    /  I   -grad\ / q \ = /0\
965:    \-div    0  / \phi/   \f/
966: */
967: 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[])
968: {
969:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
970: }

972: 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[])
973: {
974:   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
975: }

977: 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[])
978: {
979:   f0[0] += constants[SIGMA];
980:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
981: }

983: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
984: 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[])
985: {
986:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
987: }

989: 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[])
990: {
991:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
992: }

994: 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[])
995: {
996:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
997: }

999: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
1000: {
1001:   PetscFE   fephi, feq;
1002:   PetscDS   ds;
1003:   PetscBool simplex;
1004:   PetscInt  dim;

1006:   PetscFunctionBeginUser;
1007:   PetscCall(DMGetDimension(dm, &dim));
1008:   PetscCall(DMPlexIsSimplex(dm, &simplex));
1009:   if (user->em == EM_MIXED) {
1010:     DMLabel        label;
1011:     const PetscInt id = 1;

1013:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
1014:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
1015:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
1016:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1017:     PetscCall(PetscFECopyQuadrature(feq, fephi));
1018:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
1019:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
1020:     PetscCall(DMCreateDS(dm));
1021:     PetscCall(PetscFEDestroy(&fephi));
1022:     PetscCall(PetscFEDestroy(&feq));

1024:     PetscCall(DMGetLabel(dm, "marker", &label));
1025:     PetscCall(DMGetDS(dm, &ds));

1027:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
1028:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
1029:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
1030:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
1031:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

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

1035:   } else {
1036:     MatNullSpace nullsp;
1037:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
1038:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
1039:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
1040:     PetscCall(DMCreateDS(dm));
1041:     PetscCall(DMGetDS(dm, &ds));
1042:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
1043:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
1044:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
1045:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
1046:     PetscCall(MatNullSpaceDestroy(&nullsp));
1047:     PetscCall(PetscFEDestroy(&fephi));
1048:   }
1049:   PetscFunctionReturn(PETSC_SUCCESS);
1050: }

1052: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
1053: {
1054:   SNES         snes;
1055:   Mat          J;
1056:   MatNullSpace nullSpace;

1058:   PetscFunctionBeginUser;
1059:   PetscCall(CreateFEM(dm, user));
1060:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
1061:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
1062:   PetscCall(SNESSetDM(snes, dm));
1063:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
1064:   PetscCall(SNESSetFromOptions(snes));

1066:   PetscCall(DMCreateMatrix(dm, &J));
1067:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
1068:   PetscCall(MatSetNullSpace(J, nullSpace));
1069:   PetscCall(MatNullSpaceDestroy(&nullSpace));
1070:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
1071:   PetscCall(MatDestroy(&J));
1072:   if (user->em == EM_MIXED) {
1073:     const PetscInt potential = 1;

1075:     PetscCall(DMCreateSubDM(dm, 1, &potential, &user->isPot, &user->dmPot));
1076:   } else {
1077:     user->dmPot = dm;
1078:     PetscCall(PetscObjectReference((PetscObject)user->dmPot));
1079:   }
1080:   PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M));
1081:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
1082:   user->snes = snes;
1083:   PetscFunctionReturn(PETSC_SUCCESS);
1084: }

1086: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1087: {
1088:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1089:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
1090:   return PETSC_SUCCESS;
1091: }
1092: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
1093: {
1094:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
1095:   return PETSC_SUCCESS;
1096: }

1098: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1099: {
1100:   const PetscReal alpha = scale ? scale[0] : 0.0;
1101:   const PetscReal k     = scale ? scale[1] : 1.;
1102:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
1103:   return PETSC_SUCCESS;
1104: }

1106: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
1107: {
1108:   const PetscReal alpha = scale ? scale[0] : 0.;
1109:   const PetscReal k     = scale ? scale[0] : 1.;
1110:   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
1111:   return PETSC_SUCCESS;
1112: }

1114: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
1115: {
1116:   PetscFE        fe;
1117:   DMPolytopeType ct;
1118:   PetscInt       dim, cStart;
1119:   const char    *prefix = "v";

1121:   PetscFunctionBegin;
1122:   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
1123:   PetscCall(DMSetType(*vdm, DMPLEX));
1124:   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
1125:   PetscCall(DMSetFromOptions(*vdm));
1126:   PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
1127:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

1129:   PetscCall(DMGetDimension(*vdm, &dim));
1130:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
1131:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
1132:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
1133:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
1134:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
1135:   PetscCall(DMCreateDS(*vdm));
1136:   PetscCall(PetscFEDestroy(&fe));
1137:   PetscFunctionReturn(PETSC_SUCCESS);
1138: }

1140: /*
1141:   InitializeParticles_Centroid - Initialize a regular grid of particles.

1143:   Input Parameters:
1144: + sw      - The `DMSWARM`
1145: - force1D - Treat the spatial domain as 1D

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

1150:   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
1151:   and velocity meshes.
1152: */
1153: static PetscErrorCode InitializeParticles_Centroid(DM sw)
1154: {
1155:   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
1156:   DMSwarmCellDM celldm;
1157:   DM            xdm, vdm;
1158:   PetscReal     vmin[3], vmax[3];
1159:   PetscReal    *x, *v;
1160:   PetscInt     *species, *cellid;
1161:   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
1162:   PetscBool     flg;
1163:   MPI_Comm      comm;
1164:   const char   *cellidname;

1166:   PetscFunctionBegin;
1167:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

1169:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
1170:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1171:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
1172:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
1173:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
1174:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
1175:   PetscOptionsEnd();
1176:   debug = swarm->printCoords;

1178:   PetscCall(DMGetDimension(sw, &dim));
1179:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1180:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

1182:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1183:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1184:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

1186:   // One particle per centroid on the tensor product grid
1187:   Npc = (vcEnd - vcStart) * Ns;
1188:   Np  = (xcEnd - xcStart) * Npc;
1189:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
1190:   if (debug) {
1191:     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
1192:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
1193:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
1194:     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
1195:     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
1196:     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
1197:   }

1199:   // Set species and cellid
1200:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
1201:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
1202:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1203:   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
1204:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
1205:     for (PetscInt s = 0; s < Ns; ++s) {
1206:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
1207:         species[p] = s;
1208:         cellid[p]  = c;
1209:       }
1210:     }
1211:   }
1212:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1213:   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));

1215:   // Set particle coordinates
1216:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1217:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1218:   PetscCall(DMSwarmSortGetAccess(sw));
1219:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
1220:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
1221:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
1222:     const PetscInt cell = c + xcStart;
1223:     PetscInt      *pidx, Npc;
1224:     PetscReal      centroid[3], volume;

1226:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1227:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
1228:     for (PetscInt s = 0; s < Ns; ++s) {
1229:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
1230:         const PetscInt p = pidx[q * Ns + s];

1232:         for (PetscInt d = 0; d < dim; ++d) {
1233:           x[p * dim + d] = centroid[d];
1234:           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
1235:         }
1236:         if (debug > 1) {
1237:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
1238:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
1239:           for (PetscInt d = 0; d < dim; ++d) {
1240:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1241:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
1242:           }
1243:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
1244:           for (PetscInt d = 0; d < dim; ++d) {
1245:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
1246:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
1247:           }
1248:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
1249:         }
1250:       }
1251:     }
1252:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1253:   }
1254:   PetscCall(DMSwarmSortRestoreAccess(sw));
1255:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1256:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1257:   PetscFunctionReturn(PETSC_SUCCESS);
1258: }

1260: /*
1261:   InitializeWeights - Compute weight for each local particle

1263:   Input Parameters:
1264: + sw          - The `DMSwarm`
1265: . totalWeight - The sum of all particle weights
1266: . func        - The PDF for the particle spatial distribution
1267: - param       - The PDF parameters

1269:   Notes:
1270:   The PDF for velocity is assumed to be a Gaussian

1272:   The particle weights are returned in the `w_q` field of `sw`.
1273: */
1274: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
1275: {
1276:   DM               xdm, vdm;
1277:   DMSwarmCellDM    celldm;
1278:   PetscScalar     *weight;
1279:   PetscQuadrature  xquad;
1280:   const PetscReal *xq, *xwq;
1281:   const PetscInt   order = 5;
1282:   PetscReal        xi0[3];
1283:   PetscReal        xwtot = 0., pwtot = 0.;
1284:   PetscInt         xNq;
1285:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
1286:   MPI_Comm         comm;
1287:   PetscMPIInt      rank;

1289:   PetscFunctionBegin;
1290:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
1291:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1292:   PetscCall(DMGetDimension(sw, &dim));
1293:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
1294:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1295:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
1296:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
1297:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
1298:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

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

1305:   // Integrate the density function to get the weights of particles in each cell
1306:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
1307:   PetscCall(DMSwarmSortGetAccess(sw));
1308:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1309:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
1310:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
1311:     PetscInt          *pidx, Npc;
1312:     PetscInt           xNc;
1313:     const PetscScalar *xarray;
1314:     PetscScalar       *xcoords = NULL;
1315:     PetscBool          xisDG;

1317:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1318:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1319:     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);
1320:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
1321:     for (PetscInt q = 0; q < xNq; ++q) {
1322:       // Transform quadrature points from ref space to real space
1323:       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
1324:       // Get probability density at quad point
1325:       //   No need to scale xqr since PDF will be periodic
1326:       PetscCall((*func)(xqr, param, &xden));
1327:       xw += xden * (xwq[q] * xdetJ);
1328:     }
1329:     xwtot += xw;
1330:     if (debug) {
1331:       IS              globalOrdering;
1332:       const PetscInt *ordering;

1334:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
1335:       PetscCall(ISGetIndices(globalOrdering, &ordering));
1336:       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));
1337:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
1338:     }
1339:     // Set weights to be Gaussian in velocity cells
1340:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
1341:       const PetscInt     p  = pidx[vc * Ns + 0];
1342:       PetscReal          vw = 0.;
1343:       PetscInt           vNc;
1344:       const PetscScalar *varray;
1345:       PetscScalar       *vcoords = NULL;
1346:       PetscBool          visDG;

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

1354:       weight[p] = totalWeight * vw * xw;
1355:       pwtot += weight[p];
1356:       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 1: %g, %g, %g", p, xw, vw, totalWeight);
1357:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
1358:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
1359:     }
1360:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
1361:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1362:   }
1363:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1364:   PetscCall(DMSwarmSortRestoreAccess(sw));
1365:   PetscCall(PetscQuadratureDestroy(&xquad));

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

1370:     PetscCall(PetscSynchronizedFlush(comm, NULL));
1371:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1372:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
1373:   }
1374:   PetscFunctionReturn(PETSC_SUCCESS);
1375: }

1377: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
1378: {
1379:   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
1380:   PetscInt  dim;

1382:   PetscFunctionBegin;
1383:   PetscCall(DMGetDimension(sw, &dim));
1384:   PetscCall(InitializeParticles_Centroid(sw));
1385:   PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
1386:   PetscFunctionReturn(PETSC_SUCCESS);
1387: }

1389: static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
1390: {
1391:   DM         dm;
1392:   PetscInt  *species;
1393:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
1394:   PetscInt   Np, dim;

1396:   PetscFunctionBegin;
1397:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1398:   PetscCall(DMGetDimension(sw, &dim));
1399:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1400:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
1401:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1402:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1403:   for (PetscInt p = 0; p < Np; ++p) {
1404:     totalWeight += weight[p];
1405:     totalCharge += user->charges[species[p]] * weight[p];
1406:   }
1407:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1408:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1409:   {
1410:     Parameter *param;
1411:     PetscReal  Area;

1413:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1414:     switch (dim) {
1415:     case 1:
1416:       Area = (gmax[0] - gmin[0]);
1417:       break;
1418:     case 2:
1419:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1420:       break;
1421:     case 3:
1422:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1423:       break;
1424:     default:
1425:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1426:     }
1427:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1428:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1429:     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));
1430:     param->sigma = PetscAbsReal(global_charge / (Area));

1432:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
1433:     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,
1434:                           (double)param->vlasovNumber));
1435:   }
1436:   /* Setup Constants */
1437:   {
1438:     PetscDS    ds;
1439:     Parameter *param;
1440:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1441:     PetscScalar constants[NUM_CONSTANTS];
1442:     constants[SIGMA]   = param->sigma;
1443:     constants[V0]      = param->v0;
1444:     constants[T0]      = param->t0;
1445:     constants[X0]      = param->x0;
1446:     constants[M0]      = param->m0;
1447:     constants[Q0]      = param->q0;
1448:     constants[PHI0]    = param->phi0;
1449:     constants[POISSON] = param->poissonNumber;
1450:     constants[VLASOV]  = param->vlasovNumber;
1451:     PetscCall(DMGetDS(dm, &ds));
1452:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
1453:   }
1454:   PetscFunctionReturn(PETSC_SUCCESS);
1455: }

1457: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1458: {
1459:   DMSwarmCellDM celldm;
1460:   DM            vdm;
1461:   PetscReal     v0[2] = {1., 0.};
1462:   PetscInt      dim;

1464:   PetscFunctionBeginUser;
1465:   PetscCall(DMGetDimension(dm, &dim));
1466:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
1467:   PetscCall(DMSetType(*sw, DMSWARM));
1468:   PetscCall(DMSetDimension(*sw, dim));
1469:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
1470:   PetscCall(DMSetApplicationContext(*sw, user));

1472:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
1473:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
1474:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
1475:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));

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

1479:   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
1480:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1481:   PetscCall(DMSwarmCellDMDestroy(&celldm));

1483:   const char *vfieldnames[1] = {"w_q"};

1485:   PetscCall(CreateVelocityDM(*sw, &vdm));
1486:   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
1487:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1488:   PetscCall(DMSwarmCellDMDestroy(&celldm));
1489:   PetscCall(DMDestroy(&vdm));

1491:   DM mdm;

1493:   PetscCall(DMClone(dm, &mdm));
1494:   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
1495:   PetscCall(DMCopyDisc(dm, mdm));
1496:   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
1497:   PetscCall(DMDestroy(&mdm));
1498:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
1499:   PetscCall(DMSwarmCellDMDestroy(&celldm));

1501:   PetscCall(DMSetFromOptions(*sw));
1502:   PetscCall(DMSetUp(*sw));

1504:   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
1505:   user->swarm = *sw;
1506:   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
1507:   if (user->perturbed_weights) {
1508:     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
1509:   } else {
1510:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
1511:     PetscCall(DMSwarmInitializeCoordinates(*sw));
1512:     PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
1513:   }
1514:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
1515:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
1516:   PetscFunctionReturn(PETSC_SUCCESS);
1517: }

1519: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
1520: {
1521:   AppCtx     *user;
1522:   PetscReal  *coords;
1523:   PetscInt   *species, dim, Np, Ns;
1524:   PetscMPIInt size;

1526:   PetscFunctionBegin;
1527:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
1528:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
1529:   PetscCall(DMGetDimension(sw, &dim));
1530:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1531:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1532:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1534:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1535:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1536:   for (PetscInt p = 0; p < Np; ++p) {
1537:     PetscReal *pcoord = &coords[p * dim];
1538:     PetscReal  pE[3]  = {0., 0., 0.};

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

1545:       if (p == q) continue;
1546:       q_q = user->charges[species[q]] * 1.;
1547:       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1548:       r = DMPlex_NormD_Internal(dim, rpq);
1549:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
1550:       r3 = PetscPowRealInt(r, 3);
1551:       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
1552:     }
1553:     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
1554:   }
1555:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1556:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1557:   PetscFunctionReturn(PETSC_SUCCESS);
1558: }

1560: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
1561: {
1562:   DM         dm;
1563:   AppCtx    *user;
1564:   PetscDS    ds;
1565:   PetscFE    fe;
1566:   KSP        ksp;
1567:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
1568:   Vec        rho;         // Charge density, M^{-1} rhoRhs
1569:   Vec        phi, locPhi; // Potential
1570:   Vec        f;           // Particle weights
1571:   PetscReal *coords;
1572:   PetscInt   dim, cStart, cEnd, Np;

1574:   PetscFunctionBegin;
1575:   PetscCall(DMGetApplicationContext(sw, (void *)&user));
1576:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1577:   PetscCall(DMGetDimension(sw, &dim));
1578:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

1580:   PetscCall(SNESGetDM(snes, &dm));
1581:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
1582:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
1583:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
1584:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1585:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

1587:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1588:   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1589:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

1591:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1592:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1594:   // Low-pass filter rhoRhs
1595:   PetscInt window = 0;
1596:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-rho_average", &window, NULL));
1597:   if (window) {
1598:     PetscScalar *a;
1599:     PetscInt     n;
1600:     PetscReal    width = 2. * window + 1.;

1602:     // This will only work for P_1
1603:     //   This works because integration against a test function is linear
1604:     //   This only works in serial since I need the periodic values (maybe use FFT)
1605:     //   Do a real integral against weight function for higher order
1606:     PetscCall(VecGetLocalSize(rhoRhs, &n));
1607:     PetscCall(VecGetArrayWrite(rhoRhs, &a));
1608:     for (PetscInt i = 0; i < n; ++i) {
1609:       PetscScalar avg = a[i];
1610:       for (PetscInt j = 1; j <= window; ++j) avg += a[(i - j + n) % n] + a[(i + j) % n];
1611:       a[i] = avg / width;
1612:       //a[i] = (a[(i - 1 + n) % n] + a[i] + a[(i + 1) % n]) / 3.;
1613:     }
1614:     PetscCall(VecRestoreArrayWrite(rhoRhs, &a));
1615:   }

1617:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1618:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1619:   PetscCall(KSPSetOperators(ksp, user->M, user->M));
1620:   PetscCall(KSPSetFromOptions(ksp));
1621:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

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

1625:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1626:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
1627:   PetscCall(KSPDestroy(&ksp));

1629:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1630:   PetscCall(VecSet(phi, 0.0));
1631:   PetscCall(SNESSolve(snes, rhoRhs, phi));
1632:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
1633:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1635:   PetscCall(DMGetLocalVector(dm, &locPhi));
1636:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1637:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1638:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
1639:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

1641:   PetscCall(DMGetDS(dm, &ds));
1642:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1643:   PetscCall(DMSwarmSortGetAccess(sw));
1644:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1645:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1647:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1648:   PetscTabulation tab;
1649:   PetscReal      *pcoord, *refcoord;
1650:   PetscFEGeom    *chunkgeom = NULL;
1651:   PetscInt        maxNcp    = 0;

1653:   for (PetscInt c = cStart; c < cEnd; ++c) {
1654:     PetscInt Ncp;

1656:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
1657:     maxNcp = PetscMax(maxNcp, Ncp);
1658:   }
1659:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1660:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1661:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
1662:   for (PetscInt c = cStart; c < cEnd; ++c) {
1663:     PetscScalar *clPhi = NULL;
1664:     PetscInt    *points;
1665:     PetscInt     Ncp;

1667:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1668:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1669:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1670:     {
1671:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1672:       for (PetscInt i = 0; i < Ncp; ++i) {
1673:         const PetscReal x0[3] = {-1., -1., -1.};
1674:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1675:       }
1676:     }
1677:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
1678:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1679:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1680:       const PetscReal *basisDer = tab->T[1];
1681:       const PetscInt   p        = points[cp];

1683:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1684:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
1685:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
1686:     }
1687:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1688:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1689:   }
1690:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1691:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1692:   PetscCall(PetscTabulationDestroy(&tab));
1693:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1694:   PetscCall(DMSwarmSortRestoreAccess(sw));
1695:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1696:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1697:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1698:   PetscFunctionReturn(PETSC_SUCCESS);
1699: }

1701: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, Mat M_p, PetscReal E[])
1702: {
1703:   DM         dm;
1704:   AppCtx    *user;
1705:   PetscDS    ds;
1706:   PetscFE    fe;
1707:   KSP        ksp;
1708:   Vec        rhoRhs, rhoRhsFull;   // Weak charge density, \int phi_i rho, and embedding in mixed problem
1709:   Vec        rho;                  // Charge density, M^{-1} rhoRhs
1710:   Vec        phi, locPhi, phiFull; // Potential and embedding in mixed problem
1711:   Vec        f;                    // Particle weights
1712:   PetscReal *coords;
1713:   PetscInt   dim, cStart, cEnd, Np;

1715:   PetscFunctionBegin;
1716:   PetscCall(DMGetApplicationContext(sw, &user));
1717:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
1718:   PetscCall(DMGetDimension(sw, &dim));
1719:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1720:   PetscCall(SNESGetDM(snes, &dm));
1721:   PetscCall(DMGetGlobalVector(user->dmPot, &rhoRhs));
1722:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
1723:   PetscCall(DMGetGlobalVector(dm, &rhoRhsFull));
1724:   PetscCall(PetscObjectSetName((PetscObject)rhoRhsFull, "Weak charge density"));
1725:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
1726:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1727:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

1729:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1730:   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1731:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

1733:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1734:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1736:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1737:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1738:   PetscCall(KSPSetOperators(ksp, user->M, user->M));
1739:   PetscCall(KSPSetFromOptions(ksp));
1740:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

1742:   PetscCall(VecISCopy(rhoRhsFull, user->isPot, SCATTER_FORWARD, rhoRhs));
1743:   //PetscCall(VecScale(rhoRhsFull, -1.0));

1745:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1746:   PetscCall(VecViewFromOptions(rhoRhsFull, NULL, "-rho_full_view"));
1747:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
1748:   PetscCall(DMRestoreGlobalVector(user->dmPot, &rhoRhs));
1749:   PetscCall(KSPDestroy(&ksp));

1751:   PetscCall(DMGetGlobalVector(dm, &phiFull));
1752:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1753:   PetscCall(VecSet(phiFull, 0.0));
1754:   PetscCall(SNESSolve(snes, rhoRhsFull, phiFull));
1755:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhsFull));
1756:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1758:   PetscCall(VecISCopy(phiFull, user->isPot, SCATTER_REVERSE, phi));
1759:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

1761:   PetscCall(DMGetLocalVector(dm, &locPhi));
1762:   PetscCall(DMGlobalToLocalBegin(dm, phiFull, INSERT_VALUES, locPhi));
1763:   PetscCall(DMGlobalToLocalEnd(dm, phiFull, INSERT_VALUES, locPhi));
1764:   PetscCall(DMRestoreGlobalVector(dm, &phiFull));
1765:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

1767:   PetscCall(DMGetDS(dm, &ds));
1768:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1769:   PetscCall(DMSwarmSortGetAccess(sw));
1770:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1771:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1773:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1774:   PetscTabulation tab;
1775:   PetscReal      *pcoord, *refcoord;
1776:   PetscFEGeom    *chunkgeom = NULL;
1777:   PetscInt        maxNcp    = 0;

1779:   for (PetscInt c = cStart; c < cEnd; ++c) {
1780:     PetscInt Ncp;

1782:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
1783:     maxNcp = PetscMax(maxNcp, Ncp);
1784:   }
1785:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1786:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1787:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
1788:   for (PetscInt c = cStart; c < cEnd; ++c) {
1789:     PetscScalar *clPhi = NULL;
1790:     PetscInt    *points;
1791:     PetscInt     Ncp;

1793:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1794:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1795:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1796:     {
1797:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1798:       for (PetscInt i = 0; i < Ncp; ++i) {
1799:         const PetscReal x0[3] = {-1., -1., -1.};
1800:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1801:       }
1802:     }
1803:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
1804:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1805:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1806:       const PetscInt p = points[cp];

1808:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1809:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
1810:       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
1811:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
1812:     }
1813:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1814:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1815:   }
1816:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1817:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1818:   PetscCall(PetscTabulationDestroy(&tab));
1819:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1820:   PetscCall(DMSwarmSortRestoreAccess(sw));
1821:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1822:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1823:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1824:   PetscFunctionReturn(PETSC_SUCCESS);
1825: }

1827: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
1828: {
1829:   AppCtx    *user;
1830:   Mat        M_p;
1831:   PetscReal *E;
1832:   PetscInt   dim, Np;

1834:   PetscFunctionBegin;
1837:   PetscCall(DMGetDimension(sw, &dim));
1838:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1839:   PetscCall(DMGetApplicationContext(sw, &user));

1841:   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
1842:   // TODO: Could share sort context with space cellDM
1843:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
1844:   PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p));
1845:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));

1847:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1848:   PetscCall(PetscArrayzero(E, Np * dim));
1849:   user->validE = PETSC_TRUE;

1851:   switch (user->em) {
1852:   case EM_COULOMB:
1853:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1854:     break;
1855:   case EM_PRIMAL:
1856:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
1857:     break;
1858:   case EM_MIXED:
1859:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, M_p, E));
1860:     break;
1861:   case EM_NONE:
1862:     break;
1863:   default:
1864:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[user->em]);
1865:   }
1866:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1867:   PetscCall(MatDestroy(&M_p));
1868:   PetscFunctionReturn(PETSC_SUCCESS);
1869: }

1871: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1872: {
1873:   DM                 sw;
1874:   SNES               snes = ((AppCtx *)ctx)->snes;
1875:   const PetscScalar *u;
1876:   PetscScalar       *g;
1877:   PetscReal         *E, m_p = 1., q_p = -1.;
1878:   PetscInt           dim, d, Np, p;

1880:   PetscFunctionBeginUser;
1881:   PetscCall(TSGetDM(ts, &sw));
1882:   PetscCall(ComputeFieldAtParticles(snes, sw));

1884:   PetscCall(DMGetDimension(sw, &dim));
1885:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1886:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1887:   PetscCall(VecGetArrayRead(U, &u));
1888:   PetscCall(VecGetArray(G, &g));
1889:   Np /= 2 * dim;
1890:   for (p = 0; p < Np; ++p) {
1891:     for (d = 0; d < dim; ++d) {
1892:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1893:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1894:     }
1895:   }
1896:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1897:   PetscCall(VecRestoreArrayRead(U, &u));
1898:   PetscCall(VecRestoreArray(G, &g));
1899:   PetscFunctionReturn(PETSC_SUCCESS);
1900: }

1902: /* J_{ij} = dF_i/dx_j
1903:    J_p = (  0   1)
1904:          (-w^2  0)
1905:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1906:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1907: */
1908: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1909: {
1910:   DM               sw;
1911:   const PetscReal *coords, *vel;
1912:   PetscInt         dim, d, Np, p, rStart;

1914:   PetscFunctionBeginUser;
1915:   PetscCall(TSGetDM(ts, &sw));
1916:   PetscCall(DMGetDimension(sw, &dim));
1917:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1918:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1919:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1920:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1921:   Np /= 2 * dim;
1922:   for (p = 0; p < Np; ++p) {
1923:     // TODO This is not right because dv/dx has the electric field in it
1924:     PetscScalar vals[4] = {0., 1., -1, 0.};

1926:     for (d = 0; d < dim; ++d) {
1927:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1928:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1929:     }
1930:   }
1931:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1932:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1933:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1934:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1935:   PetscFunctionReturn(PETSC_SUCCESS);
1936: }

1938: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1939: {
1940:   AppCtx            *user = (AppCtx *)ctx;
1941:   DM                 sw;
1942:   const PetscScalar *v;
1943:   PetscScalar       *xres;
1944:   PetscInt           Np, p, d, dim;

1946:   PetscFunctionBeginUser;
1947:   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
1948:   PetscCall(TSGetDM(ts, &sw));
1949:   PetscCall(DMGetDimension(sw, &dim));
1950:   PetscCall(VecGetLocalSize(Xres, &Np));
1951:   PetscCall(VecGetArrayRead(V, &v));
1952:   PetscCall(VecGetArray(Xres, &xres));
1953:   Np /= dim;
1954:   for (p = 0; p < Np; ++p) {
1955:     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
1956:   }
1957:   PetscCall(VecRestoreArrayRead(V, &v));
1958:   PetscCall(VecRestoreArray(Xres, &xres));
1959:   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
1960:   PetscFunctionReturn(PETSC_SUCCESS);
1961: }

1963: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1964: {
1965:   DM                 sw;
1966:   AppCtx            *user = (AppCtx *)ctx;
1967:   SNES               snes = ((AppCtx *)ctx)->snes;
1968:   const PetscScalar *x;
1969:   PetscScalar       *vres;
1970:   PetscReal         *E, m_p, q_p;
1971:   PetscInt           Np, p, dim, d;
1972:   Parameter         *param;

1974:   PetscFunctionBeginUser;
1975:   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
1976:   PetscCall(TSGetDM(ts, &sw));
1977:   PetscCall(ComputeFieldAtParticles(snes, sw));

1979:   PetscCall(DMGetDimension(sw, &dim));
1980:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1981:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1982:   m_p = user->masses[0] * param->m0;
1983:   q_p = user->charges[0] * param->q0;
1984:   PetscCall(VecGetLocalSize(Vres, &Np));
1985:   PetscCall(VecGetArrayRead(X, &x));
1986:   PetscCall(VecGetArray(Vres, &vres));
1987:   Np /= dim;
1988:   for (p = 0; p < Np; ++p) {
1989:     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1990:   }
1991:   PetscCall(VecRestoreArrayRead(X, &x));
1992:   /*
1993:     Synchronized, ordered output for parallel/sequential test cases.
1994:     In the 1D (on the 2D mesh) case, every y component should be zero.
1995:   */
1996:   if (user->checkVRes) {
1997:     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1998:     PetscInt  step;

2000:     PetscCall(TSGetStepNumber(ts, &step));
2001:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
2002:     for (PetscInt p = 0; p < Np; ++p) {
2003:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
2004:       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]));
2005:     }
2006:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
2007:   }
2008:   PetscCall(VecRestoreArray(Vres, &vres));
2009:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2010:   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
2011:   PetscFunctionReturn(PETSC_SUCCESS);
2012: }

2014: /* Discrete Gradients Formulation: S, F, gradF (G) */
2015: PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
2016: {
2017:   PetscScalar vals[4] = {0., 1., -1., 0.};
2018:   DM          sw;
2019:   PetscInt    dim, d, Np, p, rStart;

2021:   PetscFunctionBeginUser;
2022:   PetscCall(TSGetDM(ts, &sw));
2023:   PetscCall(DMGetDimension(sw, &dim));
2024:   PetscCall(VecGetLocalSize(U, &Np));
2025:   PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
2026:   Np /= 2 * dim;
2027:   for (p = 0; p < Np; ++p) {
2028:     for (d = 0; d < dim; ++d) {
2029:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
2030:       PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
2031:     }
2032:   }
2033:   PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
2034:   PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
2035:   PetscFunctionReturn(PETSC_SUCCESS);
2036: }

2038: PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
2039: {
2040:   AppCtx            *user = (AppCtx *)ctx;
2041:   DM                 sw;
2042:   Vec                phi;
2043:   const PetscScalar *u;
2044:   PetscInt           dim, Np, cStart, cEnd;
2045:   PetscReal         *vel, *coords, m_p = 1.;

2047:   PetscFunctionBeginUser;
2048:   PetscCall(TSGetDM(ts, &sw));
2049:   PetscCall(DMGetDimension(sw, &dim));
2050:   PetscCall(DMPlexGetHeightStratum(user->dmPot, 0, &cStart, &cEnd));

2052:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
2053:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
2054:   PetscCall(computeFieldEnergy(user->dmPot, phi, F));
2055:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

2057:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2058:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2059:   PetscCall(DMSwarmSortGetAccess(sw));
2060:   PetscCall(VecGetArrayRead(U, &u));
2061:   PetscCall(VecGetLocalSize(U, &Np));
2062:   Np /= 2 * dim;
2063:   for (PetscInt c = cStart; c < cEnd; ++c) {
2064:     PetscInt *points;
2065:     PetscInt  Ncp;

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

2072:       *F += 0.5 * m_p * v2;
2073:     }
2074:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
2075:   }
2076:   PetscCall(VecRestoreArrayRead(U, &u));
2077:   PetscCall(DMSwarmSortRestoreAccess(sw));
2078:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2079:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2080:   PetscFunctionReturn(PETSC_SUCCESS);
2081: }

2083: /* dF/dx = q E   dF/dv = v */
2084: PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
2085: {
2086:   DM                 sw;
2087:   SNES               snes = ((AppCtx *)ctx)->snes;
2088:   const PetscReal   *coords, *vel, *E;
2089:   const PetscScalar *u;
2090:   PetscScalar       *g;
2091:   PetscReal          m_p = 1., q_p = -1.;
2092:   PetscInt           dim, d, Np, p;

2094:   PetscFunctionBeginUser;
2095:   PetscCall(TSGetDM(ts, &sw));
2096:   PetscCall(DMGetDimension(sw, &dim));
2097:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2098:   PetscCall(VecGetArrayRead(U, &u));
2099:   PetscCall(VecGetArray(G, &g));

2101:   PetscLogEvent COMPUTEFIELD;
2102:   PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
2103:   PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
2104:   PetscCall(ComputeFieldAtParticles(snes, sw));
2105:   PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
2106:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2107:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
2108:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
2109:   for (p = 0; p < Np; ++p) {
2110:     for (d = 0; d < dim; ++d) {
2111:       g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
2112:       g[(p * 2 + 1) * dim + d] = u[(p * 2 + 1) * dim + d];
2113:     }
2114:   }
2115:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
2116:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
2117:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
2118:   PetscCall(VecRestoreArrayRead(U, &u));
2119:   PetscCall(VecRestoreArray(G, &g));
2120:   PetscFunctionReturn(PETSC_SUCCESS);
2121: }

2123: static PetscErrorCode CreateSolution(TS ts)
2124: {
2125:   DM       sw;
2126:   Vec      u;
2127:   PetscInt dim, Np;

2129:   PetscFunctionBegin;
2130:   PetscCall(TSGetDM(ts, &sw));
2131:   PetscCall(DMGetDimension(sw, &dim));
2132:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2133:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
2134:   PetscCall(VecSetBlockSize(u, dim));
2135:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
2136:   PetscCall(VecSetUp(u));
2137:   PetscCall(TSSetSolution(ts, u));
2138:   PetscCall(VecDestroy(&u));
2139:   PetscFunctionReturn(PETSC_SUCCESS);
2140: }

2142: static PetscErrorCode SetProblem(TS ts)
2143: {
2144:   AppCtx *user;
2145:   DM      sw;

2147:   PetscFunctionBegin;
2148:   PetscCall(TSGetDM(ts, &sw));
2149:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
2150:   // Define unified system for (X, V)
2151:   {
2152:     Mat      J;
2153:     PetscInt dim, Np;

2155:     PetscCall(DMGetDimension(sw, &dim));
2156:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2157:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
2158:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
2159:     PetscCall(MatSetBlockSize(J, 2 * dim));
2160:     PetscCall(MatSetFromOptions(J));
2161:     PetscCall(MatSetUp(J));
2162:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
2163:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
2164:     PetscCall(MatDestroy(&J));
2165:   }
2166:   /* Define split system for X and V */
2167:   {
2168:     Vec             u;
2169:     IS              isx, isv, istmp;
2170:     const PetscInt *idx;
2171:     PetscInt        dim, Np, rstart;

2173:     PetscCall(TSGetSolution(ts, &u));
2174:     PetscCall(DMGetDimension(sw, &dim));
2175:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
2176:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
2177:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
2178:     PetscCall(ISGetIndices(istmp, &idx));
2179:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
2180:     PetscCall(ISRestoreIndices(istmp, &idx));
2181:     PetscCall(ISDestroy(&istmp));
2182:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
2183:     PetscCall(ISGetIndices(istmp, &idx));
2184:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
2185:     PetscCall(ISRestoreIndices(istmp, &idx));
2186:     PetscCall(ISDestroy(&istmp));
2187:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
2188:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
2189:     PetscCall(ISDestroy(&isx));
2190:     PetscCall(ISDestroy(&isv));
2191:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
2192:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
2193:   }
2194:   // Define symplectic formulation U_t = S . G, where G = grad F
2195:   {
2196:     PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, user));
2197:   }
2198:   PetscFunctionReturn(PETSC_SUCCESS);
2199: }

2201: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
2202: {
2203:   DM        sw;
2204:   Vec       u;
2205:   PetscReal t, maxt, dt;
2206:   PetscInt  n, maxn;

2208:   PetscFunctionBegin;
2209:   PetscCall(TSGetDM(ts, &sw));
2210:   PetscCall(TSGetTime(ts, &t));
2211:   PetscCall(TSGetMaxTime(ts, &maxt));
2212:   PetscCall(TSGetTimeStep(ts, &dt));
2213:   PetscCall(TSGetStepNumber(ts, &n));
2214:   PetscCall(TSGetMaxSteps(ts, &maxn));

2216:   PetscCall(TSReset(ts));
2217:   PetscCall(TSSetDM(ts, sw));
2218:   PetscCall(TSSetFromOptions(ts));
2219:   PetscCall(TSSetTime(ts, t));
2220:   PetscCall(TSSetMaxTime(ts, maxt));
2221:   PetscCall(TSSetTimeStep(ts, dt));
2222:   PetscCall(TSSetStepNumber(ts, n));
2223:   PetscCall(TSSetMaxSteps(ts, maxn));

2225:   PetscCall(CreateSolution(ts));
2226:   PetscCall(SetProblem(ts));
2227:   PetscCall(TSGetSolution(ts, &u));
2228:   PetscFunctionReturn(PETSC_SUCCESS);
2229: }

2231: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
2232: {
2233:   DM        sw, cdm;
2234:   PetscInt  Np;
2235:   PetscReal low[2], high[2];
2236:   AppCtx   *user = (AppCtx *)ctx;

2238:   sw = user->swarm;
2239:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
2240:   // Get the bounding box so we can equally space the particles
2241:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
2242:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
2243:   // shift it by h/2 so nothing is initialized directly on a boundary
2244:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
2245:   x[1] = 0.;
2246:   return PETSC_SUCCESS;
2247: }

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

2252:   Input Parameters:
2253: + ts         - The TS
2254: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

2256:   Output Parameters:
2257: . u - The initialized solution vector

2259:   Level: advanced

2261: .seealso: InitializeSolve()
2262: */
2263: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
2264: {
2265:   DM       sw;
2266:   Vec      u, gc, gv;
2267:   IS       isx, isv;
2268:   PetscInt dim;
2269:   AppCtx  *user;

2271:   PetscFunctionBeginUser;
2272:   PetscCall(TSGetDM(ts, &sw));
2273:   PetscCall(DMGetApplicationContext(sw, &user));
2274:   PetscCall(DMGetDimension(sw, &dim));
2275:   if (useInitial) {
2276:     PetscReal v0[2] = {1., 0.};
2277:     if (user->perturbed_weights) {
2278:       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
2279:     } else {
2280:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
2281:       PetscCall(DMSwarmInitializeCoordinates(sw));
2282:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
2283:     }
2284:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2285:     PetscCall(DMSwarmTSRedistribute(ts));
2286:   }
2287:   PetscCall(DMSetUp(sw));
2288:   PetscCall(TSGetSolution(ts, &u));
2289:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2290:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2291:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2292:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2293:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
2294:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
2295:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2296:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2297:   PetscFunctionReturn(PETSC_SUCCESS);
2298: }

2300: static PetscErrorCode InitializeSolve(TS ts, Vec u)
2301: {
2302:   PetscFunctionBegin;
2303:   PetscCall(TSSetSolution(ts, u));
2304:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
2305:   PetscFunctionReturn(PETSC_SUCCESS);
2306: }

2308: static PetscErrorCode MigrateParticles(TS ts)
2309: {
2310:   DM               sw, cdm;
2311:   const PetscReal *L;
2312:   AppCtx          *ctx;

2314:   PetscFunctionBeginUser;
2315:   PetscCall(TSGetDM(ts, &sw));
2316:   PetscCall(DMGetApplicationContext(sw, &ctx));
2317:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
2318:   {
2319:     Vec        u, gc, gv, position, momentum;
2320:     IS         isx, isv;
2321:     PetscReal *pos, *mom;

2323:     PetscCall(TSGetSolution(ts, &u));
2324:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
2325:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
2326:     PetscCall(VecGetSubVector(u, isx, &position));
2327:     PetscCall(VecGetSubVector(u, isv, &momentum));
2328:     PetscCall(VecGetArray(position, &pos));
2329:     PetscCall(VecGetArray(momentum, &mom));
2330:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2331:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
2332:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
2333:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

2335:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
2336:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
2337:     if ((L[0] || L[1]) >= 0.) {
2338:       PetscReal *x, *v, upper[3], lower[3];
2339:       PetscInt   Np, dim;

2341:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
2342:       PetscCall(DMGetDimension(cdm, &dim));
2343:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
2344:       PetscCall(VecGetArray(gc, &x));
2345:       PetscCall(VecGetArray(gv, &v));
2346:       for (PetscInt p = 0; p < Np; ++p) {
2347:         for (PetscInt d = 0; d < dim; ++d) {
2348:           if (pos[p * dim + d] < lower[d]) {
2349:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
2350:           } else if (pos[p * dim + d] > upper[d]) {
2351:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
2352:           } else {
2353:             x[p * dim + d] = pos[p * dim + d];
2354:           }
2355:           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]);
2356:           v[p * dim + d] = mom[p * dim + d];
2357:         }
2358:       }
2359:       PetscCall(VecRestoreArray(gc, &x));
2360:       PetscCall(VecRestoreArray(gv, &v));
2361:     }
2362:     PetscCall(VecRestoreArray(position, &pos));
2363:     PetscCall(VecRestoreArray(momentum, &mom));
2364:     PetscCall(VecRestoreSubVector(u, isx, &position));
2365:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2366:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2367:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2368:   }
2369:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2370:   PetscInt step;

2372:   PetscCall(TSGetStepNumber(ts, &step));
2373:   if (!(step % ctx->remapFreq)) {
2374:     // Monitor electric field before we destroy it
2375:     PetscReal ptime;
2376:     PetscInt  step;

2378:     PetscCall(TSGetStepNumber(ts, &step));
2379:     PetscCall(TSGetTime(ts, &ptime));
2380:     if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx));
2381:     if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx));
2382:     PetscCall(DMSwarmRemap(sw));
2383:     ctx->validE = PETSC_FALSE;
2384:   }
2385:   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
2386:   PetscCall(DMSwarmTSRedistribute(ts));
2387:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2388:   {
2389:     const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
2390:     PetscCall(DMSwarmVectorDefineFields(sw, 2, fieldnames));
2391:   }
2392:   PetscFunctionReturn(PETSC_SUCCESS);
2393: }

2395: int main(int argc, char **argv)
2396: {
2397:   DM        dm, sw;
2398:   TS        ts;
2399:   Vec       u;
2400:   PetscReal dt;
2401:   PetscInt  maxn;
2402:   AppCtx    user;

2404:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2405:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2406:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2407:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2408:   PetscCall(CreatePoisson(dm, &user));
2409:   PetscCall(CreateSwarm(dm, &user, &sw));
2410:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2411:   PetscCall(InitializeConstants(sw, &user));
2412:   PetscCall(DMSetApplicationContext(sw, &user));

2414:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2415:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2416:   PetscCall(TSSetDM(ts, sw));
2417:   PetscCall(TSSetMaxTime(ts, 0.1));
2418:   PetscCall(TSSetTimeStep(ts, 0.00001));
2419:   PetscCall(TSSetMaxSteps(ts, 100));
2420:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

2422:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2423:   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
2424:   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2425:   if (user.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2426:   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));
2427:   if (user.velocity_monitor >= 0) PetscCall(TSMonitorSet(ts, MonitorVelocity, &user, NULL));

2429:   PetscCall(TSSetFromOptions(ts));
2430:   PetscCall(TSGetTimeStep(ts, &dt));
2431:   PetscCall(TSGetMaxSteps(ts, &maxn));
2432:   user.steps    = maxn;
2433:   user.stepSize = dt;
2434:   PetscCall(SetupContext(dm, sw, &user));
2435:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2436:   PetscCall(TSSetPostStep(ts, MigrateParticles));
2437:   PetscCall(CreateSolution(ts));
2438:   PetscCall(TSGetSolution(ts, &u));
2439:   PetscCall(TSComputeInitialCondition(ts, u));
2440:   PetscCall(CheckNonNegativeWeights(sw, &user));
2441:   PetscCall(TSSolve(ts, NULL));

2443:   PetscCall(SNESDestroy(&user.snes));
2444:   PetscCall(DMDestroy(&user.dmPot));
2445:   PetscCall(ISDestroy(&user.isPot));
2446:   PetscCall(MatDestroy(&user.M));
2447:   PetscCall(PetscFEGeomDestroy(&user.fegeom));
2448:   PetscCall(TSDestroy(&ts));
2449:   PetscCall(DMDestroy(&sw));
2450:   PetscCall(DMDestroy(&dm));
2451:   PetscCall(DestroyContext(&user));
2452:   PetscCall(PetscFinalize());
2453:   return 0;
2454: }

2456: /*TEST

2458:    build:
2459:     requires: !complex double

2461:   # This tests that we can put particles in a box and compute the Coulomb force
2462:   # Recommend -draw_size 500,500
2463:    testset:
2464:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2465:      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 20 \
2466:              -dm_plex_box_lower 0 -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2467:            -vdm_plex_simplex 0 \
2468:            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2469:            -sigma 1.0e-8 -timeScale 2.0e-14 \
2470:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2471:            -output_step 50 -check_vel_res
2472:      test:
2473:        suffix: none_1d
2474:        requires:
2475:        args: -em_type none -error
2476:      test:
2477:        suffix: coulomb_1d
2478:        args: -em_type coulomb

2480:    # for viewers
2481:    #-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
2482:    testset:
2483:      nsize: {{1 2}}
2484:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2485:      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 36 \
2486:              -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2487:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2488:              -vdm_plex_box_lower -3 -vdm_plex_box_upper 3 \
2489:            -dm_swarm_num_species 1 -twostream -charges -1.,1. -sigma 1.0e-8 \
2490:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2491:            -ts_type basicsymplectic -ts_basicsymplectic_type 2 \
2492:              -ts_dt 0.01 -ts_max_time 5 -ts_max_steps 10 \
2493:            -em_snes_atol 1.e-15 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2494:            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1
2495:      test:
2496:        suffix: two_stream_c0
2497:        args: -em_type primal -petscfe_default_quadrature_order 2 -petscspace_degree 2 -em_pc_type svd
2498:      test:
2499:        suffix: two_stream_rt
2500:        requires: superlu_dist
2501:        args: -em_type mixed \
2502:                -potential_petscspace_degree 0 \
2503:                -potential_petscdualspace_lagrange_use_moments \
2504:                -potential_petscdualspace_lagrange_moment_order 2 \
2505:                -field_petscspace_degree 1 -field_petscfe_default_quadrature_order 1 \
2506:              -em_snes_error_if_not_converged \
2507:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2508:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2509:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2510:                -em_fieldsplit_field_pc_type lu \
2511:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2512:                -em_fieldsplit_potential_pc_type svd

2514:    # For an eyeball check, we use
2515:    # -ts_max_steps 1000 -dm_plex_box_faces 10,1 -vdm_plex_box_faces 2000 -efield_monitor
2516:    # For verification, we use
2517:    # -ts_max_steps 1000 -dm_plex_box_faces 100,1 -vdm_plex_box_faces 8000 -dm_swarm_num_particles 800000 -monitor_efield
2518:    # -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
2519:    testset:
2520:      nsize: {{1 2}}
2521:      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 10 \
2522:              -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2523:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2524:              -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 \
2525:              -vpetscspace_degree 2 -vdm_plex_hash_location \
2526:            -dm_swarm_num_species 1 -charges -1.,1. \
2527:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2528:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2529:              -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 1 \
2530:            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2531:            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1

2533:      test:
2534:        suffix: uniform_equilibrium_1d
2535:        args: -cosine_coefficients 0.0,0.5 -em_type primal -petscspace_degree 1 -em_pc_type svd
2536:      test:
2537:        suffix: uniform_equilibrium_1d_real
2538:        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2539:                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2540:              -cosine_coefficients 0.0 -em_type primal -petscspace_degree 1 -em_pc_type svd
2541:      test:
2542:        suffix: landau_damping_1d_c0
2543:        args: -em_type primal -petscspace_degree 1 -em_pc_type svd
2544:      test:
2545:        suffix: uniform_primal_1d_real
2546:        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2547:                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2548:              -cosine_coefficients 0.01 -em_type primal -petscspace_degree 1 -em_pc_type svd
2549:      # NOT WORKING -ftop_pc_type bjacobi -ftop_sub_pc_type lu -ftop_sub_pc_factor_shift_type nonzero
2550:      test:
2551:        suffix: uniform_primal_1d_real_pfak
2552:        nsize: 1
2553:        args: -dm_plex_dim 1 -dm_plex_simplex 1 -dm_plex_box_faces 10 \
2554:                -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2555:              -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 -remap_dm_plex_box_faces 10,10 -remap_dm_plex_box_bd periodic,none \
2556:                -remap_dm_plex_box_lower 0.,-10. -remap_dm_plex_box_upper 12.5664,10. \
2557:                -remap_petscspace_degree 2 -remap_dm_plex_hash_location \
2558:              -remap_freq 1 -dm_swarm_remap_type pfak \
2559:                -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 \
2560:                -ptof_pc_type lu \
2561:              -cosine_coefficients 0.01 -em_type primal -petscspace_degree 1 -em_pc_type svd -em_proj_pc_type lu
2562:      test:
2563:        requires: superlu_dist
2564:        suffix: landau_damping_1d_mixed
2565:        args: -em_type mixed \
2566:                -potential_petscspace_degree 0 \
2567:                -potential_petscdualspace_lagrange_use_moments \
2568:                -potential_petscdualspace_lagrange_moment_order 2 \
2569:                -field_petscspace_degree 1 \
2570:              -em_snes_error_if_not_converged \
2571:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2572:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2573:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2574:                -em_fieldsplit_field_pc_type lu \
2575:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2576:                -em_fieldsplit_potential_pc_type svd

2578:    # Same as above, with different timestepping
2579:    testset:
2580:      nsize: {{1 2}}
2581:      args: -dm_plex_dim 1 -dm_plex_simplex 0 -dm_plex_box_faces 10 \
2582:              -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 -dm_plex_box_bd periodic \
2583:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2584:              -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 \
2585:              -vpetscspace_degree 2 -vdm_plex_hash_location \
2586:            -dm_swarm_num_species 1 -charges -1.,1. \
2587:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2588:            -ts_type discgrad -ts_discgrad_type average \
2589:              -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 1 \
2590:            -snes_type qn \
2591:            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2592:            -output_step 1 -check_vel_res -dm_swarm_print_coords 1 -dm_swarm_print_weights 1

2594:      test:
2595:        suffix: landau_damping_1d_dg
2596:        args: -em_type primal -petscspace_degree 1 -em_pc_type svd

2598: TEST*/