Actual source code: ex9.c

  1: static char help[] = "Landau Damping 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: const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
 33: typedef enum {
 34:   EM_PRIMAL,
 35:   EM_MIXED,
 36:   EM_COULOMB,
 37:   EM_NONE
 38: } EMType;

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

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

106: } AppCtx;

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

151:   PetscOptionsBegin(comm, "", "Central Orbit Options", "DMSWARM");
152:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex9.c", options->error, &options->error, NULL));
153:   PetscCall(PetscOptionsBool("-monitor_efield", "Flag to show efield plot", "ex9.c", options->efield_monitor, &options->efield_monitor, NULL));
154:   PetscCall(PetscOptionsBool("-monitor_ics", "Flag to show initial condition histograms", "ex9.c", options->initial_monitor, &options->initial_monitor, NULL));
155:   PetscCall(PetscOptionsBool("-monitor_positions", "The flag to show particle positions", "ex9.c", options->monitor_positions, &options->monitor_positions, NULL));
156:   PetscCall(PetscOptionsBool("-monitor_poisson", "The flag to show charges, Efield and potential solve", "ex9.c", options->poisson_monitor, &options->poisson_monitor, NULL));
157:   PetscCall(PetscOptionsBool("-periodic", "Flag to use periodic particle boundaries", "ex9.c", options->periodic, &options->periodic, NULL));
158:   PetscCall(PetscOptionsBool("-fake_1D", "Flag to run a 1D simulation (but really in 2D)", "ex9.c", options->fake_1D, &options->fake_1D, NULL));
159:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", "ex9.c", options->perturbed_weights, &options->perturbed_weights, NULL));
160:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex9.c", options->ostep, &options->ostep, NULL));
161:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex9.c", options->timeScale, &options->timeScale, NULL));
162:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex9.c", options->initVel, &options->initVel, NULL));
163:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex9.c", options->totalWeight, &options->totalWeight, NULL));
164:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex9.c", options->cosine_coefficients, &d, NULL));
165:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex9.c", options->charges, &d, NULL));
166:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex9.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
167:   PetscOptionsEnd();
168:   PetscFunctionReturn(PETSC_SUCCESS);
169: }

171: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
172: {
173:   PetscFunctionBeginUser;
174:   if (user->efield_monitor) {
175:     PetscDrawAxis axis_ef;
176:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_efield", 0, 300, 400, 300, &user->drawef));
177:     PetscCall(PetscDrawSetSave(user->drawef, "ex9_Efield.png"));
178:     PetscCall(PetscDrawSetFromOptions(user->drawef));
179:     PetscCall(PetscDrawLGCreate(user->drawef, 1, &user->drawlg_ef));
180:     PetscCall(PetscDrawLGGetAxis(user->drawlg_ef, &axis_ef));
181:     PetscCall(PetscDrawAxisSetLabels(axis_ef, "Electron Electric Field", "time", "E_max"));
182:     PetscCall(PetscDrawLGSetLimits(user->drawlg_ef, 0., user->steps * user->stepSize, -10., 0.));
183:     PetscCall(PetscDrawAxisSetLimits(axis_ef, 0., user->steps * user->stepSize, -10., 0.));
184:   }
185:   if (user->initial_monitor) {
186:     PetscDrawAxis axis1, axis2, axis3;
187:     PetscReal     dmboxlower[2], dmboxupper[2];
188:     PetscInt      dim, cStart, cEnd;
189:     PetscCall(DMGetDimension(sw, &dim));
190:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
191:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

193:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &user->drawic_x));
194:     PetscCall(PetscDrawSetSave(user->drawic_x, "ex9_ic_x.png"));
195:     PetscCall(PetscDrawSetFromOptions(user->drawic_x));
196:     PetscCall(PetscDrawHGCreate(user->drawic_x, dim, &user->drawhgic_x));
197:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1));
198:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, cEnd - cStart));
199:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "counts"));
200:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 1500));

202:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &user->drawic_v));
203:     PetscCall(PetscDrawSetSave(user->drawic_v, "ex9_ic_v.png"));
204:     PetscCall(PetscDrawSetFromOptions(user->drawic_v));
205:     PetscCall(PetscDrawHGCreate(user->drawic_v, dim, &user->drawhgic_v));
206:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2));
207:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 1000));
208:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "counts"));
209:     PetscCall(PetscDrawAxisSetLimits(axis2, -1, 1, 0, 1500));

211:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "monitor_initial_conditions_w", 800, 300, 400, 300, &user->drawic_w));
212:     PetscCall(PetscDrawSetSave(user->drawic_w, "ex9_ic_w.png"));
213:     PetscCall(PetscDrawSetFromOptions(user->drawic_w));
214:     PetscCall(PetscDrawHGCreate(user->drawic_w, dim, &user->drawhgic_w));
215:     PetscCall(PetscDrawHGGetAxis(user->drawhgic_w, &axis3));
216:     PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_w, 10));
217:     PetscCall(PetscDrawAxisSetLabels(axis3, "Initial W Distribution", "weight", "counts"));
218:     PetscCall(PetscDrawAxisSetLimits(axis3, 0, 0.01, 0, 5000));
219:   }
220:   if (user->monitor_positions) {
221:     PetscDrawAxis axis;

223:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "position_monitor_species1", 0, 0, 400, 300, &user->positionDraw));
224:     PetscCall(PetscDrawSetFromOptions(user->positionDraw));
225:     PetscCall(PetscDrawSPCreate(user->positionDraw, 10, &user->positionDrawSP));
226:     PetscCall(PetscDrawSPSetDimension(user->positionDrawSP, 1));
227:     PetscCall(PetscDrawSPGetAxis(user->positionDrawSP, &axis));
228:     PetscCall(PetscDrawSPReset(user->positionDrawSP));
229:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
230:     PetscCall(PetscDrawSetSave(user->positionDraw, "ex9_pos.png"));
231:   }
232:   if (user->poisson_monitor) {
233:     PetscDrawAxis axis_E, axis_Rho, axis_Pot;

235:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "Efield_monitor", 0, 0, 400, 300, &user->EDraw));
236:     PetscCall(PetscDrawSetFromOptions(user->EDraw));
237:     PetscCall(PetscDrawSPCreate(user->EDraw, 10, &user->EDrawSP));
238:     PetscCall(PetscDrawSPSetDimension(user->EDrawSP, 1));
239:     PetscCall(PetscDrawSPGetAxis(user->EDrawSP, &axis_E));
240:     PetscCall(PetscDrawSPReset(user->EDrawSP));
241:     PetscCall(PetscDrawAxisSetLabels(axis_E, "Particles", "x", "E"));
242:     PetscCall(PetscDrawSetSave(user->EDraw, "ex9_E_spatial.png"));

244:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "rho_monitor", 0, 0, 400, 300, &user->RhoDraw));
245:     PetscCall(PetscDrawSetFromOptions(user->RhoDraw));
246:     PetscCall(PetscDrawSPCreate(user->RhoDraw, 10, &user->RhoDrawSP));
247:     PetscCall(PetscDrawSPSetDimension(user->RhoDrawSP, 1));
248:     PetscCall(PetscDrawSPGetAxis(user->RhoDrawSP, &axis_Rho));
249:     PetscCall(PetscDrawSPReset(user->RhoDrawSP));
250:     PetscCall(PetscDrawAxisSetLabels(axis_Rho, "Particles", "x", "rho"));
251:     PetscCall(PetscDrawSetSave(user->RhoDraw, "ex9_rho_spatial.png"));

253:     PetscCall(PetscDrawCreate(PETSC_COMM_WORLD, NULL, "potential_monitor", 0, 0, 400, 300, &user->PotDraw));
254:     PetscCall(PetscDrawSetFromOptions(user->PotDraw));
255:     PetscCall(PetscDrawSPCreate(user->PotDraw, 10, &user->PotDrawSP));
256:     PetscCall(PetscDrawSPSetDimension(user->PotDrawSP, 1));
257:     PetscCall(PetscDrawSPGetAxis(user->PotDrawSP, &axis_Pot));
258:     PetscCall(PetscDrawSPReset(user->PotDrawSP));
259:     PetscCall(PetscDrawAxisSetLabels(axis_Pot, "Particles", "x", "potential"));
260:     PetscCall(PetscDrawSetSave(user->PotDraw, "ex9_phi_spatial.png"));
261:   }
262:   PetscFunctionReturn(PETSC_SUCCESS);
263: }

265: static PetscErrorCode DestroyContext(AppCtx *user)
266: {
267:   PetscFunctionBeginUser;
268:   PetscCall(PetscDrawLGDestroy(&user->drawlg_ef));
269:   PetscCall(PetscDrawDestroy(&user->drawef));
270:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_x));
271:   PetscCall(PetscDrawDestroy(&user->drawic_x));
272:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_v));
273:   PetscCall(PetscDrawDestroy(&user->drawic_v));
274:   PetscCall(PetscDrawHGDestroy(&user->drawhgic_w));
275:   PetscCall(PetscDrawDestroy(&user->drawic_w));
276:   PetscCall(PetscDrawSPDestroy(&user->positionDrawSP));
277:   PetscCall(PetscDrawDestroy(&user->positionDraw));

279:   PetscCall(PetscDrawSPDestroy(&user->EDrawSP));
280:   PetscCall(PetscDrawDestroy(&user->EDraw));
281:   PetscCall(PetscDrawSPDestroy(&user->RhoDrawSP));
282:   PetscCall(PetscDrawDestroy(&user->RhoDraw));
283:   PetscCall(PetscDrawSPDestroy(&user->PotDrawSP));
284:   PetscCall(PetscDrawDestroy(&user->PotDraw));

