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

169:   PetscFunctionReturn(PETSC_SUCCESS);
170: }

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

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

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

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

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

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

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

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

264:   PetscFunctionReturn(PETSC_SUCCESS);
265: }

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

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

288:   PetscCall(PetscBagDestroy(&user->bag));
289:   PetscFunctionReturn(PETSC_SUCCESS);
290: }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

430:   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]));
431:   PetscCall(PetscDrawLGAddPoint(user->drawlg_ef, &t, &lgEmax));
432:   PetscCall(PetscDrawLGDraw(user->drawlg_ef));
433:   PetscCall(PetscDrawSave(user->drawef));
434:   PetscFunctionReturn(PETSC_SUCCESS);
435: }

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

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

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

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

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

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

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

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

487:     PetscCall(PetscDrawHGDraw(user->drawhgic_v));
488:     PetscCall(PetscDrawHGSave(user->drawhgic_v));

490:     PetscCall(PetscDrawHGDraw(user->drawhgic_w));
491:     PetscCall(PetscDrawHGSave(user->drawhgic_w));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

746:     PetscCall(DMGetLabel(dm, "marker", &label));
747:     PetscCall(DMGetDS(dm, &ds));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

938:       /*Transform quadrature points from real space to ideal real space (0, 2PI/k)*/
939:       if (user->fake_1D) {
940:         PetscCall(PetscPDFCosine1D(xr_x, scale, &den_x));
941:         detJ_x = J_x[0];
942:       } else PetscCall(PetscPDFCosine2D(xr_x, scale, &den_x));
943:       /*We have to transform the quadrature weights as well*/
944:       weight_x[c] += den_x * (wq_x[q] * detJ_x);
945:     }
946:     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]));
947:     totalcellweight += weight_x[c];
948:     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);

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

958:       const PetscInt p = pidx[cv];

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

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

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

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

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

1007:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
1008:     switch (dim) {
1009:     case 1:
1010:       Area = (gmax[0] - gmin[0]);
1011:       break;
1012:     case 2:
1013:       if (user->fake_1D) {
1014:         Area = (gmax[0] - gmin[0]);
1015:       } else {
1016:         Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
1017:       }
1018:       break;
1019:     case 3:
1020:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
1021:       break;
1022:     default:
1023:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
1024:     }
1025:     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));
1026:     param->sigma = PetscAbsReal(totalCharge / (Area));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1203:   KSP ksp;
1204:   Vec rho0;
1205:   /* Create the charges rho */
1206:   PetscCall(SNESGetDM(snes, &dm));

1208:   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
1209:   PetscCall(DMCreateMassMatrix(dm, dm, &M));
1210:   PetscCall(DMGetGlobalVector(dm, &rho0));
1211:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Primal Compute"));
1212:   PetscCall(DMGetGlobalVector(dm, &rho));
1213:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
1214:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));

1216:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1217:   PetscCall(MatMultTranspose(M_p, f, rho));
1218:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1219:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1220:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1221:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

