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