286:   PetscCall(PetscBagDestroy(&user->bag));
287:   PetscFunctionReturn(PETSC_SUCCESS);
288: }

290: static PetscErrorCode computeParticleMoments(DM sw, PetscReal moments[3], AppCtx *user)
291: {
292:   DM                 dm;
293:   const PetscReal   *coords;
294:   const PetscScalar *w;
295:   PetscReal          mom[3] = {0.0, 0.0, 0.0};
296:   PetscInt           cell, cStart, cEnd, dim;

298:   PetscFunctionBeginUser;
299:   PetscCall(DMGetDimension(sw, &dim));
300:   PetscCall(DMSwarmGetCellDM(sw, &dm));
301:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
302:   PetscCall(DMSwarmSortGetAccess(sw));
303:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
304:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
305:   for (cell = cStart; cell < cEnd; ++cell) {
306:     PetscInt *pidx;
307:     PetscInt  Np, p, d;

309:     PetscCall(DMSwarmSortGetPointsPerCell(sw, cell, &Np, &pidx));
310:     for (p = 0; p < Np; ++p) {
311:       const PetscInt   idx = pidx[p];
312:       const PetscReal *c   = &coords[idx * dim];

314:       mom[0] += PetscRealPart(w[idx]);
315:       mom[1] += PetscRealPart(w[idx]) * c[0];
316:       for (d = 0; d < dim; ++d) mom[2] += PetscRealPart(w[idx]) * c[d] * c[d];
317:     }
318:     PetscCall(PetscFree(pidx));
319:   }
320:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
321:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
322:   PetscCall(DMSwarmSortRestoreAccess(sw));
323:   PetscCall(MPIU_Allreduce(mom, moments, 3, MPIU_REAL, MPI_SUM, PetscObjectComm((PetscObject)sw)));
324:   PetscFunctionReturn(PETSC_SUCCESS);
325: }

327: 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[])
328: {
329:   f0[0] = u[0];
330: }

332: 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[])
333: {
334:   f0[0] = x[0] * u[0];
335: }

337: 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[])
338: {
339:   PetscInt d;

341:   f0[0] = 0.0;
342:   for (d = 0; d < dim; ++d) f0[0] += PetscSqr(x[d]) * u[0];
343: }

345: static PetscErrorCode computeFEMMoments(DM dm, Vec u, PetscReal moments[3], AppCtx *user)
346: {
347:   PetscDS     prob;
348:   PetscScalar mom;
349:   PetscInt    field = 0;

351:   PetscFunctionBeginUser;
352:   PetscCall(DMGetDS(dm, &prob));
353:   PetscCall(PetscDSSetObjective(prob, field, &f0_1));
354:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
355:   moments[0] = PetscRealPart(mom);
356:   PetscCall(PetscDSSetObjective(prob, field, &f0_x));
357:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
358:   moments[1] = PetscRealPart(mom);
359:   PetscCall(PetscDSSetObjective(prob, field, &f0_r2));
360:   PetscCall(DMPlexComputeIntegralFEM(dm, u, &mom, user));
361:   moments[2] = PetscRealPart(mom);
362:   PetscFunctionReturn(PETSC_SUCCESS);
363: }

365: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
366: {
367:   AppCtx    *user = (AppCtx *)ctx;
368:   DM         dm, sw;
369:   PetscReal *E;
370:   PetscReal  Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., temp = 0., *weight, chargesum = 0.;
371:   PetscReal *x, *v;
372:   PetscInt  *species, dim, p, d, Np, cStart, cEnd;
373:   PetscReal  pmoments[3]; /* \int f, \int x f, \int r^2 f */
374:   PetscReal  fmoments[3]; /* \int \hat f, \int x \hat f, \int r^2 \hat f */
375:   Vec        rho;

377:   PetscFunctionBeginUser;
378:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
379:   PetscCall(TSGetDM(ts, &sw));
380:   PetscCall(DMSwarmGetCellDM(sw, &dm));
381:   PetscCall(DMGetDimension(sw, &dim));
382:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
383:   PetscCall(DMSwarmSortGetAccess(sw));
384:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
385:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
386:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
387:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
388:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
389:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

391:   for (p = 0; p < Np; ++p) {
392:     for (d = 0; d < 1; ++d) {
393:       temp = PetscAbsReal(E[p * dim + d]);
394:       if (temp > Emax) Emax = temp;
395:     }
396:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
397:     sum += E[p * dim];
398:     chargesum += user->charges[0] * weight[p];
399:   }
400:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
401:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : -16.;

403:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
404:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
405:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
406:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
407:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

409:   Parameter *param;
410:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
411:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "charges", &rho));
412:   if (user->em == EM_PRIMAL) {
413:     PetscCall(computeParticleMoments(sw, pmoments, user));
414:     PetscCall(computeFEMMoments(dm, rho, fmoments, user));
415:   } else if (user->em == EM_MIXED) {
416:     DM       potential_dm;
417:     IS       potential_IS;
418:     PetscInt fields = 1;
419:     PetscCall(DMCreateSubDM(dm, 1, &fields, &potential_IS, &potential_dm));

421:     PetscCall(computeParticleMoments(sw, pmoments, user));
422:     PetscCall(computeFEMMoments(potential_dm, rho, fmoments, user));
423:     PetscCall(DMDestroy(&potential_dm));
424:     PetscCall(ISDestroy(&potential_IS));
425:   }
426:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "charges", &rho));

428:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%+e\t%e\t%f\t%f\t%f\t%f\t%f\t%f\t%f\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]));
429:   PetscCall(PetscDrawLGAddPoint(user->drawlg_ef, &t, &lgEmax));
430:   PetscCall(PetscDrawLGDraw(user->drawlg_ef));
431:   PetscCall(PetscDrawSave(user->drawef));
432:   PetscFunctionReturn(PETSC_SUCCESS);
433: }

435: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
436: {
437:   AppCtx            *user = (AppCtx *)ctx;
438:   DM                 dm, sw;
439:   const PetscScalar *u;
440:   PetscReal         *weight, *pos, *vel;
441:   PetscInt           dim, p, Np, cStart, cEnd;

443:   PetscFunctionBegin;
444:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
445:   PetscCall(TSGetDM(ts, &sw));
446:   PetscCall(DMSwarmGetCellDM(sw, &dm));
447:   PetscCall(DMGetDimension(sw, &dim));
448:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
449:   PetscCall(DMSwarmSortGetAccess(sw));
450:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

452:   if (step == 0) {
453:     PetscCall(PetscDrawHGReset(user->drawhgic_x));
454:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &user->drawic_x));
455:     PetscCall(PetscDrawClear(user->drawic_x));
456:     PetscCall(PetscDrawFlush(user->drawic_x));

458:     PetscCall(PetscDrawHGReset(user->drawhgic_v));
459:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &user->drawic_v));
460:     PetscCall(PetscDrawClear(user->drawic_v));
461:     PetscCall(PetscDrawFlush(user->drawic_v));

463:     PetscCall(PetscDrawHGReset(user->drawhgic_w));
464:     PetscCall(PetscDrawHGGetDraw(user->drawhgic_w, &user->drawic_w));
465:     PetscCall(PetscDrawClear(user->drawic_w));
466:     PetscCall(PetscDrawFlush(user->drawic_w));

468:     PetscCall(VecGetArrayRead(U, &u));
469:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
470:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
471:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));

473:     PetscCall(VecGetLocalSize(U, &Np));
474:     Np /= dim * 2;
475:     for (p = 0; p < Np; ++p) {
476:       PetscCall(PetscDrawHGAddValue(user->drawhgic_x, pos[p * dim]));
477:       PetscCall(PetscDrawHGAddValue(user->drawhgic_v, vel[p * dim]));
478:       PetscCall(PetscDrawHGAddValue(user->drawhgic_w, weight[p]));
479:     }

481:     PetscCall(VecRestoreArrayRead(U, &u));
482:     PetscCall(PetscDrawHGDraw(user->drawhgic_x));
483:     PetscCall(PetscDrawHGSave(user->drawhgic_x));

485:     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
486:     PetscCall(PetscDrawHGSave(user->drawhgic_v));

488:     PetscCall(PetscDrawHGDraw(user->drawhgic_w));
489:     PetscCall(PetscDrawHGSave(user->drawhgic_w));

491:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
492:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
493:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
494:   }
495:   PetscFunctionReturn(PETSC_SUCCESS);
496: }

498: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
499: {
500:   AppCtx         *user = (AppCtx *)ctx;
501:   DM              dm, sw;
502:   PetscScalar    *x, *v, *weight;
503:   PetscReal       lower[3], upper[3], speed;
504:   const PetscInt *s;
505:   PetscInt        dim, cStart, cEnd, c;

507:   PetscFunctionBeginUser;
508:   if (step > 0 && step % user->ostep == 0) {
509:     PetscCall(TSGetDM(ts, &sw));
510:     PetscCall(DMSwarmGetCellDM(sw, &dm));
511:     PetscCall(DMGetDimension(dm, &dim));
512:     PetscCall(DMGetBoundingBox(dm, lower, upper));
513:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
514:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
515:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
516:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
517:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
518:     PetscCall(DMSwarmSortGetAccess(sw));
519:     PetscCall(PetscDrawSPReset(user->positionDrawSP));
520:     PetscCall(PetscDrawSPSetLimits(user->positionDrawSP, lower[0], upper[0], lower[1], upper[1]));
521:     PetscCall(PetscDrawSPSetLimits(user->positionDrawSP, lower[0], upper[0], -12, 12));
522:     for (c = 0; c < cEnd - cStart; ++c) {
523:       PetscInt *pidx, Npc, q;
524:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
525:       for (q = 0; q < Npc; ++q) {
526:         const PetscInt p = pidx[q];
527:         if (s[p] == 0) {
528:           speed = PetscSqrtReal(PetscSqr(v[p * dim]) + PetscSqr(v[p * dim + 1]));
529:           if (dim == 1 || user->fake_1D) {
530:             PetscCall(PetscDrawSPAddPointColorized(user->positionDrawSP, &x[p * dim], &x[p * dim + 1], &speed));
531:           } else {
532:             PetscCall(PetscDrawSPAddPointColorized(user->positionDrawSP, &x[p * dim], &v[p * dim], &speed));
533:           }
534:         } else if (s[p] == 1) {
535:           PetscCall(PetscDrawSPAddPoint(user->positionDrawSP, &x[p * dim], &v[p * dim]));
536:         }
537:       }
538:       PetscCall(PetscFree(pidx));
539:     }
540:     PetscCall(PetscDrawSPDraw(user->positionDrawSP, PETSC_TRUE));
541:     PetscCall(PetscDrawSave(user->positionDraw));
542:     PetscCall(DMSwarmSortRestoreAccess(sw));
543:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
544:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
545:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
546:     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
547:   }
548:   PetscFunctionReturn(PETSC_SUCCESS);
549: }