1223:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1224:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1225:   PetscCall(KSPSetOperators(ksp, M, M));
1226:   PetscCall(KSPSetFromOptions(ksp));
1227:   PetscCall(KSPSolve(ksp, rho, rho0));
1228:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1230:   PetscInt           rhosize;
1231:   PetscReal         *charges;
1232:   const PetscScalar *rho_vals;
1233:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1234:   PetscCall(VecGetSize(rho0, &rhosize));
1235:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1236:   for (c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1237:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1238:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

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

1242:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1243:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1244:   PetscCall(DMRestoreGlobalVector(dm, &rho0));
1245:   PetscCall(KSPDestroy(&ksp));
1246:   PetscCall(MatDestroy(&M_p));
1247:   PetscCall(MatDestroy(&M));

1249:   PetscCall(DMGetGlobalVector(dm, &phi));
1250:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1251:   PetscCall(VecSet(phi, 0.0));
1252:   PetscCall(SNESSolve(snes, rho, phi));
1253:   PetscCall(DMRestoreGlobalVector(dm, &rho));
1254:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

1256:   PetscInt           phisize;
1257:   PetscReal         *pot;
1258:   const PetscScalar *phi_vals;
1259:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1260:   PetscCall(VecGetSize(phi, &phisize));
1261:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1262:   for (c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1263:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1264:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1266:   PetscCall(DMGetLocalVector(dm, &locPhi));
1267:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1268:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1269:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1271:   PetscCall(DMGetDS(dm, &ds));
1272:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1273:   PetscCall(DMSwarmSortGetAccess(sw));
1274:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1275:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

1277:   for (c = cStart; c < cEnd; ++c) {
1278:     PetscTabulation tab;
1279:     PetscScalar    *clPhi = NULL;
1280:     PetscReal      *pcoord, *refcoord;
1281:     PetscReal       v[3], J[9], invJ[9], detJ;
1282:     PetscInt       *points;
1283:     PetscInt        Ncp, cp;

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

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

1318: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
1319: {
1320:   AppCtx         *user;
1321:   DM              dm, potential_dm;
1322:   KSP             ksp;
1323:   IS              potential_IS;
1324:   PetscDS         ds;
1325:   PetscFE         fe;
1326:   PetscFEGeom     feGeometry;
1327:   Mat             M_p, M;
1328:   Vec             phi, locPhi, rho, f, temp_rho, rho0;
1329:   PetscQuadrature q;
1330:   PetscReal      *coords, *pot;
1331:   PetscInt        dim, d, cStart, cEnd, c, Np, fields = 1;

1333:   PetscFunctionBegin;
1334:   PetscCall(DMGetDimension(sw, &dim));
1335:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1336:   PetscCall(DMGetApplicationContext(sw, &user));

1338:   /* Create the charges rho */
1339:   PetscCall(SNESGetDM(snes, &dm));
1340:   PetscCall(DMGetGlobalVector(dm, &rho));
1341:   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));

1343:   PetscCall(DMCreateSubDM(dm, 1, &fields, &potential_IS, &potential_dm));
1344:   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
1345:   PetscCall(DMCreateMassMatrix(potential_dm, potential_dm, &M));
1346:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1347:   PetscCall(MatViewFromOptions(M, NULL, "-m_view"));
1348:   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
1349:   PetscCall(PetscObjectSetName((PetscObject)temp_rho, "Mf"));
1350:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1351:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1352:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1353:   PetscCall(MatMultTranspose(M_p, f, temp_rho));
1354:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1355:   PetscCall(DMGetGlobalVector(potential_dm, &rho0));
1356:   PetscCall(PetscObjectSetName((PetscObject)rho0, "Charge density (rho0) from Mixed Compute"));

1358:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1359:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
1360:   PetscCall(KSPSetOperators(ksp, M, M));
1361:   PetscCall(KSPSetFromOptions(ksp));
1362:   PetscCall(KSPSolve(ksp, temp_rho, rho0));
1363:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));

1365:   PetscInt           rhosize;
1366:   PetscReal         *charges;
1367:   const PetscScalar *rho_vals;
1368:   Parameter         *param;
1369:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1370:   PetscCall(DMSwarmGetField(sw, "charges", NULL, NULL, (void **)&charges));
1371:   PetscCall(VecGetSize(rho0, &rhosize));

1373:   /* Integral over reference element is size 1.  Reference element area is 4.  Scale rho0 by 1/4 because the basis function is 1/4 */
1374:   PetscCall(VecScale(rho0, 0.25));
1375:   PetscCall(VecGetArrayRead(rho0, &rho_vals));
1376:   for (c = 0; c < rhosize; ++c) charges[c] = rho_vals[c];
1377:   PetscCall(VecRestoreArrayRead(rho0, &rho_vals));
1378:   PetscCall(DMSwarmRestoreField(sw, "charges", NULL, NULL, (void **)&charges));

1380:   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
1381:   PetscCall(VecScale(rho, 0.25));
1382:   PetscCall(VecViewFromOptions(rho0, NULL, "-rho0_view"));
1383:   PetscCall(VecViewFromOptions(temp_rho, NULL, "-temprho_view"));
1384:   PetscCall(VecViewFromOptions(rho, NULL, "-rho_view"));
1385:   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
1386:   PetscCall(DMRestoreGlobalVector(potential_dm, &rho0));

