xref: /petsc/src/dm/impls/swarm/tests/ex5.c (revision 503c0ea9b45bcfbcebbb1ea5341243bbc69f0bea)
1 static char help[] = "Vlasov example of particles orbiting around a central massive point.\n";
2 
3 #include <petscdmplex.h>
4 #include <petsc/private/dmpleximpl.h>  /* For norm */
5 #include <petsc/private/petscfeimpl.h> /* For CoordinatesRefToReal() */
6 #include <petscdmswarm.h>
7 #include <petscts.h>
8 
9 typedef struct {
10   PetscInt  dim;
11   PetscInt  particlesPerCell; /* The number of partices per cell */
12   PetscReal momentTol;        /* Tolerance for checking moment conservation */
13   PetscBool monitor;          /* Flag for using the TS monitor */
14   PetscBool error;            /* Flag for printing the error */
15   PetscInt  ostep;            /* print the energy at each ostep time steps */
16 } AppCtx;
17 
18 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
19 {
20   PetscErrorCode ierr;
21 
22   PetscFunctionBeginUser;
23   options->monitor          = PETSC_FALSE;
24   options->error            = PETSC_FALSE;
25   options->particlesPerCell = 1;
26   options->momentTol        = 100.0*PETSC_MACHINE_EPSILON;
27   options->ostep            = 100;
28 
29   ierr = PetscOptionsBegin(comm, "", "Vlasov Options", "DMPLEX");PetscCall(ierr);
30   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex4.c", options->ostep, &options->ostep, PETSC_NULL));
31   PetscCall(PetscOptionsBool("-monitor", "Flag to use the TS monitor", "ex4.c", options->monitor, &options->monitor, NULL));
32   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex4.c", options->error, &options->error, NULL));
33   PetscCall(PetscOptionsInt("-particles_per_cell", "Number of particles per cell", "ex4.c", options->particlesPerCell, &options->particlesPerCell, NULL));
34   ierr = PetscOptionsEnd();PetscCall(ierr);
35 
36   PetscFunctionReturn(0);
37 }
38 
39 static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user)
40 {
41   PetscFunctionBeginUser;
42   PetscCall(DMCreate(comm, dm));
43   PetscCall(DMSetType(*dm, DMPLEX));
44   PetscCall(DMSetFromOptions(*dm));
45   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
46   PetscCall(DMGetDimension(*dm, &user->dim));
47   PetscFunctionReturn(0);
48 }
49 
50 static PetscErrorCode SetInitialCoordinates(DM dmSw)
51 {
52   DM             dm;
53   AppCtx        *user;
54   PetscRandom    rnd;
55   PetscBool      simplex;
56   PetscReal     *centroid, *coords, *xi0, *v0, *J, *invJ, detJ;
57   PetscInt       dim, d, cStart, cEnd, c, Np, p;
58 
59   PetscFunctionBeginUser;
60   PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject) dmSw), &rnd));
61   PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
62   PetscCall(PetscRandomSetFromOptions(rnd));
63 
64   PetscCall(DMGetApplicationContext(dmSw, &user));
65   Np   = user->particlesPerCell;
66   PetscCall(DMSwarmGetCellDM(dmSw, &dm));
67   PetscCall(DMGetDimension(dm, &dim));
68   PetscCall(DMPlexIsSimplex(dm, &simplex));
69   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
70   PetscCall(PetscMalloc5(dim, &centroid, dim, &xi0, dim, &v0, dim*dim, &J, dim*dim, &invJ));
71   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
72   PetscCall(DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
73   for (c = cStart; c < cEnd; ++c) {
74     if (Np == 1) {
75       PetscCall(DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL));
76       for (d = 0; d < dim; ++d) coords[c*dim+d] = centroid[d];
77     } else {
78       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ)); /* affine */
79       for (p = 0; p < Np; ++p) {
80         const PetscInt n   = c*Np + p;
81         PetscReal      sum = 0.0, refcoords[3];
82 
83         for (d = 0; d < dim; ++d) {
84           PetscCall(PetscRandomGetValueReal(rnd, &refcoords[d]));
85           sum += refcoords[d];
86         }
87         if (simplex && sum > 0.0) for (d = 0; d < dim; ++d) refcoords[d] -= PetscSqrtReal(dim)*sum;
88         CoordinatesRefToReal(dim, dim, xi0, v0, J, refcoords, &coords[n*dim]);
89       }
90     }
91   }
92   PetscCall(DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
93   PetscCall(PetscFree5(centroid, xi0, v0, J, invJ));
94   PetscCall(PetscRandomDestroy(&rnd));
95   PetscFunctionReturn(0);
96 }
97 
98 static PetscErrorCode SetInitialConditions(DM dmSw, Vec u)
99 {
100   DM             dm;
101   AppCtx         *user;
102   PetscScalar    *initialConditions;
103   PetscInt       dim, cStart, cEnd, c, Np, p;
104 
105   PetscFunctionBeginUser;
106   PetscCall(DMGetApplicationContext(dmSw, &user));
107   Np   = user->particlesPerCell;
108   PetscCall(DMSwarmGetCellDM(dmSw, &dm));
109   PetscCall(DMGetDimension(dm, &dim));
110   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
111   PetscCall(VecGetArray(u, &initialConditions));
112   for (c = cStart; c < cEnd; ++c) {
113     for (p = 0; p < Np; ++p) {
114       const PetscInt n = c*Np + p;
115 
116       initialConditions[(n*2 + 0)*dim + 0] = n+1;
117       initialConditions[(n*2 + 0)*dim + 1] = 0;
118       initialConditions[(n*2 + 1)*dim + 0] = 0;
119       initialConditions[(n*2 + 1)*dim + 1] = PetscSqrtReal(1000./(n+1.));
120     }
121   }
122   PetscCall(VecRestoreArray(u, &initialConditions));
123   PetscFunctionReturn(0);
124 }
125 
126 static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user)
127 {
128   PetscInt      *cellid;
129   PetscInt       dim, cStart, cEnd, c, Np = user->particlesPerCell, p;
130 
131   PetscFunctionBeginUser;
132   PetscCall(DMGetDimension(dm, &dim));
133   PetscCall(DMCreate(PetscObjectComm((PetscObject) dm), sw));
134   PetscCall(DMSetType(*sw, DMSWARM));
135   PetscCall(DMSetDimension(*sw, dim));
136 
137   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
138   PetscCall(DMSwarmSetCellDM(*sw, dm));
139   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", 2*dim, PETSC_REAL));
140   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
141   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
142   PetscCall(DMSwarmSetLocalSizes(*sw, (cEnd - cStart) * Np, 0));
143   PetscCall(DMSetFromOptions(*sw));
144   PetscCall(DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid));
145   for (c = cStart; c < cEnd; ++c) {
146     for (p = 0; p < Np; ++p) {
147       const PetscInt n = c*Np + p;
148 
149       cellid[n] = c;
150     }
151   }
152   PetscCall(DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid));
153   PetscCall(PetscObjectSetName((PetscObject) *sw, "Particles"));
154   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
155   PetscFunctionReturn(0);
156 }
157 
158 /* Create particle RHS Functions for gravity with G = 1 for simplification */
159 static PetscErrorCode RHSFunction1(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
160 {
161   const PetscScalar *v;
162   PetscScalar       *xres;
163   PetscInt          Np, p, dim, d;
164 
165   PetscFunctionBeginUser;
166   /* The DM is not currently pushed down to the splits */
167   dim  = ((AppCtx *) ctx)->dim;
168   PetscCall(VecGetLocalSize(Xres, &Np));
169   Np  /= dim;
170   PetscCall(VecGetArray(Xres, &xres));
171   PetscCall(VecGetArrayRead(V, &v));
172   for (p = 0; p < Np; ++p) {
173      for (d = 0; d < dim; ++d) {
174        xres[p*dim+d] = v[p*dim+d];
175      }
176   }
177   PetscCall(VecRestoreArrayRead(V,& v));
178   PetscCall(VecRestoreArray(Xres, &xres));
179   PetscFunctionReturn(0);
180 }
181 
182 static PetscErrorCode RHSFunction2(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
183 {
184   const PetscScalar *x;
185   PetscScalar       *vres;
186   PetscInt          Np, p, dim, d;
187 
188   PetscFunctionBeginUser;
189   /* The DM is not currently pushed down to the splits */
190   dim  = ((AppCtx *) ctx)->dim;
191   PetscCall(VecGetLocalSize(Vres, &Np));
192   Np  /= dim;
193   PetscCall(VecGetArray(Vres, &vres));
194   PetscCall(VecGetArrayRead(X, &x));
195   for (p = 0; p < Np; ++p) {
196     const PetscScalar rsqr = DMPlex_NormD_Internal(dim, &x[p*dim]);
197 
198     for (d = 0; d < dim; ++d) {
199       vres[p*dim+d] = -(1000./(p+1.)) * x[p*dim+d]/PetscSqr(rsqr);
200     }
201   }
202   PetscCall(VecRestoreArray(Vres, &vres));
203   PetscCall(VecRestoreArrayRead(X, &x));
204   PetscFunctionReturn(0);
205 }
206 
207 static PetscErrorCode RHSFunctionParticles(TS ts, PetscReal t , Vec U, Vec R, void *ctx)
208 {
209   DM                dm;
210   const PetscScalar *u;
211   PetscScalar       *r;
212   PetscInt          Np, p, dim, d;
213 
214   PetscFunctionBeginUser;
215   PetscCall(TSGetDM(ts, &dm));
216   PetscCall(DMGetDimension(dm, &dim));
217   PetscCall(VecGetLocalSize(U, &Np));
218   Np  /= 2*dim;
219   PetscCall(VecGetArrayRead(U, &u));
220   PetscCall(VecGetArray(R, &r));
221   for (p = 0; p < Np; ++p) {
222     const PetscScalar rsqr = DMPlex_NormD_Internal(dim, &u[p*2*dim]);
223 
224     for (d = 0; d < dim; ++d) {
225         r[p*2*dim+d]   = u[p*2*dim+d+2];
226         r[p*2*dim+d+2] = -(1000./(1.+p)) * u[p*2*dim+d]/PetscSqr(rsqr);
227     }
228   }
229   PetscCall(VecRestoreArrayRead(U, &u));
230   PetscCall(VecRestoreArray(R, &r));
231   PetscFunctionReturn(0);
232 }
233 
234 static PetscErrorCode InitializeSolve(TS ts, Vec u)
235 {
236   DM             dm;
237   AppCtx        *user;
238 
239   PetscFunctionBeginUser;
240   PetscCall(TSGetDM(ts, &dm));
241   PetscCall(DMGetApplicationContext(dm, &user));
242   PetscCall(SetInitialCoordinates(dm));
243   PetscCall(SetInitialConditions(dm, u));
244   PetscFunctionReturn(0);
245 }
246 
247 static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
248 {
249   MPI_Comm          comm;
250   DM                sdm;
251   AppCtx            *user;
252   const PetscScalar *u, *coords;
253   PetscScalar       *e;
254   PetscReal         t;
255   PetscInt          dim, Np, p;
256 
257   PetscFunctionBeginUser;
258   PetscCall(PetscObjectGetComm((PetscObject) ts, &comm));
259   PetscCall(TSGetDM(ts, &sdm));
260   PetscCall(DMGetApplicationContext(sdm, &user));
261   PetscCall(DMGetDimension(sdm, &dim));
262   PetscCall(TSGetSolveTime(ts, &t));
263   PetscCall(VecGetArray(E, &e));
264   PetscCall(VecGetArrayRead(U, &u));
265   PetscCall(VecGetLocalSize(U, &Np));
266   Np  /= 2*dim;
267   PetscCall(DMSwarmGetField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
268   for (p = 0; p < Np; ++p) {
269     const PetscScalar *x     = &u[(p*2+0)*dim];
270     const PetscScalar *v     = &u[(p*2+1)*dim];
271     const PetscReal   x0    = p+1.;
272     const PetscReal   omega = PetscSqrtReal(1000./(p+1.))/x0;
273     const PetscReal   xe[3] = { x0*PetscCosReal(omega*t),       x0*PetscSinReal(omega*t),       0.0};
274     const PetscReal   ve[3] = {-x0*omega*PetscSinReal(omega*t), x0*omega*PetscCosReal(omega*t), 0.0};
275     PetscInt          d;
276 
277     for (d = 0; d < dim; ++d) {
278       e[(p*2+0)*dim+d] = x[d] - xe[d];
279       e[(p*2+1)*dim+d] = v[d] - ve[d];
280     }
281     if (user->error) PetscCall(PetscPrintf(comm, "t %.4g: p%D error [%.2g %.2g] sol [(%.6lf %.6lf) (%.6lf %.6lf)] exact [(%.6lf %.6lf) (%.6lf %.6lf)] energy/exact energy %g / %g\n", 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], 0.5*DMPlex_NormD_Internal(dim, v), (double) (0.5*(1000./(p+1)))));
282   }
283   PetscCall(DMSwarmRestoreField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
284   PetscCall(VecRestoreArrayRead(U, &u));
285   PetscCall(VecRestoreArray(E, &e));
286   PetscFunctionReturn(0);
287 }
288 
289 int main(int argc, char **argv)
290 {
291   TS                 ts;
292   TSConvergedReason  reason;
293   DM                 dm, sw;
294   Vec                u;
295   IS                 is1, is2;
296   PetscInt          *idx1, *idx2;
297   MPI_Comm           comm;
298   AppCtx             user;
299   const PetscScalar *endVals;
300   PetscReal          ftime   = .1;
301   PetscInt           locSize, dim, d, Np, p, steps;
302 
303   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
304   comm = PETSC_COMM_WORLD;
305   PetscCall(ProcessOptions(comm, &user));
306 
307   PetscCall(CreateMesh(comm, &dm, &user));
308   PetscCall(CreateParticles(dm, &sw, &user));
309   PetscCall(DMSetApplicationContext(sw, &user));
310 
311   PetscCall(TSCreate(comm, &ts));
312   PetscCall(TSSetType(ts, TSBASICSYMPLECTIC));
313   PetscCall(TSSetDM(ts, sw));
314   PetscCall(TSSetMaxTime(ts, ftime));
315   PetscCall(TSSetTimeStep(ts, 0.0001));
316   PetscCall(TSSetMaxSteps(ts, 10));
317   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
318   PetscCall(TSSetTime(ts, 0.0));
319   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionParticles, &user));
320 
321   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &u));
322   PetscCall(DMGetDimension(sw, &dim));
323   PetscCall(VecGetLocalSize(u, &locSize));
324   Np   = locSize/(2*dim);
325   PetscCall(PetscMalloc1(locSize/2, &idx1));
326   PetscCall(PetscMalloc1(locSize/2, &idx2));
327   for (p = 0; p < Np; ++p) {
328     for (d = 0; d < dim; ++d) {
329       idx1[p*dim+d] = (p*2+0)*dim + d;
330       idx2[p*dim+d] = (p*2+1)*dim + d;
331     }
332   }
333   PetscCall(ISCreateGeneral(comm, locSize/2, idx1, PETSC_OWN_POINTER, &is1));
334   PetscCall(ISCreateGeneral(comm, locSize/2, idx2, PETSC_OWN_POINTER, &is2));
335   PetscCall(TSRHSSplitSetIS(ts, "position", is1));
336   PetscCall(TSRHSSplitSetIS(ts, "momentum", is2));
337   PetscCall(ISDestroy(&is1));
338   PetscCall(ISDestroy(&is2));
339   PetscCall(TSRHSSplitSetRHSFunction(ts,"position",NULL,RHSFunction1,&user));
340   PetscCall(TSRHSSplitSetRHSFunction(ts,"momentum",NULL,RHSFunction2,&user));
341 
342   PetscCall(TSSetFromOptions(ts));
343   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
344   PetscCall(TSSetComputeExactError(ts, ComputeError));
345   PetscCall(TSComputeInitialCondition(ts, u));
346   PetscCall(VecViewFromOptions(u, NULL, "-init_view"));
347   PetscCall(TSSolve(ts, u));
348   PetscCall(TSGetSolveTime(ts, &ftime));
349   PetscCall(TSGetConvergedReason(ts, &reason));
350   PetscCall(TSGetStepNumber(ts, &steps));
351   PetscCall(PetscPrintf(comm,"%s at time %g after %D steps\n", TSConvergedReasons[reason], (double) ftime, steps));
352 
353   PetscCall(VecGetArrayRead(u, &endVals));
354   for (p = 0; p < Np; ++p) {
355     const PetscReal norm = DMPlex_NormD_Internal(dim, &endVals[(p*2 + 1)*dim]);
356     PetscCall(PetscPrintf(comm, "Particle %D initial Energy: %g  Final Energy: %g\n", p, (double) (0.5*(1000./(p+1))), (double) 0.5*PetscSqr(norm)));
357   }
358   PetscCall(VecRestoreArrayRead(u, &endVals));
359   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &u));
360   PetscCall(TSDestroy(&ts));
361   PetscCall(DMDestroy(&sw));
362   PetscCall(DMDestroy(&dm));
363   PetscCall(PetscFinalize());
364   return 0;
365 }
366 
367 /*TEST
368 
369    build:
370      requires: triangle !single !complex
371    test:
372      suffix: bsi1
373      args: -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
374            -ts_basicsymplectic_type 1 -ts_max_time 0.1 -ts_convergence_estimate -convest_num_refine 2 \
375            -dm_view -sw_view -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain -ts_monitor_sp_swarm_phase 0
376    test:
377      suffix: bsi2
378      args: -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
379            -ts_basicsymplectic_type 2 -ts_max_time 0.1 -ts_convergence_estimate -convest_num_refine 2 \
380            -dm_view -sw_view -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain -ts_monitor_sp_swarm_phase 0
381    test:
382      suffix: euler
383      args: -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \
384            -ts_type euler -ts_max_time 0.1 -ts_convergence_estimate -convest_num_refine 2 \
385            -dm_view -sw_view -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain -ts_monitor_sp_swarm_phase 0
386 
387 TEST*/
388