xref: /petsc/src/ts/tests/ex28.c (revision bd412c90fba8895e9763ccee76c7dd2e09a982d7)
1 static char help[] = "Example application of the Bhatnagar-Gross-Krook (BGK) collision operator.\n\
2 This example is a 0D-1V setting for the kinetic equation\n\
3 https://en.wikipedia.org/wiki/Bhatnagar%E2%80%93Gross%E2%80%93Krook_operator\n";
4 
5 #include <petscdmplex.h>
6 #include <petscdmswarm.h>
7 #include <petscts.h>
8 #include <petscdraw.h>
9 #include <petscviewer.h>
10 
11 typedef struct {
12   PetscInt    particlesPerCell; /* The number of partices per cell */
13   PetscReal   momentTol;        /* Tolerance for checking moment conservation */
14   PetscBool   monitorhg;        /* Flag for using the TS histogram monitor */
15   PetscBool   monitorsp;        /* Flag for using the TS scatter monitor */
16   PetscBool   monitorks;        /* Monitor to perform KS test to the maxwellian */
17   PetscBool   error;            /* Flag for printing the error */
18   PetscInt    ostep;            /* print the energy at each ostep time steps */
19   PetscDraw   draw;             /* The draw object for histogram monitoring */
20   PetscDrawHG drawhg;           /* The histogram draw context for monitoring */
21   PetscDrawSP drawsp;           /* The scatter plot draw context for the monitor */
22   PetscDrawSP drawks;           /* Scatterplot draw context for KS test */
23 } AppCtx;
24 
25 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
26 {
27   PetscFunctionBeginUser;
28   options->monitorhg        = PETSC_FALSE;
29   options->monitorsp        = PETSC_FALSE;
30   options->monitorks        = PETSC_FALSE;
31   options->particlesPerCell = 1;
32   options->momentTol        = 100.0 * PETSC_MACHINE_EPSILON;
33   options->ostep            = 100;
34 
35   PetscOptionsBegin(comm, "", "Collision Options", "DMPLEX");
36   PetscCall(PetscOptionsBool("-monitorhg", "Flag to use the TS histogram monitor", "ex28.c", options->monitorhg, &options->monitorhg, NULL));
37   PetscCall(PetscOptionsBool("-monitorsp", "Flag to use the TS scatter plot monitor", "ex28.c", options->monitorsp, &options->monitorsp, NULL));
38   PetscCall(PetscOptionsBool("-monitorks", "Flag to plot KS test results", "ex28.c", options->monitorks, &options->monitorks, NULL));
39   PetscCall(PetscOptionsInt("-particles_per_cell", "Number of particles per cell", "ex28.c", options->particlesPerCell, &options->particlesPerCell, NULL));
40   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex28.c", options->ostep, &options->ostep, PETSC_NULL));
41   PetscOptionsEnd();
42 
43   PetscFunctionReturn(0);
44 }
45 
46 /* Create the mesh for velocity space */
47 static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user)
48 {
49   PetscFunctionBeginUser;
50   PetscCall(DMCreate(comm, dm));
51   PetscCall(DMSetType(*dm, DMPLEX));
52   PetscCall(DMSetFromOptions(*dm));
53   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
54   PetscFunctionReturn(0);
55 }
56 
57 /* Since we are putting the same number of particles in each cell, this amounts to a uniform distribution of v */
58 static PetscErrorCode SetInitialCoordinates(DM sw)
59 {
60   AppCtx        *user;
61   PetscRandom    rnd;
62   DM             dm;
63   DMPolytopeType ct;
64   PetscBool      simplex;
65   PetscReal     *centroid, *coords, *xi0, *v0, *J, *invJ, detJ, *vals;
66   PetscInt       dim, d, cStart, cEnd, c, Np, p;
67 
68   PetscFunctionBeginUser;
69   PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
70   PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
71   PetscCall(PetscRandomSetFromOptions(rnd));
72 
73   PetscCall(DMGetApplicationContext(sw, &user));
74   Np = user->particlesPerCell;
75   PetscCall(DMGetDimension(sw, &dim));
76   PetscCall(DMSwarmGetCellDM(sw, &dm));
77   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
78   PetscCall(DMPlexGetCellType(dm, cStart, &ct));
79   simplex = DMPolytopeTypeGetNumVertices(ct) == DMPolytopeTypeGetDim(ct) + 1 ? PETSC_TRUE : PETSC_FALSE;
80   PetscCall(PetscMalloc5(dim, &centroid, dim, &xi0, dim, &v0, dim * dim, &J, dim * dim, &invJ));
81   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
82   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
83   PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&vals));
84   for (c = cStart; c < cEnd; ++c) {
85     if (Np == 1) {
86       PetscCall(DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL));
87       for (d = 0; d < dim; ++d) {
88         coords[c * dim + d] = centroid[d];
89         if ((coords[c * dim + d] >= -1) && (coords[c * dim + d] <= 1)) {
90           vals[c] = 1.0;
91         } else {
92           vals[c] = 0.;
93         }
94       }
95     } else {
96       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ)); /* affine */
97       for (p = 0; p < Np; ++p) {
98         const PetscInt n   = c * Np + p;
99         PetscReal      sum = 0.0, refcoords[3];
100 
101         for (d = 0; d < dim; ++d) {
102           PetscCall(PetscRandomGetValueReal(rnd, &refcoords[d]));
103           sum += refcoords[d];
104         }
105         if (simplex && sum > 0.0)
106           for (d = 0; d < dim; ++d) refcoords[d] -= PetscSqrtReal(dim) * sum;
107         vals[n] = 1.0;
108         PetscCall(DMPlexReferenceToCoordinates(dm, c, 1, refcoords, &coords[n * dim]));
109       }
110     }
111   }
112   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
113   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&vals));
114   PetscCall(PetscFree5(centroid, xi0, v0, J, invJ));
115   PetscCall(PetscRandomDestroy(&rnd));
116   PetscFunctionReturn(0);
117 }
118 
119 /* The intiial conditions are just the initial particle weights */
120 static PetscErrorCode SetInitialConditions(DM dmSw, Vec u)
121 {
122   DM           dm;
123   AppCtx      *user;
124   PetscReal   *vals;
125   PetscScalar *initialConditions;
126   PetscInt     dim, d, cStart, cEnd, c, Np, p, n;
127 
128   PetscFunctionBeginUser;
129   PetscCall(VecGetLocalSize(u, &n));
130   PetscCall(DMGetApplicationContext(dmSw, &user));
131   Np = user->particlesPerCell;
132   PetscCall(DMSwarmGetCellDM(dmSw, &dm));
133   PetscCall(DMGetDimension(dm, &dim));
134   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
135   PetscCheck(n == (cEnd - cStart) * Np, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "TS solution local size %" PetscInt_FMT " != %" PetscInt_FMT " nm particles", n, (cEnd - cStart) * Np);
136   PetscCall(DMSwarmGetField(dmSw, "w_q", NULL, NULL, (void **)&vals));
137   PetscCall(VecGetArray(u, &initialConditions));
138   for (c = cStart; c < cEnd; ++c) {
139     for (p = 0; p < Np; ++p) {
140       const PetscInt n = c * Np + p;
141       for (d = 0; d < dim; d++) initialConditions[n] = vals[n];
142     }
143   }
144   PetscCall(VecRestoreArray(u, &initialConditions));
145   PetscCall(DMSwarmRestoreField(dmSw, "w_q", NULL, NULL, (void **)&vals));
146   PetscFunctionReturn(0);
147 }
148 
149 static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user)
150 {
151   PetscInt *cellid;
152   PetscInt  dim, cStart, cEnd, c, Np = user->particlesPerCell, p;
153 
154   PetscFunctionBeginUser;
155   PetscCall(DMGetDimension(dm, &dim));
156   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
157   PetscCall(DMSetType(*sw, DMSWARM));
158   PetscCall(DMSetDimension(*sw, dim));
159   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
160   PetscCall(DMSwarmSetCellDM(*sw, dm));
161   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", dim, PETSC_REAL));
162   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
163   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
164   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
165   PetscCall(DMSwarmSetLocalSizes(*sw, (cEnd - cStart) * Np, 0));
166   PetscCall(DMSetFromOptions(*sw));
167   PetscCall(DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
168   for (c = cStart; c < cEnd; ++c) {
169     for (p = 0; p < Np; ++p) {
170       const PetscInt n = c * Np + p;
171       cellid[n]        = c;
172     }
173   }
174   PetscCall(DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
175   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
176   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
177   PetscFunctionReturn(0);
178 }
179 
180 /*
181   f_t = 1/\tau \left( f_eq - f \right)
182   n_t = 1/\tau \left( \int f_eq - \int f \right) = 1/\tau (n - n) = 0
183   v_t = 1/\tau \left( \int v f_eq - \int v f \right) = 1/\tau (v - v) = 0
184   E_t = 1/\tau \left( \int v^2 f_eq - \int v^2 f \right) = 1/\tau (T - T) = 0
185 
186   Let's look at a single cell:
187 
188     \int_C f_t             = 1/\tau \left( \int_C f_eq - \int_C f \right)
189     \sum_{x_i \in C} w^i_t = 1/\tau \left( F_eq(C) - \sum_{x_i \in C} w_i \right)
190 */
191 
192 /* This computes the 1D Maxwellian distribution for given mass n, velocity v, and temperature T */
193 static PetscReal ComputePDF(PetscReal m, PetscReal n, PetscReal T, PetscReal v[])
194 {
195   return (n / PetscSqrtReal(2.0 * PETSC_PI * T / m)) * PetscExpReal(-0.5 * m * PetscSqr(v[0]) / T);
196 }
197 
198 /*
199   erf z = \frac{2}{\sqrt\pi} \int^z_0 dt e^{-t^2} and erf \infty = 1
200 
201   We begin with our distribution
202 
203     \sqrt{\frac{m}{2 \pi T}} e^{-m v^2/2T}
204 
205   Let t = \sqrt{m/2T} v, z = \sqrt{m/2T} w, so that we now have
206 
207       \sqrt{\frac{m}{2 \pi T}} \int^w_0 dv e^{-m v^2/2T}
208     = \sqrt{\frac{m}{2 \pi T}} \int^{\sqrt{m/2T} w}_0 \sqrt{2T/m} dt e^{-t^2}
209     = 1/\sqrt{\pi} \int^{\sqrt{m/2T} w}_0 dt e^{-t^2}
210     = 1/2 erf(\sqrt{m/2T} w)
211 */
212 static PetscReal ComputeCDF(PetscReal m, PetscReal n, PetscReal T, PetscReal va, PetscReal vb)
213 {
214   PetscReal alpha = PetscSqrtReal(0.5 * m / T);
215   PetscReal za    = alpha * va;
216   PetscReal zb    = alpha * vb;
217   PetscReal sum   = 0.0;
218 
219   sum += zb >= 0 ? erf(zb) : -erf(-zb);
220   sum -= za >= 0 ? erf(za) : -erf(-za);
221   return 0.5 * n * sum;
222 }
223 
224 static PetscErrorCode CheckDistribution(DM dm, PetscReal m, PetscReal n, PetscReal T, PetscReal v[])
225 {
226   PetscSection coordSection;
227   Vec          coordsLocal;
228   PetscReal   *xq, *wq;
229   PetscReal    vmin, vmax, neq, veq, Teq;
230   PetscInt     Nq = 100, q, cStart, cEnd, c;
231 
232   PetscFunctionBeginUser;
233   PetscCall(DMGetBoundingBox(dm, &vmin, &vmax));
234   /* Check analytic over entire line */
235   neq = ComputeCDF(m, n, T, vmin, vmax);
236   PetscCheck(PetscAbsReal(neq - n) <= PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Int f %g != %g mass (%g)", (double)neq, (double)n, (double)(neq - n));
237   /* Check analytic over cells */
238   PetscCall(DMGetCoordinatesLocal(dm, &coordsLocal));
239   PetscCall(DMGetCoordinateSection(dm, &coordSection));
240   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
241   neq = 0.0;
242   for (c = cStart; c < cEnd; ++c) {
243     PetscScalar *vcoords = NULL;
244 
245     PetscCall(DMPlexVecGetClosure(dm, coordSection, coordsLocal, c, NULL, &vcoords));
246     neq += ComputeCDF(m, n, T, vcoords[0], vcoords[1]);
247     PetscCall(DMPlexVecRestoreClosure(dm, coordSection, coordsLocal, c, NULL, &vcoords));
248   }
249   PetscCheck(PetscAbsReal(neq - n) <= PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Cell Int f %g != %g mass (%g)", (double)neq, (double)n, (double)(neq - n));
250   /* Check quadrature over entire line */
251   PetscCall(PetscMalloc2(Nq, &xq, Nq, &wq));
252   PetscCall(PetscDTGaussQuadrature(100, vmin, vmax, xq, wq));
253   neq = 0.0;
254   for (q = 0; q < Nq; ++q) neq += ComputePDF(m, n, T, &xq[q]) * wq[q];
255   PetscCheck(PetscAbsReal(neq - n) <= PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Int f %g != %g mass (%g)", (double)neq, (double)n, (double)(neq - n));
256   /* Check omemnts with quadrature */
257   veq = 0.0;
258   for (q = 0; q < Nq; ++q) veq += xq[q] * ComputePDF(m, n, T, &xq[q]) * wq[q];
259   veq /= neq;
260   PetscCheck(PetscAbsReal(veq - v[0]) <= PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Int v f %g != %g velocity (%g)", (double)veq, (double)v[0], (double)(veq - v[0]));
261   Teq = 0.0;
262   for (q = 0; q < Nq; ++q) Teq += PetscSqr(xq[q]) * ComputePDF(m, n, T, &xq[q]) * wq[q];
263   Teq = Teq * m / neq - PetscSqr(veq);
264   PetscCheck(PetscAbsReal(Teq - T) <= PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Int v^2 f %g != %g temperature (%g)", (double)Teq, (double)T, (double)(Teq - T));
265   PetscCall(PetscFree2(xq, wq));
266   PetscFunctionReturn(0);
267 }
268 
269 static PetscErrorCode RHSFunctionParticles(TS ts, PetscReal t, Vec U, Vec R, void *ctx)
270 {
271   const PetscScalar *u;
272   PetscSection       coordSection;
273   Vec                coordsLocal;
274   PetscScalar       *r, *coords;
275   PetscReal          n = 0.0, v = 0.0, E = 0.0, T = 0.0, m = 1.0, cn = 0.0, cv = 0.0, cE = 0.0, pE = 0.0, eqE = 0.0;
276   PetscInt           dim, d, Np, Ncp, p, cStart, cEnd, c;
277   DM                 dmSw, plex;
278 
279   PetscFunctionBeginUser;
280   PetscCall(VecGetLocalSize(U, &Np));
281   PetscCall(VecGetArrayRead(U, &u));
282   PetscCall(VecGetArray(R, &r));
283   PetscCall(TSGetDM(ts, &dmSw));
284   PetscCall(DMSwarmGetCellDM(dmSw, &plex));
285   PetscCall(DMGetDimension(dmSw, &dim));
286   PetscCall(DMGetCoordinatesLocal(plex, &coordsLocal));
287   PetscCall(DMGetCoordinateSection(plex, &coordSection));
288   PetscCall(DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd));
289   Np /= dim;
290   Ncp = Np / (cEnd - cStart);
291   /* Calculate moments of particle distribution, note that velocity is in the coordinate */
292   PetscCall(DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
293   for (p = 0; p < Np; ++p) {
294     PetscReal m1 = 0.0, m2 = 0.0;
295 
296     for (d = 0; d < dim; ++d) {
297       m1 += PetscRealPart(coords[p * dim + d]);
298       m2 += PetscSqr(PetscRealPart(coords[p * dim + d]));
299     }
300     n += u[p];
301     v += u[p] * m1;
302     E += u[p] * m2;
303   }
304   v /= n;
305   T = E * m / n - v * v;
306   PetscCall(PetscInfo(ts, "Time %.2f: mass %.4f velocity: %+.4f temperature: %.4f\n", (double)t, (double)n, (double)v, (double)T));
307   PetscCall(CheckDistribution(plex, m, n, T, &v));
308   /*
309      Begin cellwise evaluation of the collision operator. Essentially, penalize the weights of the particles
310      in that cell against the maxwellian for the number of particles expected to be in that cell
311   */
312   for (c = cStart; c < cEnd; ++c) {
313     PetscScalar *vcoords    = NULL;
314     PetscReal    relaxation = 1.0, neq;
315     PetscInt     sp         = c * Ncp, q;
316 
317     /* Calculate equilibrium occupation for this velocity cell */
318     PetscCall(DMPlexVecGetClosure(plex, coordSection, coordsLocal, c, NULL, &vcoords));
319     neq = ComputeCDF(m, n, T, vcoords[0], vcoords[1]);
320     PetscCall(DMPlexVecRestoreClosure(plex, coordSection, coordsLocal, c, NULL, &vcoords));
321     for (q = 0; q < Ncp; ++q) r[sp + q] = (1.0 / relaxation) * (neq - u[sp + q]);
322   }
323   /* Check update */
324   for (p = 0; p < Np; ++p) {
325     PetscReal    m1 = 0.0, m2 = 0.0;
326     PetscScalar *vcoords = NULL;
327 
328     for (d = 0; d < dim; ++d) {
329       m1 += PetscRealPart(coords[p * dim + d]);
330       m2 += PetscSqr(PetscRealPart(coords[p * dim + d]));
331     }
332     cn += r[p];
333     cv += r[p] * m1;
334     cE += r[p] * m2;
335     pE += u[p] * m2;
336     PetscCall(DMPlexVecGetClosure(plex, coordSection, coordsLocal, p, NULL, &vcoords));
337     eqE += ComputeCDF(m, n, T, vcoords[0], vcoords[1]) * m2;
338     PetscCall(DMPlexVecRestoreClosure(plex, coordSection, coordsLocal, p, NULL, &vcoords));
339   }
340   PetscCall(PetscInfo(ts, "Time %.2f: mass update %.8f velocity update: %+.8f energy update: %.8f (%.8f, %.8f)\n", (double)t, (double)cn, (double)cv, (double)cE, (double)pE, (double)eqE));
341   PetscCall(DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
342   PetscCall(VecRestoreArrayRead(U, &u));
343   PetscCall(VecRestoreArray(R, &r));
344   PetscFunctionReturn(0);
345 }
346 
347 static PetscErrorCode HGMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
348 {
349   AppCtx            *user = (AppCtx *)ctx;
350   const PetscScalar *u;
351   DM                 sw, dm;
352   PetscInt           dim, Np, p;
353 
354   PetscFunctionBeginUser;
355   if (step < 0) PetscFunctionReturn(0);
356   if (((user->ostep > 0) && (!(step % user->ostep)))) {
357     PetscDrawAxis axis;
358 
359     PetscCall(TSGetDM(ts, &sw));
360     PetscCall(DMSwarmGetCellDM(sw, &dm));
361     PetscCall(DMGetDimension(dm, &dim));
362     PetscCall(PetscDrawHGReset(user->drawhg));
363     PetscCall(PetscDrawHGGetAxis(user->drawhg, &axis));
364     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "V", "f(V)"));
365     PetscCall(PetscDrawAxisSetLimits(axis, -3, 3, 0, 100));
366     PetscCall(PetscDrawHGSetLimits(user->drawhg, -3.0, 3.0, 0, 10));
367     PetscCall(VecGetLocalSize(U, &Np));
368     Np /= dim;
369     PetscCall(VecGetArrayRead(U, &u));
370     /* get points from solution vector */
371     for (p = 0; p < Np; ++p) PetscCall(PetscDrawHGAddValue(user->drawhg, u[p]));
372     PetscCall(PetscDrawHGDraw(user->drawhg));
373     PetscCall(VecRestoreArrayRead(U, &u));
374   }
375   PetscFunctionReturn(0);
376 }
377 
378 static PetscErrorCode SPMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
379 {
380   AppCtx            *user = (AppCtx *)ctx;
381   const PetscScalar *u;
382   PetscReal         *v, *coords;
383   PetscInt           Np, p;
384   DM                 dmSw;
385 
386   PetscFunctionBeginUser;
387 
388   if (step < 0) PetscFunctionReturn(0);
389   if (((user->ostep > 0) && (!(step % user->ostep)))) {
390     PetscDrawAxis axis;
391 
392     PetscCall(TSGetDM(ts, &dmSw));
393     PetscCall(PetscDrawSPReset(user->drawsp));
394     PetscCall(PetscDrawSPGetAxis(user->drawsp, &axis));
395     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "V", "w"));
396     PetscCall(VecGetLocalSize(U, &Np));
397     PetscCall(VecGetArrayRead(U, &u));
398     /* get points from solution vector */
399     PetscCall(PetscMalloc1(Np, &v));
400     for (p = 0; p < Np; ++p) v[p] = PetscRealPart(u[p]);
401     PetscCall(DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
402     for (p = 0; p < Np - 1; ++p) PetscCall(PetscDrawSPAddPoint(user->drawsp, &coords[p], &v[p]));
403     PetscCall(PetscDrawSPDraw(user->drawsp, PETSC_TRUE));
404     PetscCall(VecRestoreArrayRead(U, &u));
405     PetscCall(DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
406     PetscCall(PetscFree(v));
407   }
408   PetscFunctionReturn(0);
409 }
410 
411 static PetscErrorCode KSConv(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
412 {
413   AppCtx            *user = (AppCtx *)ctx;
414   const PetscScalar *u;
415   PetscScalar        sup;
416   PetscReal         *v, *coords, T = 0., vel = 0., step_cast, w_sum;
417   PetscInt           dim, Np, p, cStart, cEnd;
418   DM                 sw, plex;
419 
420   PetscFunctionBeginUser;
421   if (step < 0) PetscFunctionReturn(0);
422   if (((user->ostep > 0) && (!(step % user->ostep)))) {
423     PetscDrawAxis axis;
424     PetscCall(PetscDrawSPGetAxis(user->drawks, &axis));
425     PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "t", "D_n"));
426     PetscCall(PetscDrawSPSetLimits(user->drawks, 0., 100, 0., 3.5));
427     PetscCall(TSGetDM(ts, &sw));
428     PetscCall(VecGetLocalSize(U, &Np));
429     PetscCall(VecGetArrayRead(U, &u));
430     /* get points from solution vector */
431     PetscCall(PetscMalloc1(Np, &v));
432     PetscCall(DMSwarmGetCellDM(sw, &plex));
433     PetscCall(DMGetDimension(sw, &dim));
434     PetscCall(DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd));
435     for (p = 0; p < Np; ++p) v[p] = PetscRealPart(u[p]);
436     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
437     w_sum = 0.;
438     for (p = 0; p < Np; ++p) {
439       w_sum += u[p];
440       T += u[p] * coords[p] * coords[p];
441       vel += u[p] * coords[p];
442     }
443     vel /= w_sum;
444     T   = T / w_sum - vel * vel;
445     sup = 0.0;
446     for (p = 0; p < Np; ++p) {
447       PetscReal tmp = 0.;
448       tmp           = PetscAbs(u[p] - ComputePDF(1.0, w_sum, T, &coords[p * dim]));
449       if (tmp > sup) sup = tmp;
450     }
451     step_cast = (PetscReal)step;
452     PetscCall(PetscDrawSPAddPoint(user->drawks, &step_cast, &sup));
453     PetscCall(PetscDrawSPDraw(user->drawks, PETSC_TRUE));
454     PetscCall(VecRestoreArrayRead(U, &u));
455     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
456     PetscCall(PetscFree(v));
457   }
458   PetscFunctionReturn(0);
459 }
460 
461 static PetscErrorCode InitializeSolve(TS ts, Vec u)
462 {
463   DM      dm;
464   AppCtx *user;
465 
466   PetscFunctionBeginUser;
467   PetscCall(TSGetDM(ts, &dm));
468   PetscCall(DMGetApplicationContext(dm, &user));
469   PetscCall(SetInitialCoordinates(dm));
470   PetscCall(SetInitialConditions(dm, u));
471   PetscFunctionReturn(0);
472 }
473 /*
474      A single particle is generated for each velocity space cell using the dmswarmpicfield_coor and is used to evaluate collisions in that cell.
475      0 weight ghost particles are initialized outside of a small velocity domain to ensure the tails of the amxwellian are resolved.
476    */
477 int main(int argc, char **argv)
478 {
479   TS       ts;     /* nonlinear solver */
480   DM       dm, sw; /* Velocity space mesh and Particle Swarm */
481   Vec      u, w;   /* swarm vector */
482   MPI_Comm comm;
483   AppCtx   user;
484 
485   PetscFunctionBeginUser;
486   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
487   comm = PETSC_COMM_WORLD;
488   PetscCall(ProcessOptions(comm, &user));
489   PetscCall(CreateMesh(comm, &dm, &user));
490   PetscCall(CreateParticles(dm, &sw, &user));
491   PetscCall(DMSetApplicationContext(sw, &user));
492   PetscCall(TSCreate(comm, &ts));
493   PetscCall(TSSetDM(ts, sw));
494   PetscCall(TSSetMaxTime(ts, 10.0));
495   PetscCall(TSSetTimeStep(ts, 0.01));
496   PetscCall(TSSetMaxSteps(ts, 100000));
497   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
498   if (user.monitorhg) {
499     PetscCall(PetscDrawCreate(comm, NULL, "monitor", 0, 0, 400, 300, &user.draw));
500     PetscCall(PetscDrawSetFromOptions(user.draw));
501     PetscCall(PetscDrawHGCreate(user.draw, 20, &user.drawhg));
502     PetscCall(PetscDrawHGSetColor(user.drawhg, 3));
503     PetscCall(TSMonitorSet(ts, HGMonitor, &user, NULL));
504   } else if (user.monitorsp) {
505     PetscCall(PetscDrawCreate(comm, NULL, "monitor", 0, 0, 400, 300, &user.draw));
506     PetscCall(PetscDrawSetFromOptions(user.draw));
507     PetscCall(PetscDrawSPCreate(user.draw, 1, &user.drawsp));
508     PetscCall(TSMonitorSet(ts, SPMonitor, &user, NULL));
509   } else if (user.monitorks) {
510     PetscCall(PetscDrawCreate(comm, NULL, "monitor", 0, 0, 400, 300, &user.draw));
511     PetscCall(PetscDrawSetFromOptions(user.draw));
512     PetscCall(PetscDrawSPCreate(user.draw, 1, &user.drawks));
513     PetscCall(TSMonitorSet(ts, KSConv, &user, NULL));
514   }
515   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionParticles, &user));
516   PetscCall(TSSetFromOptions(ts));
517   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
518   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &w));
519   PetscCall(VecDuplicate(w, &u));
520   PetscCall(VecCopy(w, u));
521   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &w));
522   PetscCall(TSComputeInitialCondition(ts, u));
523   PetscCall(TSSolve(ts, u));
524   if (user.monitorhg) {
525     PetscCall(PetscDrawSave(user.draw));
526     PetscCall(PetscDrawHGDestroy(&user.drawhg));
527     PetscCall(PetscDrawDestroy(&user.draw));
528   }
529   if (user.monitorsp) {
530     PetscCall(PetscDrawSave(user.draw));
531     PetscCall(PetscDrawSPDestroy(&user.drawsp));
532     PetscCall(PetscDrawDestroy(&user.draw));
533   }
534   if (user.monitorks) {
535     PetscCall(PetscDrawSave(user.draw));
536     PetscCall(PetscDrawSPDestroy(&user.drawks));
537     PetscCall(PetscDrawDestroy(&user.draw));
538   }
539   PetscCall(VecDestroy(&u));
540   PetscCall(TSDestroy(&ts));
541   PetscCall(DMDestroy(&sw));
542   PetscCall(DMDestroy(&dm));
543   PetscCall(PetscFinalize());
544   return 0;
545 }
546 
547 /*TEST
548    build:
549      requires: double !complex
550    test:
551      suffix: 1
552      args: -particles_per_cell 1 -output_step 10 -ts_type euler -dm_plex_dim 1 -dm_plex_box_faces 200 -dm_plex_box_lower -10 -dm_plex_box_upper 10 -dm_view -monitorsp
553    test:
554      suffix: 2
555      args: -particles_per_cell 1 -output_step 50 -ts_type euler -dm_plex_dim 1 -dm_plex_box_faces 200 -dm_plex_box_lower -10 -dm_plex_box_upper 10 -dm_view -monitorks
556 TEST*/
557