xref: /petsc/src/dm/impls/swarm/tests/ex4.c (revision 9566063d113dddea24716c546802770db7481bc0)
1db4653c2SDaniel Finn static char help[] = "Comparing basic symplectic, theta and discrete gradients interators on a simple hamiltonian system (harmonic oscillator) with particles\n";
2c4762a1bSJed Brown 
3c4762a1bSJed Brown #include <petscdmplex.h>
4c4762a1bSJed Brown #include <petsc/private/dmpleximpl.h>  /* For norm */
5c4762a1bSJed Brown #include <petsc/private/petscfeimpl.h> /* For CoordinatesRefToReal() */
6c4762a1bSJed Brown #include <petscdmswarm.h>
7c4762a1bSJed Brown #include <petscts.h>
8c4762a1bSJed Brown 
9c4762a1bSJed Brown typedef struct {
10db4653c2SDaniel Finn   PetscInt  dim;                          /* The topological mesh dimension */
11db4653c2SDaniel Finn   PetscBool simplex;                      /* Flag for simplices or tensor cells */
12db4653c2SDaniel Finn   char      filename[PETSC_MAX_PATH_LEN]; /* Name of the mesh filename if any */
13c4762a1bSJed Brown   PetscReal omega;                        /* Oscillation frequency omega */
14c4762a1bSJed Brown   PetscInt  particlesPerCell;             /* The number of partices per cell */
15db4653c2SDaniel Finn   PetscInt  numberOfCells;                /* Number of cells in mesh */
16c4762a1bSJed Brown   PetscReal momentTol;                    /* Tolerance for checking moment conservation */
17c4762a1bSJed Brown   PetscBool monitor;                      /* Flag for using the TS monitor */
18c4762a1bSJed Brown   PetscBool error;                        /* Flag for printing the error */
19c4762a1bSJed Brown   PetscInt  ostep;                        /* print the energy at each ostep time steps */
20c4762a1bSJed Brown } AppCtx;
21c4762a1bSJed Brown 
22c4762a1bSJed Brown static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
23c4762a1bSJed Brown {
24c4762a1bSJed Brown   PetscErrorCode ierr;
25c4762a1bSJed Brown 
26c4762a1bSJed Brown   PetscFunctionBeginUser;
27db4653c2SDaniel Finn   options->dim              = 2;
28db4653c2SDaniel Finn   options->simplex          = PETSC_TRUE;
29c4762a1bSJed Brown   options->monitor          = PETSC_FALSE;
30c4762a1bSJed Brown   options->error            = PETSC_FALSE;
31c4762a1bSJed Brown   options->particlesPerCell = 1;
32db4653c2SDaniel Finn   options->numberOfCells    = 2;
33c4762a1bSJed Brown   options->momentTol        = 100.0*PETSC_MACHINE_EPSILON;
34c4762a1bSJed Brown   options->omega            = 64.0;
35c4762a1bSJed Brown   options->ostep            = 100;
36c4762a1bSJed Brown 
37*9566063dSJacob Faibussowitsch   PetscCall(PetscStrcpy(options->filename, ""));
38db4653c2SDaniel Finn 
39*9566063dSJacob Faibussowitsch   ierr = PetscOptionsBegin(comm, "", "Harmonic Oscillator Options", "DMPLEX");PetscCall(ierr);
40*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex4.c", options->ostep, &options->ostep, PETSC_NULL));
41*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-dim", "The topological mesh dimension", "ex4.c", options->dim, &options->dim, NULL));
42*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-monitor", "Flag to use the TS monitor", "ex4.c", options->monitor, &options->monitor, NULL));
43*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex4.c", options->error, &options->error, NULL));
44*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-simplex", "The flag for simplices or tensor cells", "ex4.c", options->simplex, &options->simplex, NULL));
45*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsString("-mesh", "Name of the mesh filename if any", "ex4.c", options->filename, options->filename, sizeof(options->filename), NULL));
46*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-particles_per_cell", "Number of particles per cell", "ex4.c", options->particlesPerCell, &options->particlesPerCell, NULL));
47*9566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-omega", "Oscillator frequency", "ex4.c", options->omega, &options->omega, PETSC_NULL));
48*9566063dSJacob Faibussowitsch   ierr = PetscOptionsEnd();PetscCall(ierr);
49c4762a1bSJed Brown 
50c4762a1bSJed Brown   PetscFunctionReturn(0);
51db4653c2SDaniel Finn 
52c4762a1bSJed Brown }
53c4762a1bSJed Brown 
54c4762a1bSJed Brown static PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *user)
55c4762a1bSJed Brown {
56c4762a1bSJed Brown   PetscFunctionBeginUser;
57*9566063dSJacob Faibussowitsch   PetscCall(DMCreate(comm, dm));
58*9566063dSJacob Faibussowitsch   PetscCall(DMSetType(*dm, DMPLEX));
59*9566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(*dm));
60*9566063dSJacob Faibussowitsch   PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view"));
61db4653c2SDaniel Finn 
62c4762a1bSJed Brown   PetscFunctionReturn(0);
63c4762a1bSJed Brown }
64c4762a1bSJed Brown 
65c4762a1bSJed Brown static PetscErrorCode SetInitialCoordinates(DM dmSw)
66c4762a1bSJed Brown {
67c4762a1bSJed Brown   DM             dm;
68c4762a1bSJed Brown   AppCtx        *user;
69c4762a1bSJed Brown   PetscRandom    rnd;
70c4762a1bSJed Brown   PetscBool      simplex;
71c4762a1bSJed Brown   PetscReal     *centroid, *coords, *xi0, *v0, *J, *invJ, detJ;
72c4762a1bSJed Brown   PetscInt       dim, d, cStart, cEnd, c, Np, p;
73c4762a1bSJed Brown 
74c4762a1bSJed Brown   PetscFunctionBeginUser;
75*9566063dSJacob Faibussowitsch   PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject) dmSw), &rnd));
76*9566063dSJacob Faibussowitsch   PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
77*9566063dSJacob Faibussowitsch   PetscCall(PetscRandomSetFromOptions(rnd));
78c4762a1bSJed Brown 
79*9566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(dmSw, &user));
80c4762a1bSJed Brown   Np   = user->particlesPerCell;
81*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dmSw, &dm));
82*9566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(dm, &dim));
83*9566063dSJacob Faibussowitsch   PetscCall(DMPlexIsSimplex(dm, &simplex));
84*9566063dSJacob Faibussowitsch   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
85db4653c2SDaniel Finn   user->numberOfCells = cEnd - cStart;
86*9566063dSJacob Faibussowitsch   PetscCall(PetscMalloc5(dim, &centroid, dim, &xi0, dim, &v0, dim*dim, &J, dim*dim, &invJ));
87c4762a1bSJed Brown   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
88*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
89c4762a1bSJed Brown   for (c = cStart; c < cEnd; ++c) {
90c4762a1bSJed Brown     if (Np == 1) {
91*9566063dSJacob Faibussowitsch       PetscCall(DMPlexComputeCellGeometryFVM(dm, c, NULL, centroid, NULL));
92c4762a1bSJed Brown       for (d = 0; d < dim; ++d) coords[c*dim+d] = centroid[d];
93c4762a1bSJed Brown     } else {
94*9566063dSJacob Faibussowitsch       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ)); /* affine */
95c4762a1bSJed Brown       for (p = 0; p < Np; ++p) {
96c4762a1bSJed Brown         const PetscInt n   = c*Np + p;
97c4762a1bSJed Brown         PetscReal      sum = 0.0, refcoords[3];
98c4762a1bSJed Brown 
99c4762a1bSJed Brown         for (d = 0; d < dim; ++d) {
100*9566063dSJacob Faibussowitsch           PetscCall(PetscRandomGetValueReal(rnd, &refcoords[d]));
101c4762a1bSJed Brown           sum += refcoords[d];
102c4762a1bSJed Brown         }
103c4762a1bSJed Brown         if (simplex && sum > 0.0) for (d = 0; d < dim; ++d) refcoords[d] -= PetscSqrtReal(dim)*sum;
104c4762a1bSJed Brown         CoordinatesRefToReal(dim, dim, xi0, v0, J, refcoords, &coords[n*dim]);
105c4762a1bSJed Brown       }
106c4762a1bSJed Brown     }
107c4762a1bSJed Brown   }
108db4653c2SDaniel Finn 
109*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
110*9566063dSJacob Faibussowitsch   PetscCall(PetscFree5(centroid, xi0, v0, J, invJ));
111*9566063dSJacob Faibussowitsch   PetscCall(PetscRandomDestroy(&rnd));
112c4762a1bSJed Brown   PetscFunctionReturn(0);
113c4762a1bSJed Brown }
114c4762a1bSJed Brown 
115c4762a1bSJed Brown static PetscErrorCode SetInitialConditions(DM dmSw, Vec u)
116c4762a1bSJed Brown {
117c4762a1bSJed Brown   DM             dm;
118c4762a1bSJed Brown   AppCtx        *user;
119c4762a1bSJed Brown   PetscReal     *coords;
120c4762a1bSJed Brown   PetscScalar   *initialConditions;
121c4762a1bSJed Brown   PetscInt       dim, cStart, cEnd, c, Np, p;
122c4762a1bSJed Brown 
123c4762a1bSJed Brown   PetscFunctionBeginUser;
124*9566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(dmSw, &user));
125c4762a1bSJed Brown   Np   = user->particlesPerCell;
126*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dmSw, &dm));
127*9566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(dm, &dim));
128*9566063dSJacob Faibussowitsch   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
129*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
130*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(u, &initialConditions));
131c4762a1bSJed Brown   for (c = cStart; c < cEnd; ++c) {
132c4762a1bSJed Brown     for (p = 0; p < Np; ++p) {
133c4762a1bSJed Brown       const PetscInt n = c*Np + p;
134c4762a1bSJed Brown       initialConditions[n*2+0] = DMPlex_NormD_Internal(dim, &coords[n*dim]);
135c4762a1bSJed Brown       initialConditions[n*2+1] = 0.0;
136c4762a1bSJed Brown     }
137c4762a1bSJed Brown   }
138*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(u, &initialConditions));
139*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dmSw, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
140c4762a1bSJed Brown   PetscFunctionReturn(0);
141c4762a1bSJed Brown }
142c4762a1bSJed Brown 
143c4762a1bSJed Brown static PetscErrorCode CreateParticles(DM dm, DM *sw, AppCtx *user)
144c4762a1bSJed Brown {
145c4762a1bSJed Brown   PetscInt      *cellid;
146db4653c2SDaniel Finn   PetscInt       dim, cStart, cEnd, c, Np, p;
147c4762a1bSJed Brown 
148c4762a1bSJed Brown   PetscFunctionBeginUser;
149*9566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(dm, &dim));
150*9566063dSJacob Faibussowitsch   PetscCall(DMCreate(PetscObjectComm((PetscObject) dm), sw));
151*9566063dSJacob Faibussowitsch   PetscCall(DMSetType(*sw, DMSWARM));
152*9566063dSJacob Faibussowitsch   PetscCall(DMSetDimension(*sw, dim));
153db4653c2SDaniel Finn   Np = user->particlesPerCell;
154*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
155*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmSetCellDM(*sw, dm));
156*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "kinematics", 2, PETSC_REAL));
157*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmFinalizeFieldRegister(*sw));
158*9566063dSJacob Faibussowitsch   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
159*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmSetLocalSizes(*sw, (cEnd - cStart) * Np, 0));
160*9566063dSJacob Faibussowitsch   PetscCall(DMSetFromOptions(*sw));
161*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid));
162c4762a1bSJed Brown   for (c = cStart; c < cEnd; ++c) {
163c4762a1bSJed Brown     for (p = 0; p < Np; ++p) {
164c4762a1bSJed Brown       const PetscInt n = c*Np + p;
165c4762a1bSJed Brown 
166c4762a1bSJed Brown       cellid[n] = c;
167c4762a1bSJed Brown     }
168c4762a1bSJed Brown   }
169*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(*sw, DMSwarmPICField_cellid, NULL, NULL, (void **) &cellid));
170*9566063dSJacob Faibussowitsch   PetscCall(PetscObjectSetName((PetscObject) *sw, "Particles"));
171*9566063dSJacob Faibussowitsch   PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view"));
172c4762a1bSJed Brown   PetscFunctionReturn(0);
173c4762a1bSJed Brown }
174c4762a1bSJed Brown 
175c4762a1bSJed Brown static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx)
176c4762a1bSJed Brown {
177c4762a1bSJed Brown   AppCtx            *user  = (AppCtx *) ctx;
178c4762a1bSJed Brown   const PetscReal    omega = user->omega;
179c4762a1bSJed Brown   const PetscScalar *u;
180c4762a1bSJed Brown   MPI_Comm           comm;
181c4762a1bSJed Brown   PetscReal          dt;
182c4762a1bSJed Brown   PetscInt           Np, p;
183c4762a1bSJed Brown 
184c4762a1bSJed Brown   PetscFunctionBeginUser;
185c4762a1bSJed Brown   if (step%user->ostep == 0) {
186*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetComm((PetscObject) ts, &comm));
187*9566063dSJacob Faibussowitsch     if (!step) PetscCall(PetscPrintf(comm, "Time     Step Part     Energy Mod Energy\n"));
188*9566063dSJacob Faibussowitsch     PetscCall(TSGetTimeStep(ts, &dt));
189*9566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(U, &u));
190*9566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(U, &Np));
191c4762a1bSJed Brown     Np /= 2;
192c4762a1bSJed Brown     for (p = 0; p < Np; ++p) {
193c4762a1bSJed Brown       const PetscReal x  = PetscRealPart(u[p*2+0]);
194c4762a1bSJed Brown       const PetscReal v  = PetscRealPart(u[p*2+1]);
195c4762a1bSJed Brown       const PetscReal E  = 0.5*(v*v + PetscSqr(omega)*x*x);
196c4762a1bSJed Brown       const PetscReal mE = 0.5*(v*v + PetscSqr(omega)*x*x - PetscSqr(omega)*dt*x*v);
197*9566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(comm, "%.6lf %4D %4D %10.4lf %10.4lf\n", t, step, p, (double) E, (double) mE));
198c4762a1bSJed Brown     }
199*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(U, &u));
200c4762a1bSJed Brown   }
201c4762a1bSJed Brown   PetscFunctionReturn(0);
202c4762a1bSJed Brown }
203c4762a1bSJed Brown 
204c4762a1bSJed Brown static PetscErrorCode InitializeSolve(TS ts, Vec u)
205c4762a1bSJed Brown {
206c4762a1bSJed Brown   DM             dm;
207c4762a1bSJed Brown   AppCtx        *user;
208c4762a1bSJed Brown 
209c4762a1bSJed Brown   PetscFunctionBeginUser;
210*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
211*9566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(dm, &user));
212*9566063dSJacob Faibussowitsch   PetscCall(SetInitialCoordinates(dm));
213*9566063dSJacob Faibussowitsch   PetscCall(SetInitialConditions(dm, u));
214c4762a1bSJed Brown   PetscFunctionReturn(0);
215c4762a1bSJed Brown }
216c4762a1bSJed Brown 
217c4762a1bSJed Brown static PetscErrorCode ComputeError(TS ts, Vec U, Vec E)
218c4762a1bSJed Brown {
219c4762a1bSJed Brown   MPI_Comm           comm;
220c4762a1bSJed Brown   DM                 sdm;
221c4762a1bSJed Brown   AppCtx            *user;
222c4762a1bSJed Brown   const PetscScalar *u, *coords;
223c4762a1bSJed Brown   PetscScalar       *e;
224c4762a1bSJed Brown   PetscReal          t, omega;
225c4762a1bSJed Brown   PetscInt           dim, Np, p;
226c4762a1bSJed Brown 
227c4762a1bSJed Brown   PetscFunctionBeginUser;
228*9566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject) ts, &comm));
229*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &sdm));
230*9566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(sdm, &user));
231c4762a1bSJed Brown   omega = user->omega;
232*9566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(sdm, &dim));
233*9566063dSJacob Faibussowitsch   PetscCall(TSGetSolveTime(ts, &t));
234*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(E, &e));
235*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
236*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
237c4762a1bSJed Brown   Np  /= 2;
238*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
239c4762a1bSJed Brown   for (p = 0; p < Np; ++p) {
240c4762a1bSJed Brown     const PetscReal x  = PetscRealPart(u[p*2+0]);
241c4762a1bSJed Brown     const PetscReal v  = PetscRealPart(u[p*2+1]);
242c4762a1bSJed Brown     const PetscReal x0 = DMPlex_NormD_Internal(dim, &coords[p*dim]);
243c4762a1bSJed Brown     const PetscReal ex =  x0*PetscCosReal(omega*t);
244c4762a1bSJed Brown     const PetscReal ev = -x0*omega*PetscSinReal(omega*t);
245c4762a1bSJed Brown 
246*9566063dSJacob Faibussowitsch     if (user->error) PetscCall(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)));
247c4762a1bSJed Brown     e[p*2+0] = x - ex;
248c4762a1bSJed Brown     e[p*2+1] = v - ev;
249c4762a1bSJed Brown   }
250*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sdm, DMSwarmPICField_coor, NULL, NULL, (void **) &coords));
25140be0ff1SMatthew G. Knepley 
252*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U, &u));
253*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(E, &e));
254c4762a1bSJed Brown   PetscFunctionReturn(0);
255c4762a1bSJed Brown }
256c4762a1bSJed Brown 
25740be0ff1SMatthew G. Knepley /*---------------------Create particle RHS Functions--------------------------*/
25840be0ff1SMatthew G. Knepley static PetscErrorCode RHSFunction1(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx)
25940be0ff1SMatthew G. Knepley {
26040be0ff1SMatthew G. Knepley   const PetscScalar *v;
26140be0ff1SMatthew G. Knepley   PetscScalar       *xres;
26240be0ff1SMatthew G. Knepley   PetscInt          Np, p;
26340be0ff1SMatthew G. Knepley 
26440be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
265*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Xres, &xres));
266*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(V, &v));
267*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(Xres, &Np));
26840be0ff1SMatthew G. Knepley   for (p = 0; p < Np; ++p) {
26940be0ff1SMatthew G. Knepley      xres[p] = v[p];
27040be0ff1SMatthew G. Knepley   }
271*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(V, &v));
272*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Xres, &xres));
27340be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
27440be0ff1SMatthew G. Knepley }
27540be0ff1SMatthew G. Knepley 
27640be0ff1SMatthew G. Knepley static PetscErrorCode RHSFunction2(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx)
27740be0ff1SMatthew G. Knepley {
27840be0ff1SMatthew G. Knepley   AppCtx            *user = (AppCtx *)ctx;
27940be0ff1SMatthew G. Knepley   const PetscScalar *x;
28040be0ff1SMatthew G. Knepley   PetscScalar       *vres;
28140be0ff1SMatthew G. Knepley   PetscInt          Np, p;
28240be0ff1SMatthew G. Knepley 
28340be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
284*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Vres, &vres));
285*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(X, &x));
286*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(Vres, &Np));
28740be0ff1SMatthew G. Knepley   for (p = 0; p < Np; ++p) {
28840be0ff1SMatthew G. Knepley     vres[p] = -PetscSqr(user->omega)*x[p];
28940be0ff1SMatthew G. Knepley   }
290*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(X, &x));
291*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Vres, &vres));
29240be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
29340be0ff1SMatthew G. Knepley }
29440be0ff1SMatthew G. Knepley 
29540be0ff1SMatthew G. Knepley /*----------------------------------------------------------------------------*/
29640be0ff1SMatthew G. Knepley 
29740be0ff1SMatthew G. Knepley /*--------------------Define RHSFunction, RHSJacobian (Theta)-----------------*/
29840be0ff1SMatthew G. Knepley static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx)
29940be0ff1SMatthew G. Knepley {
30040be0ff1SMatthew G. Knepley   AppCtx            *user = (AppCtx *) ctx;
30140be0ff1SMatthew G. Knepley   DM                dm;
30240be0ff1SMatthew G. Knepley   const PetscScalar *u;
30340be0ff1SMatthew G. Knepley   PetscScalar       *g;
30440be0ff1SMatthew G. Knepley   PetscInt          Np, p;
30540be0ff1SMatthew G. Knepley 
30640be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
307*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
308*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(G, &g));
309*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
310*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
31140be0ff1SMatthew G. Knepley   Np  /= 2;
31240be0ff1SMatthew G. Knepley   for (p = 0; p < Np; ++p) {
31340be0ff1SMatthew G. Knepley     g[p*2+0] = u[p*2+1];
31440be0ff1SMatthew G. Knepley     g[p*2+1] = -PetscSqr(user->omega)*u[p*2+0];
31540be0ff1SMatthew G. Knepley   }
316*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U, &u));
317*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(G, &g));
31840be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
31940be0ff1SMatthew G. Knepley }
32040be0ff1SMatthew G. Knepley 
32140be0ff1SMatthew G. Knepley /*Ji = dFi/dxj
32240be0ff1SMatthew G. Knepley J= (0    1)
32340be0ff1SMatthew G. Knepley    (-w^2 0)
32440be0ff1SMatthew G. Knepley */
32540be0ff1SMatthew G. Knepley static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U , Mat J, Mat P, void *ctx)
32640be0ff1SMatthew G. Knepley {
32740be0ff1SMatthew G. Knepley   AppCtx             *user = (AppCtx *) ctx;
328db4653c2SDaniel Finn   PetscInt           Np;
32940be0ff1SMatthew G. Knepley   PetscInt           i, m, n;
33040be0ff1SMatthew G. Knepley   const PetscScalar *u;
33140be0ff1SMatthew G. Knepley   PetscScalar        vals[4] = {0., 1., -PetscSqr(user->omega), 0.};
33240be0ff1SMatthew G. Knepley 
333362febeeSStefano Zampini   PetscFunctionBeginUser;
334*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
335*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
336db4653c2SDaniel Finn   Np /= 2;
337*9566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(J, &m, &n));
33840be0ff1SMatthew G. Knepley   for (i = 0; i < Np; ++i) {
33940be0ff1SMatthew G. Knepley     const PetscInt rows[2] = {2*i, 2*i+1};
340*9566063dSJacob Faibussowitsch     PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES));
34140be0ff1SMatthew G. Knepley   }
342*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
343*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY));
34440be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
34540be0ff1SMatthew G. Knepley 
34640be0ff1SMatthew G. Knepley }
347db4653c2SDaniel Finn 
34840be0ff1SMatthew G. Knepley /*----------------------------------------------------------------------------*/
34940be0ff1SMatthew G. Knepley 
35040be0ff1SMatthew G. Knepley /*----------------Define S, F, G Functions (Discrete Gradients)---------------*/
35140be0ff1SMatthew G. Knepley /*
35240be0ff1SMatthew G. Knepley   u_t = S * gradF
35340be0ff1SMatthew G. Knepley     --or--
35440be0ff1SMatthew G. Knepley   u_t = S * G
35540be0ff1SMatthew G. Knepley 
35640be0ff1SMatthew G. Knepley   + Sfunc - constructor for the S matrix from the formulation
35740be0ff1SMatthew G. Knepley   . Ffunc - functional F from the formulation
35840be0ff1SMatthew G. Knepley   - Gfunc - constructor for the gradient of F from the formulation
35940be0ff1SMatthew G. Knepley */
36040be0ff1SMatthew G. Knepley 
36140be0ff1SMatthew G. Knepley PetscErrorCode Sfunc(TS ts, PetscReal t, Vec U, Mat S, void *ctx)
36240be0ff1SMatthew G. Knepley {
363db4653c2SDaniel Finn   PetscInt           Np;
36440be0ff1SMatthew G. Knepley   PetscInt           i, m, n;
36540be0ff1SMatthew G. Knepley   const PetscScalar *u;
36640be0ff1SMatthew G. Knepley   PetscScalar        vals[4] = {0., 1., -1, 0.};
36740be0ff1SMatthew G. Knepley 
36840be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
369*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
370*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
371db4653c2SDaniel Finn   Np /= 2;
372*9566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(S, &m, &n));
37340be0ff1SMatthew G. Knepley   for (i = 0; i < Np; ++i) {
37440be0ff1SMatthew G. Knepley     const PetscInt rows[2] = {2*i, 2*i+1};
375*9566063dSJacob Faibussowitsch     PetscCall(MatSetValues(S, 2, rows, 2, rows, vals, INSERT_VALUES));
37640be0ff1SMatthew G. Knepley   }
377*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(S,MAT_FINAL_ASSEMBLY));
378*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(S,MAT_FINAL_ASSEMBLY));
37940be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
380db4653c2SDaniel Finn 
38140be0ff1SMatthew G. Knepley }
38240be0ff1SMatthew G. Knepley 
38340be0ff1SMatthew G. Knepley PetscErrorCode Ffunc(TS ts, PetscReal t, Vec U, PetscScalar *F, void *ctx)
38440be0ff1SMatthew G. Knepley {
38540be0ff1SMatthew G. Knepley   AppCtx            *user = (AppCtx *) ctx;
38640be0ff1SMatthew G. Knepley   DM                 dm;
38740be0ff1SMatthew G. Knepley   const PetscScalar *u;
388db4653c2SDaniel Finn   PetscInt           Np;
38940be0ff1SMatthew G. Knepley   PetscInt           p;
39040be0ff1SMatthew G. Knepley 
39140be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
392*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
393db4653c2SDaniel Finn 
394db4653c2SDaniel Finn   /* Define F */
395*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
396*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
397db4653c2SDaniel Finn   Np /= 2;
39840be0ff1SMatthew G. Knepley   for (p = 0; p < Np; ++p) {
399db4653c2SDaniel Finn     *F += 0.5*PetscSqr(user->omega)*PetscSqr(u[p*2+0]) + 0.5*PetscSqr(u[p*2+1]);
40040be0ff1SMatthew G. Knepley   }
401*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U, &u));
40240be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
40340be0ff1SMatthew G. Knepley }
40440be0ff1SMatthew G. Knepley 
40540be0ff1SMatthew G. Knepley PetscErrorCode gradFfunc(TS ts, PetscReal t, Vec U, Vec gradF, void *ctx)
40640be0ff1SMatthew G. Knepley {
40740be0ff1SMatthew G. Knepley   AppCtx            *user = (AppCtx *) ctx;
40840be0ff1SMatthew G. Knepley   DM                dm;
40940be0ff1SMatthew G. Knepley   const PetscScalar *u;
41040be0ff1SMatthew G. Knepley   PetscScalar       *g;
411db4653c2SDaniel Finn   PetscInt          Np;
41240be0ff1SMatthew G. Knepley   PetscInt          p;
41340be0ff1SMatthew G. Knepley 
41440be0ff1SMatthew G. Knepley   PetscFunctionBeginUser;
415*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
416*9566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U, &u));
417*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(U, &Np));
418db4653c2SDaniel Finn   Np /= 2;
419db4653c2SDaniel Finn   /*Define gradF*/
420*9566063dSJacob Faibussowitsch   PetscCall(VecGetArray(gradF, &g));
42140be0ff1SMatthew G. Knepley   for (p = 0; p < Np; ++p) {
42240be0ff1SMatthew G. Knepley     g[p*2+0] = PetscSqr(user->omega)*u[p*2+0]; /*dF/dx*/
42340be0ff1SMatthew G. Knepley     g[p*2+1] = u[p*2+1]; /*dF/dv*/
42440be0ff1SMatthew G. Knepley   }
425*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U, &u));
426*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(gradF, &g));
42740be0ff1SMatthew G. Knepley   PetscFunctionReturn(0);
42840be0ff1SMatthew G. Knepley }
429db4653c2SDaniel Finn 
43040be0ff1SMatthew G. Knepley /*----------------------------------------------------------------------------*/
43140be0ff1SMatthew G. Knepley 
432c4762a1bSJed Brown int main(int argc,char **argv)
433c4762a1bSJed Brown {
434c4762a1bSJed Brown   TS             ts;     /* nonlinear solver                 */
435c4762a1bSJed Brown   DM             dm, sw; /* Mesh and particle managers       */
436c4762a1bSJed Brown   Vec            u;      /* swarm vector                     */
43740be0ff1SMatthew G. Knepley   PetscInt       n;      /* swarm vector size                */
438c4762a1bSJed Brown   IS             is1, is2;
439c4762a1bSJed Brown   MPI_Comm       comm;
44040be0ff1SMatthew G. Knepley   Mat            J;      /* Jacobian matrix                  */
441c4762a1bSJed Brown   AppCtx         user;
442c4762a1bSJed Brown 
44340be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
44440be0ff1SMatthew G. Knepley      Initialize program
44540be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
446*9566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
447c4762a1bSJed Brown   comm = PETSC_COMM_WORLD;
448*9566063dSJacob Faibussowitsch   PetscCall(ProcessOptions(comm, &user));
449c4762a1bSJed Brown 
45040be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
45140be0ff1SMatthew G. Knepley      Create Particle-Mesh
45240be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
453*9566063dSJacob Faibussowitsch   PetscCall(CreateMesh(comm, &dm, &user));
454*9566063dSJacob Faibussowitsch   PetscCall(CreateParticles(dm, &sw, &user));
455*9566063dSJacob Faibussowitsch   PetscCall(DMSetApplicationContext(sw, &user));
456c4762a1bSJed Brown 
45740be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
45840be0ff1SMatthew G. Knepley      Setup timestepping solver
45940be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
460*9566063dSJacob Faibussowitsch   PetscCall(TSCreate(comm, &ts));
461*9566063dSJacob Faibussowitsch   PetscCall(TSSetProblemType(ts,TS_NONLINEAR));
46240be0ff1SMatthew G. Knepley 
463*9566063dSJacob Faibussowitsch   PetscCall(TSSetDM(ts, sw));
464*9566063dSJacob Faibussowitsch   PetscCall(TSSetMaxTime(ts, 0.1));
465*9566063dSJacob Faibussowitsch   PetscCall(TSSetTimeStep(ts, 0.00001));
466*9566063dSJacob Faibussowitsch   PetscCall(TSSetMaxSteps(ts, 100));
467*9566063dSJacob Faibussowitsch   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
468*9566063dSJacob Faibussowitsch   if (user.monitor) PetscCall(TSMonitorSet(ts, Monitor, &user, NULL));
469c4762a1bSJed Brown 
47040be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
47140be0ff1SMatthew G. Knepley      Prepare to solve
47240be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
473*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "kinematics", &u));
474*9566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(u, &n));
475*9566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(ts));
476*9566063dSJacob Faibussowitsch   PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve));
477*9566063dSJacob Faibussowitsch   PetscCall(TSSetComputeExactError(ts, ComputeError));
47840be0ff1SMatthew G. Knepley 
47940be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
48040be0ff1SMatthew G. Knepley      Define function F(U, Udot , x , t) = G(U, x, t)
48140be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
48240be0ff1SMatthew G. Knepley 
48340be0ff1SMatthew G. Knepley   /* - - - - - - - Basic Symplectic - - - - - - - - - - - - - - - - - - - - - -*/
484*9566063dSJacob Faibussowitsch   PetscCall(ISCreateStride(comm, n/2, 0, 2, &is1));
485*9566063dSJacob Faibussowitsch   PetscCall(ISCreateStride(comm, n/2, 1, 2, &is2));
486*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitSetIS(ts, "position", is1));
487*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitSetIS(ts, "momentum", is2));
488*9566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&is1));
489*9566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&is2));
490*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunction1, &user));
491*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunction2, &user));
492c4762a1bSJed Brown 
49340be0ff1SMatthew G. Knepley   /* - - - - - - - Theta (Implicit Midpoint) - - - - - - - - - - - - - - - - -*/
494db4653c2SDaniel Finn 
495*9566063dSJacob Faibussowitsch   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, &user));
49640be0ff1SMatthew G. Knepley 
497*9566063dSJacob Faibussowitsch   PetscCall(MatCreate(PETSC_COMM_WORLD,&J));
498*9566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(J,PETSC_DECIDE,PETSC_DECIDE,n,n));
499*9566063dSJacob Faibussowitsch   PetscCall(MatSetFromOptions(J));
500*9566063dSJacob Faibussowitsch   PetscCall(MatSetUp(J));
501*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY));
502*9566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY));
503*9566063dSJacob Faibussowitsch   PetscCall(TSSetRHSJacobian(ts,J,J,RHSJacobian,&user));
50440be0ff1SMatthew G. Knepley 
50540be0ff1SMatthew G. Knepley   /* - - - - - - - Discrete Gradients - - - - - - - - - - - - - - - - - - - - */
506*9566063dSJacob Faibussowitsch   PetscCall(TSDiscGradSetFormulation(ts, Sfunc, Ffunc, gradFfunc, &user));
50740be0ff1SMatthew G. Knepley 
50840be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
50940be0ff1SMatthew G. Knepley      Solve
51040be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
511*9566063dSJacob Faibussowitsch   PetscCall(TSComputeInitialCondition(ts, u));
512*9566063dSJacob Faibussowitsch   PetscCall(TSSolve(ts, u));
513c4762a1bSJed Brown 
51440be0ff1SMatthew G. Knepley   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
51540be0ff1SMatthew G. Knepley      Clean up workspace
51640be0ff1SMatthew G. Knepley      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
517*9566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&J));
518*9566063dSJacob Faibussowitsch   PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "kinematics", &u));
519*9566063dSJacob Faibussowitsch   PetscCall(TSDestroy(&ts));
520*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&sw));
521*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dm));
522*9566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
523b122ec5aSJacob Faibussowitsch   return 0;
524c4762a1bSJed Brown }
525c4762a1bSJed Brown 
526c4762a1bSJed Brown /*TEST
527c4762a1bSJed Brown 
528c4762a1bSJed Brown    build:
529c4762a1bSJed Brown      requires: triangle !single !complex
530c4762a1bSJed Brown    test:
531c4762a1bSJed Brown      suffix: 1
532db4653c2SDaniel Finn      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_convergence_estimate -convest_num_refine 2 -dm_view  -monitor -output_step 50 -error
53340be0ff1SMatthew G. Knepley 
534c4762a1bSJed Brown    test:
535c4762a1bSJed Brown      suffix: 2
536db4653c2SDaniel Finn      args: -dm_plex_box_faces 1,1 -ts_type basicsymplectic -ts_basicsymplectic_type 2 -ts_convergence_estimate -convest_num_refine 2 -dm_view  -monitor -output_step 50 -error
53740be0ff1SMatthew G. Knepley 
538c4762a1bSJed Brown    test:
539c4762a1bSJed Brown      suffix: 3
540db4653c2SDaniel Finn      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 -monitor -output_step 50 -error
54140be0ff1SMatthew G. Knepley 
542c4762a1bSJed Brown    test:
543c4762a1bSJed Brown      suffix: 4
544db4653c2SDaniel Finn      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  -monitor -output_step 50 -error
54540be0ff1SMatthew G. Knepley 
54640be0ff1SMatthew G. Knepley    test:
54740be0ff1SMatthew G. Knepley      suffix: 5
548db4653c2SDaniel Finn      args: -dm_plex_box_faces 1,1 -ts_type theta -ts_theta_theta 0.5 -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2 -dm_view
54940be0ff1SMatthew G. Knepley 
55040be0ff1SMatthew G. Knepley    test:
55140be0ff1SMatthew G. Knepley      suffix: 6
552db4653c2SDaniel Finn      args: -dm_plex_box_faces 1,1 -ts_type discgrad -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2 -dm_view
553db4653c2SDaniel Finn 
554db4653c2SDaniel Finn    test:
555db4653c2SDaniel Finn      suffix: 7
556db4653c2SDaniel Finn      args: -dm_plex_box_faces 1,1 -ts_type discgrad -ts_discgrad_gonzalez -monitor -output_step 50 -error -ts_convergence_estimate -convest_num_refine 2 -dm_view
557c4762a1bSJed Brown 
558c4762a1bSJed Brown TEST*/
559