1388:   PetscCall(MatDestroy(&M_p));
1389:   PetscCall(MatDestroy(&M));
1390:   PetscCall(KSPDestroy(&ksp));
1391:   PetscCall(DMDestroy(&potential_dm));
1392:   PetscCall(ISDestroy(&potential_IS));

1394:   PetscCall(DMGetGlobalVector(dm, &phi));
1395:   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
1396:   PetscCall(VecSet(phi, 0.0));
1397:   PetscCall(SNESSolve(snes, rho, phi));
1398:   PetscCall(DMRestoreGlobalVector(dm, &rho));

1400:   PetscInt           phisize;
1401:   const PetscScalar *phi_vals;
1402:   PetscCall(DMSwarmGetField(sw, "potential", NULL, NULL, (void **)&pot));
1403:   PetscCall(VecGetSize(phi, &phisize));
1404:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1405:   PetscCall(VecGetArrayRead(phi, &phi_vals));
1406:   for (c = 0; c < phisize; ++c) pot[c] = phi_vals[c];
1407:   PetscCall(VecRestoreArrayRead(phi, &phi_vals));
1408:   PetscCall(DMSwarmRestoreField(sw, "potential", NULL, NULL, (void **)&pot));

1410:   PetscCall(DMGetLocalVector(dm, &locPhi));
1411:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1412:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1413:   PetscCall(DMRestoreGlobalVector(dm, &phi));

1415:   PetscCall(DMGetDS(dm, &ds));
1416:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1417:   PetscCall(DMSwarmSortGetAccess(sw));
1418:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1419:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1420:   PetscCall(PetscFEGetQuadrature(fe, &q));
1421:   PetscCall(PetscFECreateCellGeometry(fe, q, &feGeometry));
1422:   for (c = cStart; c < cEnd; ++c) {
1423:     PetscTabulation tab;
1424:     PetscScalar    *clPhi = NULL;
1425:     PetscReal      *pcoord, *refcoord;
1426:     PetscInt       *points;
1427:     PetscInt        Ncp, cp;

1429:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1430:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1431:     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1432:     for (cp = 0; cp < Ncp; ++cp)
1433:       for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1434:     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
1435:     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
1436:     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, q, feGeometry.v, feGeometry.J, feGeometry.invJ, feGeometry.detJ));
1437:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));

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

1442:       for (d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1443:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, &feGeometry, cp, &E[p * dim]));
1444:       PetscCall(PetscFEPushforward(fe, &feGeometry, 1, &E[p * dim]));
1445:       for (d = 0; d < dim; ++d) {
1446:         E[p * dim + d] *= -2.0;
1447:         if (user->fake_1D && d > 0) E[p * dim + d] = 0;
1448:       }
1449:     }
1450:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1451:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
1452:     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
1453:     PetscCall(PetscTabulationDestroy(&tab));
1454:     PetscCall(PetscFree(points));
1455:   }
1456:   PetscCall(PetscFEDestroyCellGeometry(fe, &feGeometry));
1457:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1458:   PetscCall(DMSwarmSortRestoreAccess(sw));
1459:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
1460:   PetscFunctionReturn(PETSC_SUCCESS);
1461: }

1463: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
1464: {
1465:   AppCtx  *ctx;
1466:   PetscInt dim, Np;

1468:   PetscFunctionBegin;
1472:   PetscCall(DMGetDimension(sw, &dim));
1473:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1474:   PetscCall(DMGetApplicationContext(sw, &ctx));
1475:   PetscCall(PetscArrayzero(E, Np * dim));

1477:   switch (ctx->em) {
1478:   case EM_PRIMAL:
1479:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
1480:     break;
1481:   case EM_COULOMB:
1482:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
1483:     break;
1484:   case EM_MIXED:
1485:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
1486:     break;
1487:   case EM_NONE:
1488:     break;
1489:   default:
1490:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
1491:   }

1493:   PetscFunctionReturn(PETSC_SUCCESS);
1494: }

