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