xref: /petsc/src/tao/unconstrained/tutorials/spectraladjointassimilation.c (revision d2522c19e8fa9bca20aaca277941d9a63e71db6a)
1 
2 static char help[] = "Solves a simple data assimilation problem with one dimensional advection diffusion equation using TSAdjoint\n\n";
3 
4 /*
5 
6     Not yet tested in parallel
7 
8 */
9 
10 /* ------------------------------------------------------------------------
11 
12    This program uses the one-dimensional advection-diffusion equation),
13        u_t = mu*u_xx - a u_x,
14    on the domain 0 <= x <= 1, with periodic boundary conditions
15 
16    to demonstrate solving a data assimilation problem of finding the initial conditions
17    to produce a given solution at a fixed time.
18 
19    The operators are discretized with the spectral element method
20 
21   ------------------------------------------------------------------------- */
22 
23 /*
24    Include "petscts.h" so that we can use TS solvers.  Note that this file
25    automatically includes:
26      petscsys.h       - base PETSc routines   petscvec.h  - vectors
27      petscmat.h  - matrices
28      petscis.h     - index sets            petscksp.h  - Krylov subspace methods
29      petscviewer.h - viewers               petscpc.h   - preconditioners
30      petscksp.h   - linear solvers        petscsnes.h - nonlinear solvers
31 */
32 
33 #include <petsctao.h>
34 #include <petscts.h>
35 #include <petscdt.h>
36 #include <petscdraw.h>
37 #include <petscdmda.h>
38 
39 /*
40    User-defined application context - contains data needed by the
41    application-provided call-back routines.
42 */
43 
44 typedef struct {
45   PetscInt   n;       /* number of nodes */
46   PetscReal *nodes;   /* GLL nodes */
47   PetscReal *weights; /* GLL weights */
48 } PetscGLL;
49 
50 typedef struct {
51   PetscInt  N;               /* grid points per elements*/
52   PetscInt  E;               /* number of elements */
53   PetscReal tol_L2, tol_max; /* error norms */
54   PetscInt  steps;           /* number of timesteps */
55   PetscReal Tend;            /* endtime */
56   PetscReal mu;              /* viscosity */
57   PetscReal a;               /* advection speed */
58   PetscReal L;               /* total length of domain */
59   PetscReal Le;
60   PetscReal Tadj;
61 } PetscParam;
62 
63 typedef struct {
64   Vec reference; /* desired end state */
65   Vec grid;      /* total grid */
66   Vec grad;
67   Vec ic;
68   Vec curr_sol;
69   Vec joe;
70   Vec true_solution; /* actual initial conditions for the final solution */
71 } PetscData;
72 
73 typedef struct {
74   Vec      grid;  /* total grid */
75   Vec      mass;  /* mass matrix for total integration */
76   Mat      stiff; /* stifness matrix */
77   Mat      advec;
78   Mat      keptstiff;
79   PetscGLL gll;
80 } PetscSEMOperators;
81 
82 typedef struct {
83   DM                da; /* distributed array data structure */
84   PetscSEMOperators SEMop;
85   PetscParam        param;
86   PetscData         dat;
87   TS                ts;
88   PetscReal         initial_dt;
89   PetscReal        *solutioncoefficients;
90   PetscInt          ncoeff;
91 } AppCtx;
92 
93 /*
94    User-defined routines
95 */
96 extern PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
97 extern PetscErrorCode RHSLaplacian(TS, PetscReal, Vec, Mat, Mat, void *);
98 extern PetscErrorCode RHSAdvection(TS, PetscReal, Vec, Mat, Mat, void *);
99 extern PetscErrorCode InitialConditions(Vec, AppCtx *);
100 extern PetscErrorCode ComputeReference(TS, PetscReal, Vec, AppCtx *);
101 extern PetscErrorCode MonitorError(Tao, void *);
102 extern PetscErrorCode MonitorDestroy(void **);
103 extern PetscErrorCode ComputeSolutionCoefficients(AppCtx *);
104 extern PetscErrorCode RHSFunction(TS, PetscReal, Vec, Vec, void *);
105 extern PetscErrorCode RHSJacobian(TS, PetscReal, Vec, Mat, Mat, void *);
106 
107 int main(int argc, char **argv) {
108   AppCtx       appctx; /* user-defined application context */
109   Tao          tao;
110   Vec          u; /* approximate solution vector */
111   PetscInt     i, xs, xm, ind, j, lenglob;
112   PetscReal    x, *wrk_ptr1, *wrk_ptr2;
113   MatNullSpace nsp;
114 
115   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
116      Initialize program and set problem parameters
117      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
118   PetscFunctionBegin;
119 
120   PetscFunctionBeginUser;
121   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
122 
123   /*initialize parameters */
124   appctx.param.N     = 10;      /* order of the spectral element */
125   appctx.param.E     = 8;       /* number of elements */
126   appctx.param.L     = 1.0;     /* length of the domain */
127   appctx.param.mu    = 0.00001; /* diffusion coefficient */
128   appctx.param.a     = 0.0;     /* advection speed */
129   appctx.initial_dt  = 1e-4;
130   appctx.param.steps = PETSC_MAX_INT;
131   appctx.param.Tend  = 0.01;
132   appctx.ncoeff      = 2;
133 
134   PetscCall(PetscOptionsGetInt(NULL, NULL, "-N", &appctx.param.N, NULL));
135   PetscCall(PetscOptionsGetInt(NULL, NULL, "-E", &appctx.param.E, NULL));
136   PetscCall(PetscOptionsGetInt(NULL, NULL, "-ncoeff", &appctx.ncoeff, NULL));
137   PetscCall(PetscOptionsGetReal(NULL, NULL, "-Tend", &appctx.param.Tend, NULL));
138   PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &appctx.param.mu, NULL));
139   PetscCall(PetscOptionsGetReal(NULL, NULL, "-a", &appctx.param.a, NULL));
140   appctx.param.Le = appctx.param.L / appctx.param.E;
141 
142   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
143      Create GLL data structures
144      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
145   PetscCall(PetscMalloc2(appctx.param.N, &appctx.SEMop.gll.nodes, appctx.param.N, &appctx.SEMop.gll.weights));
146   PetscCall(PetscDTGaussLobattoLegendreQuadrature(appctx.param.N, PETSCGAUSSLOBATTOLEGENDRE_VIA_LINEAR_ALGEBRA, appctx.SEMop.gll.nodes, appctx.SEMop.gll.weights));
147   appctx.SEMop.gll.n = appctx.param.N;
148   lenglob            = appctx.param.E * (appctx.param.N - 1);
149 
150   /*
151      Create distributed array (DMDA) to manage parallel grid and vectors
152      and to set up the ghost point communication pattern.  There are E*(Nl-1)+1
153      total grid values spread equally among all the processors, except first and last
154   */
155 
156   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, lenglob, 1, 1, NULL, &appctx.da));
157   PetscCall(DMSetFromOptions(appctx.da));
158   PetscCall(DMSetUp(appctx.da));
159 
160   /*
161      Extract global and local vectors from DMDA; we use these to store the
162      approximate solution.  Then duplicate these for remaining vectors that
163      have the same types.
164   */
165 
166   PetscCall(DMCreateGlobalVector(appctx.da, &u));
167   PetscCall(VecDuplicate(u, &appctx.dat.ic));
168   PetscCall(VecDuplicate(u, &appctx.dat.true_solution));
169   PetscCall(VecDuplicate(u, &appctx.dat.reference));
170   PetscCall(VecDuplicate(u, &appctx.SEMop.grid));
171   PetscCall(VecDuplicate(u, &appctx.SEMop.mass));
172   PetscCall(VecDuplicate(u, &appctx.dat.curr_sol));
173   PetscCall(VecDuplicate(u, &appctx.dat.joe));
174 
175   PetscCall(DMDAGetCorners(appctx.da, &xs, NULL, NULL, &xm, NULL, NULL));
176   PetscCall(DMDAVecGetArray(appctx.da, appctx.SEMop.grid, &wrk_ptr1));
177   PetscCall(DMDAVecGetArray(appctx.da, appctx.SEMop.mass, &wrk_ptr2));
178 
179   /* Compute function over the locally owned part of the grid */
180 
181   xs = xs / (appctx.param.N - 1);
182   xm = xm / (appctx.param.N - 1);
183 
184   /*
185      Build total grid and mass over entire mesh (multi-elemental)
186   */
187 
188   for (i = xs; i < xs + xm; i++) {
189     for (j = 0; j < appctx.param.N - 1; j++) {
190       x             = (appctx.param.Le / 2.0) * (appctx.SEMop.gll.nodes[j] + 1.0) + appctx.param.Le * i;
191       ind           = i * (appctx.param.N - 1) + j;
192       wrk_ptr1[ind] = x;
193       wrk_ptr2[ind] = .5 * appctx.param.Le * appctx.SEMop.gll.weights[j];
194       if (j == 0) wrk_ptr2[ind] += .5 * appctx.param.Le * appctx.SEMop.gll.weights[j];
195     }
196   }
197   PetscCall(DMDAVecRestoreArray(appctx.da, appctx.SEMop.grid, &wrk_ptr1));
198   PetscCall(DMDAVecRestoreArray(appctx.da, appctx.SEMop.mass, &wrk_ptr2));
199 
200   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
201    Create matrix data structure; set matrix evaluation routine.
202    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
203   PetscCall(DMSetMatrixPreallocateOnly(appctx.da, PETSC_TRUE));
204   PetscCall(DMCreateMatrix(appctx.da, &appctx.SEMop.stiff));
205   PetscCall(DMCreateMatrix(appctx.da, &appctx.SEMop.advec));
206 
207   /*
208    For linear problems with a time-dependent f(u,t) in the equation
209    u_t = f(u,t), the user provides the discretized right-hand-side
210    as a time-dependent matrix.
211    */
212   PetscCall(RHSLaplacian(appctx.ts, 0.0, u, appctx.SEMop.stiff, appctx.SEMop.stiff, &appctx));
213   PetscCall(RHSAdvection(appctx.ts, 0.0, u, appctx.SEMop.advec, appctx.SEMop.advec, &appctx));
214   PetscCall(MatAXPY(appctx.SEMop.stiff, -1.0, appctx.SEMop.advec, DIFFERENT_NONZERO_PATTERN));
215   PetscCall(MatDuplicate(appctx.SEMop.stiff, MAT_COPY_VALUES, &appctx.SEMop.keptstiff));
216 
217   /* attach the null space to the matrix, this probably is not needed but does no harm */
218   PetscCall(MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, NULL, &nsp));
219   PetscCall(MatSetNullSpace(appctx.SEMop.stiff, nsp));
220   PetscCall(MatNullSpaceTest(nsp, appctx.SEMop.stiff, NULL));
221   PetscCall(MatNullSpaceDestroy(&nsp));
222 
223   /* Create the TS solver that solves the ODE and its adjoint; set its options */
224   PetscCall(TSCreate(PETSC_COMM_WORLD, &appctx.ts));
225   PetscCall(TSSetSolutionFunction(appctx.ts, (PetscErrorCode(*)(TS, PetscReal, Vec, void *))ComputeReference, &appctx));
226   PetscCall(TSSetProblemType(appctx.ts, TS_LINEAR));
227   PetscCall(TSSetType(appctx.ts, TSRK));
228   PetscCall(TSSetDM(appctx.ts, appctx.da));
229   PetscCall(TSSetTime(appctx.ts, 0.0));
230   PetscCall(TSSetTimeStep(appctx.ts, appctx.initial_dt));
231   PetscCall(TSSetMaxSteps(appctx.ts, appctx.param.steps));
232   PetscCall(TSSetMaxTime(appctx.ts, appctx.param.Tend));
233   PetscCall(TSSetExactFinalTime(appctx.ts, TS_EXACTFINALTIME_MATCHSTEP));
234   PetscCall(TSSetTolerances(appctx.ts, 1e-7, NULL, 1e-7, NULL));
235   PetscCall(TSSetFromOptions(appctx.ts));
236   /* Need to save initial timestep user may have set with -ts_dt so it can be reset for each new TSSolve() */
237   PetscCall(TSGetTimeStep(appctx.ts, &appctx.initial_dt));
238   PetscCall(TSSetRHSFunction(appctx.ts, NULL, TSComputeRHSFunctionLinear, &appctx));
239   PetscCall(TSSetRHSJacobian(appctx.ts, appctx.SEMop.stiff, appctx.SEMop.stiff, TSComputeRHSJacobianConstant, &appctx));
240   /*  PetscCall(TSSetRHSFunction(appctx.ts,NULL,RHSFunction,&appctx));
241       PetscCall(TSSetRHSJacobian(appctx.ts,appctx.SEMop.stiff,appctx.SEMop.stiff,RHSJacobian,&appctx)); */
242 
243   /* Set random initial conditions as initial guess, compute analytic reference solution and analytic (true) initial conditions */
244   PetscCall(ComputeSolutionCoefficients(&appctx));
245   PetscCall(InitialConditions(appctx.dat.ic, &appctx));
246   PetscCall(ComputeReference(appctx.ts, appctx.param.Tend, appctx.dat.reference, &appctx));
247   PetscCall(ComputeReference(appctx.ts, 0.0, appctx.dat.true_solution, &appctx));
248 
249   /* Set up to save trajectory before TSSetFromOptions() so that TSTrajectory options can be captured */
250   PetscCall(TSSetSaveTrajectory(appctx.ts));
251   PetscCall(TSSetFromOptions(appctx.ts));
252 
253   /* Create TAO solver and set desired solution method  */
254   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
255   PetscCall(TaoSetMonitor(tao, MonitorError, &appctx, MonitorDestroy));
256   PetscCall(TaoSetType(tao, TAOBQNLS));
257   PetscCall(TaoSetSolution(tao, appctx.dat.ic));
258   /* Set routine for function and gradient evaluation  */
259   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&appctx));
260   /* Check for any TAO command line options  */
261   PetscCall(TaoSetTolerances(tao, 1e-8, PETSC_DEFAULT, PETSC_DEFAULT));
262   PetscCall(TaoSetFromOptions(tao));
263   PetscCall(TaoSolve(tao));
264 
265   PetscCall(TaoDestroy(&tao));
266   PetscCall(PetscFree(appctx.solutioncoefficients));
267   PetscCall(MatDestroy(&appctx.SEMop.advec));
268   PetscCall(MatDestroy(&appctx.SEMop.stiff));
269   PetscCall(MatDestroy(&appctx.SEMop.keptstiff));
270   PetscCall(VecDestroy(&u));
271   PetscCall(VecDestroy(&appctx.dat.ic));
272   PetscCall(VecDestroy(&appctx.dat.joe));
273   PetscCall(VecDestroy(&appctx.dat.true_solution));
274   PetscCall(VecDestroy(&appctx.dat.reference));
275   PetscCall(VecDestroy(&appctx.SEMop.grid));
276   PetscCall(VecDestroy(&appctx.SEMop.mass));
277   PetscCall(VecDestroy(&appctx.dat.curr_sol));
278   PetscCall(PetscFree2(appctx.SEMop.gll.nodes, appctx.SEMop.gll.weights));
279   PetscCall(DMDestroy(&appctx.da));
280   PetscCall(TSDestroy(&appctx.ts));
281 
282   /*
283      Always call PetscFinalize() before exiting a program.  This routine
284        - finalizes the PETSc libraries as well as MPI
285        - provides summary and diagnostic information if certain runtime
286          options are chosen (e.g., -log_summary).
287   */
288   PetscCall(PetscFinalize());
289   return 0;
290 }
291 
292 /*
293     Computes the coefficients for the analytic solution to the PDE
294 */
295 PetscErrorCode ComputeSolutionCoefficients(AppCtx *appctx) {
296   PetscRandom rand;
297   PetscInt    i;
298 
299   PetscFunctionBegin;
300   PetscCall(PetscMalloc1(appctx->ncoeff, &appctx->solutioncoefficients));
301   PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rand));
302   PetscCall(PetscRandomSetInterval(rand, .9, 1.0));
303   for (i = 0; i < appctx->ncoeff; i++) { PetscCall(PetscRandomGetValue(rand, &appctx->solutioncoefficients[i])); }
304   PetscCall(PetscRandomDestroy(&rand));
305   PetscFunctionReturn(0);
306 }
307 
308 /* --------------------------------------------------------------------- */
309 /*
310    InitialConditions - Computes the (random) initial conditions for the Tao optimization solve (these are also initial conditions for the first TSSolve()
311 
312    Input Parameter:
313    u - uninitialized solution vector (global)
314    appctx - user-defined application context
315 
316    Output Parameter:
317    u - vector with solution at initial time (global)
318 */
319 PetscErrorCode InitialConditions(Vec u, AppCtx *appctx) {
320   PetscScalar       *s;
321   const PetscScalar *xg;
322   PetscInt           i, j, lenglob;
323   PetscReal          sum, val;
324   PetscRandom        rand;
325 
326   PetscFunctionBegin;
327   PetscCall(PetscRandomCreate(PETSC_COMM_WORLD, &rand));
328   PetscCall(PetscRandomSetInterval(rand, .9, 1.0));
329   PetscCall(DMDAVecGetArray(appctx->da, u, &s));
330   PetscCall(DMDAVecGetArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
331   lenglob = appctx->param.E * (appctx->param.N - 1);
332   for (i = 0; i < lenglob; i++) {
333     s[i] = 0;
334     for (j = 0; j < appctx->ncoeff; j++) {
335       PetscCall(PetscRandomGetValue(rand, &val));
336       s[i] += val * PetscSinScalar(2 * (j + 1) * PETSC_PI * xg[i]);
337     }
338   }
339   PetscCall(PetscRandomDestroy(&rand));
340   PetscCall(DMDAVecRestoreArray(appctx->da, u, &s));
341   PetscCall(DMDAVecRestoreArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
342   /* make sure initial conditions do not contain the constant functions, since with periodic boundary conditions the constant functions introduce a null space */
343   PetscCall(VecSum(u, &sum));
344   PetscCall(VecShift(u, -sum / lenglob));
345   PetscFunctionReturn(0);
346 }
347 
348 /*
349    TrueSolution() computes the true solution for the Tao optimization solve which means they are the initial conditions for the objective function.
350 
351              InitialConditions() computes the initial conditions for the beginning of the Tao iterations
352 
353    Input Parameter:
354    u - uninitialized solution vector (global)
355    appctx - user-defined application context
356 
357    Output Parameter:
358    u - vector with solution at initial time (global)
359 */
360 PetscErrorCode TrueSolution(Vec u, AppCtx *appctx) {
361   PetscScalar       *s;
362   const PetscScalar *xg;
363   PetscInt           i, j, lenglob;
364   PetscReal          sum;
365 
366   PetscFunctionBegin;
367   PetscCall(DMDAVecGetArray(appctx->da, u, &s));
368   PetscCall(DMDAVecGetArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
369   lenglob = appctx->param.E * (appctx->param.N - 1);
370   for (i = 0; i < lenglob; i++) {
371     s[i] = 0;
372     for (j = 0; j < appctx->ncoeff; j++) { s[i] += appctx->solutioncoefficients[j] * PetscSinScalar(2 * (j + 1) * PETSC_PI * xg[i]); }
373   }
374   PetscCall(DMDAVecRestoreArray(appctx->da, u, &s));
375   PetscCall(DMDAVecRestoreArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
376   /* make sure initial conditions do not contain the constant functions, since with periodic boundary conditions the constant functions introduce a null space */
377   PetscCall(VecSum(u, &sum));
378   PetscCall(VecShift(u, -sum / lenglob));
379   PetscFunctionReturn(0);
380 }
381 /* --------------------------------------------------------------------- */
382 /*
383    Sets the desired profile for the final end time
384 
385    Input Parameters:
386    t - final time
387    obj - vector storing the desired profile
388    appctx - user-defined application context
389 
390 */
391 PetscErrorCode ComputeReference(TS ts, PetscReal t, Vec obj, AppCtx *appctx) {
392   PetscScalar       *s, tc;
393   const PetscScalar *xg;
394   PetscInt           i, j, lenglob;
395 
396   PetscFunctionBegin;
397   PetscCall(DMDAVecGetArray(appctx->da, obj, &s));
398   PetscCall(DMDAVecGetArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
399   lenglob = appctx->param.E * (appctx->param.N - 1);
400   for (i = 0; i < lenglob; i++) {
401     s[i] = 0;
402     for (j = 0; j < appctx->ncoeff; j++) {
403       tc = -appctx->param.mu * (j + 1) * (j + 1) * 4.0 * PETSC_PI * PETSC_PI * t;
404       s[i] += appctx->solutioncoefficients[j] * PetscSinScalar(2 * (j + 1) * PETSC_PI * (xg[i] + appctx->param.a * t)) * PetscExpReal(tc);
405     }
406   }
407   PetscCall(DMDAVecRestoreArray(appctx->da, obj, &s));
408   PetscCall(DMDAVecRestoreArrayRead(appctx->da, appctx->SEMop.grid, (void *)&xg));
409   PetscFunctionReturn(0);
410 }
411 
412 PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec globalin, Vec globalout, void *ctx) {
413   AppCtx *appctx = (AppCtx *)ctx;
414 
415   PetscFunctionBegin;
416   PetscCall(MatMult(appctx->SEMop.keptstiff, globalin, globalout));
417   PetscFunctionReturn(0);
418 }
419 
420 PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec globalin, Mat A, Mat B, void *ctx) {
421   AppCtx *appctx = (AppCtx *)ctx;
422 
423   PetscFunctionBegin;
424   PetscCall(MatCopy(appctx->SEMop.keptstiff, A, DIFFERENT_NONZERO_PATTERN));
425   PetscFunctionReturn(0);
426 }
427 
428 /* --------------------------------------------------------------------- */
429 
430 /*
431    RHSLaplacian -   matrix for diffusion
432 
433    Input Parameters:
434    ts - the TS context
435    t - current time  (ignored)
436    X - current solution (ignored)
437    dummy - optional user-defined context, as set by TSetRHSJacobian()
438 
439    Output Parameters:
440    AA - Jacobian matrix
441    BB - optionally different matrix from which the preconditioner is built
442    str - flag indicating matrix structure
443 
444    Scales by the inverse of the mass matrix (perhaps that should be pulled out)
445 
446 */
447 PetscErrorCode RHSLaplacian(TS ts, PetscReal t, Vec X, Mat A, Mat BB, void *ctx) {
448   PetscReal **temp;
449   PetscReal   vv;
450   AppCtx     *appctx = (AppCtx *)ctx; /* user-defined application context */
451   PetscInt    i, xs, xn, l, j;
452   PetscInt   *rowsDM;
453 
454   PetscFunctionBegin;
455   /*
456    Creates the element stiffness matrix for the given gll
457    */
458   PetscCall(PetscGaussLobattoLegendreElementLaplacianCreate(appctx->SEMop.gll.n, appctx->SEMop.gll.nodes, appctx->SEMop.gll.weights, &temp));
459 
460   /* scale by the size of the element */
461   for (i = 0; i < appctx->param.N; i++) {
462     vv = -appctx->param.mu * 2.0 / appctx->param.Le;
463     for (j = 0; j < appctx->param.N; j++) temp[i][j] = temp[i][j] * vv;
464   }
465 
466   PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE));
467   PetscCall(DMDAGetCorners(appctx->da, &xs, NULL, NULL, &xn, NULL, NULL));
468 
469   PetscCheck(appctx->param.N - 1 >= 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Polynomial order must be at least 2");
470   xs = xs / (appctx->param.N - 1);
471   xn = xn / (appctx->param.N - 1);
472 
473   PetscCall(PetscMalloc1(appctx->param.N, &rowsDM));
474   /*
475    loop over local elements
476    */
477   for (j = xs; j < xs + xn; j++) {
478     for (l = 0; l < appctx->param.N; l++) { rowsDM[l] = 1 + (j - xs) * (appctx->param.N - 1) + l; }
479     PetscCall(MatSetValuesLocal(A, appctx->param.N, rowsDM, appctx->param.N, rowsDM, &temp[0][0], ADD_VALUES));
480   }
481   PetscCall(PetscFree(rowsDM));
482   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
483   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
484   PetscCall(VecReciprocal(appctx->SEMop.mass));
485   PetscCall(MatDiagonalScale(A, appctx->SEMop.mass, 0));
486   PetscCall(VecReciprocal(appctx->SEMop.mass));
487 
488   PetscCall(PetscGaussLobattoLegendreElementLaplacianDestroy(appctx->SEMop.gll.n, appctx->SEMop.gll.nodes, appctx->SEMop.gll.weights, &temp));
489   PetscFunctionReturn(0);
490 }
491 
492 /*
493     Almost identical to Laplacian
494 
495     Note that the element matrix is NOT scaled by the size of element like the Laplacian term.
496  */
497 PetscErrorCode RHSAdvection(TS ts, PetscReal t, Vec X, Mat A, Mat BB, void *ctx) {
498   PetscReal **temp;
499   PetscReal   vv;
500   AppCtx     *appctx = (AppCtx *)ctx; /* user-defined application context */
501   PetscInt    i, xs, xn, l, j;
502   PetscInt   *rowsDM;
503 
504   PetscFunctionBegin;
505   /*
506    Creates the element stiffness matrix for the given gll
507    */
508   PetscCall(PetscGaussLobattoLegendreElementAdvectionCreate(appctx->SEMop.gll.n, appctx->SEMop.gll.nodes, appctx->SEMop.gll.weights, &temp));
509 
510   /* scale by the size of the element */
511   for (i = 0; i < appctx->param.N; i++) {
512     vv = -appctx->param.a;
513     for (j = 0; j < appctx->param.N; j++) temp[i][j] = temp[i][j] * vv;
514   }
515 
516   PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE));
517   PetscCall(DMDAGetCorners(appctx->da, &xs, NULL, NULL, &xn, NULL, NULL));
518 
519   PetscCheck(appctx->param.N - 1 >= 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Polynomial order must be at least 2");
520   xs = xs / (appctx->param.N - 1);
521   xn = xn / (appctx->param.N - 1);
522 
523   PetscCall(PetscMalloc1(appctx->param.N, &rowsDM));
524   /*
525    loop over local elements
526    */
527   for (j = xs; j < xs + xn; j++) {
528     for (l = 0; l < appctx->param.N; l++) { rowsDM[l] = 1 + (j - xs) * (appctx->param.N - 1) + l; }
529     PetscCall(MatSetValuesLocal(A, appctx->param.N, rowsDM, appctx->param.N, rowsDM, &temp[0][0], ADD_VALUES));
530   }
531   PetscCall(PetscFree(rowsDM));
532   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
533   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
534   PetscCall(VecReciprocal(appctx->SEMop.mass));
535   PetscCall(MatDiagonalScale(A, appctx->SEMop.mass, 0));
536   PetscCall(VecReciprocal(appctx->SEMop.mass));
537 
538   PetscCall(PetscGaussLobattoLegendreElementAdvectionDestroy(appctx->SEMop.gll.n, appctx->SEMop.gll.nodes, appctx->SEMop.gll.weights, &temp));
539   PetscFunctionReturn(0);
540 }
541 
542 /* ------------------------------------------------------------------ */
543 /*
544    FormFunctionGradient - Evaluates the function and corresponding gradient.
545 
546    Input Parameters:
547    tao - the Tao context
548    ic   - the input vector
549    ctx - optional user-defined context, as set when calling TaoSetObjectiveAndGradient()
550 
551    Output Parameters:
552    f   - the newly evaluated function
553    G   - the newly evaluated gradient
554 
555    Notes:
556 
557           The forward equation is
558               M u_t = F(U)
559           which is converted to
560                 u_t = M^{-1} F(u)
561           in the user code since TS has no direct way of providing a mass matrix. The Jacobian of this is
562                  M^{-1} J
563           where J is the Jacobian of F. Now the adjoint equation is
564                 M v_t = J^T v
565           but TSAdjoint does not solve this since it can only solve the transposed system for the
566           Jacobian the user provided. Hence TSAdjoint solves
567                  w_t = J^T M^{-1} w  (where w = M v)
568           since there is no way to indicate the mass matrix as a separate entity to TS. Thus one
569           must be careful in initializing the "adjoint equation" and using the result. This is
570           why
571               G = -2 M(u(T) - u_d)
572           below (instead of -2(u(T) - u_d)
573 
574 */
575 PetscErrorCode FormFunctionGradient(Tao tao, Vec ic, PetscReal *f, Vec G, void *ctx) {
576   AppCtx *appctx = (AppCtx *)ctx; /* user-defined application context */
577   Vec     temp;
578 
579   PetscFunctionBegin;
580   PetscCall(TSSetTime(appctx->ts, 0.0));
581   PetscCall(TSSetStepNumber(appctx->ts, 0));
582   PetscCall(TSSetTimeStep(appctx->ts, appctx->initial_dt));
583   PetscCall(VecCopy(ic, appctx->dat.curr_sol));
584 
585   PetscCall(TSSolve(appctx->ts, appctx->dat.curr_sol));
586   PetscCall(VecCopy(appctx->dat.curr_sol, appctx->dat.joe));
587 
588   /*     Compute the difference between the current ODE solution and target ODE solution */
589   PetscCall(VecWAXPY(G, -1.0, appctx->dat.curr_sol, appctx->dat.reference));
590 
591   /*     Compute the objective/cost function   */
592   PetscCall(VecDuplicate(G, &temp));
593   PetscCall(VecPointwiseMult(temp, G, G));
594   PetscCall(VecDot(temp, appctx->SEMop.mass, f));
595   PetscCall(VecDestroy(&temp));
596 
597   /*     Compute initial conditions for the adjoint integration. See Notes above  */
598   PetscCall(VecScale(G, -2.0));
599   PetscCall(VecPointwiseMult(G, G, appctx->SEMop.mass));
600   PetscCall(TSSetCostGradients(appctx->ts, 1, &G, NULL));
601 
602   PetscCall(TSAdjointSolve(appctx->ts));
603   /* PetscCall(VecPointwiseDivide(G,G,appctx->SEMop.mass));*/
604   PetscFunctionReturn(0);
605 }
606 
607 PetscErrorCode MonitorError(Tao tao, void *ctx) {
608   AppCtx   *appctx = (AppCtx *)ctx;
609   Vec       temp, grad;
610   PetscReal nrm;
611   PetscInt  its;
612   PetscReal fct, gnorm;
613 
614   PetscFunctionBegin;
615   PetscCall(VecDuplicate(appctx->dat.ic, &temp));
616   PetscCall(VecWAXPY(temp, -1.0, appctx->dat.ic, appctx->dat.true_solution));
617   PetscCall(VecPointwiseMult(temp, temp, temp));
618   PetscCall(VecDot(temp, appctx->SEMop.mass, &nrm));
619   nrm = PetscSqrtReal(nrm);
620   PetscCall(TaoGetGradient(tao, &grad, NULL, NULL));
621   PetscCall(VecPointwiseMult(temp, temp, temp));
622   PetscCall(VecDot(temp, appctx->SEMop.mass, &gnorm));
623   gnorm = PetscSqrtReal(gnorm);
624   PetscCall(VecDestroy(&temp));
625   PetscCall(TaoGetIterationNumber(tao, &its));
626   PetscCall(TaoGetSolutionStatus(tao, NULL, &fct, NULL, NULL, NULL, NULL));
627   if (!its) {
628     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%% Iteration Error Objective Gradient-norm\n"));
629     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "history = [\n"));
630   }
631   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%3" PetscInt_FMT " %g %g %g\n", its, (double)nrm, (double)fct, (double)gnorm));
632   PetscFunctionReturn(0);
633 }
634 
635 PetscErrorCode MonitorDestroy(void **ctx) {
636   PetscFunctionBegin;
637   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "];\n"));
638   PetscFunctionReturn(0);
639 }
640 
641 /*TEST
642 
643    build:
644      requires: !complex
645 
646    test:
647      requires: !single
648      args:  -ts_adapt_dt_max 3.e-3 -E 10 -N 8 -ncoeff 5 -tao_bqnls_mat_lmvm_scale_type none
649 
650    test:
651      suffix: cn
652      requires: !single
653      args:  -ts_type cn -ts_dt .003 -pc_type lu -E 10 -N 8 -ncoeff 5 -tao_bqnls_mat_lmvm_scale_type none
654 
655    test:
656      suffix: 2
657      requires: !single
658      args:  -ts_adapt_dt_max 3.e-3 -E 10 -N 8 -ncoeff 5  -a .1 -tao_bqnls_mat_lmvm_scale_type none
659 
660 TEST*/
661