1496: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1497: {
1498:   DM                 sw;
1499:   SNES               snes = ((AppCtx *)ctx)->snes;
1500:   const PetscReal   *coords, *vel;
1501:   const PetscScalar *u;
1502:   PetscScalar       *g;
1503:   PetscReal         *E, m_p = 1., q_p = -1.;
1504:   PetscInt           dim, d, Np, p;

1506:   PetscFunctionBeginUser;
1507:   PetscCall(TSGetDM(ts, &sw));
1508:   PetscCall(DMGetDimension(sw, &dim));
1509:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1510:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1511:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1512:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1513:   PetscCall(VecGetArrayRead(U, &u));
1514:   PetscCall(VecGetArray(G, &g));

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

1518:   Np /= 2 * dim;
1519:   for (p = 0; p < Np; ++p) {
1520:     for (d = 0; d < dim; ++d) {
1521:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1522:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1523:     }
1524:   }
1525:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1526:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1527:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1528:   PetscCall(VecRestoreArrayRead(U, &u));
1529:   PetscCall(VecRestoreArray(G, &g));
1530:   PetscFunctionReturn(PETSC_SUCCESS);
1531: }

1533: /* J_{ij} = dF_i/dx_j
1534:    J_p = (  0   1)
1535:          (-w^2  0)
1536:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1537:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1538: */
1539: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1540: {
1541:   DM               sw;
1542:   const PetscReal *coords, *vel;
1543:   PetscInt         dim, d, Np, p, rStart;

1545:   PetscFunctionBeginUser;
1546:   PetscCall(TSGetDM(ts, &sw));
1547:   PetscCall(DMGetDimension(sw, &dim));
1548:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1549:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1550:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1551:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1552:   Np /= 2 * dim;
1553:   for (p = 0; p < Np; ++p) {
1554:     const PetscReal x0      = coords[p * dim + 0];
1555:     const PetscReal vy0     = vel[p * dim + 1];
1556:     const PetscReal omega   = vy0 / x0;
1557:     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};

1559:     for (d = 0; d < dim; ++d) {
1560:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1561:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1562:     }
1563:   }
1564:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1565:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1566:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1567:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1568:   PetscFunctionReturn(PETSC_SUCCESS);
1569: }

1571: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1572: {
1573:   AppCtx            *user = (AppCtx *)ctx;
1574:   DM                 sw;
1575:   const PetscScalar *v;
1576:   PetscScalar       *xres;
1577:   PetscInt           Np, p, d, dim;

1579:   PetscFunctionBeginUser;
1580:   PetscCall(TSGetDM(ts, &sw));
1581:   PetscCall(DMGetDimension(sw, &dim));
1582:   PetscCall(VecGetLocalSize(Xres, &Np));
1583:   PetscCall(VecGetArrayRead(V, &v));
1584:   PetscCall(VecGetArray(Xres, &xres));
1585:   Np /= dim;
1586:   for (p = 0; p < Np; ++p) {
1587:     for (d = 0; d < dim; ++d) {
1588:       xres[p * dim + d] = v[p * dim + d];
1589:       if (user->fake_1D && d > 0) xres[p * dim + d] = 0;
1590:     }
1591:   }
1592:   PetscCall(VecRestoreArrayRead(V, &v));
1593:   PetscCall(VecRestoreArray(Xres, &xres));
1594:   PetscFunctionReturn(PETSC_SUCCESS);
1595: }