551: static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
552: {
553:   AppCtx      *user = (AppCtx *)ctx;
554:   DM           dm, sw;
555:   PetscScalar *x, *E, *weight, *pot, *charges;
556:   PetscReal    lower[3], upper[3], xval;
557:   PetscInt     dim, cStart, cEnd, c;

559:   PetscFunctionBeginUser;
560:   if (step > 0 && step % user->ostep == 0) {
561:     PetscCall(TSGetDM(ts, &sw));
562:     PetscCall(DMSwarmGetCellDM(sw, &dm));
563:     PetscCall(DMGetDimension(dm, &dim));
564:     PetscCall(DMGetBoundingBox(dm, lower, upper));
565:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

567:     PetscCall(PetscDrawSPReset(user->RhoDrawSP));
568:     PetscCall(PetscDrawSPReset(user->EDrawSP));
569:     PetscCall(PetscDrawSPReset(user->PotDrawSP));
570:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
571:     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
572:     PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
573:     PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
574:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

576:     PetscCall(DMSwarmSortGetAccess(sw));
577:     for (c = 0; c < cEnd - cStart; ++c) {
578:       PetscReal Esum = 0.0;
579:       PetscInt *pidx, Npc, q;
580:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
581:       for (q = 0; q < Npc; ++q) {
582:         const PetscInt p = pidx[q];
583:         Esum += E[p * dim];
584:       }
585:       xval = (c + 0.5) * ((upper - lower) / (cEnd - cStart));
586:       PetscCall(PetscDrawSPAddPoint(user->EDrawSP, &xval, &Esum));
587:       PetscCall(PetscFree(pidx));
588:     }
589:     for (c = 0; c < (cEnd - cStart); ++c) {
590:       xval = (c + 0.5) * ((upper - lower) / (cEnd - cStart));
591:       PetscCall(PetscDrawSPAddPoint(user->RhoDrawSP, &xval, &charges[c]));
592:       PetscCall(PetscDrawSPAddPoint(user->PotDrawSP, &xval, &pot[c]));
593:     }
594:     PetscCall(PetscDrawSPDraw(user->RhoDrawSP, PETSC_TRUE));
595:     PetscCall(PetscDrawSave(user->RhoDraw));
596:     PetscCall(PetscDrawSPDraw(user->EDrawSP, PETSC_TRUE));
597:     PetscCall(PetscDrawSave(user->EDraw));
598:     PetscCall(PetscDrawSPDraw(user->PotDrawSP, PETSC_TRUE));
599:     PetscCall(PetscDrawSave(user->PotDraw));
600:     PetscCall(DMSwarmSortRestoreAccess(sw));
601:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
602:     PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));
603:     PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));
604:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
605:     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
606:   }
607:   PetscFunctionReturn(PETSC_SUCCESS);
608: }

610: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
611: {
612:   PetscBag   bag;
613:   Parameter *p;

615:   PetscFunctionBeginUser;
616:   /* setup PETSc parameter bag */
617:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
618:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
619:   bag = ctx->bag;
620:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
621:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
622:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
623:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
624:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
625:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
626:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
627:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

629:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
630:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
631:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
632:   PetscCall(PetscBagSetFromOptions(bag));
633:   {
634:     PetscViewer       viewer;
635:     PetscViewerFormat format;
636:     PetscBool         flg;

638:     PetscCall(PetscOptionsGetViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
639:     if (flg) {
640:       PetscCall(PetscViewerPushFormat(viewer, format));
641:       PetscCall(PetscBagView(bag, viewer));
642:       PetscCall(PetscViewerFlush(viewer));
643:       PetscCall(PetscViewerPopFormat(viewer));
644:       PetscCall(PetscOptionsRestoreViewer(&viewer));
645:     }
646:   }
647:   PetscFunctionReturn(PETSC_SUCCESS);
648: }

650: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
651: {
652:   PetscFunctionBeginUser;
653:   PetscCall(DMCreate(comm, dm));
654:   PetscCall(DMSetType(*dm, DMPLEX));
655:   PetscCall(DMSetFromOptions(*dm));
656:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
657:   PetscFunctionReturn(PETSC_SUCCESS);
658: }

660: 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[])
661: {
662:   f0[0] = -constants[SIGMA];
663: }

665: 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[])
666: {
667:   PetscInt d;
668:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
669: }

671: 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[])
672: {
673:   PetscInt d;
674:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
675: }

677: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
678: {
679:   *u = 0.0;
680:   return PETSC_SUCCESS;
681: }

683: /*
684:    /  I   -grad\ / q \ = /0\
685:    \-div    0  / \phi/   \f/
686: */
687: 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[])
688: {
689:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
690: }

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

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

703: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
704: 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[])
705: {
706:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
707: }

709: 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[])
710: {
711:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
712: }

714: 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[])
715: {
716:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
717: }

719: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
720: {
721:   PetscFE   fephi, feq;
722:   PetscDS   ds;
723:   PetscBool simplex;
724:   PetscInt  dim;

726:   PetscFunctionBeginUser;
727:   PetscCall(DMGetDimension(dm, &dim));
728:   PetscCall(DMPlexIsSimplex(dm, &simplex));
729:   if (user->em == EM_MIXED) {
730:     DMLabel        label;
731:     const PetscInt id = 1;

733:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
734:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
735:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
736:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
737:     PetscCall(PetscFECopyQuadrature(feq, fephi));
738:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
739:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
740:     PetscCall(DMCreateDS(dm));
741:     PetscCall(PetscFEDestroy(&fephi));
742:     PetscCall(PetscFEDestroy(&feq));

744:     PetscCall(DMGetLabel(dm, "marker", &label));
745:     PetscCall(DMGetDS(dm, &ds));

747:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
748:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
749:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
750:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
751:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

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

755:   } else if (user->em == EM_PRIMAL) {
756:     MatNullSpace nullsp;
757:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
758:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
759:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
760:     PetscCall(DMCreateDS(dm));
761:     PetscCall(DMGetDS(dm, &ds));
762:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
763:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
764:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
765:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
766:     PetscCall(MatNullSpaceDestroy(&nullsp));
767:     PetscCall(PetscFEDestroy(&fephi));
768:   }
769:   PetscFunctionReturn(PETSC_SUCCESS);
770: }

772: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
773: {
774:   SNES         snes;
775:   Mat          J;
776:   MatNullSpace nullSpace;

778:   PetscFunctionBeginUser;
779:   PetscCall(CreateFEM(dm, user));
780:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
781:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
782:   PetscCall(SNESSetDM(snes, dm));
783:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
784:   PetscCall(SNESSetFromOptions(snes));

786:   PetscCall(DMCreateMatrix(dm, &J));
787:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
788:   PetscCall(MatSetNullSpace(J, nullSpace));
789:   PetscCall(MatNullSpaceDestroy(&nullSpace));
790:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
791:   PetscCall(MatDestroy(&J));
792:   user->snes = snes;
793:   PetscFunctionReturn(PETSC_SUCCESS);
794: }

796: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
797: {
798:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
799:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
800:   return PETSC_SUCCESS;
801: }
802: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
803: {
804:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
805:   return PETSC_SUCCESS;
806: }

808: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
809: {
810:   const PetscReal alpha = scale ? scale[0] : 0.0;
811:   const PetscReal k     = scale ? scale[1] : 1.;
812:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
813:   return PETSC_SUCCESS;
814: }

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

824: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
825: {
826:   DM           vdm, dm;
827:   PetscScalar *weight;
828:   PetscReal   *x, *v, vmin[3], vmax[3], gmin[3], gmax[3], xi0[3];
829:   PetscInt    *N, Ns, dim, *cellid, *species, Np, cStart, cEnd, Npc, n;
830:   PetscInt     p, q, s, c, d, cv;
831:   PetscBool    flg;
832:   PetscMPIInt  size, rank;
833:   Parameter   *param;

835:   PetscFunctionBegin;
836:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
837:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
838:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
839:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
840:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
841:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
842:   PetscCall(PetscCalloc1(Ns, &N));
843:   n = Ns;
844:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
845:   PetscOptionsEnd();

847:   Np = N[0];
848:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Np = %" PetscInt_FMT "\n", Np));
849:   PetscCall(DMGetDimension(sw, &dim));
850:   PetscCall(DMSwarmGetCellDM(sw, &dm));
851:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

853:   PetscCall(DMCreate(PETSC_COMM_WORLD, &vdm));
854:   PetscCall(DMSetType(vdm, DMPLEX));
855:   PetscCall(DMPlexSetOptionsPrefix(vdm, "v"));
856:   PetscCall(DMSetFromOptions(vdm));
857:   PetscCall(DMViewFromOptions(vdm, NULL, "-vdm_view"));

859:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
860:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
861:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
862:   Npc = Np / (cEnd - cStart);
863:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
864:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
865:     for (s = 0; s < Ns; ++s) {
866:       for (q = 0; q < Npc; ++q, ++p) cellid[p] = c;
867:     }
868:   }
869:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
870:   PetscCall(PetscFree(N));

872:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
873:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
874:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
875:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));

877:   PetscCall(DMSwarmSortGetAccess(sw));
878:   PetscInt vStart, vEnd;
879:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vStart, &vEnd));
880:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
881:   for (c = 0; c < cEnd - cStart; ++c) {
882:     const PetscInt cell = c + cStart;
883:     PetscInt      *pidx, Npc;
884:     PetscReal      centroid[3], volume;

886:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
887:     PetscCall(DMPlexComputeCellGeometryFVM(dm, cell, &volume, centroid, NULL));
888:     for (q = 0; q < Npc; ++q) {
889:       const PetscInt p = pidx[q];

891:       for (d = 0; d < dim; ++d) {
892:         x[p * dim + d] = centroid[d];
893:         v[p * dim + d] = vmin[0] + (q + 0.5) * (vmax[0] - vmin[0]) / Npc;
894:         if (user->fake_1D && d > 0) v[p * dim + d] = 0;
895:       }
896:     }
897:     PetscCall(PetscFree(pidx));
898:   }
899:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));

901:   /* Setup Quadrature for spatial and velocity weight calculations*/
902:   PetscQuadrature  quad_x;
903:   PetscInt         Nq_x;
904:   const PetscReal *wq_x, *xq_x;
905:   PetscReal       *xq_x_extended;
906:   PetscReal        weightsum = 0., totalcellweight = 0., *weight_x, *weight_v;
907:   PetscReal        scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};

909:   PetscCall(PetscCalloc2(cEnd - cStart, &weight_x, Np, &weight_v));
910:   if (user->fake_1D) PetscCall(PetscDTGaussTensorQuadrature(1, 1, 5, -1.0, 1.0, &quad_x));
911:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad_x));
912:   PetscCall(PetscQuadratureGetData(quad_x, NULL, NULL, &Nq_x, &xq_x, &wq_x));
913:   if (user->fake_1D) {
914:     PetscCall(PetscCalloc1(Nq_x * dim, &xq_x_extended));
915:     for (PetscInt i = 0; i < Nq_x; ++i) xq_x_extended[i * dim] = xq_x[i];
916:   }
917:   /* Integrate the density function to get the weights of particles in each cell */
918:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
919:   for (c = cStart; c < cEnd; ++c) {
920:     PetscReal          v0_x[3], J_x[9], invJ_x[9], detJ_x, xr_x[3], den_x;
921:     PetscInt          *pidx, Npc, q;
922:     PetscInt           Ncx;
923:     const PetscScalar *array_x;
924:     PetscScalar       *coords_x = NULL;
925:     PetscBool          isDGx;
926:     weight_x[c] = 0.;

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

936:       /*Transform quadrature points from real space to ideal real space (0, 2PI/k)*/
937:       if (user->fake_1D) {
938:         PetscCall(PetscPDFCosine1D(xr_x, scale, &den_x));
939:         detJ_x = J_x[0];
940:       } else PetscCall(PetscPDFCosine2D(xr_x, scale, &den_x));
941:       /*We have to transform the quadrature weights as well*/
942:       weight_x[c] += den_x * (wq_x[q] * detJ_x);
943:     }
944:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", c, (double)PetscRealPart(coords_x[0]), (double)PetscRealPart(coords_x[2]), (double)weight_x[c]));
945:     totalcellweight += weight_x[c];
946:     PetscCheck(Npc / size == 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);

948:     /* Set weights to be gaussian in velocity cells (using exact solution) */
949:     for (cv = 0; cv < vEnd - vStart; ++cv) {
950:       PetscInt           Nc;
951:       const PetscScalar *array_v;
952:       PetscScalar       *coords_v = NULL;
953:       PetscBool          isDG;
954:       PetscCall(DMPlexGetCellCoordinates(vdm, cv, &isDG, &Nc, &array_v, &coords_v));

956:       const PetscInt p = pidx[cv];

958:       weight_v[p] = 0.5 * (PetscErfReal(coords_v[1] / PetscSqrtReal(2.)) - PetscErfReal(coords_v[0] / PetscSqrtReal(2.)));

960:       weight[p] = user->totalWeight * weight_v[p] * weight_x[c];
961:       weightsum += weight[p];

963:       PetscCall(DMPlexRestoreCellCoordinates(vdm, cv, &isDG, &Nc, &array_v, &coords_v));
964:     }
965:     PetscCall(DMPlexRestoreCellCoordinates(dm, c, &isDGx, &Ncx, &array_x, &coords_x));
966:     PetscCall(PetscFree(pidx));
967:   }
968:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)totalcellweight, (double)weightsum));
969:   if (user->fake_1D) PetscCall(PetscFree(xq_x_extended));
970:   PetscCall(PetscFree2(weight_x, weight_v));
971:   PetscCall(PetscQuadratureDestroy(&quad_x));
972:   PetscCall(DMSwarmSortRestoreAccess(sw));
973:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
974:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
975:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
976:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
977:   PetscCall(DMDestroy(&vdm));
978:   PetscFunctionReturn(PETSC_SUCCESS);
979: }

981: static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
982: {
983:   DM         dm;
984:   PetscInt  *species;
985:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3];
986:   PetscInt   Np, p, dim;

988:   PetscFunctionBegin;
989:   PetscCall(DMSwarmGetCellDM(sw, &dm));
990:   PetscCall(DMGetDimension(sw, &dim));
991:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
992:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
993:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
994:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
995:   for (p = 0; p < Np; ++p) {
996:     totalWeight += weight[p];
997:     totalCharge += user->charges[species[p]] * weight[p];
998:   }
999:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1000:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1001:   {
1002:     Parameter *param;
1003:     PetscReal  Area;

1005:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1006:     switch (dim) {
1007:     case 1:
1008:       Area = (gmax[0] - gmin[0]);
1009:       break;
1010:     case 2:
1011:       if (user->fake_1D) {
1012:         Area = (gmax[0] - gmin[0]);
1013:       } else {
1014:         Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1015:       }
1016:       break;
1017:     case 3:
1018:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1019:       break;
1020:     default:
1021:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1022:     }
1023:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[p]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)totalWeight, (double)user->charges[0], (double)totalCharge, (double)Area));
1024:     param->sigma = PetscAbsReal(totalCharge / (Area));

1026:     PetscCall(PetscPrintf(PETSC_COMM_SELF, "sigma: %g\n", (double)param->sigma));
1027:     PetscCall(PetscPrintf(PETSC_COMM_SELF, "(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,
1028:                           (double)param->vlasovNumber));
1029:   }
1030:   /* Setup Constants */
1031:   {
1032:     PetscDS    ds;
1033:     Parameter *param;
1034:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1035:     PetscScalar constants[NUM_CONSTANTS];
1036:     constants[SIGMA]   = param->sigma;
1037:     constants[V0]      = param->v0;
1038:     constants[T0]      = param->t0;
1039:     constants[X0]      = param->x0;
1040:     constants[M0]      = param->m0;
1041:     constants[Q0]      = param->q0;
1042:     constants[PHI0]    = param->phi0;
1043:     constants[POISSON] = param->poissonNumber;
1044:     constants[VLASOV]  = param->vlasovNumber;
1045:     PetscCall(DMGetDS(dm, &ds));
1046:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
1047:   }
1048:   PetscFunctionReturn(PETSC_SUCCESS);
1049: }

1051: static PetscErrorCode InitializeVelocites_Fake1D(DM sw, AppCtx *user)
1052: {
1053:   DM         dm;
1054:   PetscReal *v;
1055:   PetscInt  *species, cStart, cEnd;
1056:   PetscInt   dim, Np, p;

1058:   PetscFunctionBegin;
1059:   PetscCall(DMGetDimension(sw, &dim));
1060:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1061:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1062:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1063:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1064:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1065:   PetscRandom rnd;
1066:   PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1067:   PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1068:   PetscCall(PetscRandomSetFromOptions(rnd));

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

1073:     PetscCall(PetscRandomGetValueReal(rnd, &a[0]));
1074:     if (user->perturbed_weights) {
1075:       PetscCall(PetscPDFSampleConstant1D(a, NULL, vel));
1076:     } else {
1077:       PetscCall(PetscPDFSampleGaussian1D(a, NULL, vel));
1078:     }
1079:     v[p * dim] = vel[0];
1080:   }
1081:   PetscCall(PetscRandomDestroy(&rnd));
1082:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1083:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1084:   PetscFunctionReturn(PETSC_SUCCESS);
1085: }

1087: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
1088: {
1089:   PetscReal v0[2] = {1., 0.};
1090:   PetscInt  dim;

1092:   PetscFunctionBeginUser;
1093:   PetscCall(DMGetDimension(dm, &dim));
1094:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
1095:   PetscCall(DMSetType(*sw, DMSWARM));
1096:   PetscCall(DMSetDimension(*sw, dim));
1097:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
1098:   PetscCall(DMSwarmSetCellDM(*sw, dm));
1099:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
1100:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
1101:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
1102:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initCoordinates", dim, PETSC_REAL));
1103:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initVelocity", dim, PETSC_REAL));
1104:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
1105:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "potential", dim, PETSC_REAL));
1106:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "charges", dim, PETSC_REAL));
1107:   PetscCall(DMSwarmFinalizeFieldRegister(*sw));

