xref: /petsc/src/dm/impls/swarm/tests/ex4.c (revision d57bb9db46f99bff4aaa690e4e9a7e55805ae415)
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   PetscFunctionBeginUser;
226   ierr = PetscObjectGetComm((PetscObject) ts, &comm);CHKERRQ(ierr);
227   ierr = TSGetDM(ts, &sdm);CHKERRQ(ierr);
228   ierr = DMGetApplicationContext(sdm, (void **) &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   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
359   //ierr = PetscPrintf(PETSC_COMM_WORLD, "# Particles (Np) = %d \n" , Np);CHKERRQ(ierr);
360 
361   ierr = MatGetOwnershipRange(J, &m, &n);CHKERRQ(ierr);
362   for (i = 0; i < Np; ++i) {
363     const PetscInt rows[2] = {2*i, 2*i+1};
364     ierr = MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
365   }
366   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
367   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
368   //ierr = MatView(J,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
369 
370   PetscFunctionReturn(0);
371 
372 }
373 /*----------------------------------------------------------------------------*/
374 
375 /*----------------Define S, F, G Functions (Discrete Gradients)---------------*/
376 /*
377   u_t = S * gradF
378     --or--
379   u_t = S * G
380 
381   + Sfunc - constructor for the S matrix from the formulation
382   . Ffunc - functional F from the formulation
383   - Gfunc - constructor for the gradient of F from the formulation
384 */
385 
386 PetscErrorCode Sfunc(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
387 {
388   AppCtx            *user = (AppCtx *) ctx;
389   PetscInt           Np = user->dim * user->particlesPerCell;
390   PetscInt           i, m, n;
391   const PetscScalar *u;
392   PetscScalar        vals[4] = {0., 1., -1, 0.};
393   PetscErrorCode     ierr;
394 
395   PetscFunctionBeginUser;
396   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
397   ierr = MatGetOwnershipRange(S, &m, &n);CHKERRQ(ierr);
398   for (i = 0; i < Np; ++i) {
399     const PetscInt rows[2] = {2*i, 2*i+1};
400     ierr = MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES);CHKERRQ(ierr);
401   }
402   ierr = MatAssemblyBegin(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
403   ierr = MatAssemblyEnd(S,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
404 
405   PetscFunctionReturn(0);
406 }
407 
408 PetscErrorCode Ffunc(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
409 {
410   AppCtx            *user = (AppCtx *) ctx;
411   DM                 dm;
412   const PetscScalar *u;
413   PetscInt           Np = user->dim * user->particlesPerCell;
414   PetscInt           p;
415   PetscErrorCode     ierr;
416 
417   PetscFunctionBeginUser;
418   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
419   //Define F
420   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
421   for (p = 0; p < Np; ++p) {
422     *F += (1/2)*PetscSqr(user->omega)*PetscSqr(u[p*2+0]) + (1/2)*PetscSqr(u[p*2+1]);
423   }
424   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
425 
426   PetscFunctionReturn(0);
427 }
428 
429 PetscErrorCode gradFfunc(TS ts, PetscReal t, Vec U, Vec gradF, void *ctx)
430 {
431   AppCtx            *user = (AppCtx *) ctx;
432   DM                 dm;
433   const PetscScalar *u;
434   PetscScalar       *g;
435   PetscInt           Np = user->dim * user->particlesPerCell;
436   PetscInt           p;
437   PetscErrorCode     ierr;
438 
439   PetscFunctionBeginUser;
440   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
441   ierr = VecGetArrayRead(U, &u);CHKERRQ(ierr);
442 
443   //Define gradF
444   ierr = VecGetArray(gradF, &g);CHKERRQ(ierr);
445   for (p = 0; p < Np; ++p) {
446     g[p*2+0] = PetscSqr(user->omega)*u[p*2+0]; /*dF/dx*/
447     g[p*2+1] = u[p*2+1]; /*dF/dv*/
448   }
449   ierr = VecRestoreArrayRead(U, &u);CHKERRQ(ierr);
450   ierr = VecRestoreArray(gradF, &g);CHKERRQ(ierr);
451 
452   PetscFunctionReturn(0);
453 }
454 /*----------------------------------------------------------------------------*/
455 
456 int main(int argc,char **argv)
457 {
458   TS             ts;     /* nonlinear solver                 */
459   DM             dm, sw; /* Mesh and particle managers       */
460   Vec            u;      /* swarm vector                     */
461   PetscInt       n;      /* swarm vector size                */
462   IS             is1, is2;
463   MPI_Comm       comm;
464   Mat            J;      /* Jacobian matrix                  */
465   AppCtx         user;
466   PetscErrorCode ierr;
467 
468   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
469      Initialize program
470      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
471   ierr = PetscInitialize(&argc, &argv, NULL, help);if (ierr) return ierr;
472   comm = PETSC_COMM_WORLD;
473   ierr = ProcessOptions(comm, &user);CHKERRQ(ierr);
474 
475   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
476      Create Particle-Mesh
477      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
478   ierr = CreateMesh(comm, &dm, &user);CHKERRQ(ierr);
479   ierr = CreateParticles(dm, &sw, &user);CHKERRQ(ierr);
480   ierr = DMSetApplicationContext(sw, &user);CHKERRQ(ierr);
481 
482   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
483      Setup timestepping solver
484      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
485   ierr = TSCreate(comm, &ts);CHKERRQ(ierr);
486   ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
487 
488   ierr = TSSetDM(ts, sw);CHKERRQ(ierr);
489   ierr = TSSetMaxTime(ts, 0.1);CHKERRQ(ierr);
490   ierr = TSSetTimeStep(ts, 0.00001);CHKERRQ(ierr);
491   ierr = TSSetMaxSteps(ts, 100);CHKERRQ(ierr);
492   ierr = TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
493   if (user.monitor) {ierr = TSMonitorSet(ts, Monitor, &user, NULL);CHKERRQ(ierr);}
494 
495   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
496      Prepare to solve
497      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
498   ierr = DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
499   ierr = VecGetLocalSize(u, &n);CHKERRQ(ierr);
500   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
501   ierr = TSSetComputeInitialCondition(ts, InitializeSolve);CHKERRQ(ierr);
502   ierr = TSSetComputeExactError(ts, ComputeError);CHKERRQ(ierr);
503 
504   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
505      Define function F(U, Udot , x , t) = G(U, x, t)
506      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
507 
508   /* - - - - - - - Basic Symplectic - - - - - - - - - - - - - - - - - - - - - -*/
509   ierr = ISCreateStride(comm, n/2, 0, 2, &is1);CHKERRQ(ierr);
510   ierr = ISCreateStride(comm, n/2, 1, 2, &is2);CHKERRQ(ierr);
511   ierr = TSRHSSplitSetIS(ts, "position", is1);CHKERRQ(ierr);
512   ierr = TSRHSSplitSetIS(ts, "momentum", is2);CHKERRQ(ierr);
513   ierr = ISDestroy(&is1);CHKERRQ(ierr);
514   ierr = ISDestroy(&is2);CHKERRQ(ierr);
515   ierr = TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunction1, &user);CHKERRQ(ierr);
516   ierr = TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunction2, &user);CHKERRQ(ierr);
517 
518   /* - - - - - - - Theta (Implicit Midpoint) - - - - - - - - - - - - - - - - -*/
519   ierr = TSSetRHSFunction(ts, NULL, RHSFunction, &user);CHKERRQ(ierr);
520 
521   ierr = MatCreate(PETSC_COMM_WORLD,&J);CHKERRQ(ierr);
522   ierr = MatSetSizes(J,PETSC_DECIDE,PETSC_DECIDE,n,n);CHKERRQ(ierr);
523   ierr = MatSetFromOptions(J);CHKERRQ(ierr);
524   ierr = MatSetUp(J);CHKERRQ(ierr);
525   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
526   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
527   ierr = PetscPrintf(comm, "n = %d\n",n);CHKERRQ(ierr);//Check number of particles
528   ierr = TSSetRHSJacobian(ts,J,J,RHSJacobian,&user);CHKERRQ(ierr);
529 
530   /* - - - - - - - Discrete Gradients - - - - - - - - - - - - - - - - - - - - */
531   ierr = TSDiscGradSetFormulation(ts, Sfunc, Ffunc, gradFfunc, &user);CHKERRQ(ierr);
532 
533   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
534      Solve
535      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
536 
537   ierr = TSComputeInitialCondition(ts, u);CHKERRQ(ierr);
538   ierr = TSSolve(ts, u);CHKERRQ(ierr);
539 
540   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
541      Clean up workspace
542      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
543   ierr = MatDestroy(&J);CHKERRQ(ierr);
544   ierr = DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &u);CHKERRQ(ierr);
545   ierr = TSDestroy(&ts);CHKERRQ(ierr);
546   ierr = DMDestroy(&sw);CHKERRQ(ierr);
547   ierr = DMDestroy(&dm);CHKERRQ(ierr);
548   ierr = PetscFinalize();
549   return ierr;
550 }
551 
552 /*TEST
553 
554    build:
555      requires: triangle !single !complex
556    test:
557      suffix: 1
558      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
559 
560    test:
561      suffix: 2
562      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
563 
564    test:
565      suffix: 3
566      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
567 
568    test:
569      suffix: 4
570      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
571 
572    test:
573      suffix: 5
574      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
575 
576    test:
577      suffix: 6
578      args: -dm_plex_box_faces 1,1 -ts_type discgrad -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2
579 
580 TEST*/
581