1597: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1598: {
1599:   DM                 sw;
1600:   AppCtx            *user = (AppCtx *)ctx;
1601:   SNES               snes = ((AppCtx *)ctx)->snes;
1602:   const PetscScalar *x;
1603:   const PetscReal   *coords, *vel;
1604:   PetscReal         *E, m_p, q_p;
1605:   PetscScalar       *vres;
1606:   PetscInt           Np, p, dim, d;
1607:   Parameter         *param;

1609:   PetscFunctionBeginUser;
1610:   PetscCall(TSGetDM(ts, &sw));
1611:   PetscCall(DMGetDimension(sw, &dim));
1612:   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1613:   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1614:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1615:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
1616:   m_p = user->masses[0] * param->m0;
1617:   q_p = user->charges[0] * param->q0;
1618:   PetscCall(VecGetLocalSize(Vres, &Np));
1619:   PetscCall(VecGetArrayRead(X, &x));
1620:   PetscCall(VecGetArray(Vres, &vres));
1621:   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2");
1622:   PetscCall(ComputeFieldAtParticles(snes, sw, E));

1624:   Np /= dim;
1625:   for (p = 0; p < Np; ++p) {
1626:     for (d = 0; d < dim; ++d) {
1627:       vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1628:       if (user->fake_1D && d > 0) vres[p * dim + d] = 0.;
1629:     }
1630:   }
1631:   PetscCall(VecRestoreArrayRead(X, &x));
1632:   PetscCall(VecRestoreArray(Vres, &vres));
1633:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1634:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1635:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1636:   PetscFunctionReturn(PETSC_SUCCESS);
1637: }

1639: static PetscErrorCode CreateSolution(TS ts)
1640: {
1641:   DM       sw;
1642:   Vec      u;
1643:   PetscInt dim, Np;

1645:   PetscFunctionBegin;
1646:   PetscCall(TSGetDM(ts, &sw));
1647:   PetscCall(DMGetDimension(sw, &dim));
1648:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1649:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1650:   PetscCall(VecSetBlockSize(u, dim));
1651:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1652:   PetscCall(VecSetUp(u));
1653:   PetscCall(TSSetSolution(ts, u));
1654:   PetscCall(VecDestroy(&u));
1655:   PetscFunctionReturn(PETSC_SUCCESS);
1656: }

1658: static PetscErrorCode SetProblem(TS ts)
1659: {
1660:   AppCtx *user;
1661:   DM      sw;

1663:   PetscFunctionBegin;
1664:   PetscCall(TSGetDM(ts, &sw));
1665:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
1666:   // Define unified system for (X, V)
1667:   {
1668:     Mat      J;
1669:     PetscInt dim, Np;

1671:     PetscCall(DMGetDimension(sw, &dim));
1672:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1673:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1674:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1675:     PetscCall(MatSetBlockSize(J, 2 * dim));
1676:     PetscCall(MatSetFromOptions(J));
1677:     PetscCall(MatSetUp(J));
1678:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1679:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1680:     PetscCall(MatDestroy(&J));
1681:   }
1682:   /* Define split system for X and V */
1683:   {
1684:     Vec             u;
1685:     IS              isx, isv, istmp;
1686:     const PetscInt *idx;
1687:     PetscInt        dim, Np, rstart;

1689:     PetscCall(TSGetSolution(ts, &u));
1690:     PetscCall(DMGetDimension(sw, &dim));
1691:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1692:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1693:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1694:     PetscCall(ISGetIndices(istmp, &idx));
1695:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1696:     PetscCall(ISRestoreIndices(istmp, &idx));
1697:     PetscCall(ISDestroy(&istmp));
1698:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1699:     PetscCall(ISGetIndices(istmp, &idx));
1700:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1701:     PetscCall(ISRestoreIndices(istmp, &idx));
1702:     PetscCall(ISDestroy(&istmp));
1703:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1704:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1705:     PetscCall(ISDestroy(&isx));
1706:     PetscCall(ISDestroy(&isv));
1707:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1708:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1709:   }
1710:   PetscFunctionReturn(PETSC_SUCCESS);
1711: }