1109:   if (user->perturbed_weights) {
1110:     PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
1111:   } else {
1112:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
1113:     PetscCall(DMSwarmInitializeCoordinates(*sw));
1114:     if (user->fake_1D) {
1115:       PetscCall(InitializeVelocites_Fake1D(*sw, user));
1116:     } else {
1117:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
1118:     }
1119:   }
1120:   PetscCall(DMSetFromOptions(*sw));
1121:   PetscCall(DMSetApplicationContext(*sw, user));
1122:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
1123:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
1124:   {
1125:     Vec gc, gc0, gv, gv0;

1127:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1128:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1129:     PetscCall(VecCopy(gc, gc0));
1130:     PetscCall(VecViewFromOptions(gc, NULL, "-ic_x_view"));
1131:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
1132:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0));
1133:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv));
1134:     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0));
1135:     PetscCall(VecCopy(gv, gv0));
1136:     PetscCall(VecViewFromOptions(gv, NULL, "-ic_v_view"));
1137:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv));
1138:     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0));
1139:   }
1140:   PetscFunctionReturn(PETSC_SUCCESS);
1141: }

1143: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
1144: {
1145:   AppCtx     *user;
1146:   PetscReal  *coords;
1147:   PetscInt   *species, dim, d, Np, p, q, Ns;
1148:   PetscMPIInt size;

1150:   PetscFunctionBegin;
1151:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
1152:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
1153:   PetscCall(DMGetDimension(sw, &dim));
1154:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1155:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1156:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1158:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1159:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1160:   for (p = 0; p < Np; ++p) {
1161:     PetscReal *pcoord = &coords[p * dim];
1162:     PetscReal  pE[3]  = {0., 0., 0.};

1164:     /* Calculate field at particle p due to particle q */
1165:     for (q = 0; q < Np; ++q) {
1166:       PetscReal *qcoord = &coords[q * dim];
1167:       PetscReal  rpq[3], r, r3, q_q;

1169:       if (p == q) continue;
1170:       q_q = user->charges[species[q]] * 1.;
1171:       for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1172:       r = DMPlex_NormD_Internal(dim, rpq);
1173:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
1174:       r3 = PetscPowRealInt(r, 3);
1175:       for (d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
1176:     }
1177:     for (d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
1178:   }
1179:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1180:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1181:   PetscFunctionReturn(PETSC_SUCCESS);
1182: }

1184: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[])
1185: {
1186:   DM              dm;
1187:   AppCtx         *user;
1188:   PetscDS         ds;
1189:   PetscFE         fe;
1190:   Mat             M_p, M;
1191:   Vec             phi, locPhi, rho, f;
1192:   PetscReal      *coords;
1193:   PetscInt        dim, d, cStart, cEnd, c, Np;
1194:   PetscQuadrature q;

1196:   PetscFunctionBegin;
1197:   PetscCall(DMGetDimension(sw, &dim));
1198:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1199:   PetscCall(DMGetApplicationContext(sw, (void *)&user));

1201:   KSP         ksp;
1202:   Vec         rho0;
1203:   char        oldField[PETSC_MAX_PATH_LEN];
1204:   const char *tmp;

1206:   /* Create the charges rho */
1207:   PetscCall(SNESGetDM(snes, &dm));
1208:   PetscCall(DMSwarmVectorGetField(sw, &tmp));
1209:   PetscCall(PetscStrncpy(oldField, tmp, PETSC_MAX_PATH_LEN));
1210:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1211:   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
1212:   PetscCall(DMSwarmVectorDefineField(sw, oldField));

1214:   PetscCall(DMCreateMassMatrix(dm, dm, &M));
1215:   PetscCall(DMGetGlobalVector(dm, &rho0));
1216:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Primal Compute"));
1217:   PetscCall(DMGetGlobalVector(dm, &rho));
1218:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
1219:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));

1221:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1222:   PetscCall(MatMultTranspose(M_p, f, rho));
1223:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1224:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1225:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1226:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1228:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1229:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1230:   PetscCall(KSPSetOperators(ksp, M, M));
1231:   PetscCall(KSPSetFromOptions(ksp));
1232:   PetscCall(KSPSolve(ksp, rho, rho0));
1233:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1235:   PetscInt           rhosize;
1236:   PetscReal         *charges;
1237:   const PetscScalar *rho_vals;
1238:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1239:   PetscCall(VecGetSize(rho0, &rhosize));
1240:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1241:   for (c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1242:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1243:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

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

1247:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1248:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1249:   PetscCall(DMRestoreGlobalVector(dm, &rho0));
1250:   PetscCall(KSPDestroy(&ksp));
1251:   PetscCall(MatDestroy(&M_p));
1252:   PetscCall(MatDestroy(&M));

1254:   PetscCall(DMGetGlobalVector(dm, &phi));
1255:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1256:   PetscCall(VecSet(phi, 0.0));
1257:   PetscCall(SNESSolve(snes, rho, phi));
1258:   PetscCall(DMRestoreGlobalVector(dm, &rho));
1259:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1261:   PetscInt           phisize;
1262:   PetscReal         *pot;
1263:   const PetscScalar *phi_vals;
1264:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1265:   PetscCall(VecGetSize(phi, &phisize));
1266:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1267:   for (c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1268:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1269:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1271:   PetscCall(DMGetLocalVector(dm, &locPhi));
1272:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1273:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1274:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1276:   PetscCall(DMGetDS(dm, &ds));
1277:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1278:   PetscCall(DMSwarmSortGetAccess(sw));
1279:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1280:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1282:   for (c = cStart; c < cEnd; ++c) {
1283:     PetscTabulation tab;
1284:     PetscScalar    *clPhi = NULL;
1285:     PetscReal      *pcoord, *refcoord;
1286:     PetscReal       v[3], J[9], invJ[9], detJ;
1287:     PetscInt       *points;
1288:     PetscInt        Ncp, cp;

1290:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1291:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1292:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1293:     for (cp = 0; cp < Ncp; ++cp)
1294:       for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1295:     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
1296:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1297:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ));
1298:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1299:     for (cp = 0; cp < Ncp; ++cp) {
1300:       const PetscReal *basisDer = tab->T[1];
1301:       const PetscInt   p        = points[cp];

1303:       for (d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1304:       PetscCall(PetscFEGetQuadrature(fe, &q));
1305:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, invJ, NULL, cp, &E[p * dim]));
1306:       for (d = 0; d < dim; ++d) {
1307:         E[p * dim + d] *= -1.0;
1308:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1309:       }
1310:     }
1311:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1312:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1313:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1314:     PetscCall(PetscTabulationDestroy(&tab));
1315:     PetscCall(PetscFree(points));
1316:   }
1317:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1318:   PetscCall(DMSwarmSortRestoreAccess(sw));
1319:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1320:   PetscFunctionReturn(PETSC_SUCCESS);
1321: }

1323: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
1324: {
1325:   AppCtx         *user;
1326:   DM              dm, potential_dm;
1327:   KSP             ksp;
1328:   IS              potential_IS;
1329:   PetscDS         ds;
1330:   PetscFE         fe;
1331:   PetscFEGeom     feGeometry;
1332:   Mat             M_p, M;
1333:   Vec             phi, locPhi, rho, f, temp_rho, rho0;
1334:   PetscQuadrature q;
1335:   PetscReal      *coords, *pot;
1336:   PetscInt        dim, d, cStart, cEnd, c, Np, fields = 1;
1337:   char            oldField[PETSC_MAX_PATH_LEN];
1338:   const char     *tmp;

1340:   PetscFunctionBegin;
1341:   PetscCall(DMGetDimension(sw, &dim));
1342:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1343:   PetscCall(DMGetApplicationContext(sw, &user));

1345:   /* Create the charges rho */
1346:   PetscCall(SNESGetDM(snes, &dm));
1347:   PetscCall(DMGetGlobalVector(dm, &rho));
1348:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));

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

1352:   PetscCall(DMSwarmVectorGetField(sw, &tmp));
1353:   PetscCall(PetscStrncpy(oldField, tmp, PETSC_MAX_PATH_LEN));
1354:   PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
1355:   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
1356:   PetscCall(DMSwarmVectorDefineField(sw, oldField));

1358:   PetscCall(DMCreateMassMatrix(potential_dm, potential_dm, &M));
1359:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1360:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1361:   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
1362:   PetscCall(PetscObjectSetName((PetscObject)temp_rho, "Mf"));
1363:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1364:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1365:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1366:   PetscCall(MatMultTranspose(M_p, f, temp_rho));
1367:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1368:   PetscCall(DMGetGlobalVector(potential_dm, &rho0));
1369:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Mixed Compute"));

1371:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1372:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1373:   PetscCall(KSPSetOperators(ksp, M, M));
1374:   PetscCall(KSPSetFromOptions(ksp));
1375:   PetscCall(KSPSolve(ksp, temp_rho, rho0));
1376:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1378:   PetscInt           rhosize;
1379:   PetscReal         *charges;
1380:   const PetscScalar *rho_vals;
1381:   Parameter         *param;
1382:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1383:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1384:   PetscCall(VecGetSize(rho0, &rhosize));

