Actual source code: ex2.c

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

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

  7:   To visualize the efield use

  9:     -monitor_efield

 11:   To visualize the swarm distribution use

 13:     -ts_monitor_hg_swarm

 15:   To visualize the particles, we can use

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

 19: */
 20: #include <petscts.h>
 21: #include <petscdmplex.h>
 22: #include <petscdmswarm.h>
 23: #include <petscfe.h>
 24: #include <petscds.h>
 25: #include <petscbag.h>
 26: #include <petscdraw.h>
 27: #include <petsc/private/dmpleximpl.h>
 28: #include <petsc/private/petscfeimpl.h>
 29: #include "petscdm.h"
 30: #include "petscdmlabel.h"

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

 35: const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
 36: typedef enum {
 37:   EM_PRIMAL,
 38:   EM_MIXED,
 39:   EM_COULOMB,
 40:   EM_NONE
 41: } EMType;

 43: typedef enum {
 44:   V0,
 45:   X0,
 46:   T0,
 47:   M0,
 48:   Q0,
 49:   PHI0,
 50:   POISSON,
 51:   VLASOV,
 52:   SIGMA,
 53:   NUM_CONSTANTS
 54: } ConstantType;
 55: typedef struct {
 56:   PetscScalar v0; /* Velocity scale, often the thermal velocity */
 57:   PetscScalar t0; /* Time scale */
 58:   PetscScalar x0; /* Space scale */
 59:   PetscScalar m0; /* Mass scale */
 60:   PetscScalar q0; /* Charge scale */
 61:   PetscScalar kb;
 62:   PetscScalar epsi0;
 63:   PetscScalar phi0;          /* Potential scale */
 64:   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
 65:   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
 66:   PetscReal   sigma;         /* Nondimensional charge per length in x */
 67: } Parameter;

 69: typedef struct {
 70:   PetscBag    bag;            /* Problem parameters */
 71:   PetscBool   error;          /* Flag for printing the error */
 72:   PetscBool   efield_monitor; /* Flag to show electric field monitor */
 73:   PetscBool   initial_monitor;
 74:   PetscBool   fake_1D;           /* Run simulation in 2D but zeroing second dimension */
 75:   PetscBool   perturbed_weights; /* Uniformly sample x,v space with gaussian weights */
 76:   PetscBool   poisson_monitor;
 77:   PetscInt    ostep; /* print the energy at each ostep time steps */
 78:   PetscInt    numParticles;
 79:   PetscReal   timeScale;              /* Nondimensionalizing time scale */
 80:   PetscReal   charges[2];             /* The charges of each species */
 81:   PetscReal   masses[2];              /* The masses of each species */
 82:   PetscReal   thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
 83:   PetscReal   cosine_coefficients[2]; /*(alpha, k)*/
 84:   PetscReal   totalWeight;
 85:   PetscReal   stepSize;
 86:   PetscInt    steps;
 87:   PetscReal   initVel;
 88:   EMType      em; /* Type of electrostatic model */
 89:   SNES        snes;
 90:   PetscDraw   drawef;
 91:   PetscDrawLG drawlg_ef;
 92:   PetscDraw   drawic_x;
 93:   PetscDraw   drawic_v;
 94:   PetscDraw   drawic_w;
 95:   PetscDrawHG drawhgic_x;
 96:   PetscDrawHG drawhgic_v;
 97:   PetscDrawHG drawhgic_w;
 98:   PetscDraw   EDraw;
 99:   PetscDraw   RhoDraw;
100:   PetscDraw   PotDraw;
101:   PetscDrawSP EDrawSP;
102:   PetscDrawSP RhoDrawSP;
103:   PetscDrawSP PotDrawSP;
104:   PetscBool   monitor_positions; /* Flag to show particle positins at each time step */
105:   PetscDraw   positionDraw;
106:   PetscDrawSP positionDrawSP;
107:   DM          swarm;
108:   PetscRandom random;
109:   PetscBool   twostream;
110:   PetscBool   checkweights;
111:   PetscInt    checkVRes; /* Flag to check/output velocity residuals for nightly tests */
112: } AppCtx;

114: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
115: {
116:   PetscFunctionBeginUser;
117:   PetscInt d                      = 2;
118:   PetscInt maxSpecies             = 2;
119:   options->error                  = PETSC_FALSE;
120:   options->efield_monitor         = PETSC_FALSE;
121:   options->initial_monitor        = PETSC_FALSE;
122:   options->fake_1D                = PETSC_FALSE;
123:   options->perturbed_weights      = PETSC_FALSE;
124:   options->poisson_monitor        = PETSC_FALSE;
125:   options->ostep                  = 100;
126:   options->timeScale              = 2.0e-14;
127:   options->charges[0]             = -1.0;
128:   options->charges[1]             = 1.0;
129:   options->masses[0]              = 1.0;
130:   options->masses[1]              = 1000.0;
131:   options->thermal_energy[0]      = 1.0;
132:   options->thermal_energy[1]      = 1.0;
133:   options->cosine_coefficients[0] = 0.01;
134:   options->cosine_coefficients[1] = 0.5;
135:   options->initVel                = 1;
136:   options->totalWeight            = 1.0;
137:   options->drawef                 = NULL;
138:   options->drawlg_ef              = NULL;
139:   options->drawic_x               = NULL;
140:   options->drawic_v               = NULL;
141:   options->drawic_w               = NULL;
142:   options->drawhgic_x             = NULL;
143:   options->drawhgic_v             = NULL;
144:   options->drawhgic_w             = NULL;
145:   options->EDraw                  = NULL;
146:   options->RhoDraw                = NULL;
147:   options->PotDraw                = NULL;
148:   options->EDrawSP                = NULL;
149:   options->RhoDrawSP              = NULL;
150:   options->PotDrawSP              = NULL;
151:   options->em                     = EM_COULOMB;
152:   options->numParticles           = 32768;
153:   options->monitor_positions      = PETSC_FALSE;
154:   options->positionDraw           = NULL;
155:   options->positionDrawSP         = NULL;
156:   options->twostream              = PETSC_FALSE;
157:   options->checkweights           = PETSC_FALSE;
158:   options->checkVRes              = 0;

160:   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
161:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex2.c", options->error, &options->error, NULL));
162:   PetscCall(PetscOptionsBool("-monitor_efield", "Flag to show efield plot", "ex2.c", options->efield_monitor, &options->efield_monitor, NULL));
163:   PetscCall(PetscOptionsBool("-monitor_ics", "Flag to show initial condition histograms", "ex2.c", options->initial_monitor, &options->initial_monitor, NULL));
164:   PetscCall(PetscOptionsBool("-monitor_positions", "The flag to show particle positions", "ex2.c", options->monitor_positions, &options->monitor_positions, NULL));
165:   PetscCall(PetscOptionsBool("-monitor_poisson", "The flag to show charges, Efield and potential solve", "ex2.c", options->poisson_monitor, &options->poisson_monitor, NULL));
166:   PetscCall(PetscOptionsBool("-fake_1D", "Flag to run a 1D simulation (but really in 2D)", "ex2.c", options->fake_1D, &options->fake_1D, NULL));
167:   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", "ex2.c", options->twostream, &options->twostream, NULL));
168:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", "ex2.c", options->perturbed_weights, &options->perturbed_weights, NULL));
169:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex2.c", options->checkweights, &options->checkweights, NULL));
170:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex2.c", options->checkVRes, &options->checkVRes, NULL));
171:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex2.c", options->ostep, &options->ostep, NULL));
172:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex2.c", options->timeScale, &options->timeScale, NULL));
173:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex2.c", options->initVel, &options->initVel, NULL));
174:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex2.c", options->totalWeight, &options->totalWeight, NULL));
175:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex2.c", options->cosine_coefficients, &d, NULL));
176:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex2.c", options->charges, &maxSpecies, NULL));
177:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex2.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
178:   PetscOptionsEnd();
179:   PetscFunctionReturn(PETSC_SUCCESS);
180: }

182: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
183: {
184:   PetscFunctionBeginUser;
185:   if (user->efield_monitor) {
186:     PetscDrawAxis axis_ef;
187:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_efield", 0, 300, 400, 300, &user->drawef));
188:     PetscCall(PetscDrawSetSave(user->drawef, "ex9_Efield.png"));
189:     PetscCall(PetscDrawSetFromOptions(user->drawef));
190:     PetscCall(PetscDrawLGCreate(user->drawef, 1, &user->drawlg_ef));
191:     PetscCall(PetscDrawLGGetAxis(user->drawlg_ef, &axis_ef));
192:     PetscCall(PetscDrawAxisSetLabels(axis_ef, "Electron Electric Field", "time", "E_max"));
193:     PetscCall(PetscDrawLGSetLimits(user->drawlg_ef, 0., user->steps * user->stepSize, -10., 0.));
194:     PetscCall(PetscDrawAxisSetLimits(axis_ef, 0., user->steps * user->stepSize, -10., 0.));
195:   }
196:   if (user->initial_monitor) {
197:     PetscDrawAxis axis1, axis2, axis3;
198:     PetscReal     dmboxlower[2], dmboxupper[2];
199:     PetscInt      dim, cStart, cEnd;
200:     PetscCall(DMGetDimension(sw, &dim));
201:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
202:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

204:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &user->drawic_x));
205:     PetscCall(PetscDrawSetSave(user->drawic_x, "ex9_ic_x.png"));
206:     PetscCall(PetscDrawSetFromOptions(user->drawic_x));
207:     PetscCall(PetscDrawHGCreate(user->drawic_x, (int)dim, &user->drawhgic_x));
208:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
209:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart)));
210:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "counts"));
211:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 1500));

213:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &user->drawic_v));
214:     PetscCall(PetscDrawSetSave(user->drawic_v, "ex9_ic_v.png"));
215:     PetscCall(PetscDrawSetFromOptions(user->drawic_v));
216:     PetscCall(PetscDrawHGCreate(user->drawic_v, (int)dim, &user->drawhgic_v));
217:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
218:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 1000));
219:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "counts"));
220:     PetscCall(PetscDrawAxisSetLimits(axis2, -1, 1, 0, 1500));

222:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_w", 800, 300, 400, 300, &user->drawic_w));
223:     PetscCall(PetscDrawSetSave(user->drawic_w, "ex9_ic_w.png"));
224:     PetscCall(PetscDrawSetFromOptions(user->drawic_w));
225:     PetscCall(PetscDrawHGCreate(user->drawic_w, (int)dim, &user->drawhgic_w));
226:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_w, &axis3));
227:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_w, 10));
228:     PetscCall(PetscDrawAxisSetLabels(axis3, "Initial W Distribution", "weight", "counts"));
229:     PetscCall(PetscDrawAxisSetLimits(axis3, 0, 0.01, 0, 5000));
230:   }
231:   if (user->monitor_positions) {
232:     PetscDrawAxis axis;

234:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "position_monitor_species1", 0, 0, 400, 300, &user->positionDraw));
235:     PetscCall(PetscDrawSetFromOptions(user->positionDraw));
236:     PetscCall(PetscDrawSPCreate(user->positionDraw, 10, &user->positionDrawSP));
237:     PetscCall(PetscDrawSPSetDimension(user->positionDrawSP, 1));
238:     PetscCall(PetscDrawSPGetAxis(user->positionDrawSP, &axis));
239:     PetscCall(PetscDrawSPReset(user->positionDrawSP));
240:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
241:     PetscCall(PetscDrawSetSave(user->positionDraw, "ex9_pos.png"));
242:   }
243:   if (user->poisson_monitor) {
244:     PetscDrawAxis axis_E, axis_Rho, axis_Pot;

246:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Efield_monitor", 0, 0, 400, 300, &user->EDraw));
247:     PetscCall(PetscDrawSetFromOptions(user->EDraw));
248:     PetscCall(PetscDrawSPCreate(user->EDraw, 10, &user->EDrawSP));
249:     PetscCall(PetscDrawSPSetDimension(user->EDrawSP, 1));
250:     PetscCall(PetscDrawSPGetAxis(user->EDrawSP, &axis_E));
251:     PetscCall(PetscDrawSPReset(user->EDrawSP));
252:     PetscCall(PetscDrawAxisSetLabels(axis_E, "Particles", "x", "E"));
253:     PetscCall(PetscDrawSetSave(user->EDraw, "ex9_E_spatial.png"));

255:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "rho_monitor", 0, 0, 400, 300, &user->RhoDraw));
256:     PetscCall(PetscDrawSetFromOptions(user->RhoDraw));
257:     PetscCall(PetscDrawSPCreate(user->RhoDraw, 10, &user->RhoDrawSP));
258:     PetscCall(PetscDrawSPSetDimension(user->RhoDrawSP, 1));
259:     PetscCall(PetscDrawSPGetAxis(user->RhoDrawSP, &axis_Rho));
260:     PetscCall(PetscDrawSPReset(user->RhoDrawSP));
261:     PetscCall(PetscDrawAxisSetLabels(axis_Rho, "Particles", "x", "rho"));
262:     PetscCall(PetscDrawSetSave(user->RhoDraw, "ex9_rho_spatial.png"));

264:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "potential_monitor", 0, 0, 400, 300, &user->PotDraw));
265:     PetscCall(PetscDrawSetFromOptions(user->PotDraw));
266:     PetscCall(PetscDrawSPCreate(user->PotDraw, 10, &user->PotDrawSP));
267:     PetscCall(PetscDrawSPSetDimension(user->PotDrawSP, 1));
268:     PetscCall(PetscDrawSPGetAxis(user->PotDrawSP, &axis_Pot));
269:     PetscCall(PetscDrawSPReset(user->PotDrawSP));
270:     PetscCall(PetscDrawAxisSetLabels(axis_Pot, "Particles", "x", "potential"));
271:     PetscCall(PetscDrawSetSave(user->PotDraw, "ex9_phi_spatial.png"));
272:   }
273:   PetscFunctionReturn(PETSC_SUCCESS);
274: }

276: static PetscErrorCode DestroyContext(AppCtx *user)
277: {
278:   PetscFunctionBeginUser;
279:   PetscCall(PetscDrawLGDestroy(&user->drawlg_ef));
280:   PetscCall(PetscDrawDestroy(&user->drawef));
281:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
282:   PetscCall(PetscDrawDestroy(&user->drawic_x));
283:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
284:   PetscCall(PetscDrawDestroy(&user->drawic_v));
285:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_w));
286:   PetscCall(PetscDrawDestroy(&user->drawic_w));
287:   PetscCall(PetscDrawSPDestroy(&user->positionDrawSP));
288:   PetscCall(PetscDrawDestroy(&user->positionDraw));

290:   PetscCall(PetscDrawSPDestroy(&user->EDrawSP));
291:   PetscCall(PetscDrawDestroy(&user->EDraw));
292:   PetscCall(PetscDrawSPDestroy(&user->RhoDrawSP));
293:   PetscCall(PetscDrawDestroy(&user->RhoDraw));
294:   PetscCall(PetscDrawSPDestroy(&user->PotDrawSP));
295:   PetscCall(PetscDrawDestroy(&user->PotDraw));

297:   PetscCall(PetscBagDestroy(&user->bag));
298:   PetscFunctionReturn(PETSC_SUCCESS);
299: }

301: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
302: {
303:   const PetscScalar *w;
304:   PetscInt           Np;

306:   PetscFunctionBeginUser;
307:   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
308:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
309:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
310:   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]);
311:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
312:   PetscFunctionReturn(PETSC_SUCCESS);
313: }

315: static PetscErrorCode computeParticleMoments(DM sw, PetscReal moments[3], AppCtx *user)
316: {
317:   DM                 dm;
318:   const PetscReal   *coords;
319:   const PetscScalar *w;
320:   PetscReal          mom[3] = {0.0, 0.0, 0.0};
321:   PetscInt           cell, cStart, cEnd, dim;

323:   PetscFunctionBeginUser;
324:   PetscCall(DMGetDimension(sw, &dim));
325:   PetscCall(DMSwarmGetCellDM(sw, &dm));
326:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
327:   PetscCall(DMSwarmSortGetAccess(sw));
328:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&coords));
329:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
330:   for (cell = cStart; cell < cEnd; ++cell) {
331:     PetscInt *pidx;
332:     PetscInt  Np, p, d;

334:     PetscCall(DMSwarmSortGetPointsPerCell(sw, cell, &Np, &pidx));
335:     for (p = 0; p < Np; ++p) {
336:       const PetscInt   idx = pidx[p];
337:       const PetscReal *c   = &coords[idx * dim];

339:       mom[0] += PetscRealPart(w[idx]);
340:       mom[1] += PetscRealPart(w[idx]) * c[0];
341:       for (d = 0; d < dim; ++d) mom[2] += PetscRealPart(w[idx]) * c[d] * c[d];
342:       //if (w[idx] < 0. ) PetscPrintf(PETSC_COMM_WORLD, "error, negative weight %" PetscInt_FMT " \n", idx);
343:     }
344:     PetscCall(PetscFree(pidx));
345:   }
346:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&coords));
347:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
348:   PetscCall(DMSwarmSortRestoreAccess(sw));
349:   PetscCallMPI(MPIU_Allreduce(mom, moments, 3, MPIU_REAL, MPI_SUM, PetscObjectComm((PetscObject)sw)));
350:   PetscFunctionReturn(PETSC_SUCCESS);
351: }

353: static void f0_1(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[])
354: {
355:   f0[0] = u[0];
356: }

358: static void f0_x(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[])
359: {
360:   f0[0] = x[0] * u[0];
361: }

363: static void f0_r2(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[])
364: {
365:   PetscInt d;

367:   f0[0] = 0.0;
368:   for (d = 0; d < dim; ++d) f0[0] += PetscSqr(x[d]) * u[0];
369: }

371: static PetscErrorCode computeFEMMoments(DM dm, Vec u, PetscReal moments[3], AppCtx *user)
372: {
373:   PetscDS     prob;
374:   PetscScalar mom;
375:   PetscInt    field = 0;

377:   PetscFunctionBeginUser;
378:   PetscCall(DMGetDS(dm, &prob));
379:   PetscCall(PetscDSSetObjective(prob, field, &f0_1));
380:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
381:   moments[0] = PetscRealPart(mom);
382:   PetscCall(PetscDSSetObjective(prob, field, &f0_x));
383:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
384:   moments[1] = PetscRealPart(mom);
385:   PetscCall(PetscDSSetObjective(prob, field, &f0_r2));
386:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
387:   moments[2] = PetscRealPart(mom);
388:   PetscFunctionReturn(PETSC_SUCCESS);
389: }

391: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
392: {
393:   AppCtx    *user = (AppCtx *)ctx;
394:   DM         dm, sw;
395:   PetscReal *E;
396:   PetscReal  Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., temp = 0., *weight, chargesum = 0.;
397:   PetscReal *x, *v;
398:   PetscInt  *species, dim, p, d, Np, cStart, cEnd;
399:   PetscReal  pmoments[3]; /* \int f, \int x f, \int r^2 f */
400:   PetscReal  fmoments[3]; /* \int \hat f, \int x \hat f, \int r^2 \hat f */
401:   Vec        rho;

403:   PetscFunctionBeginUser;
404:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
405:   PetscCall(TSGetDM(ts, &sw));
406:   PetscCall(DMSwarmGetCellDM(sw, &dm));
407:   PetscCall(DMGetDimension(sw, &dim));
408:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
409:   PetscCall(DMSwarmSortGetAccess(sw));
410:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
411:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
412:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
413:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
414:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
415:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

417:   for (p = 0; p < Np; ++p) {
418:     for (d = 0; d < 1; ++d) {
419:       temp = PetscAbsReal(E[p * dim + d]);
420:       if (temp > Emax) Emax = temp;
421:     }
422:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
423:     sum += E[p * dim];
424:     chargesum += user->charges[0] * weight[p];
425:   }
426:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
427:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : -16.;

429:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
430:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
431:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
432:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
433:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

435:   Parameter *param;
436:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
437:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "charges", &rho));
438:   if (user->em == EM_PRIMAL) {
439:     PetscCall(computeParticleMoments(sw, pmoments, user));
440:     PetscCall(computeFEMMoments(dm, rho, fmoments, user));
441:   } else if (user->em == EM_MIXED) {
442:     DM       potential_dm;
443:     IS       potential_IS;
444:     PetscInt fields = 1;
445:     PetscCall(DMCreateSubDM(dm, 1, &fields, &potential_IS, &potential_dm));

447:     PetscCall(computeParticleMoments(sw, pmoments, user));
448:     PetscCall(computeFEMMoments(potential_dm, rho, fmoments, user));
449:     PetscCall(DMDestroy(&potential_dm));
450:     PetscCall(ISDestroy(&potential_IS));
451:   }
452:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "charges", &rho));

454:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%+e\t%e\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[2], (double)fmoments[0], (double)fmoments[1], (double)fmoments[2]));
455:   PetscCall(PetscDrawLGAddPoint(user->drawlg_ef, &t, &lgEmax));
456:   PetscCall(PetscDrawLGDraw(user->drawlg_ef));
457:   PetscCall(PetscDrawSave(user->drawef));
458:   PetscFunctionReturn(PETSC_SUCCESS);
459: }

