xref: /petsc/src/dm/impls/swarm/tests/ex4.c (revision dfd57a172ac9fa6c7b5fe6de6ab5df85cefc2996)
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, &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, &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, &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   PetscFunctionBeginUser;
226   ierr = PetscObjectGetComm((PetscObject) ts, &comm);CHKERRQ(ierr);
227   ierr = TSGetDM(ts, &sdm);CHKERRQ(ierr);
228   ierr = DMGetApplicationContext(sdm, &user);CHKERRQ(ierr);
229   omega = user->omega;
230   ierr = DMGetDimension(sdm, &dim);CHKERRQ(ierr);
231   ierr = TSGetSolveTime(ts, &t);CHKERRQ(ierr);
232   ierr = VecGetArray(E, &e);CHKERRQ(ierr);
233   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
234   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
235   Np  /= 2;
236   ierr = DMSwarmGetField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
237   for (p = 0; p < Np; ++p) {
238     const PetscReal x  = PetscRealPart(u[p*2+0]);
239     const PetscReal v  = PetscRealPart(u[p*2+1]);
240     const PetscReal x0 = DMPlex_NormD_Internal(dim, &coords[p*dim]);
241     const PetscReal ex =  x0*PetscCosReal(omega*t);
242     const PetscReal ev = -x0*omega*PetscSinReal(omega*t);
243 
244     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));}
245     e[p*2+0] = x - ex;
246     e[p*2+1] = v - ev;
247   }
248   ierr = DMSwarmRestoreField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords);CHKERRQ(ierr);
249 
250   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
251   ierr = VecRestoreArray(E, &e);CHKERRQ(ierr);
252   PetscFunctionReturn(0);
253 }
254 
255 /*---------------------Create particle RHS Functions--------------------------*/
256 static PetscErrorCode RHSFunction1(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
257 {
258   const PetscScalar *v;
259   PetscScalar       *xres;
260   PetscInt           Np, p;
261   PetscErrorCode     ierr;
262 
263   PetscFunctionBeginUser;
264   ierr = VecGetArray(Xres, &xres);CHKERRQ(ierr);
265   ierr = VecGetArrayRead(V, &v);CHKERRQ(ierr);
266   ierr = VecGetLocalSize(Xres, &Np);CHKERRQ(ierr);
267   for (p = 0; p < Np; ++p) {
268      xres[p] = v[p];
269   }
270   ierr = VecRestoreArrayRead(V, &v);CHKERRQ(ierr);
271   ierr = VecRestoreArray(Xres, &xres);CHKERRQ(ierr);
272   PetscFunctionReturn(0);
273 }
274 
275 static PetscErrorCode RHSFunction2(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
276 {
277   AppCtx            *user = (AppCtx *)ctx;
278   const PetscScalar *x;
279   PetscScalar       *vres;
280   PetscInt           Np, p;
281   PetscErrorCode     ierr;
282 
283   PetscFunctionBeginUser;
284   ierr = VecGetArray(Vres, &vres);CHKERRQ(ierr);
285   ierr = VecGetArrayRead(X, &x);CHKERRQ(ierr);
286   ierr = VecGetLocalSize(Vres, &Np);CHKERRQ(ierr);
287   for (p = 0; p < Np; ++p) {
288     vres[p] = -PetscSqr(user->omega)*x[p];
289   }
290   ierr = VecRestoreArrayRead(X, &x);CHKERRQ(ierr);
291   ierr = VecRestoreArray(Vres, &vres);CHKERRQ(ierr);
292   PetscFunctionReturn(0);
293 }
294 
295 // static PetscErrorCode RHSFunctionParticles(TS ts, PetscReal t, Vec U, Vec R, void *ctx)
296 // {
297 //   AppCtx            *user = (AppCtx *) ctx;
298 //   DM                 dm;
299 //   const PetscScalar *u;
300 //   PetscScalar       *r;
301 //   PetscInt           Np, p;
302 //   PetscErrorCode     ierr;
303 //
304 //   PetscFunctionBeginUser;
305 //   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
306 //   ierr = VecGetArray(R, &r);CHKERRQ(ierr);
307 //   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
308 //   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
309 //   Np  /= 2;
310 //   for (p = 0; p < Np; ++p) {
311 //     r[p*2+0] = u[p*2+1];
312 //     r[p*2+1] = -PetscSqr(user->omega)*u[p*2+0];
313 //   }
314 //   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
315 //   ierr = VecRestoreArray(R, &r);CHKERRQ(ierr);
316 //   PetscFunctionReturn(0);
317 // }
318 /*----------------------------------------------------------------------------*/
319 
320 /*--------------------Define RHSFunction, RHSJacobian (Theta)-----------------*/
321 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
322 {
323   AppCtx            *user = (AppCtx *) ctx;
324   DM                 dm;
325   const PetscScalar *u;
326   PetscScalar       *g;
327   PetscInt           Np, p;
328   PetscErrorCode     ierr;
329 
330   PetscFunctionBeginUser;
331   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
332   ierr = VecGetArray(G, &g);CHKERRQ(ierr);
333   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
334   ierr = VecGetLocalSize(U, &Np);CHKERRQ(ierr);
335   Np  /= 2;
336   for (p = 0; p < Np; ++p) {
337     g[p*2+0] = u[p*2+1];
338     g[p*2+1] = -PetscSqr(user->omega)*u[p*2+0];
339   }
340   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
341   ierr = VecRestoreArray(G, &g);CHKERRQ(ierr);
342   PetscFunctionReturn(0);
343 }
344 
345 /*Ji = dFi/dxj
346 J= (0    1)
347    (-w^2 0)
348 */
349 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U , Mat J, Mat P, void *ctx)
350 {
351   AppCtx            *user = (AppCtx *) ctx;
352   PetscInt           Np = user->dim * user->particlesPerCell;
353   PetscInt           i, m, n;
354   const PetscScalar *u;
355   PetscScalar        vals[4] = {0., 1., -PetscSqr(user->omega), 0.};
356   PetscErrorCode     ierr;
357 
358   PetscFunctionBeginUser;
359   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
360   //ierr = PetscPrintf(PETSC_COMM_WORLD, "# Particles (Np) = %d \n" , Np);CHKERRQ(ierr);
361 
362   ierr = MatGetOwnershipRange(J, &m, &n);CHKERRQ(ierr);
363   for (i = 0; i < Np; ++i) {
364     const PetscInt rows[2] = {2*i, 2*i+1};
365     ierr = MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
366   }
367   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
368   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
369   //ierr = MatView(J,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
370 
371   PetscFunctionReturn(0);
372 
373 }
374 /*----------------------------------------------------------------------------*/
375 
376 /*----------------Define S, F, G Functions (Discrete Gradients)---------------*/
377 /*
378   u_t = S * gradF
379     --or--
380   u_t = S * G
381 
382   + Sfunc - constructor for the S matrix from the formulation
383   . Ffunc - functional F from the formulation
384   - Gfunc - constructor for the gradient of F from the formulation
385 */
386 
387 PetscErrorCode Sfunc(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
388 {
389   AppCtx            *user = (AppCtx *) ctx;
390   PetscInt           Np = user->dim * user->particlesPerCell;
391   PetscInt           i, m, n;
392   const PetscScalar *u;
393   PetscScalar        vals[4] = {0., 1., -1, 0.};
394   PetscErrorCode     ierr;
395 
396   PetscFunctionBeginUser;
397   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
398   ierr = MatGetOwnershipRange(S, &m, &n);CHKERRQ(ierr);
399   for (i = 0; i < Np; ++i) {
400     const PetscInt rows[2] = {2*i, 2*i+1};
401     ierr = MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
402   }
403   ierr = MatAssemblyBegin(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
404   ierr = MatAssemblyEnd(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
405 
406   PetscFunctionReturn(0);
407 }
408 
409 PetscErrorCode Ffunc(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
410 {
411   AppCtx            *user = (AppCtx *) ctx;
412   DM                 dm;
413   const PetscScalar *u;
414   PetscInt           Np = user->dim * user->particlesPerCell;
415   PetscInt           p;
416   PetscErrorCode     ierr;
417 
418   PetscFunctionBeginUser;
419   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
420   //Define F
421   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
422   for (p = 0; p < Np; ++p) {
423     *F += (1/2)*PetscSqr(user->omega)*PetscSqr(u[p*2+0]) + (1/2)*PetscSqr(u[p*2+1]);
424   }
425   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
426 
427   PetscFunctionReturn(0);
428 }
429 
430 PetscErrorCode gradFfunc(TS ts, PetscReal t, Vec U, Vec gradF, void *ctx)
431 {
432   AppCtx            *user = (AppCtx *) ctx;
433   DM                 dm;
434   const PetscScalar *u;
435   PetscScalar       *g;
436   PetscInt           Np = user->dim * user->particlesPerCell;
437   PetscInt           p;
438   PetscErrorCode     ierr;
439 
440   PetscFunctionBeginUser;
441   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
442   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
443 
444   //Define gradF
445   ierr = VecGetArray(gradF, &g);CHKERRQ(ierr);
446   for (p = 0; p < Np; ++p) {
447     g[p*2+0] = PetscSqr(user->omega)*u[p*2+0]; /*dF/dx*/
448     g[p*2+1] = u[p*2+1]; /*dF/dv*/
449   }
450   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
451   ierr = VecRestoreArray(gradF, &g);CHKERRQ(ierr);
452 
453   PetscFunctionReturn(0);
454 }
455 /*----------------------------------------------------------------------------*/
456 
457 int main(int argc,char **argv)
458 {
459   TS             ts;     /* nonlinear solver                 */
460   DM             dm, sw; /* Mesh and particle managers       */
461   Vec            u;      /* swarm vector                     */
462   PetscInt       n;      /* swarm vector size                */
463   IS             is1, is2;
464   MPI_Comm       comm;
465   Mat            J;      /* Jacobian matrix                  */
466   AppCtx         user;
467   PetscErrorCode ierr;
468 
469   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
470      Initialize program
471      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
472   ierr = PetscInitialize(&argc, &argv, NULL, help);if (ierr) return ierr;
473   comm = PETSC_COMM_WORLD;
474   ierr = ProcessOptions(comm, &user);CHKERRQ(ierr);
475 
476   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
477      Create Particle-Mesh
478      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
479   ierr = CreateMesh(comm, &dm, &user);CHKERRQ(ierr);
480   ierr = CreateParticles(dm, &sw, &user);CHKERRQ(ierr);
481   ierr = DMSetApplicationContext(sw, &user);CHKERRQ(ierr);
482 
483   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
484      Setup timestepping solver
485      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
486   ierr = TSCreate(comm, &ts);CHKERRQ(ierr);
487   ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
488 
489   ierr = TSSetDM(ts, sw);CHKERRQ(ierr);
490   ierr = TSSetMaxTime(ts, 0.1);CHKERRQ(ierr);
491   ierr = TSSetTimeStep(ts, 0.00001);CHKERRQ(ierr);
492   ierr = TSSetMaxSteps(ts, 100);CHKERRQ(ierr);
493   ierr = TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
494   if (user.monitor) {ierr = TSMonitorSet(ts, Monitor, &user, NULL);CHKERRQ(ierr);}
495 
496   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
497      Prepare to solve
498      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
499   ierr = DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
500   ierr = VecGetLocalSize(u, &n);CHKERRQ(ierr);
501   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
502   ierr = TSSetComputeInitialCondition(ts, InitializeSolve);CHKERRQ(ierr);
503   ierr = TSSetComputeExactError(ts, ComputeError);CHKERRQ(ierr);
504 
505   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
506      Define function F(U, Udot , x , t) = G(U, x, t)
507      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
508 
509   /* - - - - - - - Basic Symplectic - - - - - - - - - - - - - - - - - - - - - -*/
510   ierr = ISCreateStride(comm, n/2, 0, 2, &is1);CHKERRQ(ierr);
511   ierr = ISCreateStride(comm, n/2, 1, 2, &is2);CHKERRQ(ierr);
512   ierr = TSRHSSplitSetIS(ts, "position", is1);CHKERRQ(ierr);
513   ierr = TSRHSSplitSetIS(ts, "momentum", is2);CHKERRQ(ierr);
514   ierr = ISDestroy(&is1);CHKERRQ(ierr);
515   ierr = ISDestroy(&is2);CHKERRQ(ierr);
516   ierr = TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunction1, &user);CHKERRQ(ierr);
517   ierr = TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunction2, &user);CHKERRQ(ierr);
518 
519   /* - - - - - - - Theta (Implicit Midpoint) - - - - - - - - - - - - - - - - -*/
520   ierr = TSSetRHSFunction(ts, NULL, RHSFunction, &user);CHKERRQ(ierr);
521 
522   ierr = MatCreate(PETSC_COMM_WORLD,&J);CHKERRQ(ierr);
523   ierr = MatSetSizes(J,PETSC_DECIDE,PETSC_DECIDE,n,n);CHKERRQ(ierr);
524   ierr = MatSetFromOptions(J);CHKERRQ(ierr);
525   ierr = MatSetUp(J);CHKERRQ(ierr);
526   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
527   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
528   ierr = PetscPrintf(comm, "n = %d\n",n);CHKERRQ(ierr);//Check number of particles
529   ierr = TSSetRHSJacobian(ts,J,J,RHSJacobian,&user);CHKERRQ(ierr);
530 
531   /* - - - - - - - Discrete Gradients - - - - - - - - - - - - - - - - - - - - */
532   ierr = TSDiscGradSetFormulation(ts, Sfunc, Ffunc, gradFfunc, &user);CHKERRQ(ierr);
533 
534   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
535      Solve
536      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
537 
538   ierr = TSComputeInitialCondition(ts, u);CHKERRQ(ierr);
539   ierr = TSSolve(ts, u);CHKERRQ(ierr);
540 
541   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
542      Clean up workspace
543      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
544   ierr = MatDestroy(&J);CHKERRQ(ierr);
545   ierr = DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
546   ierr = TSDestroy(&ts);CHKERRQ(ierr);
547   ierr = DMDestroy(&sw);CHKERRQ(ierr);
548   ierr = DMDestroy(&dm);CHKERRQ(ierr);
549   ierr = PetscFinalize();
550   return ierr;
551 }
552 
553 /*TEST
554 
555    build:
556      requires: triangle !single !complex
557    test:
558      suffix: 1
559      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
560 
561    test:
562      suffix: 2
563      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
564 
565    test:
566      suffix: 3
567      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
568 
569    test:
570      suffix: 4
571      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
572 
573    test:
574      suffix: 5
575      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
576 
577    test:
578      suffix: 6
579      args: -dm_plex_box_faces 1,1 -ts_type discgrad -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2
580 
581 TEST*/
582