xref: /petsc/src/ts/tutorials/hamiltonian/ex2.c (revision 5fedec97950c19de564efaecd0f125b1a6cb2b20)
1 static char help[] = "Two stream instability from Birdsal and Langdon with DMSwarm and TS basic symplectic integrators\n";
2 
3 #include <petscdmplex.h>
4 #include <petscfe.h>
5 #include <petscdmswarm.h>
6 #include <petscds.h>
7 #include <petscksp.h>
8 #include <petsc/private/petscfeimpl.h>
9 #include <petsc/private/tsimpl.h>
10 #include <petscts.h>
11 #include <petscmath.h>
12 
13 typedef struct {
14   PetscInt       dim;                              /* The topological mesh dimension */
15   PetscBool      simplex;                          /* Flag for simplices or tensor cells */
16   PetscBool      bdm;                              /* Flag for mixed form poisson */
17   PetscBool      monitor;                          /* Flag for use of the TS monitor */
18   PetscBool      uniform;                          /* Flag to uniformly space particles in x */
19   char           meshFilename[PETSC_MAX_PATH_LEN]; /* Name of the mesh filename if any */
20   PetscReal      sigma;                            /* Linear charge per box length */
21   PetscReal      timeScale;                        /* Nondimensionalizing time scaling */
22   PetscInt       particlesPerCell;                 /* The number of partices per cell */
23   PetscReal      particleRelDx;                    /* Relative particle position perturbation compared to average cell diameter h */
24   PetscInt       k;                                /* Mode number for test function */
25   PetscReal      momentTol;                        /* Tolerance for checking moment conservation */
26   SNES           snes;                             /* SNES object */
27   PetscInt       steps;                            /* TS iterations */
28   PetscReal      stepSize;                         /* Time stepper step size */
29   PetscErrorCode (*func)(PetscInt, PetscReal, const PetscReal [], PetscInt, PetscScalar *, void *);
30 } AppCtx;
31 
32 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
33 {
34   PetscErrorCode ierr;
35 
36   PetscFunctionBeginUser;
37   options->dim              = 2;
38   options->simplex          = PETSC_TRUE;
39   options->monitor          = PETSC_TRUE;
40   options->particlesPerCell = 1;
41   options->k                = 1;
42   options->particleRelDx    = 1.e-20;
43   options->momentTol        = 100.*PETSC_MACHINE_EPSILON;
44   options->sigma            = 1.;
45   options->timeScale        = 1.0e-6;
46   options->uniform          = PETSC_FALSE;
47   options->steps            = 1;
48   options->stepSize         = 0.01;
49   options->bdm              = PETSC_FALSE;
50 
51   ierr = PetscOptionsBegin(comm, "", "Two Stream options", "DMPLEX");CHKERRQ(ierr);
52   ierr = PetscStrcpy(options->meshFilename, "");CHKERRQ(ierr);
53   ierr = PetscOptionsInt("-dim", "The topological mesh dimension", "ex2.c", options->dim, &options->dim, NULL);CHKERRQ(ierr);
54   ierr = PetscOptionsInt("-steps", "TS steps to take", "ex2.c", options->steps, &options->steps, NULL);CHKERRQ(ierr);
55   ierr = PetscOptionsBool("-monitor", "To use the TS monitor or not", "ex2.c", options->monitor, &options->monitor, NULL);CHKERRQ(ierr);
56   ierr = PetscOptionsBool("-simplex", "The flag for simplices or tensor cells", "ex2.c", options->simplex, &options->simplex, NULL);CHKERRQ(ierr);
57   ierr = PetscOptionsBool("-uniform", "Uniform particle spacing", "ex2.c", options->uniform, &options->uniform, NULL);CHKERRQ(ierr);
58   ierr = PetscOptionsBool("-bdm", "Use H1 instead of C0", "ex2.c", options->bdm, &options->bdm, NULL);CHKERRQ(ierr);
59   ierr = PetscOptionsString("-mesh", "Name of the mesh filename if any", "ex2.c", options->meshFilename, options->meshFilename, PETSC_MAX_PATH_LEN, NULL);CHKERRQ(ierr);
60   ierr = PetscOptionsInt("-k", "Mode number of test", "ex5.c", options->k, &options->k, NULL);CHKERRQ(ierr);
61   ierr = PetscOptionsInt("-particlesPerCell", "Number of particles per cell", "ex2.c", options->particlesPerCell, &options->particlesPerCell, NULL);CHKERRQ(ierr);
62   ierr = PetscOptionsReal("-sigma","parameter","<1>",options->sigma,&options->sigma,PETSC_NULL);CHKERRQ(ierr);
63   ierr = PetscOptionsReal("-stepSize","parameter","<1e-2>",options->stepSize,&options->stepSize,PETSC_NULL);CHKERRQ(ierr);
64   ierr = PetscOptionsReal("-timeScale","parameter","<1>",options->timeScale,&options->timeScale,PETSC_NULL);CHKERRQ(ierr);
65   ierr = PetscOptionsReal("-particle_perturbation", "Relative perturbation of particles (0,1)", "ex2.c", options->particleRelDx, &options->particleRelDx, NULL);CHKERRQ(ierr);
66   ierr = PetscOptionsEnd();CHKERRQ(ierr);
67   PetscFunctionReturn(0);
68 }
69 
70 static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user)
71 {
72   PetscErrorCode ierr;
73 
74   PetscFunctionBeginUser;
75   ierr = DMCreate(comm, dm);CHKERRQ(ierr);
76   ierr = DMSetType(*dm, DMPLEX);CHKERRQ(ierr);
77   ierr = DMSetFromOptions(*dm);CHKERRQ(ierr);
78   ierr = DMViewFromOptions(*dm, NULL, "-dm_view");CHKERRQ(ierr);
79   PetscFunctionReturn(0);
80 }
81 
82 static void laplacian_f1(PetscInt dim, PetscInt Nf, PetscInt NfAux,
83                          const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
84                          const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
85                          PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[])
86 {
87   PetscInt d;
88   for (d = 0; d < dim; ++d) {f1[d] = u_x[d];}
89 }
90 
91 static void laplacian(PetscInt dim, PetscInt Nf, PetscInt NfAux,
92                       const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
93                       const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
94                       PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g3[])
95 {
96   PetscInt d;
97   for (d = 0; d < dim; ++d) {g3[d*dim+d] = 1.0;}
98 }
99 
100 static PetscErrorCode CreateFEM(DM dm, AppCtx *user)
101 {
102   PetscFE        fe;
103   PetscDS        ds;
104   DMPolytopeType ct;
105   PetscBool      simplex;
106   PetscInt       dim, cStart;
107   PetscErrorCode ierr;
108 
109   PetscFunctionBeginUser;
110   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
111   ierr = DMPlexGetHeightStratum(dm, 0, &cStart, NULL);CHKERRQ(ierr);
112   ierr = DMPlexGetCellType(dm, cStart, &ct);CHKERRQ(ierr);
113   simplex = DMPolytopeTypeGetNumVertices(ct) == DMPolytopeTypeGetDim(ct)+1 ? PETSC_TRUE : PETSC_FALSE;
114   ierr = PetscFECreateDefault(PetscObjectComm((PetscObject) dm), dim, 1, simplex, NULL, -1, &fe);CHKERRQ(ierr);
115   ierr = PetscObjectSetName((PetscObject) fe, "potential");CHKERRQ(ierr);
116   ierr = DMSetField(dm, 0, NULL, (PetscObject) fe);CHKERRQ(ierr);
117   ierr = DMCreateDS(dm);CHKERRQ(ierr);
118   ierr = PetscFEDestroy(&fe);CHKERRQ(ierr);
119   ierr = DMGetDS(dm, &ds);CHKERRQ(ierr);
120   ierr = PetscDSSetResidual(ds, 0, NULL, laplacian_f1);CHKERRQ(ierr);
121   ierr = PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian);CHKERRQ(ierr);
122   PetscFunctionReturn(0);
123 }
124 
125 /*
126   Initialize particle coordinates uniformly and with opposing velocities
127 */
128 static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user)
129 {
130   PetscRandom    rnd, rndp;
131   PetscReal      interval = user->particleRelDx;
132   PetscScalar    value, *vals;
133   PetscReal     *centroid, *coords, *xi0, *v0, *J, *invJ, detJ, *initialConditions, normalized_vel;
134   PetscInt      *cellid, cStart;
135   PetscInt       Ncell, Np = user->particlesPerCell, p, c, dim, d;
136   PetscErrorCode ierr;
137 
138   PetscFunctionBeginUser;
139   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
140   ierr = DMCreate(PetscObjectComm((PetscObject) dm), sw);CHKERRQ(ierr);
141   ierr = DMSetType(*sw, DMSWARM);CHKERRQ(ierr);
142   ierr = DMSetDimension(*sw, dim);CHKERRQ(ierr);
143   ierr = PetscRandomCreate(PetscObjectComm((PetscObject) dm), &rnd);CHKERRQ(ierr);
144   ierr = PetscRandomSetInterval(rnd, 0.0, 1.0);CHKERRQ(ierr);
145   ierr = PetscRandomSetFromOptions(rnd);CHKERRQ(ierr);
146   ierr = PetscRandomCreate(PetscObjectComm((PetscObject) dm), &rndp);CHKERRQ(ierr);
147   ierr = PetscRandomSetInterval(rndp, -interval, interval);CHKERRQ(ierr);
148   ierr = PetscRandomSetFromOptions(rndp);CHKERRQ(ierr);
149   ierr = DMSwarmSetType(*sw, DMSWARM_PIC);CHKERRQ(ierr);
150   ierr = DMSwarmSetCellDM(*sw, dm);CHKERRQ(ierr);
151   ierr = DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR);CHKERRQ(ierr);
152   ierr = DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", dim, PETSC_REAL);CHKERRQ(ierr);
153   ierr = DMSwarmFinalizeFieldRegister(*sw);CHKERRQ(ierr);
154   ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &Ncell);CHKERRQ(ierr);
155   ierr = DMSwarmSetLocalSizes(*sw, Ncell * Np, 0);CHKERRQ(ierr);
156   ierr = DMSetFromOptions(*sw);CHKERRQ(ierr);
157   ierr = DMSwarmGetField(*sw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
158   ierr = DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid);CHKERRQ(ierr);
159   ierr = DMSwarmGetField(*sw, "w_q", NULL, NULL, (void **) &vals);CHKERRQ(ierr);
160   ierr = DMSwarmGetField(*sw, "kinematics", NULL, NULL, (void **) &initialConditions);CHKERRQ(ierr);
161   ierr = PetscMalloc5(dim, &centroid, dim, &xi0, dim, &v0, dim*dim, &J, dim*dim, &invJ);CHKERRQ(ierr);
162   for (c = cStart; c < Ncell; c++) {
163     if (Np == 1) {
164       ierr = DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL);CHKERRQ(ierr);
165       cellid[c] = c;
166       for (d = 0; d < dim; ++d) coords[c*dim+d] = centroid[d];
167     } else {
168       for (d = 0; d < dim; ++d) xi0[d] = -1.0;
169       ierr = DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ);CHKERRQ(ierr); /* affine */
170       for (p = 0; p < Np; ++p) {
171         const PetscInt n   = c*Np + p;
172         PetscReal      refcoords[3], spacing;
173 
174         cellid[n] = c;
175         if (user->uniform) {
176           spacing = 2./Np;
177           ierr = PetscRandomGetValue(rnd, &value);
178           for (d=0; d<dim; ++d) refcoords[d] = d == 0 ? -1. + spacing/2. + p*spacing + value/100. : 0.;
179         }
180         else{
181           for (d = 0; d < dim; ++d) {ierr = PetscRandomGetValue(rnd, &value);CHKERRQ(ierr); refcoords[d] = d == 0 ? PetscRealPart(value) : 0. ;}
182         }
183         CoordinatesRefToReal(dim, dim, xi0, v0, J, refcoords, &coords[n*dim]);
184         /* constant particle weights */
185         for (d = 0; d < dim; ++d) vals[n] = user->sigma/Np;
186       }
187     }
188   }
189   ierr = PetscFree5(centroid, xi0, v0, J, invJ);CHKERRQ(ierr);
190   normalized_vel = 1.;
191   for (c = 0; c < Ncell; ++c) {
192     for (p = 0; p < Np; ++p) {
193       if (p%2 == 0) {
194         for (d = 0; d < dim; ++d) initialConditions[(c*Np + p)*dim + d] = d == 0 ? normalized_vel : 0.;
195       }
196       else {
197         for (d = 0; d < dim; ++d) initialConditions[(c*Np + p)*dim + d] = d == 0 ? -(normalized_vel) : 0.;
198       }
199     }
200   }
201   ierr = DMSwarmRestoreField(*sw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
202   ierr = DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid);CHKERRQ(ierr);
203   ierr = DMSwarmRestoreField(*sw, "w_q", NULL, NULL, (void **) &vals);CHKERRQ(ierr);
204   ierr = DMSwarmRestoreField(*sw, "kinematics", NULL, NULL, (void **) &initialConditions);CHKERRQ(ierr);
205   ierr = PetscRandomDestroy(&rnd);CHKERRQ(ierr);
206   ierr = PetscRandomDestroy(&rndp);CHKERRQ(ierr);
207   ierr = PetscObjectSetName((PetscObject) *sw, "Particles");CHKERRQ(ierr);
208   ierr = DMViewFromOptions(*sw, NULL, "-sw_view");CHKERRQ(ierr);
209   ierr = DMLocalizeCoordinates(*sw);CHKERRQ(ierr);
210   PetscFunctionReturn(0);
211 }
212 
213 /* Solve for particle position updates */
214 static PetscErrorCode RHSFunction1(TS ts,PetscReal t,Vec V,Vec Posres,void *ctx)
215 {
216   const PetscScalar *v;
217   PetscScalar       *posres;
218   PetscInt          Np, p, dim, d;
219   DM                dm;
220   PetscErrorCode    ierr;
221 
222   PetscFunctionBeginUser;
223   ierr = VecGetLocalSize(Posres, &Np);CHKERRQ(ierr);
224   ierr = VecGetArray(Posres,&posres);CHKERRQ(ierr);
225   ierr = VecGetArrayRead(V,&v);CHKERRQ(ierr);
226   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
227   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
228   Np  /= dim;
229   for (p = 0; p < Np; ++p) {
230      for (d = 0; d < dim; ++d) {
231        posres[p*dim+d] = v[p*dim+d];
232      }
233   }
234   ierr = VecRestoreArrayRead(V,&v);CHKERRQ(ierr);
235   ierr = VecRestoreArray(Posres,&posres);CHKERRQ(ierr);
236   PetscFunctionReturn(0);
237 
238 }
239 
240 /*
241   Solve for the gradient of the electric field and apply force to particles.
242  */
243 static PetscErrorCode RHSFunction2(TS ts,PetscReal t,Vec X,Vec Vres,void *ctx)
244 {
245  AppCtx             *user = (AppCtx *) ctx;
246   DM                 dm, plex;
247   PetscDS            prob;
248   PetscFE            fe;
249   Mat                M_p;
250   Vec                phi, locPhi, rho, f;
251   const PetscScalar *x;
252   PetscScalar       *vres;
253   PetscReal         *coords, phi_0;
254   PetscInt           dim, d, cStart, cEnd, cell, cdim;
255   PetscErrorCode     ierr;
256   PetscReal          m_e = 9.11e-31, q_e = 1.60e-19, epsi_0 = 8.85e-12;
257 
258   PetscFunctionBeginUser;
259   PetscObjectSetName((PetscObject) X, "rhsf2 position");
260   VecViewFromOptions(X, NULL, "-rhsf2_x_view");
261   ierr = VecGetArrayRead(X,&x);CHKERRQ(ierr);
262   ierr = VecGetArray(Vres,&vres);CHKERRQ(ierr);
263   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
264   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
265   ierr = SNESGetDM(user->snes, &plex);CHKERRQ(ierr);
266   ierr = DMGetCoordinateDim(plex, &cdim);CHKERRQ(ierr);
267   ierr = DMGetDS(plex, &prob);CHKERRQ(ierr);
268   ierr = PetscDSGetDiscretization(prob, 0, (PetscObject *) &fe);CHKERRQ(ierr);
269   ierr = DMGetGlobalVector(plex, &phi);CHKERRQ(ierr);
270   ierr = DMGetLocalVector(plex, &locPhi);CHKERRQ(ierr);
271   ierr = DMCreateMassMatrix(dm, plex, &M_p);CHKERRQ(ierr);
272   ierr = MatViewFromOptions(M_p, NULL, "-mp_view");
273   ierr = DMGetGlobalVector(plex, &rho);CHKERRQ(ierr);
274   ierr = DMSwarmCreateGlobalVectorFromField(dm, "w_q", &f);CHKERRQ(ierr);
275   ierr = PetscObjectSetName((PetscObject) f, "weights vector");
276   ierr = VecViewFromOptions(f, NULL, "-weights_view");
277   ierr = MatMultTranspose(M_p, f, rho);CHKERRQ(ierr);
278   ierr = DMSwarmDestroyGlobalVectorFromField(dm, "w_q", &f);CHKERRQ(ierr);
279   ierr = PetscObjectSetName((PetscObject) rho, "rho");CHKERRQ(ierr);
280   ierr = VecViewFromOptions(rho, NULL, "-poisson_rho_view");CHKERRQ(ierr);
281   /* Take nullspace out of rhs */
282   {
283     PetscScalar sum;
284     PetscInt    n;
285     phi_0 = (user->sigma*user->sigma*user->sigma)*(user->timeScale*user->timeScale)/(m_e*q_e*epsi_0);
286 
287     ierr = VecGetSize(rho, &n);CHKERRQ(ierr);
288     ierr = VecSum(rho, &sum);CHKERRQ(ierr);
289     ierr = VecShift(rho, -sum/n);CHKERRQ(ierr);
290 
291     ierr = VecSum(rho, &sum);CHKERRQ(ierr);
292     if (PetscAbsScalar(sum) > 1.0e-10) SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Charge should have no DC component %g", sum);
293     ierr = VecScale(rho, phi_0);CHKERRQ(ierr);
294   }
295   ierr = VecSet(phi, 0.0);CHKERRQ(ierr);
296   ierr = SNESSolve(user->snes, rho, phi);CHKERRQ(ierr);
297   ierr = VecViewFromOptions(phi, NULL, "-phi_view");CHKERRQ(ierr);
298   ierr = DMRestoreGlobalVector(plex, &rho);CHKERRQ(ierr);
299   ierr = MatDestroy(&M_p);CHKERRQ(ierr);
300   ierr = DMGlobalToLocalBegin(plex, phi, INSERT_VALUES, locPhi);CHKERRQ(ierr);
301   ierr = DMGlobalToLocalEnd(plex, phi, INSERT_VALUES, locPhi);CHKERRQ(ierr);
302   ierr = DMSwarmSortGetAccess(dm);CHKERRQ(ierr);
303   ierr = DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
304   ierr = DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd);CHKERRQ(ierr);
305   for (cell = cStart; cell < cEnd; ++cell) {
306     PetscTabulation tab;
307     PetscReal    v[3], J[9], invJ[9], detJ;
308     PetscScalar *ph       = PETSC_NULL;
309     PetscReal   *pcoord   = PETSC_NULL;
310     PetscReal   *refcoord = PETSC_NULL;
311     PetscInt    *points   = PETSC_NULL, Ncp, cp;
312     PetscScalar  gradPhi[3];
313 
314     ierr = DMPlexComputeCellGeometryFEM(plex, cell, NULL, v, J, invJ, &detJ);CHKERRQ(ierr);
315     ierr = DMSwarmSortGetPointsPerCell(dm, cell, &Ncp, &points);CHKERRQ(ierr);
316     ierr = DMGetWorkArray(dm, Ncp*cdim, MPIU_REAL, &pcoord);CHKERRQ(ierr);
317     ierr = DMGetWorkArray(dm, Ncp*cdim, MPIU_REAL, &refcoord);CHKERRQ(ierr);
318     for (cp = 0; cp < Ncp; ++cp) {
319       for (d = 0; d < cdim; ++d) {
320         pcoord[cp*cdim+d] = coords[points[cp]*cdim+d];
321       }
322     }
323     ierr = DMPlexCoordinatesToReference(plex, cell, Ncp, pcoord, refcoord);CHKERRQ(ierr);
324     PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab);CHKERRQ(ierr);
325     ierr = DMPlexVecGetClosure(plex, NULL, locPhi, cell, NULL, &ph);CHKERRQ(ierr);
326     for (cp = 0; cp < Ncp; ++cp) {
327       const PetscInt p = points[cp];
328       gradPhi[0] = 0.0;
329       gradPhi[1] = 0.0;
330       gradPhi[2] = 0.0;
331       const PetscReal *basisDer = tab->T[1];
332 
333       ierr = PetscFEFreeInterpolateGradient_Static(fe, basisDer, ph, cdim, invJ, NULL, cp, gradPhi);CHKERRQ(ierr);
334       for (d = 0; d < cdim; ++d) {
335         vres[p*cdim+d] = d == 0 ? gradPhi[d] : 0.;
336       }
337     }
338     ierr = DMPlexVecRestoreClosure(plex, NULL, locPhi, cell, NULL, &ph);CHKERRQ(ierr);
339     ierr = PetscTabulationDestroy(&tab);CHKERRQ(ierr);
340     ierr = DMRestoreWorkArray(dm, Ncp*cdim, MPIU_REAL, &pcoord);CHKERRQ(ierr);
341     ierr = DMRestoreWorkArray(dm, Ncp*cdim, MPIU_REAL, &refcoord);CHKERRQ(ierr);
342     ierr = PetscFree(points);CHKERRQ(ierr);
343   }
344   ierr = DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
345   ierr = DMSwarmSortRestoreAccess(dm);CHKERRQ(ierr);
346   ierr = DMRestoreLocalVector(plex, &locPhi);CHKERRQ(ierr);
347   ierr = DMRestoreGlobalVector(plex, &phi);CHKERRQ(ierr);
348   ierr = VecRestoreArray(Vres,&vres);CHKERRQ(ierr);
349   ierr = VecRestoreArrayRead(X,&x);CHKERRQ(ierr);
350   ierr = VecViewFromOptions(Vres, NULL, "-vel_res_view");CHKERRQ(ierr);
351   PetscFunctionReturn(0);
352 }
353 
354 int main(int argc,char **argv)
355 {
356   PetscInt          i, par;
357   PetscInt          locSize, p, d, dim, Np, step, *idx1, *idx2;
358   TS                ts;
359   DM                dm, sw;
360   AppCtx            user;
361   MPI_Comm          comm;
362   PetscErrorCode    ierr;
363   Vec               coorVec, kinVec, probVec, solution, position, momentum;
364   const PetscScalar *coorArr, *kinArr;
365   PetscReal         ftime   = 10., *probArr, *probVecArr;
366   IS                is1,is2;
367   PetscReal         *coor, *kin, *pos, *mom;
368 
369   ierr = PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;
370   comm = PETSC_COMM_WORLD;
371   ierr = ProcessOptions(comm, &user);CHKERRQ(ierr);
372   /* Create dm and particles */
373   ierr = CreateMesh(comm, &dm, &user);CHKERRQ(ierr);
374   ierr = CreateFEM(dm, &user);CHKERRQ(ierr);
375   ierr = CreateParticles(dm, &sw, &user);CHKERRQ(ierr);
376   ierr = SNESCreate(comm, &user.snes);CHKERRQ(ierr);
377   ierr = SNESSetDM(user.snes, dm);CHKERRQ(ierr);
378   ierr = DMPlexSetSNESLocalFEM(dm,&user,&user,&user);CHKERRQ(ierr);
379   ierr = SNESSetFromOptions(user.snes);CHKERRQ(ierr);
380   {
381     Mat          J;
382     MatNullSpace nullSpace;
383 
384     ierr = DMCreateMatrix(dm, &J);CHKERRQ(ierr);
385     ierr = MatNullSpaceCreate(PetscObjectComm((PetscObject) dm), PETSC_TRUE, 0, NULL, &nullSpace);CHKERRQ(ierr);
386     ierr = MatSetNullSpace(J, nullSpace);CHKERRQ(ierr);
387     ierr = MatNullSpaceDestroy(&nullSpace);CHKERRQ(ierr);
388     ierr = SNESSetJacobian(user.snes, J, J, NULL, NULL);CHKERRQ(ierr);
389     ierr = MatDestroy(&J);CHKERRQ(ierr);
390   }
391   /* Place TSSolve in a loop to handle resetting the TS at every manual call of TSStep() */
392   ierr = TSCreate(comm, &ts);CHKERRQ(ierr);
393   ierr = TSSetMaxTime(ts,ftime);CHKERRQ(ierr);
394   ierr = TSSetTimeStep(ts,user.stepSize);CHKERRQ(ierr);
395   ierr = TSSetMaxSteps(ts,100000);CHKERRQ(ierr);
396   ierr = TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
397   for (step = 0; step < user.steps ; ++step){
398 
399     ierr = DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &kinVec);CHKERRQ(ierr);
400     ierr = DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &coorVec);CHKERRQ(ierr);
401     ierr = VecViewFromOptions(kinVec, NULL, "-ic_vec_view");
402     ierr = DMGetDimension(sw, &dim);CHKERRQ(ierr);
403     ierr = VecGetLocalSize(kinVec, &locSize);CHKERRQ(ierr);
404     ierr = PetscMalloc1(locSize, &idx1);CHKERRQ(ierr);
405     ierr = PetscMalloc1(locSize, &idx2);CHKERRQ(ierr);
406     ierr = PetscMalloc1(2*locSize, &probArr);CHKERRQ(ierr);
407     Np = locSize/dim;
408     ierr = VecGetArrayRead(kinVec, &kinArr);CHKERRQ(ierr);
409     ierr = VecGetArrayRead(coorVec, &coorArr);CHKERRQ(ierr);
410     for (p=0; p<Np; ++p){
411         for (d=0; d<dim;++d) {
412             probArr[p*2*dim + d] = coorArr[p*dim+d];
413             probArr[(p*2+1)*dim + d] = kinArr[p*dim+d];
414         }
415     }
416     ierr = VecRestoreArrayRead(kinVec, &kinArr);CHKERRQ(ierr);
417     ierr = VecRestoreArrayRead(coorVec, &coorArr);CHKERRQ(ierr);
418     /* Allocate for IS Strides that will contain x, y and vx, vy */
419     for (p = 0; p < Np; ++p) {
420       for (d = 0; d < dim; ++d) {
421         idx1[p*dim+d] = (p*2+0)*dim + d;
422         idx2[p*dim+d] = (p*2+1)*dim + d;
423       }
424     }
425 
426     ierr = ISCreateGeneral(comm, locSize, idx1, PETSC_OWN_POINTER, &is1);CHKERRQ(ierr);
427     ierr = ISCreateGeneral(comm, locSize, idx2, PETSC_OWN_POINTER, &is2);CHKERRQ(ierr);
428     /* DM needs to be set before splits so it propogates to sub TSs */
429     ierr = TSSetDM(ts, sw);CHKERRQ(ierr);
430     ierr = TSSetType(ts,TSBASICSYMPLECTIC);CHKERRQ(ierr);
431     ierr = TSRHSSplitSetIS(ts,"position",is1);CHKERRQ(ierr);
432     ierr = TSRHSSplitSetIS(ts,"momentum",is2);CHKERRQ(ierr);
433     ierr = TSRHSSplitSetRHSFunction(ts,"position",NULL,RHSFunction1,&user);CHKERRQ(ierr);
434     ierr = TSRHSSplitSetRHSFunction(ts,"momentum",NULL,RHSFunction2,&user);CHKERRQ(ierr);
435     ierr = TSSetTime(ts, step*user.stepSize);CHKERRQ(ierr);
436     if (step == 0) {
437       ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
438     }
439     /* Compose vector from array for TS solve with all kinematic variables */
440     ierr = VecCreate(comm,&probVec);CHKERRQ(ierr);
441     ierr = VecSetBlockSize(probVec,1);CHKERRQ(ierr);
442     ierr = VecSetSizes(probVec,PETSC_DECIDE,2*locSize);CHKERRQ(ierr);
443     ierr = VecSetUp(probVec);CHKERRQ(ierr);
444     ierr = VecGetArray(probVec,&probVecArr);
445     for (i=0; i < 2*locSize; ++i) {
446       probVecArr[i] = probArr[i];
447     }
448     ierr = VecRestoreArray(probVec,&probVecArr);CHKERRQ(ierr);
449     ierr = TSSetSolution(ts, probVec);CHKERRQ(ierr);
450     ierr = PetscFree(probArr);CHKERRQ(ierr);
451     ierr = VecViewFromOptions(kinVec, NULL, "-ic_view");
452     ierr = DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &kinVec);CHKERRQ(ierr);
453     ierr = DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &coorVec);CHKERRQ(ierr);
454     ierr = TSMonitor(ts, step, ts->ptime, ts->vec_sol);CHKERRQ(ierr);
455     if (!ts->steprollback) {
456       ierr = TSPreStep(ts);CHKERRQ(ierr);
457     }
458     ierr = TSStep(ts);CHKERRQ(ierr);
459     if (ts->steprollback) {
460       ierr = TSPostEvaluate(ts);CHKERRQ(ierr);
461     }
462     if (!ts->steprollback) {
463 
464       TSPostStep(ts);
465       ierr = DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **) &coor);CHKERRQ(ierr);
466       ierr = DMSwarmGetField(sw, "kinematics", NULL, NULL, (void **) &kin);CHKERRQ(ierr);
467       ierr = TSGetSolution(ts, &solution);CHKERRQ(ierr);
468       ierr = VecGetSubVector(solution,is1,&position);CHKERRQ(ierr);
469       ierr = VecGetSubVector(solution,is2,&momentum);CHKERRQ(ierr);
470       ierr = VecGetArray(position, &pos);CHKERRQ(ierr);
471       ierr = VecGetArray(momentum, &mom);CHKERRQ(ierr);
472       for (par = 0; par < Np; ++par){
473         for (d=0; d<dim; ++d){
474           if (pos[par*dim+d] < 0.) {
475             coor[par*dim+d] = pos[par*dim+d] + 2.*PETSC_PI;
476           }
477           else if (pos[par*dim+d] > 2.*PETSC_PI) {
478             coor[par*dim+d] = pos[par*dim+d] - 2.*PETSC_PI;
479           }
480           else{
481             coor[par*dim+d] = pos[par*dim+d];
482           }
483           kin[par*dim+d] = mom[par*dim+d];
484         }
485       }
486       ierr = VecRestoreArray(position, &pos);CHKERRQ(ierr);
487       ierr = VecRestoreArray(momentum, &mom);CHKERRQ(ierr);
488       ierr = VecRestoreSubVector(solution,is1,&position);CHKERRQ(ierr);
489       ierr = VecRestoreSubVector(solution,is2,&momentum);CHKERRQ(ierr);
490       ierr = DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **) &coor);CHKERRQ(ierr);
491       ierr = DMSwarmRestoreField(sw, "kinematics", NULL, NULL, (void **) &kin);CHKERRQ(ierr);
492     }
493     ierr = DMSwarmMigrate(sw, PETSC_TRUE);CHKERRQ(ierr);
494     ierr = DMLocalizeCoordinates(sw);CHKERRQ(ierr);
495     ierr = TSReset(ts);CHKERRQ(ierr);
496     ierr = VecDestroy(&probVec);CHKERRQ(ierr);
497     ierr = ISDestroy(&is1);CHKERRQ(ierr);
498     ierr = ISDestroy(&is2);CHKERRQ(ierr);
499   }
500   ierr = SNESDestroy(&user.snes);CHKERRQ(ierr);
501   ierr = TSDestroy(&ts);CHKERRQ(ierr);
502   ierr = DMDestroy(&sw);CHKERRQ(ierr);
503   ierr = DMDestroy(&dm);CHKERRQ(ierr);
504   ierr = PetscFinalize();
505   return ierr;
506 }
507 
508 /*TEST
509 
510    build:
511      requires: triangle !single !complex
512    test:
513      suffix: bsi1q3
514      args: -particlesPerCell 200\
515       -petscspace_degree 2\
516       -petscfe_default_quadrature_order 3\
517       -ts_basicsymplectic_type 1\
518       -pc_type svd\
519       -uniform\
520       -sigma 1.0e-8\
521       -timeScale 2.0e-14\
522       -stepSize 1.0e-2\
523       -ts_monitor_sp_swarm\
524       -steps 10\
525       -dm_view\
526       -dm_plex_simplex 0 -dm_plex_dim 2\
527       -dm_plex_box_lower 0,-1\
528       -dm_plex_box_upper 6.283185307179586,1\
529       -dm_plex_box_bd periodic,none\
530       -dm_plex_box_faces 4,1
531    test:
532      suffix: bsi2q3
533      args: -particlesPerCell 200\
534       -petscspace_degree 2\
535       -petscfe_default_quadrature_order 3\
536       -ts_basicsymplectic_type 2\
537       -pc_type svd\
538       -uniform\
539       -sigma 1.0e-8\
540       -timeScale 2.0e-14\
541       -stepSize 1.0e-2\
542       -ts_monitor_sp_swarm\
543       -steps 10\
544       -dm_view\
545       -dm_plex_simplex 0 -dm_plex_dim 2\
546       -dm_plex_box_lower 0,-1\
547       -dm_plex_box_upper 6.283185307179586,1\
548       -dm_plex_box_bd periodic,none\
549       -dm_plex_box_faces 4,1
550 TEST*/
551