461: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
462: {
463:   AppCtx            *user = (AppCtx *)ctx;
464:   DM                 dm, sw;
465:   const PetscScalar *u;
466:   PetscReal         *weight, *pos, *vel;
467:   PetscInt           dim, p, Np, cStart, cEnd;

469:   PetscFunctionBegin;
470:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
471:   PetscCall(TSGetDM(ts, &sw));
472:   PetscCall(DMSwarmGetCellDM(sw, &dm));
473:   PetscCall(DMGetDimension(sw, &dim));
474:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
475:   PetscCall(DMSwarmSortGetAccess(sw));
476:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

478:   if (step == 0) {
479:     PetscCall(PetscDrawHGReset(user->drawhgic_x));
480:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &user->drawic_x));
481:     PetscCall(PetscDrawClear(user->drawic_x));
482:     PetscCall(PetscDrawFlush(user->drawic_x));

484:     PetscCall(PetscDrawHGReset(user->drawhgic_v));
485:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &user->drawic_v));
486:     PetscCall(PetscDrawClear(user->drawic_v));
487:     PetscCall(PetscDrawFlush(user->drawic_v));

489:     PetscCall(PetscDrawHGReset(user->drawhgic_w));
490:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_w, &user->drawic_w));
491:     PetscCall(PetscDrawClear(user->drawic_w));
492:     PetscCall(PetscDrawFlush(user->drawic_w));

494:     PetscCall(VecGetArrayRead(U, &u));
495:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
496:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
497:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));

499:     PetscCall(VecGetLocalSize(U, &Np));
500:     Np /= dim * 2;
501:     for (p = 0; p < Np; ++p) {
502:       PetscCall(PetscDrawHGAddValue(user->drawhgic_x, pos[p * dim]));
503:       PetscCall(PetscDrawHGAddValue(user->drawhgic_v, vel[p * dim]));
504:       PetscCall(PetscDrawHGAddValue(user->drawhgic_w, weight[p]));
505:     }

507:     PetscCall(VecRestoreArrayRead(U, &u));
508:     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
509:     PetscCall(PetscDrawHGSave(user->drawhgic_x));

511:     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
512:     PetscCall(PetscDrawHGSave(user->drawhgic_v));

514:     PetscCall(PetscDrawHGDraw(user->drawhgic_w));
515:     PetscCall(PetscDrawHGSave(user->drawhgic_w));

517:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
518:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
519:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
520:   }
521:   PetscFunctionReturn(PETSC_SUCCESS);
522: }

524: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
525: {
526:   AppCtx         *user = (AppCtx *)ctx;
527:   DM              dm, sw;
528:   PetscScalar    *x, *v, *weight;
529:   PetscReal       lower[3], upper[3], speed;
530:   const PetscInt *s;
531:   PetscInt        dim, cStart, cEnd, c;

533:   PetscFunctionBeginUser;
534:   if (step > 0 && step % user->ostep == 0) {
535:     PetscCall(TSGetDM(ts, &sw));
536:     PetscCall(DMSwarmGetCellDM(sw, &dm));
537:     PetscCall(DMGetDimension(dm, &dim));
538:     PetscCall(DMGetBoundingBox(dm, lower, upper));
539:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
540:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
541:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
542:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
543:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
544:     PetscCall(DMSwarmSortGetAccess(sw));
545:     PetscCall(PetscDrawSPReset(user->positionDrawSP));
546:     PetscCall(PetscDrawSPSetLimits(user->positionDrawSP, lower[0], upper[0], lower[1], upper[1]));
547:     PetscCall(PetscDrawSPSetLimits(user->positionDrawSP, lower[0], upper[0], -12, 12));
548:     for (c = 0; c < cEnd - cStart; ++c) {
549:       PetscInt *pidx, Npc, q;
550:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
551:       for (q = 0; q < Npc; ++q) {
552:         const PetscInt p = pidx[q];
553:         if (s[p] == 0) {
554:           speed = PetscSqrtReal(PetscSqr(v[p * dim]) + PetscSqr(v[p * dim + 1]));
555:           if (dim == 1 || user->fake_1D) {
556:             PetscCall(PetscDrawSPAddPointColorized(user->positionDrawSP, &x[p * dim], &v[p * dim], &speed));
557:           } else {
558:             PetscCall(PetscDrawSPAddPointColorized(user->positionDrawSP, &x[p * dim], &x[p * dim + 1], &speed));
559:           }
560:         } else if (s[p] == 1) {
561:           PetscCall(PetscDrawSPAddPoint(user->positionDrawSP, &x[p * dim], &v[p * dim]));
562:         }
563:       }
564:       PetscCall(PetscFree(pidx));
565:     }
566:     PetscCall(PetscDrawSPDraw(user->positionDrawSP, PETSC_TRUE));
567:     PetscCall(PetscDrawSave(user->positionDraw));
568:     PetscCall(DMSwarmSortRestoreAccess(sw));
569:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
570:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
571:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
572:     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
573:   }
574:   PetscFunctionReturn(PETSC_SUCCESS);
575: }

577: static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
578: {
579:   AppCtx      *user = (AppCtx *)ctx;
580:   DM           dm, sw;
581:   PetscScalar *x, *E, *weight, *pot, *charges;
582:   PetscReal    lower[3], upper[3], xval;
583:   PetscInt     dim, cStart, cEnd, c;

585:   PetscFunctionBeginUser;
586:   if (step > 0 && step % user->ostep == 0) {
587:     PetscCall(TSGetDM(ts, &sw));
588:     PetscCall(DMSwarmGetCellDM(sw, &dm));
589:     PetscCall(DMGetDimension(dm, &dim));
590:     PetscCall(DMGetBoundingBox(dm, lower, upper));
591:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

593:     PetscCall(PetscDrawSPReset(user->RhoDrawSP));
594:     PetscCall(PetscDrawSPReset(user->EDrawSP));
595:     PetscCall(PetscDrawSPReset(user->PotDrawSP));
596:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
597:     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
598:     PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
599:     PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
600:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

602:     PetscCall(DMSwarmSortGetAccess(sw));
603:     for (c = 0; c < cEnd - cStart; ++c) {
604:       PetscReal Esum = 0.0;
605:       PetscInt *pidx, Npc, q;
606:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
607:       for (q = 0; q < Npc; ++q) {
608:         const PetscInt p = pidx[q];
609:         Esum += E[p * dim];
610:       }
611:       xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
612:       PetscCall(PetscDrawSPAddPoint(user->EDrawSP, &xval, &Esum));
613:       PetscCall(PetscFree(pidx));
614:     }
615:     for (c = 0; c < (cEnd - cStart); ++c) {
616:       xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
617:       PetscCall(PetscDrawSPAddPoint(user->RhoDrawSP, &xval, &charges[c]));
618:       PetscCall(PetscDrawSPAddPoint(user->PotDrawSP, &xval, &pot[c]));
619:     }
620:     PetscCall(PetscDrawSPDraw(user->RhoDrawSP, PETSC_TRUE));
621:     PetscCall(PetscDrawSave(user->RhoDraw));
622:     PetscCall(PetscDrawSPDraw(user->EDrawSP, PETSC_TRUE));
623:     PetscCall(PetscDrawSave(user->EDraw));
624:     PetscCall(PetscDrawSPDraw(user->PotDrawSP, PETSC_TRUE));
625:     PetscCall(PetscDrawSave(user->PotDraw));
626:     PetscCall(DMSwarmSortRestoreAccess(sw));
627:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
628:     PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));
629:     PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));
630:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
631:     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
632:   }
633:   PetscFunctionReturn(PETSC_SUCCESS);
634: }

636: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
637: {
638:   PetscBag   bag;
639:   Parameter *p;

641:   PetscFunctionBeginUser;
642:   /* setup PETSc parameter bag */
643:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
644:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
645:   bag = ctx->bag;
646:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
647:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
648:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
649:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
650:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
651:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
652:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
653:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

655:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
656:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
657:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
658:   PetscCall(PetscBagSetFromOptions(bag));
659:   {
660:     PetscViewer       viewer;
661:     PetscViewerFormat format;
662:     PetscBool         flg;

664:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
665:     if (flg) {
666:       PetscCall(PetscViewerPushFormat(viewer, format));
667:       PetscCall(PetscBagView(bag, viewer));
668:       PetscCall(PetscViewerFlush(viewer));
669:       PetscCall(PetscViewerPopFormat(viewer));
670:       PetscCall(PetscViewerDestroy(&viewer));
671:     }
672:   }
673:   PetscFunctionReturn(PETSC_SUCCESS);
674: }

676: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
677: {
678:   PetscFunctionBeginUser;
679:   PetscCall(DMCreate(comm, dm));
680:   PetscCall(DMSetType(*dm, DMPLEX));
681:   PetscCall(DMSetFromOptions(*dm));
682:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
683:   PetscFunctionReturn(PETSC_SUCCESS);
684: }

686: 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[])
687: {
688:   f0[0] = -constants[SIGMA];
689: }

691: 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[])
692: {
693:   PetscInt d;
694:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
695: }

697: 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[])
698: {
699:   PetscInt d;
700:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
701: }

703: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
704: {
705:   *u = 0.0;
706:   return PETSC_SUCCESS;
707: }

709: /*
710:    /  I   -grad\ / q \ = /0\
711:    \-div    0  / \phi/   \f/
712: */
713: 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[])
714: {
715:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
716: }

718: 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[])
719: {
720:   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
721: }

723: 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[])
724: {
725:   f0[0] += constants[SIGMA];
726:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
727: }

729: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
730: 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[])
731: {
732:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
733: }

735: 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[])
736: {
737:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
738: }

740: 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[])
741: {
742:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
743: }

745: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
746: {
747:   PetscFE   fephi, feq;
748:   PetscDS   ds;
749:   PetscBool simplex;
750:   PetscInt  dim;

752:   PetscFunctionBeginUser;
753:   PetscCall(DMGetDimension(dm, &dim));
754:   PetscCall(DMPlexIsSimplex(dm, &simplex));
755:   if (user->em == EM_MIXED) {
756:     DMLabel        label;
757:     const PetscInt id = 1;

759:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
760:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
761:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
762:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
763:     PetscCall(PetscFECopyQuadrature(feq, fephi));
764:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
765:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
766:     PetscCall(DMCreateDS(dm));
767:     PetscCall(PetscFEDestroy(&fephi));
768:     PetscCall(PetscFEDestroy(&feq));

770:     PetscCall(DMGetLabel(dm, "marker", &label));
771:     PetscCall(DMGetDS(dm, &ds));

773:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
774:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
775:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
776:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
777:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

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

781:   } else if (user->em == EM_PRIMAL) {
782:     MatNullSpace nullsp;
783:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
784:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
785:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
786:     PetscCall(DMCreateDS(dm));
787:     PetscCall(DMGetDS(dm, &ds));
788:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
789:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
790:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
791:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
792:     PetscCall(MatNullSpaceDestroy(&nullsp));
793:     PetscCall(PetscFEDestroy(&fephi));
794:   }
795:   PetscFunctionReturn(PETSC_SUCCESS);
796: }