1386:   /* Integral over reference element is size 1.  Reference element area is 4.  Scale rho0 by 1/4 because the basis function is 1/4 */
1387:   PetscCall(VecScale(rho0, 0.25));
1388:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1389:   for (c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1390:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1391:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

1393:   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
1394:   PetscCall(VecScale(rho, 0.25));
1395:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1396:   PetscCall(VecViewFromOptions(temp_rho, NULL, "-temprho_view"));
1397:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1398:   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
1399:   PetscCall(DMRestoreGlobalVector(potential_dm, &rho0));

1401:   PetscCall(MatDestroy(&M_p));
1402:   PetscCall(MatDestroy(&M));
1403:   PetscCall(KSPDestroy(&ksp));
1404:   PetscCall(DMDestroy(&potential_dm));
1405:   PetscCall(ISDestroy(&potential_IS));

1407:   PetscCall(DMGetGlobalVector(dm, &phi));
1408:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1409:   PetscCall(VecSet(phi, 0.0));
1410:   PetscCall(SNESSolve(snes, rho, phi));
1411:   PetscCall(DMRestoreGlobalVector(dm, &rho));

1413:   PetscInt           phisize;
1414:   const PetscScalar *phi_vals;
1415:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1416:   PetscCall(VecGetSize(phi, &phisize));
1417:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1418:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1419:   for (c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1420:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1421:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1423:   PetscCall(DMGetLocalVector(dm, &locPhi));
1424:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1425:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1426:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1428:   PetscCall(DMGetDS(dm, &ds));
1429:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1430:   PetscCall(DMSwarmSortGetAccess(sw));
1431:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1432:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1433:   PetscCall(PetscFEGetQuadrature(fe, &q));
1434:   PetscCall(PetscFECreateCellGeometry(fe, q, &feGeometry));
1435:   for (c = cStart; c < cEnd; ++c) {
1436:     PetscTabulation tab;
1437:     PetscScalar    *clPhi = NULL;
1438:     PetscReal      *pcoord, *refcoord;
1439:     PetscInt       *points;
1440:     PetscInt        Ncp, cp;

1442:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1443:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1444:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1445:     for (cp = 0; cp < Ncp; ++cp)
1446:       for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1447:     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
1448:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1449:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, q, feGeometry.v, feGeometry.J, feGeometry.invJ, feGeometry.detJ));
1450:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));

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

1455:       for (d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1456:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, &feGeometry, cp, &E[p * dim]));
1457:       PetscCall(PetscFEPushforward(fe, &feGeometry, 1, &E[p * dim]));
1458:       for (d = 0; d < dim; ++d) {
1459:         E[p * dim + d] *= -2.0;
1460:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1461:       }
1462:     }
1463:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1464:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1465:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1466:     PetscCall(PetscTabulationDestroy(&tab));
1467:     PetscCall(PetscFree(points));
1468:   }
1469:   PetscCall(PetscFEDestroyCellGeometry(fe, &feGeometry));
1470:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1471:   PetscCall(DMSwarmSortRestoreAccess(sw));
1472:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1473:   PetscFunctionReturn(PETSC_SUCCESS);
1474: }

1476: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
1477: {
1478:   AppCtx  *ctx;
1479:   PetscInt dim, Np;

1481:   PetscFunctionBegin;
1484:   PetscAssertPointer(E, 3);
1485:   PetscCall(DMGetDimension(sw, &dim));
1486:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1487:   PetscCall(DMGetApplicationContext(sw, &ctx));
1488:   PetscCall(PetscArrayzero(E, Np * dim));

1490:   switch (ctx->em) {
1491:   case EM_PRIMAL:
1492:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
1493:     break;
1494:   case EM_COULOMB:
1495:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1496:     break;
1497:   case EM_MIXED:
1498:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
1499:     break;
1500:   case EM_NONE:
1501:     break;
1502:   default:
1503:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
1504:   }
1505:   PetscFunctionReturn(PETSC_SUCCESS);
1506: }

1508: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1509: {
1510:   DM                 sw;
1511:   SNES               snes = ((AppCtx *)ctx)->snes;
1512:   const PetscReal   *coords, *vel;
1513:   const PetscScalar *u;
1514:   PetscScalar       *g;
1515:   PetscReal         *E, m_p = 1., q_p = -1.;
1516:   PetscInt           dim, d, Np, p;

1518:   PetscFunctionBeginUser;
1519:   PetscCall(TSGetDM(ts, &sw));
1520:   PetscCall(DMGetDimension(sw, &dim));
1521:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1522:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1523:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1524:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1525:   PetscCall(VecGetArrayRead(U, &u));
1526:   PetscCall(VecGetArray(G, &g));

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

1530:   Np /= 2 * dim;
1531:   for (p = 0; p < Np; ++p) {
1532:     for (d = 0; d < dim; ++d) {
1533:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1534:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1535:     }
1536:   }
1537:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1538:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1539:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1540:   PetscCall(VecRestoreArrayRead(U, &u));
1541:   PetscCall(VecRestoreArray(G, &g));
1542:   PetscFunctionReturn(PETSC_SUCCESS);
1543: }

1545: /* J_{ij} = dF_i/dx_j
1546:    J_p = (  0   1)
1547:          (-w^2  0)
1548:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1549:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1550: */
1551: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1552: {
1553:   DM               sw;
1554:   const PetscReal *coords, *vel;
1555:   PetscInt         dim, d, Np, p, rStart;

1557:   PetscFunctionBeginUser;
1558:   PetscCall(TSGetDM(ts, &sw));
1559:   PetscCall(DMGetDimension(sw, &dim));
1560:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1561:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1562:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1563:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1564:   Np /= 2 * dim;
1565:   for (p = 0; p < Np; ++p) {
1566:     const PetscReal x0      = coords[p * dim + 0];
1567:     const PetscReal vy0     = vel[p * dim + 1];
1568:     const PetscReal omega   = vy0 / x0;
1569:     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};

1571:     for (d = 0; d < dim; ++d) {
1572:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1573:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1574:     }
1575:   }
1576:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1577:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1578:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1579:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1580:   PetscFunctionReturn(PETSC_SUCCESS);
1581: }

1583: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1584: {
1585:   AppCtx            *user = (AppCtx *)ctx;
1586:   DM                 sw;
1587:   const PetscScalar *v;
1588:   PetscScalar       *xres;
1589:   PetscInt           Np, p, d, dim;

1591:   PetscFunctionBeginUser;
1592:   PetscCall(TSGetDM(ts, &sw));
1593:   PetscCall(DMGetDimension(sw, &dim));
1594:   PetscCall(VecGetLocalSize(Xres, &Np));
1595:   PetscCall(VecGetArrayRead(V, &v));
1596:   PetscCall(VecGetArray(Xres, &xres));
1597:   Np /= dim;
1598:   for (p = 0; p < Np; ++p) {
1599:     for (d = 0; d < dim; ++d) {
1600:       xres[p * dim + d] = v[p * dim + d];
1601:       if (user->fake_1D && d > 0) xres[p * dim + d] = 0;
1602:     }
1603:   }
1604:   PetscCall(VecRestoreArrayRead(V, &v));
1605:   PetscCall(VecRestoreArray(Xres, &xres));
1606:   PetscFunctionReturn(PETSC_SUCCESS);
1607: }

1609: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1610: {
1611:   DM                 sw;
1612:   AppCtx            *user = (AppCtx *)ctx;
1613:   SNES               snes = ((AppCtx *)ctx)->snes;
1614:   const PetscScalar *x;
1615:   const PetscReal   *coords, *vel;
1616:   PetscReal         *E, m_p, q_p;
1617:   PetscScalar       *vres;
1618:   PetscInt           Np, p, dim, d;
1619:   Parameter         *param;

1621:   PetscFunctionBeginUser;
1622:   PetscCall(TSGetDM(ts, &sw));
1623:   PetscCall(DMGetDimension(sw, &dim));
1624:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1625:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1626:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1627:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1628:   m_p = user->masses[0] * param->m0;
1629:   q_p = user->charges[0] * param->q0;
1630:   PetscCall(VecGetLocalSize(Vres, &Np));
1631:   PetscCall(VecGetArrayRead(X, &x));
1632:   PetscCall(VecGetArray(Vres, &vres));
1633:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2");
1634:   PetscCall(ComputeFieldAtParticles(snes, sw, E));

1636:   Np /= dim;
1637:   for (p = 0; p < Np; ++p) {
1638:     for (d = 0; d < dim; ++d) {
1639:       vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1640:       if (user->fake_1D && d > 0) vres[p * dim + d] = 0.;
1641:     }
1642:   }
1643:   PetscCall(VecRestoreArrayRead(X, &x));
1644:   PetscCall(VecRestoreArray(Vres, &vres));
1645:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1646:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1647:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1648:   PetscFunctionReturn(PETSC_SUCCESS);
1649: }

1651: static PetscErrorCode CreateSolution(TS ts)
1652: {
1653:   DM       sw;
1654:   Vec      u;
1655:   PetscInt dim, Np;

1657:   PetscFunctionBegin;
1658:   PetscCall(TSGetDM(ts, &sw));
1659:   PetscCall(DMGetDimension(sw, &dim));
1660:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1661:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1662:   PetscCall(VecSetBlockSize(u, dim));
1663:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1664:   PetscCall(VecSetUp(u));
1665:   PetscCall(TSSetSolution(ts, u));
1666:   PetscCall(VecDestroy(&u));
1667:   PetscFunctionReturn(PETSC_SUCCESS);
1668: }

