xref: /petsc/src/dm/impls/swarm/tests/ex5.c (revision 2e65eb737b5b07432530db55f6b2a145ebc548b2)
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 coodinates 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   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
533   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
534   PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm));
535   PetscCall(CreateSwarm(dm, &user, &sw));
536   PetscCall(DMSetApplicationContext(sw, &user));
537 
538   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
539   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
540   PetscCall(TSSetDM(ts, sw));
541   PetscCall(TSSetMaxTime(ts, 0.1));
542   PetscCall(TSSetTimeStep(ts, 0.00001));
543   PetscCall(TSSetMaxSteps(ts, 100));
544   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
545   PetscCall(TSMonitorSet(ts, EnergyMonitor, &user, NULL));
546   PetscCall(TSSetFromOptions(ts));
547   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
548   PetscCall(TSSetComputeExactError(ts, ComputeError));
549   PetscCall(TSSetPostStep(ts, MigrateParticles));
550 
551   PetscCall(CreateSolution(ts));
552   PetscCall(TSGetSolution(ts, &u));
553   PetscCall(TSComputeInitialCondition(ts, u));
554   PetscCall(TSSolve(ts, NULL));
555 
556   PetscCall(TSDestroy(&ts));
557   PetscCall(DMDestroy(&sw));
558   PetscCall(DMDestroy(&dm));
559   PetscCall(PetscFinalize());
560   return 0;
561 }
562 
563 /*TEST
564 
565    build:
566      requires: double !complex
567 
568    testset:
569      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
570      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 \
571            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
572            -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \
573            -dm_view -output_step 50 -error
574      test:
575        suffix: bsi_2d_1
576        args: -ts_basicsymplectic_type 1
577      test:
578        suffix: bsi_2d_2
579        args: -ts_basicsymplectic_type 2
580      test:
581        suffix: bsi_2d_3
582        args: -ts_basicsymplectic_type 3
583      test:
584        suffix: bsi_2d_4
585        args: -ts_basicsymplectic_type 4 -ts_dt 0.0001
586 
587    testset:
588      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
589      args: -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
590            -ts_type theta -ts_theta_theta 0.5 -ts_convergence_estimate -convest_num_refine 2 \
591              -mat_type baij -ksp_error_if_not_converged -pc_type lu \
592            -dm_view -output_step 50 -error
593      test:
594        suffix: im_2d_0
595        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
596 
597    testset:
598      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
599      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 \
600            -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \
601            -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \
602            -dm_view -output_step 50
603      test:
604        suffix: bsi_2d_mesh_1
605        args: -ts_basicsymplectic_type 1 -error
606      test:
607        suffix: bsi_2d_mesh_1_par_2
608        nsize: 2
609        args: -ts_basicsymplectic_type 1
610      test:
611        suffix: bsi_2d_mesh_1_par_3
612        nsize: 3
613        args: -ts_basicsymplectic_type 1
614      test:
615        suffix: bsi_2d_mesh_1_par_4
616        nsize: 4
617        args: -ts_basicsymplectic_type 1 -dm_swarm_num_particles 0,0,2,0
618 
619    testset:
620      requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT)
621      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 \
622            -dm_swarm_num_particles 10 -dm_swarm_coordinate_function circleMultipleX -dm_swarm_velocity_function circleMultipleV \
623            -ts_convergence_estimate -convest_num_refine 2 \
624            -dm_view -output_step 50 -error
625      test:
626        suffix: bsi_2d_multiple_1
627        args: -ts_type basicsymplectic -ts_basicsymplectic_type 1
628      test:
629        suffix: bsi_2d_multiple_2
630        args: -ts_type basicsymplectic -ts_basicsymplectic_type 2
631      test:
632        suffix: bsi_2d_multiple_3
633        args: -ts_type basicsymplectic -ts_basicsymplectic_type 3 -ts_dt 0.001
634      test:
635        suffix: im_2d_multiple_0
636        args: -ts_type theta -ts_theta_theta 0.5 \
637                -mat_type baij -ksp_error_if_not_converged -pc_type lu
638 
639 TEST*/
640