xref: /petsc/src/dm/impls/swarm/tests/ex5.c (revision df4cd43f92eaa320656440c40edb1046daee8f75)
1 static char help[] = "Vlasov 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 
18 PETSC_EXTERN PetscErrorCode circleSingleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
19 PETSC_EXTERN PetscErrorCode circleSingleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
20 PETSC_EXTERN PetscErrorCode circleMultipleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
21 PETSC_EXTERN PetscErrorCode circleMultipleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
22 
23 typedef struct {
24   PetscBool error; /* Flag for printing the error */
25   PetscInt  ostep; /* print the energy at each ostep time steps */
26 } AppCtx;
27 
28 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
29 {
30   PetscFunctionBeginUser;
31   options->error = PETSC_FALSE;
32   options->ostep = 100;
33 
34   PetscOptionsBegin(comm, "", "Central Orbit Options", "DMSWARM");
35   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex5.c", options->error, &options->error, NULL));
36   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex5.c", options->ostep, &options->ostep, NULL));
37   PetscOptionsEnd();
38   PetscFunctionReturn(PETSC_SUCCESS);
39 }
40 
41 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
42 {
43   PetscFunctionBeginUser;
44   PetscCall(DMCreate(comm, dm));
45   PetscCall(DMSetType(*dm, DMPLEX));
46   PetscCall(DMSetFromOptions(*dm));
47   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
48   PetscFunctionReturn(PETSC_SUCCESS);
49 }
50 
51 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw)
52 {
53   PetscReal v0[1] = {1.};
54   PetscInt  dim;
55 
56   PetscFunctionBeginUser;
57   PetscCall(DMGetDimension(dm, &dim));
58   PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw));
59   PetscCall(DMSetType(*sw, DMSWARM));
60   PetscCall(DMSetDimension(*sw, dim));
61   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
62   PetscCall(DMSwarmSetCellDM(*sw, dm));
63   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR));
64   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL));
65   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT));
66   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initCoordinates", dim, PETSC_REAL));
67   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initVelocity", dim, PETSC_REAL));
68   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
69   PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw));
70   PetscCall(DMSwarmInitializeCoordinates(*sw));
71   PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0));
72   PetscCall(DMSetFromOptions(*sw));
73   PetscCall(DMSetApplicationContext(*sw, user));
74   PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles"));
75   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
76   {
77     Vec gc, gc0, gv, gv0;
78 
79     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
80     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0));
81     PetscCall(VecCopy(gc, gc0));
82     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc));
83     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0));
84     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv));
85     PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0));
86     PetscCall(VecCopy(gv, gv0));
87     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv));
88     PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0));
89   }
90   PetscFunctionReturn(PETSC_SUCCESS);
91 }
92 
93 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
94 {
95   DM                 sw;
96   const PetscReal   *coords, *vel;
97   const PetscScalar *u;
98   PetscScalar       *g;
99   PetscInt           dim, d, Np, p;
100 
101   PetscFunctionBeginUser;
102   PetscCall(TSGetDM(ts, &sw));
103   PetscCall(DMGetDimension(sw, &dim));
104   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
105   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
106   PetscCall(VecGetLocalSize(U, &Np));
107   PetscCall(VecGetArrayRead(U, &u));
108   PetscCall(VecGetArray(G, &g));
109   Np /= 2 * dim;
110   for (p = 0; p < Np; ++p) {
111     const PetscReal x0    = coords[p * dim + 0];
112     const PetscReal vy0   = vel[p * dim + 1];
113     const PetscReal omega = vy0 / x0;
114 
115     for (d = 0; d < dim; ++d) {
116       g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d];
117       g[(p * 2 + 1) * dim + d] = -PetscSqr(omega) * u[(p * 2 + 0) * dim + d];
118     }
119   }
120   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
121   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
122   PetscCall(VecRestoreArrayRead(U, &u));
123   PetscCall(VecRestoreArray(G, &g));
124   PetscFunctionReturn(PETSC_SUCCESS);
125 }
126 
127 /* J_{ij} = dF_i/dx_j
128    J_p = (  0   1)
129          (-w^2  0)
130 */
131 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx)
132 {
133   DM               sw;
134   const PetscReal *coords, *vel;
135   PetscInt         dim, d, Np, p, rStart;
136 
137   PetscFunctionBeginUser;
138   PetscCall(TSGetDM(ts, &sw));
139   PetscCall(DMGetDimension(sw, &dim));
140   PetscCall(VecGetLocalSize(U, &Np));
141   PetscCall(MatGetOwnershipRange(J, &rStart, NULL));
142   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
143   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
144   Np /= 2 * dim;
145   for (p = 0; p < Np; ++p) {
146     const PetscReal x0      = coords[p * dim + 0];
147     const PetscReal vy0     = vel[p * dim + 1];
148     const PetscReal omega   = vy0 / x0;
149     PetscScalar     vals[4] = {0., 1., -PetscSqr(omega), 0.};
150 
151     for (d = 0; d < dim; ++d) {
152       const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart};
153       PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
154     }
155   }
156   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
157   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
158   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
159   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
160   PetscFunctionReturn(PETSC_SUCCESS);
161 }
162 
163 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
164 {
165   DM                 sw;
166   const PetscScalar *v;
167   PetscScalar       *xres;
168   PetscInt           Np, p, dim, d;
169 
170   PetscFunctionBeginUser;
171   PetscCall(TSGetDM(ts, &sw));
172   PetscCall(DMGetDimension(sw, &dim));
173   PetscCall(VecGetLocalSize(Xres, &Np));
174   PetscCall(VecGetArrayRead(V, &v));
175   PetscCall(VecGetArray(Xres, &xres));
176   Np /= dim;
177   for (p = 0; p < Np; ++p)
178     for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d];
179   PetscCall(VecRestoreArrayRead(V, &v));
180   PetscCall(VecRestoreArray(Xres, &xres));
181   PetscFunctionReturn(PETSC_SUCCESS);
182 }
183 
184 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
185 {
186   DM                 sw;
187   const PetscScalar *x;
188   const PetscReal   *coords, *vel;
189   PetscScalar       *vres;
190   PetscInt           Np, p, dim, d;
191 
192   PetscFunctionBeginUser;
193   PetscCall(TSGetDM(ts, &sw));
194   PetscCall(DMGetDimension(sw, &dim));
195   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
196   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
197   PetscCall(VecGetLocalSize(Vres, &Np));
198   PetscCall(VecGetArrayRead(X, &x));
199   PetscCall(VecGetArray(Vres, &vres));
200   PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2");
201   Np /= dim;
202   for (p = 0; p < Np; ++p) {
203     const PetscReal x0    = coords[p * dim + 0];
204     const PetscReal vy0   = vel[p * dim + 1];
205     const PetscReal omega = vy0 / x0;
206 
207     for (d = 0; d < dim; ++d) vres[p * dim + d] = -PetscSqr(omega) * x[p * dim + d];
208   }
209   PetscCall(VecRestoreArrayRead(X, &x));
210   PetscCall(VecRestoreArray(Vres, &vres));
211   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
212   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
213   PetscFunctionReturn(PETSC_SUCCESS);
214 }
215 
216 static PetscErrorCode CreateSolution(TS ts)
217 {
218   DM       sw;
219   Vec      u;
220   PetscInt dim, Np;
221 
222   PetscFunctionBegin;
223   PetscCall(TSGetDM(ts, &sw));
224   PetscCall(DMGetDimension(sw, &dim));
225   PetscCall(DMSwarmGetLocalSize(sw, &Np));
226   PetscCall(VecCreate(PETSC_COMM_WORLD, &u));
227   PetscCall(VecSetBlockSize(u, dim));
228   PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE));
229   PetscCall(VecSetUp(u));
230   PetscCall(TSSetSolution(ts, u));
231   PetscCall(VecDestroy(&u));
232   PetscFunctionReturn(PETSC_SUCCESS);
233 }
234 
235 static PetscErrorCode SetProblem(TS ts)
236 {
237   AppCtx *user;
238   DM      sw;
239 
240   PetscFunctionBegin;
241   PetscCall(TSGetDM(ts, &sw));
242   PetscCall(DMGetApplicationContext(sw, (void **)&user));
243   // Define unified system for (X, V)
244   {
245     Mat      J;
246     PetscInt dim, Np;
247 
248     PetscCall(DMGetDimension(sw, &dim));
249     PetscCall(DMSwarmGetLocalSize(sw, &Np));
250     PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
251     PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE));
252     PetscCall(MatSetBlockSize(J, 2 * dim));
253     PetscCall(MatSetFromOptions(J));
254     PetscCall(MatSetUp(J));
255     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user));
256     PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user));
257     PetscCall(MatDestroy(&J));
258   }
259   // Define split system for X and V
260   {
261     Vec             u;
262     IS              isx, isv, istmp;
263     const PetscInt *idx;
264     PetscInt        dim, Np, rstart;
265 
266     PetscCall(TSGetSolution(ts, &u));
267     PetscCall(DMGetDimension(sw, &dim));
268     PetscCall(DMSwarmGetLocalSize(sw, &Np));
269     PetscCall(VecGetOwnershipRange(u, &rstart, NULL));
270     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp));
271     PetscCall(ISGetIndices(istmp, &idx));
272     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx));
273     PetscCall(ISRestoreIndices(istmp, &idx));
274     PetscCall(ISDestroy(&istmp));
275     PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp));
276     PetscCall(ISGetIndices(istmp, &idx));
277     PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv));
278     PetscCall(ISRestoreIndices(istmp, &idx));
279     PetscCall(ISDestroy(&istmp));
280     PetscCall(TSRHSSplitSetIS(ts, "position", isx));
281     PetscCall(TSRHSSplitSetIS(ts, "momentum", isv));
282     PetscCall(ISDestroy(&isx));
283     PetscCall(ISDestroy(&isv));
284     PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user));
285     PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user));
286   }
287   // Define symplectic formulation U_t = S . G, where G = grad F
288   {
289     //PetscCall(TSDiscGradSetFormulation(ts, RHSJacobianS, RHSObjectiveF, RHSFunctionG, user));
290   }
291   PetscFunctionReturn(PETSC_SUCCESS);
292 }
293 
294 static PetscErrorCode DMSwarmTSRedistribute(TS ts)
295 {
296   DM        sw;
297   Vec       u;
298   PetscReal t, maxt, dt;
299   PetscInt  n, maxn;
300 
301   PetscFunctionBegin;
302   PetscCall(TSGetDM(ts, &sw));
303   PetscCall(TSGetTime(ts, &t));
304   PetscCall(TSGetMaxTime(ts, &maxt));
305   PetscCall(TSGetTimeStep(ts, &dt));
306   PetscCall(TSGetStepNumber(ts, &n));
307   PetscCall(TSGetMaxSteps(ts, &maxn));
308 
309   PetscCall(TSReset(ts));
310   PetscCall(TSSetDM(ts, sw));
311   /* TODO Check whether TS was set from options */
312   PetscCall(TSSetFromOptions(ts));
313   PetscCall(TSSetTime(ts, t));
314   PetscCall(TSSetMaxTime(ts, maxt));
315   PetscCall(TSSetTimeStep(ts, dt));
316   PetscCall(TSSetStepNumber(ts, n));
317   PetscCall(TSSetMaxSteps(ts, maxn));
318 
319   PetscCall(CreateSolution(ts));
320   PetscCall(SetProblem(ts));
321   PetscCall(TSGetSolution(ts, &u));
322   PetscFunctionReturn(PETSC_SUCCESS);
323 }
324 
325 PetscErrorCode circleSingleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
326 {
327   x[0] = p + 1.;
328   x[1] = 0.;
329   return PETSC_SUCCESS;
330 }
331 
332 PetscErrorCode circleSingleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx)
333 {
334   v[0] = 0.;
335   v[1] = PetscSqrtReal(1000. / (p + 1.));
336   return PETSC_SUCCESS;
337 }
338 
339 /* Put 5 particles into each circle */
340 PetscErrorCode circleMultipleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx)
341 {
342   const PetscInt  n   = 5;
343   const PetscReal r0  = (p / n) + 1.;
344   const PetscReal th0 = (2. * PETSC_PI * (p % n)) / n;
345 
346   x[0] = r0 * PetscCosReal(th0);
347   x[1] = r0 * PetscSinReal(th0);
348   return PETSC_SUCCESS;
349 }
350 
351 /* Put 5 particles into each circle */
352 PetscErrorCode circleMultipleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx)
353 {
354   const PetscInt  n     = 5;
355   const PetscReal r0    = (p / n) + 1.;
356   const PetscReal th0   = (2. * PETSC_PI * (p % n)) / n;
357   const PetscReal omega = PetscSqrtReal(1000. / r0) / r0;
358 
359   v[0] = -r0 * omega * PetscSinReal(th0);
360   v[1] = r0 * omega * PetscCosReal(th0);
361   return PETSC_SUCCESS;
362 }
363 
364 /*
365   InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values.
366 
367   Input Parameters:
368 + ts         - The TS
369 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem
370 
371   Output Parameters:
372 . u - The initialized solution vector
373 
374   Level: advanced
375 
376 .seealso: InitializeSolve()
377 */
378 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial)
379 {
380   DM      sw;
381   Vec     u, gc, gv, gc0, gv0;
382   IS      isx, isv;
383   AppCtx *user;
384 
385   PetscFunctionBeginUser;
386   PetscCall(TSGetDM(ts, &sw));
387   PetscCall(DMGetApplicationContext(sw, &user));
388   if (useInitial) {
389     PetscReal v0[1] = {1.};
390 
391     PetscCall(DMSwarmInitializeCoordinates(sw));
392     PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0));
393     PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
394     PetscCall(DMSwarmTSRedistribute(ts));
395   }
396   PetscCall(TSGetSolution(ts, &u));
397   PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
398   PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
399   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
400   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0));
401   if (useInitial) PetscCall(VecCopy(gc, gc0));
402   PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc));
403   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
404   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0));
405   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
406   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0));
407   if (useInitial) PetscCall(VecCopy(gv, gv0));
408   PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv));
409   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
410   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0));
411   PetscFunctionReturn(PETSC_SUCCESS);
412 }
413 
414 static PetscErrorCode InitializeSolve(TS ts, Vec u)
415 {
416   PetscFunctionBegin;
417   PetscCall(TSSetSolution(ts, u));
418   PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE));
419   PetscFunctionReturn(PETSC_SUCCESS);
420 }
421 
422 static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
423 {
424   MPI_Comm           comm;
425   DM                 sw;
426   AppCtx            *user;
427   const PetscScalar *u;
428   const PetscReal   *coords, *vel;
429   PetscScalar       *e;
430   PetscReal          t;
431   PetscInt           dim, Np, p;
432 
433   PetscFunctionBeginUser;
434   PetscCall(PetscObjectGetComm((PetscObject)ts, &comm));
435   PetscCall(TSGetDM(ts, &sw));
436   PetscCall(DMGetApplicationContext(sw, &user));
437   PetscCall(DMGetDimension(sw, &dim));
438   PetscCall(TSGetSolveTime(ts, &t));
439   PetscCall(VecGetArray(E, &e));
440   PetscCall(VecGetArrayRead(U, &u));
441   PetscCall(VecGetLocalSize(U, &Np));
442   PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
443   PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel));
444   Np /= 2 * dim;
445   for (p = 0; p < Np; ++p) {
446     /* TODO generalize initial conditions and project into plane instead of assuming x-y */
447     const PetscReal    r0    = DMPlex_NormD_Internal(dim, &coords[p * dim]);
448     const PetscReal    th0   = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]);
449     const PetscReal    v0    = DMPlex_NormD_Internal(dim, &vel[p * dim]);
450     const PetscReal    omega = v0 / r0;
451     const PetscReal    ct    = PetscCosReal(omega * t + th0);
452     const PetscReal    st    = PetscSinReal(omega * t + th0);
453     const PetscScalar *x     = &u[(p * 2 + 0) * dim];
454     const PetscScalar *v     = &u[(p * 2 + 1) * dim];
455     const PetscReal    xe[3] = {r0 * ct, r0 * st, 0.0};
456     const PetscReal    ve[3] = {-v0 * st, v0 * ct, 0.0};
457     PetscInt           d;
458 
459     for (d = 0; d < dim; ++d) {
460       e[(p * 2 + 0) * dim + d] = x[d] - xe[d];
461       e[(p * 2 + 1) * dim + d] = v[d] - ve[d];
462     }
463     if (user->error) {
464       const PetscReal en   = 0.5 * DMPlex_DotRealD_Internal(dim, v, v);
465       const PetscReal exen = 0.5 * PetscSqr(v0);
466       PetscCall(PetscPrintf(comm, "t %.4g: p%" PetscInt_FMT " error [%.2g %.2g] 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)));
467     }
468   }
469   PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords));
470   PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel));
471   PetscCall(VecRestoreArrayRead(U, &u));
472   PetscCall(VecRestoreArray(E, &e));
473   PetscFunctionReturn(PETSC_SUCCESS);
474 }
475 
476 static PetscErrorCode EnergyMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
477 {
478   const PetscInt     ostep = ((AppCtx *)ctx)->ostep;
479   DM                 sw;
480   const PetscScalar *u;
481   PetscInt           dim, Np, p;
482 
483   PetscFunctionBeginUser;
484   if (step % ostep == 0) {
485     PetscCall(TSGetDM(ts, &sw));
486     PetscCall(DMGetDimension(sw, &dim));
487     PetscCall(VecGetArrayRead(U, &u));
488     PetscCall(VecGetLocalSize(U, &Np));
489     Np /= 2 * dim;
490     if (!step) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)ts), "Time     Step Part     Energy\n"));
491     for (p = 0; p < Np; ++p) {
492       const PetscReal v2 = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]);
493 
494       PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " %4" PetscInt_FMT " %10.4lf\n", (double)t, step, p, (double)(0.5 * v2)));
495     }
496     PetscCall(PetscSynchronizedFlush(PetscObjectComm((PetscObject)ts), NULL));
497     PetscCall(VecRestoreArrayRead(U, &u));
498   }
499   PetscFunctionReturn(PETSC_SUCCESS);
500 }
501 
502 static PetscErrorCode MigrateParticles(TS ts)
503 {
504   DM sw;
505 
506   PetscFunctionBeginUser;
507   PetscCall(TSGetDM(ts, &sw));
508   PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre"));
509   {
510     Vec u, gc, gv;
511     IS  isx, isv;
512 
513     PetscCall(TSGetSolution(ts, &u));
514     PetscCall(TSRHSSplitGetIS(ts, "position", &isx));
515     PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv));
516     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
517     PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc));
518     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc));
519     PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv));
520     PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv));
521     PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv));
522   }
523   PetscCall(DMSwarmMigrate(sw, PETSC_TRUE));
524   PetscCall(DMSwarmTSRedistribute(ts));
525   PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE));
526   PetscFunctionReturn(PETSC_SUCCESS);
527 }
528 
529 int main(int argc, char **argv)
530 {
531   DM     dm, sw;
532   TS     ts;
533   Vec    u;
534   AppCtx user;
535 
536   PetscFunctionBeginUser;
537   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
538   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
539   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
540   PetscCall(CreateSwarm(dm, &user, &sw));
541   PetscCall(DMSetApplicationContext(sw, &user));
542 
543   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
544   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
545   PetscCall(TSSetDM(ts, sw));
546   PetscCall(TSSetMaxTime(ts, 0.1));
547   PetscCall(TSSetTimeStep(ts, 0.00001));
548   PetscCall(TSSetMaxSteps(ts, 100));
549   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
550   PetscCall(TSMonitorSet(ts, EnergyMonitor, &user, NULL));
551   PetscCall(TSSetFromOptions(ts));
552   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
553   PetscCall(TSSetComputeExactError(ts, ComputeError));
554   PetscCall(TSSetPostStep(ts, MigrateParticles));
555 
556   PetscCall(CreateSolution(ts));
557   PetscCall(TSGetSolution(ts, &u));
558   PetscCall(TSComputeInitialCondition(ts, u));
559   PetscCall(TSSolve(ts, NULL));
560 
561   PetscCall(TSDestroy(&ts));
562   PetscCall(DMDestroy(&sw));
563   PetscCall(DMDestroy(&dm));
564   PetscCall(PetscFinalize());
565   return 0;
566 }
567 
568 /*TEST
569 
570    build:
571      requires: double !complex
572 
573    testset:
574      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
575      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 \
576            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
577            -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \
578            -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
579      test:
580        suffix: bsi_2d_1
581        args: -ts_basicsymplectic_type 1
582      test:
583        suffix: bsi_2d_2
584        args: -ts_basicsymplectic_type 2
585      test:
586        suffix: bsi_2d_3
587        args: -ts_basicsymplectic_type 3
588      test:
589        suffix: bsi_2d_4
590        args: -ts_basicsymplectic_type 4 -ts_dt 0.0001
591 
592    testset:
593      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
594      args: -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
595            -ts_type theta -ts_theta_theta 0.5 -ts_convergence_estimate -convest_num_refine 2 \
596              -mat_type baij -ksp_error_if_not_converged -pc_type lu \
597            -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
598      test:
599        suffix: im_2d_0
600        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
601 
602    testset:
603      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
604      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 \
605            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
606            -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \
607            -dm_view -output_step 50 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10
608      test:
609        suffix: bsi_2d_mesh_1
610        args: -ts_basicsymplectic_type 1 -error
611      test:
612        suffix: bsi_2d_mesh_1_par_2
613        nsize: 2
614        args: -ts_basicsymplectic_type 1
615      test:
616        suffix: bsi_2d_mesh_1_par_3
617        nsize: 3
618        args: -ts_basicsymplectic_type 1
619      test:
620        suffix: bsi_2d_mesh_1_par_4
621        nsize: 4
622        args: -ts_basicsymplectic_type 1 -dm_swarm_num_particles 0,0,2,0
623 
624    testset:
625      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
626      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 \
627            -dm_swarm_num_particles 10 -dm_swarm_coordinate_function circleMultipleX -dm_swarm_velocity_function circleMultipleV \
628            -ts_convergence_estimate -convest_num_refine 2 \
629            -dm_view -output_step 50 -error
630      test:
631        suffix: bsi_2d_multiple_1
632        args: -ts_type basicsymplectic -ts_basicsymplectic_type 1
633      test:
634        suffix: bsi_2d_multiple_2
635        args: -ts_type basicsymplectic -ts_basicsymplectic_type 2
636      test:
637        suffix: bsi_2d_multiple_3
638        args: -ts_type basicsymplectic -ts_basicsymplectic_type 3 -ts_dt 0.001
639      test:
640        suffix: im_2d_multiple_0
641        args: -ts_type theta -ts_theta_theta 0.5 \
642                -mat_type baij -ksp_error_if_not_converged -pc_type lu
643 
644 TEST*/
645