xref: /petsc/src/dm/impls/swarm/tests/ex4.c (revision d0295fc027abbea29f13fa82c19cf92da8e9ba99)
1 static char help[] = "Example of simple hamiltonian system (harmonic oscillator) with particles and a basic symplectic integrator\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;                          /* The topological mesh dimension */
11   PetscBool simplex;                      /* Flag for simplices or tensor cells */
12   char      filename[PETSC_MAX_PATH_LEN]; /* Name of the mesh filename if any */
13   PetscReal omega;                        /* Oscillation frequency omega */
14   PetscInt  particlesPerCell;             /* The number of partices per cell */
15   PetscReal momentTol;                    /* Tolerance for checking moment conservation */
16   PetscBool monitor;                      /* Flag for using the TS monitor */
17   PetscBool error;                        /* Flag for printing the error */
18   PetscInt  ostep;                        /* print the energy at each ostep time steps */
19 } AppCtx;
20 
21 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
22 {
23   PetscErrorCode ierr;
24 
25   PetscFunctionBeginUser;
26   options->dim              = 2;
27   options->simplex          = PETSC_TRUE;
28   options->monitor          = PETSC_FALSE;
29   options->error            = PETSC_FALSE;
30   options->particlesPerCell = 1;
31   options->momentTol        = 100.0*PETSC_MACHINE_EPSILON;
32   options->omega            = 64.0;
33   options->ostep            = 100;
34 
35   ierr = PetscStrcpy(options->filename, "");CHKERRQ(ierr);
36 
37   ierr = PetscOptionsBegin(comm, "", "Harmonic Oscillator Options", "DMPLEX");CHKERRQ(ierr);
38   ierr = PetscOptionsInt("-output_step", "Number of time steps between output", "ex4.c", options->ostep, &options->ostep, PETSC_NULL);CHKERRQ(ierr);
39   ierr = PetscOptionsInt("-dim", "The topological mesh dimension", "ex4.c", options->dim, &options->dim, NULL);CHKERRQ(ierr);
40   ierr = PetscOptionsBool("-monitor", "Flag to use the TS monitor", "ex4.c", options->monitor, &options->monitor, NULL);CHKERRQ(ierr);
41   ierr = PetscOptionsBool("-error", "Flag to print the error", "ex4.c", options->error, &options->error, NULL);CHKERRQ(ierr);
42   ierr = PetscOptionsBool("-simplex", "The flag for simplices or tensor cells", "ex4.c", options->simplex, &options->simplex, NULL);CHKERRQ(ierr);
43   ierr = PetscOptionsString("-mesh", "Name of the mesh filename if any", "ex4.c", options->filename, options->filename, sizeof(options->filename), NULL);CHKERRQ(ierr);
44   ierr = PetscOptionsInt("-particles_per_cell", "Number of particles per cell", "ex4.c", options->particlesPerCell, &options->particlesPerCell, NULL);CHKERRQ(ierr);
45   ierr = PetscOptionsReal("-omega", "Oscillator frequency", "ex4.c", options->omega, &options->omega, PETSC_NULL);CHKERRQ(ierr);
46   ierr = PetscOptionsEnd();CHKERRQ(ierr);
47 
48   PetscFunctionReturn(0);
49 }
50 
51 static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user)
52 {
53   PetscBool      flg;
54   PetscErrorCode ierr;
55 
56   PetscFunctionBeginUser;
57   ierr = PetscStrcmp(user->filename, "", &flg);CHKERRQ(ierr);
58   if (flg) {
59     ierr = DMPlexCreateBoxMesh(comm, user->dim, user->simplex, NULL, NULL, NULL, NULL, PETSC_TRUE, dm);CHKERRQ(ierr);
60   } else {
61     ierr = DMPlexCreateFromFile(comm, user->filename, PETSC_TRUE, dm);CHKERRQ(ierr);
62     ierr = DMGetDimension(*dm, &user->dim);CHKERRQ(ierr);
63   }
64   {
65     DM distributedMesh = NULL;
66 
67     ierr = DMPlexDistribute(*dm, 0, NULL, &distributedMesh);CHKERRQ(ierr);
68     if (distributedMesh) {
69       ierr = DMDestroy(dm);CHKERRQ(ierr);
70       *dm  = distributedMesh;
71     }
72   }
73   ierr = DMLocalizeCoordinates(*dm);CHKERRQ(ierr); /* needed for periodic */
74   ierr = DMSetFromOptions(*dm);CHKERRQ(ierr);
75   ierr = PetscObjectSetName((PetscObject) *dm, "Mesh");CHKERRQ(ierr);
76   ierr = DMViewFromOptions(*dm, NULL, "-dm_view");CHKERRQ(ierr);
77   PetscFunctionReturn(0);
78 }
79 
80 static PetscErrorCode SetInitialCoordinates(DM dmSw)
81 {
82   DM             dm;
83   AppCtx        *user;
84   PetscRandom    rnd;
85   PetscBool      simplex;
86   PetscReal     *centroid, *coords, *xi0, *v0, *J, *invJ, detJ;
87   PetscInt       dim, d, cStart, cEnd, c, Np, p;
88   PetscErrorCode ierr;
89 
90   PetscFunctionBeginUser;
91   ierr = PetscRandomCreate(PetscObjectComm((PetscObject) dmSw), &rnd);CHKERRQ(ierr);
92   ierr = PetscRandomSetInterval(rnd, -1.0, 1.0);CHKERRQ(ierr);
93   ierr = PetscRandomSetFromOptions(rnd);CHKERRQ(ierr);
94 
95   ierr = DMGetApplicationContext(dmSw, (void **) &user);CHKERRQ(ierr);
96   simplex = user->simplex;
97   Np   = user->particlesPerCell;
98   ierr = DMSwarmGetCellDM(dmSw, &dm);CHKERRQ(ierr);
99   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
100   ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);CHKERRQ(ierr);
101   ierr = PetscMalloc5(dim, &centroid, dim, &xi0, dim, &v0, dim*dim, &J, dim*dim, &invJ);CHKERRQ(ierr);
102   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
103   ierr = DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
104   for (c = cStart; c < cEnd; ++c) {
105     if (Np == 1) {
106       ierr = DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL);CHKERRQ(ierr);
107       for (d = 0; d < dim; ++d) coords[c*dim+d] = centroid[d];
108     } else {
109       ierr = DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ);CHKERRQ(ierr); /* affine */
110       for (p = 0; p < Np; ++p) {
111         const PetscInt n   = c*Np + p;
112         PetscReal      sum = 0.0, refcoords[3];
113 
114         for (d = 0; d < dim; ++d) {
115           ierr = PetscRandomGetValueReal(rnd, &refcoords[d]);CHKERRQ(ierr);
116           sum += refcoords[d];
117         }
118         if (simplex && sum > 0.0) for (d = 0; d < dim; ++d) refcoords[d] -= PetscSqrtReal(dim)*sum;
119         CoordinatesRefToReal(dim, dim, xi0, v0, J, refcoords, &coords[n*dim]);
120       }
121     }
122   }
123   ierr = DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
124   ierr = PetscFree5(centroid, xi0, v0, J, invJ);CHKERRQ(ierr);
125   ierr = PetscRandomDestroy(&rnd);CHKERRQ(ierr);
126   PetscFunctionReturn(0);
127 }
128 
129 static PetscErrorCode SetInitialConditions(DM dmSw, Vec u)
130 {
131   DM             dm;
132   AppCtx        *user;
133   PetscReal     *coords;
134   PetscScalar   *initialConditions;
135   PetscInt       dim, cStart, cEnd, c, Np, p;
136   PetscErrorCode ierr;
137 
138   PetscFunctionBeginUser;
139   ierr = DMGetApplicationContext(dmSw, (void **) &user);CHKERRQ(ierr);
140   Np   = user->particlesPerCell;
141   ierr = DMSwarmGetCellDM(dmSw, &dm);CHKERRQ(ierr);
142   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
143   ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);CHKERRQ(ierr);
144   ierr = DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
145   ierr = VecGetArray(u, &initialConditions);CHKERRQ(ierr);
146   for (c = cStart; c < cEnd; ++c) {
147     for (p = 0; p < Np; ++p) {
148       const PetscInt n = c*Np + p;
149 
150       initialConditions[n*2+0] = DMPlex_NormD_Internal(dim, &coords[n*dim]);
151       initialConditions[n*2+1] = 0.0;
152     }
153   }
154   ierr = VecRestoreArray(u, &initialConditions);CHKERRQ(ierr);
155   ierr = DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
156 
157   ierr = VecView(u, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
158   PetscFunctionReturn(0);
159 }
160 
161 static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user)
162 {
163   PetscInt      *cellid;
164   PetscInt       dim, cStart, cEnd, c, Np = user->particlesPerCell, p;
165   PetscErrorCode ierr;
166 
167   PetscFunctionBeginUser;
168   ierr = DMGetDimension(dm, &dim);CHKERRQ(ierr);
169   ierr = DMCreate(PetscObjectComm((PetscObject) dm), sw);CHKERRQ(ierr);
170   ierr = DMSetType(*sw, DMSWARM);CHKERRQ(ierr);
171   ierr = DMSetDimension(*sw, dim);CHKERRQ(ierr);
172 
173   ierr = DMSwarmSetType(*sw, DMSWARM_PIC);CHKERRQ(ierr);
174   ierr = DMSwarmSetCellDM(*sw, dm);CHKERRQ(ierr);
175   ierr = DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", 2, PETSC_REAL);CHKERRQ(ierr);
176   ierr = DMSwarmFinalizeFieldRegister(*sw);CHKERRQ(ierr);
177   ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);CHKERRQ(ierr);
178   ierr = DMSwarmSetLocalSizes(*sw, (cEnd - cStart) * Np, 0);CHKERRQ(ierr);
179   ierr = DMSetFromOptions(*sw);CHKERRQ(ierr);
180   ierr = DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid);CHKERRQ(ierr);
181   for (c = cStart; c < cEnd; ++c) {
182     for (p = 0; p < Np; ++p) {
183       const PetscInt n = c*Np + p;
184 
185       cellid[n] = c;
186     }
187   }
188   ierr = DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid);CHKERRQ(ierr);
189   ierr = PetscObjectSetName((PetscObject) *sw, "Particles");CHKERRQ(ierr);
190   ierr = DMViewFromOptions(*sw, NULL, "-sw_view");CHKERRQ(ierr);
191   PetscFunctionReturn(0);
192 }
193 
194 static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
195 {
196   AppCtx            *user  = (AppCtx *) ctx;
197   const PetscReal    omega = user->omega;
198   const PetscScalar *u;
199   MPI_Comm           comm;
200   PetscReal          dt;
201   PetscInt           Np, p;
202   PetscErrorCode     ierr;
203 
204   PetscFunctionBeginUser;
205   if (step%user->ostep == 0) {
206     ierr = PetscObjectGetComm((PetscObject) ts, &comm);CHKERRQ(ierr);
207     if (!step) {ierr = PetscPrintf(comm, "Time     Step Part     Energy Mod Energy\n");CHKERRQ(ierr);}
208     ierr = TSGetTimeStep(ts, &dt);CHKERRQ(ierr);
209     ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
210     ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
211     Np /= 2;
212     for (p = 0; p < Np; ++p) {
213       const PetscReal x  = PetscRealPart(u[p*2+0]);
214       const PetscReal v  = PetscRealPart(u[p*2+1]);
215       const PetscReal E  = 0.5*(v*v + PetscSqr(omega)*x*x);
216       const PetscReal mE = 0.5*(v*v + PetscSqr(omega)*x*x - PetscSqr(omega)*dt*x*v);
217 
218       ierr = PetscPrintf(comm, "%.6lf %4D %4D %10.4lf %10.4lf\n", t, step, p, (double) E, (double) mE);CHKERRQ(ierr);
219     }
220     ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
221   }
222   PetscFunctionReturn(0);
223 }
224 
225 static PetscErrorCode InitializeSolve(TS ts, Vec u)
226 {
227   DM             dm;
228   AppCtx        *user;
229   PetscErrorCode ierr;
230 
231   PetscFunctionBeginUser;
232   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
233   ierr = DMGetApplicationContext(dm, (void **) &user);CHKERRQ(ierr);
234   ierr = SetInitialCoordinates(dm);CHKERRQ(ierr);
235   ierr = SetInitialConditions(dm, u);CHKERRQ(ierr);
236   PetscFunctionReturn(0);
237 }
238 
239 static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
240 {
241   MPI_Comm           comm;
242   DM                 sdm;
243   AppCtx            *user;
244   const PetscScalar *u, *coords;
245   PetscScalar       *e;
246   PetscReal          t, omega;
247   PetscInt           dim, Np, p;
248   PetscErrorCode     ierr;
249 
250 
251   PetscFunctionBeginUser;
252   ierr = PetscObjectGetComm((PetscObject) ts, &comm);CHKERRQ(ierr);
253   ierr = TSGetDM(ts, &sdm);CHKERRQ(ierr);
254   ierr = DMGetApplicationContext(sdm, (void **) &user);CHKERRQ(ierr);
255   omega = user->omega;
256   ierr = DMGetDimension(sdm, &dim);CHKERRQ(ierr);
257   ierr = TSGetSolveTime(ts, &t);CHKERRQ(ierr);
258   ierr = VecGetArray(E, &e);CHKERRQ(ierr);
259   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
260   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
261   Np  /= 2;
262   ierr = DMSwarmGetField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
263   for (p = 0; p < Np; ++p) {
264     const PetscReal x  = PetscRealPart(u[p*2+0]);
265     const PetscReal v  = PetscRealPart(u[p*2+1]);
266     const PetscReal x0 = DMPlex_NormD_Internal(dim, &coords[p*dim]);
267     const PetscReal ex =  x0*PetscCosReal(omega*t);
268     const PetscReal ev = -x0*omega*PetscSinReal(omega*t);
269 
270     if (user->error) {ierr = PetscPrintf(comm, "p%D error [%.2g %.2g] sol [%.6lf %.6lf] exact [%.6lf %.6lf] energy/exact energy %g / %g\n", p, (double) PetscAbsReal(x-ex), (double) PetscAbsReal(v-ev), (double) x, (double) v, (double) ex, (double) ev, 0.5*(v*v + PetscSqr(omega)*x*x), (double) 0.5*PetscSqr(omega*x0));}
271     e[p*2+0] = x - ex;
272     e[p*2+1] = v - ev;
273   }
274   ierr = DMSwarmRestoreField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
275 
276   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
277   ierr = VecRestoreArray(E, &e);CHKERRQ(ierr);
278   PetscFunctionReturn(0);
279 }
280 
281 
282 /*---------------------Create particle RHS Functions--------------------------*/
283 static PetscErrorCode RHSFunction1(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
284 {
285   const PetscScalar *v;
286   PetscScalar       *xres;
287   PetscInt           Np, p;
288   PetscErrorCode     ierr;
289 
290   PetscFunctionBeginUser;
291   ierr = VecGetArray(Xres, &xres);CHKERRQ(ierr);
292   ierr = VecGetArrayRead(V, &v);CHKERRQ(ierr);
293   ierr = VecGetLocalSize(Xres, &Np);CHKERRQ(ierr);
294   for (p = 0; p < Np; ++p) {
295      xres[p] = v[p];
296   }
297   ierr = VecRestoreArrayRead(V, &v);CHKERRQ(ierr);
298   ierr = VecRestoreArray(Xres, &xres);CHKERRQ(ierr);
299   PetscFunctionReturn(0);
300 }
301 
302 static PetscErrorCode RHSFunction2(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
303 {
304   AppCtx            *user = (AppCtx *)ctx;
305   const PetscScalar *x;
306   PetscScalar       *vres;
307   PetscInt           Np, p;
308   PetscErrorCode     ierr;
309 
310   PetscFunctionBeginUser;
311   ierr = VecGetArray(Vres, &vres);CHKERRQ(ierr);
312   ierr = VecGetArrayRead(X, &x);CHKERRQ(ierr);
313   ierr = VecGetLocalSize(Vres, &Np);CHKERRQ(ierr);
314   for (p = 0; p < Np; ++p) {
315     vres[p] = -PetscSqr(user->omega)*x[p];
316   }
317   ierr = VecRestoreArrayRead(X, &x);CHKERRQ(ierr);
318   ierr = VecRestoreArray(Vres, &vres);CHKERRQ(ierr);
319   PetscFunctionReturn(0);
320 }
321 
322 // static PetscErrorCode RHSFunctionParticles(TS ts, PetscReal t, Vec U, Vec R, void *ctx)
323 // {
324 //   AppCtx            *user = (AppCtx *) ctx;
325 //   DM                 dm;
326 //   const PetscScalar *u;
327 //   PetscScalar       *r;
328 //   PetscInt           Np, p;
329 //   PetscErrorCode     ierr;
330 //
331 //   PetscFunctionBeginUser;
332 //   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
333 //   ierr = VecGetArray(R, &r);CHKERRQ(ierr);
334 //   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
335 //   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
336 //   Np  /= 2;
337 //   for (p = 0; p < Np; ++p) {
338 //     r[p*2+0] = u[p*2+1];
339 //     r[p*2+1] = -PetscSqr(user->omega)*u[p*2+0];
340 //   }
341 //   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
342 //   ierr = VecRestoreArray(R, &r);CHKERRQ(ierr);
343 //   PetscFunctionReturn(0);
344 // }
345 /*----------------------------------------------------------------------------*/
346 
347 /*--------------------Define RHSFunction, RHSJacobian (Theta)-----------------*/
348 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
349 {
350   AppCtx            *user = (AppCtx *) ctx;
351   DM                 dm;
352   const PetscScalar *u;
353   PetscScalar       *g;
354   PetscInt           Np, p;
355   PetscErrorCode     ierr;
356 
357   PetscFunctionBeginUser;
358   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
359   ierr = VecGetArray(G, &g);CHKERRQ(ierr);
360   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
361   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
362   Np  /= 2;
363   for (p = 0; p < Np; ++p) {
364     g[p*2+0] = u[p*2+1];
365     g[p*2+1] = -PetscSqr(user->omega)*u[p*2+0];
366   }
367   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
368   ierr = VecRestoreArray(G, &g);CHKERRQ(ierr);
369   PetscFunctionReturn(0);
370 }
371 
372 /*Ji = dFi/dxj
373 J= (0    1)
374    (-w^2 0)
375 */
376 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U , Mat J, Mat P, void *ctx)
377 {
378   AppCtx            *user = (AppCtx *) ctx;
379   PetscInt           Np = user->dim * user->particlesPerCell;
380   PetscInt           i, m, n;
381   const PetscScalar *u;
382   PetscScalar        vals[4] = {0., 1., -PetscSqr(user->omega), 0.};
383   PetscErrorCode     ierr;
384 
385 
386   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
387   //ierr = PetscPrintf(PETSC_COMM_WORLD, "# Particles (Np) = %d \n" , Np);CHKERRQ(ierr);
388 
389   ierr = MatGetOwnershipRange(J, &m, &n);CHKERRQ(ierr);
390   for (i = 0; i < Np; ++i) {
391     const PetscInt rows[2] = {2*i, 2*i+1};
392     ierr = MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
393   }
394   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
395   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
396   //ierr = MatView(J,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
397 
398   PetscFunctionReturn(0);
399 
400 }
401 /*----------------------------------------------------------------------------*/
402 
403 /*----------------Define S, F, G Functions (Discrete Gradients)---------------*/
404 /*
405   u_t = S * gradF
406     --or--
407   u_t = S * G
408 
409   + Sfunc - constructor for the S matrix from the formulation
410   . Ffunc - functional F from the formulation
411   - Gfunc - constructor for the gradient of F from the formulation
412 */
413 
414 PetscErrorCode Sfunc(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
415 {
416   AppCtx            *user = (AppCtx *) ctx;
417   PetscInt           Np = user->dim * user->particlesPerCell;
418   PetscInt           i, m, n;
419   const PetscScalar *u;
420   PetscScalar        vals[4] = {0., 1., -1, 0.};
421   PetscErrorCode     ierr;
422 
423   PetscFunctionBeginUser;
424   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
425   ierr = MatGetOwnershipRange(S, &m, &n);CHKERRQ(ierr);
426   for (i = 0; i < Np; ++i) {
427     const PetscInt rows[2] = {2*i, 2*i+1};
428     ierr = MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
429   }
430   ierr = MatAssemblyBegin(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
431   ierr = MatAssemblyEnd(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
432 
433   PetscFunctionReturn(0);
434 }
435 
436 PetscErrorCode Ffunc(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
437 {
438   AppCtx            *user = (AppCtx *) ctx;
439   DM                 dm;
440   const PetscScalar *u;
441   PetscInt           Np = user->dim * user->particlesPerCell;
442   PetscInt           p;
443   PetscErrorCode     ierr;
444 
445   PetscFunctionBeginUser;
446   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
447   //Define F
448   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
449   for (p = 0; p < Np; ++p) {
450     *F += (1/2)*PetscSqr(user->omega)*PetscSqr(u[p*2+0]) + (1/2)*PetscSqr(u[p*2+1]);
451   }
452   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
453 
454   PetscFunctionReturn(0);
455 }
456 
457 PetscErrorCode gradFfunc(TS ts, PetscReal t, Vec U, Vec gradF, void *ctx)
458 {
459   AppCtx            *user = (AppCtx *) ctx;
460   DM                 dm;
461   const PetscScalar *u;
462   PetscScalar       *g;
463   PetscInt           Np = user->dim * user->particlesPerCell;
464   PetscInt           p;
465   PetscErrorCode     ierr;
466 
467   PetscFunctionBeginUser;
468   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
469   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
470 
471   //Define gradF
472   ierr = VecGetArray(gradF, &g);CHKERRQ(ierr);
473   for (p = 0; p < Np; ++p) {
474     g[p*2+0] = PetscSqr(user->omega)*u[p*2+0]; /*dF/dx*/
475     g[p*2+1] = u[p*2+1]; /*dF/dv*/
476   }
477   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
478   ierr = VecRestoreArray(gradF, &g);CHKERRQ(ierr);
479 
480   PetscFunctionReturn(0);
481 }
482 /*----------------------------------------------------------------------------*/
483 
484 int main(int argc,char **argv)
485 {
486   TS             ts;     /* nonlinear solver                 */
487   DM             dm, sw; /* Mesh and particle managers       */
488   Vec            u;      /* swarm vector                     */
489   PetscInt       n;      /* swarm vector size                */
490   IS             is1, is2;
491   MPI_Comm       comm;
492   Mat            J;      /* Jacobian matrix                  */
493   AppCtx         user;
494   PetscErrorCode ierr;
495 
496   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
497      Initialize program
498      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
499   ierr = PetscInitialize(&argc, &argv, NULL, help);if (ierr) return ierr;
500   comm = PETSC_COMM_WORLD;
501   ierr = ProcessOptions(comm, &user);CHKERRQ(ierr);
502 
503 
504   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
505      Create Particle-Mesh
506      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
507   ierr = CreateMesh(comm, &dm, &user);CHKERRQ(ierr);
508   ierr = CreateParticles(dm, &sw, &user);CHKERRQ(ierr);
509   ierr = DMSetApplicationContext(sw, &user);CHKERRQ(ierr);
510 
511 
512   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
513      Setup timestepping solver
514      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
515   ierr = TSCreate(comm, &ts);CHKERRQ(ierr);
516   ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
517 
518   ierr = TSSetDM(ts, sw);CHKERRQ(ierr);
519   ierr = TSSetMaxTime(ts, 0.1);CHKERRQ(ierr);
520   ierr = TSSetTimeStep(ts, 0.00001);CHKERRQ(ierr);
521   ierr = TSSetMaxSteps(ts, 100);CHKERRQ(ierr);
522   ierr = TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
523   if (user.monitor) {ierr = TSMonitorSet(ts, Monitor, &user, NULL);CHKERRQ(ierr);}
524 
525 
526   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
527      Prepare to solve
528      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
529   ierr = DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
530   ierr = VecGetLocalSize(u, &n);CHKERRQ(ierr);
531   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
532   ierr = TSSetComputeInitialCondition(ts, InitializeSolve);CHKERRQ(ierr);
533   ierr = TSSetComputeExactError(ts, ComputeError);CHKERRQ(ierr);
534 
535   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
536      Define function F(U, Udot , x , t) = G(U, x, t)
537      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
538 
539   /* - - - - - - - Basic Symplectic - - - - - - - - - - - - - - - - - - - - - -*/
540   ierr = ISCreateStride(comm, n/2, 0, 2, &is1);CHKERRQ(ierr);
541   ierr = ISCreateStride(comm, n/2, 1, 2, &is2);CHKERRQ(ierr);
542   ierr = TSRHSSplitSetIS(ts, "position", is1);CHKERRQ(ierr);
543   ierr = TSRHSSplitSetIS(ts, "momentum", is2);CHKERRQ(ierr);
544   ierr = ISDestroy(&is1);CHKERRQ(ierr);
545   ierr = ISDestroy(&is2);CHKERRQ(ierr);
546   ierr = TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunction1, &user);CHKERRQ(ierr);
547   ierr = TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunction2, &user);CHKERRQ(ierr);
548 
549   /* - - - - - - - Theta (Implicit Midpoint) - - - - - - - - - - - - - - - - -*/
550   ierr = TSSetRHSFunction(ts, NULL, RHSFunction, &user);CHKERRQ(ierr);
551 
552   ierr = MatCreate(PETSC_COMM_WORLD,&J);CHKERRQ(ierr);
553   ierr = MatSetSizes(J,PETSC_DECIDE,PETSC_DECIDE,n,n);CHKERRQ(ierr);
554   ierr = MatSetFromOptions(J);CHKERRQ(ierr);
555   ierr = MatSetUp(J);CHKERRQ(ierr);
556   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
557   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
558   ierr = PetscPrintf(comm, "n = %d\n",n);CHKERRQ(ierr);//Check number of particles
559   ierr = TSSetRHSJacobian(ts,J,J,RHSJacobian,&user);CHKERRQ(ierr);
560 
561   /* - - - - - - - Discrete Gradients - - - - - - - - - - - - - - - - - - - - */
562   ierr = TSDiscGradSetFormulation(ts, Sfunc, Ffunc, gradFfunc, &user);CHKERRQ(ierr);
563 
564   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
565      Solve
566      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
567 
568   ierr = TSComputeInitialCondition(ts, u);CHKERRQ(ierr);
569   ierr = TSSolve(ts, u);CHKERRQ(ierr);
570 
571 
572   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
573      Clean up workspace
574      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
575   ierr = MatDestroy(&J);CHKERRQ(ierr);
576   ierr = DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
577   ierr = TSDestroy(&ts);CHKERRQ(ierr);
578   ierr = DMDestroy(&sw);CHKERRQ(ierr);
579   ierr = DMDestroy(&dm);CHKERRQ(ierr);
580   ierr = PetscFinalize();
581   return ierr;
582 }
583 
584 /*TEST
585 
586    build:
587      requires: triangle !single !complex
588    test:
589      suffix: 1
590      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_convergence_estimate -convest_num_refine 2 -dm_view -sw_view -monitor -output_step 50 -error
591 
592    test:
593      suffix: 2
594      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 2 -ts_convergence_estimate -convest_num_refine 2 -dm_view -sw_view -monitor -output_step 50 -error
595 
596    test:
597      suffix: 3
598      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 3 -ts_convergence_estimate -convest_num_refine 2 -ts_dt 0.0001 -dm_view -sw_view -monitor -output_step 50 -error
599 
600    test:
601      suffix: 4
602      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 4 -ts_convergence_estimate -convest_num_refine 2 -ts_dt 0.0001 -dm_view -sw_view -monitor -output_step 50 -error
603 
604    test:
605      suffix: 5
606      args: -dm_plex_box_faces 1,1 -ts_type theta -ts_theta_theta 0.5 -ts_convergence_estimate -convest_num_refine 2 -dm_view -sw_view -monitor -output_step 50 -error
607 
608    test:
609      suffix: 6
610      args: -dm_plex_box_faces 1,1 -ts_type discgrad -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2
611 
612 TEST*/
613