1713: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1714: {
1715:   DM        sw;
1716:   Vec       u;
1717:   PetscReal t, maxt, dt;
1718:   PetscInt  n, maxn;

1720:   PetscFunctionBegin;
1721:   PetscCall(TSGetDM(ts, &sw));
1722:   PetscCall(TSGetTime(ts, &t));
1723:   PetscCall(TSGetMaxTime(ts, &maxt));
1724:   PetscCall(TSGetTimeStep(ts, &dt));
1725:   PetscCall(TSGetStepNumber(ts, &n));
1726:   PetscCall(TSGetMaxSteps(ts, &maxn));

1728:   PetscCall(TSReset(ts));
1729:   PetscCall(TSSetDM(ts, sw));
1730:   PetscCall(TSSetFromOptions(ts));
1731:   PetscCall(TSSetTime(ts, t));
1732:   PetscCall(TSSetMaxTime(ts, maxt));
1733:   PetscCall(TSSetTimeStep(ts, dt));
1734:   PetscCall(TSSetStepNumber(ts, n));
1735:   PetscCall(TSSetMaxSteps(ts, maxn));

1737:   PetscCall(CreateSolution(ts));
1738:   PetscCall(SetProblem(ts));
1739:   PetscCall(TSGetSolution(ts, &u));
1740:   PetscFunctionReturn(PETSC_SUCCESS);
1741: }

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

1746:   Input Parameters:
1747: + ts         - The TS
1748: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

1750:   Output Parameter:
1751: . u - The initialized solution vector

1753:   Level: advanced

1755: .seealso: InitializeSolve()
1756: */
1757: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1758: {
1759:   DM       sw;
1760:   Vec      u, gc, gv, gc0, gv0;
1761:   IS       isx, isv;
1762:   PetscInt dim;
1763:   AppCtx  *user;

1765:   PetscFunctionBeginUser;
1766:   PetscCall(TSGetDM(ts, &sw));
1767:   PetscCall(DMGetApplicationContext(sw, &user));
1768:   PetscCall(DMGetDimension(sw, &dim));
1769:   if (useInitial) {
1770:     PetscReal v0[2] = {1., 0.};
1771:     if (user->perturbed_weights) {
1772:       PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1773:     } else {
1774:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
1775:       PetscCall(DMSwarmInitializeCoordinates(sw));
1776:       if (user->fake_1D) {
1777:         PetscCall(InitializeVelocites_Fake1D(sw, user));
1778:       } else {
1779:         PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
1780:       }
1781:     }
1782:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1783:     PetscCall(DMSwarmTSRedistribute(ts));
1784:   }
1785:   PetscCall(TSGetSolution(ts, &u));
1786:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1787:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1788:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1789:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
1790:   if (useInitial) PetscCall(VecCopy(gc, gc0));
1791:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1792:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1793:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
1794:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1795:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
1796:   if (useInitial) PetscCall(VecCopy(gv, gv0));
1797:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1798:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1799:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
1800:   PetscFunctionReturn(PETSC_SUCCESS);
1801: }

1803: static PetscErrorCode InitializeSolve(TS ts, Vec u)
1804: {
1805:   PetscFunctionBegin;
1806:   PetscCall(TSSetSolution(ts, u));
1807:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1808:   PetscFunctionReturn(PETSC_SUCCESS);
1809: }

1811: static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
1812: {
1813:   MPI_Comm           comm;
1814:   DM                 sw;
1815:   AppCtx            *user;
1816:   const PetscScalar *u;
1817:   const PetscReal   *coords, *vel;
1818:   PetscScalar       *e;
1819:   PetscReal          t;
1820:   PetscInt           dim, Np, p;

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

1848:     for (d = 0; d < dim; ++d) {
1849:       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
1850:       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
1851:     }
1852:     if (user->error) {
1853:       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
1854:       const PetscReal exen = 0.5 * PetscSqr(v0);
1855:       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)));
1856:     }
1857:   }
1858:   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
1859:   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
1860:   PetscCall(VecRestoreArrayRead(U, &u));
1861:   PetscCall(VecRestoreArray(E, &e));
1862:   PetscFunctionReturn(PETSC_SUCCESS);
1863: }