798: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
799: {
800:   SNES         snes;
801:   Mat          J;
802:   MatNullSpace nullSpace;

804:   PetscFunctionBeginUser;
805:   PetscCall(CreateFEM(dm, user));
806:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
807:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
808:   PetscCall(SNESSetDM(snes, dm));
809:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
810:   PetscCall(SNESSetFromOptions(snes));

812:   PetscCall(DMCreateMatrix(dm, &J));
813:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
814:   PetscCall(MatSetNullSpace(J, nullSpace));
815:   PetscCall(MatNullSpaceDestroy(&nullSpace));
816:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
817:   PetscCall(MatDestroy(&J));
818:   user->snes = snes;
819:   PetscFunctionReturn(PETSC_SUCCESS);
820: }

822: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
823: {
824:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
825:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
826:   return PETSC_SUCCESS;
827: }
828: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
829: {
830:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
831:   return PETSC_SUCCESS;
832: }

834: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
835: {
836:   const PetscReal alpha = scale ? scale[0] : 0.0;
837:   const PetscReal k     = scale ? scale[1] : 1.;
838:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
839:   return PETSC_SUCCESS;
840: }

842: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
843: {
844:   const PetscReal alpha = scale ? scale[0] : 0.;
845:   const PetscReal k     = scale ? scale[0] : 1.;
846:   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
847:   return PETSC_SUCCESS;
848: }

850: PetscErrorCode PetscPDFCosine1D_TwoStream(const PetscReal x[], const PetscReal scale[], PetscReal p[])
851: {
852:   const PetscReal alpha = scale ? scale[0] : 0.0;
853:   const PetscReal k     = scale ? scale[1] : 1.;
854:   p[0]                  = (1. + alpha * PetscCosReal(k * x[0]));
855:   return PETSC_SUCCESS;
856: }

858: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
859: {
860:   DM           vdm, dm;
861:   PetscScalar *weight;
862:   PetscReal   *x, *v, vmin[3], vmax[3], gmin[3], gmax[3], xi0[3];
863:   PetscInt    *N, Ns, dim, *cellid, *species, Np, cStart, cEnd, Npc, n;
864:   PetscInt     Np_global, p, q, s, c, d, cv;
865:   PetscBool    flg;
866:   PetscMPIInt  size, rank;
867:   Parameter   *param;

869:   PetscFunctionBegin;
870:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
871:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
872:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
873:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
874:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
875:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
876:   PetscCall(PetscCalloc1(Ns, &N));
877:   n = Ns;
878:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
879:   PetscOptionsEnd();

881:   PetscCall(DMGetDimension(sw, &dim));
882:   PetscCall(DMSwarmGetCellDM(sw, &dm));
883:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

885:   PetscCall(DMCreate(PETSC_COMM_SELF, &vdm));
886:   PetscCall(DMSetType(vdm, DMPLEX));
887:   PetscCall(DMPlexSetOptionsPrefix(vdm, "v"));
888:   PetscCall(DMSetFromOptions(vdm));
889:   PetscCall(DMViewFromOptions(vdm, NULL, "-vdm_view"));

891:   PetscInt vStart, vEnd;
892:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vStart, &vEnd));
893:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));

895:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
896:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
897:   Np = (cEnd - cStart) * (vEnd - vStart);
898:   PetscCallMPI(MPIU_Allreduce(&Np, &Np_global, 1, MPIU_INT, MPIU_SUM, PETSC_COMM_WORLD));
899:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Global Np = %" PetscInt_FMT "\n", Np_global));
900:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
901:   Npc = Np / (cEnd - cStart);
902:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
903:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
904:     for (s = 0; s < Ns; ++s) {
905:       for (q = 0; q < Npc; ++q, ++p) cellid[p] = c;
906:     }
907:   }
908:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
909:   PetscCall(PetscFree(N));

911:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
912:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
913:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
914:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));

916:   PetscCall(DMSwarmSortGetAccess(sw));
917:   for (c = 0; c < cEnd - cStart; ++c) {
918:     const PetscInt cell = c + cStart;
919:     PetscInt      *pidx, Npc;
920:     PetscReal      centroid[3], volume;

922:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
923:     PetscCall(DMPlexComputeCellGeometryFVM(dm, cell, &volume, centroid, NULL));
924:     for (q = 0; q < Npc; ++q) {
925:       const PetscInt p = pidx[q];

927:       for (d = 0; d < dim; ++d) {
928:         x[p * dim + d] = centroid[d];
929:         v[p * dim + d] = vmin[0] + (q + 0.5) * (vmax[0] - vmin[0]) / Npc;
930:         if (user->fake_1D && d > 0) v[p * dim + d] = 0;
931:       }
932:     }
933:     PetscCall(PetscFree(pidx));
934:   }
935:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));

937:   /* Setup Quadrature for spatial and velocity weight calculations*/
938:   PetscQuadrature  quad_x;
939:   PetscInt         Nq_x;
940:   const PetscReal *wq_x, *xq_x;
941:   PetscReal       *xq_x_extended;
942:   PetscReal        weightsum = 0., totalcellweight = 0., *weight_x, *weight_v;
943:   PetscReal        scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};

945:   PetscCall(PetscCalloc2(cEnd - cStart, &weight_x, Np, &weight_v));
946:   if (user->fake_1D) PetscCall(PetscDTGaussTensorQuadrature(1, 1, 5, -1.0, 1.0, &quad_x));
947:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad_x));
948:   PetscCall(PetscQuadratureGetData(quad_x, NULL, NULL, &Nq_x, &xq_x, &wq_x));
949:   if (user->fake_1D) {
950:     PetscCall(PetscCalloc1(Nq_x * dim, &xq_x_extended));
951:     for (PetscInt i = 0; i < Nq_x; ++i) xq_x_extended[i * dim] = xq_x[i];
952:   }
953:   /* Integrate the density function to get the weights of particles in each cell */
954:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
955:   for (c = cStart; c < cEnd; ++c) {
956:     PetscReal          v0_x[3], J_x[9], invJ_x[9], detJ_x, xr_x[3], den_x;
957:     PetscInt          *pidx, Npc, q;
958:     PetscInt           Ncx;
959:     const PetscScalar *array_x;
960:     PetscScalar       *coords_x = NULL;
961:     PetscBool          isDGx;
962:     weight_x[c] = 0.;

964:     PetscCall(DMPlexGetCellCoordinates(dm, c, &isDGx, &Ncx, &array_x, &coords_x));
965:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
966:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0_x, J_x, invJ_x, &detJ_x));
967:     for (q = 0; q < Nq_x; ++q) {
968:       /*Transform quadrature points from ref space to real space (0,12.5664)*/
969:       if (user->fake_1D) CoordinatesRefToReal(dim, dim, xi0, v0_x, J_x, &xq_x_extended[q * dim], xr_x);
970:       else CoordinatesRefToReal(dim, dim, xi0, v0_x, J_x, &xq_x[q * dim], xr_x);

972:       /*Transform quadrature points from real space to ideal real space (0, 2PI/k)*/
973:       if (user->fake_1D) {
974:         if (user->twostream) PetscCall(PetscPDFCosine1D_TwoStream(xr_x, scale, &den_x));
975:         else PetscCall(PetscPDFCosine1D(xr_x, scale, &den_x));
976:         detJ_x = J_x[0];
977:       } else PetscCall(PetscPDFCosine2D(xr_x, scale, &den_x));
978:       /*We have to transform the quadrature weights as well*/
979:       weight_x[c] += den_x * (wq_x[q] * detJ_x);
980:     }
981:     // Get the cell numbering for consistent output between sequential and distributed tests
982:     IS              globalOrdering;
983:     const PetscInt *ordering;
984:     PetscCall(DMPlexGetCellNumbering(dm, &globalOrdering));
985:     PetscCall(ISGetIndices(globalOrdering, &ordering));
986:     PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(coords_x[0]), (double)PetscRealPart(coords_x[2]), (double)weight_x[c]));
987:     PetscCall(ISRestoreIndices(globalOrdering, &ordering));
988:     totalcellweight += weight_x[c];
989:     // Confirm the number of particles per spatial cell conforms to the size of the velocity grid
990:     PetscCheck(Npc == vEnd - vStart, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Number of particles %" PetscInt_FMT " in cell (rank %d/%d) != %" PetscInt_FMT " number of velocity vertices", Npc, rank, size, vEnd - vStart);

992:     /* Set weights to be gaussian in velocity cells (using exact solution) */
993:     for (cv = 0; cv < vEnd - vStart; ++cv) {
994:       PetscInt           Nc;
995:       const PetscScalar *array_v;
996:       PetscScalar       *coords_v = NULL;
997:       PetscBool          isDG;
998:       PetscCall(DMPlexGetCellCoordinates(vdm, cv, &isDG, &Nc, &array_v, &coords_v));

1000:       const PetscInt p = pidx[cv];
1001:       // Two stream function from 1/2pi v^2 e^(-v^2/2)
1002:       if (user->twostream)
1003:         weight_v[p] = 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.)));
1004:       else weight_v[p] = 0.5 * (PetscErfReal(coords_v[1] / PetscSqrtReal(2.)) - PetscErfReal(coords_v[0] / PetscSqrtReal(2.)));

1006:       weight[p] = user->totalWeight * weight_v[p] * weight_x[c];
1007:       if (weight[p] > 1.) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "weights: %g, %g, %g\n", user->totalWeight, weight_v[p], weight_x[c]));
1008:       //PetscPrintf(PETSC_COMM_WORLD, "particle %"PetscInt_FMT": %g, weight_v: %g weight_x: %g\n", p, weight[p], weight_v[p], weight_x[p]);
1009:       weightsum += weight[p];

1011:       PetscCall(DMPlexRestoreCellCoordinates(vdm, cv, &isDG, &Nc, &array_v, &coords_v));
1012:     }
1013:     PetscCall(DMPlexRestoreCellCoordinates(dm, c, &isDGx, &Ncx, &array_x, &coords_x));
1014:     PetscCall(PetscFree(pidx));
1015:   }
1016:   PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1017:   PetscReal global_cellweight, global_weightsum;
1018:   PetscCallMPI(MPIU_Allreduce(&totalcellweight, &global_cellweight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1019:   PetscCallMPI(MPIU_Allreduce(&weightsum, &global_weightsum, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1020:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)global_cellweight, (double)global_weightsum));
1021:   if (user->fake_1D) PetscCall(PetscFree(xq_x_extended));
1022:   PetscCall(PetscFree2(weight_x, weight_v));
1023:   PetscCall(PetscQuadratureDestroy(&quad_x));
1024:   PetscCall(DMSwarmSortRestoreAccess(sw));
1025:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
1026:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1027:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1028:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1029:   PetscCall(DMDestroy(&vdm));
1030:   PetscFunctionReturn(PETSC_SUCCESS);
1031: }