1670: static PetscErrorCode SetProblem(TS ts)
1671: {
1672:   AppCtx *user;
1673:   DM      sw;

1675:   PetscFunctionBegin;
1676:   PetscCall(TSGetDM(ts, &sw));
1677:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1678:   // Define unified system for (X, V)
1679:   {
1680:     Mat      J;
1681:     PetscInt dim, Np;

1683:     PetscCall(DMGetDimension(sw, &dim));
1684:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1685:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1686:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1687:     PetscCall(MatSetBlockSize(J, 2 * dim));
1688:     PetscCall(MatSetFromOptions(J));
1689:     PetscCall(MatSetUp(J));
1690:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1691:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1692:     PetscCall(MatDestroy(&J));
1693:   }
1694:   /* Define split system for X and V */
1695:   {
1696:     Vec             u;
1697:     IS              isx, isv, istmp;
1698:     const PetscInt *idx;
1699:     PetscInt        dim, Np, rstart;

1701:     PetscCall(TSGetSolution(ts, &u));
1702:     PetscCall(DMGetDimension(sw, &dim));
1703:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1704:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1705:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1706:     PetscCall(ISGetIndices(istmp, &idx));
1707:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1708:     PetscCall(ISRestoreIndices(istmp, &idx));
1709:     PetscCall(ISDestroy(&istmp));
1710:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1711:     PetscCall(ISGetIndices(istmp, &idx));
1712:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1713:     PetscCall(ISRestoreIndices(istmp, &idx));
1714:     PetscCall(ISDestroy(&istmp));
1715:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1716:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1717:     PetscCall(ISDestroy(&isx));
1718:     PetscCall(ISDestroy(&isv));
1719:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1720:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1721:   }
1722:   PetscFunctionReturn(PETSC_SUCCESS);
1723: }

1725: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1726: {
1727:   DM        sw;
1728:   Vec       u;
1729:   PetscReal t, maxt, dt;
1730:   PetscInt  n, maxn;

1732:   PetscFunctionBegin;
1733:   PetscCall(TSGetDM(ts, &sw));
1734:   PetscCall(TSGetTime(ts, &t));
1735:   PetscCall(TSGetMaxTime(ts, &maxt));
1736:   PetscCall(TSGetTimeStep(ts, &dt));
1737:   PetscCall(TSGetStepNumber(ts, &n));
1738:   PetscCall(TSGetMaxSteps(ts, &maxn));

1740:   PetscCall(TSReset(ts));
1741:   PetscCall(TSSetDM(ts, sw));
1742:   PetscCall(TSSetFromOptions(ts));
1743:   PetscCall(TSSetTime(ts, t));
1744:   PetscCall(TSSetMaxTime(ts, maxt));
1745:   PetscCall(TSSetTimeStep(ts, dt));
1746:   PetscCall(TSSetStepNumber(ts, n));
1747:   PetscCall(TSSetMaxSteps(ts, maxn));

1749:   PetscCall(CreateSolution(ts));
1750:   PetscCall(SetProblem(ts));
1751:   PetscCall(TSGetSolution(ts, &u));
1752:   PetscFunctionReturn(PETSC_SUCCESS);
1753: }

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

1758:   Input Parameters:
1759: + ts         - The TS
1760: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

1762:   Output Parameter:
1763: . u - The initialized solution vector

1765:   Level: advanced

1767: .seealso: InitializeSolve()
1768: */
1769: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1770: {
1771:   DM       sw;
1772:   Vec      u, gc, gv, gc0, gv0;
1773:   IS       isx, isv;
1774:   PetscInt dim;
1775:   AppCtx  *user;

1777:   PetscFunctionBeginUser;
1778:   PetscCall(TSGetDM(ts, &sw));
1779:   PetscCall(DMGetApplicationContext(sw, &user));
1780:   PetscCall(DMGetDimension(sw, &dim));
1781:   if (useInitial) {
1782:     PetscReal v0[2] = {1., 0.};
1783:     if (user->perturbed_weights) {
1784:       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1785:     } else {
1786:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
1787:       PetscCall(DMSwarmInitializeCoordinates(sw));
1788:       if (user->fake_1D) {
1789:         PetscCall(InitializeVelocites_Fake1D(sw, user));
1790:       } else {
1791:         PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
1792:       }
1793:     }
1794:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1795:     PetscCall(DMSwarmTSRedistribute(ts));
1796:   }
1797:   PetscCall(TSGetSolution(ts, &u));
1798:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1799:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1800:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1801:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
1802:   if (useInitial) PetscCall(VecCopy(gc, gc0));
1803:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1804:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1805:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
1806:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1807:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
1808:   if (useInitial) PetscCall(VecCopy(gv, gv0));
1809:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1810:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1811:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
1812:   PetscFunctionReturn(PETSC_SUCCESS);
1813: }

1815: static PetscErrorCode InitializeSolve(TS ts, Vec u)
1816: {
1817:   PetscFunctionBegin;
1818:   PetscCall(TSSetSolution(ts, u));
1819:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1820:   PetscFunctionReturn(PETSC_SUCCESS);
1821: }

1823: static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
1824: {
1825:   MPI_Comm           comm;
1826:   DM                 sw;
1827:   AppCtx            *user;
1828:   const PetscScalar *u;
1829:   const PetscReal   *coords, *vel;
1830:   PetscScalar       *e;
1831:   PetscReal          t;
1832:   PetscInt           dim, Np, p;

1834:   PetscFunctionBeginUser;
1835:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
1836:   PetscCall(TSGetDM(ts, &sw));
1837:   PetscCall(DMGetApplicationContext(sw, &user));
1838:   PetscCall(DMGetDimension(sw, &dim));
1839:   PetscCall(TSGetSolveTime(ts, &t));
1840:   PetscCall(VecGetArray(E, &e));
1841:   PetscCall(VecGetArrayRead(U, &u));
1842:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1843:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1844:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1845:   Np /= 2 * dim;
1846:   for (p = 0; p < Np; ++p) {
1847:     /* TODO generalize initial conditions and project into plane instead of assuming x-y */
1848:     const PetscReal    r0    = DMPlex_NormD_Internal(dim, &coords[p * dim]);
1849:     const PetscReal    th0   = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]);
1850:     const PetscReal    v0    = DMPlex_NormD_Internal(dim, &vel[p * dim]);
1851:     const PetscReal    omega = v0 / r0;
1852:     const PetscReal    ct    = PetscCosReal(omega * t + th0);
1853:     const PetscReal    st    = PetscSinReal(omega * t + th0);
1854:     const PetscScalar *x     = &u[(p * 2 + 0) * dim];
1855:     const PetscScalar *v     = &u[(p * 2 + 1) * dim];
1856:     const PetscReal    xe[3] = {r0 * ct, r0 * st, 0.0};
1857:     const PetscReal    ve[3] = {-v0 * st, v0 * ct, 0.0};
1858:     PetscInt           d;

1860:     for (d = 0; d < dim; ++d) {
1861:       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
1862:       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
1863:     }
1864:     if (user->error) {
1865:       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
1866:       const PetscReal exen = 0.5 * PetscSqr(v0);
1867:       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)));
1868:     }
1869:   }
1870:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1871:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1872:   PetscCall(VecRestoreArrayRead(U, &u));
1873:   PetscCall(VecRestoreArray(E, &e));
1874:   PetscFunctionReturn(PETSC_SUCCESS);
1875: }

1877: #if 0
1878: static PetscErrorCode EnergyMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
1879: {
1880:   const PetscInt     ostep = ((AppCtx *)ctx)->ostep;
1881:   const EMType       em    = ((AppCtx *)ctx)->em;
1882:   DM                 sw;
1883:   const PetscScalar *u;
1884:   PetscReal         *coords, *E;
1885:   PetscReal          enKin = 0., enEM = 0.;
1886:   PetscInt           dim, d, Np, p, q;

1888:   PetscFunctionBeginUser;
1889:   if (step % ostep == 0) {
1890:     PetscCall(TSGetDM(ts, &sw));
1891:     PetscCall(DMGetDimension(sw, &dim));
1892:     PetscCall(VecGetArrayRead(U, &u));
1893:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1894:     Np /= 2 * dim;
1895:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1896:     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1897:     if (!step) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)ts), "Time     Step Part     Energy\n"));
1898:     for (p = 0; p < Np; ++p) {
1899:       const PetscReal v2     = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);
1900:       PetscReal      *pcoord = &coords[p * dim];

1902:       PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4D %5D %10.4lf\n", t, step, p, (double)0.5 * v2));
1903:       enKin += 0.5 * v2;
1904:       if (em == EM_NONE) {
1905:         continue;
1906:       } else if (em == EM_COULOMB) {
1907:         for (q = p + 1; q < Np; ++q) {
1908:           PetscReal *qcoord = &coords[q * dim];
1909:           PetscReal  rpq[3], r;
1910:           for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1911:           r = DMPlex_NormD_Internal(dim, rpq);
1912:           enEM += 1. / r;
1913:         }
1914:       } else if (em == EM_PRIMAL) {
1915:         for (d = 0; d < dim; ++d) enEM += E[p * dim + d];
1916:       }
1917:     }
1918:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 2\t    %10.4lf\n", t, step, (double)enKin));
1919:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 3\t    %10.4lf\n", t, step, (double)enEM));
1920:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 4\t    %10.4lf\n", t, step, (double)enKin + enEM));
1921:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1922:     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1923:     PetscCall(PetscSynchronizedFlush(PetscObjectComm((PetscObject)ts), NULL));
1924:     PetscCall(VecRestoreArrayRead(U, &u));
1925:   }
1926:   PetscFunctionReturn(PETSC_SUCCESS);
1927: }
1928: #endif

1930: static PetscErrorCode MigrateParticles(TS ts)
1931: {
1932:   DM sw;

1934:   PetscFunctionBeginUser;
1935:   PetscCall(TSGetDM(ts, &sw));
1936:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1937:   {
1938:     Vec u, gc, gv;
1939:     IS  isx, isv;

1941:     PetscCall(TSGetSolution(ts, &u));
1942:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1943:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1944:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1945:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1946:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1947:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1948:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
1949:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1950:   }
1951:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1952:   PetscCall(DMSwarmTSRedistribute(ts));
1953:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
1954:   PetscFunctionReturn(PETSC_SUCCESS);
1955: }