1865: #if 0
1866: static PetscErrorCode EnergyMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
1867: {
1868:   const PetscInt     ostep = ((AppCtx *)ctx)->ostep;
1869:   const EMType       em    = ((AppCtx *)ctx)->em;
1870:   DM                 sw;
1871:   const PetscScalar *u;
1872:   PetscReal         *coords, *E;
1873:   PetscReal          enKin = 0., enEM = 0.;
1874:   PetscInt           dim, d, Np, p, q;

1876:   PetscFunctionBeginUser;
1877:   if (step % ostep == 0) {
1878:     PetscCall(TSGetDM(ts, &sw));
1879:     PetscCall(DMGetDimension(sw, &dim));
1880:     PetscCall(VecGetArrayRead(U, &u));
1881:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1882:     Np /= 2 * dim;
1883:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1884:     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1885:     if (!step) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)ts), "Time     Step Part     Energy\n"));
1886:     for (p = 0; p < Np; ++p) {
1887:       const PetscReal v2     = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);
1888:       PetscReal      *pcoord = &coords[p * dim];

1890:       PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4D %5D %10.4lf\n", t, step, p, (double)0.5 * v2));
1891:       enKin += 0.5 * v2;
1892:       if (em == EM_NONE) {
1893:         continue;
1894:       } else if (em == EM_COULOMB) {
1895:         for (q = p + 1; q < Np; ++q) {
1896:           PetscReal *qcoord = &coords[q * dim];
1897:           PetscReal  rpq[3], r;
1898:           for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
1899:           r = DMPlex_NormD_Internal(dim, rpq);
1900:           enEM += 1. / r;
1901:         }
1902:       } else if (em == EM_PRIMAL) {
1903:         for (d = 0; d < dim; ++d) enEM += E[p * dim + d];
1904:       }
1905:     }
1906:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 2\t    %10.4lf\n", t, step, (double)enKin));
1907:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 3\t    %10.4lf\n", t, step, (double)enEM));
1908:     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " 4\t    %10.4lf\n", t, step, (double)enKin + enEM));
1909:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1910:     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1911:     PetscCall(PetscSynchronizedFlush(PetscObjectComm((PetscObject)ts), NULL));
1912:     PetscCall(VecRestoreArrayRead(U, &u));
1913:   }
1914:   PetscFunctionReturn(PETSC_SUCCESS);
1915: }
1916: #endif

1918: static PetscErrorCode MigrateParticles(TS ts)
1919: {
1920:   DM sw;

1922:   PetscFunctionBeginUser;
1923:   PetscCall(TSGetDM(ts, &sw));
1924:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1925:   {
1926:     Vec u, gc, gv;
1927:     IS  isx, isv;

1929:     PetscCall(TSGetSolution(ts, &u));
1930:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1931:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1932:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1933:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1934:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1935:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1936:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
1937:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1938:   }
1939:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1940:   PetscCall(DMSwarmTSRedistribute(ts));
1941:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
1942:   PetscFunctionReturn(PETSC_SUCCESS);
1943: }

1945: static PetscErrorCode MigrateParticles_Periodic(TS ts)
1946: {
1947:   DM       sw, dm;
1948:   PetscInt dim;

1950:   PetscFunctionBeginUser;
1951:   PetscCall(TSGetDM(ts, &sw));
1952:   PetscCall(DMGetDimension(sw, &dim));
1953:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1954:   {
1955:     Vec        u, position, momentum, gc, gv;
1956:     IS         isx, isv;
1957:     PetscReal *pos, *mom, *x, *v;
1958:     PetscReal  lower_bound[3], upper_bound[3];
1959:     PetscInt   p, d, Np;

1961:     PetscCall(TSGetSolution(ts, &u));
1962:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
1963:     PetscCall(DMSwarmGetCellDM(sw, &dm));
1964:     PetscCall(DMGetBoundingBox(dm, lower_bound, upper_bound));
1965:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1966:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1967:     PetscCall(VecGetSubVector(u, isx, &position));
1968:     PetscCall(VecGetSubVector(u, isv, &momentum));
1969:     PetscCall(VecGetArray(position, &pos));
1970:     PetscCall(VecGetArray(momentum, &mom));

1972:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1973:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1974:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1975:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

1977:     PetscCall(VecGetArray(gc, &x));
1978:     PetscCall(VecGetArray(gv, &v));
1979:     for (p = 0; p < Np; ++p) {
1980:       for (d = 0; d < dim; ++d) {
1981:         if (pos[p * dim + d] < lower_bound[d]) {
1982:           x[p * dim + d] = pos[p * dim + d] + (upper_bound[d] - lower_bound[d]);
1983:         } else if (pos[p * dim + d] > upper_bound[d]) {
1984:           x[p * dim + d] = pos[p * dim + d] - (upper_bound[d] - lower_bound[d]);
1985:         } else {
1986:           x[p * dim + d] = pos[p * dim + d];
1987:         }
1988:         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]);
1989:         v[p * dim + d] = mom[p * dim + d];
1990:       }
1991:     }
1992:     PetscCall(VecRestoreArray(gc, &x));
1993:     PetscCall(VecRestoreArray(gv, &v));
1994:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1995:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));