1033: static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
1034: {
1035:   DM         dm;
1036:   PetscInt  *species;
1037:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
1038:   PetscInt   Np, dim;

1040:   PetscFunctionBegin;
1041:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1042:   PetscCall(DMGetDimension(sw, &dim));
1043:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1044:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
1045:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
1046:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1047:   for (PetscInt p = 0; p < Np; ++p) {
1048:     totalWeight += weight[p];
1049:     totalCharge += user->charges[species[p]] * weight[p];
1050:   }
1051:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1052:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1053:   {
1054:     Parameter *param;
1055:     PetscReal  Area;

1057:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1058:     switch (dim) {
1059:     case 1:
1060:       Area = (gmax[0] - gmin[0]);
1061:       break;
1062:     case 2:
1063:       if (user->fake_1D) {
1064:         Area = (gmax[0] - gmin[0]);
1065:       } else {
1066:         Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1067:       }
1068:       break;
1069:     case 3:
1070:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1071:       break;
1072:     default:
1073:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1074:     }
1075:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1076:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
1077:     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));
1078:     param->sigma = PetscAbsReal(global_charge / (Area));

1080:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
1081:     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,
1082:                           (double)param->vlasovNumber));
1083:   }
1084:   /* Setup Constants */
1085:   {
1086:     PetscDS    ds;
1087:     Parameter *param;
1088:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1089:     PetscScalar constants[NUM_CONSTANTS];
1090:     constants[SIGMA]   = param->sigma;
1091:     constants[V0]      = param->v0;
1092:     constants[T0]      = param->t0;
1093:     constants[X0]      = param->x0;
1094:     constants[M0]      = param->m0;
1095:     constants[Q0]      = param->q0;
1096:     constants[PHI0]    = param->phi0;
1097:     constants[POISSON] = param->poissonNumber;
1098:     constants[VLASOV]  = param->vlasovNumber;
1099:     PetscCall(DMGetDS(dm, &ds));
1100:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
1101:   }
1102:   PetscFunctionReturn(PETSC_SUCCESS);
1103: }

1105: static PetscErrorCode InitializeVelocities_Fake1D(DM sw, AppCtx *user)
1106: {
1107:   DM         dm;
1108:   PetscReal *v;
1109:   PetscInt  *species, cStart, cEnd;
1110:   PetscInt   dim, Np;

1112:   PetscFunctionBegin;
1113:   PetscCall(DMGetDimension(sw, &dim));
1114:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1115:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1116:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1117:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1118:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1119:   PetscRandom rnd;
1120:   PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1121:   PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1122:   PetscCall(PetscRandomSetFromOptions(rnd));

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

1127:     PetscCall(PetscRandomGetValueReal(rnd, &a[0]));
1128:     if (user->perturbed_weights) {
1129:       PetscCall(PetscPDFSampleConstant1D(a, NULL, vel));
1130:     } else {
1131:       PetscCall(PetscPDFSampleGaussian1D(a, NULL, vel));
1132:     }
1133:     v[p * dim] = vel[0];
1134:   }
1135:   PetscCall(PetscRandomDestroy(&rnd));
1136:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1137:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1138:   PetscFunctionReturn(PETSC_SUCCESS);
1139: }

1141: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1142: {
1143:   PetscReal v0[2] = {1., 0.};
1144:   PetscInt  dim;

1146:   PetscFunctionBeginUser;
1147:   PetscCall(DMGetDimension(dm, &dim));
1148:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
1149:   PetscCall(DMSetType(*sw, DMSWARM));
1150:   PetscCall(DMSetDimension(*sw, dim));
1151:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
1152:   PetscCall(DMSwarmSetCellDM(*sw, dm));
1153:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
1154:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
1155:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
1156:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initCoordinates", dim, PETSC_REAL));
1157:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initVelocity", dim, PETSC_REAL));
1158:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
1159:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "potential", dim, PETSC_REAL));
1160:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "charges", dim, PETSC_REAL));
1161:   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
1162:   PetscCall(DMSetApplicationContext(*sw, user));
1163:   PetscCall(DMSetFromOptions(*sw));
1164:   user->swarm = *sw;
1165:   if (user->perturbed_weights) {
1166:     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
1167:   } else {
1168:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
1169:     PetscCall(DMSwarmInitializeCoordinates(*sw));
1170:     if (user->fake_1D) {
1171:       PetscCall(InitializeVelocities_Fake1D(*sw, user));
1172:     } else {
1173:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
1174:     }
1175:   }
1176:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
1177:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
1178:   {
1179:     Vec gc, gc0, gv, gv0;

1181:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1182:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1183:     PetscCall(VecCopy(gc, gc0));
1184:     PetscCall(VecViewFromOptions(gc, NULL, "-ic_x_view"));
1185:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1186:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1187:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv));
1188:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0));
1189:     PetscCall(VecCopy(gv, gv0));
1190:     PetscCall(VecViewFromOptions(gv, NULL, "-ic_v_view"));
1191:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv));
1192:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0));
1193:   }
1194:   PetscFunctionReturn(PETSC_SUCCESS);
1195: }

1197: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
1198: {
1199:   AppCtx     *user;
1200:   PetscReal  *coords;
1201:   PetscInt   *species, dim, Np, Ns;
1202:   PetscMPIInt size;

1204:   PetscFunctionBegin;
1205:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
1206:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
1207:   PetscCall(DMGetDimension(sw, &dim));
1208:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1209:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1210:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1212:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1213:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1214:   for (PetscInt p = 0; p < Np; ++p) {
1215:     PetscReal *pcoord = &coords[p * dim];
1216:     PetscReal  pE[3]  = {0., 0., 0.};

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

1223:       if (p == q) continue;
1224:       q_q = user->charges[species[q]] * 1.;
1225:       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1226:       r = DMPlex_NormD_Internal(dim, rpq);
1227:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
1228:       r3 = PetscPowRealInt(r, 3);
1229:       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
1230:     }
1231:     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
1232:   }
1233:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1234:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1235:   PetscFunctionReturn(PETSC_SUCCESS);
1236: }

1238: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[])
1239: {
1240:   DM              dm;
1241:   AppCtx         *user;
1242:   PetscDS         ds;
1243:   PetscFE         fe;
1244:   Mat             M_p, M;
1245:   Vec             phi, locPhi, rho, f;
1246:   PetscReal      *coords;
1247:   PetscInt        dim, cStart, cEnd, Np;
1248:   PetscQuadrature q;

1250:   PetscFunctionBegin;
1251:   PetscCall(DMGetDimension(sw, &dim));
1252:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1253:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1255:   KSP         ksp;
1256:   Vec         rho0;
1257:   char        oldField[PETSC_MAX_PATH_LEN];
1258:   const char *tmp;

1260:   /* Create the charges rho */
1261:   PetscCall(SNESGetDM(snes, &dm));
1262:   PetscCall(DMSwarmVectorGetField(sw, &tmp));
1263:   PetscCall(PetscStrncpy(oldField, tmp, PETSC_MAX_PATH_LEN));
1264:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1265:   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
1266:   PetscCall(DMSwarmVectorDefineField(sw, oldField));

1268:   PetscCall(DMCreateMassMatrix(dm, dm, &M));
1269:   PetscCall(DMGetGlobalVector(dm, &rho0));
1270:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Primal Compute"));
1271:   PetscCall(DMGetGlobalVector(dm, &rho));
1272:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
1273:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));

1275:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1276:   PetscCall(MatMultTranspose(M_p, f, rho));
1277:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1278:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1279:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1280:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1282:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1283:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1284:   PetscCall(KSPSetOperators(ksp, M, M));
1285:   PetscCall(KSPSetFromOptions(ksp));
1286:   PetscCall(KSPSolve(ksp, rho, rho0));
1287:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1289:   PetscInt           rhosize;
1290:   PetscReal         *charges;
1291:   const PetscScalar *rho_vals;
1292:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1293:   PetscCall(VecGetLocalSize(rho0, &rhosize));
1294:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1295:   for (PetscInt c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1296:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1297:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

1299:   PetscCall(VecScale(rho, -1.0));

1301:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1302:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1303:   PetscCall(DMRestoreGlobalVector(dm, &rho0));
1304:   PetscCall(KSPDestroy(&ksp));
1305:   PetscCall(MatDestroy(&M_p));
1306:   PetscCall(MatDestroy(&M));

1308:   PetscCall(DMGetGlobalVector(dm, &phi));
1309:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1310:   PetscCall(VecSet(phi, 0.0));
1311:   PetscCall(SNESSolve(snes, rho, phi));
1312:   PetscCall(DMRestoreGlobalVector(dm, &rho));
1313:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1315:   PetscInt           phisize;
1316:   PetscReal         *pot;
1317:   const PetscScalar *phi_vals;
1318:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1319:   PetscCall(VecGetLocalSize(phi, &phisize));
1320:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1321:   for (PetscInt c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1322:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1323:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1325:   PetscCall(DMGetLocalVector(dm, &locPhi));
1326:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1327:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1328:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1330:   PetscCall(DMGetDS(dm, &ds));
1331:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1332:   PetscCall(DMSwarmSortGetAccess(sw));
1333:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1334:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1336:   for (PetscInt c = cStart; c < cEnd; ++c) {
1337:     PetscTabulation tab;
1338:     PetscScalar    *clPhi = NULL;
1339:     PetscReal      *pcoord, *refcoord;
1340:     PetscReal       v[3], J[9], invJ[9], detJ;
1341:     PetscInt       *points;
1342:     PetscInt        Ncp;

1344:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1345:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1346:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1347:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1348:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1349:     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
1350:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1351:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ));
1352:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1353:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
1354:       const PetscReal *basisDer = tab->T[1];
1355:       const PetscInt   p        = points[cp];

1357:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1358:       PetscCall(PetscFEGetQuadrature(fe, &q));
1359:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, invJ, NULL, cp, &E[p * dim]));
1360:       for (PetscInt d = 0; d < dim; ++d) {
1361:         E[p * dim + d] *= -1.0;
1362:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1363:       }
1364:     }
1365:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1366:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1367:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1368:     PetscCall(PetscTabulationDestroy(&tab));
1369:     PetscCall(PetscFree(points));
1370:   }
1371:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1372:   PetscCall(DMSwarmSortRestoreAccess(sw));
1373:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1374:   PetscFunctionReturn(PETSC_SUCCESS);
1375: }