1957: static PetscErrorCode MigrateParticles_Periodic(TS ts)
1958: {
1959:   DM       sw, dm;
1960:   PetscInt dim;

1962:   PetscFunctionBeginUser;
1963:   PetscCall(TSGetDM(ts, &sw));
1964:   PetscCall(DMGetDimension(sw, &dim));
1965:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1966:   {
1967:     Vec        u, position, momentum, gc, gv;
1968:     IS         isx, isv;
1969:     PetscReal *pos, *mom, *x, *v;
1970:     PetscReal  lower_bound[3], upper_bound[3];
1971:     PetscInt   p, d, Np;

1973:     PetscCall(TSGetSolution(ts, &u));
1974:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1975:     PetscCall(DMSwarmGetCellDM(sw, &dm));
1976:     PetscCall(DMGetBoundingBox(dm, lower_bound, upper_bound));
1977:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1978:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1979:     PetscCall(VecGetSubVector(u, isx, &position));
1980:     PetscCall(VecGetSubVector(u, isv, &momentum));
1981:     PetscCall(VecGetArray(position, &pos));
1982:     PetscCall(VecGetArray(momentum, &mom));

1984:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1985:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1986:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1987:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

1989:     PetscCall(VecGetArray(gc, &x));
1990:     PetscCall(VecGetArray(gv, &v));
1991:     for (p = 0; p < Np; ++p) {
1992:       for (d = 0; d < dim; ++d) {
1993:         if (pos[p * dim + d] < lower_bound[d]) {
1994:           x[p * dim + d] = pos[p * dim + d] + (upper_bound[d] - lower_bound[d]);
1995:         } else if (pos[p * dim + d] > upper_bound[d]) {
1996:           x[p * dim + d] = pos[p * dim + d] - (upper_bound[d] - lower_bound[d]);
1997:         } else {
1998:           x[p * dim + d] = pos[p * dim + d];
1999:         }
2000:         PetscCheck(x[p * dim + d] >= lower_bound[d] && x[p * dim + d] <= upper_bound[d], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "p: %" PetscInt_FMT "x[%" PetscInt_FMT "] %g", p, d, (double)x[p * dim + d]);
2001:         v[p * dim + d] = mom[p * dim + d];
2002:       }
2003:     }
2004:     PetscCall(VecRestoreArray(gc, &x));
2005:     PetscCall(VecRestoreArray(gv, &v));
2006:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
2007:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));

2009:     PetscCall(VecRestoreArray(position, &pos));
2010:     PetscCall(VecRestoreArray(momentum, &mom));
2011:     PetscCall(VecRestoreSubVector(u, isx, &position));
2012:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2013:   }
2014:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2015:   PetscCall(DMSwarmTSRedistribute(ts));
2016:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2017:   PetscFunctionReturn(PETSC_SUCCESS);
2018: }

2020: int main(int argc, char **argv)
2021: {
2022:   DM     dm, sw;
2023:   TS     ts;
2024:   Vec    u;
2025:   AppCtx user;

2027:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2028:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2029:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2030:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2031:   PetscCall(CreatePoisson(dm, &user));
2032:   PetscCall(CreateSwarm(dm, &user, &sw));
2033:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2034:   PetscCall(InitializeConstants(sw, &user));
2035:   PetscCall(DMSetApplicationContext(sw, &user));

2037:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2038:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2039:   PetscCall(TSSetDM(ts, sw));
2040:   PetscCall(TSSetMaxTime(ts, 0.1));
2041:   PetscCall(TSSetTimeStep(ts, 0.00001));
2042:   PetscCall(TSSetMaxSteps(ts, 100));
2043:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

2045:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2046:   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2047:   if (user.monitor_positions) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2048:   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));

2050:   PetscCall(TSSetFromOptions(ts));
2051:   PetscReal dt;
2052:   PetscInt  maxn;
2053:   PetscCall(TSGetTimeStep(ts, &dt));
2054:   PetscCall(TSGetMaxSteps(ts, &maxn));
2055:   user.steps    = maxn;
2056:   user.stepSize = dt;
2057:   PetscCall(SetupContext(dm, sw, &user));

2059:   PetscCall(DMSwarmVectorDefineField(sw, "velocity"));
2060:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2061:   PetscCall(TSSetComputeExactError(ts, ComputeError));
2062:   if (user.periodic) {
2063:     PetscCall(TSSetPostStep(ts, MigrateParticles_Periodic));
2064:   } else {
2065:     PetscCall(TSSetPostStep(ts, MigrateParticles));
2066:   }
2067:   PetscCall(CreateSolution(ts));
2068:   PetscCall(TSGetSolution(ts, &u));
2069:   PetscCall(TSComputeInitialCondition(ts, u));

2071:   PetscCall(TSSolve(ts, NULL));

2073:   PetscCall(SNESDestroy(&user.snes));
2074:   PetscCall(TSDestroy(&ts));
2075:   PetscCall(DMDestroy(&sw));
2076:   PetscCall(DMDestroy(&dm));
2077:   PetscCall(DestroyContext(&user));
2078:   PetscCall(PetscFinalize());
2079:   return 0;
2080: }

2082: /*TEST

2084:    build:
2085:      requires: double !complex

2087:     # Recommend -draw_size 500,500
2088:    testset:
2089:      args: -dm_plex_dim 2 -fake_1D -dm_plex_simplex 0 -dm_plex_box_faces 20,1 -dm_plex_box_lower 0,-1 -dm_plex_box_upper 12.5664,1 \
2090:            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2091:            -dm_plex_box_bd periodic,none -periodic -ts_type basicsymplectic -ts_basicsymplectic_type 1\
2092:            -dm_view -output_step 50 -sigma 1.0e-8 -timeScale 2.0e-14\
2093:            -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 0
2094:      test:
2095:        suffix: none_1d
2096:        args: -em_type none -error
2097:      test:
2098:        suffix: coulomb_1d
2099:        args: -em_type coulomb

2101:    # For verification, we use
2102:    # -dm_plex_box_faces 100,1 -vdm_plex_box_faces 8000 -dm_swarm_num_particles 800000
2103:    # -ts_monitor_sp_swarm_multi_species 0 -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
2104:    testset:
2105:      args: -dm_plex_dim 2 -dm_plex_box_bd periodic,none -dm_plex_simplex 0 -dm_plex_box_faces 10,1 -dm_plex_box_lower 0,-0.5 -dm_plex_box_upper 12.5664,0.5\
2106:            -ts_dt 0.03 -ts_max_time 500 -ts_max_steps 500 -ts_type basicsymplectic -ts_basicsymplectic_type 1\
2107:            -em_snes_atol 1.e-12 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged\
2108:            -dm_swarm_num_species 1 -dm_swarm_num_particles 100 -dm_view\
2109:            -vdm_plex_dim 1 -vdm_plex_box_lower -10 -vdm_plex_box_upper 10 -vdm_plex_simplex 0 -vdm_plex_box_faces 10\
2110:            -output_step 1 -fake_1D -perturbed_weights -periodic -cosine_coefficients 0.01,0.5 -charges -1.0,1.0 -total_weight 1.0
2111:      test:
2112:        suffix: uniform_equilibrium_1d
2113:        args: -cosine_coefficients 0.0,0.5 -em_type primal -petscspace_degree 1 -em_pc_type svd
2114:      test:
2115:        suffix: uniform_primal_1d
2116:        args: -em_type primal -petscspace_degree 1 -em_pc_type svd
2117:      test:
2118:        suffix: uniform_none_1d
2119:        args: -em_type none
2120:      test:
2121:        suffix: uniform_mixed_1d
2122:        args: -em_type mixed\
2123:              -ksp_rtol 1e-10\
2124:              -em_ksp_type preonly\
2125:              -em_ksp_error_if_not_converged\
2126:              -em_snes_error_if_not_converged\
2127:              -em_pc_type fieldsplit\
2128:              -em_fieldsplit_field_pc_type lu \
2129:              -em_fieldsplit_potential_pc_type svd\
2130:              -em_pc_fieldsplit_type schur\
2131:              -em_pc_fieldsplit_schur_fact_type full\
2132:              -em_pc_fieldsplit_schur_precondition full\
2133:              -potential_petscspace_degree 0 \
2134:              -potential_petscdualspace_lagrange_use_moments \
2135:              -potential_petscdualspace_lagrange_moment_order 2 \
2136:              -field_petscspace_degree 2\
2137:              -field_petscfe_default_quadrature_order 1\
2138:              -field_petscspace_type sum \
2139:              -field_petscspace_variables 2 \
2140:              -field_petscspace_components 2 \
2141:              -field_petscspace_sum_spaces 2 \
2142:              -field_petscspace_sum_concatenate true \
2143:              -field_sumcomp_0_petscspace_variables 2 \
2144:              -field_sumcomp_0_petscspace_type tensor \
2145:              -field_sumcomp_0_petscspace_tensor_spaces 2 \
2146:              -field_sumcomp_0_petscspace_tensor_uniform false \
2147:              -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
2148:              -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
2149:              -field_sumcomp_1_petscspace_variables 2 \
2150:              -field_sumcomp_1_petscspace_type tensor \
2151:              -field_sumcomp_1_petscspace_tensor_spaces 2 \
2152:              -field_sumcomp_1_petscspace_tensor_uniform false \
2153:              -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
2154:              -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
2155:              -field_petscdualspace_form_degree -1 \
2156:              -field_petscdualspace_order 1 \
2157:              -field_petscdualspace_lagrange_trimmed true \
2158:              -ksp_gmres_restart 500

2160: TEST*/