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