1377: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
1378: {
1379:   AppCtx         *user;
1380:   DM              dm, potential_dm;
1381:   KSP             ksp;
1382:   IS              potential_IS;
1383:   PetscDS         ds;
1384:   PetscFE         fe;
1385:   PetscFEGeom     feGeometry;
1386:   Mat             M_p, M;
1387:   Vec             phi, locPhi, rho, f, temp_rho, rho0;
1388:   PetscQuadrature q;
1389:   PetscReal      *coords, *pot;
1390:   PetscInt        dim, cStart, cEnd, Np, fields = 1;
1391:   char            oldField[PETSC_MAX_PATH_LEN];
1392:   const char     *tmp;

1394:   PetscFunctionBegin;
1395:   PetscCall(DMGetDimension(sw, &dim));
1396:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1397:   PetscCall(DMGetApplicationContext(sw, &user));

1399:   /* Create the charges rho */
1400:   PetscCall(SNESGetDM(snes, &dm));
1401:   PetscCall(DMGetGlobalVector(dm, &rho));
1402:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));

1404:   PetscCall(DMCreateSubDM(dm, 1, &fields, &potential_IS, &potential_dm));

1406:   PetscCall(DMSwarmVectorGetField(sw, &tmp));
1407:   PetscCall(PetscStrncpy(oldField, tmp, PETSC_MAX_PATH_LEN));
1408:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1409:   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
1410:   PetscCall(DMSwarmVectorDefineField(sw, oldField));

1412:   PetscCall(DMCreateMassMatrix(potential_dm, potential_dm, &M));
1413:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1414:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1415:   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
1416:   PetscCall(PetscObjectSetName((PetscObject)temp_rho, "Mf"));
1417:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1418:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1419:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1420:   PetscCall(MatMultTranspose(M_p, f, temp_rho));
1421:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1422:   PetscCall(DMGetGlobalVector(potential_dm, &rho0));
1423:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Mixed Compute"));

1425:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1426:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1427:   PetscCall(KSPSetOperators(ksp, M, M));
1428:   PetscCall(KSPSetFromOptions(ksp));
1429:   PetscCall(KSPSolve(ksp, temp_rho, rho0));
1430:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1432:   PetscInt           rhosize;
1433:   PetscReal         *charges;
1434:   const PetscScalar *rho_vals;
1435:   Parameter         *param;
1436:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1437:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1438:   PetscCall(VecGetLocalSize(rho0, &rhosize));

1440:   /* Integral over reference element is size 1.  Reference element area is 4.  Scale rho0 by 1/4 because the basis function is 1/4 */
1441:   PetscCall(VecScale(rho0, 0.25));
1442:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1443:   for (PetscInt c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1444:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1445:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

1447:   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
1448:   PetscCall(VecScale(rho, 0.25));
1449:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1450:   PetscCall(VecViewFromOptions(temp_rho, NULL, "-temprho_view"));
1451:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1452:   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
1453:   PetscCall(DMRestoreGlobalVector(potential_dm, &rho0));

1455:   PetscCall(MatDestroy(&M_p));
1456:   PetscCall(MatDestroy(&M));
1457:   PetscCall(KSPDestroy(&ksp));
1458:   PetscCall(DMDestroy(&potential_dm));
1459:   PetscCall(ISDestroy(&potential_IS));

1461:   PetscCall(DMGetGlobalVector(dm, &phi));
1462:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1463:   PetscCall(VecSet(phi, 0.0));
1464:   PetscCall(SNESSolve(snes, rho, phi));
1465:   PetscCall(DMRestoreGlobalVector(dm, &rho));

1467:   PetscInt           phisize;
1468:   const PetscScalar *phi_vals;
1469:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1470:   PetscCall(VecGetLocalSize(phi, &phisize));
1471:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1472:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1473:   for (PetscInt c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1474:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1475:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1477:   PetscCall(DMGetLocalVector(dm, &locPhi));
1478:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1479:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1480:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1482:   PetscCall(DMGetDS(dm, &ds));
1483:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1484:   PetscCall(DMSwarmSortGetAccess(sw));
1485:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1486:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1487:   PetscCall(PetscFEGetQuadrature(fe, &q));
1488:   PetscCall(PetscFECreateCellGeometry(fe, q, &feGeometry));
1489:   for (PetscInt c = cStart; c < cEnd; ++c) {
1490:     PetscTabulation tab;
1491:     PetscScalar    *clPhi = NULL;
1492:     PetscReal      *pcoord, *refcoord;
1493:     PetscInt       *points;
1494:     PetscInt        Ncp;

1496:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1497:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1498:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1499:     for (PetscInt cp = 0; cp < Ncp; ++cp)
1500:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1501:     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
1502:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1503:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, q, feGeometry.v, feGeometry.J, feGeometry.invJ, feGeometry.detJ));
1504:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));

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

1509:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1510:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, &feGeometry, cp, &E[p * dim]));
1511:       PetscCall(PetscFEPushforward(fe, &feGeometry, 1, &E[p * dim]));
1512:       for (PetscInt d = 0; d < dim; ++d) {
1513:         E[p * dim + d] *= -2.0;
1514:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1515:       }
1516:     }
1517:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1518:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1519:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1520:     PetscCall(PetscTabulationDestroy(&tab));
1521:     PetscCall(PetscFree(points));
1522:   }
1523:   PetscCall(PetscFEDestroyCellGeometry(fe, &feGeometry));
1524:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1525:   PetscCall(DMSwarmSortRestoreAccess(sw));
1526:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1527:   PetscFunctionReturn(PETSC_SUCCESS);
1528: }

1530: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
1531: {
1532:   AppCtx  *ctx;
1533:   PetscInt dim, Np;

1535:   PetscFunctionBegin;
1538:   PetscAssertPointer(E, 3);
1539:   PetscCall(DMGetDimension(sw, &dim));
1540:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1541:   PetscCall(DMGetApplicationContext(sw, &ctx));
1542:   PetscCall(PetscArrayzero(E, Np * dim));

1544:   switch (ctx->em) {
1545:   case EM_PRIMAL:
1546:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
1547:     break;
1548:   case EM_COULOMB:
1549:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1550:     break;
1551:   case EM_MIXED:
1552:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
1553:     break;
1554:   case EM_NONE:
1555:     break;
1556:   default:
1557:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
1558:   }
1559:   PetscFunctionReturn(PETSC_SUCCESS);
1560: }

1562: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1563: {
1564:   DM                 sw;
1565:   SNES               snes = ((AppCtx *)ctx)->snes;
1566:   const PetscReal   *coords, *vel;
1567:   const PetscScalar *u;
1568:   PetscScalar       *g;
1569:   PetscReal         *E, m_p = 1., q_p = -1.;
1570:   PetscInt           dim, d, Np, p;

1572:   PetscFunctionBeginUser;
1573:   PetscCall(TSGetDM(ts, &sw));
1574:   PetscCall(DMGetDimension(sw, &dim));
1575:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1576:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1577:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1578:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1579:   PetscCall(VecGetArrayRead(U, &u));
1580:   PetscCall(VecGetArray(G, &g));

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

1584:   Np /= 2 * dim;
1585:   for (p = 0; p < Np; ++p) {
1586:     for (d = 0; d < dim; ++d) {
1587:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1588:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1589:     }
1590:   }
1591:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1592:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1593:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1594:   PetscCall(VecRestoreArrayRead(U, &u));
1595:   PetscCall(VecRestoreArray(G, &g));
1596:   PetscFunctionReturn(PETSC_SUCCESS);
1597: }

1599: /* J_{ij} = dF_i/dx_j
1600:    J_p = (  0   1)
1601:          (-w^2  0)
1602:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1603:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1604: */
1605: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1606: {
1607:   DM               sw;
1608:   const PetscReal *coords, *vel;
1609:   PetscInt         dim, d, Np, p, rStart;

1611:   PetscFunctionBeginUser;
1612:   PetscCall(TSGetDM(ts, &sw));
1613:   PetscCall(DMGetDimension(sw, &dim));
1614:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1615:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1616:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1617:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1618:   Np /= 2 * dim;
1619:   for (p = 0; p < Np; ++p) {
1620:     const PetscReal x0      = coords[p * dim + 0];
1621:     const PetscReal vy0     = vel[p * dim + 1];
1622:     const PetscReal omega   = vy0 / x0;
1623:     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};

1625:     for (d = 0; d < dim; ++d) {
1626:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1627:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1628:     }
1629:   }
1630:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1631:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1632:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1633:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1634:   PetscFunctionReturn(PETSC_SUCCESS);
1635: }

1637: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1638: {
1639:   AppCtx            *user = (AppCtx *)ctx;
1640:   DM                 sw;
1641:   const PetscScalar *v;
1642:   PetscScalar       *xres;
1643:   PetscInt           Np, p, d, dim;

1645:   PetscFunctionBeginUser;
1646:   PetscCall(TSGetDM(ts, &sw));
1647:   PetscCall(DMGetDimension(sw, &dim));
1648:   PetscCall(VecGetLocalSize(Xres, &Np));
1649:   PetscCall(VecGetArrayRead(V, &v));
1650:   PetscCall(VecGetArray(Xres, &xres));
1651:   Np /= dim;
1652:   for (p = 0; p < Np; ++p) {
1653:     for (d = 0; d < dim; ++d) {
1654:       xres[p * dim + d] = v[p * dim + d];
1655:       if (user->fake_1D && d > 0) xres[p * dim + d] = 0;
1656:     }
1657:   }
1658:   PetscCall(VecRestoreArrayRead(V, &v));
1659:   PetscCall(VecRestoreArray(Xres, &xres));
1660:   PetscFunctionReturn(PETSC_SUCCESS);
1661: }

1663: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1664: {
1665:   DM                 sw;
1666:   AppCtx            *user = (AppCtx *)ctx;
1667:   SNES               snes = ((AppCtx *)ctx)->snes;
1668:   const PetscScalar *x;
1669:   const PetscReal   *coords, *vel;
1670:   PetscReal         *E, m_p, q_p;
1671:   PetscScalar       *vres;
1672:   PetscInt           Np, p, dim, d;
1673:   Parameter         *param;

1675:   PetscFunctionBeginUser;
1676:   PetscCall(TSGetDM(ts, &sw));
1677:   PetscCall(DMGetDimension(sw, &dim));
1678:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1679:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1680:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1681:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1682:   m_p = user->masses[0] * param->m0;
1683:   q_p = user->charges[0] * param->q0;
1684:   PetscCall(VecGetLocalSize(Vres, &Np));
1685:   PetscCall(VecGetArrayRead(X, &x));
1686:   PetscCall(VecGetArray(Vres, &vres));
1687:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2");
1688:   PetscCall(ComputeFieldAtParticles(snes, sw, E));

1690:   Np /= dim;
1691:   for (p = 0; p < Np; ++p) {
1692:     for (d = 0; d < dim; ++d) {
1693:       vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1694:       if (user->fake_1D && d > 0) vres[p * dim + d] = 0.;
1695:     }
1696:   }
1697:   PetscCall(VecRestoreArrayRead(X, &x));
1698:   /*
1699:     Synchronized, ordered output for parallel/sequential test cases.
1700:     In the 1D (on the 2D mesh) case, every y component should be zero.
1701:   */
1702:   if (user->checkVRes) {
1703:     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1704:     PetscInt  step;

1706:     PetscCall(TSGetStepNumber(ts, &step));
1707:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
1708:     for (PetscInt p = 0; p < Np; ++p) {
1709:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
1710:       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]));
1711:     }
1712:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1713:   }
1714:   PetscCall(VecRestoreArray(Vres, &vres));
1715:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1716:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1717:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1718:   PetscFunctionReturn(PETSC_SUCCESS);
1719: }

