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