Date: 2026-03-31T05:46:37+00:00

Actual source code: ex3.c

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

      3: /*
      4:   This example solves the Vlasov-Poisson system for Landau damping (6X + 6V).
      5:   The system is solved using a Particle-In-Cell (PIC) method with DMSwarm for particles and DMPlex/PetscFE for the field.
      6:   This particular example uses the velocity mesh from DMPlexLandauCreateVelocitySpace for 3D velocity space.

      8:   Options:
      9:   -particle_monitor [prefix] : Output particle data (x, v, w, E) to binary files at each output step.
     10:                                Optional prefix for filenames (default: "particles").

     12: */
     13: #include <petscts.h>
     14: #include <petscdmplex.h>
     15: #include <petscdmswarm.h>
     16: #include <petscfe.h>
     17: #include <petscds.h>
     18: #include <petscbag.h>
     19: #include <petscdraw.h>
     20: #include <petscviewer.h>
     21: #include <petsclandau.h>
     22: #include <petscdmcomposite.h>
     23: #include <petsc/private/dmpleximpl.h>
     24: #include <petsc/private/petscfeimpl.h>
     25: #include <petsc/private/dmswarmimpl.h>
     26: #include "petscdm.h"
     27: #include "petscdmlabel.h"

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

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

     58: typedef struct {
     59:   PetscBag      bag;            // Problem parameters
     60:   PetscBool     error;          // Flag for printing the error
     61:   PetscBool     efield_monitor; // Flag to show electric field monitor
     62:   PetscBool     moment_monitor; // Flag to show distribution moment monitor
     63:   char          particle_monitor_prefix[PETSC_MAX_PATH_LEN];
     64:   PetscBool     particle_monitor; // Flag to output particle data
     65:   PetscInt      ostep;            // Print the energy at each ostep time steps
     66:   PetscInt      numParticles;
     67:   PetscReal     timeScale;              /* Nondimensionalizing time scale */
     68:   PetscReal     charges[2];             /* The charges of each species */
     69:   PetscReal     masses[2];              /* The masses of each species */
     70:   PetscReal     thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
     71:   PetscReal     cosine_coefficients[2]; /*(alpha, k)*/
     72:   PetscReal     totalWeight;
     73:   PetscReal     stepSize;
     74:   PetscInt      steps;
     75:   PetscReal     initVel;
     76:   SNES          snes;       // EM solver
     77:   DM            dmPot;      // The DM for potential
     78:   Mat           M;          // The finite element mass matrix for potential
     79:   PetscFEGeom  *fegeom;     // Geometric information for the DM cells
     80:   PetscBool     validE;     // Flag to indicate E-field in swarm is valid
     81:   PetscReal     drawlgEmin; // The minimum lg(E) to plot
     82:   PetscDrawLG   drawlgE;    // Logarithm of maximum electric field
     83:   DM            swarm;
     84:   PetscBool     checkweights;
     85:   PetscInt      checkVRes; // Flag to check/output velocity residuals for nightly tests
     86:   DM            landau_pack;
     87:   PetscBool     use_landau_velocity_space;
     88:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
     89: } AppCtx;

     91: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
     92: {
     93:   PetscFunctionBeginUser;
     94:   PetscInt d                         = 2;
     95:   PetscInt maxSpecies                = 2;
     96:   options->error                     = PETSC_FALSE;
     97:   options->efield_monitor            = PETSC_FALSE;
     98:   options->moment_monitor            = PETSC_FALSE;
     99:   options->particle_monitor          = PETSC_FALSE;
    100:   options->ostep                     = 100;
    101:   options->timeScale                 = 2.0e-14;
    102:   options->charges[0]                = -1.0;
    103:   options->charges[1]                = 1.0;
    104:   options->masses[0]                 = 1.0;
    105:   options->masses[1]                 = 1000.0;
    106:   options->thermal_energy[0]         = 1.0;
    107:   options->thermal_energy[1]         = 1.0;
    108:   options->cosine_coefficients[0]    = 0.01;
    109:   options->cosine_coefficients[1]    = 0.5;
    110:   options->initVel                   = 1;
    111:   options->totalWeight               = 1.0;
    112:   options->validE                    = PETSC_FALSE;
    113:   options->drawlgEmin                = -6;
    114:   options->drawlgE                   = NULL;
    115:   options->snes                      = NULL;
    116:   options->dmPot                     = NULL;
    117:   options->M                         = NULL;
    118:   options->numParticles              = 32768;
    119:   options->checkweights              = PETSC_FALSE;
    120:   options->checkVRes                 = 0;
    121:   options->landau_pack               = NULL;
    122:   options->use_landau_velocity_space = PETSC_FALSE;

    124:   PetscOptionsBegin(comm, "", "Landau Damping options", "DMSWARM");
    125:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex3.c", options->error, &options->error, NULL));
    126:   PetscCall(PetscOptionsBool("-efield_monitor", "Flag to plot log(max E) over time", "ex3.c", options->efield_monitor, &options->efield_monitor, NULL));
    127:   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", "ex3.c", options->drawlgEmin, &options->drawlgEmin, NULL));
    128:   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", "ex3.c", options->moment_monitor, &options->moment_monitor, NULL));
    129:   PetscCall(PetscOptionsString("-particle_monitor", "Prefix for particle data files", "ex3.c", options->particle_monitor_prefix, options->particle_monitor_prefix, sizeof(options->particle_monitor_prefix), &options->particle_monitor));
    130:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex3.c", options->checkweights, &options->checkweights, NULL));
    131:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex3.c", options->checkVRes, &options->checkVRes, NULL));
    132:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex3.c", options->ostep, &options->ostep, NULL));
    133:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex3.c", options->timeScale, &options->timeScale, NULL));
    134:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex3.c", options->initVel, &options->initVel, NULL));
    135:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex3.c", options->totalWeight, &options->totalWeight, NULL));
    136:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex3.c", options->cosine_coefficients, &d, NULL));
    137:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex3.c", options->charges, &maxSpecies, NULL));
    138:   PetscCall(PetscOptionsBool("-use_landau_velocity_space", "Use Landau velocity space", "ex3.c", options->use_landau_velocity_space, &options->use_landau_velocity_space, NULL));
    139:   PetscOptionsEnd();

    141:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
    142:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
    143:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
    144:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
    145:   PetscFunctionReturn(PETSC_SUCCESS);
    146: }

    148: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
    149: {
    150:   MPI_Comm comm;

    152:   PetscFunctionBeginUser;
    153:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
    154:   if (user->efield_monitor) {
    155:     PetscDraw     draw;
    156:     PetscDrawAxis axis;

    158:     PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
    159:     PetscCall(PetscDrawSetSave(draw, "ex3_Efield"));
    160:     PetscCall(PetscDrawSetFromOptions(draw));
    161:     PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
    162:     PetscCall(PetscDrawDestroy(&draw));
    163:     PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
    164:     PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
    165:     PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
    166:     PetscCall(PetscDrawLGSetUseMarkers(user->drawlgE, PETSC_FALSE));
    167:   }
    168:   PetscFunctionReturn(PETSC_SUCCESS);
    169: }

    171: static PetscErrorCode DestroyContext(AppCtx *user)
    172: {
    173:   PetscFunctionBeginUser;
    174:   if (user->landau_pack) PetscCall(DMPlexLandauDestroyVelocitySpace(&user->landau_pack));
    175:   PetscCall(PetscDrawLGDestroy(&user->drawlgE));
    176:   PetscCall(PetscBagDestroy(&user->bag));
    177:   PetscFunctionReturn(PETSC_SUCCESS);
    178: }

    180: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
    181: {
    182:   const PetscScalar *w;
    183:   PetscInt           Np;

    185:   PetscFunctionBeginUser;
    186:   if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
    187:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
    188:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    189:   for (PetscInt p = 0; p < Np; ++p) PetscCheck(w[p] >= 0.0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " has negative weight %g", p, w[p]);
    190:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
    191:   PetscFunctionReturn(PETSC_SUCCESS);
    192: }

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

    199: static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
    200: {
    201:   PetscDS        ds;
    202:   const PetscInt field = 0;
    203:   PetscInt       Nf;
    204:   void          *ctx;

    206:   PetscFunctionBegin;
    207:   PetscCall(DMGetApplicationContext(dm, &ctx));
    208:   PetscCall(DMGetDS(dm, &ds));
    209:   PetscCall(PetscDSGetNumFields(ds, &Nf));
    210:   PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
    211:   PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
    212:   PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
    213:   PetscFunctionReturn(PETSC_SUCCESS);
    214: }

    216: static void f0_grad_phi2(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
    217: {
    218:   f0[0] = 0.;
    219:   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d]); // + d * dim  cause segfault
    220: }

    222: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
    223: {
    224:   AppCtx     *user = (AppCtx *)ctx;
    225:   DM          sw;
    226:   PetscScalar intESq;
    227:   PetscReal  *E, *x, *weight;
    228:   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
    229:   PetscReal   pmoments[16]; /* \int f, \int v f, \int v^2 f */
    230:   PetscInt   *species, dim, Np, gNp;
    231:   MPI_Comm    comm;

    233:   PetscFunctionBeginUser;
    234:   if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
    235:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
    236:   PetscCall(TSGetDM(ts, &sw));
    237:   PetscCall(DMGetDimension(sw, &dim));
    238:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    239:   PetscCall(DMSwarmGetSize(sw, &gNp));
    240:   PetscCall(DMSwarmSortGetAccess(sw));
    241:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    242:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    243:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    244:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

    246:   for (PetscInt p = 0; p < Np; ++p) {
    247:     PetscReal Emag = 0.;
    248:     for (PetscInt d = 0; d < dim; ++d) {
    249:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
    250:       if (temp > Emax) Emax = temp;
    251:       Emag += PetscSqr(E[p * dim + d]);
    252:     }
    253:     Enorm += PetscSqrtReal(Emag);
    254:     sum += E[p * dim];
    255:     chargesum += user->charges[0] * weight[p];
    256:   }
    257:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
    258:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Enorm, 1, MPIU_REAL, MPIU_SUM, comm));
    259:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &sum, 1, MPIU_REAL, MPIU_SUM, comm));
    260:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &chargesum, 1, MPIU_REAL, MPIU_SUM, comm));
    261:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
    262:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;
    263:   if (lgEmax < user->drawlgEmin) lgEmax = user->drawlgEmin;

    265:   PetscDS ds;
    266:   Vec     phi;

    268:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
    269:   PetscCall(DMGetDS(user->dmPot, &ds));
    270:   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
    271:   PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user));
    272:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

    274:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    275:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    276:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    277:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    278:   PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
    279:   PetscCall(PetscDrawLGDraw(user->drawlgE));
    280:   PetscDraw draw;
    281:   PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
    282:   PetscCall(PetscDrawSave(draw));

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

    290: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
    291: {
    292:   DM        sw;
    293:   PetscReal pmoments[16]; /* \int f, \int v f, \int v^2 f */

    295:   PetscFunctionBeginUser;
    296:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
    297:   PetscCall(TSGetDM(ts, &sw));

    299:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));

    301:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3]));
    302:   PetscFunctionReturn(PETSC_SUCCESS);
    303: }

    305: static PetscErrorCode MonitorParticles(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
    306: {
    307:   AppCtx            *user = (AppCtx *)ctx;
    308:   DM                 sw;
    309:   PetscInt           Np, dim;
    310:   const PetscReal   *x, *v, *E;
    311:   const PetscScalar *w;
    312:   PetscViewer        viewer;
    313:   char               filename[64];
    314:   MPI_Comm           comm;

    316:   PetscFunctionBeginUser;
    317:   if (step < 0 || step % user->ostep != 0) PetscFunctionReturn(PETSC_SUCCESS);
    318:   PetscCall(TSGetDM(ts, &sw));
    319:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
    320:   PetscCall(DMGetDimension(sw, &dim));
    321:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

    323:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    324:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
    325:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
    326:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));

    328:   if (user->particle_monitor_prefix[0]) {
    329:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "%s_step_%d.bin", user->particle_monitor_prefix, (int)step));
    330:   } else {
    331:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "particles_step_%d.bin", (int)step));
    332:   }
    333:   PetscCall(PetscViewerBinaryOpen(comm, filename, FILE_MODE_WRITE, &viewer));

    335:   {
    336:     Vec vx, vv, vw, vE;
    337:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, x, &vx));
    338:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, v, &vv));
    339:     PetscCall(VecCreateMPIWithArray(comm, 1, Np, PETSC_DECIDE, w, &vw));
    340:     PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, E, &vE));

    342:     PetscCall(VecView(vx, viewer));
    343:     PetscCall(VecView(vv, viewer));
    344:     PetscCall(VecView(vw, viewer));
    345:     PetscCall(VecView(vE, viewer));

    347:     PetscCall(VecDestroy(&vx));
    348:     PetscCall(VecDestroy(&vv));
    349:     PetscCall(VecDestroy(&vw));
    350:     PetscCall(VecDestroy(&vE));
    351:   }
    352:   PetscCall(PetscViewerDestroy(&viewer));

    354:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    355:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
    356:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
    357:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    358:   PetscFunctionReturn(PETSC_SUCCESS);
    359: }

    361: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
    362: {
    363:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (4. * PETSC_PI * 4. * PETSC_PI);
    364:   return PETSC_SUCCESS;
    365: }
    366: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
    367: {
    368:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
    369:   return PETSC_SUCCESS;
    370: }

    372: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
    373: {
    374:   const PetscReal alpha = scale ? scale[0] : 0.0;
    375:   const PetscReal k     = scale ? scale[1] : 1.;
    376:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
    377:   return PETSC_SUCCESS;
    378: }

    380: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
    381: {
    382:   const PetscReal alpha = scale ? scale[0] : 0.;
    383:   const PetscReal k     = scale ? scale[1] : 1.;
    384:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
    385:   return PETSC_SUCCESS;
    386: }

    388: PetscErrorCode PetscPDFCosine3D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
    389: {
    390:   const PetscReal alpha = scale ? scale[0] : 0.;
    391:   const PetscReal k     = scale ? scale[1] : 1.;
    392:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
    393:   return PETSC_SUCCESS;
    394: }

    396: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
    397: {
    398:   PetscFE        fe;
    399:   DMPolytopeType ct;
    400:   PetscInt       dim, cStart;
    401:   const char    *prefix = "v";
    402:   AppCtx        *user;

    404:   PetscFunctionBegin;
    405:   PetscCall(DMGetDimension(sw, &dim));
    406:   PetscCall(DMGetApplicationContext(sw, &user));
    407:   if (dim == 3 && user->use_landau_velocity_space) {
    408:     LandauCtx *ctx;
    409:     Vec        X;
    410:     Mat        J;

    412:     PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, prefix, &X, &J, &user->landau_pack));
    413:     PetscCall(DMGetApplicationContext(user->landau_pack, &ctx));
    414:     *vdm = ctx->plex[0];
    415:     PetscCall(PetscObjectReference((PetscObject)*vdm));
    416:     PetscCall(VecDestroy(&X));
    417:     PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
    418:   } else {
    419:     PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
    420:     PetscCall(DMSetType(*vdm, DMPLEX));
    421:     PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
    422:     PetscCall(DMSetFromOptions(*vdm));
    423:     PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
    424:   }
    425:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

    427:   PetscCall(DMGetDimension(*vdm, &dim));
    428:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
    429:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
    430:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
    431:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
    432:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
    433:   PetscCall(DMCreateDS(*vdm));
    434:   PetscCall(PetscFEDestroy(&fe));
    435:   PetscFunctionReturn(PETSC_SUCCESS);
    436: }

    438: /*
    439:   InitializeParticles_Centroid - Initialize a regular grid of particles.

    441:   Input Parameters:
    442: + sw      - The `DMSWARM`
    443: - force1D - Treat the spatial domain as 1D

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

    448:   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
    449:   and velocity meshes.
    450: */
    451: static PetscErrorCode InitializeParticles_Centroid(DM sw)
    452: {
    453:   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
    454:   DMSwarmCellDM celldm;
    455:   DM            xdm, vdm;
    456:   PetscReal     vmin[3], vmax[3];
    457:   PetscReal    *x, *v;
    458:   PetscInt     *species, *cellid;
    459:   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
    460:   PetscBool     flg;
    461:   MPI_Comm      comm;
    462:   const char   *cellidname;

    464:   PetscFunctionBegin;
    465:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

    467:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
    468:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
    469:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
    470:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
    471:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
    472:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
    473:   PetscOptionsEnd();
    474:   debug = swarm->printCoords;

    476:   PetscCall(DMGetDimension(sw, &dim));
    477:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
    478:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

    480:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
    481:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    482:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

    484:   // One particle per centroid on the tensor product grid
    485:   Npc = (vcEnd - vcStart) * Ns;
    486:   Np  = (xcEnd - xcStart) * Npc;
    487:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
    488:   if (debug) {
    489:     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
    490:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
    491:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
    492:     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
    493:     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
    494:     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
    495:   }

    497:   // Set species and cellid
    498:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
    499:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
    500:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    501:   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
    502:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
    503:     for (PetscInt s = 0; s < Ns; ++s) {
    504:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
    505:         species[p] = s;
    506:         cellid[p]  = c;
    507:       }
    508:     }
    509:   }
    510:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    511:   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));

    513:   // Set particle coordinates
    514:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    515:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
    516:   PetscCall(DMSwarmSortGetAccess(sw));
    517:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
    518:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
    519:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
    520:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
    521:     const PetscInt cell = c + xcStart;
    522:     PetscInt      *pidx, Npc;
    523:     PetscReal      centroid[3], volume;

    525:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    526:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
    527:     for (PetscInt s = 0; s < Ns; ++s) {
    528:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
    529:         const PetscInt p  = pidx[q * Ns + s];
    530:         const PetscInt vc = vcStart + q;
    531:         PetscReal      vcentroid[3], vvolume;

    533:         PetscCall(DMPlexComputeCellGeometryFVM(vdm, vc, &vvolume, vcentroid, NULL));
    534:         for (PetscInt d = 0; d < dim; ++d) {
    535:           x[p * dim + d] = centroid[d];
    536:           v[p * dim + d] = vcentroid[d];
    537:         }

    539:         if (debug > 1) {
    540:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
    541:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
    542:           for (PetscInt d = 0; d < dim; ++d) {
    543:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    544:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
    545:           }
    546:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
    547:           for (PetscInt d = 0; d < dim; ++d) {
    548:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    549:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
    550:           }
    551:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
    552:         }
    553:       }
    554:     }
    555:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    556:   }
    557:   PetscCall(DMSwarmSortRestoreAccess(sw));
    558:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    559:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
    560:   PetscFunctionReturn(PETSC_SUCCESS);
    561: }

    563: /*
    564:   InitializeWeights - Compute weight for each local particle

    566:   Input Parameters:
    567: + sw          - The `DMSwarm`
    568: . totalWeight - The sum of all particle weights
    569: . func        - The PDF for the particle spatial distribution
    570: - param       - The PDF parameters

    572:   Notes:
    573:   The PDF for velocity is assumed to be a Gaussian

    575:   The particle weights are returned in the `w_q` field of `sw`.
    576: */
    577: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
    578: {
    579:   DM               xdm, vdm;
    580:   DMSwarmCellDM    celldm;
    581:   PetscScalar     *weight;
    582:   PetscQuadrature  xquad;
    583:   const PetscReal *xq, *xwq;
    584:   const PetscInt   order = 5;
    585:   PetscReal        xi0[3];
    586:   PetscReal        xwtot = 0., pwtot = 0.;
    587:   PetscInt         xNq;
    588:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
    589:   MPI_Comm         comm;
    590:   PetscMPIInt      rank;

    592:   PetscFunctionBegin;
    593:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
    594:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
    595:   PetscCall(DMGetDimension(sw, &dim));
    596:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
    597:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
    598:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
    599:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
    600:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    601:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

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

    608:   // Integrate the density function to get the weights of particles in each cell
    609:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
    610:   PetscCall(DMSwarmSortGetAccess(sw));
    611:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    612:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
    613:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
    614:     PetscInt          *pidx, Npc;
    615:     PetscInt           xNc;
    616:     const PetscScalar *xarray;
    617:     PetscScalar       *xcoords = NULL;
    618:     PetscBool          xisDG;

    620:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
    621:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    622:     PetscCheck(Npc == (vcEnd - vcStart) * Ns, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Number of particles %" PetscInt_FMT " in cell (rank %d) != %" PetscInt_FMT " number of velocity vertices", Npc, rank, (vcEnd - vcStart) * Ns);
    623:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
    624:     for (PetscInt q = 0; q < xNq; ++q) {
    625:       // Transform quadrature points from ref space to real space
    626:       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
    627:       // Get probability density at quad point
    628:       //   No need to scale xqr since PDF will be periodic
    629:       PetscCall((*func)(xqr, param, &xden));
    630:       xw += xden * (xwq[q] * xdetJ);
    631:     }
    632:     xwtot += xw;
    633:     if (debug) {
    634:       IS              globalOrdering;
    635:       const PetscInt *ordering;

    637:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
    638:       PetscCall(ISGetIndices(globalOrdering, &ordering));
    639:       PetscCall(PetscSynchronizedPrintf(comm, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(xcoords[0]), (double)PetscRealPart(xcoords[0 + dim]), (double)xw));
    640:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
    641:     }
    642:     // Set weights to be Gaussian in velocity cells
    643:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
    644:       const PetscInt     p  = pidx[vc * Ns + 0];
    645:       PetscReal          vw = 0.;
    646:       PetscInt           vNc;
    647:       const PetscScalar *varray;
    648:       PetscScalar       *vcoords = NULL;
    649:       PetscBool          visDG;

    651:       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
    652:       PetscCheck(vNc > 0 && vNc % dim == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Velocity cell %" PetscInt_FMT " has invalid coordinates (vNc=%" PetscInt_FMT ", dim=%" PetscInt_FMT ")", vc, vNc, dim);
    653:       {
    654:         PetscReal vmin[3] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
    655:         PetscReal vmax[3] = {-PETSC_MAX_REAL, -PETSC_MAX_REAL, -PETSC_MAX_REAL};
    656:         PetscInt  numVert = vNc / dim;
    657:         for (PetscInt i = 0; i < numVert; ++i) {
    658:           for (PetscInt d = 0; d < dim; ++d) {
    659:             PetscReal coord = PetscRealPart(vcoords[i * dim + d]);
    660:             vmin[d]         = PetscMin(vmin[d], coord);
    661:             vmax[d]         = PetscMax(vmax[d], coord);
    662:           }
    663:         }
    664:         vw = 1.0;
    665:         for (PetscInt d = 0; d < dim; ++d) vw *= 0.5 * (PetscErfReal(vmax[d] / PetscSqrtReal(2.)) - PetscErfReal(vmin[d] / PetscSqrtReal(2.)));
    666:         PetscCheck(PetscIsNormalReal(vw), PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " velocity weight is not normal: vw=%g, vmin=(%g,%g,%g), vmax=(%g,%g,%g)", p, vw, vmin[0], vmin[1], vmin[2], vmax[0], vmax[1], vmax[2]);
    667:       }

    669:       weight[p] = totalWeight * vw * xw;
    670:       pwtot += weight[p];
    671:       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 10: weight=%g, xw=%g, vw=%g, totalWeight=%g", p, weight[p], xw, vw, totalWeight);
    672:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
    673:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
    674:     }
    675:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
    676:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    677:   }
    678:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    679:   PetscCall(DMSwarmSortRestoreAccess(sw));
    680:   PetscCall(PetscQuadratureDestroy(&xquad));

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

    685:     PetscCall(PetscSynchronizedFlush(comm, NULL));
    686:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    687:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
    688:   }
    689:   PetscFunctionReturn(PETSC_SUCCESS);
    690: }

    692: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
    693: {
    694:   PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
    695:   PetscInt  dim;

    697:   PetscFunctionBegin;
    698:   PetscCall(DMGetDimension(sw, &dim));
    699:   PetscCall(InitializeParticles_Centroid(sw));
    700:   PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : (dim == 2 ? PetscPDFCosine2D : PetscPDFCosine3D), scale));
    701:   PetscFunctionReturn(PETSC_SUCCESS);
    702: }

    704: static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
    705: {
    706:   DM         dm;
    707:   PetscInt  *species;
    708:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
    709:   PetscInt   Np, dim;

    711:   PetscFunctionBegin;
    712:   PetscCall(DMSwarmGetCellDM(sw, &dm));
    713:   PetscCall(DMGetDimension(sw, &dim));
    714:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    715:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
    716:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    717:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    718:   for (PetscInt p = 0; p < Np; ++p) {
    719:     totalWeight += weight[p];
    720:     totalCharge += user->charges[species[p]] * weight[p];
    721:   }
    722:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    723:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    724:   {
    725:     Parameter *param;
    726:     PetscReal  Area;

    728:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
    729:     switch (dim) {
    730:     case 1:
    731:       Area = (gmax[0] - gmin[0]);
    732:       break;
    733:     case 2:
    734:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
    735:       break;
    736:     case 3:
    737:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
    738:       break;
    739:     default:
    740:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
    741:     }
    742:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    743:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    744:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)user->charges[0], (double)global_charge, (double)Area));
    745:     param->sigma = PetscAbsReal(global_charge / (Area));

    747:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
    748:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "(x0,v0,t0,m0,q0,phi0): (%e, %e, %e, %e, %e, %e) - (P, V) = (%e, %e)\n", (double)param->x0, (double)param->v0, (double)param->t0, (double)param->m0, (double)param->q0, (double)param->phi0, (double)param->poissonNumber,
    749:                           (double)param->vlasovNumber));
    750:   }
    751:   /* Setup Constants */
    752:   {
    753:     PetscDS    ds;
    754:     Parameter *param;
    755:     PetscCall(PetscBagGetData(user->bag, (void **)&param));
    756:     PetscScalar constants[NUM_CONSTANTS];
    757:     constants[SIGMA]   = param->sigma;
    758:     constants[V0]      = param->v0;
    759:     constants[T0]      = param->t0;
    760:     constants[X0]      = param->x0;
    761:     constants[M0]      = param->m0;
    762:     constants[Q0]      = param->q0;
    763:     constants[PHI0]    = param->phi0;
    764:     constants[POISSON] = param->poissonNumber;
    765:     constants[VLASOV]  = param->vlasovNumber;
    766:     PetscCall(DMGetDS(dm, &ds));
    767:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
    768:   }
    769:   PetscFunctionReturn(PETSC_SUCCESS);
    770: }

    772: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
    773: {
    774:   PetscBag   bag;
    775:   Parameter *p;

    777:   PetscFunctionBeginUser;
    778:   /* setup PETSc parameter bag */
    779:   PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
    780:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
    781:   bag = ctx->bag;
    782:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
    783:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
    784:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
    785:   PetscCall(PetscBagRegisterScalar(bag, &p->phi0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
    786:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
    787:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
    788:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
    789:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

    791:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
    792:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
    793:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
    794:   PetscCall(PetscBagSetFromOptions(bag));
    795:   {
    796:     PetscViewer       viewer;
    797:     PetscViewerFormat format;
    798:     PetscBool         flg;

    800:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
    801:     if (flg) {
    802:       PetscCall(PetscViewerPushFormat(viewer, format));
    803:       PetscCall(PetscBagView(bag, viewer));
    804:       PetscCall(PetscViewerFlush(viewer));
    805:       PetscCall(PetscViewerPopFormat(viewer));
    806:       PetscCall(PetscViewerDestroy(&viewer));
    807:     }
    808:   }
    809:   PetscFunctionReturn(PETSC_SUCCESS);
    810: }

    812: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
    813: {
    814:   const char *prefix = "x";

    816:   PetscFunctionBeginUser;
    817:   PetscCall(DMCreate(comm, dm));
    818:   PetscCall(DMSetType(*dm, DMPLEX));
    819:   PetscCall(DMPlexSetOptionsPrefix(*dm, prefix));
    820:   PetscCall(DMSetFromOptions(*dm));
    821:   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
    822:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

    824:   // Cache the mesh geometry
    825:   DMField         coordField;
    826:   IS              cellIS;
    827:   PetscQuadrature quad;
    828:   PetscReal      *wt, *pt;
    829:   PetscInt        cdim, cStart, cEnd;

    831:   PetscCall(DMGetCoordinateField(*dm, &coordField));
    832:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
    833:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
    834:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
    835:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
    836:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
    837:   PetscCall(PetscMalloc1(1, &wt));
    838:   PetscCall(PetscMalloc1(cdim, &pt));
    839:   wt[0] = 1.;
    840:   for (PetscInt d = 0; d < cdim; ++d) pt[d] = -1.;
    841:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
    842:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom));
    843:   PetscCall(PetscQuadratureDestroy(&quad));
    844:   PetscCall(ISDestroy(&cellIS));
    845:   PetscFunctionReturn(PETSC_SUCCESS);
    846: }

    848: 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[])
    849: {
    850:   f0[0] = -constants[SIGMA];
    851: }

    853: 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[])
    854: {
    855:   PetscInt d;
    856:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
    857: }

    859: 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[])
    860: {
    861:   PetscInt d;
    862:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
    863: }

    865: static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
    866: {
    867:   PetscFE      fephi;
    868:   PetscDS      ds;
    869:   PetscBool    simplex;
    870:   PetscInt     dim;
    871:   MatNullSpace nullsp;

    873:   PetscFunctionBeginUser;
    874:   PetscCall(DMGetDimension(dm, &dim));
    875:   PetscCall(DMPlexIsSimplex(dm, &simplex));

    877:   PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
    878:   PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
    879:   PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
    880:   PetscCall(DMCreateDS(dm));
    881:   PetscCall(DMGetDS(dm, &ds));
    882:   PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
    883:   PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
    884:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
    885:   PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
    886:   PetscCall(MatNullSpaceDestroy(&nullsp));
    887:   PetscCall(PetscFEDestroy(&fephi));
    888:   PetscFunctionReturn(PETSC_SUCCESS);
    889: }

    891: static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
    892: {
    893:   SNES         snes;
    894:   Mat          J;
    895:   MatNullSpace nullSpace;

    897:   PetscFunctionBeginUser;
    898:   PetscCall(CreateFEM(dm, user));
    899:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
    900:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
    901:   PetscCall(SNESSetDM(snes, dm));
    902:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
    903:   PetscCall(SNESSetFromOptions(snes));

    905:   PetscCall(DMCreateMatrix(dm, &J));
    906:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
    907:   PetscCall(MatSetNullSpace(J, nullSpace));
    908:   PetscCall(MatNullSpaceDestroy(&nullSpace));
    909:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
    910:   PetscCall(MatDestroy(&J));

    912:   user->dmPot = dm;
    913:   PetscCall(PetscObjectReference((PetscObject)user->dmPot));

    915:   PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M));
    916:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
    917:   user->snes = snes;
    918:   PetscFunctionReturn(PETSC_SUCCESS);
    919: }

    921: static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
    922: {
    923:   DMSwarmCellDM celldm;
    924:   PetscInt      dim;

    926:   PetscFunctionBeginUser;
    927:   PetscCall(DMGetDimension(dm, &dim));
    928:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
    929:   PetscCall(DMSetType(*sw, DMSWARM));
    930:   PetscCall(DMSetDimension(*sw, dim));
    931:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
    932:   PetscCall(DMSetApplicationContext(*sw, user));

    934:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
    935:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
    936:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
    937:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));

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

    941:   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
    942:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    943:   PetscCall(DMSwarmCellDMDestroy(&celldm));

    945:   const char *vfieldnames[2] = {"w_q"};
    946:   DM          vdm;

    948:   PetscCall(CreateVelocityDM(*sw, &vdm));
    949:   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
    950:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    951:   PetscCall(DMSwarmCellDMDestroy(&celldm));
    952:   PetscCall(DMDestroy(&vdm));

    954:   DM mdm;

    956:   PetscCall(DMClone(dm, &mdm));
    957:   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
    958:   PetscCall(DMCopyDisc(dm, mdm));
    959:   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
    960:   PetscCall(DMDestroy(&mdm));
    961:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    962:   PetscCall(DMSwarmCellDMDestroy(&celldm));

    964:   PetscCall(DMSetFromOptions(*sw));
    965:   PetscCall(DMSetUp(*sw));

    967:   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
    968:   user->swarm = *sw;
    969:   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
    970:   PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
    971:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
    972:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
    973:   PetscFunctionReturn(PETSC_SUCCESS);
    974: }

    976: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
    977: {
    978:   DM         dm;
    979:   AppCtx    *user;
    980:   PetscDS    ds;
    981:   PetscFE    fe;
    982:   KSP        ksp;
    983:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
    984:   Vec        rho;         // Charge density, M^{-1} rhoRhs
    985:   Vec        phi, locPhi; // Potential
    986:   Vec        f;           // Particle weights
    987:   PetscReal *coords;
    988:   PetscInt   dim, cStart, cEnd, Np;

    990:   PetscFunctionBegin;
    991:   PetscCall(DMGetApplicationContext(sw, (void *)&user));
    992:   PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
    993:   PetscCall(DMGetDimension(sw, &dim));
    994:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

    996:   PetscCall(SNESGetDM(snes, &dm));
    997:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
    998:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
    999:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
    1000:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
    1001:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

    1003:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
    1004:   PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
    1005:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

    1007:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
    1008:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

    1010:   // Low-pass filter rhoRhs
    1011:   PetscInt window = 0;
    1012:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-rho_average", &window, NULL));
    1013:   if (window) {
    1014:     PetscScalar *a;
    1015:     PetscInt     n;
    1016:     PetscReal    width = 2. * window + 1.;

    1018:     // This will only work for P_1
    1019:     //   This works because integration against a test function is linear
    1020:     //   Do a real integral against weight function for higher order
    1021:     PetscCall(VecGetLocalSize(rhoRhs, &n));
    1022:     PetscCall(VecGetArrayWrite(rhoRhs, &a));
    1023:     for (PetscInt i = 0; i < n; ++i) {
    1024:       PetscScalar avg = a[i];
    1025:       for (PetscInt j = 1; j <= window; ++j) avg += a[(i - j + n) % n] + a[(i + j) % n];
    1026:       a[i] = avg / width;
    1027:       //a[i] = (a[(i - 1 + n) % n] + a[i] + a[(i + 1) % n]) / 3.;
    1028:     }
    1029:     PetscCall(VecRestoreArrayWrite(rhoRhs, &a));
    1030:   }

    1032:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
    1033:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
    1034:   PetscCall(KSPSetOperators(ksp, user->M, user->M));
    1035:   PetscCall(KSPSetFromOptions(ksp));
    1036:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

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

    1040:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
    1041:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
    1042:   PetscCall(KSPDestroy(&ksp));

    1044:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
    1045:   PetscCall(VecSet(phi, 0.0));
    1046:   PetscCall(SNESSolve(snes, rhoRhs, phi));
    1047:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
    1048:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

    1050:   PetscCall(DMGetLocalVector(dm, &locPhi));
    1051:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
    1052:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
    1053:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
    1054:   PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));

    1056:   PetscCall(DMGetDS(dm, &ds));
    1057:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
    1058:   PetscCall(DMSwarmSortGetAccess(sw));
    1059:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
    1060:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

    1062:   PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
    1063:   PetscTabulation tab;
    1064:   PetscReal      *pcoord, *refcoord;
    1065:   PetscFEGeom    *chunkgeom = NULL;
    1066:   PetscInt        maxNcp    = 0;

    1068:   for (PetscInt c = cStart; c < cEnd; ++c) {
    1069:     PetscInt Ncp;

    1071:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
    1072:     maxNcp = PetscMax(maxNcp, Ncp);
    1073:   }
    1074:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    1075:   PetscCall(PetscArrayzero(refcoord, maxNcp * dim));
    1076:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    1077:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
    1078:   for (PetscInt c = cStart; c < cEnd; ++c) {
    1079:     PetscScalar *clPhi = NULL;
    1080:     PetscInt    *points;
    1081:     PetscInt     Ncp;

    1083:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
    1084:     for (PetscInt cp = 0; cp < Ncp; ++cp)
    1085:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
    1086:     {
    1087:       PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
    1088:       for (PetscInt i = 0; i < Ncp; ++i) {
    1089:         const PetscReal x0[3] = {-1., -1., -1.};
    1090:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
    1091:       }
    1092:     }
    1093:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
    1094:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    1095:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
    1096:       const PetscReal *basisDer = tab->T[1];
    1097:       const PetscInt   p        = points[cp];

    1099:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
    1100:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
    1101:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
    1102:     }
    1103:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    1104:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
    1105:   }
    1106:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    1107:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    1108:   PetscCall(PetscTabulationDestroy(&tab));
    1109:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1110:   PetscCall(DMSwarmSortRestoreAccess(sw));
    1111:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
    1112:   PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
    1113:   PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
    1114:   PetscFunctionReturn(PETSC_SUCCESS);
    1115: }

    1117: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
    1118: {
    1119:   AppCtx    *user;
    1120:   Mat        M_p;
    1121:   PetscReal *E;
    1122:   PetscInt   dim, Np;

    1124:   PetscFunctionBegin;
    1127:   PetscCall(DMGetDimension(sw, &dim));
    1128:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1129:   PetscCall(DMGetApplicationContext(sw, &user));

    1131:   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
    1132:   // TODO: Could share sort context with space cellDM
    1133:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
    1134:   PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p));
    1135:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));

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

    1141:   PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
    1142:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    1143:   PetscCall(MatDestroy(&M_p));
    1144:   PetscFunctionReturn(PETSC_SUCCESS);
    1145: }

    1147: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
    1148: {
    1149:   DM                 sw;
    1150:   SNES               snes = ((AppCtx *)ctx)->snes;
    1151:   const PetscScalar *u;
    1152:   PetscScalar       *g;
    1153:   PetscReal         *E, m_p = 1., q_p = -1.;
    1154:   PetscInt           dim, d, Np, p;

    1156:   PetscFunctionBeginUser;
    1157:   PetscCall(TSGetDM(ts, &sw));
    1158:   PetscCall(ComputeFieldAtParticles(snes, sw));

    1160:   PetscCall(DMGetDimension(sw, &dim));
    1161:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1162:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    1163:   PetscCall(VecGetArrayRead(U, &u));
    1164:   PetscCall(VecGetArray(G, &g));
    1165:   Np /= 2 * dim;
    1166:   for (p = 0; p < Np; ++p) {
    1167:     for (d = 0; d < dim; ++d) {
    1168:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
    1169:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
    1170:     }
    1171:   }
    1172:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    1173:   PetscCall(VecRestoreArrayRead(U, &u));
    1174:   PetscCall(VecRestoreArray(G, &g));
    1175:   PetscFunctionReturn(PETSC_SUCCESS);
    1176: }

    1178: /* J_{ij} = dF_i/dx_j
    1179:    J_p = (  0   1)
    1180:          (-w^2  0)
    1181:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
    1182:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
    1183: */
    1184: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
    1185: {
    1186:   DM               sw;
    1187:   const PetscReal *coords, *vel;
    1188:   PetscInt         dim, d, Np, p, rStart;

    1190:   PetscFunctionBeginUser;
    1191:   PetscCall(TSGetDM(ts, &sw));
    1192:   PetscCall(DMGetDimension(sw, &dim));
    1193:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1194:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
    1195:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1196:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    1197:   Np /= 2 * dim;
    1198:   for (p = 0; p < Np; ++p) {
    1199:     // TODO This is not right because dv/dx has the electric field in it
    1200:     PetscScalar vals[4] = {0., 1., -1, 0.};

    1202:     for (d = 0; d < dim; ++d) {
    1203:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
    1204:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
    1205:     }
    1206:   }
    1207:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1208:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    1209:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
    1210:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
    1211:   PetscFunctionReturn(PETSC_SUCCESS);
    1212: }

    1214: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
    1215: {
    1216:   AppCtx            *user = (AppCtx *)ctx;
    1217:   DM                 sw;
    1218:   const PetscScalar *v;
    1219:   PetscScalar       *xres;
    1220:   PetscInt           Np, p, d, dim;

    1222:   PetscFunctionBeginUser;
    1223:   PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
    1224:   PetscCall(TSGetDM(ts, &sw));
    1225:   PetscCall(DMGetDimension(sw, &dim));
    1226:   PetscCall(VecGetLocalSize(Xres, &Np));
    1227:   PetscCall(VecGetArrayRead(V, &v));
    1228:   PetscCall(VecGetArray(Xres, &xres));
    1229:   Np /= dim;
    1230:   for (p = 0; p < Np; ++p) {
    1231:     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
    1232:   }
    1233:   PetscCall(VecRestoreArrayRead(V, &v));
    1234:   PetscCall(VecRestoreArray(Xres, &xres));
    1235:   PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
    1236:   PetscFunctionReturn(PETSC_SUCCESS);
    1237: }

    1239: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
    1240: {
    1241:   DM                 sw;
    1242:   AppCtx            *user = (AppCtx *)ctx;
    1243:   SNES               snes = ((AppCtx *)ctx)->snes;
    1244:   const PetscScalar *x;
    1245:   PetscScalar       *vres;
    1246:   PetscReal         *E, m_p, q_p;
    1247:   PetscInt           Np, p, dim, d;
    1248:   Parameter         *param;

    1250:   PetscFunctionBeginUser;
    1251:   PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
    1252:   PetscCall(TSGetDM(ts, &sw));
    1253:   PetscCall(ComputeFieldAtParticles(snes, sw));

    1255:   PetscCall(DMGetDimension(sw, &dim));
    1256:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    1257:   PetscCall(PetscBagGetData(user->bag, (void **)&param));
    1258:   m_p = user->masses[0] * param->m0;
    1259:   q_p = user->charges[0] * param->q0;
    1260:   PetscCall(VecGetLocalSize(Vres, &Np));
    1261:   PetscCall(VecGetArrayRead(X, &x));
    1262:   PetscCall(VecGetArray(Vres, &vres));
    1263:   Np /= dim;
    1264:   for (p = 0; p < Np; ++p) {
    1265:     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
    1266:   }
    1267:   PetscCall(VecRestoreArrayRead(X, &x));
    1268:   /*
    1269:     Synchronized, ordered output for parallel/sequential test cases.
    1270:     In the 1D (on the 2D mesh) case, every y component should be zero.
    1271:   */
    1272:   if (user->checkVRes) {
    1273:     PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
    1274:     PetscInt  step;

    1276:     PetscCall(TSGetStepNumber(ts, &step));
    1277:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
    1278:     for (PetscInt p = 0; p < Np; ++p) {
    1279:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
    1280:       PetscCheck(PetscAbsScalar(vres[p * dim + 1]) < PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Y velocity should be 0., not %g", (double)PetscRealPart(vres[p * dim + 1]));
    1281:     }
    1282:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
    1283:   }
    1284:   PetscCall(VecRestoreArray(Vres, &vres));
    1285:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    1286:   PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
    1287:   PetscFunctionReturn(PETSC_SUCCESS);
    1288: }

    1290: /* Discrete Gradients Formulation: S, F, gradF (G) */
    1291: PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
    1292: {
    1293:   PetscScalar vals[4] = {0., 1., -1., 0.};
    1294:   DM          sw;
    1295:   PetscInt    dim, d, Np, p, rStart;

    1297:   PetscFunctionBeginUser;
    1298:   PetscCall(TSGetDM(ts, &sw));
    1299:   PetscCall(DMGetDimension(sw, &dim));
    1300:   PetscCall(VecGetLocalSize(U, &Np));
    1301:   PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
    1302:   Np /= 2 * dim;
    1303:   for (p = 0; p < Np; ++p) {
    1304:     for (d = 0; d < dim; ++d) {
    1305:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
    1306:       PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
    1307:     }
    1308:   }
    1309:   PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
    1310:   PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
    1311:   PetscFunctionReturn(PETSC_SUCCESS);
    1312: }

    1314: PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
    1315: {
    1316:   AppCtx            *user = (AppCtx *)ctx;
    1317:   DM                 sw;
    1318:   Vec                phi;
    1319:   const PetscScalar *u;
    1320:   PetscInt           dim, Np, cStart, cEnd;
    1321:   PetscReal         *vel, *coords, m_p = 1.;

    1323:   PetscFunctionBeginUser;
    1324:   PetscCall(TSGetDM(ts, &sw));
    1325:   PetscCall(DMGetDimension(sw, &dim));
    1326:   PetscCall(DMPlexGetHeightStratum(user->dmPot, 0, &cStart, &cEnd));

    1328:   PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
    1329:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
    1330:   PetscCall(computeFieldEnergy(user->dmPot, phi, F));
    1331:   PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));

    1333:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1334:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    1335:   PetscCall(DMSwarmSortGetAccess(sw));
    1336:   PetscCall(VecGetArrayRead(U, &u));
    1337:   PetscCall(VecGetLocalSize(U, &Np));
    1338:   Np /= 2 * dim;
    1339:   for (PetscInt c = cStart; c < cEnd; ++c) {
    1340:     PetscInt *points;
    1341:     PetscInt  Ncp;

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

    1348:       *F += 0.5 * m_p * v2;
    1349:     }
    1350:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
    1351:   }
    1352:   PetscCall(VecRestoreArrayRead(U, &u));
    1353:   PetscCall(DMSwarmSortRestoreAccess(sw));
    1354:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1355:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    1356:   PetscFunctionReturn(PETSC_SUCCESS);
    1357: }

    1359: /* dF/dx = q E   dF/dv = v */
    1360: PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
    1361: {
    1362:   DM                 sw;
    1363:   SNES               snes = ((AppCtx *)ctx)->snes;
    1364:   const PetscReal   *coords, *vel, *E;
    1365:   const PetscScalar *u;
    1366:   PetscScalar       *g;
    1367:   PetscReal          m_p = 1., q_p = -1.;
    1368:   PetscInt           dim, d, Np, p;

    1370:   PetscFunctionBeginUser;
    1371:   PetscCall(TSGetDM(ts, &sw));
    1372:   PetscCall(DMGetDimension(sw, &dim));
    1373:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1374:   PetscCall(VecGetArrayRead(U, &u));
    1375:   PetscCall(VecGetArray(G, &g));

    1377:   PetscLogEvent COMPUTEFIELD;
    1378:   PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
    1379:   PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
    1380:   PetscCall(ComputeFieldAtParticles(snes, sw));
    1381:   PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
    1382:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1383:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    1384:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    1385:   for (p = 0; p < Np; ++p) {
    1386:     for (d = 0; d < dim; ++d) {
    1387:       g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
    1388:       g[(p * 2 + 1) * dim + d] = u[(p * 2 + 1) * dim + d];
    1389:     }
    1390:   }
    1391:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    1392:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    1393:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    1394:   PetscCall(VecRestoreArrayRead(U, &u));
    1395:   PetscCall(VecRestoreArray(G, &g));
    1396:   PetscFunctionReturn(PETSC_SUCCESS);
    1397: }

    1399: static PetscErrorCode CreateSolution(TS ts)
    1400: {
    1401:   DM       sw;
    1402:   Vec      u;
    1403:   PetscInt dim, Np;

    1405:   PetscFunctionBegin;
    1406:   PetscCall(TSGetDM(ts, &sw));
    1407:   PetscCall(DMGetDimension(sw, &dim));
    1408:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1409:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
    1410:   PetscCall(VecSetBlockSize(u, dim));
    1411:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
    1412:   PetscCall(VecSetUp(u));
    1413:   PetscCall(TSSetSolution(ts, u));
    1414:   PetscCall(VecDestroy(&u));
    1415:   PetscFunctionReturn(PETSC_SUCCESS);
    1416: }

    1418: static PetscErrorCode SetProblem(TS ts)
    1419: {
    1420:   AppCtx *user;
    1421:   DM      sw;

    1423:   PetscFunctionBegin;
    1424:   PetscCall(TSGetDM(ts, &sw));
    1425:   PetscCall(DMGetApplicationContext(sw, (void **)&user));
    1426:   // Define unified system for (X, V)
    1427:   {
    1428:     Mat      J;
    1429:     PetscInt dim, Np;

    1431:     PetscCall(DMGetDimension(sw, &dim));
    1432:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1433:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
    1434:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
    1435:     PetscCall(MatSetBlockSize(J, 2 * dim));
    1436:     PetscCall(MatSetFromOptions(J));
    1437:     PetscCall(MatSetUp(J));
    1438:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
    1439:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
    1440:     PetscCall(MatDestroy(&J));
    1441:   }
    1442:   /* Define split system for X and V */
    1443:   {
    1444:     Vec             u;
    1445:     IS              isx, isv, istmp;
    1446:     const PetscInt *idx;
    1447:     PetscInt        dim, Np, rstart;

    1449:     PetscCall(TSGetSolution(ts, &u));
    1450:     PetscCall(DMGetDimension(sw, &dim));
    1451:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1452:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
    1453:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
    1454:     PetscCall(ISGetIndices(istmp, &idx));
    1455:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
    1456:     PetscCall(ISRestoreIndices(istmp, &idx));
    1457:     PetscCall(ISDestroy(&istmp));
    1458:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
    1459:     PetscCall(ISGetIndices(istmp, &idx));
    1460:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
    1461:     PetscCall(ISRestoreIndices(istmp, &idx));
    1462:     PetscCall(ISDestroy(&istmp));
    1463:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
    1464:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
    1465:     PetscCall(ISDestroy(&isx));
    1466:     PetscCall(ISDestroy(&isv));
    1467:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
    1468:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
    1469:   }
    1470:   // Define symplectic formulation U_t = S . G, where G = grad F
    1471:   {
    1472:     PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, user));
    1473:   }
    1474:   PetscFunctionReturn(PETSC_SUCCESS);
    1475: }

    1477: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
    1478: {
    1479:   DM        sw;
    1480:   Vec       u;
    1481:   PetscReal t, maxt, dt;
    1482:   PetscInt  n, maxn;

    1484:   PetscFunctionBegin;
    1485:   PetscCall(TSGetDM(ts, &sw));
    1486:   PetscCall(TSGetTime(ts, &t));
    1487:   PetscCall(TSGetMaxTime(ts, &maxt));
    1488:   PetscCall(TSGetTimeStep(ts, &dt));
    1489:   PetscCall(TSGetStepNumber(ts, &n));
    1490:   PetscCall(TSGetMaxSteps(ts, &maxn));

    1492:   PetscCall(TSReset(ts));
    1493:   PetscCall(TSSetDM(ts, sw));
    1494:   PetscCall(TSSetFromOptions(ts));
    1495:   PetscCall(TSSetTime(ts, t));
    1496:   PetscCall(TSSetMaxTime(ts, maxt));
    1497:   PetscCall(TSSetTimeStep(ts, dt));
    1498:   PetscCall(TSSetStepNumber(ts, n));
    1499:   PetscCall(TSSetMaxSteps(ts, maxn));

    1501:   PetscCall(CreateSolution(ts));
    1502:   PetscCall(SetProblem(ts));
    1503:   PetscCall(TSGetSolution(ts, &u));
    1504:   PetscFunctionReturn(PETSC_SUCCESS);
    1505: }

    1507: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
    1508: {
    1509:   DM        sw, cdm;
    1510:   PetscInt  Np;
    1511:   PetscReal low[2], high[2];
    1512:   AppCtx   *user = (AppCtx *)ctx;

    1514:   sw = user->swarm;
    1515:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
    1516:   // Get the bounding box so we can equally space the particles
    1517:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
    1518:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1519:   // shift it by h/2 so nothing is initialized directly on a boundary
    1520:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
    1521:   x[1] = 0.;
    1522:   return PETSC_SUCCESS;
    1523: }

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

    1528:   Input Parameters:
    1529: + ts         - The TS
    1530: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

    1532:   Output Parameters:
    1533: . u - The initialized solution vector

    1535:   Level: advanced

    1537: .seealso: InitializeSolve()
    1538: */
    1539: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
    1540: {
    1541:   DM       sw;
    1542:   Vec      u, gc, gv;
    1543:   IS       isx, isv;
    1544:   PetscInt dim;
    1545:   AppCtx  *user;

    1547:   PetscFunctionBeginUser;
    1548:   PetscCall(TSGetDM(ts, &sw));
    1549:   PetscCall(DMGetApplicationContext(sw, &user));
    1550:   PetscCall(DMGetDimension(sw, &dim));
    1551:   if (useInitial) {
    1552:     PetscCall(InitializeParticles_PerturbedWeights(sw, user));
    1553:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
    1554:     PetscCall(DMSwarmTSRedistribute(ts));
    1555:   }
    1556:   PetscCall(DMSetUp(sw));
    1557:   PetscCall(TSGetSolution(ts, &u));
    1558:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
    1559:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
    1560:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    1561:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
    1562:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
    1563:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
    1564:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    1565:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
    1566:   PetscFunctionReturn(PETSC_SUCCESS);
    1567: }

    1569: static PetscErrorCode InitializeSolve(TS ts, Vec u)
    1570: {
    1571:   PetscFunctionBegin;
    1572:   PetscCall(TSSetSolution(ts, u));
    1573:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
    1574:   PetscFunctionReturn(PETSC_SUCCESS);
    1575: }

    1577: static PetscErrorCode MigrateParticles(TS ts)
    1578: {
    1579:   DM               sw, cdm;
    1580:   const PetscReal *L;
    1581:   AppCtx          *ctx;

    1583:   PetscFunctionBeginUser;
    1584:   PetscCall(TSGetDM(ts, &sw));
    1585:   PetscCall(DMGetApplicationContext(sw, &ctx));
    1586:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
    1587:   {
    1588:     Vec        u, gc, gv, position, momentum;
    1589:     IS         isx, isv;
    1590:     PetscReal *pos, *mom;

    1592:     PetscCall(TSGetSolution(ts, &u));
    1593:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
    1594:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
    1595:     PetscCall(VecGetSubVector(u, isx, &position));
    1596:     PetscCall(VecGetSubVector(u, isv, &momentum));
    1597:     PetscCall(VecGetArray(position, &pos));
    1598:     PetscCall(VecGetArray(momentum, &mom));
    1599:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    1600:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
    1601:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
    1602:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

    1604:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
    1605:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
    1606:     if ((L[0] || L[1]) >= 0.) {
    1607:       PetscReal *x, *v, upper[3], lower[3];
    1608:       PetscInt   Np, dim;

    1610:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
    1611:       PetscCall(DMGetDimension(cdm, &dim));
    1612:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
    1613:       PetscCall(VecGetArray(gc, &x));
    1614:       PetscCall(VecGetArray(gv, &v));
    1615:       for (PetscInt p = 0; p < Np; ++p) {
    1616:         for (PetscInt d = 0; d < dim; ++d) {
    1617:           if (pos[p * dim + d] < lower[d]) {
    1618:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
    1619:           } else if (pos[p * dim + d] > upper[d]) {
    1620:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
    1621:           } else {
    1622:             x[p * dim + d] = pos[p * dim + d];
    1623:           }
    1624:           PetscCheck(x[p * dim + d] >= lower[d] && x[p * dim + d] <= upper[d], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "p: %" PetscInt_FMT "x[%" PetscInt_FMT "] %g", p, d, (double)x[p * dim + d]);
    1625:           v[p * dim + d] = mom[p * dim + d];
    1626:         }
    1627:       }
    1628:       PetscCall(VecRestoreArray(gc, &x));
    1629:       PetscCall(VecRestoreArray(gv, &v));
    1630:     }
    1631:     PetscCall(VecRestoreArray(position, &pos));
    1632:     PetscCall(VecRestoreArray(momentum, &mom));
    1633:     PetscCall(VecRestoreSubVector(u, isx, &position));
    1634:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
    1635:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
    1636:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    1637:   }
    1638:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
    1639:   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
    1640:   PetscCall(DMSwarmTSRedistribute(ts));
    1641:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
    1642:   {
    1643:     const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
    1644:     PetscCall(DMSwarmVectorDefineFields(sw, 2, fieldnames));
    1645:   }
    1646:   PetscFunctionReturn(PETSC_SUCCESS);
    1647: }

    1649: int main(int argc, char **argv)
    1650: {
    1651:   DM        dm, sw;
    1652:   TS        ts;
    1653:   Vec       u;
    1654:   PetscReal dt;
    1655:   PetscInt  maxn;
    1656:   AppCtx    user;

    1658:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
    1659:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
    1660:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
    1661:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
    1662:   PetscCall(CreatePoisson(dm, &user));
    1663:   PetscCall(CreateSwarm(dm, &user, &sw));
    1664:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
    1665:   PetscCall(InitializeConstants(sw, &user));
    1666:   PetscCall(DMSetApplicationContext(sw, &user));

    1668:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
    1669:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
    1670:   PetscCall(TSSetDM(ts, sw));
    1671:   PetscCall(TSSetMaxTime(ts, 0.1));
    1672:   PetscCall(TSSetTimeStep(ts, 0.00001));
    1673:   PetscCall(TSSetMaxSteps(ts, 100));
    1674:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

    1676:   if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
    1677:   if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
    1678:   if (user.particle_monitor) PetscCall(TSMonitorSet(ts, MonitorParticles, &user, NULL));

    1680:   PetscCall(TSSetFromOptions(ts));
    1681:   PetscCall(TSGetTimeStep(ts, &dt));
    1682:   PetscCall(TSGetMaxSteps(ts, &maxn));
    1683:   user.steps    = maxn;
    1684:   user.stepSize = dt;
    1685:   PetscCall(SetupContext(dm, sw, &user));
    1686:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
    1687:   PetscCall(TSSetPostStep(ts, MigrateParticles));
    1688:   PetscCall(CreateSolution(ts));
    1689:   PetscCall(TSGetSolution(ts, &u));
    1690:   PetscCall(TSComputeInitialCondition(ts, u));
    1691:   PetscCall(CheckNonNegativeWeights(sw, &user));
    1692:   PetscCall(TSSolve(ts, NULL));

    1694:   PetscCall(SNESDestroy(&user.snes));
    1695:   PetscCall(DMDestroy(&user.dmPot));
    1696:   PetscCall(MatDestroy(&user.M));
    1697:   PetscCall(PetscFEGeomDestroy(&user.fegeom));
    1698:   PetscCall(TSDestroy(&ts));
    1699:   PetscCall(DMDestroy(&sw));
    1700:   PetscCall(DMDestroy(&dm));
    1701:   PetscCall(DestroyContext(&user));
    1702:   PetscCall(PetscFinalize());
    1703:   return 0;
    1704: }

    1706: /*TEST

    1708:    build:
    1709:     requires: !complex double

    1711:    testset:
    1712:      nsize: 2
    1713:      args: -cosine_coefficients 0.01 -charges -1. -total_weight 1. -xdm_plex_hash_location -vpetscspace_degree 2 -petscspace_degree 1 -em_snes_atol 1.e-12 \
    1714:      -em_snes_error_if_not_converged -em_ksp_error_if_not_converged -em_ksp_type cg -em_pc_type gamg -em_mg_coarse_ksp_type preonly -em_mg_coarse_pc_type svd -em_proj_ksp_type cg \
    1715:      -em_proj_pc_type gamg -em_proj_mg_coarse_ksp_type preonly -em_proj_mg_coarse_pc_type svd -ts_time_step 0.03 -xdm_plex_simplex 0 \
    1716:      -ts_max_time 100 -output_step 1 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -check_weights -ts_max_steps 60

    1718:      test:
    1719:        suffix: landau_damping_1d
    1720:        args: -xdm_plex_dim 1 -xdm_plex_box_faces 60 -xdm_plex_box_lower 0. -xdm_plex_box_upper 12.5664 -xdm_plex_box_bd periodic -vdm_plex_dim 1 -vdm_plex_box_faces 60 \
    1721:        -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none

    1723:      test:
    1724:        suffix: landau_damping_2d
    1725:        args: -xdm_plex_dim 2 -xdm_plex_box_bd periodic,periodic -vdm_plex_dim 2 -xdm_plex_box_lower 0.,-.5 -vdm_plex_box_lower -6,-6 -vdm_plex_box_upper 6,6 -xdm_plex_box_faces 6,3 \
    1726:        -xdm_plex_box_upper 12.5664,.5 -vdm_plex_box_faces 15,15 -vdm_plex_box_bd none,none -vdm_plex_hash_location -vdm_plex_simplex 0

    1728:      test:
    1729:        suffix: landau_damping_3d
    1730:        requires: !defined(PETSCTEST_VALGRIND)
    1731:        args: -xdm_plex_dim 3 -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_plex_dim 3 -vdm_plex_box_faces 4,4,4 \
    1732:        -vdm_plex_box_lower -6,-6,-6 -vdm_plex_box_upper 6,6,6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none,none,none

    1734:      test:
    1735:        requires: !defined(PETSC_USE_DMLANDAU_2D)
    1736:        suffix: sphere_3d
    1737:        nsize: 1
    1738:        args: -use_landau_velocity_space -xdm_plex_dim 3 -vdm_landau_thermal_temps 1 -vdm_landau_device_type cpu -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 \
    1739:        -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_landau_verbose 2 -vdm_landau_sphere -vdm_landau_map_sphere \
    1740:        -vdm_landau_domain_radius 6,6,6 -vdm_landau_sphere_inner_radius_90degree_scale .35 -vdm_refine 1

    1742: TEST*/