1721: static PetscErrorCode CreateSolution(TS ts)
1722: {
1723:   DM       sw;
1724:   Vec      u;
1725:   PetscInt dim, Np;

1727:   PetscFunctionBegin;
1728:   PetscCall(TSGetDM(ts, &sw));
1729:   PetscCall(DMGetDimension(sw, &dim));
1730:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1731:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1732:   PetscCall(VecSetBlockSize(u, dim));
1733:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1734:   PetscCall(VecSetUp(u));
1735:   PetscCall(TSSetSolution(ts, u));
1736:   PetscCall(VecDestroy(&u));
1737:   PetscFunctionReturn(PETSC_SUCCESS);
1738: }

1740: static PetscErrorCode SetProblem(TS ts)
1741: {
1742:   AppCtx *user;
1743:   DM      sw;

1745:   PetscFunctionBegin;
1746:   PetscCall(TSGetDM(ts, &sw));
1747:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1748:   // Define unified system for (X, V)
1749:   {
1750:     Mat      J;
1751:     PetscInt dim, Np;

1753:     PetscCall(DMGetDimension(sw, &dim));
1754:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1755:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1756:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1757:     PetscCall(MatSetBlockSize(J, 2 * dim));
1758:     PetscCall(MatSetFromOptions(J));
1759:     PetscCall(MatSetUp(J));
1760:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1761:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1762:     PetscCall(MatDestroy(&J));
1763:   }
1764:   /* Define split system for X and V */
1765:   {
1766:     Vec             u;
1767:     IS              isx, isv, istmp;
1768:     const PetscInt *idx;
1769:     PetscInt        dim, Np, rstart;

1771:     PetscCall(TSGetSolution(ts, &u));
1772:     PetscCall(DMGetDimension(sw, &dim));
1773:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1774:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1775:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1776:     PetscCall(ISGetIndices(istmp, &idx));
1777:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1778:     PetscCall(ISRestoreIndices(istmp, &idx));
1779:     PetscCall(ISDestroy(&istmp));
1780:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1781:     PetscCall(ISGetIndices(istmp, &idx));
1782:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1783:     PetscCall(ISRestoreIndices(istmp, &idx));
1784:     PetscCall(ISDestroy(&istmp));
1785:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1786:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1787:     PetscCall(ISDestroy(&isx));
1788:     PetscCall(ISDestroy(&isv));
1789:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1790:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1791:   }
1792:   PetscFunctionReturn(PETSC_SUCCESS);
1793: }

1795: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1796: {
1797:   DM        sw;
1798:   Vec       u;
1799:   PetscReal t, maxt, dt;
1800:   PetscInt  n, maxn;

1802:   PetscFunctionBegin;
1803:   PetscCall(TSGetDM(ts, &sw));
1804:   PetscCall(TSGetTime(ts, &t));
1805:   PetscCall(TSGetMaxTime(ts, &maxt));
1806:   PetscCall(TSGetTimeStep(ts, &dt));
1807:   PetscCall(TSGetStepNumber(ts, &n));
1808:   PetscCall(TSGetMaxSteps(ts, &maxn));

1810:   PetscCall(TSReset(ts));
1811:   PetscCall(TSSetDM(ts, sw));
1812:   PetscCall(TSSetFromOptions(ts));
1813:   PetscCall(TSSetTime(ts, t));
1814:   PetscCall(TSSetMaxTime(ts, maxt));
1815:   PetscCall(TSSetTimeStep(ts, dt));
1816:   PetscCall(TSSetStepNumber(ts, n));
1817:   PetscCall(TSSetMaxSteps(ts, maxn));

1819:   PetscCall(CreateSolution(ts));
1820:   PetscCall(SetProblem(ts));
1821:   PetscCall(TSGetSolution(ts, &u));
1822:   PetscFunctionReturn(PETSC_SUCCESS);
1823: }

1825: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
1826: {
1827:   DM        sw, cdm;
1828:   PetscInt  Np;
1829:   PetscReal low[2], high[2];
1830:   AppCtx   *user = (AppCtx *)ctx;

1832:   sw = user->swarm;
1833:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
1834:   // Get the bounding box so we can equally space the particles
1835:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
1836:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1837:   // shift it by h/2 so nothing is initialized directly on a boundary
1838:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
1839:   x[1] = 0.;
1840:   return PETSC_SUCCESS;
1841: }

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

1846:   Input Parameters:
1847: + ts         - The TS
1848: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

1850:   Output Parameters:
1851: . u - The initialized solution vector

1853:   Level: advanced

1855: .seealso: InitializeSolve()
1856: */
1857: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1858: {
1859:   DM       sw;
1860:   Vec      u, gc, gv, gc0, gv0;
1861:   IS       isx, isv;
1862:   PetscInt dim;
1863:   AppCtx  *user;

1865:   PetscFunctionBeginUser;
1866:   PetscCall(TSGetDM(ts, &sw));
1867:   PetscCall(DMGetApplicationContext(sw, &user));
1868:   PetscCall(DMGetDimension(sw, &dim));
1869:   if (useInitial) {
1870:     PetscReal v0[2] = {1., 0.};
1871:     if (user->perturbed_weights) {
1872:       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1873:     } else {
1874:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
1875:       PetscCall(DMSwarmInitializeCoordinates(sw));
1876:       if (user->fake_1D) {
1877:         PetscCall(InitializeVelocities_Fake1D(sw, user));
1878:       } else {
1879:         PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
1880:       }
1881:     }
1882:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1883:     PetscCall(DMSwarmTSRedistribute(ts));
1884:   }
1885:   PetscCall(TSGetSolution(ts, &u));
1886:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1887:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1888:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1889:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
1890:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1891:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
1892:   if (useInitial) {
1893:     PetscCall(VecCopy(gc, gc0));
1894:     PetscCall(VecCopy(gv, gv0));
1895:   }
1896:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1897:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1898:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1899:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
1900:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1901:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
1902:   PetscFunctionReturn(PETSC_SUCCESS);
1903: }

1905: static PetscErrorCode InitializeSolve(TS ts, Vec u)
1906: {
1907:   PetscFunctionBegin;
1908:   PetscCall(TSSetSolution(ts, u));
1909:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1910:   PetscFunctionReturn(PETSC_SUCCESS);
1911: }

1913: static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
1914: {
1915:   MPI_Comm           comm;
1916:   DM                 sw;
1917:   AppCtx            *user;
1918:   const PetscScalar *u;
1919:   const PetscReal   *coords, *vel;
1920:   PetscScalar       *e;
1921:   PetscReal          t;
1922:   PetscInt           dim, Np, p;

1924:   PetscFunctionBeginUser;
1925:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
1926:   PetscCall(TSGetDM(ts, &sw));
1927:   PetscCall(DMGetApplicationContext(sw, &user));
1928:   PetscCall(DMGetDimension(sw, &dim));
1929:   PetscCall(TSGetSolveTime(ts, &t));
1930:   PetscCall(VecGetArray(E, &e));
1931:   PetscCall(VecGetArrayRead(U, &u));
1932:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1933:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1934:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1935:   Np /= 2 * dim;
1936:   for (p = 0; p < Np; ++p) {
1937:     /* TODO generalize initial conditions and project into plane instead of assuming x-y */
1938:     const PetscReal    r0    = DMPlex_NormD_Internal(dim, &coords[p * dim]);
1939:     const PetscReal    th0   = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]);
1940:     const PetscReal    v0    = DMPlex_NormD_Internal(dim, &vel[p * dim]);
1941:     const PetscReal    omega = v0 / r0;
1942:     const PetscReal    ct    = PetscCosReal(omega * t + th0);
1943:     const PetscReal    st    = PetscSinReal(omega * t + th0);
1944:     const PetscScalar *x     = &u[(p * 2 + 0) * dim];
1945:     const PetscScalar *v     = &u[(p * 2 + 1) * dim];
1946:     const PetscReal    xe[3] = {r0 * ct, r0 * st, 0.0};
1947:     const PetscReal    ve[3] = {-v0 * st, v0 * ct, 0.0};
1948:     PetscInt           d;

1950:     for (d = 0; d < dim; ++d) {
1951:       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
1952:       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
1953:     }
1954:     if (user->error) {
1955:       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
1956:       const PetscReal exen = 0.5 * PetscSqr(v0);
1957:       PetscCall(PetscPrintf(comm, "t %.4g: p%" PetscInt_FMT " error [%.2g %.2g] sol [(%.6lf %.6lf) (%.6lf %.6lf)] exact [(%.6lf %.6lf) (%.6lf %.6lf)] energy/exact energy %g / %g (%.10lf%%)\n", (double)t, p, (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 0) * dim]), (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 1) * dim]), (double)x[0], (double)x[1], (double)v[0], (double)v[1], (double)xe[0], (double)xe[1], (double)ve[0], (double)ve[1], (double)en, (double)exen, (double)(PetscAbsReal(exen - en) * 100. / exen)));
1958:     }
1959:   }
1960:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1961:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1962:   PetscCall(VecRestoreArrayRead(U, &u));
1963:   PetscCall(VecRestoreArray(E, &e));
1964:   PetscFunctionReturn(PETSC_SUCCESS);
1965: }

