xref: /petsc/src/dm/impls/swarm/tests/ex6.c (revision 66af8762ec03dbef0e079729eb2a1734a35ed7ff)
1 static char help[] = "Vlasov-Poisson example of central orbits\n";
2 
3 /*
4   To visualize the orbit, we can used
5 
6     -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain -1 -ts_monitor_sp_swarm_phase 0 -draw_size 500,500
7 
8   and we probably want it to run fast and not check convergence
9 
10     -convest_num_refine 0 -ts_dt 0.01 -ts_max_steps 100 -ts_max_time 100 -output_step 10
11 */
12 
13 #include <petscts.h>
14 #include <petscdmplex.h>
15 #include <petscdmswarm.h>
16 #include <petsc/private/dmpleximpl.h> /* For norm and dot */
17 #include <petscfe.h>
18 #include <petscds.h>
19 #include <petsc/private/petscfeimpl.h> /* For interpolation */
20 #include <petscksp.h>
21 #include "petscsnes.h"
22 
23 PETSC_EXTERN PetscErrorCode circleSingleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
24 PETSC_EXTERN PetscErrorCode circleSingleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
25 PETSC_EXTERN PetscErrorCode circleMultipleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
26 PETSC_EXTERN PetscErrorCode circleMultipleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
27 
28 const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL};
29 typedef enum {
30   EM_PRIMAL,
31   EM_MIXED,
32   EM_COULOMB,
33   EM_NONE
34 } EMType;
35 
36 typedef struct {
37   PetscBool error;     /* Flag for printing the error */
38   PetscInt  ostep;     /* print the energy at each ostep time steps */
39   PetscReal timeScale; /* Nondimensionalizing time scale */
40   PetscReal sigma;     /* Linear charge per box length */
41   EMType    em;        /* Type of electrostatic model */
42   SNES      snes;
43 } AppCtx;
44 
45 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
46 {
47   PetscFunctionBeginUser;
48   options->error     = PETSC_FALSE;
49   options->ostep     = 100;
50   options->timeScale = 1.0e-6;
51   options->sigma     = 1.;
52   options->em        = EM_COULOMB;
53 
54   PetscOptionsBegin(comm, "", "Central Orbit Options", "DMSWARM");
55   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex6.c", options->error, &options->error, NULL));
56   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex6.c", options->ostep, &options->ostep, NULL));
57   PetscCall(PetscOptionsReal("-sigma", "Linear charge per box length", "ex6.c", options->sigma, &options->sigma, NULL));
58   PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex6.c", options->timeScale, &options->timeScale, NULL));
59   PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex6.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL));
60   PetscOptionsEnd();
61   PetscFunctionReturn(PETSC_SUCCESS);
62 }
63 
64 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
65 {
66   PetscFunctionBeginUser;
67   PetscCall(DMCreate(comm, dm));
68   PetscCall(DMSetType(*dm, DMPLEX));
69   PetscCall(DMSetFromOptions(*dm));
70   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
71   PetscFunctionReturn(PETSC_SUCCESS);
72 }
73 
74 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[])
75 {
76   PetscInt d;
77   for (d = 0; d < dim; ++d) f1[d] = u_x[d];
78 }
79 
80 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[])
81 {
82   PetscInt d;
83   for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0;
84 }
85 
86 /*
87    /  I   grad\ /q\ = /0\
88    \-div    0 / \u/   \f/
89 */
90 static void f0_q(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
91 {
92   for (PetscInt c = 0; c < dim; ++c) f0[c] += u[uOff[0] + c];
93 }
94 
95 static void f1_q(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[])
96 {
97   for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] += u[uOff[1]];
98 }
99 
100 /* <t, q> */
101 static void g0_qq(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g0[])
102 {
103   PetscInt c;
104   for (c = 0; c < dim; ++c) g0[c * dim + c] += 1.0;
105 }
106 
107 static void g2_qu(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g2[])
108 {
109   for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] += 1.0;
110 }
111 
112 static void g1_uq(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g1[])
113 {
114   for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] += 1.0;
115 }
116 
117 static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
118 {
119   PetscFE   feu, feq;
120   PetscDS   ds;
121   PetscBool simplex;
122   PetscInt  dim;
123 
124   PetscFunctionBeginUser;
125   PetscCall(DMGetDimension(dm, &dim));
126   PetscCall(DMPlexIsSimplex(dm, &simplex));
127   if (user->em == EM_MIXED) {
128     DMLabel label;
129 
130     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq));
131     PetscCall(PetscObjectSetName((PetscObject)feq, "field"));
132     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &feu));
133     PetscCall(PetscObjectSetName((PetscObject)feu, "potential"));
134     PetscCall(PetscFECopyQuadrature(feq, feu));
135     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq));
136     PetscCall(DMSetField(dm, 1, NULL, (PetscObject)feu));
137     PetscCall(DMCreateDS(dm));
138     PetscCall(PetscFEDestroy(&feu));
139     PetscCall(PetscFEDestroy(&feq));
140 
141     PetscCall(DMGetLabel(dm, "marker", &label));
142     PetscCall(DMGetDS(dm, &ds));
143     PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q));
144     PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL));
145     PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qu, NULL));
146     PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_uq, NULL, NULL));
147   } else if (user->em == EM_PRIMAL) {
148     PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &feu));
149     PetscCall(PetscObjectSetName((PetscObject)feu, "potential"));
150     PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feu));
151     PetscCall(DMCreateDS(dm));
152     PetscCall(PetscFEDestroy(&feu));
153     PetscCall(DMGetDS(dm, &ds));
154     PetscCall(PetscDSSetResidual(ds, 0, NULL, laplacian_f1));
155     PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3));
156   }
157   PetscFunctionReturn(PETSC_SUCCESS);
158 }
159 
160 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user)
161 {
162   SNES         snes;
163   Mat          J;
164   MatNullSpace nullSpace;
165 
166   PetscFunctionBeginUser;
167   PetscCall(CreateFEM(dm, user));
168   PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes));
169   PetscCall(SNESSetOptionsPrefix(snes, "em_"));
170   PetscCall(SNESSetDM(snes, dm));
171   PetscCall(DMPlexSetSNESLocalFEM(dm, user, user, user));
172   PetscCall(SNESSetFromOptions(snes));
173 
174   PetscCall(DMCreateMatrix(dm, &J));
175   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace));
176   PetscCall(MatSetNullSpace(J, nullSpace));
177   PetscCall(MatNullSpaceDestroy(&nullSpace));
178   PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL));
179   PetscCall(MatDestroy(&J));
180   user->snes = snes;
181   PetscFunctionReturn(PETSC_SUCCESS);
182 }
183 
184 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
185 {
186   PetscReal v0[1] = {1.};
187   PetscInt  dim;
188 
189   PetscFunctionBeginUser;
190   PetscCall(DMGetDimension(dm, &dim));
191   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
192   PetscCall(DMSetType(*sw, DMSWARM));
193   PetscCall(DMSetDimension(*sw, dim));
194   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
195   PetscCall(DMSwarmSetCellDM(*sw, dm));
196   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
197   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
198   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
199   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initCoordinates", dim, PETSC_REAL));
200   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initVelocity", dim, PETSC_REAL));
201   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL));
202   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
203   PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
204   PetscCall(DMSwarmInitializeCoordinates(*sw));
205   PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
206   PetscCall(DMSetFromOptions(*sw));
207   PetscCall(DMSetApplicationContext(*sw, user));
208   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
209   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
210   {
211     Vec gc, gc0, gv, gv0;
212 
213     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
214     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0));
215     PetscCall(VecCopy(gc, gc0));
216     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
217     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0));
218     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv));
219     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0));
220     PetscCall(VecCopy(gv, gv0));
221     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv));
222     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0));
223   }
224   PetscFunctionReturn(PETSC_SUCCESS);
225 }
226 
227 static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[])
228 {
229   PetscReal  *coords;
230   PetscInt    dim, d, Np, p, q;
231   PetscMPIInt size;
232 
233   PetscFunctionBegin;
234   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size));
235   PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial");
236   PetscCall(DMGetDimension(sw, &dim));
237   PetscCall(DMSwarmGetLocalSize(sw, &Np));
238 
239   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
240   for (p = 0; p < Np; ++p) {
241     PetscReal *pcoord = &coords[p * dim];
242     PetscReal *pE     = &E[p * dim];
243     /* Calculate field at particle p due to particle q */
244     for (q = 0; q < Np; ++q) {
245       PetscReal *qcoord = &coords[q * dim];
246       PetscReal  rpq[3], r;
247 
248       if (p == q) continue;
249       for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
250       r = DMPlex_NormD_Internal(dim, rpq);
251       for (d = 0; d < dim; ++d) pE[d] += rpq[d] / PetscPowRealInt(r, 3);
252     }
253   }
254   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
255   PetscFunctionReturn(PETSC_SUCCESS);
256 }
257 
258 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[])
259 {
260   DM         dm;
261   PetscDS    ds;
262   PetscFE    fe;
263   Mat        M_p;
264   Vec        phi, locPhi, rho, f;
265   PetscReal *coords, chargeTol = 1e-13;
266   PetscInt   dim, d, cStart, cEnd, c, Np;
267 
268   PetscFunctionBegin;
269   PetscCall(DMGetDimension(sw, &dim));
270   PetscCall(DMSwarmGetLocalSize(sw, &Np));
271 
272   /* Create the charges rho */
273   PetscCall(SNESGetDM(snes, &dm));
274   PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
275   PetscCall(DMGetGlobalVector(dm, &rho));
276   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
277   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
278   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
279   PetscCall(MatMultTranspose(M_p, f, rho));
280   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
281   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
282   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
283   PetscCall(MatDestroy(&M_p));
284   {
285     PetscScalar sum;
286     PetscInt    n;
287     PetscReal   phi_0 = 1.; /* (sigma*sigma*sigma)*(timeScale*timeScale)/(m_e*q_e*epsi_0)*/
288 
289     /* Remove constant from rho */
290     PetscCall(VecGetSize(rho, &n));
291     PetscCall(VecSum(rho, &sum));
292     PetscCall(VecShift(rho, -sum / n));
293     PetscCall(VecSum(rho, &sum));
294     PetscCheck(PetscAbsScalar(sum) < chargeTol, PetscObjectComm((PetscObject)sw), PETSC_ERR_PLIB, "Charge should have no DC component: %g", (double)PetscRealPart(sum));
295     /* Nondimensionalize rho */
296     PetscCall(VecScale(rho, phi_0));
297   }
298   PetscCall(VecViewFromOptions(rho, NULL, "-poisson_rho_view"));
299 
300   PetscCall(DMGetGlobalVector(dm, &phi));
301   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
302   PetscCall(VecSet(phi, 0.0));
303   PetscCall(SNESSolve(snes, rho, phi));
304   PetscCall(DMRestoreGlobalVector(dm, &rho));
305   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
306 
307   PetscCall(DMGetLocalVector(dm, &locPhi));
308   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
309   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
310   PetscCall(DMRestoreGlobalVector(dm, &phi));
311 
312   PetscCall(DMGetDS(dm, &ds));
313   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
314   PetscCall(DMSwarmSortGetAccess(sw));
315   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
316   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
317   for (c = cStart; c < cEnd; ++c) {
318     PetscTabulation tab;
319     PetscScalar    *clPhi = NULL;
320     PetscReal      *pcoord, *refcoord;
321     PetscReal       v[3], J[9], invJ[9], detJ;
322     PetscInt       *points;
323     PetscInt        Ncp, cp;
324 
325     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
326     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
327     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
328     for (cp = 0; cp < Ncp; ++cp)
329       for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
330     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
331     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
332     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ));
333     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
334     for (cp = 0; cp < Ncp; ++cp) {
335       const PetscReal *basisDer = tab->T[1];
336       const PetscInt   p        = points[cp];
337 
338       for (d = 0; d < dim; ++d) E[p * dim + d] = 0.;
339       PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, invJ, NULL, cp, &E[p * dim]));
340       for (d = 0; d < dim; ++d) E[p * dim + d] *= -1.0;
341     }
342     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
343     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
344     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
345     PetscCall(PetscTabulationDestroy(&tab));
346     PetscCall(PetscFree(points));
347   }
348   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
349   PetscCall(DMSwarmSortRestoreAccess(sw));
350   PetscCall(DMRestoreLocalVector(dm, &locPhi));
351   PetscFunctionReturn(PETSC_SUCCESS);
352 }
353 
354 static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[])
355 {
356   DM              dm, potential_dm;
357   IS              potential_IS;
358   PetscDS         ds;
359   PetscFE         fe;
360   PetscFEGeom     feGeometry;
361   Mat             M_p;
362   Vec             phi, locPhi, rho, f, temp_rho;
363   PetscQuadrature q;
364   PetscReal      *coords, chargeTol = 1e-13;
365   PetscInt        dim, d, cStart, cEnd, c, Np, fields = 1;
366 
367   PetscFunctionBegin;
368   PetscCall(DMGetDimension(sw, &dim));
369   PetscCall(DMSwarmGetLocalSize(sw, &Np));
370 
371   /* Create the charges rho */
372   PetscCall(SNESGetDM(snes, &dm));
373   PetscCall(DMGetGlobalVector(dm, &rho));
374   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
375 
376   PetscCall(DMCreateSubDM(dm, 1, &fields, &potential_IS, &potential_dm));
377   PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p));
378   PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view"));
379   PetscCall(DMGetGlobalVector(potential_dm, &temp_rho));
380   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
381   PetscCall(PetscObjectSetName((PetscObject)f, "particle weight"));
382   PetscCall(VecViewFromOptions(f, NULL, "-weights_view"));
383   PetscCall(MatMultTranspose(M_p, f, temp_rho));
384   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
385   PetscCall(MatDestroy(&M_p));
386   PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
387   PetscCall(VecViewFromOptions(rho, NULL, "-poisson_rho_view"));
388   PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho));
389   PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho));
390   PetscCall(DMDestroy(&potential_dm));
391   PetscCall(ISDestroy(&potential_IS));
392   {
393     PetscScalar sum;
394     PetscInt    n;
395     PetscReal   phi_0 = 1.; /*(sigma*sigma*sigma)*(timeScale*timeScale)/(m_e*q_e*epsi_0);*/
396 
397     /*Remove constant from rho*/
398     PetscCall(VecGetSize(rho, &n));
399     PetscCall(VecSum(rho, &sum));
400     PetscCall(VecShift(rho, -sum / n));
401     PetscCall(VecSum(rho, &sum));
402     PetscCheck(PetscAbsScalar(sum) < chargeTol, PetscObjectComm((PetscObject)sw), PETSC_ERR_PLIB, "Charge should have no DC component: %g", (double)PetscRealPart(sum));
403     /* Nondimensionalize rho */
404     PetscCall(VecScale(rho, phi_0));
405   }
406   PetscCall(DMGetGlobalVector(dm, &phi));
407   PetscCall(PetscObjectSetName((PetscObject)phi, "potential"));
408   PetscCall(VecSet(phi, 0.0));
409   PetscCall(SNESSolve(snes, NULL, phi));
410   PetscCall(DMRestoreGlobalVector(dm, &rho));
411   PetscCall(VecViewFromOptions(phi, NULL, "-phi_view"));
412 
413   PetscCall(DMGetLocalVector(dm, &locPhi));
414   PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi));
415   PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi));
416   PetscCall(DMRestoreGlobalVector(dm, &phi));
417 
418   PetscCall(DMGetDS(dm, &ds));
419   PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe));
420   PetscCall(DMSwarmSortGetAccess(sw));
421   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
422   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
423   for (c = cStart; c < cEnd; ++c) {
424     PetscTabulation tab;
425     PetscScalar    *clPhi = NULL;
426     PetscReal      *pcoord, *refcoord;
427     PetscReal       v[3], J[9], invJ[9], detJ;
428     PetscInt       *points;
429     PetscInt        Ncp, cp;
430 
431     PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points));
432     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
433     PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
434     for (cp = 0; cp < Ncp; ++cp)
435       for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d];
436     PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord));
437     PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab));
438     PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ));
439     PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi));
440     for (cp = 0; cp < Ncp; ++cp) {
441       const PetscInt p = points[cp];
442 
443       for (d = 0; d < dim; ++d) E[p * dim + d] = 0.;
444       PetscCall(PetscFEGetQuadrature(fe, &q));
445       PetscCall(PetscFECreateCellGeometry(fe, q, &feGeometry));
446       PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, &feGeometry, cp, &E[p * dim]));
447       PetscCall(PetscFEDestroyCellGeometry(fe, &feGeometry));
448     }
449     PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi));
450     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord));
451     PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord));
452     PetscCall(PetscTabulationDestroy(&tab));
453     PetscCall(PetscFree(points));
454   }
455   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
456   PetscCall(DMSwarmSortRestoreAccess(sw));
457   PetscCall(DMRestoreLocalVector(dm, &locPhi));
458   PetscFunctionReturn(PETSC_SUCCESS);
459 }
460 
461 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[])
462 {
463   AppCtx  *ctx;
464   PetscInt dim, Np;
465 
466   PetscFunctionBegin;
467   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
468   PetscValidHeaderSpecific(sw, DM_CLASSID, 2);
469   PetscAssertPointer(E, 3);
470   PetscCall(DMGetDimension(sw, &dim));
471   PetscCall(DMSwarmGetLocalSize(sw, &Np));
472   PetscCall(DMGetApplicationContext(sw, &ctx));
473   PetscCall(PetscArrayzero(E, Np * dim));
474 
475   switch (ctx->em) {
476   case EM_PRIMAL:
477     PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E));
478     break;
479   case EM_COULOMB:
480     PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E));
481     break;
482   case EM_MIXED:
483     PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E));
484     break;
485   case EM_NONE:
486     break;
487   default:
488     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]);
489   }
490 
491   PetscFunctionReturn(PETSC_SUCCESS);
492 }
493 
494 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
495 {
496   DM                 sw;
497   SNES               snes = ((AppCtx *)ctx)->snes;
498   const PetscReal   *coords, *vel;
499   const PetscScalar *u;
500   PetscScalar       *g;
501   PetscReal         *E;
502   PetscInt           dim, d, Np, p;
503 
504   PetscFunctionBeginUser;
505   PetscCall(TSGetDM(ts, &sw));
506   PetscCall(DMGetDimension(sw, &dim));
507   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
508   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
509   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
510   PetscCall(VecGetLocalSize(U, &Np));
511   PetscCall(VecGetArrayRead(U, &u));
512   PetscCall(VecGetArray(G, &g));
513 
514   PetscCall(ComputeFieldAtParticles(snes, sw, E));
515 
516   Np /= 2 * dim;
517   for (p = 0; p < Np; ++p) {
518     const PetscReal x0    = coords[p * dim + 0];
519     const PetscReal vy0   = vel[p * dim + 1];
520     const PetscReal omega = vy0 / x0;
521 
522     for (d = 0; d < dim; ++d) {
523       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
524       g[(p * 2 + 1) * dim + d] = E[p * dim + d] - PetscSqr(omega) * u[(p * 2 + 0) * dim + d];
525     }
526   }
527   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
528   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
529   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
530   PetscCall(VecRestoreArrayRead(U, &u));
531   PetscCall(VecRestoreArray(G, &g));
532   PetscFunctionReturn(PETSC_SUCCESS);
533 }
534 
535 /* J_{ij} = dF_i/dx_j
536    J_p = (  0   1)
537          (-w^2  0)
538    TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator.
539         Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb
540 */
541 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
542 {
543   DM               sw;
544   const PetscReal *coords, *vel;
545   PetscInt         dim, d, Np, p, rStart;
546 
547   PetscFunctionBeginUser;
548   PetscCall(TSGetDM(ts, &sw));
549   PetscCall(DMGetDimension(sw, &dim));
550   PetscCall(VecGetLocalSize(U, &Np));
551   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
552   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
553   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
554   Np /= 2 * dim;
555   for (p = 0; p < Np; ++p) {
556     const PetscReal x0      = coords[p * dim + 0];
557     const PetscReal vy0     = vel[p * dim + 1];
558     const PetscReal omega   = vy0 / x0;
559     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};
560 
561     for (d = 0; d < dim; ++d) {
562       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
563       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
564     }
565   }
566   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
567   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
568   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
569   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
570   PetscFunctionReturn(PETSC_SUCCESS);
571 }
572 
573 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
574 {
575   DM                 sw;
576   const PetscScalar *v;
577   PetscScalar       *xres;
578   PetscInt           Np, p, dim, d;
579 
580   PetscFunctionBeginUser;
581   PetscCall(TSGetDM(ts, &sw));
582   PetscCall(DMGetDimension(sw, &dim));
583   PetscCall(VecGetLocalSize(Xres, &Np));
584   Np /= dim;
585   PetscCall(VecGetArrayRead(V, &v));
586   PetscCall(VecGetArray(Xres, &xres));
587   for (p = 0; p < Np; ++p) {
588     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
589   }
590   PetscCall(VecRestoreArrayRead(V, &v));
591   PetscCall(VecRestoreArray(Xres, &xres));
592   PetscFunctionReturn(PETSC_SUCCESS);
593 }
594 
595 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
596 {
597   DM                 sw;
598   SNES               snes = ((AppCtx *)ctx)->snes;
599   const PetscScalar *x;
600   const PetscReal   *coords, *vel;
601   PetscScalar       *vres;
602   PetscReal         *E;
603   PetscInt           Np, p, dim, d;
604 
605   PetscFunctionBeginUser;
606   PetscCall(TSGetDM(ts, &sw));
607   PetscCall(DMGetDimension(sw, &dim));
608   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
609   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
610   PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
611   PetscCall(VecGetLocalSize(Vres, &Np));
612   PetscCall(VecGetArrayRead(X, &x));
613   PetscCall(VecGetArray(Vres, &vres));
614   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2");
615 
616   PetscCall(ComputeFieldAtParticles(snes, sw, E));
617 
618   Np /= dim;
619   for (p = 0; p < Np; ++p) {
620     const PetscReal x0    = coords[p * dim + 0];
621     const PetscReal vy0   = vel[p * dim + 1];
622     const PetscReal omega = vy0 / x0;
623 
624     for (d = 0; d < dim; ++d) vres[p * dim + d] = E[p * dim + d] - PetscSqr(omega) * x[p * dim + d];
625   }
626   PetscCall(VecRestoreArrayRead(X, &x));
627   PetscCall(VecRestoreArray(Vres, &vres));
628   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
629   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
630   PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
631   PetscFunctionReturn(PETSC_SUCCESS);
632 }
633 
634 static PetscErrorCode CreateSolution(TS ts)
635 {
636   DM       sw;
637   Vec      u;
638   PetscInt dim, Np;
639 
640   PetscFunctionBegin;
641   PetscCall(TSGetDM(ts, &sw));
642   PetscCall(DMGetDimension(sw, &dim));
643   PetscCall(DMSwarmGetLocalSize(sw, &Np));
644   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
645   PetscCall(VecSetBlockSize(u, dim));
646   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
647   PetscCall(VecSetUp(u));
648   PetscCall(TSSetSolution(ts, u));
649   PetscCall(VecDestroy(&u));
650   PetscFunctionReturn(PETSC_SUCCESS);
651 }
652 
653 static PetscErrorCode SetProblem(TS ts)
654 {
655   AppCtx *user;
656   DM      sw;
657 
658   PetscFunctionBegin;
659   PetscCall(TSGetDM(ts, &sw));
660   PetscCall(DMGetApplicationContext(sw, (void **)&user));
661   // Define unified system for (X, V)
662   {
663     Mat      J;
664     PetscInt dim, Np;
665 
666     PetscCall(DMGetDimension(sw, &dim));
667     PetscCall(DMSwarmGetLocalSize(sw, &Np));
668     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
669     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
670     PetscCall(MatSetBlockSize(J, 2 * dim));
671     PetscCall(MatSetFromOptions(J));
672     PetscCall(MatSetUp(J));
673     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
674     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
675     PetscCall(MatDestroy(&J));
676   }
677   /* Define split system for X and V */
678   {
679     Vec             u;
680     IS              isx, isv, istmp;
681     const PetscInt *idx;
682     PetscInt        dim, Np, rstart;
683 
684     PetscCall(TSGetSolution(ts, &u));
685     PetscCall(DMGetDimension(sw, &dim));
686     PetscCall(DMSwarmGetLocalSize(sw, &Np));
687     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
688     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
689     PetscCall(ISGetIndices(istmp, &idx));
690     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
691     PetscCall(ISRestoreIndices(istmp, &idx));
692     PetscCall(ISDestroy(&istmp));
693     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
694     PetscCall(ISGetIndices(istmp, &idx));
695     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
696     PetscCall(ISRestoreIndices(istmp, &idx));
697     PetscCall(ISDestroy(&istmp));
698     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
699     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
700     PetscCall(ISDestroy(&isx));
701     PetscCall(ISDestroy(&isv));
702     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
703     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
704   }
705   PetscFunctionReturn(PETSC_SUCCESS);
706 }
707 
708 PetscErrorCode circleSingleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
709 {
710   x[0] = p + 1.;
711   x[1] = 0.;
712   return PETSC_SUCCESS;
713 }
714 
715 PetscErrorCode circleSingleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx)
716 {
717   v[0] = 0.;
718   v[1] = PetscSqrtReal(1000. / (p + 1.));
719   return PETSC_SUCCESS;
720 }
721 
722 /* Put 5 particles into each circle */
723 PetscErrorCode circleMultipleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
724 {
725   const PetscInt  n   = 5;
726   const PetscReal r0  = (p / n) + 1.;
727   const PetscReal th0 = (2. * PETSC_PI * (p % n)) / n;
728 
729   x[0] = r0 * PetscCosReal(th0);
730   x[1] = r0 * PetscSinReal(th0);
731   return PETSC_SUCCESS;
732 }
733 
734 /* Put 5 particles into each circle */
735 PetscErrorCode circleMultipleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx)
736 {
737   const PetscInt  n     = 5;
738   const PetscReal r0    = (p / n) + 1.;
739   const PetscReal th0   = (2. * PETSC_PI * (p % n)) / n;
740   const PetscReal omega = PetscSqrtReal(1000. / r0) / r0;
741 
742   v[0] = -r0 * omega * PetscSinReal(th0);
743   v[1] = r0 * omega * PetscCosReal(th0);
744   return PETSC_SUCCESS;
745 }
746 
747 /*
748   InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.
749 
750   Input Parameters:
751 + ts         - The TS
752 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem
753 
754   Output Parameter:
755 . u - The initialized solution vector
756 
757   Level: advanced
758 
759 .seealso: InitializeSolve()
760 */
761 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
762 {
763   DM      sw;
764   Vec     u, gc, gv, gc0, gv0;
765   IS      isx, isv;
766   AppCtx *user;
767 
768   PetscFunctionBeginUser;
769   PetscCall(TSGetDM(ts, &sw));
770   PetscCall(DMGetApplicationContext(sw, &user));
771   if (useInitial) {
772     PetscReal v0[1] = {1.};
773 
774     PetscCall(DMSwarmInitializeCoordinates(sw));
775     PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
776     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
777     PetscCall(TSReset(ts));
778     PetscCall(CreateSolution(ts));
779     PetscCall(SetProblem(ts));
780   }
781   PetscCall(TSGetSolution(ts, &u));
782   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
783   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
784   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
785   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
786   if (useInitial) PetscCall(VecCopy(gc, gc0));
787   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
788   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
789   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
790   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
791   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
792   if (useInitial) PetscCall(VecCopy(gv, gv0));
793   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
794   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
795   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
796   PetscFunctionReturn(PETSC_SUCCESS);
797 }
798 
799 static PetscErrorCode InitializeSolve(TS ts, Vec u)
800 {
801   PetscFunctionBegin;
802   PetscCall(TSSetSolution(ts, u));
803   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
804   PetscFunctionReturn(PETSC_SUCCESS);
805 }
806 
807 static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
808 {
809   MPI_Comm           comm;
810   DM                 sw;
811   AppCtx            *user;
812   const PetscScalar *u;
813   const PetscReal   *coords, *vel;
814   PetscScalar       *e;
815   PetscReal          t;
816   PetscInt           dim, Np, p;
817 
818   PetscFunctionBeginUser;
819   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
820   PetscCall(TSGetDM(ts, &sw));
821   PetscCall(DMGetApplicationContext(sw, &user));
822   PetscCall(DMGetDimension(sw, &dim));
823   PetscCall(TSGetSolveTime(ts, &t));
824   PetscCall(VecGetArray(E, &e));
825   PetscCall(VecGetArrayRead(U, &u));
826   PetscCall(VecGetLocalSize(U, &Np));
827   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
828   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
829   Np /= 2 * dim;
830   for (p = 0; p < Np; ++p) {
831     /* TODO generalize initial conditions and project into plane instead of assuming x-y */
832     const PetscReal    r0    = DMPlex_NormD_Internal(dim, &coords[p * dim]);
833     const PetscReal    th0   = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]);
834     const PetscReal    v0    = DMPlex_NormD_Internal(dim, &vel[p * dim]);
835     const PetscReal    omega = v0 / r0;
836     const PetscReal    ct    = PetscCosReal(omega * t + th0);
837     const PetscReal    st    = PetscSinReal(omega * t + th0);
838     const PetscScalar *x     = &u[(p * 2 + 0) * dim];
839     const PetscScalar *v     = &u[(p * 2 + 1) * dim];
840     const PetscReal    xe[3] = {r0 * ct, r0 * st, 0.0};
841     const PetscReal    ve[3] = {-v0 * st, v0 * ct, 0.0};
842     PetscInt           d;
843 
844     for (d = 0; d < dim; ++d) {
845       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
846       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
847     }
848     if (user->error) {
849       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
850       const PetscReal exen = 0.5 * PetscSqr(v0);
851       PetscCall(PetscPrintf(comm, "t %.4g: p%" PetscInt_FMT " error [%.2f %.2f] sol [(%.6lf %.6lf) (%.6lf %.6lf)] exact [(%.6lf %.6lf) (%.6lf %.6lf)] energy/exact energy %g / %g (%.10lf%%)\n", (double)t, p, (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 0) * dim]), (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 1) * dim]), (double)x[0], (double)x[1], (double)v[0], (double)v[1], (double)xe[0], (double)xe[1], (double)ve[0], (double)ve[1], (double)en, (double)exen, (double)(PetscAbsReal(exen - en) * 100. / exen)));
852     }
853   }
854   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
855   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
856   PetscCall(VecRestoreArrayRead(U, &u));
857   PetscCall(VecRestoreArray(E, &e));
858   PetscFunctionReturn(PETSC_SUCCESS);
859 }
860 
861 static PetscErrorCode EnergyMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
862 {
863   const PetscInt     ostep = ((AppCtx *)ctx)->ostep;
864   const EMType       em    = ((AppCtx *)ctx)->em;
865   DM                 sw;
866   const PetscScalar *u;
867   PetscReal         *coords, *E;
868   PetscReal          enKin = 0., enEM = 0.;
869   PetscInt           dim, d, Np, p, q;
870 
871   PetscFunctionBeginUser;
872   if (step % ostep == 0) {
873     PetscCall(TSGetDM(ts, &sw));
874     PetscCall(DMGetDimension(sw, &dim));
875     PetscCall(VecGetArrayRead(U, &u));
876     PetscCall(VecGetLocalSize(U, &Np));
877     Np /= 2 * dim;
878     PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
879     PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E));
880     if (!step) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)ts), "Time     Step Part     Energy\n"));
881     for (p = 0; p < Np; ++p) {
882       const PetscReal v2     = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);
883       PetscReal      *pcoord = &coords[p * dim];
884 
885       PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " %5" PetscInt_FMT " %10.4lf\n", (double)t, step, p, (double)(0.5 * v2)));
886       enKin += 0.5 * v2;
887       if (em == EM_NONE) {
888         continue;
889       } else if (em == EM_COULOMB) {
890         for (q = p + 1; q < Np; ++q) {
891           PetscReal *qcoord = &coords[q * dim];
892           PetscReal  rpq[3], r;
893           for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d];
894           r = DMPlex_NormD_Internal(dim, rpq);
895           enEM += 1. / r;
896         }
897       } else if (em == EM_PRIMAL || em == EM_MIXED) {
898         for (d = 0; d < dim; ++d) enEM += E[p * dim + d];
899       }
900     }
901     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " KE\t    %10.4lf\n", (double)t, step, (double)enKin));
902     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " PE\t    %1.10g\n", (double)t, step, (double)enEM));
903     PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " E\t    %10.4lf\n", (double)t, step, (double)(enKin + enEM)));
904     PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords));
905     PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E));
906     PetscCall(PetscSynchronizedFlush(PetscObjectComm((PetscObject)ts), NULL));
907     PetscCall(VecRestoreArrayRead(U, &u));
908   }
909   PetscFunctionReturn(PETSC_SUCCESS);
910 }
911 
912 static PetscErrorCode SetUpMigrateParticles(TS ts, PetscInt n, PetscReal t, Vec x, PetscBool *flg, void *ctx)
913 {
914   DM sw;
915 
916   PetscFunctionBeginUser;
917   *flg = PETSC_TRUE;
918   PetscCall(TSGetDM(ts, &sw));
919   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
920   {
921     Vec u, gc, gv;
922     IS  isx, isv;
923 
924     PetscCall(TSGetSolution(ts, &u));
925     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
926     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
927     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
928     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
929     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
930     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
931     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
932     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
933   }
934   PetscFunctionReturn(PETSC_SUCCESS);
935 }
936 
937 static PetscErrorCode MigrateParticles(TS ts, PetscInt nv, Vec vecsin[], Vec vecsout[], void *ctx)
938 {
939   DM sw;
940 
941   PetscFunctionBeginUser;
942   PetscCall(TSGetDM(ts, &sw));
943   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
944   PetscCall(CreateSolution(ts));
945   PetscCall(SetProblem(ts));
946   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
947   PetscFunctionReturn(PETSC_SUCCESS);
948 }
949 
950 int main(int argc, char **argv)
951 {
952   DM     dm, sw;
953   TS     ts;
954   Vec    u;
955   AppCtx user;
956 
957   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
958   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
959   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
960   PetscCall(CreatePoisson(dm, &user));
961   PetscCall(CreateSwarm(dm, &user, &sw));
962   PetscCall(DMSetApplicationContext(sw, &user));
963 
964   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
965   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
966   PetscCall(TSSetDM(ts, sw));
967   PetscCall(TSSetMaxTime(ts, 0.1));
968   PetscCall(TSSetTimeStep(ts, 0.00001));
969   PetscCall(TSSetMaxSteps(ts, 100));
970   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
971   PetscCall(TSMonitorSet(ts, EnergyMonitor, &user, NULL));
972   PetscCall(TSSetFromOptions(ts));
973   PetscCall(DMSwarmVectorDefineField(sw, "velocity"));
974   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
975   PetscCall(TSSetComputeExactError(ts, ComputeError));
976   PetscCall(TSSetResize(ts, SetUpMigrateParticles, MigrateParticles, NULL));
977 
978   PetscCall(CreateSolution(ts));
979   PetscCall(TSGetSolution(ts, &u));
980   PetscCall(TSComputeInitialCondition(ts, u));
981   PetscCall(TSSolve(ts, NULL));
982 
983   PetscCall(SNESDestroy(&user.snes));
984   PetscCall(TSDestroy(&ts));
985   PetscCall(DMDestroy(&sw));
986   PetscCall(DMDestroy(&dm));
987   PetscCall(PetscFinalize());
988   return 0;
989 }
990 
991 /*TEST
992 
993    build:
994      requires: double !complex
995 
996    testset:
997      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
998      args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
999            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
1000            -ts_type basicsymplectic\
1001            -dm_view -output_step 50 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
1002      test:
1003        suffix: none_bsi_2d_1
1004        args: -ts_basicsymplectic_type 1 -em_type none -error
1005      test:
1006        suffix: none_bsi_2d_2
1007        args: -ts_basicsymplectic_type 2 -em_type none -error
1008      test:
1009        suffix: none_bsi_2d_3
1010        args: -ts_basicsymplectic_type 3 -em_type none -error
1011      test:
1012        suffix: none_bsi_2d_4
1013        args: -ts_basicsymplectic_type 4 -em_type none -error
1014      test:
1015        suffix: coulomb_bsi_2d_1
1016        args: -ts_basicsymplectic_type 1
1017      test:
1018        suffix: coulomb_bsi_2d_2
1019        args: -ts_basicsymplectic_type 2
1020      test:
1021        suffix: coulomb_bsi_2d_3
1022        args: -ts_basicsymplectic_type 3
1023      test:
1024        suffix: coulomb_bsi_2d_4
1025        args: -ts_basicsymplectic_type 4
1026 
1027    testset:
1028      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
1029      args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
1030            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
1031            -ts_type basicsymplectic\
1032            -em_type primal -em_pc_type svd\
1033            -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10\
1034            -petscspace_degree 2 -petscfe_default_quadrature_order 3 -sigma 1.0e-8 -timeScale 2.0e-14
1035      test:
1036        suffix: poisson_bsi_2d_1
1037        args: -ts_basicsymplectic_type 1
1038      test:
1039        suffix: poisson_bsi_2d_2
1040        args: -ts_basicsymplectic_type 2
1041      test:
1042        suffix: poisson_bsi_2d_3
1043        args: -ts_basicsymplectic_type 3
1044      test:
1045        suffix: poisson_bsi_2d_4
1046        args: -ts_basicsymplectic_type 4
1047 
1048    testset:
1049      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
1050      args: -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
1051            -ts_type theta -ts_theta_theta 0.5 -ts_convergence_estimate -convest_num_refine 2 \
1052              -mat_type baij -ksp_error_if_not_converged -em_pc_type svd\
1053            -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10\
1054            -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14
1055      test:
1056        suffix: im_2d_0
1057        args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5
1058 
1059    testset:
1060      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
1061      args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 10,10 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 -petscpartitioner_type simple \
1062            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV -dm_swarm_num_species 1\
1063            -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \
1064            -em_snes_type ksponly -em_pc_type svd -em_type primal -petscspace_degree 1\
1065            -dm_view -output_step 50\
1066            -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 1.0 -ts_max_steps 10
1067      test:
1068        suffix: bsi_2d_mesh_1
1069        args: -ts_basicsymplectic_type 4
1070      test:
1071        suffix: bsi_2d_mesh_1_par_2
1072        nsize: 2
1073        args: -ts_basicsymplectic_type 4
1074      test:
1075        suffix: bsi_2d_mesh_1_par_3
1076        nsize: 3
1077        args: -ts_basicsymplectic_type 4
1078      test:
1079        suffix: bsi_2d_mesh_1_par_4
1080        nsize: 4
1081        args: -ts_basicsymplectic_type 4 -dm_swarm_num_particles 0,0,2,0
1082 
1083    testset:
1084      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
1085      args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 10,10 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
1086            -dm_swarm_num_particles 10 -dm_swarm_coordinate_function circleMultipleX -dm_swarm_velocity_function circleMultipleV \
1087            -ts_convergence_estimate -convest_num_refine 2 \
1088              -em_pc_type lu\
1089            -dm_view -output_step 50 -error\
1090            -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
1091      test:
1092        suffix: bsi_2d_multiple_1
1093        args: -ts_type basicsymplectic -ts_basicsymplectic_type 1
1094      test:
1095        suffix: bsi_2d_multiple_2
1096        args: -ts_type basicsymplectic -ts_basicsymplectic_type 2
1097      test:
1098        suffix: bsi_2d_multiple_3
1099        args: -ts_type basicsymplectic -ts_basicsymplectic_type 3 -ts_dt 0.001
1100      test:
1101        suffix: im_2d_multiple_0
1102        args: -ts_type theta -ts_theta_theta 0.5 \
1103                -mat_type baij -ksp_error_if_not_converged -em_pc_type lu
1104 
1105    testset:
1106      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
1107      args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 2,2 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
1108            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
1109            -em_pc_type fieldsplit -ksp_rtol 1e-10 -em_ksp_type preonly -em_type mixed -em_ksp_error_if_not_converged\
1110            -dm_view -output_step 50 -error -dm_refine 0\
1111            -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
1112      test:
1113        suffix: bsi_4_rt_1
1114        args: -ts_type basicsymplectic -ts_basicsymplectic_type 4\
1115              -pc_fieldsplit_detect_saddle_point\
1116              -pc_type fieldsplit\
1117              -pc_fieldsplit_type schur\
1118              -pc_fieldsplit_schur_precondition full \
1119              -field_petscspace_degree 2\
1120              -field_petscfe_default_quadrature_order 1\
1121              -field_petscspace_type sum \
1122              -field_petscspace_variables 2 \
1123              -field_petscspace_components 2 \
1124              -field_petscspace_sum_spaces 2 \
1125              -field_petscspace_sum_concatenate true \
1126              -field_sumcomp_0_petscspace_variables 2 \
1127              -field_sumcomp_0_petscspace_type tensor \
1128              -field_sumcomp_0_petscspace_tensor_spaces 2 \
1129              -field_sumcomp_0_petscspace_tensor_uniform false \
1130              -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \
1131              -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \
1132              -field_sumcomp_1_petscspace_variables 2 \
1133              -field_sumcomp_1_petscspace_type tensor \
1134              -field_sumcomp_1_petscspace_tensor_spaces 2 \
1135              -field_sumcomp_1_petscspace_tensor_uniform false \
1136              -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \
1137              -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \
1138              -field_petscdualspace_form_degree -1 \
1139              -field_petscdualspace_order 1 \
1140              -field_petscdualspace_lagrange_trimmed true\
1141              -ksp_gmres_restart 500
1142 
1143 TEST*/
1144