Date: 2026-05-01T06:12:11+00:00

Actual source code: ex4.c

      1: static char help[] = "Two-level system for Landau Damping using Vlasov-Poisson equations\n";

      3: /*
      4:   Moment Equations:

      6:     We will discretize the moment equations using finite elements, and we will project the moments into the finite element space We will use the PFAK method, which guarantees that our FE approximation is weakly equivalent to the true moment. The first moment, number density, is given by

      8:       \int dx \phi_i n_f = \int dx \phi_i n_p
      9:       \int dx \phi_i n_f = \int dx \phi_i \int dv f
     10:       \int dx \phi_i n_f = \int dx \phi_i \int dv \sum_p w_p \delta(x - x_p) \delta(v - v_p)
     11:       \int dx \phi_i n_f = \int dx \phi_i \sum_p w_p \delta(x - x_p)
     12:                    M n_F = M_p w_p

     14:     where

     16:       (M_p){ip} = \phi_i(x_p)

     18:     which is just a scaled version of the charge density. The second moment, momentum density, is given by

     20:       \int dx \phi_i p_f = m \int dx \phi_i \int dv v f
     21:       \int dx \phi_i p_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) v_p
     22:                    M p_F = M_p v_p w_p

     24:     And finally the third moment, pressure, is given by

     26:       \int dx \phi_i pr_f = m \int dx \phi_i \int dv (v - u)^2 f
     27:       \int dx \phi_i pr_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) (v_p - u)^2
     28:                    M pr_F = M_p (v_p - u)^2 w_p
     29:                           = M_p (v_p - p_F(x_p) / m n_F(x_p))^2 w_p
     30:                           = M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p

     32:     Here we need all FEM basis functions \phi_i that see that particle p.

     34:   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"
     35:   According to Lukas, good damping results come at ~16k particles per cell

     37:   Swarm CellDMs
     38:   =============
     39:   Name: "space"
     40:   Fields: DMSwarmPICField_coor, "velocity"
     41:   Coordinates: DMSwarmPICField_coor

     43:   Name: "velocity"
     44:   Fields: "w_q"
     45:   Coordinates: "velocity"

     47:   Name: "moments"
     48:   Fields: "w_q"
     49:   Coordinates: DMSwarmPICField_coor

     51:   Name: "moment fields"
     52:   Fields: "velocity"
     53:   Coordinates: DMSwarmPICField_coor

     55:   To visualize the maximum electric field use

     57:     -efield_monitor

     59:   To monitor velocity moments of the distribution use

     61:     -ptof_pc_type lu -moments_monitor

     63:   To monitor the particle positions in phase space use

     65:     -positions_monitor

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

     69:     -poisson_monitor

     71:   To monitor the remapping field use

     73:     -remap_uf_view draw

     75:   To visualize the swarm distribution use

     77:     -ts_monitor_hg_swarm

     79:   To visualize the particles, we can use

     81:     -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500
     82: */
     83: #include <petsctao.h>
     84: #include <petscts.h>
     85: #include <petscdmplex.h>
     86: #include <petscdmswarm.h>
     87: #include <petscfe.h>
     88: #include <petscds.h>
     89: #include <petscbag.h>
     90: #include <petscdraw.h>
     91: #include <petsc/private/petscfeimpl.h>
     92: #include <petsc/private/dmswarmimpl.h>
     93: #include "petscdm.h"
     94: #include "petscdmlabel.h"

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

     99: const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
    100: typedef enum {
    101:   EM_PRIMAL,
    102:   EM_MIXED,
    103:   EM_COULOMB,
    104:   EM_NONE
    105: } EMType;

    107: typedef enum {
    108:   V0,
    109:   X0,
    110:   T0,
    111:   M0,
    112:   Q0,
    113:   PHI0,
    114:   POISSON,
    115:   VLASOV,
    116:   SIGMA,
    117:   NUM_CONSTANTS
    118: } ConstantType;

    120: typedef enum {
    121:   E_MONITOR_NONE,
    122:   E_MONITOR_FULL,
    123:   E_MONITOR_QUIET
    124: } EMonitorType;
    125: const char *const EMonitorTypes[] = {"NONE", "FULL", "QUIET", "EMonitorType", "E_MONITOR_", NULL};

    127: typedef struct {
    128:   PetscScalar v0; /* Velocity scale, often the thermal velocity */
    129:   PetscScalar t0; /* Time scale */
    130:   PetscScalar x0; /* Space scale */
    131:   PetscScalar m0; /* Mass scale */
    132:   PetscScalar q0; /* Charge scale */
    133:   PetscScalar kb;
    134:   PetscScalar epsi0;
    135:   PetscScalar phi0;          /* Potential scale */
    136:   PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
    137:   PetscScalar vlasovNumber;  /* Non-Dimensional Vlasov Number */
    138:   PetscReal   sigma;         /* Nondimensional charge per length in x */
    139: } Parameter;

    141: typedef struct {
    142:   PetscInt         s;    // Starting sample (we ignore some in the beginning)
    143:   PetscInt         e;    // Ending sample
    144:   PetscInt         per;  // Period of fitting
    145:   const PetscReal *t;    // Time for each sample
    146:   const PetscReal *Emax; // Emax for each sample
    147: } EmaxCtx;

    149: typedef struct {
    150:   PetscBag     bag;                  // Problem parameters
    151:   PetscBool    error;                // Flag for printing the error
    152:   PetscInt     remapFreq;            // Number of timesteps between remapping
    153:   EMonitorType efield_monitor;       // Flag to show electric field monitor
    154:   PetscBool    moment_monitor;       // Flag to show distribution moment monitor
    155:   PetscBool    moment_field_monitor; // Flag to show moment field monitor
    156:   PetscBool    positions_monitor;    // Flag to show particle positins at each time step
    157:   PetscBool    poisson_monitor;      // Flag to display charge, E field, and potential at each solve
    158:   PetscBool    initial_monitor;      // Flag to monitor the initial conditions
    159:   PetscInt     velocity_monitor;     // Cell to monitor the velocity distribution for
    160:   PetscBool    perturbed_weights;    // Uniformly sample x,v space with gaussian weights
    161:   PetscInt     ostep;                // Print the energy at each ostep time steps
    162:   PetscInt     numParticles;
    163:   PetscReal    timeScale;              /* Nondimensionalizing time scale */
    164:   PetscReal    charges[2];             /* The charges of each species */
    165:   PetscReal    masses[2];              /* The masses of each species */
    166:   PetscReal    thermal_energy[2];      /* Thermal Energy (used to get other constants)*/
    167:   PetscReal    cosine_coefficients[2]; /*(alpha, k)*/
    168:   PetscReal    totalWeight;
    169:   PetscReal    stepSize;
    170:   PetscInt     steps;
    171:   PetscReal    initVel;
    172:   EMType       em;           // Type of electrostatic model
    173:   SNES         snes;         // EM solver
    174:   DM           dmMom;        // The DM for moment fields
    175:   DM           dmN;          // The DM for number density fields
    176:   IS           isN;          // The IS mapping dmN into dmMom
    177:   Mat          MN;           // The finite element mass matrix for number density
    178:   DM           dmP;          // The DM for momentum density fields
    179:   IS           isP;          // The IS mapping dmP into dmMom
    180:   Mat          MP;           // The finite element mass matrix for momentum density
    181:   DM           dmE;          // The DM for energy density (pressure) fields
    182:   IS           isE;          // The IS mapping dmE into dmMom
    183:   Mat          ME;           // The finite element mass matrix for energy density (pressure)
    184:   DM           dmPot;        // The DM for potential
    185:   Mat          fftPot;       // Fourier Transform operator for the potential
    186:   Vec          fftX, fftY;   //   FFT vectors with phases added (complex parts)
    187:   IS           fftReal;      //   The indices for real parts
    188:   IS           isPot;        // The IS for potential, or NULL in primal
    189:   Mat          M;            // The finite element mass matrix for potential
    190:   PetscFEGeom *fegeom;       // Geometric information for the DM cells
    191:   PetscDrawHG  drawhgic_x;   // Histogram of the particle weight in each X cell
    192:   PetscDrawHG  drawhgic_v;   // Histogram of the particle weight in each X cell
    193:   PetscDrawHG  drawhgcell_v; // Histogram of the particle weight in a given cell
    194:   PetscBool    validE;       // Flag to indicate E-field in swarm is valid
    195:   PetscReal    drawlgEmin;   // The minimum lg(E) to plot
    196:   PetscDrawLG  drawlgE;      // Logarithm of maximum electric field
    197:   PetscDrawSP  drawspE;      // Electric field at particle positions
    198:   PetscDrawSP  drawspX;      // Particle positions
    199:   PetscViewer  viewerRho;    // Charge density viewer
    200:   PetscViewer  viewerRhoHat; // Charge density Fourier Transform viewer
    201:   PetscViewer  viewerPhi;    // Potential viewer
    202:   PetscViewer  viewerN;      // Number density viewer
    203:   PetscViewer  viewerP;      // Momentum density viewer
    204:   PetscViewer  viewerE;      // Energy density (pressure) viewer
    205:   PetscViewer  viewerNRes;   // Number density residual viewer
    206:   PetscViewer  viewerPRes;   // Momentum density residual viewer
    207:   PetscViewer  viewerERes;   // Energy density (pressure) residual viewer
    208:   PetscDrawLG  drawlgMomRes; // Residuals for the moment equations
    209:   DM           swarm;        // The particle swarm
    210:   PetscRandom  random;       // Used for particle perturbations
    211:   PetscBool    twostream;    // Flag for activating 2-stream setup
    212:   PetscBool    checkweights; // Check weight normalization
    213:   PetscInt     checkVRes;    // Flag to check/output velocity residuals for nightly tests
    214:   PetscBool    checkLandau;  // Check the Landau damping result
    215:   EmaxCtx      emaxCtx;      // Information for fit to decay profile
    216:   PetscReal    gamma;        // The damping rate for Landau damping
    217:   PetscReal    omega;        // The perturbed oscillation frequency for Landau damping

    219:   PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
    220: } AppCtx;

    222: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
    223: {
    224:   PetscFunctionBeginUser;
    225:   PetscInt d                      = 2;
    226:   PetscInt maxSpecies             = 2;
    227:   options->error                  = PETSC_FALSE;
    228:   options->remapFreq              = 1;
    229:   options->efield_monitor         = E_MONITOR_NONE;
    230:   options->moment_monitor         = PETSC_FALSE;
    231:   options->moment_field_monitor   = PETSC_FALSE;
    232:   options->initial_monitor        = PETSC_FALSE;
    233:   options->perturbed_weights      = PETSC_FALSE;
    234:   options->poisson_monitor        = PETSC_FALSE;
    235:   options->positions_monitor      = PETSC_FALSE;
    236:   options->velocity_monitor       = -1;
    237:   options->ostep                  = 100;
    238:   options->timeScale              = 2.0e-14;
    239:   options->charges[0]             = -1.0;
    240:   options->charges[1]             = 1.0;
    241:   options->masses[0]              = 1.0;
    242:   options->masses[1]              = 1000.0;
    243:   options->thermal_energy[0]      = 1.0;
    244:   options->thermal_energy[1]      = 1.0;
    245:   options->cosine_coefficients[0] = 0.01;
    246:   options->cosine_coefficients[1] = 0.5;
    247:   options->initVel                = 1;
    248:   options->totalWeight            = 1.0;
    249:   options->drawhgic_x             = NULL;
    250:   options->drawhgic_v             = NULL;
    251:   options->drawhgcell_v           = NULL;
    252:   options->validE                 = PETSC_FALSE;
    253:   options->drawlgEmin             = -6;
    254:   options->drawlgE                = NULL;
    255:   options->drawspE                = NULL;
    256:   options->drawspX                = NULL;
    257:   options->viewerRho              = NULL;
    258:   options->viewerRhoHat           = NULL;
    259:   options->viewerPhi              = NULL;
    260:   options->viewerN                = NULL;
    261:   options->viewerP                = NULL;
    262:   options->viewerE                = NULL;
    263:   options->viewerNRes             = NULL;
    264:   options->viewerPRes             = NULL;
    265:   options->viewerERes             = NULL;
    266:   options->drawlgMomRes           = NULL;
    267:   options->em                     = EM_COULOMB;
    268:   options->snes                   = NULL;
    269:   options->dmMom                  = NULL;
    270:   options->dmN                    = NULL;
    271:   options->MN                     = NULL;
    272:   options->dmP                    = NULL;
    273:   options->MP                     = NULL;
    274:   options->dmE                    = NULL;
    275:   options->ME                     = NULL;
    276:   options->dmPot                  = NULL;
    277:   options->fftPot                 = NULL;
    278:   options->fftX                   = NULL;
    279:   options->fftY                   = NULL;
    280:   options->fftReal                = NULL;
    281:   options->isPot                  = NULL;
    282:   options->M                      = NULL;
    283:   options->numParticles           = 32768;
    284:   options->twostream              = PETSC_FALSE;
    285:   options->checkweights           = PETSC_FALSE;
    286:   options->checkVRes              = 0;
    287:   options->checkLandau            = PETSC_FALSE;
    288:   options->emaxCtx.s              = 50;
    289:   options->emaxCtx.per            = 100;

    291:   PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM");
    292:   PetscCall(PetscOptionsBool("-error", "Flag to print the error", __FILE__, options->error, &options->error, NULL));
    293:   PetscCall(PetscOptionsInt("-remap_freq", "Number", __FILE__, options->remapFreq, &options->remapFreq, NULL));
    294:   PetscCall(PetscOptionsEnum("-efield_monitor", "Flag to record and plot log(max E) over time", __FILE__, EMonitorTypes, (PetscEnum)options->efield_monitor, (PetscEnum *)&options->efield_monitor, NULL));
    295:   PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", __FILE__, options->drawlgEmin, &options->drawlgEmin, NULL));
    296:   PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", __FILE__, options->moment_monitor, &options->moment_monitor, NULL));
    297:   PetscCall(PetscOptionsBool("-moment_field_monitor", "Flag to show moment fields", __FILE__, options->moment_field_monitor, &options->moment_field_monitor, NULL));
    298:   PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", __FILE__, options->initial_monitor, &options->initial_monitor, NULL));
    299:   PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", __FILE__, options->positions_monitor, &options->positions_monitor, NULL));
    300:   PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", __FILE__, options->poisson_monitor, &options->poisson_monitor, NULL));
    301:   PetscCall(PetscOptionsInt("-velocity_monitor", "Cell to show velocity histograms", __FILE__, options->velocity_monitor, &options->velocity_monitor, NULL));
    302:   PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", __FILE__, options->twostream, &options->twostream, NULL));
    303:   PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", __FILE__, options->perturbed_weights, &options->perturbed_weights, NULL));
    304:   PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", __FILE__, options->checkweights, &options->checkweights, NULL));
    305:   PetscCall(PetscOptionsBool("-check_landau", "Check the decay from Landau damping", __FILE__, options->checkLandau, &options->checkLandau, NULL));
    306:   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", __FILE__, options->ostep, &options->ostep, NULL));
    307:   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", __FILE__, options->timeScale, &options->timeScale, NULL));
    308:   PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", __FILE__, options->checkVRes, &options->checkVRes, NULL));
    309:   PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", __FILE__, options->initVel, &options->initVel, NULL));
    310:   PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", __FILE__, options->totalWeight, &options->totalWeight, NULL));
    311:   PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", __FILE__, options->cosine_coefficients, &d, NULL));
    312:   PetscCall(PetscOptionsRealArray("-charges", "Species charges", __FILE__, options->charges, &maxSpecies, NULL));
    313:   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", __FILE__, EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
    314:   PetscCall(PetscOptionsInt("-emax_start_step", "First time step to use for Emax fits", __FILE__, options->emaxCtx.s, &options->emaxCtx.s, NULL));
    315:   PetscCall(PetscOptionsInt("-emax_solve_step", "Number of time steps between Emax fits", __FILE__, options->emaxCtx.per, &options->emaxCtx.per, NULL));
    316:   PetscOptionsEnd();

    318:   PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
    319:   PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
    320:   PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
    321:   PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
    322:   PetscFunctionReturn(PETSC_SUCCESS);
    323: }

    325: static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *ctx)
    326: {
    327:   MPI_Comm comm;

    329:   PetscFunctionBeginUser;
    330:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
    331:   if (ctx->efield_monitor) {
    332:     PetscDraw     draw;
    333:     PetscDrawAxis axis;

    335:     if (ctx->efield_monitor == E_MONITOR_FULL) {
    336:       PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 0, 400, 300, &draw));
    337:       PetscCall(PetscDrawSetSave(draw, "ex2_Efield"));
    338:       PetscCall(PetscDrawSetFromOptions(draw));
    339:     } else {
    340:       PetscCall(PetscDrawOpenNull(comm, &draw));
    341:     }
    342:     PetscCall(PetscDrawLGCreate(draw, 1, &ctx->drawlgE));
    343:     PetscCall(PetscDrawDestroy(&draw));
    344:     PetscCall(PetscDrawLGGetAxis(ctx->drawlgE, &axis));
    345:     PetscCall(PetscDrawAxisSetLabels(axis, "Max Electric Field", "time", "E_max"));
    346:     PetscCall(PetscDrawLGSetLimits(ctx->drawlgE, 0., ctx->steps * ctx->stepSize, ctx->drawlgEmin, 0.));
    347:   }

    349:   if (ctx->initial_monitor) {
    350:     PetscDraw     drawic_x, drawic_v;
    351:     PetscDrawAxis axis1, axis2;
    352:     PetscReal     dmboxlower[2], dmboxupper[2];
    353:     PetscInt      dim, cStart, cEnd;

    355:     PetscCall(DMGetDimension(sw, &dim));
    356:     PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper));
    357:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

    359:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &drawic_x));
    360:     PetscCall(PetscDrawSetSave(drawic_x, "ex2_ic_x"));
    361:     PetscCall(PetscDrawSetFromOptions(drawic_x));
    362:     PetscCall(PetscDrawHGCreate(drawic_x, (int)dim, &ctx->drawhgic_x));
    363:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgic_x, PETSC_TRUE));
    364:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgic_x, &axis1));
    365:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgic_x, (int)(cEnd - cStart)));
    366:     PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "weight"));
    367:     PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 0));
    368:     PetscCall(PetscDrawDestroy(&drawic_x));

    370:     PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &drawic_v));
    371:     PetscCall(PetscDrawSetSave(drawic_v, "ex9_ic_v"));
    372:     PetscCall(PetscDrawSetFromOptions(drawic_v));
    373:     PetscCall(PetscDrawHGCreate(drawic_v, (int)dim, &ctx->drawhgic_v));
    374:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgic_v, PETSC_TRUE));
    375:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgic_v, &axis2));
    376:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgic_v, 21));
    377:     PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "weight"));
    378:     PetscCall(PetscDrawAxisSetLimits(axis2, -6, 6, 0, 0));
    379:     PetscCall(PetscDrawDestroy(&drawic_v));
    380:   }

    382:   if (ctx->velocity_monitor >= 0) {
    383:     DM            vdm;
    384:     DMSwarmCellDM celldm;
    385:     PetscDraw     drawcell_v;
    386:     PetscDrawAxis axis;
    387:     PetscReal     dmboxlower[2], dmboxupper[2];
    388:     PetscInt      dim;
    389:     char          title[PETSC_MAX_PATH_LEN];

    391:     PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
    392:     PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    393:     PetscCall(DMGetDimension(vdm, &dim));
    394:     PetscCall(DMGetBoundingBox(vdm, dmboxlower, dmboxupper));

    396:     PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Cell %" PetscInt_FMT ": Velocity Distribution", ctx->velocity_monitor));
    397:     PetscCall(PetscDrawCreate(comm, NULL, title, 400, 300, 400, 300, &drawcell_v));
    398:     PetscCall(PetscDrawSetSave(drawcell_v, "ex2_cell_v"));
    399:     PetscCall(PetscDrawSetFromOptions(drawcell_v));
    400:     PetscCall(PetscDrawHGCreate(drawcell_v, (int)dim, &ctx->drawhgcell_v));
    401:     PetscCall(PetscDrawHGCalcStats(ctx->drawhgcell_v, PETSC_TRUE));
    402:     PetscCall(PetscDrawHGGetAxis(ctx->drawhgcell_v, &axis));
    403:     PetscCall(PetscDrawHGSetNumberBins(ctx->drawhgcell_v, 21));
    404:     PetscCall(PetscDrawAxisSetLabels(axis, "V_x Distribution", "V", "weight"));
    405:     PetscCall(PetscDrawAxisSetLimits(axis, dmboxlower[0], dmboxupper[0], 0, 0));
    406:     PetscCall(PetscDrawDestroy(&drawcell_v));
    407:   }

    409:   if (ctx->positions_monitor) {
    410:     PetscDraw     draw;
    411:     PetscDrawAxis axis;

    413:     PetscCall(PetscDrawCreate(comm, NULL, "Particle Position", 0, 0, 400, 300, &draw));
    414:     PetscCall(PetscDrawSetSave(draw, "ex9_pos"));
    415:     PetscCall(PetscDrawSetFromOptions(draw));
    416:     PetscCall(PetscDrawSPCreate(draw, 10, &ctx->drawspX));
    417:     PetscCall(PetscDrawDestroy(&draw));
    418:     PetscCall(PetscDrawSPSetDimension(ctx->drawspX, 1));
    419:     PetscCall(PetscDrawSPGetAxis(ctx->drawspX, &axis));
    420:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v"));
    421:     PetscCall(PetscDrawSPReset(ctx->drawspX));
    422:   }
    423:   if (ctx->poisson_monitor) {
    424:     Vec           rho, rhohat, phi;
    425:     PetscDraw     draw;
    426:     PetscDrawAxis axis;

    428:     PetscCall(PetscDrawCreate(comm, NULL, "Electric_Field", 0, 0, 400, 300, &draw));
    429:     PetscCall(PetscDrawSetFromOptions(draw));
    430:     PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial"));
    431:     PetscCall(PetscDrawSPCreate(draw, 10, &ctx->drawspE));
    432:     PetscCall(PetscDrawDestroy(&draw));
    433:     PetscCall(PetscDrawSPSetDimension(ctx->drawspE, 1));
    434:     PetscCall(PetscDrawSPGetAxis(ctx->drawspE, &axis));
    435:     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E"));
    436:     PetscCall(PetscDrawSPReset(ctx->drawspE));

    438:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Charge Density", 0, 0, 400, 300, &ctx->viewerRho));
    439:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerRho, "rho_"));
    440:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerRho, 0, &draw));
    441:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial"));
    442:     PetscCall(PetscViewerSetFromOptions(ctx->viewerRho));
    443:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
    444:     PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density"));
    445:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));

    447:     PetscInt dim, N;

    449:     PetscCall(DMGetDimension(ctx->dmPot, &dim));
    450:     if (dim == 1) {
    451:       PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
    452:       PetscCall(VecGetSize(rhohat, &N));
    453:       PetscCall(MatCreateFFT(comm, dim, &N, MATFFTW, &ctx->fftPot));
    454:       PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
    455:       PetscCall(MatCreateVecs(ctx->fftPot, &ctx->fftX, &ctx->fftY));
    456:       PetscCall(ISCreateStride(PETSC_COMM_SELF, N, 0, 1, &ctx->fftReal));
    457:     }

    459:     PetscCall(PetscViewerDrawOpen(comm, NULL, "rhohat: Charge Density FT", 0, 0, 400, 300, &ctx->viewerRhoHat));
    460:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerRhoHat, "rhohat_"));
    461:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerRhoHat, 0, &draw));
    462:     PetscCall(PetscDrawSetSave(draw, "ex9_rho_ft"));
    463:     PetscCall(PetscViewerSetFromOptions(ctx->viewerRhoHat));
    464:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
    465:     PetscCall(PetscObjectSetName((PetscObject)rhohat, "charge_density_ft"));
    466:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));

    468:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Potential", 400, 0, 400, 300, &ctx->viewerPhi));
    469:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerPhi, "phi_"));
    470:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerPhi, 0, &draw));
    471:     PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial"));
    472:     PetscCall(PetscViewerSetFromOptions(ctx->viewerPhi));
    473:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    474:     PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
    475:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
    476:   }
    477:   if (ctx->moment_field_monitor) {
    478:     Vec       n, p, e;
    479:     Vec       nres, pres, eres;
    480:     PetscDraw draw;

    482:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density", 400, 0, 400, 300, &ctx->viewerN));
    483:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerN, "n_"));
    484:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerN, 0, &draw));
    485:     PetscCall(PetscDrawSetSave(draw, "ex4_n_spatial"));
    486:     PetscCall(PetscViewerSetFromOptions(ctx->viewerN));
    487:     PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
    488:     PetscCall(PetscObjectSetName((PetscObject)n, "Number Density"));
    489:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));

    491:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density", 800, 0, 400, 300, &ctx->viewerP));
    492:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerP, "p_"));
    493:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerP, 0, &draw));
    494:     PetscCall(PetscDrawSetSave(draw, "ex4_p_spatial"));
    495:     PetscCall(PetscViewerSetFromOptions(ctx->viewerP));
    496:     PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
    497:     PetscCall(PetscObjectSetName((PetscObject)p, "Momentum Density"));
    498:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));

    500:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Emergy Density (Pressure)", 1200, 0, 400, 300, &ctx->viewerE));
    501:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerE, "e_"));
    502:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerE, 0, &draw));
    503:     PetscCall(PetscDrawSetSave(draw, "ex4_e_spatial"));
    504:     PetscCall(PetscViewerSetFromOptions(ctx->viewerE));
    505:     PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));
    506:     PetscCall(PetscObjectSetName((PetscObject)e, "Energy Density (Pressure)"));
    507:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));

    509:     PetscDrawAxis axis;

    511:     PetscCall(PetscDrawCreate(comm, NULL, "Moment Residual", 0, 320, 400, 300, &draw));
    512:     PetscCall(PetscDrawSetSave(draw, "ex4_moment_res"));
    513:     PetscCall(PetscDrawSetFromOptions(draw));
    514:     PetscCall(PetscDrawLGCreate(draw, 3, &ctx->drawlgMomRes));
    515:     PetscCall(PetscDrawDestroy(&draw));
    516:     PetscCall(PetscDrawLGGetAxis(ctx->drawlgMomRes, &axis));
    517:     PetscCall(PetscDrawAxisSetLabels(axis, "Moment Residial", "time", "Residual Norm"));
    518:     PetscCall(PetscDrawLGSetLimits(ctx->drawlgMomRes, 0., ctx->steps * ctx->stepSize, -8, 0));

    520:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density Residual", 400, 300, 400, 300, &ctx->viewerNRes));
    521:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerNRes, "nres_"));
    522:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerNRes, 0, &draw));
    523:     PetscCall(PetscDrawSetSave(draw, "ex4_nres_spatial"));
    524:     PetscCall(PetscViewerSetFromOptions(ctx->viewerNRes));
    525:     PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
    526:     PetscCall(PetscObjectSetName((PetscObject)nres, "Number Density Residual"));
    527:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));

    529:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density Residual", 800, 300, 400, 300, &ctx->viewerPRes));
    530:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerPRes, "pres_"));
    531:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerPRes, 0, &draw));
    532:     PetscCall(PetscDrawSetSave(draw, "ex4_pres_spatial"));
    533:     PetscCall(PetscViewerSetFromOptions(ctx->viewerPRes));
    534:     PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
    535:     PetscCall(PetscObjectSetName((PetscObject)pres, "Momentum Density Residual"));
    536:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));

    538:     PetscCall(PetscViewerDrawOpen(comm, NULL, "Energy Density Residual", 1200, 300, 400, 300, &ctx->viewerERes));
    539:     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)ctx->viewerERes, "eres_"));
    540:     PetscCall(PetscViewerDrawGetDraw(ctx->viewerERes, 0, &draw));
    541:     PetscCall(PetscDrawSetSave(draw, "ex4_eres_spatial"));
    542:     PetscCall(PetscViewerSetFromOptions(ctx->viewerERes));
    543:     PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
    544:     PetscCall(PetscObjectSetName((PetscObject)eres, "Energy Density Residual"));
    545:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
    546:   }
    547:   PetscFunctionReturn(PETSC_SUCCESS);
    548: }

    550: static PetscErrorCode DestroyContext(AppCtx *ctx)
    551: {
    552:   PetscFunctionBeginUser;
    553:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgic_x));
    554:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgic_v));
    555:   PetscCall(PetscDrawHGDestroy(&ctx->drawhgcell_v));

    557:   PetscCall(PetscDrawLGDestroy(&ctx->drawlgE));
    558:   PetscCall(PetscDrawSPDestroy(&ctx->drawspE));
    559:   PetscCall(PetscDrawSPDestroy(&ctx->drawspX));
    560:   PetscCall(PetscViewerDestroy(&ctx->viewerRho));
    561:   PetscCall(PetscViewerDestroy(&ctx->viewerRhoHat));
    562:   PetscCall(MatDestroy(&ctx->fftPot));
    563:   PetscCall(VecDestroy(&ctx->fftX));
    564:   PetscCall(VecDestroy(&ctx->fftY));
    565:   PetscCall(ISDestroy(&ctx->fftReal));
    566:   PetscCall(PetscViewerDestroy(&ctx->viewerPhi));
    567:   PetscCall(PetscViewerDestroy(&ctx->viewerN));
    568:   PetscCall(PetscViewerDestroy(&ctx->viewerP));
    569:   PetscCall(PetscViewerDestroy(&ctx->viewerE));
    570:   PetscCall(PetscViewerDestroy(&ctx->viewerNRes));
    571:   PetscCall(PetscViewerDestroy(&ctx->viewerPRes));
    572:   PetscCall(PetscViewerDestroy(&ctx->viewerERes));
    573:   PetscCall(PetscDrawLGDestroy(&ctx->drawlgMomRes));

    575:   PetscCall(PetscBagDestroy(&ctx->bag));
    576:   PetscFunctionReturn(PETSC_SUCCESS);
    577: }

    579: static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *ctx)
    580: {
    581:   const PetscScalar *w;
    582:   PetscInt           Np;

    584:   PetscFunctionBeginUser;
    585:   if (!ctx->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
    586:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
    587:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    588:   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]);
    589:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
    590:   PetscFunctionReturn(PETSC_SUCCESS);
    591: }

    593: 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[])
    594: {
    595:   for (PetscInt d = 0; d < dim; ++d) f0[0] += 0.5 * PetscSqr(u_x[d]);
    596: }

    598: static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
    599: {
    600:   PetscDS        ds;
    601:   const PetscInt field = 0;
    602:   PetscInt       Nf;
    603:   void          *ctx;

    605:   PetscFunctionBegin;
    606:   PetscCall(DMGetApplicationContext(dm, &ctx));
    607:   PetscCall(DMGetDS(dm, &ds));
    608:   PetscCall(PetscDSGetNumFields(ds, &Nf));
    609:   PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
    610:   PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
    611:   PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
    612:   PetscFunctionReturn(PETSC_SUCCESS);
    613: }

    615: static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *ctx)
    616: {
    617:   DMSwarmCellDM celldm;
    618:   DM            vdm;
    619:   Vec           u[1];
    620:   const char   *fields[1] = {"w_q"};

    622:   PetscFunctionBegin;
    623:   PetscCall(DMSwarmSetCellDMActive(sw, "velocity"));
    624:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
    625:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    626:   PetscCall(DMGetGlobalVector(vdm, &u[0]));
    627:   PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD));
    628:   PetscCall(DMPlexComputeMoments(vdm, u[0], moments));
    629:   PetscCall(DMRestoreGlobalVector(vdm, &u[0]));
    630:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
    631:   PetscFunctionReturn(PETSC_SUCCESS);
    632: }

    634: 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[])
    635: {
    636:   f0[0] = 0.;
    637:   for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d * dim + d]);
    638: }

    640: // Our model is E_max(t) = C e^{-gamma t} |cos(omega t - phi)|
    641: static PetscErrorCode ComputeEmaxResidual(Tao tao, Vec x, Vec res, void *Ctx)
    642: {
    643:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
    644:   const PetscScalar *a;
    645:   PetscScalar       *F;
    646:   PetscReal          C, gamma, omega, phi;

    648:   PetscFunctionBegin;
    649:   PetscCall(VecGetArrayRead(x, &a));
    650:   PetscCall(VecGetArray(res, &F));
    651:   C     = PetscRealPart(a[0]);
    652:   gamma = PetscRealPart(a[1]);
    653:   omega = PetscRealPart(a[2]);
    654:   phi   = PetscRealPart(a[3]);
    655:   PetscCall(VecRestoreArrayRead(x, &a));
    656:   for (PetscInt i = ctx->s; i < ctx->e; ++i) F[i - ctx->s] = PetscPowReal(10., ctx->Emax[i]) - C * PetscExpReal(-gamma * ctx->t[i]) * PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi));
    657:   PetscCall(VecRestoreArray(res, &F));
    658:   PetscFunctionReturn(PETSC_SUCCESS);
    659: }

    661: // The Jacobian of the residual J = dr(x)/dx
    662: static PetscErrorCode ComputeEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *Ctx)
    663: {
    664:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
    665:   const PetscScalar *a;
    666:   PetscScalar       *jac;
    667:   PetscReal          C, gamma, omega, phi;
    668:   const PetscInt     n = ctx->e - ctx->s;

    670:   PetscFunctionBegin;
    671:   PetscCall(VecGetArrayRead(x, &a));
    672:   C     = PetscRealPart(a[0]);
    673:   gamma = PetscRealPart(a[1]);
    674:   omega = PetscRealPart(a[2]);
    675:   phi   = PetscRealPart(a[3]);
    676:   PetscCall(VecRestoreArrayRead(x, &a));
    677:   PetscCall(MatDenseGetArray(J, &jac));
    678:   for (PetscInt i = 0; i < n; ++i) {
    679:     const PetscInt k = i + ctx->s;

    681:     jac[i * 4 + 0] = -PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
    682:     jac[i * 4 + 1] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi));
    683:     jac[i * 4 + 2] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi);
    684:     jac[i * 4 + 3] = -C * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi);
    685:   }
    686:   PetscCall(MatDenseRestoreArray(J, &jac));
    687:   PetscFunctionReturn(PETSC_SUCCESS);
    688: }

    690: // Our model is log_10 E_max(t) = log_10 C - gamma t log_10 e + log_10 |cos(omega t - phi)|
    691: static PetscErrorCode ComputeLogEmaxResidual(Tao tao, Vec x, Vec res, void *Ctx)
    692: {
    693:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
    694:   const PetscScalar *a;
    695:   PetscScalar       *F;
    696:   PetscReal          C, gamma, omega, phi;

    698:   PetscFunctionBegin;
    699:   PetscCall(VecGetArrayRead(x, &a));
    700:   PetscCall(VecGetArray(res, &F));
    701:   C     = PetscRealPart(a[0]);
    702:   gamma = PetscRealPart(a[1]);
    703:   omega = PetscRealPart(a[2]);
    704:   phi   = PetscRealPart(a[3]);
    705:   PetscCall(VecRestoreArrayRead(x, &a));
    706:   for (PetscInt i = ctx->s; i < ctx->e; ++i) {
    707:     if (C < 0) {
    708:       F[i - ctx->s] = 1e10;
    709:       continue;
    710:     }
    711:     F[i - ctx->s] = ctx->Emax[i] - (PetscLog10Real(C) - gamma * ctx->t[i] * PetscLog10Real(PETSC_E) + PetscLog10Real(PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi))));
    712:   }
    713:   PetscCall(VecRestoreArray(res, &F));
    714:   PetscFunctionReturn(PETSC_SUCCESS);
    715: }

    717: // The Jacobian of the residual J = dr(x)/dx
    718: static PetscErrorCode ComputeLogEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *Ctx)
    719: {
    720:   EmaxCtx           *ctx = (EmaxCtx *)Ctx;
    721:   const PetscScalar *a;
    722:   PetscScalar       *jac;
    723:   PetscReal          C, omega, phi;
    724:   const PetscInt     n = ctx->e - ctx->s;

    726:   PetscFunctionBegin;
    727:   PetscCall(VecGetArrayRead(x, &a));
    728:   C     = PetscRealPart(a[0]);
    729:   omega = PetscRealPart(a[2]);
    730:   phi   = PetscRealPart(a[3]);
    731:   PetscCall(VecRestoreArrayRead(x, &a));
    732:   PetscCall(MatDenseGetArray(J, &jac));
    733:   for (PetscInt i = 0; i < n; ++i) {
    734:     const PetscInt k = i + ctx->s;

    736:     jac[0 * n + i] = -1. / (PetscLog10Real(PETSC_E) * C);
    737:     jac[1 * n + i] = ctx->t[k] * PetscLog10Real(PETSC_E);
    738:     jac[2 * n + i] = (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * ctx->t[k] * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)));
    739:     jac[3 * n + i] = -(PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)));
    740:   }
    741:   PetscCall(MatDenseRestoreArray(J, &jac));
    742:   PetscCall(MatViewFromOptions(J, NULL, "-emax_jac_view"));
    743:   PetscFunctionReturn(PETSC_SUCCESS);
    744: }

    746: static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    747: {
    748:   AppCtx     *ctx = (AppCtx *)Ctx;
    749:   DM          sw;
    750:   PetscScalar intESq;
    751:   PetscReal  *E, *x, *weight;
    752:   PetscReal   Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
    753:   PetscReal   pmoments[4]; /* \int f, \int v f, \int v^2 f */
    754:   PetscInt   *species, dim, Np, gNp;
    755:   MPI_Comm    comm;
    756:   PetscMPIInt rank;

    758:   PetscFunctionBeginUser;
    759:   if (step < 0 || !ctx->validE) PetscFunctionReturn(PETSC_SUCCESS);
    760:   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
    761:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
    762:   PetscCall(TSGetDM(ts, &sw));
    763:   PetscCall(DMGetDimension(sw, &dim));
    764:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    765:   PetscCall(DMSwarmGetSize(sw, &gNp));
    766:   PetscCall(DMSwarmSortGetAccess(sw));
    767:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    768:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    769:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    770:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

    772:   for (PetscInt p = 0; p < Np; ++p) {
    773:     for (PetscInt d = 0; d < 1; ++d) {
    774:       PetscReal temp = PetscAbsReal(E[p * dim + d]);
    775:       if (temp > Emax) Emax = temp;
    776:     }
    777:     Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]);
    778:     sum += E[p * dim];
    779:     chargesum += ctx->charges[0] * weight[p];
    780:   }
    781:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
    782:   lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
    783:   lgEmax  = Emax != 0 ? PetscLog10Real(Emax) : ctx->drawlgEmin;

    785:   PetscDS ds;
    786:   Vec     phi;

    788:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    789:   PetscCall(DMGetDS(ctx->dmPot, &ds));
    790:   PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
    791:   PetscCall(DMPlexComputeIntegralFEM(ctx->dmPot, phi, &intESq, ctx));
    792:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

    794:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    795:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    796:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    797:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    798:   PetscCall(PetscDrawLGAddPoint(ctx->drawlgE, &t, &lgEmax));
    799:   if (ctx->efield_monitor == E_MONITOR_FULL) {
    800:     PetscDraw draw;

    802:     PetscCall(PetscDrawLGDraw(ctx->drawlgE));
    803:     PetscCall(PetscDrawLGGetDraw(ctx->drawlgE, &draw));
    804:     PetscCall(PetscDrawSave(draw));

    806:     PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
    807:     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));
    808:     PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
    809:   }

    811:   // Compute decay rate and frequency
    812:   PetscCall(PetscDrawLGGetData(ctx->drawlgE, NULL, &ctx->emaxCtx.e, &ctx->emaxCtx.t, &ctx->emaxCtx.Emax));
    813:   if (!rank && !(ctx->emaxCtx.e % ctx->emaxCtx.per)) {
    814:     Tao          tao;
    815:     Mat          J;
    816:     Vec          x, r;
    817:     PetscScalar *a;
    818:     PetscBool    fitLog = PETSC_TRUE, debug = PETSC_FALSE;

    820:     PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
    821:     PetscCall(TaoSetOptionsPrefix(tao, "emax_"));
    822:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, 4, &x));
    823:     PetscCall(TaoSetSolution(tao, x));
    824:     PetscCall(VecCreateSeq(PETSC_COMM_SELF, ctx->emaxCtx.e - ctx->emaxCtx.s, &r));
    825:     if (fitLog) PetscCall(TaoSetResidualRoutine(tao, r, ComputeLogEmaxResidual, &ctx->emaxCtx));
    826:     else PetscCall(TaoSetResidualRoutine(tao, r, ComputeEmaxResidual, &ctx->emaxCtx));
    827:     PetscCall(VecDestroy(&r));
    828:     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, ctx->emaxCtx.e - ctx->emaxCtx.s, 4, NULL, &J));
    829:     if (fitLog) PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeLogEmaxJacobian, &ctx->emaxCtx));
    830:     else PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeEmaxJacobian, &ctx->emaxCtx));
    831:     PetscCall(MatDestroy(&J));
    832:     PetscCall(TaoSetFromOptions(tao));
    833:     PetscCall(VecGetArray(x, &a));
    834:     a[0] = 0.02;
    835:     a[1] = 0.15;
    836:     a[2] = 1.4;
    837:     a[3] = 0.45;
    838:     PetscCall(VecRestoreArray(x, &a));
    839:     PetscCall(TaoSolve(tao));
    840:     if (debug) {
    841:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "t = ["));
    842:       for (PetscInt i = 0; i < ctx->emaxCtx.e; ++i) {
    843:         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    844:         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ctx->emaxCtx.t[i]));
    845:       }
    846:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
    847:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax = ["));
    848:       for (PetscInt i = 0; i < ctx->emaxCtx.e; ++i) {
    849:         if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    850:         PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ctx->emaxCtx.Emax[i]));
    851:       }
    852:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n"));
    853:     }
    854:     PetscDraw     draw;
    855:     PetscDrawAxis axis;
    856:     char          title[PETSC_MAX_PATH_LEN];

    858:     PetscCall(VecGetArray(x, &a));
    859:     ctx->gamma = a[1];
    860:     ctx->omega = a[2];
    861:     if (ctx->efield_monitor == E_MONITOR_FULL) {
    862:       PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax Fit: gamma %g omega %g C %g phi %g\n", a[1], a[2], a[0], a[3]));
    863:       PetscCall(PetscDrawLGGetDraw(ctx->drawlgE, &draw));
    864:       PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Max Electric Field gamma %.4g omega %.4g", a[1], a[2]));
    865:       PetscCall(PetscDrawSetTitle(draw, title));
    866:       PetscCall(PetscDrawLGGetAxis(ctx->drawlgE, &axis));
    867:       PetscCall(PetscDrawAxisSetLabels(axis, title, "time", "E_max"));
    868:     }
    869:     PetscCall(VecRestoreArray(x, &a));
    870:     PetscCall(VecDestroy(&x));
    871:     PetscCall(TaoDestroy(&tao));
    872:   }
    873:   PetscFunctionReturn(PETSC_SUCCESS);
    874: }

    876: static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    877: {
    878:   AppCtx   *ctx = (AppCtx *)Ctx;
    879:   DM        sw;
    880:   PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */

    882:   PetscFunctionBeginUser;
    883:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
    884:   PetscCall(TSGetDM(ts, &sw));

    886:   PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
    887:   PetscCall(computeVelocityFEMMoments(sw, fmoments, ctx));

    889:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3], (double)fmoments[0], (double)fmoments[1], (double)fmoments[2]));
    890:   PetscFunctionReturn(PETSC_SUCCESS);
    891: }

    893: static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, PetscCtx ctx)
    894: {
    895:   u[0] = 0.0;
    896:   return PETSC_SUCCESS;
    897: }

    899: /*
    900:     M_p w_p
    901:       - Make M_p with "moments"
    902:       - Get w_p from Swarm
    903:     M_p v_p w_p
    904:       - Get v_p from Swarm
    905:       - pointwise multiply v_p and w_p
    906:     M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p
    907:       - ProjectField(sw, {n, p} U, {v_p} A, tmp_p)
    908:       - pointwise multiply tmp_p and w_p

    910:   Projection works fpr swarms
    911:     Fields are FE from the CellDM, and aux fields are the swarm fields
    912: */
    913: static PetscErrorCode ComputeMomentFields(TS ts)
    914: {
    915:   AppCtx   *ctx;
    916:   DM        sw;
    917:   KSP       ksp;
    918:   Mat       M_p, D_p;
    919:   Vec       f, v, E, tmpMom;
    920:   Vec       m, mold, mfluxold, mres, n, nrhs, nflux, nres, p, prhs, pflux, pres, e, erhs, eflux, eres;
    921:   PetscReal dt, t;
    922:   PetscInt  Nts;

    924:   PetscFunctionBegin;
    925:   PetscCall(TSGetStepNumber(ts, &Nts));
    926:   PetscCall(TSGetTimeStep(ts, &dt));
    927:   PetscCall(TSGetTime(ts, &t));
    928:   PetscCall(TSGetDM(ts, &sw));
    929:   PetscCall(DMGetApplicationContext(sw, &ctx));
    930:   PetscCall(DMSwarmSetCellDMActive(sw, "moment fields"));
    931:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
    932:   // TODO In higher dimensions, we will have to create different M_p and D_p for each field
    933:   PetscCall(DMCreateMassMatrix(sw, ctx->dmN, &M_p));
    934:   PetscCall(DMCreateGradientMatrix(sw, ctx->dmN, &D_p));
    935:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
    936:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &v));
    937:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "E_field", &E));
    938:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

    940:   PetscCall(MatViewFromOptions(ctx->MN, NULL, "-mn_view"));
    941:   PetscCall(MatViewFromOptions(ctx->MP, NULL, "-mp_view"));
    942:   PetscCall(MatViewFromOptions(ctx->ME, NULL, "-me_view"));
    943:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

    945:   PetscCall(DMGetGlobalVector(ctx->dmN, &nrhs));
    946:   PetscCall(DMGetGlobalVector(ctx->dmN, &nflux));
    947:   PetscCall(PetscObjectSetName((PetscObject)nrhs, "Weak number density"));
    948:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
    949:   PetscCall(DMGetGlobalVector(ctx->dmP, &prhs));
    950:   PetscCall(DMGetGlobalVector(ctx->dmP, &pflux));
    951:   PetscCall(PetscObjectSetName((PetscObject)prhs, "Weak momentum density"));
    952:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
    953:   PetscCall(DMGetGlobalVector(ctx->dmE, &erhs));
    954:   PetscCall(DMGetGlobalVector(ctx->dmE, &eflux));
    955:   PetscCall(PetscObjectSetName((PetscObject)erhs, "Weak energy density (pressure)"));
    956:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));

    958:   // Compute moments and fluxes
    959:   PetscCall(VecDuplicate(f, &tmpMom));

    961:   PetscCall(MatMultTranspose(M_p, f, nrhs));

    963:   PetscCall(VecPointwiseMult(tmpMom, f, v));
    964:   PetscCall(MatMultTranspose(M_p, tmpMom, prhs));
    965:   PetscCall(MatMultTranspose(D_p, tmpMom, nflux));

    967:   PetscCall(VecPointwiseMult(tmpMom, tmpMom, v));
    968:   PetscCall(MatMultTranspose(M_p, tmpMom, erhs));
    969:   PetscCall(MatMultTranspose(D_p, tmpMom, pflux));

    971:   PetscCall(VecPointwiseMult(tmpMom, tmpMom, v));
    972:   PetscCall(MatMultTranspose(D_p, tmpMom, eflux));

    974:   PetscCall(VecPointwiseMult(tmpMom, f, E));
    975:   PetscCall(MatMultTransposeAdd(M_p, tmpMom, pflux, pflux));

    977:   PetscCall(VecPointwiseMult(tmpMom, v, E));
    978:   PetscCall(VecScale(tmpMom, 2.));
    979:   PetscCall(MatMultTransposeAdd(M_p, tmpMom, eflux, eflux));

    981:   PetscCall(VecDestroy(&tmpMom));
    982:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &v));
    983:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
    984:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "E_field", &E));

    986:   PetscCall(MatDestroy(&M_p));
    987:   PetscCall(MatDestroy(&D_p));

    989:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)sw), &ksp));
    990:   PetscCall(KSPSetOptionsPrefix(ksp, "mom_proj_"));
    991:   PetscCall(KSPSetOperators(ksp, ctx->MN, ctx->MN));
    992:   PetscCall(KSPSetFromOptions(ksp));
    993:   PetscCall(KSPSolve(ksp, nrhs, n));
    994:   PetscCall(KSPSetOperators(ksp, ctx->MP, ctx->MP));
    995:   PetscCall(KSPSetFromOptions(ksp));
    996:   PetscCall(KSPSolve(ksp, prhs, p));
    997:   PetscCall(KSPSetOperators(ksp, ctx->ME, ctx->ME));
    998:   PetscCall(KSPSetFromOptions(ksp));
    999:   PetscCall(KSPSolve(ksp, erhs, e));
    1000:   PetscCall(KSPDestroy(&ksp));
    1001:   PetscCall(DMRestoreGlobalVector(ctx->dmN, &nrhs));
    1002:   PetscCall(DMRestoreGlobalVector(ctx->dmP, &prhs));
    1003:   PetscCall(DMRestoreGlobalVector(ctx->dmE, &erhs));

    1005:   // Check moment residual
    1006:   // TODO Fix global2local here
    1007:   PetscReal res[3], logres[3];

    1009:   PetscCall(DMGetGlobalVector(ctx->dmMom, &m));
    1010:   PetscCall(VecISCopy(m, ctx->isN, SCATTER_FORWARD, n));
    1011:   PetscCall(VecISCopy(m, ctx->isP, SCATTER_FORWARD, p));
    1012:   PetscCall(VecISCopy(m, ctx->isE, SCATTER_FORWARD, e));
    1013:   PetscCall(DMGetNamedGlobalVector(ctx->dmMom, "mold", &mold));
    1014:   PetscCall(DMGetNamedGlobalVector(ctx->dmMom, "mfluxold", &mfluxold));
    1015:   if (!Nts) goto end;

    1017:   // e = \Tr{\tau}
    1018:   // M_p w^{k+1} - M_p w^k - \Delta t D_p (w^k \vb{v}^k) = 0
    1019:   // M_p \vb{p}^{k+1} - M_p \vb{p}^k - \Delta t D_p \tau - e \Delta t M_p \left( n \vb{E} \right) = 0
    1020:   // M_p e^{k+1} - M_p e^k - \Delta t D_p \vb{Q} - 2 e \Delta t M_p \left( \vb{p} \cdot \vb{E} \right) = 0
    1021:   PetscCall(DMGetGlobalVector(ctx->dmMom, &mres));
    1022:   PetscCall(VecCopy(mfluxold, mres));
    1023:   PetscCall(VecAXPBYPCZ(mres, 1. / dt, -1. / dt, -1., m, mold));

    1025:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
    1026:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
    1027:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
    1028:   PetscCall(VecISCopy(mres, ctx->isN, SCATTER_REVERSE, nres));
    1029:   PetscCall(VecISCopy(mres, ctx->isP, SCATTER_REVERSE, pres));
    1030:   PetscCall(VecISCopy(mres, ctx->isE, SCATTER_REVERSE, eres));
    1031:   PetscCall(VecNorm(nres, NORM_2, &res[0]));
    1032:   PetscCall(VecNorm(pres, NORM_2, &res[1]));
    1033:   PetscCall(VecNorm(eres, NORM_2, &res[2]));
    1034:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Mass Residual: %g\n", (double)res[0]));
    1035:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Momentum Residual: %g\n", (double)res[1]));
    1036:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Energy Residual: %g\n", (double)res[2]));
    1037:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));
    1038:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));
    1039:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
    1040:   PetscCall(DMRestoreGlobalVector(ctx->dmMom, &mres));

    1042:   for (PetscInt i = 0; i < 3; ++i) logres[i] = PetscLog10Real(res[i]);
    1043:   PetscCall(PetscDrawLGAddCommonPoint(ctx->drawlgMomRes, t, logres));
    1044:   PetscCall(PetscDrawLGDraw(ctx->drawlgMomRes));
    1045:   {
    1046:     PetscDraw draw;

    1048:     PetscCall(PetscDrawLGGetDraw(ctx->drawlgMomRes, &draw));
    1049:     PetscCall(PetscDrawSave(draw));
    1050:   }

    1052: end:
    1053:   PetscCall(VecCopy(m, mold));
    1054:   PetscCall(DMRestoreGlobalVector(ctx->dmMom, &m));
    1055:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmMom, "mold", &mold));
    1056:   PetscCall(VecISCopy(mfluxold, ctx->isN, SCATTER_FORWARD, nflux));
    1057:   PetscCall(VecISCopy(mfluxold, ctx->isP, SCATTER_FORWARD, pflux));
    1058:   PetscCall(VecISCopy(mfluxold, ctx->isE, SCATTER_FORWARD, eflux));
    1059:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmMom, "mfluxold", &mfluxold));

    1061:   PetscCall(DMRestoreGlobalVector(ctx->dmN, &nflux));
    1062:   PetscCall(DMRestoreGlobalVector(ctx->dmP, &pflux));
    1063:   PetscCall(DMRestoreGlobalVector(ctx->dmE, &eflux));
    1064:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));
    1065:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));
    1066:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));
    1067:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));
    1068:   PetscFunctionReturn(PETSC_SUCCESS);
    1069: }

    1071: static PetscErrorCode MonitorMomentFields(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    1072: {
    1073:   AppCtx *ctx = (AppCtx *)Ctx;
    1074:   Vec     n, p, e;
    1075:   Vec     nres, pres, eres;

    1077:   PetscFunctionBeginUser;
    1078:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
    1079:   PetscCall(ComputeMomentFields(ts));

    1081:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "n", &n));
    1082:   PetscCall(VecView(n, ctx->viewerN));
    1083:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "n", &n));

    1085:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "p", &p));
    1086:   PetscCall(VecView(p, ctx->viewerP));
    1087:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "p", &p));

    1089:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "e", &e));
    1090:   PetscCall(VecView(e, ctx->viewerE));
    1091:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "e", &e));

    1093:   PetscCall(DMGetNamedGlobalVector(ctx->dmN, "nres", &nres));
    1094:   PetscCall(VecView(nres, ctx->viewerNRes));
    1095:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmN, "nres", &nres));

    1097:   PetscCall(DMGetNamedGlobalVector(ctx->dmP, "pres", &pres));
    1098:   PetscCall(VecView(pres, ctx->viewerPRes));
    1099:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmP, "pres", &pres));

    1101:   PetscCall(DMGetNamedGlobalVector(ctx->dmE, "eres", &eres));
    1102:   PetscCall(VecView(eres, ctx->viewerERes));
    1103:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmE, "eres", &eres));
    1104:   PetscFunctionReturn(PETSC_SUCCESS);
    1105: }

    1107: PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    1108: {
    1109:   AppCtx    *ctx = (AppCtx *)Ctx;
    1110:   DM         sw;
    1111:   PetscDraw  drawic_x, drawic_v;
    1112:   PetscReal *weight, *pos, *vel;
    1113:   PetscInt   dim, Np;

    1115:   PetscFunctionBegin;
    1116:   if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */
    1117:   if (step == 0) {
    1118:     PetscCall(TSGetDM(ts, &sw));
    1119:     PetscCall(DMGetDimension(sw, &dim));
    1120:     PetscCall(DMSwarmGetLocalSize(sw, &Np));

    1122:     PetscCall(PetscDrawHGReset(ctx->drawhgic_x));
    1123:     PetscCall(PetscDrawHGGetDraw(ctx->drawhgic_x, &drawic_x));
    1124:     PetscCall(PetscDrawClear(drawic_x));
    1125:     PetscCall(PetscDrawFlush(drawic_x));

    1127:     PetscCall(PetscDrawHGReset(ctx->drawhgic_v));
    1128:     PetscCall(PetscDrawHGGetDraw(ctx->drawhgic_v, &drawic_v));
    1129:     PetscCall(PetscDrawClear(drawic_v));
    1130:     PetscCall(PetscDrawFlush(drawic_v));

    1132:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
    1133:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    1134:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    1135:     for (PetscInt p = 0; p < Np; ++p) {
    1136:       PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgic_x, pos[p * dim], weight[p]));
    1137:       PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgic_v, vel[p * dim], weight[p]));
    1138:     }
    1139:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos));
    1140:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    1141:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));

    1143:     PetscCall(PetscDrawHGDraw(ctx->drawhgic_x));
    1144:     PetscCall(PetscDrawHGSave(ctx->drawhgic_x));
    1145:     PetscCall(PetscDrawHGDraw(ctx->drawhgic_v));
    1146:     PetscCall(PetscDrawHGSave(ctx->drawhgic_v));
    1147:   }
    1148:   PetscFunctionReturn(PETSC_SUCCESS);
    1149: }

    1151: // Right now, make the complete velocity histogram
    1152: PetscErrorCode MonitorVelocity(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    1153: {
    1154:   AppCtx      *ctx = (AppCtx *)Ctx;
    1155:   DM           sw, dm;
    1156:   Vec          ks;
    1157:   PetscProbFn *cdf;
    1158:   PetscDraw    drawcell_v;
    1159:   PetscScalar *ksa;
    1160:   PetscReal   *weight, *vel;
    1161:   PetscInt    *pidx;
    1162:   PetscInt     dim, Npc, cStart, cEnd, cell = ctx->velocity_monitor;

    1164:   PetscFunctionBegin;
    1165:   PetscCall(TSGetDM(ts, &sw));
    1166:   PetscCall(DMGetDimension(sw, &dim));

    1168:   PetscCall(DMSwarmGetCellDM(sw, &dm));
    1169:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
    1170:   PetscCall(VecCreate(PetscObjectComm((PetscObject)dm), &ks));
    1171:   PetscCall(PetscObjectSetName((PetscObject)ks, "KS Statistic by Cell"));
    1172:   PetscCall(VecSetSizes(ks, cEnd - cStart, PETSC_DETERMINE));
    1173:   PetscCall(VecSetFromOptions(ks));
    1174:   switch (dim) {
    1175:   case 1:
    1176:     //cdf = PetscCDFMaxwellBoltzmann1D;
    1177:     cdf = PetscCDFGaussian1D;
    1178:     break;
    1179:   case 2:
    1180:     cdf = PetscCDFMaxwellBoltzmann2D;
    1181:     break;
    1182:   case 3:
    1183:     cdf = PetscCDFMaxwellBoltzmann3D;
    1184:     break;
    1185:   default:
    1186:     SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported", dim);
    1187:   }

    1189:   PetscCall(PetscDrawHGReset(ctx->drawhgcell_v));
    1190:   PetscCall(PetscDrawHGGetDraw(ctx->drawhgcell_v, &drawcell_v));
    1191:   PetscCall(PetscDrawClear(drawcell_v));
    1192:   PetscCall(PetscDrawFlush(drawcell_v));

    1194:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    1195:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    1196:   PetscCall(DMSwarmSortGetAccess(sw));
    1197:   PetscCall(VecGetArrayWrite(ks, &ksa));
    1198:   for (PetscInt c = cStart; c < cEnd; ++c) {
    1199:     Vec          cellv, cellw;
    1200:     PetscScalar *cella, *cellaw;
    1201:     PetscReal    totWgt = 0.;

    1203:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    1204:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellv));
    1205:     PetscCall(VecSetBlockSize(cellv, dim));
    1206:     PetscCall(VecSetSizes(cellv, Npc * dim, Npc));
    1207:     PetscCall(VecSetFromOptions(cellv));
    1208:     PetscCall(VecCreate(PETSC_COMM_SELF, &cellw));
    1209:     PetscCall(VecSetSizes(cellw, Npc, Npc));
    1210:     PetscCall(VecSetFromOptions(cellw));
    1211:     PetscCall(VecGetArrayWrite(cellv, &cella));
    1212:     PetscCall(VecGetArrayWrite(cellw, &cellaw));
    1213:     for (PetscInt q = 0; q < Npc; ++q) {
    1214:       const PetscInt p = pidx[q];
    1215:       if (c == cell) PetscCall(PetscDrawHGAddWeightedValue(ctx->drawhgcell_v, vel[p * dim], weight[p]));
    1216:       for (PetscInt d = 0; d < dim; ++d) cella[q * dim + d] = vel[p * dim + d];
    1217:       cellaw[q] = weight[p];
    1218:       totWgt += weight[p];
    1219:     }
    1220:     PetscCall(VecRestoreArrayWrite(cellv, &cella));
    1221:     PetscCall(VecRestoreArrayWrite(cellw, &cellaw));
    1222:     PetscCall(VecScale(cellw, 1. / totWgt));
    1223:     PetscCall(PetscProbComputeKSStatisticWeighted(cellv, cellw, cdf, &ksa[c - cStart]));
    1224:     PetscCall(VecDestroy(&cellv));
    1225:     PetscCall(VecDestroy(&cellw));
    1226:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    1227:   }
    1228:   PetscCall(VecRestoreArrayWrite(ks, &ksa));
    1229:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    1230:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    1231:   PetscCall(DMSwarmSortRestoreAccess(sw));

    1233:   PetscReal minalpha, maxalpha;
    1234:   PetscInt  mincell, maxcell;

    1236:   PetscCall(VecFilter(ks, PETSC_SMALL));
    1237:   PetscCall(VecMin(ks, &mincell, &minalpha));
    1238:   PetscCall(VecMax(ks, &maxcell, &maxalpha));
    1239:   PetscCall(PetscPrintf(PetscObjectComm((PetscObject)dm), "Step %" PetscInt_FMT ": Min/Max KS statistic %g/%g in cell %" PetscInt_FMT "/%" PetscInt_FMT "\n", step, minalpha, maxalpha, mincell, maxcell));
    1240:   PetscCall(VecViewFromOptions(ks, NULL, "-ks_view"));
    1241:   PetscCall(VecDestroy(&ks));

    1243:   PetscCall(PetscDrawHGDraw(ctx->drawhgcell_v));
    1244:   PetscCall(PetscDrawHGSave(ctx->drawhgcell_v));
    1245:   PetscFunctionReturn(PETSC_SUCCESS);
    1246: }

    1248: static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    1249: {
    1250:   AppCtx         *ctx = (AppCtx *)Ctx;
    1251:   DM              dm, sw;
    1252:   PetscDrawAxis   axis;
    1253:   char            title[1024];
    1254:   PetscScalar    *x, *v, *weight;
    1255:   PetscReal       lower[3], upper[3], speed;
    1256:   const PetscInt *s;
    1257:   PetscInt        dim, cStart, cEnd, c;

    1259:   PetscFunctionBeginUser;
    1260:   if (step > 0 && step % ctx->ostep == 0) {
    1261:     PetscCall(TSGetDM(ts, &sw));
    1262:     PetscCall(DMSwarmGetCellDM(sw, &dm));
    1263:     PetscCall(DMGetDimension(dm, &dim));
    1264:     PetscCall(DMGetBoundingBox(dm, lower, upper));
    1265:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
    1266:     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1267:     PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
    1268:     PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    1269:     PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s));
    1270:     PetscCall(DMSwarmSortGetAccess(sw));
    1271:     PetscCall(PetscDrawSPReset(ctx->drawspX));
    1272:     PetscCall(PetscDrawSPGetAxis(ctx->drawspX, &axis));
    1273:     PetscCall(PetscSNPrintf(title, 1024, "Step %" PetscInt_FMT " Time: %g", step, (double)t));
    1274:     PetscCall(PetscDrawAxisSetLabels(axis, title, "x", "v"));
    1275:     PetscCall(PetscDrawSPSetLimits(ctx->drawspX, lower[0], upper[0], lower[1], upper[1]));
    1276:     PetscCall(PetscDrawSPSetLimits(ctx->drawspX, lower[0], upper[0], -12, 12));
    1277:     for (c = 0; c < cEnd - cStart; ++c) {
    1278:       PetscInt *pidx, Npc, q;
    1279:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    1280:       for (q = 0; q < Npc; ++q) {
    1281:         const PetscInt p = pidx[q];
    1282:         if (s[p] == 0) {
    1283:           speed = 0.;
    1284:           for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]);
    1285:           speed = PetscSqrtReal(speed);
    1286:           if (dim == 1) {
    1287:             PetscCall(PetscDrawSPAddPointColorized(ctx->drawspX, &x[p * dim], &v[p * dim], &speed));
    1288:           } else {
    1289:             PetscCall(PetscDrawSPAddPointColorized(ctx->drawspX, &x[p * dim], &x[p * dim + 1], &speed));
    1290:           }
    1291:         } else if (s[p] == 1) {
    1292:           PetscCall(PetscDrawSPAddPoint(ctx->drawspX, &x[p * dim], &v[p * dim]));
    1293:         }
    1294:       }
    1295:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    1296:     }
    1297:     PetscCall(PetscDrawSPDraw(ctx->drawspX, PETSC_TRUE));
    1298:     PetscDraw draw;
    1299:     PetscCall(PetscDrawSPGetDraw(ctx->drawspX, &draw));
    1300:     PetscCall(PetscDrawSave(draw));
    1301:     PetscCall(DMSwarmSortRestoreAccess(sw));
    1302:     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1303:     PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    1304:     PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
    1305:     PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s));
    1306:   }
    1307:   PetscFunctionReturn(PETSC_SUCCESS);
    1308: }

    1310: static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *Ctx)
    1311: {
    1312:   AppCtx *ctx = (AppCtx *)Ctx;
    1313:   DM      dm, sw;

    1315:   PetscFunctionBeginUser;
    1316:   if (step > 0 && step % ctx->ostep == 0) {
    1317:     PetscCall(TSGetDM(ts, &sw));
    1318:     PetscCall(DMSwarmGetCellDM(sw, &dm));

    1320:     if (ctx->validE) {
    1321:       PetscScalar *x, *E, *weight;
    1322:       PetscReal    lower[3], upper[3], xval;
    1323:       PetscDraw    draw;
    1324:       PetscInt     dim, cStart, cEnd;

    1326:       PetscCall(DMGetDimension(dm, &dim));
    1327:       PetscCall(DMGetBoundingBox(dm, lower, upper));
    1328:       PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));

    1330:       PetscCall(PetscDrawSPReset(ctx->drawspE));
    1331:       PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1332:       PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    1333:       PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));

    1335:       PetscCall(DMSwarmSortGetAccess(sw));
    1336:       for (PetscInt c = 0; c < cEnd - cStart; ++c) {
    1337:         PetscReal Eavg = 0.0;
    1338:         PetscInt *pidx, Npc;

    1340:         PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    1341:         for (PetscInt q = 0; q < Npc; ++q) {
    1342:           const PetscInt p = pidx[q];
    1343:           Eavg += E[p * dim];
    1344:         }
    1345:         Eavg /= Npc;
    1346:         xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart));
    1347:         PetscCall(PetscDrawSPAddPoint(ctx->drawspE, &xval, &Eavg));
    1348:         PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    1349:       }
    1350:       PetscCall(PetscDrawSPDraw(ctx->drawspE, PETSC_TRUE));
    1351:       PetscCall(PetscDrawSPGetDraw(ctx->drawspE, &draw));
    1352:       PetscCall(PetscDrawSave(draw));
    1353:       PetscCall(DMSwarmSortRestoreAccess(sw));
    1354:       PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1355:       PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    1356:       PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    1357:     }

    1359:     Vec rho, rhohat, phi;

    1361:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
    1362:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));
    1363:     PetscCall(VecView(rho, ctx->viewerRho));
    1364:     PetscCall(VecISCopy(ctx->fftX, ctx->fftReal, SCATTER_FORWARD, rho));
    1365:     PetscCall(MatMult(ctx->fftPot, ctx->fftX, ctx->fftY));
    1366:     PetscCall(VecFilter(ctx->fftY, PETSC_SMALL));
    1367:     PetscCall(VecViewFromOptions(ctx->fftX, NULL, "-real_view"));
    1368:     PetscCall(VecViewFromOptions(ctx->fftY, NULL, "-fft_view"));
    1369:     PetscCall(VecISCopy(ctx->fftY, ctx->fftReal, SCATTER_REVERSE, rhohat));
    1370:     PetscCall(VecSetValue(rhohat, 0, 0., INSERT_VALUES)); // Remove large DC component
    1371:     PetscCall(VecView(rhohat, ctx->viewerRhoHat));
    1372:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
    1373:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rhohat", &rhohat));

    1375:     PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    1376:     PetscCall(VecView(phi, ctx->viewerPhi));
    1377:     PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
    1378:   }
    1379:   PetscFunctionReturn(PETSC_SUCCESS);
    1380: }

    1382: static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
    1383: {
    1384:   PetscBag   bag;
    1385:   Parameter *p;

    1387:   PetscFunctionBeginUser;
    1388:   /* setup PETSc parameter bag */
    1389:   PetscCall(PetscBagGetData(ctx->bag, &p));
    1390:   PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
    1391:   bag = ctx->bag;
    1392:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
    1393:   PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
    1394:   PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
    1395:   PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
    1396:   PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
    1397:   PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
    1398:   PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
    1399:   PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));

    1401:   PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
    1402:   PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
    1403:   PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
    1404:   PetscCall(PetscBagSetFromOptions(bag));
    1405:   {
    1406:     PetscViewer       viewer;
    1407:     PetscViewerFormat format;
    1408:     PetscBool         flg;

    1410:     PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
    1411:     if (flg) {
    1412:       PetscCall(PetscViewerPushFormat(viewer, format));
    1413:       PetscCall(PetscBagView(bag, viewer));
    1414:       PetscCall(PetscViewerFlush(viewer));
    1415:       PetscCall(PetscViewerPopFormat(viewer));
    1416:       PetscCall(PetscViewerDestroy(&viewer));
    1417:     }
    1418:   }
    1419:   PetscFunctionReturn(PETSC_SUCCESS);
    1420: }

    1422: static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *ctx, DM *dm)
    1423: {
    1424:   DMField         coordField;
    1425:   IS              cellIS;
    1426:   PetscQuadrature quad;
    1427:   PetscReal      *wt, *pt;
    1428:   PetscInt        cdim, cStart, cEnd;

    1430:   PetscFunctionBeginUser;
    1431:   PetscCall(DMCreate(comm, dm));
    1432:   PetscCall(DMSetType(*dm, DMPLEX));
    1433:   PetscCall(DMSetFromOptions(*dm));
    1434:   PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
    1435:   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));

    1437:   // Cache the mesh geometry
    1438:   PetscCall(DMGetCoordinateField(*dm, &coordField));
    1439:   PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
    1440:   PetscCall(DMGetCoordinateDim(*dm, &cdim));
    1441:   PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
    1442:   PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
    1443:   PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
    1444:   PetscCall(PetscMalloc1(1, &wt));
    1445:   PetscCall(PetscMalloc1(2, &pt));
    1446:   wt[0] = 1.;
    1447:   pt[0] = -1.;
    1448:   pt[1] = -1.;
    1449:   PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
    1450:   PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &ctx->fegeom));
    1451:   PetscCall(PetscQuadratureDestroy(&quad));
    1452:   PetscCall(ISDestroy(&cellIS));
    1453:   PetscFunctionReturn(PETSC_SUCCESS);
    1454: }

    1456: 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[])
    1457: {
    1458:   f0[0] = -constants[SIGMA];
    1459: }

    1461: 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[])
    1462: {
    1463:   PetscInt d;
    1464:   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
    1465: }

    1467: 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[])
    1468: {
    1469:   PetscInt d;
    1470:   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
    1471: }

    1473: /*
    1474:    /  I   -grad\ / q \ = /0\
    1475:    \-div    0  / \phi/   \f/
    1476: */
    1477: 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[])
    1478: {
    1479:   for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d];
    1480: }

    1482: 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[])
    1483: {
    1484:   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]];
    1485: }

    1487: 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[])
    1488: {
    1489:   f0[0] += constants[SIGMA];
    1490:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d];
    1491: }

    1493: /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */
    1494: 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[])
    1495: {
    1496:   for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0;
    1497: }

    1499: 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[])
    1500: {
    1501:   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0;
    1502: }

    1504: 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[])
    1505: {
    1506:   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0;
    1507: }

    1509: static PetscErrorCode CreateFEM(DM dm, AppCtx *ctx)
    1510: {
    1511:   PetscFE   fephi, feq;
    1512:   PetscDS   ds;
    1513:   PetscBool simplex;
    1514:   PetscInt  dim;

    1516:   PetscFunctionBeginUser;
    1517:   PetscCall(DMGetDimension(dm, &dim));
    1518:   PetscCall(DMPlexIsSimplex(dm, &simplex));
    1519:   if (ctx->em == EM_MIXED) {
    1520:     DMLabel        label;
    1521:     const PetscInt id = 1;

    1523:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
    1524:     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
    1525:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi));
    1526:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
    1527:     PetscCall(PetscFECopyQuadrature(feq, fephi));
    1528:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
    1529:     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi));
    1530:     PetscCall(DMCreateDS(dm));
    1531:     PetscCall(PetscFEDestroy(&fephi));
    1532:     PetscCall(PetscFEDestroy(&feq));

    1534:     PetscCall(DMGetLabel(dm, "marker", &label));
    1535:     PetscCall(DMGetDS(dm, &ds));

    1537:     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
    1538:     PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL));
    1539:     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
    1540:     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL));
    1541:     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL));

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

    1545:   } else {
    1546:     MatNullSpace nullsp;
    1547:     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
    1548:     PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
    1549:     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
    1550:     PetscCall(DMCreateDS(dm));
    1551:     PetscCall(DMGetDS(dm, &ds));
    1552:     PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
    1553:     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
    1554:     PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
    1555:     PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
    1556:     PetscCall(MatNullSpaceDestroy(&nullsp));
    1557:     PetscCall(PetscFEDestroy(&fephi));
    1558:   }
    1559:   PetscFunctionReturn(PETSC_SUCCESS);
    1560: }

    1562: static PetscErrorCode CreatePoisson(DM dm, AppCtx *ctx)
    1563: {
    1564:   SNES         snes;
    1565:   Mat          J;
    1566:   MatNullSpace nullSpace;

    1568:   PetscFunctionBeginUser;
    1569:   PetscCall(CreateFEM(dm, ctx));
    1570:   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
    1571:   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
    1572:   PetscCall(SNESSetDM(snes, dm));
    1573:   PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, ctx));
    1574:   PetscCall(SNESSetFromOptions(snes));

    1576:   PetscCall(DMCreateMatrix(dm, &J));
    1577:   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
    1578:   PetscCall(MatSetNullSpace(J, nullSpace));
    1579:   PetscCall(MatNullSpaceDestroy(&nullSpace));
    1580:   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
    1581:   PetscCall(MatDestroy(&J));
    1582:   if (ctx->em == EM_MIXED) {
    1583:     const PetscInt potential = 1;

    1585:     PetscCall(DMCreateSubDM(dm, 1, &potential, &ctx->isPot, &ctx->dmPot));
    1586:   } else {
    1587:     ctx->dmPot = dm;
    1588:     PetscCall(PetscObjectReference((PetscObject)ctx->dmPot));
    1589:   }
    1590:   PetscCall(DMCreateMassMatrix(ctx->dmPot, ctx->dmPot, &ctx->M));
    1591:   PetscCall(DMPlexCreateClosureIndex(dm, NULL));
    1592:   ctx->snes = snes;
    1593:   PetscFunctionReturn(PETSC_SUCCESS);
    1594: }

    1596: // Conservation of mass (m = 1.0)
    1597: // n_t + 1/ m p_x = 0
    1598: static void f0_mass(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[])
    1599: {
    1600:   for (PetscInt d = 0; d < dim; ++d) f0[0] += u_t[uOff[0]] + u_x[uOff_x[1] + d * dim + d];
    1601: }

    1603: // Conservation of momentum (m = 1, e = 1)
    1604: // p_t + (u p)_x = -pr_x + e n E
    1605: // p_t + (div u) p + u . grad p = -pr_x + e n E
    1606: // p_t + (div p) p / n - (p . grad n) p / n^2 + p / n . grad p = -pr_x + e n E
    1607: static void f0_momentum(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[])
    1608: {
    1609:   const PetscScalar n = u[uOff[0]];

    1611:   for (PetscInt d = 0; d < dim; ++d) {
    1612:     PetscReal divp = 0.;

    1614:     f0[d] += u_t[uOff[1] + d];
    1615:     for (PetscInt e = 0; e < dim; ++e) {
    1616:       f0[d] += u[uOff[1] + e] * u_x[uOff_x[1] + d * dim + e] / n;                    // p / n . grad p
    1617:       f0[d] -= (u[uOff[1] + e] * u_x[uOff_x[0] + e]) * u[uOff[1] + d] / PetscSqr(n); // -(p . grad n) p / n^2
    1618:       divp += u_x[uOff_x[1] + e * dim + e];
    1619:     }
    1620:     f0[d] += divp * u[uOff[1] + d] / n; // (div p) p / n
    1621:     f0[d] += u_x[uOff_x[2] + d];        // pr_x
    1622:     f0[d] -= n * a[d];                  // -e n E
    1623:   }
    1624: }

    1626: // Conservation of energy
    1627: // pr_t + (u pr)_x = -3 pr u_x - q_x
    1628: // pr_t + (div u) pr + u . grad pr = -3 pr (div u) - q_x
    1629: // pr_t + 4 (div u) pr + u . grad pr = -q_x
    1630: // pr_t + 4 div p pr / n - 4 (p . grad n) pr / n^2 + p . grad pr / n = -q_x
    1631: static void f0_energy(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[])
    1632: {
    1633:   const PetscScalar n    = u[uOff[0]];
    1634:   const PetscScalar pr   = u[uOff[2]];
    1635:   PetscReal         divp = 0.;

    1637:   f0[0] += u_t[uOff[2]];
    1638:   for (PetscInt d = 0; d < dim; ++d) {
    1639:     f0[0] += u[uOff[1] + d] * u_x[uOff_x[2] + d] / n;                     // p . grad pr / n
    1640:     f0[0] -= 4. * u[uOff[1] + d] * u_x[uOff_x[0] + d] * pr / PetscSqr(n); // -4 (p . grad n) pr / n^2
    1641:     divp += u_x[uOff_x[1] + d * dim + d];
    1642:   }
    1643:   f0[0] += 4. * divp * pr / n; // 4 div p pr / n
    1644: }

    1646: static PetscErrorCode SetupMomentProblem(DM dm, AppCtx *ctx)
    1647: {
    1648:   PetscDS ds;

    1650:   PetscFunctionBegin;
    1651:   PetscCall(DMGetDS(dm, &ds));
    1652:   PetscCall(PetscDSSetResidual(ds, 0, f0_mass, NULL));
    1653:   PetscCall(PetscDSSetResidual(ds, 1, f0_momentum, NULL));
    1654:   PetscCall(PetscDSSetResidual(ds, 2, f0_energy, NULL));
    1655:   //PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_mass_uu, NULL, NULL, NULL));
    1656:   PetscFunctionReturn(PETSC_SUCCESS);
    1657: }

    1659: static PetscErrorCode CreateMomentFields(DM odm, AppCtx *ctx)
    1660: {
    1661:   DM             dm;
    1662:   PetscFE        fe;
    1663:   DMPolytopeType ct;
    1664:   PetscInt       dim, cStart;

    1666:   PetscFunctionBeginUser;
    1667:   PetscCall(DMClone(odm, &dm));
    1668:   PetscCall(DMGetDimension(dm, &dim));
    1669:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, NULL));
    1670:   PetscCall(DMPlexGetCellType(dm, cStart, &ct));
    1671:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
    1672:   PetscCall(PetscObjectSetName((PetscObject)fe, "number density"));
    1673:   PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fe));
    1674:   PetscCall(PetscFEDestroy(&fe));
    1675:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, dim, ct, NULL, PETSC_DETERMINE, &fe));
    1676:   PetscCall(PetscObjectSetName((PetscObject)fe, "momentum density"));
    1677:   PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fe));
    1678:   PetscCall(PetscFEDestroy(&fe));
    1679:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe));
    1680:   PetscCall(PetscObjectSetName((PetscObject)fe, "energy density"));
    1681:   PetscCall(DMSetField(dm, 2, NULL, (PetscObject)fe));
    1682:   PetscCall(PetscFEDestroy(&fe));
    1683:   PetscCall(DMCreateDS(dm));
    1684:   PetscCall(SetupMomentProblem(dm, ctx));
    1685:   ctx->dmMom = dm;
    1686:   PetscInt field;

    1688:   field = 0;
    1689:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isN, &ctx->dmN));
    1690:   PetscCall(DMCreateMassMatrix(ctx->dmN, ctx->dmN, &ctx->MN));
    1691:   field = 1;
    1692:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isP, &ctx->dmP));
    1693:   PetscCall(DMCreateMassMatrix(ctx->dmP, ctx->dmP, &ctx->MP));
    1694:   field = 2;
    1695:   PetscCall(DMCreateSubDM(ctx->dmMom, 1, &field, &ctx->isE, &ctx->dmE));
    1696:   PetscCall(DMCreateMassMatrix(ctx->dmE, ctx->dmE, &ctx->ME));
    1697:   PetscFunctionReturn(PETSC_SUCCESS);
    1698: }

    1700: PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
    1701: {
    1702:   p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
    1703:   p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI);
    1704:   return PETSC_SUCCESS;
    1705: }
    1706: PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
    1707: {
    1708:   p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
    1709:   return PETSC_SUCCESS;
    1710: }

    1712: PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
    1713: {
    1714:   const PetscReal alpha = scale ? scale[0] : 0.0;
    1715:   const PetscReal k     = scale ? scale[1] : 1.;
    1716:   p[0]                  = (1 + alpha * PetscCosReal(k * x[0]));
    1717:   return PETSC_SUCCESS;
    1718: }

    1720: PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
    1721: {
    1722:   const PetscReal alpha = scale ? scale[0] : 0.;
    1723:   const PetscReal k     = scale ? scale[0] : 1.;
    1724:   p[0]                  = (1 + alpha * PetscCosReal(k * (x[0] + x[1])));
    1725:   return PETSC_SUCCESS;
    1726: }

    1728: static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
    1729: {
    1730:   PetscFE        fe;
    1731:   DMPolytopeType ct;
    1732:   PetscInt       dim, cStart;
    1733:   const char    *prefix = "v";

    1735:   PetscFunctionBegin;
    1736:   PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
    1737:   PetscCall(DMSetType(*vdm, DMPLEX));
    1738:   PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
    1739:   PetscCall(DMSetFromOptions(*vdm));
    1740:   PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
    1741:   PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));

    1743:   PetscCall(DMGetDimension(*vdm, &dim));
    1744:   PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
    1745:   PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
    1746:   PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
    1747:   PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
    1748:   PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
    1749:   PetscCall(DMCreateDS(*vdm));
    1750:   PetscCall(PetscFEDestroy(&fe));
    1751:   PetscFunctionReturn(PETSC_SUCCESS);
    1752: }

    1754: /*
    1755:   InitializeParticles_Centroid - Initialize a regular grid of particles.

    1757:   Input Parameters:
    1758: + sw      - The `DMSWARM`
    1759: - force1D - Treat the spatial domain as 1D

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

    1764:   It places one particle in the centroid of each cell in the implicit tensor product of the spatial
    1765:   and velocity meshes.
    1766: */
    1767: static PetscErrorCode InitializeParticles_Centroid(DM sw)
    1768: {
    1769:   DM_Swarm     *swarm = (DM_Swarm *)sw->data;
    1770:   DMSwarmCellDM celldm;
    1771:   DM            xdm, vdm;
    1772:   PetscReal     vmin[3], vmax[3];
    1773:   PetscReal    *x, *v;
    1774:   PetscInt     *species, *cellid;
    1775:   PetscInt      dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
    1776:   PetscBool     flg;
    1777:   MPI_Comm      comm;
    1778:   const char   *cellidname;

    1780:   PetscFunctionBegin;
    1781:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));

    1783:   PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
    1784:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
    1785:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
    1786:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
    1787:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
    1788:   PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
    1789:   PetscOptionsEnd();
    1790:   debug = swarm->printCoords;

    1792:   PetscCall(DMGetDimension(sw, &dim));
    1793:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
    1794:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));

    1796:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
    1797:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    1798:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

    1800:   // One particle per centroid on the tensor product grid
    1801:   Npc = (vcEnd - vcStart) * Ns;
    1802:   Np  = (xcEnd - xcStart) * Npc;
    1803:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
    1804:   if (debug) {
    1805:     PetscInt gNp, gNc, Nc = xcEnd - xcStart;
    1806:     PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
    1807:     PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
    1808:     PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
    1809:     PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
    1810:     PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
    1811:   }

    1813:   // Set species and cellid
    1814:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
    1815:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
    1816:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    1817:   PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
    1818:   for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
    1819:     for (PetscInt s = 0; s < Ns; ++s) {
    1820:       for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
    1821:         species[p] = s;
    1822:         cellid[p]  = c;
    1823:       }
    1824:     }
    1825:   }
    1826:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    1827:   PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));

    1829:   // Set particle coordinates
    1830:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1831:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
    1832:   PetscCall(DMSwarmSortGetAccess(sw));
    1833:   PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
    1834:   PetscCall(DMGetCoordinatesLocalSetUp(xdm));
    1835:   for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
    1836:     const PetscInt cell = c + xcStart;
    1837:     PetscInt      *pidx, Npc;
    1838:     PetscReal      centroid[3], volume;

    1840:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    1841:     PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
    1842:     for (PetscInt s = 0; s < Ns; ++s) {
    1843:       for (PetscInt q = 0; q < Npc / Ns; ++q) {
    1844:         const PetscInt p = pidx[q * Ns + s];

    1846:         for (PetscInt d = 0; d < dim; ++d) {
    1847:           x[p * dim + d] = centroid[d];
    1848:           v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns));
    1849:         }
    1850:         if (debug > 1) {
    1851:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
    1852:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "  x: ("));
    1853:           for (PetscInt d = 0; d < dim; ++d) {
    1854:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    1855:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
    1856:           }
    1857:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
    1858:           for (PetscInt d = 0; d < dim; ++d) {
    1859:             if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
    1860:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
    1861:           }
    1862:           PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
    1863:         }
    1864:       }
    1865:     }
    1866:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    1867:   }
    1868:   PetscCall(DMSwarmSortRestoreAccess(sw));
    1869:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
    1870:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
    1871:   PetscFunctionReturn(PETSC_SUCCESS);
    1872: }

    1874: /*
    1875:   InitializeWeights - Compute weight for each local particle

    1877:   Input Parameters:
    1878: + sw          - The `DMSwarm`
    1879: . totalWeight - The sum of all particle weights
    1880: . func        - The PDF for the particle spatial distribution
    1881: - param       - The PDF parameters

    1883:   Notes:
    1884:   The PDF for velocity is assumed to be a Gaussian

    1886:   The particle weights are returned in the `w_q` field of `sw`.
    1887: */
    1888: static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
    1889: {
    1890:   DM               xdm, vdm;
    1891:   DMSwarmCellDM    celldm;
    1892:   PetscScalar     *weight;
    1893:   PetscQuadrature  xquad;
    1894:   const PetscReal *xq, *xwq;
    1895:   const PetscInt   order = 5;
    1896:   PetscReal        xi0[3];
    1897:   PetscReal        xwtot = 0., pwtot = 0.;
    1898:   PetscInt         xNq;
    1899:   PetscInt         dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
    1900:   MPI_Comm         comm;
    1901:   PetscMPIInt      rank;

    1903:   PetscFunctionBegin;
    1904:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
    1905:   PetscCallMPI(MPI_Comm_rank(comm, &rank));
    1906:   PetscCall(DMGetDimension(sw, &dim));
    1907:   PetscCall(DMSwarmGetCellDM(sw, &xdm));
    1908:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
    1909:   PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
    1910:   PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
    1911:   PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
    1912:   PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));

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

    1919:   // Integrate the density function to get the weights of particles in each cell
    1920:   PetscCall(DMGetCoordinatesLocalSetUp(vdm));
    1921:   PetscCall(DMSwarmSortGetAccess(sw));
    1922:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    1923:   for (PetscInt c = xcStart; c < xcEnd; ++c) {
    1924:     PetscReal          xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
    1925:     PetscInt          *pidx, Npc;
    1926:     PetscInt           xNc;
    1927:     const PetscScalar *xarray;
    1928:     PetscScalar       *xcoords = NULL;
    1929:     PetscBool          xisDG;

    1931:     PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
    1932:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
    1933:     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);
    1934:     PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
    1935:     for (PetscInt q = 0; q < xNq; ++q) {
    1936:       // Transform quadrature points from ref space to real space
    1937:       CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
    1938:       // Get probability density at quad point
    1939:       //   No need to scale xqr since PDF will be periodic
    1940:       PetscCall((*func)(xqr, param, &xden));
    1941:       xw += xden * (xwq[q] * xdetJ);
    1942:     }
    1943:     xwtot += xw;
    1944:     if (debug) {
    1945:       IS              globalOrdering;
    1946:       const PetscInt *ordering;

    1948:       PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
    1949:       PetscCall(ISGetIndices(globalOrdering, &ordering));
    1950:       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));
    1951:       PetscCall(ISRestoreIndices(globalOrdering, &ordering));
    1952:     }
    1953:     // Set weights to be Gaussian in velocity cells
    1954:     for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
    1955:       const PetscInt     p  = pidx[vc * Ns + 0];
    1956:       PetscReal          vw = 0.;
    1957:       PetscInt           vNc;
    1958:       const PetscScalar *varray;
    1959:       PetscScalar       *vcoords = NULL;
    1960:       PetscBool          visDG;

    1962:       PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
    1963:       // TODO: Fix 2 stream Ask Joe
    1964:       //   Two stream function from 1/2pi v^2 e^(-v^2/2)
    1965:       //   vw = 1. / (PetscSqrtReal(2 * PETSC_PI)) * (((coords_v[0] * PetscExpReal(-PetscSqr(coords_v[0]) / 2.)) - (coords_v[1] * PetscExpReal(-PetscSqr(coords_v[1]) / 2.)))) - 0.5 * PetscErfReal(coords_v[0] / PetscSqrtReal(2.)) + 0.5 * (PetscErfReal(coords_v[1] / PetscSqrtReal(2.)));
    1966:       vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.)));

    1968:       weight[p] = totalWeight * vw * xw;
    1969:       pwtot += weight[p];
    1970:       PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 1: %g, %g, %g", p, xw, vw, totalWeight);
    1971:       PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
    1972:       if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
    1973:     }
    1974:     PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
    1975:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
    1976:   }
    1977:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    1978:   PetscCall(DMSwarmSortRestoreAccess(sw));
    1979:   PetscCall(PetscQuadratureDestroy(&xquad));

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

    1984:     PetscCall(PetscSynchronizedFlush(comm, NULL));
    1985:     PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    1986:     PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
    1987:   }
    1988:   PetscFunctionReturn(PETSC_SUCCESS);
    1989: }

    1991: static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *ctx)
    1992: {
    1993:   PetscReal scale[2] = {ctx->cosine_coefficients[0], ctx->cosine_coefficients[1]};
    1994:   PetscInt  dim;

    1996:   PetscFunctionBegin;
    1997:   PetscCall(DMGetDimension(sw, &dim));
    1998:   PetscCall(InitializeParticles_Centroid(sw));
    1999:   PetscCall(InitializeWeights(sw, ctx->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale));
    2000:   PetscFunctionReturn(PETSC_SUCCESS);
    2001: }

    2003: static PetscErrorCode InitializeConstants(DM sw, AppCtx *ctx)
    2004: {
    2005:   DM         dm;
    2006:   PetscInt  *species;
    2007:   PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
    2008:   PetscInt   Np, dim;

    2010:   PetscFunctionBegin;
    2011:   PetscCall(DMSwarmGetCellDM(sw, &dm));
    2012:   PetscCall(DMGetDimension(sw, &dim));
    2013:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2014:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
    2015:   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
    2016:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    2017:   for (PetscInt p = 0; p < Np; ++p) {
    2018:     totalWeight += weight[p];
    2019:     totalCharge += ctx->charges[species[p]] * weight[p];
    2020:   }
    2021:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
    2022:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    2023:   {
    2024:     Parameter *param;
    2025:     PetscReal  Area;

    2027:     PetscCall(PetscBagGetData(ctx->bag, &param));
    2028:     switch (dim) {
    2029:     case 1:
    2030:       Area = (gmax[0] - gmin[0]);
    2031:       break;
    2032:     case 2:
    2033:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
    2034:       break;
    2035:     case 3:
    2036:       Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
    2037:       break;
    2038:     default:
    2039:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
    2040:     }
    2041:     PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    2042:     PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
    2043:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, ctx->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)ctx->charges[0], (double)global_charge, (double)Area));
    2044:     param->sigma = PetscAbsReal(global_charge / (Area));

    2046:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
    2047:     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,
    2048:                           (double)param->vlasovNumber));
    2049:   }
    2050:   /* Setup Constants */
    2051:   {
    2052:     PetscDS    ds;
    2053:     Parameter *param;
    2054:     PetscCall(PetscBagGetData(ctx->bag, &param));
    2055:     PetscScalar constants[NUM_CONSTANTS];
    2056:     constants[SIGMA]   = param->sigma;
    2057:     constants[V0]      = param->v0;
    2058:     constants[T0]      = param->t0;
    2059:     constants[X0]      = param->x0;
    2060:     constants[M0]      = param->m0;
    2061:     constants[Q0]      = param->q0;
    2062:     constants[PHI0]    = param->phi0;
    2063:     constants[POISSON] = param->poissonNumber;
    2064:     constants[VLASOV]  = param->vlasovNumber;
    2065:     PetscCall(DMGetDS(dm, &ds));
    2066:     PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
    2067:   }
    2068:   PetscFunctionReturn(PETSC_SUCCESS);
    2069: }

    2071: static PetscErrorCode CreateSwarm(DM dm, AppCtx *ctx, DM *sw)
    2072: {
    2073:   DMSwarmCellDM celldm;
    2074:   DM            vdm;
    2075:   PetscReal     v0[2] = {1., 0.};
    2076:   PetscInt      dim;

    2078:   PetscFunctionBeginUser;
    2079:   PetscCall(DMGetDimension(dm, &dim));
    2080:   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
    2081:   PetscCall(DMSetType(*sw, DMSWARM));
    2082:   PetscCall(DMSetDimension(*sw, dim));
    2083:   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
    2084:   PetscCall(DMSetApplicationContext(*sw, ctx));

    2086:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
    2087:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
    2088:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
    2089:   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));

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

    2093:   PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
    2094:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    2095:   PetscCall(DMSwarmCellDMDestroy(&celldm));

    2097:   const char *vfieldnames[2] = {"w_q"};

    2099:   PetscCall(CreateVelocityDM(*sw, &vdm));
    2100:   PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
    2101:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    2102:   PetscCall(DMSwarmCellDMDestroy(&celldm));
    2103:   PetscCall(DMDestroy(&vdm));

    2105:   DM mdm;

    2107:   PetscCall(DMClone(dm, &mdm));
    2108:   PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
    2109:   PetscCall(DMCopyDisc(dm, mdm));
    2110:   PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
    2111:   PetscCall(DMDestroy(&mdm));
    2112:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    2113:   PetscCall(DMSwarmCellDMDestroy(&celldm));

    2115:   DM mfdm;

    2117:   PetscCall(DMClone(dm, &mfdm));
    2118:   PetscCall(PetscObjectSetName((PetscObject)mfdm, "moment fields"));
    2119:   PetscCall(DMCopyDisc(dm, mfdm));
    2120:   PetscCall(DMSwarmCellDMCreate(mfdm, 1, &fieldnames[1], 1, fieldnames, &celldm));
    2121:   PetscCall(DMDestroy(&mfdm));
    2122:   PetscCall(DMSwarmAddCellDM(*sw, celldm));
    2123:   PetscCall(DMSwarmCellDMDestroy(&celldm));

    2125:   PetscCall(DMSetFromOptions(*sw));
    2126:   PetscCall(DMSetUp(*sw));

    2128:   PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
    2129:   ctx->swarm = *sw;
    2130:   // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
    2131:   if (ctx->perturbed_weights) {
    2132:     PetscCall(InitializeParticles_PerturbedWeights(*sw, ctx));
    2133:   } else {
    2134:     PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
    2135:     PetscCall(DMSwarmInitializeCoordinates(*sw));
    2136:     PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
    2137:   }
    2138:   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
    2139:   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
    2140:   PetscFunctionReturn(PETSC_SUCCESS);
    2141: }

    2143: static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
    2144: {
    2145:   AppCtx     *ctx;
    2146:   PetscReal  *coords;
    2147:   PetscInt   *species, dim, Np, Ns;
    2148:   PetscMPIInt size;

    2150:   PetscFunctionBegin;
    2151:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
    2152:   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
    2153:   PetscCall(DMGetDimension(sw, &dim));
    2154:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2155:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
    2156:   PetscCall(DMGetApplicationContext(sw, &ctx));

    2158:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2159:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
    2160:   for (PetscInt p = 0; p < Np; ++p) {
    2161:     PetscReal *pcoord = &coords[p * dim];
    2162:     PetscReal  pE[3]  = {0., 0., 0.};

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

    2169:       if (p == q) continue;
    2170:       q_q = ctx->charges[species[q]] * 1.;
    2171:       for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
    2172:       r = DMPlex_NormD_Internal(dim, rpq);
    2173:       if (r < PETSC_SQRT_MACHINE_EPSILON) continue;
    2174:       r3 = PetscPowRealInt(r, 3);
    2175:       for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3;
    2176:     }
    2177:     for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d];
    2178:   }
    2179:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
    2180:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2181:   PetscFunctionReturn(PETSC_SUCCESS);
    2182: }

    2184: static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
    2185: {
    2186:   DM         dm;
    2187:   AppCtx    *ctx;
    2188:   PetscDS    ds;
    2189:   PetscFE    fe;
    2190:   KSP        ksp;
    2191:   Vec        rhoRhs;      // Weak charge density, \int phi_i rho
    2192:   Vec        rho;         // Charge density, M^{-1} rhoRhs
    2193:   Vec        phi, locPhi; // Potential
    2194:   Vec        f;           // Particle weights
    2195:   PetscReal *coords;
    2196:   PetscInt   dim, cStart, cEnd, Np;

    2198:   PetscFunctionBegin;
    2199:   PetscCall(DMGetApplicationContext(sw, &ctx));
    2200:   PetscCall(PetscLogEventBegin(ctx->ESolveEvent, snes, sw, 0, 0));
    2201:   PetscCall(DMGetDimension(sw, &dim));
    2202:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

    2204:   PetscCall(SNESGetDM(snes, &dm));
    2205:   PetscCall(DMGetGlobalVector(dm, &rhoRhs));
    2206:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
    2207:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
    2208:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
    2209:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

    2211:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
    2212:   PetscCall(MatViewFromOptions(ctx->M, NULL, "-m_view"));
    2213:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

    2215:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
    2216:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

    2218:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
    2219:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
    2220:   PetscCall(KSPSetOperators(ksp, ctx->M, ctx->M));
    2221:   PetscCall(KSPSetFromOptions(ksp));
    2222:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

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

    2226:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
    2227:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
    2228:   PetscCall(KSPDestroy(&ksp));

    2230:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    2231:   PetscCall(VecSet(phi, 0.0));
    2232:   PetscCall(SNESSolve(snes, rhoRhs, phi));
    2233:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
    2234:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

    2236:   PetscCall(DMGetLocalVector(dm, &locPhi));
    2237:   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
    2238:   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
    2239:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));
    2240:   PetscCall(PetscLogEventEnd(ctx->ESolveEvent, snes, sw, 0, 0));

    2242:   PetscCall(DMGetDS(dm, &ds));
    2243:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
    2244:   PetscCall(DMSwarmSortGetAccess(sw));
    2245:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
    2246:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

    2248:   PetscCall(PetscLogEventBegin(ctx->ETabEvent, snes, sw, 0, 0));
    2249:   PetscTabulation tab;
    2250:   PetscReal      *pcoord, *refcoord;
    2251:   PetscFEGeom    *chunkgeom = NULL;
    2252:   PetscInt        maxNcp    = 0;

    2254:   for (PetscInt c = cStart; c < cEnd; ++c) {
    2255:     PetscInt Ncp;

    2257:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
    2258:     maxNcp = PetscMax(maxNcp, Ncp);
    2259:   }
    2260:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    2261:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    2262:   // This can raise an FP_INEXACT in the dgemm inside
    2263:   PetscCall(PetscFPTrapPush(PETSC_FP_TRAP_OFF));
    2264:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
    2265:   PetscCall(PetscFPTrapPop());
    2266:   for (PetscInt c = cStart; c < cEnd; ++c) {
    2267:     PetscScalar *clPhi = NULL;
    2268:     PetscInt    *points;
    2269:     PetscInt     Ncp;

    2271:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
    2272:     for (PetscInt cp = 0; cp < Ncp; ++cp)
    2273:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
    2274:     {
    2275:       PetscCall(PetscFEGeomGetChunk(ctx->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
    2276:       for (PetscInt i = 0; i < Ncp; ++i) {
    2277:         const PetscReal x0[3] = {-1., -1., -1.};
    2278:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
    2279:       }
    2280:     }
    2281:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
    2282:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    2283:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
    2284:       const PetscReal *basisDer = tab->T[1];
    2285:       const PetscInt   p        = points[cp];

    2287:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
    2288:       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
    2289:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
    2290:     }
    2291:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    2292:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
    2293:   }
    2294:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    2295:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    2296:   PetscCall(PetscTabulationDestroy(&tab));
    2297:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2298:   PetscCall(DMSwarmSortRestoreAccess(sw));
    2299:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
    2300:   PetscCall(PetscFEGeomRestoreChunk(ctx->fegeom, 0, 1, &chunkgeom));
    2301:   PetscCall(PetscLogEventEnd(ctx->ETabEvent, snes, sw, 0, 0));
    2302:   PetscFunctionReturn(PETSC_SUCCESS);
    2303: }

    2305: static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, Mat M_p, PetscReal E[])
    2306: {
    2307:   DM         dm;
    2308:   AppCtx    *ctx;
    2309:   PetscDS    ds;
    2310:   PetscFE    fe;
    2311:   KSP        ksp;
    2312:   Vec        rhoRhs, rhoRhsFull;   // Weak charge density, \int phi_i rho, and embedding in mixed problem
    2313:   Vec        rho;                  // Charge density, M^{-1} rhoRhs
    2314:   Vec        phi, locPhi, phiFull; // Potential and embedding in mixed problem
    2315:   Vec        f;                    // Particle weights
    2316:   PetscReal *coords;
    2317:   PetscInt   dim, cStart, cEnd, Np;

    2319:   PetscFunctionBegin;
    2320:   PetscCall(DMGetApplicationContext(sw, &ctx));
    2321:   PetscCall(PetscLogEventBegin(ctx->ESolveEvent, snes, sw, 0, 0));
    2322:   PetscCall(DMGetDimension(sw, &dim));
    2323:   PetscCall(DMSwarmGetLocalSize(sw, &Np));

    2325:   PetscCall(SNESGetDM(snes, &dm));
    2326:   PetscCall(DMGetGlobalVector(ctx->dmPot, &rhoRhs));
    2327:   PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
    2328:   PetscCall(DMGetGlobalVector(dm, &rhoRhsFull));
    2329:   PetscCall(PetscObjectSetName((PetscObject)rhoRhsFull, "Weak charge density"));
    2330:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "rho", &rho));
    2331:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
    2332:   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));

    2334:   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
    2335:   PetscCall(MatViewFromOptions(ctx->M, NULL, "-m_view"));
    2336:   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));

    2338:   PetscCall(MatMultTranspose(M_p, f, rhoRhs));
    2339:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));

    2341:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
    2342:   PetscCall(KSPSetOptionsPrefix(ksp, "em_proj"));
    2343:   PetscCall(KSPSetOperators(ksp, ctx->M, ctx->M));
    2344:   PetscCall(KSPSetFromOptions(ksp));
    2345:   PetscCall(KSPSolve(ksp, rhoRhs, rho));

    2347:   PetscCall(VecISCopy(rhoRhsFull, ctx->isPot, SCATTER_FORWARD, rhoRhs));
    2348:   //PetscCall(VecScale(rhoRhsFull, -1.0));

    2350:   PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
    2351:   PetscCall(VecViewFromOptions(rhoRhsFull, NULL, "-rho_full_view"));
    2352:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "rho", &rho));
    2353:   PetscCall(DMRestoreGlobalVector(ctx->dmPot, &rhoRhs));
    2354:   PetscCall(KSPDestroy(&ksp));

    2356:   PetscCall(DMGetGlobalVector(dm, &phiFull));
    2357:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    2358:   PetscCall(VecSet(phiFull, 0.0));
    2359:   PetscCall(SNESSolve(snes, rhoRhsFull, phiFull));
    2360:   PetscCall(DMRestoreGlobalVector(dm, &rhoRhsFull));
    2361:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));

    2363:   PetscCall(VecISCopy(phiFull, ctx->isPot, SCATTER_REVERSE, phi));
    2364:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

    2366:   PetscCall(DMGetLocalVector(dm, &locPhi));
    2367:   PetscCall(DMGlobalToLocalBegin(dm, phiFull, INSERT_VALUES, locPhi));
    2368:   PetscCall(DMGlobalToLocalEnd(dm, phiFull, INSERT_VALUES, locPhi));
    2369:   PetscCall(DMRestoreGlobalVector(dm, &phiFull));
    2370:   PetscCall(PetscLogEventEnd(ctx->ESolveEvent, snes, sw, 0, 0));

    2372:   PetscCall(DMGetDS(dm, &ds));
    2373:   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
    2374:   PetscCall(DMSwarmSortGetAccess(sw));
    2375:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
    2376:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));

    2378:   PetscCall(PetscLogEventBegin(ctx->ETabEvent, snes, sw, 0, 0));
    2379:   PetscTabulation tab;
    2380:   PetscReal      *pcoord, *refcoord;
    2381:   PetscFEGeom    *chunkgeom = NULL;
    2382:   PetscInt        maxNcp    = 0;

    2384:   for (PetscInt c = cStart; c < cEnd; ++c) {
    2385:     PetscInt Ncp;

    2387:     PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
    2388:     maxNcp = PetscMax(maxNcp, Ncp);
    2389:   }
    2390:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    2391:   PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    2392:   PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
    2393:   for (PetscInt c = cStart; c < cEnd; ++c) {
    2394:     PetscScalar *clPhi = NULL;
    2395:     PetscInt    *points;
    2396:     PetscInt     Ncp;

    2398:     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
    2399:     for (PetscInt cp = 0; cp < Ncp; ++cp)
    2400:       for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
    2401:     {
    2402:       PetscCall(PetscFEGeomGetChunk(ctx->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
    2403:       for (PetscInt i = 0; i < Ncp; ++i) {
    2404:         const PetscReal x0[3] = {-1., -1., -1.};
    2405:         CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
    2406:       }
    2407:     }
    2408:     PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
    2409:     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    2410:     for (PetscInt cp = 0; cp < Ncp; ++cp) {
    2411:       const PetscInt p = points[cp];

    2413:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
    2414:       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim]));
    2415:       PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim]));
    2416:       for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
    2417:     }
    2418:     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
    2419:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
    2420:   }
    2421:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
    2422:   PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
    2423:   PetscCall(PetscTabulationDestroy(&tab));
    2424:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2425:   PetscCall(DMSwarmSortRestoreAccess(sw));
    2426:   PetscCall(DMRestoreLocalVector(dm, &locPhi));
    2427:   PetscCall(PetscFEGeomRestoreChunk(ctx->fegeom, 0, 1, &chunkgeom));
    2428:   PetscCall(PetscLogEventEnd(ctx->ETabEvent, snes, sw, 0, 0));
    2429:   PetscFunctionReturn(PETSC_SUCCESS);
    2430: }

    2432: static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
    2433: {
    2434:   AppCtx    *ctx;
    2435:   Mat        M_p;
    2436:   PetscReal *E;
    2437:   PetscInt   dim, Np;

    2439:   PetscFunctionBegin;
    2442:   PetscCall(DMGetDimension(sw, &dim));
    2443:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2444:   PetscCall(DMGetApplicationContext(sw, &ctx));

    2446:   PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
    2447:   // TODO: Could share sort context with space cellDM
    2448:   PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
    2449:   PetscCall(DMCreateMassMatrix(sw, ctx->dmPot, &M_p));
    2450:   PetscCall(DMSwarmSetCellDMActive(sw, "space"));

    2452:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    2453:   PetscCall(PetscArrayzero(E, Np * dim));
    2454:   ctx->validE = PETSC_TRUE;

    2456:   switch (ctx->em) {
    2457:   case EM_COULOMB:
    2458:     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
    2459:     break;
    2460:   case EM_PRIMAL:
    2461:     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
    2462:     break;
    2463:   case EM_MIXED:
    2464:     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, M_p, E));
    2465:     break;
    2466:   case EM_NONE:
    2467:     break;
    2468:   default:
    2469:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
    2470:   }
    2471:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    2472:   PetscCall(MatDestroy(&M_p));
    2473:   PetscFunctionReturn(PETSC_SUCCESS);
    2474: }

    2476: static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, PetscCtx ctx)
    2477: {
    2478:   DM                 sw;
    2479:   SNES               snes = ((AppCtx *)ctx)->snes;
    2480:   const PetscScalar *u;
    2481:   PetscScalar       *g;
    2482:   PetscReal         *E, m_p = 1., q_p = -1.;
    2483:   PetscInt           dim, d, Np, p;

    2485:   PetscFunctionBeginUser;
    2486:   PetscCall(TSGetDM(ts, &sw));
    2487:   PetscCall(ComputeFieldAtParticles(snes, sw));

    2489:   PetscCall(DMGetDimension(sw, &dim));
    2490:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2491:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    2492:   PetscCall(VecGetArrayRead(U, &u));
    2493:   PetscCall(VecGetArray(G, &g));
    2494:   Np /= 2 * dim;
    2495:   for (p = 0; p < Np; ++p) {
    2496:     for (d = 0; d < dim; ++d) {
    2497:       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
    2498:       g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
    2499:     }
    2500:   }
    2501:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    2502:   PetscCall(VecRestoreArrayRead(U, &u));
    2503:   PetscCall(VecRestoreArray(G, &g));
    2504:   PetscFunctionReturn(PETSC_SUCCESS);
    2505: }

    2507: /* J_{ij} = dF_i/dx_j
    2508:    J_p = (  0   1)
    2509:          (-w^2  0)
    2510:    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
    2511:         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
    2512: */
    2513: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, PetscCtx ctx)
    2514: {
    2515:   DM               sw;
    2516:   const PetscReal *coords, *vel;
    2517:   PetscInt         dim, d, Np, p, rStart;

    2519:   PetscFunctionBeginUser;
    2520:   PetscCall(TSGetDM(ts, &sw));
    2521:   PetscCall(DMGetDimension(sw, &dim));
    2522:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2523:   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
    2524:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2525:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    2526:   Np /= 2 * dim;
    2527:   for (p = 0; p < Np; ++p) {
    2528:     // TODO This is not right because dv/dx has the electric field in it
    2529:     PetscScalar vals[4] = {0., 1., -1., 0.};

    2531:     for (d = 0; d < dim; ++d) {
    2532:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
    2533:       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
    2534:     }
    2535:   }
    2536:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2537:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    2538:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
    2539:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
    2540:   PetscFunctionReturn(PETSC_SUCCESS);
    2541: }

    2543: static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *Ctx)
    2544: {
    2545:   AppCtx            *ctx = (AppCtx *)Ctx;
    2546:   DM                 sw;
    2547:   const PetscScalar *v;
    2548:   PetscScalar       *xres;
    2549:   PetscInt           Np, p, d, dim;

    2551:   PetscFunctionBeginUser;
    2552:   PetscCall(PetscLogEventBegin(ctx->RhsXEvent, ts, 0, 0, 0));
    2553:   PetscCall(TSGetDM(ts, &sw));
    2554:   PetscCall(DMGetDimension(sw, &dim));
    2555:   PetscCall(VecGetLocalSize(Xres, &Np));
    2556:   PetscCall(VecGetArrayRead(V, &v));
    2557:   PetscCall(VecGetArray(Xres, &xres));
    2558:   Np /= dim;
    2559:   for (p = 0; p < Np; ++p) {
    2560:     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
    2561:   }
    2562:   PetscCall(VecRestoreArrayRead(V, &v));
    2563:   PetscCall(VecRestoreArray(Xres, &xres));
    2564:   PetscCall(PetscLogEventEnd(ctx->RhsXEvent, ts, 0, 0, 0));
    2565:   PetscFunctionReturn(PETSC_SUCCESS);
    2566: }

    2568: static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *Ctx)
    2569: {
    2570:   DM                 sw;
    2571:   AppCtx            *ctx  = (AppCtx *)Ctx;
    2572:   SNES               snes = ((AppCtx *)ctx)->snes;
    2573:   const PetscScalar *x;
    2574:   PetscScalar       *vres;
    2575:   PetscReal         *E, m_p, q_p;
    2576:   PetscInt           Np, p, dim, d;
    2577:   Parameter         *param;

    2579:   PetscFunctionBeginUser;
    2580:   PetscCall(PetscLogEventBegin(ctx->RhsVEvent, ts, 0, 0, 0));
    2581:   PetscCall(TSGetDM(ts, &sw));
    2582:   PetscCall(ComputeFieldAtParticles(snes, sw));

    2584:   PetscCall(DMGetDimension(sw, &dim));
    2585:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    2586:   PetscCall(PetscBagGetData(ctx->bag, &param));
    2587:   m_p = ctx->masses[0] * param->m0;
    2588:   q_p = ctx->charges[0] * param->q0;
    2589:   PetscCall(VecGetLocalSize(Vres, &Np));
    2590:   PetscCall(VecGetArrayRead(X, &x));
    2591:   PetscCall(VecGetArray(Vres, &vres));
    2592:   Np /= dim;
    2593:   for (p = 0; p < Np; ++p) {
    2594:     for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
    2595:   }
    2596:   PetscCall(VecRestoreArrayRead(X, &x));
    2597:   /*
    2598:     Synchronized, ordered output for parallel/sequential test cases.
    2599:     In the 1D (on the 2D mesh) case, every y component should be zero.
    2600:   */
    2601:   if (ctx->checkVRes) {
    2602:     PetscBool pr = ctx->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
    2603:     PetscInt  step;

    2605:     PetscCall(TSGetStepNumber(ts, &step));
    2606:     if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
    2607:     for (PetscInt p = 0; p < Np; ++p) {
    2608:       if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
    2609:       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]));
    2610:     }
    2611:     if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
    2612:   }
    2613:   PetscCall(VecRestoreArray(Vres, &vres));
    2614:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    2615:   PetscCall(PetscLogEventEnd(ctx->RhsVEvent, ts, 0, 0, 0));
    2616:   PetscFunctionReturn(PETSC_SUCCESS);
    2617: }

    2619: /* Discrete Gradients Formulation: S, F, gradF (G) */
    2620: PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, PetscCtx ctx)
    2621: {
    2622:   PetscScalar vals[4] = {0., 1., -1., 0.};
    2623:   DM          sw;
    2624:   PetscInt    dim, d, Np, p, rStart;

    2626:   PetscFunctionBeginUser;
    2627:   PetscCall(TSGetDM(ts, &sw));
    2628:   PetscCall(DMGetDimension(sw, &dim));
    2629:   PetscCall(VecGetLocalSize(U, &Np));
    2630:   PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
    2631:   Np /= 2 * dim;
    2632:   for (p = 0; p < Np; ++p) {
    2633:     for (d = 0; d < dim; ++d) {
    2634:       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
    2635:       PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
    2636:     }
    2637:   }
    2638:   PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
    2639:   PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
    2640:   PetscFunctionReturn(PETSC_SUCCESS);
    2641: }

    2643: PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *Ctx)
    2644: {
    2645:   AppCtx            *ctx = (AppCtx *)Ctx;
    2646:   DM                 sw;
    2647:   Vec                phi;
    2648:   const PetscScalar *u;
    2649:   PetscInt           dim, Np, cStart, cEnd;
    2650:   PetscReal         *vel, *coords, m_p = 1.;

    2652:   PetscFunctionBeginUser;
    2653:   PetscCall(TSGetDM(ts, &sw));
    2654:   PetscCall(DMGetDimension(sw, &dim));
    2655:   PetscCall(DMPlexGetHeightStratum(ctx->dmPot, 0, &cStart, &cEnd));

    2657:   PetscCall(DMGetNamedGlobalVector(ctx->dmPot, "phi", &phi));
    2658:   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
    2659:   PetscCall(computeFieldEnergy(ctx->dmPot, phi, F));
    2660:   PetscCall(DMRestoreNamedGlobalVector(ctx->dmPot, "phi", &phi));

    2662:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2663:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    2664:   PetscCall(DMSwarmSortGetAccess(sw));
    2665:   PetscCall(VecGetArrayRead(U, &u));
    2666:   PetscCall(VecGetLocalSize(U, &Np));
    2667:   Np /= 2 * dim;
    2668:   for (PetscInt c = cStart; c < cEnd; ++c) {
    2669:     PetscInt *points;
    2670:     PetscInt  Ncp;

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

    2677:       *F += 0.5 * m_p * v2;
    2678:     }
    2679:     PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
    2680:   }
    2681:   PetscCall(VecRestoreArrayRead(U, &u));
    2682:   PetscCall(DMSwarmSortRestoreAccess(sw));
    2683:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2684:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    2685:   PetscFunctionReturn(PETSC_SUCCESS);
    2686: }

    2688: /* dF/dx = q E   dF/dv = v */
    2689: PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, PetscCtx ctx)
    2690: {
    2691:   DM                 sw;
    2692:   SNES               snes = ((AppCtx *)ctx)->snes;
    2693:   const PetscReal   *coords, *vel, *E;
    2694:   const PetscScalar *u;
    2695:   PetscScalar       *g;
    2696:   PetscReal          m_p = 1., q_p = -1.;
    2697:   PetscInt           dim, d, Np, p;

    2699:   PetscFunctionBeginUser;
    2700:   PetscCall(TSGetDM(ts, &sw));
    2701:   PetscCall(DMGetDimension(sw, &dim));
    2702:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2703:   PetscCall(VecGetArrayRead(U, &u));
    2704:   PetscCall(VecGetArray(G, &g));

    2706:   PetscLogEvent COMPUTEFIELD;
    2707:   PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
    2708:   PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
    2709:   PetscCall(ComputeFieldAtParticles(snes, sw));
    2710:   PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
    2711:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2712:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
    2713:   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
    2714:   for (p = 0; p < Np; ++p) {
    2715:     for (d = 0; d < dim; ++d) {
    2716:       g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
    2717:       g[(p * 2 + 1) * dim + d] = m_p * u[(p * 2 + 1) * dim + d];
    2718:     }
    2719:   }
    2720:   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
    2721:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
    2722:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
    2723:   PetscCall(VecRestoreArrayRead(U, &u));
    2724:   PetscCall(VecRestoreArray(G, &g));
    2725:   PetscFunctionReturn(PETSC_SUCCESS);
    2726: }

    2728: static PetscErrorCode CreateSolution(TS ts)
    2729: {
    2730:   DM       sw;
    2731:   Vec      u;
    2732:   PetscInt dim, Np;

    2734:   PetscFunctionBegin;
    2735:   PetscCall(TSGetDM(ts, &sw));
    2736:   PetscCall(DMGetDimension(sw, &dim));
    2737:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2738:   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
    2739:   PetscCall(VecSetBlockSize(u, dim));
    2740:   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
    2741:   PetscCall(VecSetUp(u));
    2742:   PetscCall(TSSetSolution(ts, u));
    2743:   PetscCall(VecDestroy(&u));
    2744:   PetscFunctionReturn(PETSC_SUCCESS);
    2745: }

    2747: static PetscErrorCode SetProblem(TS ts)
    2748: {
    2749:   AppCtx *ctx;
    2750:   DM      sw;

    2752:   PetscFunctionBegin;
    2753:   PetscCall(TSGetDM(ts, &sw));
    2754:   PetscCall(DMGetApplicationContext(sw, &ctx));
    2755:   // Define unified system for (X, V)
    2756:   {
    2757:     Mat      J;
    2758:     PetscInt dim, Np;

    2760:     PetscCall(DMGetDimension(sw, &dim));
    2761:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2762:     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
    2763:     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
    2764:     PetscCall(MatSetBlockSize(J, 2 * dim));
    2765:     PetscCall(MatSetFromOptions(J));
    2766:     PetscCall(MatSetUp(J));
    2767:     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, ctx));
    2768:     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, ctx));
    2769:     PetscCall(MatDestroy(&J));
    2770:   }
    2771:   /* Define split system for X and V */
    2772:   {
    2773:     Vec             u;
    2774:     IS              isx, isv, istmp;
    2775:     const PetscInt *idx;
    2776:     PetscInt        dim, Np, rstart;

    2778:     PetscCall(TSGetSolution(ts, &u));
    2779:     PetscCall(DMGetDimension(sw, &dim));
    2780:     PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2781:     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
    2782:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
    2783:     PetscCall(ISGetIndices(istmp, &idx));
    2784:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
    2785:     PetscCall(ISRestoreIndices(istmp, &idx));
    2786:     PetscCall(ISDestroy(&istmp));
    2787:     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
    2788:     PetscCall(ISGetIndices(istmp, &idx));
    2789:     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
    2790:     PetscCall(ISRestoreIndices(istmp, &idx));
    2791:     PetscCall(ISDestroy(&istmp));
    2792:     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
    2793:     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
    2794:     PetscCall(ISDestroy(&isx));
    2795:     PetscCall(ISDestroy(&isv));
    2796:     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, ctx));
    2797:     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, ctx));
    2798:   }
    2799:   // Define symplectic formulation U_t = S . G, where G = grad F
    2800:   {
    2801:     PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, ctx));
    2802:   }
    2803:   PetscFunctionReturn(PETSC_SUCCESS);
    2804: }

    2806: static PetscErrorCode DMSwarmTSRedistribute(TS ts)
    2807: {
    2808:   DM        sw;
    2809:   Vec       u;
    2810:   PetscReal t, maxt, dt;
    2811:   PetscInt  n, maxn;

    2813:   PetscFunctionBegin;
    2814:   PetscCall(TSGetDM(ts, &sw));
    2815:   PetscCall(TSGetTime(ts, &t));
    2816:   PetscCall(TSGetMaxTime(ts, &maxt));
    2817:   PetscCall(TSGetTimeStep(ts, &dt));
    2818:   PetscCall(TSGetStepNumber(ts, &n));
    2819:   PetscCall(TSGetMaxSteps(ts, &maxn));

    2821:   PetscCall(TSReset(ts));
    2822:   PetscCall(TSSetDM(ts, sw));
    2823:   PetscCall(TSSetFromOptions(ts));
    2824:   PetscCall(TSSetTime(ts, t));
    2825:   PetscCall(TSSetMaxTime(ts, maxt));
    2826:   PetscCall(TSSetTimeStep(ts, dt));
    2827:   PetscCall(TSSetStepNumber(ts, n));
    2828:   PetscCall(TSSetMaxSteps(ts, maxn));

    2830:   PetscCall(CreateSolution(ts));
    2831:   PetscCall(SetProblem(ts));
    2832:   PetscCall(TSGetSolution(ts, &u));
    2833:   PetscFunctionReturn(PETSC_SUCCESS);
    2834: }

    2836: PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *Ctx)
    2837: {
    2838:   DM        sw, cdm;
    2839:   PetscInt  Np;
    2840:   PetscReal low[2], high[2];
    2841:   AppCtx   *ctx = (AppCtx *)Ctx;

    2843:   sw = ctx->swarm;
    2844:   PetscCall(DMSwarmGetCellDM(sw, &cdm));
    2845:   // Get the bounding box so we can equally space the particles
    2846:   PetscCall(DMGetLocalBoundingBox(cdm, low, high));
    2847:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2848:   // shift it by h/2 so nothing is initialized directly on a boundary
    2849:   x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
    2850:   x[1] = 0.;
    2851:   return PETSC_SUCCESS;
    2852: }

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

    2857:   Input Parameters:
    2858: + ts         - The TS
    2859: - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem

    2861:   Output Parameters:
    2862: . u - The initialized solution vector

    2864:   Level: advanced

    2866: .seealso: InitializeSolve()
    2867: */
    2868: static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
    2869: {
    2870:   DM       sw;
    2871:   Vec      u, gc, gv;
    2872:   IS       isx, isv;
    2873:   PetscInt dim;
    2874:   AppCtx  *ctx;

    2876:   PetscFunctionBeginUser;
    2877:   PetscCall(TSGetDM(ts, &sw));
    2878:   PetscCall(DMGetApplicationContext(sw, &ctx));
    2879:   PetscCall(DMGetDimension(sw, &dim));
    2880:   if (useInitial) {
    2881:     PetscReal v0[2] = {1., 0.};
    2882:     if (ctx->perturbed_weights) {
    2883:       PetscCall(InitializeParticles_PerturbedWeights(sw, ctx));
    2884:     } else {
    2885:       PetscCall(DMSwarmComputeLocalSizeFromOptions(sw));
    2886:       PetscCall(DMSwarmInitializeCoordinates(sw));
    2887:       PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
    2888:     }
    2889:     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
    2890:     PetscCall(DMSwarmTSRedistribute(ts));
    2891:   }
    2892:   PetscCall(DMSetUp(sw));
    2893:   PetscCall(TSGetSolution(ts, &u));
    2894:   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
    2895:   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
    2896:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    2897:   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
    2898:   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
    2899:   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
    2900:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    2901:   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
    2902:   PetscFunctionReturn(PETSC_SUCCESS);
    2903: }

    2905: static PetscErrorCode InitializeSolve(TS ts, Vec u)
    2906: {
    2907:   PetscFunctionBegin;
    2908:   PetscCall(TSSetSolution(ts, u));
    2909:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
    2910:   PetscFunctionReturn(PETSC_SUCCESS);
    2911: }

    2913: static PetscErrorCode MigrateParticles(TS ts)
    2914: {
    2915:   DM               sw, cdm;
    2916:   const PetscReal *L;
    2917:   AppCtx          *ctx;

    2919:   PetscFunctionBeginUser;
    2920:   PetscCall(TSGetDM(ts, &sw));
    2921:   PetscCall(DMGetApplicationContext(sw, &ctx));
    2922:   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
    2923:   {
    2924:     Vec        u, gc, gv, position, momentum;
    2925:     IS         isx, isv;
    2926:     PetscReal *pos, *mom;

    2928:     PetscCall(TSGetSolution(ts, &u));
    2929:     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
    2930:     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
    2931:     PetscCall(VecGetSubVector(u, isx, &position));
    2932:     PetscCall(VecGetSubVector(u, isv, &momentum));
    2933:     PetscCall(VecGetArray(position, &pos));
    2934:     PetscCall(VecGetArray(momentum, &mom));
    2935:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    2936:     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
    2937:     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
    2938:     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));

    2940:     PetscCall(DMSwarmGetCellDM(sw, &cdm));
    2941:     PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
    2942:     PetscCheck(L, PetscObjectComm((PetscObject)cdm), PETSC_ERR_ARG_WRONG, "Mesh must be periodic");
    2943:     if ((L[0] || L[1]) >= 0.) {
    2944:       PetscReal *x, *v, upper[3], lower[3];
    2945:       PetscInt   Np, dim;

    2947:       PetscCall(DMSwarmGetLocalSize(sw, &Np));
    2948:       PetscCall(DMGetDimension(cdm, &dim));
    2949:       PetscCall(DMGetBoundingBox(cdm, lower, upper));
    2950:       PetscCall(VecGetArray(gc, &x));
    2951:       PetscCall(VecGetArray(gv, &v));
    2952:       for (PetscInt p = 0; p < Np; ++p) {
    2953:         for (PetscInt d = 0; d < dim; ++d) {
    2954:           if (pos[p * dim + d] < lower[d]) {
    2955:             x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
    2956:           } else if (pos[p * dim + d] > upper[d]) {
    2957:             x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
    2958:           } else {
    2959:             x[p * dim + d] = pos[p * dim + d];
    2960:           }
    2961:           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]);
    2962:           v[p * dim + d] = mom[p * dim + d];
    2963:         }
    2964:       }
    2965:       PetscCall(VecRestoreArray(gc, &x));
    2966:       PetscCall(VecRestoreArray(gv, &v));
    2967:     }
    2968:     PetscCall(VecRestoreArray(position, &pos));
    2969:     PetscCall(VecRestoreArray(momentum, &mom));
    2970:     PetscCall(VecRestoreSubVector(u, isx, &position));
    2971:     PetscCall(VecRestoreSubVector(u, isv, &momentum));
    2972:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
    2973:     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
    2974:   }
    2975:   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
    2976:   PetscInt step;

    2978:   PetscCall(TSGetStepNumber(ts, &step));
    2979:   if (!(step % ctx->remapFreq)) {
    2980:     // Monitor electric field before we destroy it
    2981:     PetscReal ptime;
    2982:     PetscInt  step;

    2984:     PetscCall(TSGetStepNumber(ts, &step));
    2985:     PetscCall(TSGetTime(ts, &ptime));
    2986:     if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx));
    2987:     if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx));
    2988:     PetscCall(DMSwarmRemap(sw));
    2989:     ctx->validE = PETSC_FALSE;
    2990:   }
    2991:   // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
    2992:   PetscCall(DMSwarmTSRedistribute(ts));
    2993:   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
    2994:   PetscFunctionReturn(PETSC_SUCCESS);
    2995: }

    2997: int main(int argc, char **argv)
    2998: {
    2999:   DM        dm, sw;
    3000:   TS        ts;
    3001:   Vec       u;
    3002:   PetscReal dt;
    3003:   PetscInt  maxn;
    3004:   AppCtx    ctx;

    3006:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
    3007:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &ctx));
    3008:   PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &ctx.bag));
    3009:   PetscCall(CreateMesh(PETSC_COMM_WORLD, &ctx, &dm));
    3010:   PetscCall(CreatePoisson(dm, &ctx));
    3011:   PetscCall(CreateMomentFields(dm, &ctx));
    3012:   PetscCall(CreateSwarm(dm, &ctx, &sw));
    3013:   PetscCall(SetupParameters(PETSC_COMM_WORLD, &ctx));
    3014:   PetscCall(InitializeConstants(sw, &ctx));
    3015:   PetscCall(DMSetApplicationContext(sw, &ctx));

    3017:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
    3018:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
    3019:   PetscCall(TSSetDM(ts, sw));
    3020:   PetscCall(TSSetMaxTime(ts, 0.1));
    3021:   PetscCall(TSSetTimeStep(ts, 0.00001));
    3022:   PetscCall(TSSetMaxSteps(ts, 100));
    3023:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

    3025:   if (ctx.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &ctx, NULL));
    3026:   if (ctx.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &ctx, NULL));
    3027:   if (ctx.moment_field_monitor) PetscCall(TSMonitorSet(ts, MonitorMomentFields, &ctx, NULL));
    3028:   if (ctx.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &ctx, NULL));
    3029:   if (ctx.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &ctx, NULL));
    3030:   if (ctx.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &ctx, NULL));
    3031:   if (ctx.velocity_monitor >= 0) PetscCall(TSMonitorSet(ts, MonitorVelocity, &ctx, NULL));

    3033:   PetscCall(TSSetFromOptions(ts));
    3034:   PetscCall(TSGetTimeStep(ts, &dt));
    3035:   PetscCall(TSGetMaxSteps(ts, &maxn));
    3036:   ctx.steps    = maxn;
    3037:   ctx.stepSize = dt;
    3038:   PetscCall(SetupContext(dm, sw, &ctx));
    3039:   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
    3040:   PetscCall(TSSetPostStep(ts, MigrateParticles));
    3041:   PetscCall(CreateSolution(ts));
    3042:   PetscCall(TSGetSolution(ts, &u));
    3043:   PetscCall(TSComputeInitialCondition(ts, u));
    3044:   PetscCall(CheckNonNegativeWeights(sw, &ctx));
    3045:   PetscCall(TSSolve(ts, NULL));

    3047:   if (ctx.checkLandau) {
    3048:     // We should get a lookup table based on charge density and \hat k
    3049:     const PetscReal gammaEx = -0.15336;
    3050:     const PetscReal omegaEx = 1.4156;
    3051:     const PetscReal tol     = 1e-2;

    3053:     PetscCheck(PetscAbsReal((ctx.gamma - gammaEx) / gammaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau gamma %g != %g", ctx.gamma, gammaEx);
    3054:     PetscCheck(PetscAbsReal((ctx.omega - omegaEx) / omegaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau omega %g != %g", ctx.omega, omegaEx);
    3055:   }

    3057:   PetscCall(SNESDestroy(&ctx.snes));
    3058:   PetscCall(DMDestroy(&ctx.dmN));
    3059:   PetscCall(ISDestroy(&ctx.isN));
    3060:   PetscCall(MatDestroy(&ctx.MN));
    3061:   PetscCall(DMDestroy(&ctx.dmP));
    3062:   PetscCall(ISDestroy(&ctx.isP));
    3063:   PetscCall(MatDestroy(&ctx.MP));
    3064:   PetscCall(DMDestroy(&ctx.dmE));
    3065:   PetscCall(ISDestroy(&ctx.isE));
    3066:   PetscCall(MatDestroy(&ctx.ME));
    3067:   PetscCall(DMDestroy(&ctx.dmMom));
    3068:   PetscCall(DMDestroy(&ctx.dmPot));
    3069:   PetscCall(ISDestroy(&ctx.isPot));
    3070:   PetscCall(MatDestroy(&ctx.M));
    3071:   PetscCall(PetscFEGeomDestroy(&ctx.fegeom));
    3072:   PetscCall(TSDestroy(&ts));
    3073:   PetscCall(DMDestroy(&sw));
    3074:   PetscCall(DMDestroy(&dm));
    3075:   PetscCall(DestroyContext(&ctx));
    3076:   PetscCall(PetscFinalize());
    3077:   return 0;
    3078: }

    3080: /*TEST

    3082:   build:
    3083:     requires: !complex double

    3085:   # This tests that we can compute the correct decay rate and frequency
    3086:   #   For gold runs, use -dm_plex_box_faces 160 -vdm_plex_box_faces 450 -remap_dm_plex_box_faces 80,150 -ts_max_steps 1000
    3087:   #                      -remap_freq 100 -emax_start_step 50 -emax_solve_step 100
    3088:   testset:
    3089:     args: -cosine_coefficients 0.01 -charges -1. -perturbed_weights -total_weight 1. \
    3090:           -dm_plex_dim 1 -dm_plex_box_faces 80 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 \
    3091:             -dm_plex_box_bd periodic -dm_plex_hash_location \
    3092:           -vdm_plex_dim 1 -vdm_plex_box_faces 220 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 \
    3093:             -vpetscspace_degree 2 -vdm_plex_hash_location \
    3094:           -remap_freq 1 -dm_swarm_remap_type pfak -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 \
    3095:             -remap_dm_plex_box_faces 40,110 -remap_dm_plex_box_bd periodic,none \
    3096:             -remap_dm_plex_box_lower 0.,-6. -remap_dm_plex_box_upper 12.5664,6. \
    3097:             -remap_petscspace_degree 1 -remap_dm_plex_hash_location \
    3098:             -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu \
    3099:           -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged \
    3100:             -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu \
    3101:           -ts_time_step 0.03 -ts_max_steps 2 -ts_max_time 100 \
    3102:           -emax_tao_type brgn -emax_tao_max_it 100 -emax_tao_brgn_regularization_type l2pure \
    3103:             -emax_tao_brgn_regularizer_weight 1e-5 -emax_tao_brgn_subsolver_tao_bnk_ksp_rtol 1e-12 \
    3104:             -emax_start_step 1 -emax_solve_step 1 \
    3105:           -output_step 1 -efield_monitor quiet

    3107:     test:
    3108:       suffix: landau_damping_1d_bs
    3109:       args: -ts_type basicsymplectic -ts_basicsymplectic_type 1

    3111:     test:
    3112:       suffix: landau_damping_1d_dg
    3113:       args: -ts_type discgrad -ts_discgrad_type average -snes_type qn

    3115: TEST*/