1967: static PetscErrorCode MigrateParticles(TS ts)
1968: {
1969:   DM               sw, cdm;
1970:   const PetscReal *L;

1972:   PetscFunctionBeginUser;
1973:   PetscCall(TSGetDM(ts, &sw));
1974:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1975:   {
1976:     Vec        u, gc, gv, position, momentum;
1977:     IS         isx, isv;
1978:     PetscReal *pos, *mom;

1980:     PetscCall(TSGetSolution(ts, &u));
1981:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1982:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1983:     PetscCall(VecGetSubVector(u, isx, &position));
1984:     PetscCall(VecGetSubVector(u, isv, &momentum));
1985:     PetscCall(VecGetArray(position, &pos));
1986:     PetscCall(VecGetArray(momentum, &mom));
1987:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1988:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1989:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1990:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

1992:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
1993:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
1994:     if ((L[0] || L[1]) >= 0.) {
1995:       PetscReal *x, *v, upper[3], lower[3];
1996:       PetscInt   Np, dim;

1998:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
1999:       PetscCall(DMGetDimension(cdm, &dim));
2000:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
2001:       PetscCall(VecGetArray(gc, &x));
2002:       PetscCall(VecGetArray(gv, &v));
2003:       for (PetscInt p = 0; p < Np; ++p) {
2004:         for (PetscInt d = 0; d < dim; ++d) {
2005:           if (pos[p * dim + d] < lower[d]) {
2006:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
2007:           } else if (pos[p * dim + d] > upper[d]) {
2008:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
2009:           } else {
2010:             x[p * dim + d] = pos[p * dim + d];
2011:           }
2012:           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]);
2013:           v[p * dim + d] = mom[p * dim + d];
2014:         }
2015:       }
2016:       PetscCall(VecRestoreArray(gc, &x));
2017:       PetscCall(VecRestoreArray(gv, &v));
2018:     }
2019:     PetscCall(VecRestoreArray(position, &pos));
2020:     PetscCall(VecRestoreArray(momentum, &mom));
2021:     PetscCall(VecRestoreSubVector(u, isx, &position));
2022:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2023:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
2024:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2025:   }
2026:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2027:   PetscCall(DMSwarmTSRedistribute(ts));
2028:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2029:   PetscFunctionReturn(PETSC_SUCCESS);
2030: }

2032: int main(int argc, char **argv)
2033: {
2034:   DM        dm, sw;
2035:   TS        ts;
2036:   Vec       u;
2037:   PetscReal dt;
2038:   PetscInt  maxn;
2039:   AppCtx    user;

2041:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2042:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2043:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2044:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2045:   PetscCall(CreatePoisson(dm, &user));
2046:   PetscCall(CreateSwarm(dm, &user, &sw));
2047:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2048:   PetscCall(InitializeConstants(sw, &user));
2049:   PetscCall(DMSetApplicationContext(sw, &user));

2051:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2052:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2053:   PetscCall(TSSetDM(ts, sw));
2054:   PetscCall(TSSetMaxTime(ts, 0.1));
2055:   PetscCall(TSSetTimeStep(ts, 0.00001));
2056:   PetscCall(TSSetMaxSteps(ts, 100));
2057:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

2059:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2060:   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2061:   if (user.monitor_positions) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2062:   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));

2064:   PetscCall(TSSetFromOptions(ts));
2065:   PetscCall(TSGetTimeStep(ts, &dt));
2066:   PetscCall(TSGetMaxSteps(ts, &maxn));
2067:   user.steps    = maxn;
2068:   user.stepSize = dt;
2069:   PetscCall(SetupContext(dm, sw, &user));
2070:   PetscCall(DMSwarmVectorDefineField(sw, "velocity"));
2071:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2072:   PetscCall(TSSetComputeExactError(ts, ComputeError));
2073:   PetscCall(TSSetPostStep(ts, MigrateParticles));
2074:   PetscCall(CreateSolution(ts));
2075:   PetscCall(TSGetSolution(ts, &u));
2076:   PetscCall(TSComputeInitialCondition(ts, u));
2077:   PetscCall(CheckNonNegativeWeights(sw, &user));
2078:   PetscCall(TSSolve(ts, NULL));

2080:   PetscCall(SNESDestroy(&user.snes));
2081:   PetscCall(TSDestroy(&ts));
2082:   PetscCall(DMDestroy(&sw));
2083:   PetscCall(DMDestroy(&dm));
2084:   PetscCall(DestroyContext(&user));
2085:   PetscCall(PetscFinalize());
2086:   return 0;
2087: }

2089: /*TEST

2091:    build:
2092:     requires: !complex double

2094:   # This tests that we can put particles in a box and compute the Coulomb force
2095:   # Recommend -draw_size 500,500
2096:    testset:
2097:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2098:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 20,1 \
2099:              -dm_plex_box_lower 0,-1 -dm_plex_box_upper 12.5664,1 \
2100:              -dm_plex_box_bd periodic,none \
2101:            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2102:            -sigma 1.0e-8 -timeScale 2.0e-14 \
2103:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2104:              -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 0 \
2105:            -output_step 50 -check_vel_res
2106:      test:
2107:        suffix: none_1d
2108:        args: -em_type none -error
2109:      test:
2110:        suffix: coulomb_1d
2111:        args: -em_type coulomb

2113:    # for viewers
2114:    #-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
2115:    testset:
2116:      nsize: {{1 2}}
2117:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2118:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 36,1 \
2119:              -dm_plex_box_lower 0.,-0.5 -dm_plex_box_upper 12.5664,0.5 \
2120:              -dm_plex_box_bd periodic,none \
2121:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2122:              -vdm_plex_box_lower -3 -vdm_plex_box_upper 3 \
2123:            -dm_swarm_num_species 1 -dm_swarm_num_particles 360 \
2124:            -twostream -charges -1.,1. -sigma 1.0e-8 \
2125:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2126:            -ts_type basicsymplectic -ts_basicsymplectic_type 2 \
2127:              -ts_dt 0.01 -ts_max_time 5 -ts_max_steps 10 \
2128:            -em_snes_atol 1.e-15 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2129:            -output_step 1 -check_vel_res
2130:      test:
2131:        suffix: two_stream_c0
2132:        args: -em_type primal -petscfe_default_quadrature_order 2 -petscspace_degree 2 -em_pc_type svd
2133:      test:
2134:        suffix: two_stream_rt
2135:        requires: superlu_dist
2136:        args: -em_type mixed \
2137:                -potential_petscspace_degree 0 \
2138:                -potential_petscdualspace_lagrange_use_moments \
2139:                -potential_petscdualspace_lagrange_moment_order 2 \
2140:                -field_petscspace_degree 2 -field_petscfe_default_quadrature_order 1 \
2141:                -field_petscspace_type sum \
2142:                  -field_petscspace_variables 2 \
2143:                  -field_petscspace_components 2 \
2144:                  -field_petscspace_sum_spaces 2 \
2145:                  -field_petscspace_sum_concatenate true \
2146:                  -field_sumcomp_0_petscspace_variables 2 \
2147:                  -field_sumcomp_0_petscspace_type tensor \
2148:                  -field_sumcomp_0_petscspace_tensor_spaces 2 \
2149:                  -field_sumcomp_0_petscspace_tensor_uniform false \
2150:                  -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
2151:                  -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
2152:                  -field_sumcomp_1_petscspace_variables 2 \
2153:                  -field_sumcomp_1_petscspace_type tensor \
2154:                  -field_sumcomp_1_petscspace_tensor_spaces 2 \
2155:                  -field_sumcomp_1_petscspace_tensor_uniform false \
2156:                  -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
2157:                  -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
2158:                -field_petscdualspace_form_degree -1 \
2159:                -field_petscdualspace_order 1 \
2160:                -field_petscdualspace_lagrange_trimmed true \
2161:              -em_snes_error_if_not_converged \
2162:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2163:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2164:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2165:                -em_fieldsplit_field_pc_type lu \
2166:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2167:                -em_fieldsplit_potential_pc_type svd

2169:    # For verification, we use
2170:    # -dm_plex_box_faces 100,1 -vdm_plex_box_faces 8000 -dm_swarm_num_particles 800000
2171:    # -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
2172:    testset:
2173:      nsize: {{1 2}}
2174:      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
2175:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 10,1 \
2176:              -dm_plex_box_lower 0.,-0.5 -dm_plex_box_upper 12.5664,0.5 \
2177:              -dm_plex_box_bd periodic,none \
2178:            -vdm_plex_dim 1 -vdm_plex_simplex 0 -vdm_plex_box_faces 10 \
2179:              -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 \
2180:            -dm_swarm_num_species 1 -dm_swarm_num_particles 100 \
2181:            -charges -1.,1. \
2182:              -cosine_coefficients 0.01,0.5 -perturbed_weights -total_weight 1. \
2183:            -ts_type basicsymplectic -ts_basicsymplectic_type 1 \
2184:              -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 1 \
2185:            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged \
2186:            -output_step 1 -check_vel_res

2188:      test:
2189:        suffix: uniform_equilibrium_1d
2190:        args: -cosine_coefficients 0.0,0.5 -em_type primal -petscspace_degree 1 -em_pc_type svd
2191:      test:
2192:        suffix: uniform_primal_1d
2193:        args: -em_type primal -petscspace_degree 1 -em_pc_type svd
2194:      test:
2195:        requires: superlu_dist
2196:        suffix: uniform_mixed_1d
2197:        args: -em_type mixed \
2198:                -potential_petscspace_degree 0 \
2199:                -potential_petscdualspace_lagrange_use_moments \
2200:                -potential_petscdualspace_lagrange_moment_order 2 \
2201:                -field_petscspace_degree 2 -field_petscfe_default_quadrature_order 1 \
2202:                -field_petscspace_type sum \
2203:                  -field_petscspace_variables 2 \
2204:                  -field_petscspace_components 2 \
2205:                  -field_petscspace_sum_spaces 2 \
2206:                  -field_petscspace_sum_concatenate true \
2207:                  -field_sumcomp_0_petscspace_variables 2 \
2208:                  -field_sumcomp_0_petscspace_type tensor \
2209:                  -field_sumcomp_0_petscspace_tensor_spaces 2 \
2210:                  -field_sumcomp_0_petscspace_tensor_uniform false \
2211:                  -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
2212:                  -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
2213:                  -field_sumcomp_1_petscspace_variables 2 \
2214:                  -field_sumcomp_1_petscspace_type tensor \
2215:                  -field_sumcomp_1_petscspace_tensor_spaces 2 \
2216:                  -field_sumcomp_1_petscspace_tensor_uniform false \
2217:                  -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
2218:                  -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
2219:                -field_petscdualspace_form_degree -1 \
2220:                -field_petscdualspace_order 1 \
2221:                -field_petscdualspace_lagrange_trimmed true \
2222:              -em_snes_error_if_not_converged \
2223:              -em_ksp_type preonly -em_ksp_error_if_not_converged \
2224:              -em_pc_type fieldsplit -em_pc_fieldsplit_type schur \
2225:                -em_pc_fieldsplit_schur_fact_type full -em_pc_fieldsplit_schur_precondition full \
2226:                -em_fieldsplit_field_pc_type lu \
2227:                  -em_fieldsplit_field_pc_factor_mat_solver_type superlu_dist \
2228:                -em_fieldsplit_potential_pc_type svd

2230: TEST*/