Date: 2026-07-30T21:14:31+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:     Discretize the moment equations using finite elements and project the moments into the finite element space. Use DMSWARM_REMAP_PFAK, which guarantees that the 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*/