1997:     PetscCall(VecRestoreArray(position, &pos));
1998:     PetscCall(VecRestoreArray(momentum, &mom));
1999:     PetscCall(VecRestoreSubVector(u, isx, &position));
2000:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
2001:   }
2002:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
2003:   PetscCall(DMSwarmTSRedistribute(ts));
2004:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
2005:   PetscFunctionReturn(PETSC_SUCCESS);
2006: }

2008: int main(int argc, char **argv)
2009: {
2010:   DM     dm, sw;
2011:   TS     ts;
2012:   Vec    u;
2013:   AppCtx user;

2015:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
2016:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
2017:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
2018:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
2019:   PetscCall(CreatePoisson(dm, &user));
2020:   PetscCall(CreateSwarm(dm, &user, &sw));
2021:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
2022:   PetscCall(InitializeConstants(sw, &user));
2023:   PetscCall(DMSetApplicationContext(sw, &user));

2025:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
2026:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
2027:   PetscCall(TSSetDM(ts, sw));
2028:   PetscCall(TSSetMaxTime(ts, 0.1));
2029:   PetscCall(TSSetTimeStep(ts, 0.00001));
2030:   PetscCall(TSSetMaxSteps(ts, 100));
2031:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

2033:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
2034:   if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL));
2035:   if (user.monitor_positions) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL));
2036:   if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL));

2038:   PetscCall(TSSetFromOptions(ts));
2039:   PetscReal dt;
2040:   PetscInt  maxn;
2041:   PetscCall(TSGetTimeStep(ts, &dt));
2042:   PetscCall(TSGetMaxSteps(ts, &maxn));
2043:   user.steps    = maxn;
2044:   user.stepSize = dt;
2045:   PetscCall(SetupContext(dm, sw, &user));

2047:   PetscCall(DMSwarmVectorDefineField(sw, "velocity"));
2048:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
2049:   PetscCall(TSSetComputeExactError(ts, ComputeError));
2050:   if (user.periodic) {
2051:     PetscCall(TSSetPostStep(ts, MigrateParticles_Periodic));
2052:   } else {
2053:     PetscCall(TSSetPostStep(ts, MigrateParticles));
2054:   }
2055:   PetscCall(CreateSolution(ts));
2056:   PetscCall(TSGetSolution(ts, &u));
2057:   PetscCall(TSComputeInitialCondition(ts, u));

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

2061:   PetscCall(SNESDestroy(&user.snes));
2062:   PetscCall(TSDestroy(&ts));
2063:   PetscCall(DMDestroy(&sw));
2064:   PetscCall(DMDestroy(&dm));
2065:   PetscCall(DestroyContext(&user));
2066:   PetscCall(PetscFinalize());
2067:   return 0;
2068: }

2070: /*TEST

2072:    build:
2073:      requires: double !complex

2075:     # Recommend -draw_size 500,500
2076:    testset:
2077:      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 \
2078:            -dm_swarm_coordinate_density constant -dm_swarm_num_particles 100 \
2079:            -dm_plex_box_bd periodic,none -periodic -ts_type basicsymplectic -ts_basicsymplectic_type 1\
2080:            -dm_view -output_step 50 -sigma 1.0e-8 -timeScale 2.0e-14\
2081:            -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 0
2082:      test:
2083:        suffix: none_1d
2084:        args: -em_type none -error
2085:      test:
2086:        suffix: coulomb_1d
2087:        args: -em_type coulomb

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

2148: TEST*/