1 static char help[] = "Landau Damping test using Vlasov-Poisson equations\n";
2
3 /*
4 This example solves the Vlasov-Poisson system for Landau damping (6X + 6V).
5 The system is solved using a Particle-In-Cell (PIC) method with DMSwarm for particles and DMPlex/PetscFE for the field.
6 This particular example uses the velocity mesh from DMPlexLandauCreateVelocitySpace for 3D velocity space.
7
8 Options:
9 -particle_monitor [prefix] : Output particle data (x, v, w, E) to binary files at each output step.
10 Optional prefix for filenames (default: "particles").
11
12 */
13 #include <petscts.h>
14 #include <petscdmplex.h>
15 #include <petscdmswarm.h>
16 #include <petscfe.h>
17 #include <petscds.h>
18 #include <petscbag.h>
19 #include <petscdraw.h>
20 #include <petscviewer.h>
21 #include <petsclandau.h>
22 #include <petscdmcomposite.h>
23 #include <petsc/private/dmpleximpl.h> /* For norm and dot */
24 #include <petsc/private/petscfeimpl.h> /* For interpolation */
25 #include <petsc/private/dmswarmimpl.h> /* For swarm debugging */
26 #include "petscdm.h"
27 #include "petscdmlabel.h"
28
29 PETSC_EXTERN PetscErrorCode stream(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
30 PETSC_EXTERN PetscErrorCode line(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
31
32 typedef enum {
33 V0,
34 X0,
35 T0,
36 M0,
37 Q0,
38 PHI0,
39 POISSON,
40 VLASOV,
41 SIGMA,
42 NUM_CONSTANTS
43 } ConstantType;
44 typedef struct {
45 PetscScalar v0; /* Velocity scale, often the thermal velocity */
46 PetscScalar t0; /* Time scale */
47 PetscScalar x0; /* Space scale */
48 PetscScalar m0; /* Mass scale */
49 PetscScalar q0; /* Charge scale */
50 PetscScalar kb;
51 PetscScalar epsi0;
52 PetscScalar phi0; /* Potential scale */
53 PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */
54 PetscScalar vlasovNumber; /* Non-Dimensional Vlasov Number */
55 PetscReal sigma; /* Nondimensional charge per length in x */
56 } Parameter;
57
58 typedef struct {
59 PetscBag bag; // Problem parameters
60 PetscBool error; // Flag for printing the error
61 PetscBool efield_monitor; // Flag to show electric field monitor
62 PetscBool moment_monitor; // Flag to show distribution moment monitor
63 char particle_monitor_prefix[PETSC_MAX_PATH_LEN];
64 PetscBool particle_monitor; // Flag to output particle data
65 PetscInt ostep; // Print the energy at each ostep time steps
66 PetscInt numParticles;
67 PetscReal timeScale; /* Nondimensionalizing time scale */
68 PetscReal charges[2]; /* The charges of each species */
69 PetscReal masses[2]; /* The masses of each species */
70 PetscReal thermal_energy[2]; /* Thermal Energy (used to get other constants)*/
71 PetscReal cosine_coefficients[2]; /*(alpha, k)*/
72 PetscReal totalWeight;
73 PetscReal stepSize;
74 PetscInt steps;
75 PetscReal initVel;
76 SNES snes; // EM solver
77 DM dmPot; // The DM for potential
78 Mat M; // The finite element mass matrix for potential
79 PetscFEGeom *fegeom; // Geometric information for the DM cells
80 PetscBool validE; // Flag to indicate E-field in swarm is valid
81 PetscReal drawlgEmin; // The minimum lg(E) to plot
82 PetscDrawLG drawlgE; // Logarithm of maximum electric field
83 DM swarm;
84 PetscBool checkweights;
85 PetscInt checkVRes; // Flag to check/output velocity residuals for nightly tests
86 DM landau_pack;
87 PetscBool use_landau_velocity_space;
88 PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent;
89 } AppCtx;
90
ProcessOptions(MPI_Comm comm,AppCtx * options)91 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
92 {
93 PetscFunctionBeginUser;
94 PetscInt d = 2;
95 PetscInt maxSpecies = 2;
96 options->error = PETSC_FALSE;
97 options->efield_monitor = PETSC_FALSE;
98 options->moment_monitor = PETSC_FALSE;
99 options->particle_monitor = PETSC_FALSE;
100 options->ostep = 100;
101 options->timeScale = 2.0e-14;
102 options->charges[0] = -1.0;
103 options->charges[1] = 1.0;
104 options->masses[0] = 1.0;
105 options->masses[1] = 1000.0;
106 options->thermal_energy[0] = 1.0;
107 options->thermal_energy[1] = 1.0;
108 options->cosine_coefficients[0] = 0.01;
109 options->cosine_coefficients[1] = 0.5;
110 options->initVel = 1;
111 options->totalWeight = 1.0;
112 options->validE = PETSC_FALSE;
113 options->drawlgEmin = -6;
114 options->drawlgE = NULL;
115 options->snes = NULL;
116 options->dmPot = NULL;
117 options->M = NULL;
118 options->numParticles = 32768;
119 options->checkweights = PETSC_FALSE;
120 options->checkVRes = 0;
121 options->landau_pack = NULL;
122 options->use_landau_velocity_space = PETSC_FALSE;
123
124 PetscOptionsBegin(comm, "", "Landau Damping options", "DMSWARM");
125 PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex3.c", options->error, &options->error, NULL));
126 PetscCall(PetscOptionsBool("-efield_monitor", "Flag to plot log(max E) over time", "ex3.c", options->efield_monitor, &options->efield_monitor, NULL));
127 PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", "ex3.c", options->drawlgEmin, &options->drawlgEmin, NULL));
128 PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", "ex3.c", options->moment_monitor, &options->moment_monitor, NULL));
129 PetscCall(PetscOptionsString("-particle_monitor", "Prefix for particle data files", "ex3.c", options->particle_monitor_prefix, options->particle_monitor_prefix, sizeof(options->particle_monitor_prefix), &options->particle_monitor));
130 PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", "ex3.c", options->checkweights, &options->checkweights, NULL));
131 PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", "ex3.c", options->checkVRes, &options->checkVRes, NULL));
132 PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex3.c", options->ostep, &options->ostep, NULL));
133 PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex3.c", options->timeScale, &options->timeScale, NULL));
134 PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", "ex3.c", options->initVel, &options->initVel, NULL));
135 PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", "ex3.c", options->totalWeight, &options->totalWeight, NULL));
136 PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", "ex3.c", options->cosine_coefficients, &d, NULL));
137 PetscCall(PetscOptionsRealArray("-charges", "Species charges", "ex3.c", options->charges, &maxSpecies, NULL));
138 PetscCall(PetscOptionsBool("-use_landau_velocity_space", "Use Landau velocity space", "ex3.c", options->use_landau_velocity_space, &options->use_landau_velocity_space, NULL));
139 PetscOptionsEnd();
140
141 PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent));
142 PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent));
143 PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent));
144 PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent));
145 PetscFunctionReturn(PETSC_SUCCESS);
146 }
147
SetupContext(DM dm,DM sw,AppCtx * user)148 static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user)
149 {
150 MPI_Comm comm;
151
152 PetscFunctionBeginUser;
153 PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
154 if (user->efield_monitor) {
155 PetscDraw draw;
156 PetscDrawAxis axis;
157
158 PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 300, 400, 300, &draw));
159 PetscCall(PetscDrawSetSave(draw, "ex3_Efield"));
160 PetscCall(PetscDrawSetFromOptions(draw));
161 PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE));
162 PetscCall(PetscDrawDestroy(&draw));
163 PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis));
164 PetscCall(PetscDrawAxisSetLabels(axis, "Electron Electric Field", "time", "E_max"));
165 PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.));
166 PetscCall(PetscDrawLGSetUseMarkers(user->drawlgE, PETSC_FALSE));
167 }
168 PetscFunctionReturn(PETSC_SUCCESS);
169 }
170
DestroyContext(AppCtx * user)171 static PetscErrorCode DestroyContext(AppCtx *user)
172 {
173 PetscFunctionBeginUser;
174 if (user->landau_pack) PetscCall(DMPlexLandauDestroyVelocitySpace(&user->landau_pack));
175 PetscCall(PetscDrawLGDestroy(&user->drawlgE));
176 PetscCall(PetscBagDestroy(&user->bag));
177 PetscFunctionReturn(PETSC_SUCCESS);
178 }
179
CheckNonNegativeWeights(DM sw,AppCtx * user)180 static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user)
181 {
182 const PetscScalar *w;
183 PetscInt Np;
184
185 PetscFunctionBeginUser;
186 if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS);
187 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
188 PetscCall(DMSwarmGetLocalSize(sw, &Np));
189 for (PetscInt p = 0; p < Np; ++p) PetscCheck(w[p] >= 0.0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " has negative weight %g", p, w[p]);
190 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
191 PetscFunctionReturn(PETSC_SUCCESS);
192 }
193
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[])194 static void f0_Dirichlet(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
195 {
196 for (PetscInt d = 0; d < dim; ++d) f0[0] += 0.5 * PetscSqr(u_x[d]);
197 }
198
computeFieldEnergy(DM dm,Vec u,PetscReal * En)199 static PetscErrorCode computeFieldEnergy(DM dm, Vec u, PetscReal *En)
200 {
201 PetscDS ds;
202 const PetscInt field = 0;
203 PetscInt Nf;
204 void *ctx;
205
206 PetscFunctionBegin;
207 PetscCall(DMGetApplicationContext(dm, &ctx));
208 PetscCall(DMGetDS(dm, &ds));
209 PetscCall(PetscDSGetNumFields(ds, &Nf));
210 PetscCheck(Nf == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_WRONG, "We currently only support 1 field, not %" PetscInt_FMT, Nf);
211 PetscCall(PetscDSSetObjective(ds, field, &f0_Dirichlet));
212 PetscCall(DMPlexComputeIntegralFEM(dm, u, En, ctx));
213 PetscFunctionReturn(PETSC_SUCCESS);
214 }
215
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[])216 static void f0_grad_phi2(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
217 {
218 f0[0] = 0.;
219 for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d]); // + d * dim cause segfault
220 }
221
MonitorEField(TS ts,PetscInt step,PetscReal t,Vec U,void * ctx)222 static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
223 {
224 AppCtx *user = (AppCtx *)ctx;
225 DM sw;
226 PetscScalar intESq;
227 PetscReal *E, *x, *weight;
228 PetscReal Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.;
229 PetscReal pmoments[16]; /* \int f, \int v f, \int v^2 f */
230 PetscInt *species, dim, Np, gNp;
231 MPI_Comm comm;
232
233 PetscFunctionBeginUser;
234 if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS);
235 PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
236 PetscCall(TSGetDM(ts, &sw));
237 PetscCall(DMGetDimension(sw, &dim));
238 PetscCall(DMSwarmGetLocalSize(sw, &Np));
239 PetscCall(DMSwarmGetSize(sw, &gNp));
240 PetscCall(DMSwarmSortGetAccess(sw));
241 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
242 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
243 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
244 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
245
246 for (PetscInt p = 0; p < Np; ++p) {
247 PetscReal Emag = 0.;
248 for (PetscInt d = 0; d < dim; ++d) {
249 PetscReal temp = PetscAbsReal(E[p * dim + d]);
250 if (temp > Emax) Emax = temp;
251 Emag += PetscSqr(E[p * dim + d]);
252 }
253 Enorm += PetscSqrtReal(Emag);
254 sum += E[p * dim];
255 chargesum += user->charges[0] * weight[p];
256 }
257 PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm));
258 PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Enorm, 1, MPIU_REAL, MPIU_SUM, comm));
259 PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &sum, 1, MPIU_REAL, MPIU_SUM, comm));
260 PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &chargesum, 1, MPIU_REAL, MPIU_SUM, comm));
261 lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.;
262 lgEmax = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin;
263 if (lgEmax < user->drawlgEmin) lgEmax = user->drawlgEmin;
264
265 PetscDS ds;
266 Vec phi;
267
268 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
269 PetscCall(DMGetDS(user->dmPot, &ds));
270 PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2));
271 PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user));
272 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
273
274 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
275 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
276 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
277 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
278 PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax));
279 PetscCall(PetscDrawLGDraw(user->drawlgE));
280 PetscDraw draw;
281 PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw));
282 PetscCall(PetscDrawSave(draw));
283
284 PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
285 PetscCall(PetscPrintf(comm, "E: %f\t%+e\t%e\t%f\t%20.15e\t%f\t%f\t%f\t%20.15e\t%20.15e\t%20.15e\t%" PetscInt_FMT "\t(%" PetscInt_FMT ")\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[1 + dim], (double)PetscSqrtReal(intESq), gNp, step));
286 PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view"));
287 PetscFunctionReturn(PETSC_SUCCESS);
288 }
289
MonitorMoments(TS ts,PetscInt step,PetscReal t,Vec U,void * ctx)290 static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
291 {
292 DM sw;
293 PetscReal pmoments[16]; /* \int f, \int v f, \int v^2 f */
294
295 PetscFunctionBeginUser;
296 if (step < 0) PetscFunctionReturn(PETSC_SUCCESS);
297 PetscCall(TSGetDM(ts, &sw));
298
299 PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments));
300
301 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3]));
302 PetscFunctionReturn(PETSC_SUCCESS);
303 }
304
MonitorParticles(TS ts,PetscInt step,PetscReal t,Vec U,void * ctx)305 static PetscErrorCode MonitorParticles(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
306 {
307 AppCtx *user = (AppCtx *)ctx;
308 DM sw;
309 PetscInt Np, dim;
310 const PetscReal *x, *v, *E;
311 const PetscScalar *w;
312 PetscViewer viewer;
313 char filename[64];
314 MPI_Comm comm;
315
316 PetscFunctionBeginUser;
317 if (step < 0 || step % user->ostep != 0) PetscFunctionReturn(PETSC_SUCCESS);
318 PetscCall(TSGetDM(ts, &sw));
319 PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
320 PetscCall(DMGetDimension(sw, &dim));
321 PetscCall(DMSwarmGetLocalSize(sw, &Np));
322
323 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
324 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
325 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w));
326 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
327
328 if (user->particle_monitor_prefix[0]) {
329 PetscCall(PetscSNPrintf(filename, sizeof(filename), "%s_step_%d.bin", user->particle_monitor_prefix, (int)step));
330 } else {
331 PetscCall(PetscSNPrintf(filename, sizeof(filename), "particles_step_%d.bin", (int)step));
332 }
333 PetscCall(PetscViewerBinaryOpen(comm, filename, FILE_MODE_WRITE, &viewer));
334
335 {
336 Vec vx, vv, vw, vE;
337 PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, x, &vx));
338 PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, v, &vv));
339 PetscCall(VecCreateMPIWithArray(comm, 1, Np, PETSC_DECIDE, w, &vw));
340 PetscCall(VecCreateMPIWithArray(comm, dim, Np * dim, PETSC_DECIDE, E, &vE));
341
342 PetscCall(VecView(vx, viewer));
343 PetscCall(VecView(vv, viewer));
344 PetscCall(VecView(vw, viewer));
345 PetscCall(VecView(vE, viewer));
346
347 PetscCall(VecDestroy(&vx));
348 PetscCall(VecDestroy(&vv));
349 PetscCall(VecDestroy(&vw));
350 PetscCall(VecDestroy(&vE));
351 }
352 PetscCall(PetscViewerDestroy(&viewer));
353
354 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
355 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
356 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w));
357 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
358 PetscFunctionReturn(PETSC_SUCCESS);
359 }
360
PetscPDFPertubedConstant2D(const PetscReal x[],const PetscReal dummy[],PetscReal p[])361 PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
362 {
363 p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (4. * PETSC_PI * 4. * PETSC_PI);
364 return PETSC_SUCCESS;
365 }
PetscPDFPertubedConstant1D(const PetscReal x[],const PetscReal dummy[],PetscReal p[])366 PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[])
367 {
368 p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI);
369 return PETSC_SUCCESS;
370 }
371
PetscPDFCosine1D(const PetscReal x[],const PetscReal scale[],PetscReal p[])372 PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
373 {
374 const PetscReal alpha = scale ? scale[0] : 0.0;
375 const PetscReal k = scale ? scale[1] : 1.;
376 p[0] = (1 + alpha * PetscCosReal(k * x[0]));
377 return PETSC_SUCCESS;
378 }
379
PetscPDFCosine2D(const PetscReal x[],const PetscReal scale[],PetscReal p[])380 PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
381 {
382 const PetscReal alpha = scale ? scale[0] : 0.;
383 const PetscReal k = scale ? scale[1] : 1.;
384 p[0] = (1 + alpha * PetscCosReal(k * x[0]));
385 return PETSC_SUCCESS;
386 }
387
PetscPDFCosine3D(const PetscReal x[],const PetscReal scale[],PetscReal p[])388 PetscErrorCode PetscPDFCosine3D(const PetscReal x[], const PetscReal scale[], PetscReal p[])
389 {
390 const PetscReal alpha = scale ? scale[0] : 0.;
391 const PetscReal k = scale ? scale[1] : 1.;
392 p[0] = (1 + alpha * PetscCosReal(k * x[0]));
393 return PETSC_SUCCESS;
394 }
395
CreateVelocityDM(DM sw,DM * vdm)396 static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm)
397 {
398 PetscFE fe;
399 DMPolytopeType ct;
400 PetscInt dim, cStart;
401 const char *prefix = "v";
402 AppCtx *user;
403
404 PetscFunctionBegin;
405 PetscCall(DMGetDimension(sw, &dim));
406 PetscCall(DMGetApplicationContext(sw, &user));
407 if (dim == 3 && user->use_landau_velocity_space) {
408 LandauCtx *ctx;
409 Vec X;
410 Mat J;
411
412 PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, prefix, &X, &J, &user->landau_pack));
413 PetscCall(DMGetApplicationContext(user->landau_pack, &ctx));
414 *vdm = ctx->plex[0];
415 PetscCall(PetscObjectReference((PetscObject)*vdm));
416 PetscCall(VecDestroy(&X));
417 PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
418 } else {
419 PetscCall(DMCreate(PETSC_COMM_SELF, vdm));
420 PetscCall(DMSetType(*vdm, DMPLEX));
421 PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix));
422 PetscCall(DMSetFromOptions(*vdm));
423 PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity"));
424 }
425 PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view"));
426
427 PetscCall(DMGetDimension(*vdm, &dim));
428 PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL));
429 PetscCall(DMPlexGetCellType(*vdm, cStart, &ct));
430 PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe));
431 PetscCall(PetscObjectSetName((PetscObject)fe, "distribution"));
432 PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe));
433 PetscCall(DMCreateDS(*vdm));
434 PetscCall(PetscFEDestroy(&fe));
435 PetscFunctionReturn(PETSC_SUCCESS);
436 }
437
438 /*
439 InitializeParticles_Centroid - Initialize a regular grid of particles.
440
441 Input Parameters:
442 + sw - The `DMSWARM`
443 - force1D - Treat the spatial domain as 1D
444
445 Notes:
446 This functions sets the species, cellid, spatial coordinate, and velocity fields for all particles.
447
448 It places one particle in the centroid of each cell in the implicit tensor product of the spatial
449 and velocity meshes.
450 */
InitializeParticles_Centroid(DM sw)451 static PetscErrorCode InitializeParticles_Centroid(DM sw)
452 {
453 DM_Swarm *swarm = (DM_Swarm *)sw->data;
454 DMSwarmCellDM celldm;
455 DM xdm, vdm;
456 PetscReal vmin[3], vmax[3];
457 PetscReal *x, *v;
458 PetscInt *species, *cellid;
459 PetscInt dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug;
460 PetscBool flg;
461 MPI_Comm comm;
462 const char *cellidname;
463
464 PetscFunctionBegin;
465 PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
466
467 PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM");
468 PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
469 PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
470 if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
471 PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0));
472 PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0));
473 PetscOptionsEnd();
474 debug = swarm->printCoords;
475
476 PetscCall(DMGetDimension(sw, &dim));
477 PetscCall(DMSwarmGetCellDM(sw, &xdm));
478 PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
479
480 PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
481 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
482 PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
483
484 // One particle per centroid on the tensor product grid
485 Npc = (vcEnd - vcStart) * Ns;
486 Np = (xcEnd - xcStart) * Npc;
487 PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
488 if (debug) {
489 PetscInt gNp, gNc, Nc = xcEnd - xcStart;
490 PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm));
491 PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp));
492 PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm));
493 PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc));
494 PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart));
495 }
496
497 // Set species and cellid
498 PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
499 PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname));
500 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
501 PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid));
502 for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) {
503 for (PetscInt s = 0; s < Ns; ++s) {
504 for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) {
505 species[p] = s;
506 cellid[p] = c;
507 }
508 }
509 }
510 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
511 PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid));
512
513 // Set particle coordinates
514 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
515 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
516 PetscCall(DMSwarmSortGetAccess(sw));
517 PetscCall(DMGetBoundingBox(vdm, vmin, vmax));
518 PetscCall(DMGetCoordinatesLocalSetUp(xdm));
519 PetscCall(DMGetCoordinatesLocalSetUp(vdm));
520 for (PetscInt c = 0; c < xcEnd - xcStart; ++c) {
521 const PetscInt cell = c + xcStart;
522 PetscInt *pidx, Npc;
523 PetscReal centroid[3], volume;
524
525 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
526 PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL));
527 for (PetscInt s = 0; s < Ns; ++s) {
528 for (PetscInt q = 0; q < Npc / Ns; ++q) {
529 const PetscInt p = pidx[q * Ns + s];
530 const PetscInt vc = vcStart + q;
531 PetscReal vcentroid[3], vvolume;
532
533 PetscCall(DMPlexComputeCellGeometryFVM(vdm, vc, &vvolume, vcentroid, NULL));
534 for (PetscInt d = 0; d < dim; ++d) {
535 x[p * dim + d] = centroid[d];
536 v[p * dim + d] = vcentroid[d];
537 }
538
539 if (debug > 1) {
540 PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p));
541 PetscCall(PetscPrintf(PETSC_COMM_SELF, " x: ("));
542 for (PetscInt d = 0; d < dim; ++d) {
543 if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
544 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d]));
545 }
546 PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:("));
547 for (PetscInt d = 0; d < dim; ++d) {
548 if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", "));
549 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d]));
550 }
551 PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n"));
552 }
553 }
554 }
555 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
556 }
557 PetscCall(DMSwarmSortRestoreAccess(sw));
558 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
559 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
560 PetscFunctionReturn(PETSC_SUCCESS);
561 }
562
563 /*
564 InitializeWeights - Compute weight for each local particle
565
566 Input Parameters:
567 + sw - The `DMSwarm`
568 . totalWeight - The sum of all particle weights
569 . func - The PDF for the particle spatial distribution
570 - param - The PDF parameters
571
572 Notes:
573 The PDF for velocity is assumed to be a Gaussian
574
575 The particle weights are returned in the `w_q` field of `sw`.
576 */
InitializeWeights(DM sw,PetscReal totalWeight,PetscProbFn * func,const PetscReal param[])577 static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[])
578 {
579 DM xdm, vdm;
580 DMSwarmCellDM celldm;
581 PetscScalar *weight;
582 PetscQuadrature xquad;
583 const PetscReal *xq, *xwq;
584 const PetscInt order = 5;
585 PetscReal xi0[3];
586 PetscReal xwtot = 0., pwtot = 0.;
587 PetscInt xNq;
588 PetscInt dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights;
589 MPI_Comm comm;
590 PetscMPIInt rank;
591
592 PetscFunctionBegin;
593 PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
594 PetscCallMPI(MPI_Comm_rank(comm, &rank));
595 PetscCall(DMGetDimension(sw, &dim));
596 PetscCall(DMSwarmGetCellDM(sw, &xdm));
597 PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
598 PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd));
599 PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm));
600 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm));
601 PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd));
602
603 // Setup Quadrature for spatial and velocity weight calculations
604 PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad));
605 PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq));
606 for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0;
607
608 // Integrate the density function to get the weights of particles in each cell
609 PetscCall(DMGetCoordinatesLocalSetUp(vdm));
610 PetscCall(DMSwarmSortGetAccess(sw));
611 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
612 for (PetscInt c = xcStart; c < xcEnd; ++c) {
613 PetscReal xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.;
614 PetscInt *pidx, Npc;
615 PetscInt xNc;
616 const PetscScalar *xarray;
617 PetscScalar *xcoords = NULL;
618 PetscBool xisDG;
619
620 PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
621 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
622 PetscCheck(Npc == (vcEnd - vcStart) * Ns, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Number of particles %" PetscInt_FMT " in cell (rank %d) != %" PetscInt_FMT " number of velocity vertices", Npc, rank, (vcEnd - vcStart) * Ns);
623 PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ));
624 for (PetscInt q = 0; q < xNq; ++q) {
625 // Transform quadrature points from ref space to real space
626 CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr);
627 // Get probability density at quad point
628 // No need to scale xqr since PDF will be periodic
629 PetscCall((*func)(xqr, param, &xden));
630 xw += xden * (xwq[q] * xdetJ);
631 }
632 xwtot += xw;
633 if (debug) {
634 IS globalOrdering;
635 const PetscInt *ordering;
636
637 PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering));
638 PetscCall(ISGetIndices(globalOrdering, &ordering));
639 PetscCall(PetscSynchronizedPrintf(comm, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(xcoords[0]), (double)PetscRealPart(xcoords[0 + dim]), (double)xw));
640 PetscCall(ISRestoreIndices(globalOrdering, &ordering));
641 }
642 // Set weights to be Gaussian in velocity cells
643 for (PetscInt vc = vcStart; vc < vcEnd; ++vc) {
644 const PetscInt p = pidx[vc * Ns + 0];
645 PetscReal vw = 0.;
646 PetscInt vNc;
647 const PetscScalar *varray;
648 PetscScalar *vcoords = NULL;
649 PetscBool visDG;
650
651 PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
652 PetscCheck(vNc > 0 && vNc % dim == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Velocity cell %" PetscInt_FMT " has invalid coordinates (vNc=%" PetscInt_FMT ", dim=%" PetscInt_FMT ")", vc, vNc, dim);
653 {
654 PetscReal vmin[3] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
655 PetscReal vmax[3] = {-PETSC_MAX_REAL, -PETSC_MAX_REAL, -PETSC_MAX_REAL};
656 PetscInt numVert = vNc / dim;
657 for (PetscInt i = 0; i < numVert; ++i) {
658 for (PetscInt d = 0; d < dim; ++d) {
659 PetscReal coord = PetscRealPart(vcoords[i * dim + d]);
660 vmin[d] = PetscMin(vmin[d], coord);
661 vmax[d] = PetscMax(vmax[d], coord);
662 }
663 }
664 vw = 1.0;
665 for (PetscInt d = 0; d < dim; ++d) vw *= 0.5 * (PetscErfReal(vmax[d] / PetscSqrtReal(2.)) - PetscErfReal(vmin[d] / PetscSqrtReal(2.)));
666 PetscCheck(PetscIsNormalReal(vw), PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " velocity weight is not normal: vw=%g, vmin=(%g,%g,%g), vmax=(%g,%g,%g)", p, vw, vmin[0], vmin[1], vmin[2], vmax[0], vmax[1], vmax[2]);
667 }
668
669 weight[p] = totalWeight * vw * xw;
670 pwtot += weight[p];
671 PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 10: weight=%g, xw=%g, vw=%g, totalWeight=%g", p, weight[p], xw, vw, totalWeight);
672 PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords));
673 if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw));
674 }
675 PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords));
676 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
677 }
678 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
679 PetscCall(DMSwarmSortRestoreAccess(sw));
680 PetscCall(PetscQuadratureDestroy(&xquad));
681
682 if (debug) {
683 PetscReal wtot[2] = {pwtot, xwtot}, gwtot[2];
684
685 PetscCall(PetscSynchronizedFlush(comm, NULL));
686 PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
687 PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1]));
688 }
689 PetscFunctionReturn(PETSC_SUCCESS);
690 }
691
InitializeParticles_PerturbedWeights(DM sw,AppCtx * user)692 static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user)
693 {
694 PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]};
695 PetscInt dim;
696
697 PetscFunctionBegin;
698 PetscCall(DMGetDimension(sw, &dim));
699 PetscCall(InitializeParticles_Centroid(sw));
700 PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : (dim == 2 ? PetscPDFCosine2D : PetscPDFCosine3D), scale));
701 PetscFunctionReturn(PETSC_SUCCESS);
702 }
703
InitializeConstants(DM sw,AppCtx * user)704 static PetscErrorCode InitializeConstants(DM sw, AppCtx *user)
705 {
706 DM dm;
707 PetscInt *species;
708 PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight;
709 PetscInt Np, dim;
710
711 PetscFunctionBegin;
712 PetscCall(DMSwarmGetCellDM(sw, &dm));
713 PetscCall(DMGetDimension(sw, &dim));
714 PetscCall(DMSwarmGetLocalSize(sw, &Np));
715 PetscCall(DMGetBoundingBox(dm, gmin, gmax));
716 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight));
717 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
718 for (PetscInt p = 0; p < Np; ++p) {
719 totalWeight += weight[p];
720 totalCharge += user->charges[species[p]] * weight[p];
721 }
722 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
723 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
724 {
725 Parameter *param;
726 PetscReal Area;
727
728 PetscCall(PetscBagGetData(user->bag, (void **)¶m));
729 switch (dim) {
730 case 1:
731 Area = (gmax[0] - gmin[0]);
732 break;
733 case 2:
734 Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]);
735 break;
736 case 3:
737 Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]);
738 break;
739 default:
740 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim);
741 }
742 PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
743 PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD));
744 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)user->charges[0], (double)global_charge, (double)Area));
745 param->sigma = PetscAbsReal(global_charge / (Area));
746
747 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma));
748 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "(x0,v0,t0,m0,q0,phi0): (%e, %e, %e, %e, %e, %e) - (P, V) = (%e, %e)\n", (double)param->x0, (double)param->v0, (double)param->t0, (double)param->m0, (double)param->q0, (double)param->phi0, (double)param->poissonNumber,
749 (double)param->vlasovNumber));
750 }
751 /* Setup Constants */
752 {
753 PetscDS ds;
754 Parameter *param;
755 PetscCall(PetscBagGetData(user->bag, (void **)¶m));
756 PetscScalar constants[NUM_CONSTANTS];
757 constants[SIGMA] = param->sigma;
758 constants[V0] = param->v0;
759 constants[T0] = param->t0;
760 constants[X0] = param->x0;
761 constants[M0] = param->m0;
762 constants[Q0] = param->q0;
763 constants[PHI0] = param->phi0;
764 constants[POISSON] = param->poissonNumber;
765 constants[VLASOV] = param->vlasovNumber;
766 PetscCall(DMGetDS(dm, &ds));
767 PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants));
768 }
769 PetscFunctionReturn(PETSC_SUCCESS);
770 }
771
SetupParameters(MPI_Comm comm,AppCtx * ctx)772 static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx)
773 {
774 PetscBag bag;
775 Parameter *p;
776
777 PetscFunctionBeginUser;
778 /* setup PETSc parameter bag */
779 PetscCall(PetscBagGetData(ctx->bag, (void **)&p));
780 PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters"));
781 bag = ctx->bag;
782 PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s"));
783 PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s"));
784 PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m"));
785 PetscCall(PetscBagRegisterScalar(bag, &p->phi0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3"));
786 PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s"));
787 PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg"));
788 PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg"));
789 PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1"));
790
791 PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3"));
792 PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number"));
793 PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number"));
794 PetscCall(PetscBagSetFromOptions(bag));
795 {
796 PetscViewer viewer;
797 PetscViewerFormat format;
798 PetscBool flg;
799
800 PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg));
801 if (flg) {
802 PetscCall(PetscViewerPushFormat(viewer, format));
803 PetscCall(PetscBagView(bag, viewer));
804 PetscCall(PetscViewerFlush(viewer));
805 PetscCall(PetscViewerPopFormat(viewer));
806 PetscCall(PetscViewerDestroy(&viewer));
807 }
808 }
809 PetscFunctionReturn(PETSC_SUCCESS);
810 }
811
CreateMesh(MPI_Comm comm,AppCtx * user,DM * dm)812 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
813 {
814 const char *prefix = "x";
815
816 PetscFunctionBeginUser;
817 PetscCall(DMCreate(comm, dm));
818 PetscCall(DMSetType(*dm, DMPLEX));
819 PetscCall(DMPlexSetOptionsPrefix(*dm, prefix));
820 PetscCall(DMSetFromOptions(*dm));
821 PetscCall(PetscObjectSetName((PetscObject)*dm, "space"));
822 PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
823
824 // Cache the mesh geometry
825 DMField coordField;
826 IS cellIS;
827 PetscQuadrature quad;
828 PetscReal *wt, *pt;
829 PetscInt cdim, cStart, cEnd;
830
831 PetscCall(DMGetCoordinateField(*dm, &coordField));
832 PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field");
833 PetscCall(DMGetCoordinateDim(*dm, &cdim));
834 PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd));
835 PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS));
836 PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad));
837 PetscCall(PetscMalloc1(1, &wt));
838 PetscCall(PetscMalloc1(cdim, &pt));
839 wt[0] = 1.;
840 for (PetscInt d = 0; d < cdim; ++d) pt[d] = -1.;
841 PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt));
842 PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom));
843 PetscCall(PetscQuadratureDestroy(&quad));
844 PetscCall(ISDestroy(&cellIS));
845 PetscFunctionReturn(PETSC_SUCCESS);
846 }
847
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[])848 static void ion_f0(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
849 {
850 f0[0] = -constants[SIGMA];
851 }
852
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[])853 static void laplacian_f1(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[])
854 {
855 PetscInt d;
856 for (d = 0; d < dim; ++d) f1[d] = u_x[d];
857 }
858
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[])859 static void laplacian_g3(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g3[])
860 {
861 PetscInt d;
862 for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
863 }
864
CreateFEM(DM dm,AppCtx * user)865 static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
866 {
867 PetscFE fephi;
868 PetscDS ds;
869 PetscBool simplex;
870 PetscInt dim;
871 MatNullSpace nullsp;
872
873 PetscFunctionBeginUser;
874 PetscCall(DMGetDimension(dm, &dim));
875 PetscCall(DMPlexIsSimplex(dm, &simplex));
876
877 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi));
878 PetscCall(PetscObjectSetName((PetscObject)fephi, "potential"));
879 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi));
880 PetscCall(DMCreateDS(dm));
881 PetscCall(DMGetDS(dm, &ds));
882 PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1));
883 PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
884 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp));
885 PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp));
886 PetscCall(MatNullSpaceDestroy(&nullsp));
887 PetscCall(PetscFEDestroy(&fephi));
888 PetscFunctionReturn(PETSC_SUCCESS);
889 }
890
CreatePoisson(DM dm,AppCtx * user)891 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
892 {
893 SNES snes;
894 Mat J;
895 MatNullSpace nullSpace;
896
897 PetscFunctionBeginUser;
898 PetscCall(CreateFEM(dm, user));
899 PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
900 PetscCall(SNESSetOptionsPrefix(snes, "em_"));
901 PetscCall(SNESSetDM(snes, dm));
902 PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user));
903 PetscCall(SNESSetFromOptions(snes));
904
905 PetscCall(DMCreateMatrix(dm, &J));
906 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
907 PetscCall(MatSetNullSpace(J, nullSpace));
908 PetscCall(MatNullSpaceDestroy(&nullSpace));
909 PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
910 PetscCall(MatDestroy(&J));
911
912 user->dmPot = dm;
913 PetscCall(PetscObjectReference((PetscObject)user->dmPot));
914
915 PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M));
916 PetscCall(DMPlexCreateClosureIndex(dm, NULL));
917 user->snes = snes;
918 PetscFunctionReturn(PETSC_SUCCESS);
919 }
920
CreateSwarm(DM dm,AppCtx * user,DM * sw)921 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
922 {
923 DMSwarmCellDM celldm;
924 PetscInt dim;
925
926 PetscFunctionBeginUser;
927 PetscCall(DMGetDimension(dm, &dim));
928 PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
929 PetscCall(DMSetType(*sw, DMSWARM));
930 PetscCall(DMSetDimension(*sw, dim));
931 PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
932 PetscCall(DMSetApplicationContext(*sw, user));
933
934 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
935 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
936 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
937 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
938
939 const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
940
941 PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm));
942 PetscCall(DMSwarmAddCellDM(*sw, celldm));
943 PetscCall(DMSwarmCellDMDestroy(&celldm));
944
945 const char *vfieldnames[2] = {"w_q"};
946 DM vdm;
947
948 PetscCall(CreateVelocityDM(*sw, &vdm));
949 PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm));
950 PetscCall(DMSwarmAddCellDM(*sw, celldm));
951 PetscCall(DMSwarmCellDMDestroy(&celldm));
952 PetscCall(DMDestroy(&vdm));
953
954 DM mdm;
955
956 PetscCall(DMClone(dm, &mdm));
957 PetscCall(PetscObjectSetName((PetscObject)mdm, "moments"));
958 PetscCall(DMCopyDisc(dm, mdm));
959 PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm));
960 PetscCall(DMDestroy(&mdm));
961 PetscCall(DMSwarmAddCellDM(*sw, celldm));
962 PetscCall(DMSwarmCellDMDestroy(&celldm));
963
964 PetscCall(DMSetFromOptions(*sw));
965 PetscCall(DMSetUp(*sw));
966
967 PetscCall(DMSwarmSetCellDMActive(*sw, "space"));
968 user->swarm = *sw;
969 // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set
970 PetscCall(InitializeParticles_PerturbedWeights(*sw, user));
971 PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
972 PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
973 PetscFunctionReturn(PETSC_SUCCESS);
974 }
975
ComputeFieldAtParticles_Primal(SNES snes,DM sw,Mat M_p,PetscReal E[])976 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[])
977 {
978 DM dm;
979 AppCtx *user;
980 PetscDS ds;
981 PetscFE fe;
982 KSP ksp;
983 Vec rhoRhs; // Weak charge density, \int phi_i rho
984 Vec rho; // Charge density, M^{-1} rhoRhs
985 Vec phi, locPhi; // Potential
986 Vec f; // Particle weights
987 PetscReal *coords;
988 PetscInt dim, cStart, cEnd, Np;
989
990 PetscFunctionBegin;
991 PetscCall(DMGetApplicationContext(sw, (void *)&user));
992 PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0));
993 PetscCall(DMGetDimension(sw, &dim));
994 PetscCall(DMSwarmGetLocalSize(sw, &Np));
995
996 PetscCall(SNESGetDM(snes, &dm));
997 PetscCall(DMGetGlobalVector(dm, &rhoRhs));
998 PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density"));
999 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho));
1000 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
1001 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
1002
1003 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
1004 PetscCall(MatViewFromOptions(user->M, NULL, "-m_view"));
1005 PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
1006
1007 PetscCall(MatMultTranspose(M_p, f, rhoRhs));
1008 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
1009
1010 // Low-pass filter rhoRhs
1011 PetscInt window = 0;
1012 PetscCall(PetscOptionsGetInt(NULL, NULL, "-rho_average", &window, NULL));
1013 if (window) {
1014 PetscScalar *a;
1015 PetscInt n;
1016 PetscReal width = 2. * window + 1.;
1017
1018 // This will only work for P_1
1019 // This works because integration against a test function is linear
1020 // Do a real integral against weight function for higher order
1021 PetscCall(VecGetLocalSize(rhoRhs, &n));
1022 PetscCall(VecGetArrayWrite(rhoRhs, &a));
1023 for (PetscInt i = 0; i < n; ++i) {
1024 PetscScalar avg = a[i];
1025 for (PetscInt j = 1; j <= window; ++j) avg += a[(i - j + n) % n] + a[(i + j) % n];
1026 a[i] = avg / width;
1027 //a[i] = (a[(i - 1 + n) % n] + a[i] + a[(i + 1) % n]) / 3.;
1028 }
1029 PetscCall(VecRestoreArrayWrite(rhoRhs, &a));
1030 }
1031
1032 PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
1033 PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_"));
1034 PetscCall(KSPSetOperators(ksp, user->M, user->M));
1035 PetscCall(KSPSetFromOptions(ksp));
1036 PetscCall(KSPSolve(ksp, rhoRhs, rho));
1037
1038 PetscCall(VecScale(rhoRhs, -1.0));
1039
1040 PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view"));
1041 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho));
1042 PetscCall(KSPDestroy(&ksp));
1043
1044 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1045 PetscCall(VecSet(phi, 0.0));
1046 PetscCall(SNESSolve(snes, rhoRhs, phi));
1047 PetscCall(DMRestoreGlobalVector(dm, &rhoRhs));
1048 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
1049
1050 PetscCall(DMGetLocalVector(dm, &locPhi));
1051 PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
1052 PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
1053 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
1054 PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0));
1055
1056 PetscCall(DMGetDS(dm, &ds));
1057 PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
1058 PetscCall(DMSwarmSortGetAccess(sw));
1059 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1060 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1061
1062 PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0));
1063 PetscTabulation tab;
1064 PetscReal *pcoord, *refcoord;
1065 PetscFEGeom *chunkgeom = NULL;
1066 PetscInt maxNcp = 0;
1067
1068 for (PetscInt c = cStart; c < cEnd; ++c) {
1069 PetscInt Ncp;
1070
1071 PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp));
1072 maxNcp = PetscMax(maxNcp, Ncp);
1073 }
1074 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1075 PetscCall(PetscArrayzero(refcoord, maxNcp * dim));
1076 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1077 PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab));
1078 for (PetscInt c = cStart; c < cEnd; ++c) {
1079 PetscScalar *clPhi = NULL;
1080 PetscInt *points;
1081 PetscInt Ncp;
1082
1083 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1084 for (PetscInt cp = 0; cp < Ncp; ++cp)
1085 for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
1086 {
1087 PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom));
1088 for (PetscInt i = 0; i < Ncp; ++i) {
1089 const PetscReal x0[3] = {-1., -1., -1.};
1090 CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]);
1091 }
1092 }
1093 PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab));
1094 PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1095 for (PetscInt cp = 0; cp < Ncp; ++cp) {
1096 const PetscReal *basisDer = tab->T[1];
1097 const PetscInt p = points[cp];
1098
1099 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.;
1100 PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim]));
1101 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
1102 }
1103 PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
1104 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1105 }
1106 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord));
1107 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord));
1108 PetscCall(PetscTabulationDestroy(&tab));
1109 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1110 PetscCall(DMSwarmSortRestoreAccess(sw));
1111 PetscCall(DMRestoreLocalVector(dm, &locPhi));
1112 PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom));
1113 PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0));
1114 PetscFunctionReturn(PETSC_SUCCESS);
1115 }
1116
ComputeFieldAtParticles(SNES snes,DM sw)1117 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw)
1118 {
1119 AppCtx *user;
1120 Mat M_p;
1121 PetscReal *E;
1122 PetscInt dim, Np;
1123
1124 PetscFunctionBegin;
1125 PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
1126 PetscValidHeaderSpecific(sw, DM_CLASSID, 2);
1127 PetscCall(DMGetDimension(sw, &dim));
1128 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1129 PetscCall(DMGetApplicationContext(sw, &user));
1130
1131 PetscCall(DMSwarmSetCellDMActive(sw, "moments"));
1132 // TODO: Could share sort context with space cellDM
1133 PetscCall(DMSwarmMigrate(sw, PETSC_FALSE));
1134 PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p));
1135 PetscCall(DMSwarmSetCellDMActive(sw, "space"));
1136
1137 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1138 PetscCall(PetscArrayzero(E, Np * dim));
1139 user->validE = PETSC_TRUE;
1140
1141 PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E));
1142 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1143 PetscCall(MatDestroy(&M_p));
1144 PetscFunctionReturn(PETSC_SUCCESS);
1145 }
1146
RHSFunction(TS ts,PetscReal t,Vec U,Vec G,void * ctx)1147 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1148 {
1149 DM sw;
1150 SNES snes = ((AppCtx *)ctx)->snes;
1151 const PetscScalar *u;
1152 PetscScalar *g;
1153 PetscReal *E, m_p = 1., q_p = -1.;
1154 PetscInt dim, d, Np, p;
1155
1156 PetscFunctionBeginUser;
1157 PetscCall(TSGetDM(ts, &sw));
1158 PetscCall(ComputeFieldAtParticles(snes, sw));
1159
1160 PetscCall(DMGetDimension(sw, &dim));
1161 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1162 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1163 PetscCall(VecGetArrayRead(U, &u));
1164 PetscCall(VecGetArray(G, &g));
1165 Np /= 2 * dim;
1166 for (p = 0; p < Np; ++p) {
1167 for (d = 0; d < dim; ++d) {
1168 g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
1169 g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p;
1170 }
1171 }
1172 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1173 PetscCall(VecRestoreArrayRead(U, &u));
1174 PetscCall(VecRestoreArray(G, &g));
1175 PetscFunctionReturn(PETSC_SUCCESS);
1176 }
1177
1178 /* J_{ij} = dF_i/dx_j
1179 J_p = ( 0 1)
1180 (-w^2 0)
1181 TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
1182 Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
1183 */
RHSJacobian(TS ts,PetscReal t,Vec U,Mat J,Mat P,void * ctx)1184 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
1185 {
1186 DM sw;
1187 const PetscReal *coords, *vel;
1188 PetscInt dim, d, Np, p, rStart;
1189
1190 PetscFunctionBeginUser;
1191 PetscCall(TSGetDM(ts, &sw));
1192 PetscCall(DMGetDimension(sw, &dim));
1193 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1194 PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
1195 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1196 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1197 Np /= 2 * dim;
1198 for (p = 0; p < Np; ++p) {
1199 // TODO This is not right because dv/dx has the electric field in it
1200 PetscScalar vals[4] = {0., 1., -1, 0.};
1201
1202 for (d = 0; d < dim; ++d) {
1203 const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1204 PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
1205 }
1206 }
1207 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1208 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1209 PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
1210 PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
1211 PetscFunctionReturn(PETSC_SUCCESS);
1212 }
1213
RHSFunctionX(TS ts,PetscReal t,Vec V,Vec Xres,void * ctx)1214 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
1215 {
1216 AppCtx *user = (AppCtx *)ctx;
1217 DM sw;
1218 const PetscScalar *v;
1219 PetscScalar *xres;
1220 PetscInt Np, p, d, dim;
1221
1222 PetscFunctionBeginUser;
1223 PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0));
1224 PetscCall(TSGetDM(ts, &sw));
1225 PetscCall(DMGetDimension(sw, &dim));
1226 PetscCall(VecGetLocalSize(Xres, &Np));
1227 PetscCall(VecGetArrayRead(V, &v));
1228 PetscCall(VecGetArray(Xres, &xres));
1229 Np /= dim;
1230 for (p = 0; p < Np; ++p) {
1231 for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
1232 }
1233 PetscCall(VecRestoreArrayRead(V, &v));
1234 PetscCall(VecRestoreArray(Xres, &xres));
1235 PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0));
1236 PetscFunctionReturn(PETSC_SUCCESS);
1237 }
1238
RHSFunctionV(TS ts,PetscReal t,Vec X,Vec Vres,void * ctx)1239 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
1240 {
1241 DM sw;
1242 AppCtx *user = (AppCtx *)ctx;
1243 SNES snes = ((AppCtx *)ctx)->snes;
1244 const PetscScalar *x;
1245 PetscScalar *vres;
1246 PetscReal *E, m_p, q_p;
1247 PetscInt Np, p, dim, d;
1248 Parameter *param;
1249
1250 PetscFunctionBeginUser;
1251 PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0));
1252 PetscCall(TSGetDM(ts, &sw));
1253 PetscCall(ComputeFieldAtParticles(snes, sw));
1254
1255 PetscCall(DMGetDimension(sw, &dim));
1256 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1257 PetscCall(PetscBagGetData(user->bag, (void **)¶m));
1258 m_p = user->masses[0] * param->m0;
1259 q_p = user->charges[0] * param->q0;
1260 PetscCall(VecGetLocalSize(Vres, &Np));
1261 PetscCall(VecGetArrayRead(X, &x));
1262 PetscCall(VecGetArray(Vres, &vres));
1263 Np /= dim;
1264 for (p = 0; p < Np; ++p) {
1265 for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p;
1266 }
1267 PetscCall(VecRestoreArrayRead(X, &x));
1268 /*
1269 Synchronized, ordered output for parallel/sequential test cases.
1270 In the 1D (on the 2D mesh) case, every y component should be zero.
1271 */
1272 if (user->checkVRes) {
1273 PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE;
1274 PetscInt step;
1275
1276 PetscCall(TSGetStepNumber(ts, &step));
1277 if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step));
1278 for (PetscInt p = 0; p < Np; ++p) {
1279 if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1])));
1280 PetscCheck(PetscAbsScalar(vres[p * dim + 1]) < PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Y velocity should be 0., not %g", (double)PetscRealPart(vres[p * dim + 1]));
1281 }
1282 if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT));
1283 }
1284 PetscCall(VecRestoreArray(Vres, &vres));
1285 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1286 PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0));
1287 PetscFunctionReturn(PETSC_SUCCESS);
1288 }
1289
1290 /* Discrete Gradients Formulation: S, F, gradF (G) */
RHSJacobianS(TS ts,PetscReal t,Vec U,Mat S,void * ctx)1291 PetscErrorCode RHSJacobianS(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
1292 {
1293 PetscScalar vals[4] = {0., 1., -1., 0.};
1294 DM sw;
1295 PetscInt dim, d, Np, p, rStart;
1296
1297 PetscFunctionBeginUser;
1298 PetscCall(TSGetDM(ts, &sw));
1299 PetscCall(DMGetDimension(sw, &dim));
1300 PetscCall(VecGetLocalSize(U, &Np));
1301 PetscCall(MatGetOwnershipRange(S, &rStart, NULL));
1302 Np /= 2 * dim;
1303 for (p = 0; p < Np; ++p) {
1304 for (d = 0; d < dim; ++d) {
1305 const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
1306 PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
1307 }
1308 }
1309 PetscCall(MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY));
1310 PetscCall(MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY));
1311 PetscFunctionReturn(PETSC_SUCCESS);
1312 }
1313
RHSObjectiveF(TS ts,PetscReal t,Vec U,PetscScalar * F,void * ctx)1314 PetscErrorCode RHSObjectiveF(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
1315 {
1316 AppCtx *user = (AppCtx *)ctx;
1317 DM sw;
1318 Vec phi;
1319 const PetscScalar *u;
1320 PetscInt dim, Np, cStart, cEnd;
1321 PetscReal *vel, *coords, m_p = 1.;
1322
1323 PetscFunctionBeginUser;
1324 PetscCall(TSGetDM(ts, &sw));
1325 PetscCall(DMGetDimension(sw, &dim));
1326 PetscCall(DMPlexGetHeightStratum(user->dmPot, 0, &cStart, &cEnd));
1327
1328 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi));
1329 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view_dg"));
1330 PetscCall(computeFieldEnergy(user->dmPot, phi, F));
1331 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi));
1332
1333 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1334 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1335 PetscCall(DMSwarmSortGetAccess(sw));
1336 PetscCall(VecGetArrayRead(U, &u));
1337 PetscCall(VecGetLocalSize(U, &Np));
1338 Np /= 2 * dim;
1339 for (PetscInt c = cStart; c < cEnd; ++c) {
1340 PetscInt *points;
1341 PetscInt Ncp;
1342
1343 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
1344 for (PetscInt cp = 0; cp < Ncp; ++cp) {
1345 const PetscInt p = points[cp];
1346 const PetscReal v2 = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);
1347
1348 *F += 0.5 * m_p * v2;
1349 }
1350 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points));
1351 }
1352 PetscCall(VecRestoreArrayRead(U, &u));
1353 PetscCall(DMSwarmSortRestoreAccess(sw));
1354 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1355 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1356 PetscFunctionReturn(PETSC_SUCCESS);
1357 }
1358
1359 /* dF/dx = q E dF/dv = v */
RHSFunctionG(TS ts,PetscReal t,Vec U,Vec G,void * ctx)1360 PetscErrorCode RHSFunctionG(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
1361 {
1362 DM sw;
1363 SNES snes = ((AppCtx *)ctx)->snes;
1364 const PetscReal *coords, *vel, *E;
1365 const PetscScalar *u;
1366 PetscScalar *g;
1367 PetscReal m_p = 1., q_p = -1.;
1368 PetscInt dim, d, Np, p;
1369
1370 PetscFunctionBeginUser;
1371 PetscCall(TSGetDM(ts, &sw));
1372 PetscCall(DMGetDimension(sw, &dim));
1373 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1374 PetscCall(VecGetArrayRead(U, &u));
1375 PetscCall(VecGetArray(G, &g));
1376
1377 PetscLogEvent COMPUTEFIELD;
1378 PetscCall(PetscLogEventRegister("COMPFIELDATPART", TS_CLASSID, &COMPUTEFIELD));
1379 PetscCall(PetscLogEventBegin(COMPUTEFIELD, 0, 0, 0, 0));
1380 PetscCall(ComputeFieldAtParticles(snes, sw));
1381 PetscCall(PetscLogEventEnd(COMPUTEFIELD, 0, 0, 0, 0));
1382 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1383 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel));
1384 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
1385 for (p = 0; p < Np; ++p) {
1386 for (d = 0; d < dim; ++d) {
1387 g[(p * 2 + 0) * dim + d] = -(q_p / m_p) * E[p * dim + d];
1388 g[(p * 2 + 1) * dim + d] = u[(p * 2 + 1) * dim + d];
1389 }
1390 }
1391 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
1392 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
1393 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel));
1394 PetscCall(VecRestoreArrayRead(U, &u));
1395 PetscCall(VecRestoreArray(G, &g));
1396 PetscFunctionReturn(PETSC_SUCCESS);
1397 }
1398
CreateSolution(TS ts)1399 static PetscErrorCode CreateSolution(TS ts)
1400 {
1401 DM sw;
1402 Vec u;
1403 PetscInt dim, Np;
1404
1405 PetscFunctionBegin;
1406 PetscCall(TSGetDM(ts, &sw));
1407 PetscCall(DMGetDimension(sw, &dim));
1408 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1409 PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
1410 PetscCall(VecSetBlockSize(u, dim));
1411 PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
1412 PetscCall(VecSetUp(u));
1413 PetscCall(TSSetSolution(ts, u));
1414 PetscCall(VecDestroy(&u));
1415 PetscFunctionReturn(PETSC_SUCCESS);
1416 }
1417
SetProblem(TS ts)1418 static PetscErrorCode SetProblem(TS ts)
1419 {
1420 AppCtx *user;
1421 DM sw;
1422
1423 PetscFunctionBegin;
1424 PetscCall(TSGetDM(ts, &sw));
1425 PetscCall(DMGetApplicationContext(sw, (void **)&user));
1426 // Define unified system for (X, V)
1427 {
1428 Mat J;
1429 PetscInt dim, Np;
1430
1431 PetscCall(DMGetDimension(sw, &dim));
1432 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1433 PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
1434 PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
1435 PetscCall(MatSetBlockSize(J, 2 * dim));
1436 PetscCall(MatSetFromOptions(J));
1437 PetscCall(MatSetUp(J));
1438 PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
1439 PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
1440 PetscCall(MatDestroy(&J));
1441 }
1442 /* Define split system for X and V */
1443 {
1444 Vec u;
1445 IS isx, isv, istmp;
1446 const PetscInt *idx;
1447 PetscInt dim, Np, rstart;
1448
1449 PetscCall(TSGetSolution(ts, &u));
1450 PetscCall(DMGetDimension(sw, &dim));
1451 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1452 PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
1453 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
1454 PetscCall(ISGetIndices(istmp, &idx));
1455 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
1456 PetscCall(ISRestoreIndices(istmp, &idx));
1457 PetscCall(ISDestroy(&istmp));
1458 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
1459 PetscCall(ISGetIndices(istmp, &idx));
1460 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
1461 PetscCall(ISRestoreIndices(istmp, &idx));
1462 PetscCall(ISDestroy(&istmp));
1463 PetscCall(TSRHSSplitSetIS(ts, "position", isx));
1464 PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
1465 PetscCall(ISDestroy(&isx));
1466 PetscCall(ISDestroy(&isv));
1467 PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
1468 PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
1469 }
1470 // Define symplectic formulation U_t = S . G, where G = grad F
1471 {
1472 PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, user));
1473 }
1474 PetscFunctionReturn(PETSC_SUCCESS);
1475 }
1476
DMSwarmTSRedistribute(TS ts)1477 static PetscErrorCode DMSwarmTSRedistribute(TS ts)
1478 {
1479 DM sw;
1480 Vec u;
1481 PetscReal t, maxt, dt;
1482 PetscInt n, maxn;
1483
1484 PetscFunctionBegin;
1485 PetscCall(TSGetDM(ts, &sw));
1486 PetscCall(TSGetTime(ts, &t));
1487 PetscCall(TSGetMaxTime(ts, &maxt));
1488 PetscCall(TSGetTimeStep(ts, &dt));
1489 PetscCall(TSGetStepNumber(ts, &n));
1490 PetscCall(TSGetMaxSteps(ts, &maxn));
1491
1492 PetscCall(TSReset(ts));
1493 PetscCall(TSSetDM(ts, sw));
1494 PetscCall(TSSetFromOptions(ts));
1495 PetscCall(TSSetTime(ts, t));
1496 PetscCall(TSSetMaxTime(ts, maxt));
1497 PetscCall(TSSetTimeStep(ts, dt));
1498 PetscCall(TSSetStepNumber(ts, n));
1499 PetscCall(TSSetMaxSteps(ts, maxn));
1500
1501 PetscCall(CreateSolution(ts));
1502 PetscCall(SetProblem(ts));
1503 PetscCall(TSGetSolution(ts, &u));
1504 PetscFunctionReturn(PETSC_SUCCESS);
1505 }
1506
line(PetscInt dim,PetscReal time,const PetscReal dummy[],PetscInt p,PetscScalar x[],void * ctx)1507 PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
1508 {
1509 DM sw, cdm;
1510 PetscInt Np;
1511 PetscReal low[2], high[2];
1512 AppCtx *user = (AppCtx *)ctx;
1513
1514 sw = user->swarm;
1515 PetscCall(DMSwarmGetCellDM(sw, &cdm));
1516 // Get the bounding box so we can equally space the particles
1517 PetscCall(DMGetLocalBoundingBox(cdm, low, high));
1518 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1519 // shift it by h/2 so nothing is initialized directly on a boundary
1520 x[0] = ((high[0] - low[0]) / Np) * (p + 0.5);
1521 x[1] = 0.;
1522 return PETSC_SUCCESS;
1523 }
1524
1525 /*
1526 InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.
1527
1528 Input Parameters:
1529 + ts - The TS
1530 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem
1531
1532 Output Parameters:
1533 . u - The initialized solution vector
1534
1535 Level: advanced
1536
1537 .seealso: InitializeSolve()
1538 */
InitializeSolveAndSwarm(TS ts,PetscBool useInitial)1539 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
1540 {
1541 DM sw;
1542 Vec u, gc, gv;
1543 IS isx, isv;
1544 PetscInt dim;
1545 AppCtx *user;
1546
1547 PetscFunctionBeginUser;
1548 PetscCall(TSGetDM(ts, &sw));
1549 PetscCall(DMGetApplicationContext(sw, &user));
1550 PetscCall(DMGetDimension(sw, &dim));
1551 if (useInitial) {
1552 PetscCall(InitializeParticles_PerturbedWeights(sw, user));
1553 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1554 PetscCall(DMSwarmTSRedistribute(ts));
1555 }
1556 PetscCall(DMSetUp(sw));
1557 PetscCall(TSGetSolution(ts, &u));
1558 PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1559 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1560 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1561 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1562 PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
1563 PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
1564 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1565 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1566 PetscFunctionReturn(PETSC_SUCCESS);
1567 }
1568
InitializeSolve(TS ts,Vec u)1569 static PetscErrorCode InitializeSolve(TS ts, Vec u)
1570 {
1571 PetscFunctionBegin;
1572 PetscCall(TSSetSolution(ts, u));
1573 PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
1574 PetscFunctionReturn(PETSC_SUCCESS);
1575 }
1576
MigrateParticles(TS ts)1577 static PetscErrorCode MigrateParticles(TS ts)
1578 {
1579 DM sw, cdm;
1580 const PetscReal *L;
1581 AppCtx *ctx;
1582
1583 PetscFunctionBeginUser;
1584 PetscCall(TSGetDM(ts, &sw));
1585 PetscCall(DMGetApplicationContext(sw, &ctx));
1586 PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
1587 {
1588 Vec u, gc, gv, position, momentum;
1589 IS isx, isv;
1590 PetscReal *pos, *mom;
1591
1592 PetscCall(TSGetSolution(ts, &u));
1593 PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
1594 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
1595 PetscCall(VecGetSubVector(u, isx, &position));
1596 PetscCall(VecGetSubVector(u, isv, &momentum));
1597 PetscCall(VecGetArray(position, &pos));
1598 PetscCall(VecGetArray(momentum, &mom));
1599 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1600 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
1601 PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
1602 PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
1603
1604 PetscCall(DMSwarmGetCellDM(sw, &cdm));
1605 PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L));
1606 if ((L[0] || L[1]) >= 0.) {
1607 PetscReal *x, *v, upper[3], lower[3];
1608 PetscInt Np, dim;
1609
1610 PetscCall(DMSwarmGetLocalSize(sw, &Np));
1611 PetscCall(DMGetDimension(cdm, &dim));
1612 PetscCall(DMGetBoundingBox(cdm, lower, upper));
1613 PetscCall(VecGetArray(gc, &x));
1614 PetscCall(VecGetArray(gv, &v));
1615 for (PetscInt p = 0; p < Np; ++p) {
1616 for (PetscInt d = 0; d < dim; ++d) {
1617 if (pos[p * dim + d] < lower[d]) {
1618 x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]);
1619 } else if (pos[p * dim + d] > upper[d]) {
1620 x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]);
1621 } else {
1622 x[p * dim + d] = pos[p * dim + d];
1623 }
1624 PetscCheck(x[p * dim + d] >= lower[d] && x[p * dim + d] <= upper[d], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "p: %" PetscInt_FMT "x[%" PetscInt_FMT "] %g", p, d, (double)x[p * dim + d]);
1625 v[p * dim + d] = mom[p * dim + d];
1626 }
1627 }
1628 PetscCall(VecRestoreArray(gc, &x));
1629 PetscCall(VecRestoreArray(gv, &v));
1630 }
1631 PetscCall(VecRestoreArray(position, &pos));
1632 PetscCall(VecRestoreArray(momentum, &mom));
1633 PetscCall(VecRestoreSubVector(u, isx, &position));
1634 PetscCall(VecRestoreSubVector(u, isv, &momentum));
1635 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
1636 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
1637 }
1638 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
1639 // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm
1640 PetscCall(DMSwarmTSRedistribute(ts));
1641 PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
1642 {
1643 const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"};
1644 PetscCall(DMSwarmVectorDefineFields(sw, 2, fieldnames));
1645 }
1646 PetscFunctionReturn(PETSC_SUCCESS);
1647 }
1648
main(int argc,char ** argv)1649 int main(int argc, char **argv)
1650 {
1651 DM dm, sw;
1652 TS ts;
1653 Vec u;
1654 PetscReal dt;
1655 PetscInt maxn;
1656 AppCtx user;
1657
1658 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
1659 PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
1660 PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag));
1661 PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
1662 PetscCall(CreatePoisson(dm, &user));
1663 PetscCall(CreateSwarm(dm, &user, &sw));
1664 PetscCall(SetupParameters(PETSC_COMM_WORLD, &user));
1665 PetscCall(InitializeConstants(sw, &user));
1666 PetscCall(DMSetApplicationContext(sw, &user));
1667
1668 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
1669 PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
1670 PetscCall(TSSetDM(ts, sw));
1671 PetscCall(TSSetMaxTime(ts, 0.1));
1672 PetscCall(TSSetTimeStep(ts, 0.00001));
1673 PetscCall(TSSetMaxSteps(ts, 100));
1674 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
1675
1676 if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL));
1677 if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL));
1678 if (user.particle_monitor) PetscCall(TSMonitorSet(ts, MonitorParticles, &user, NULL));
1679
1680 PetscCall(TSSetFromOptions(ts));
1681 PetscCall(TSGetTimeStep(ts, &dt));
1682 PetscCall(TSGetMaxSteps(ts, &maxn));
1683 user.steps = maxn;
1684 user.stepSize = dt;
1685 PetscCall(SetupContext(dm, sw, &user));
1686 PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
1687 PetscCall(TSSetPostStep(ts, MigrateParticles));
1688 PetscCall(CreateSolution(ts));
1689 PetscCall(TSGetSolution(ts, &u));
1690 PetscCall(TSComputeInitialCondition(ts, u));
1691 PetscCall(CheckNonNegativeWeights(sw, &user));
1692 PetscCall(TSSolve(ts, NULL));
1693
1694 PetscCall(SNESDestroy(&user.snes));
1695 PetscCall(DMDestroy(&user.dmPot));
1696 PetscCall(MatDestroy(&user.M));
1697 PetscCall(PetscFEGeomDestroy(&user.fegeom));
1698 PetscCall(TSDestroy(&ts));
1699 PetscCall(DMDestroy(&sw));
1700 PetscCall(DMDestroy(&dm));
1701 PetscCall(DestroyContext(&user));
1702 PetscCall(PetscFinalize());
1703 return 0;
1704 }
1705
1706 /*TEST
1707
1708 build:
1709 requires: !complex double
1710
1711 testset:
1712 nsize: 2
1713 args: -cosine_coefficients 0.01 -charges -1. -total_weight 1. -xdm_plex_hash_location -vpetscspace_degree 2 -petscspace_degree 1 -em_snes_atol 1.e-12 \
1714 -em_snes_error_if_not_converged -em_ksp_error_if_not_converged -em_ksp_type cg -em_pc_type gamg -em_mg_coarse_ksp_type preonly -em_mg_coarse_pc_type svd -em_proj_ksp_type cg \
1715 -em_proj_pc_type gamg -em_proj_mg_coarse_ksp_type preonly -em_proj_mg_coarse_pc_type svd -ts_time_step 0.03 -xdm_plex_simplex 0 \
1716 -ts_max_time 100 -output_step 1 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -check_weights -ts_max_steps 60
1717
1718 test:
1719 suffix: landau_damping_1d
1720 args: -xdm_plex_dim 1 -xdm_plex_box_faces 60 -xdm_plex_box_lower 0. -xdm_plex_box_upper 12.5664 -xdm_plex_box_bd periodic -vdm_plex_dim 1 -vdm_plex_box_faces 60 \
1721 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none
1722
1723 test:
1724 suffix: landau_damping_2d
1725 args: -xdm_plex_dim 2 -xdm_plex_box_bd periodic,periodic -vdm_plex_dim 2 -xdm_plex_box_lower 0.,-.5 -vdm_plex_box_lower -6,-6 -vdm_plex_box_upper 6,6 -xdm_plex_box_faces 6,3 \
1726 -xdm_plex_box_upper 12.5664,.5 -vdm_plex_box_faces 15,15 -vdm_plex_box_bd none,none -vdm_plex_hash_location -vdm_plex_simplex 0
1727
1728 test:
1729 suffix: landau_damping_3d
1730 args: -xdm_plex_dim 3 -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_plex_dim 3 -vdm_plex_box_faces 4,4,4 \
1731 -vdm_plex_box_lower -6,-6,-6 -vdm_plex_box_upper 6,6,6 -vdm_plex_hash_location -vdm_plex_simplex 0 -vdm_plex_box_bd none,none,none
1732
1733 test:
1734 requires: !defined(PETSC_USE_DMLANDAU_2D)
1735 suffix: sphere_3d
1736 nsize: 1
1737 args: -use_landau_velocity_space -xdm_plex_dim 3 -vdm_landau_thermal_temps 1 -vdm_landau_device_type cpu -xdm_plex_box_faces 6,3,3 -xdm_plex_box_lower 0.,-1,-1 \
1738 -xdm_plex_box_upper 12.5664,1,1 -xdm_plex_box_bd periodic,periodic,periodic -vdm_landau_verbose 2 -vdm_landau_sphere -vdm_landau_map_sphere \
1739 -vdm_landau_domain_radius 6,6,6 -vdm_landau_sphere_inner_radius_90degree_scale .35 -vdm_refine 1
1740
1741 TEST*/
1742