xref: /petsc/src/ts/tutorials/ex20opt_p.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a) !
1 
2 static char help[] = "Solves the van der Pol equation.\n\
3 Input parameters include:\n";
4 
5 /* ------------------------------------------------------------------------
6 
7   Notes:
8   This code demonstrates how to solve a DAE-constrained optimization problem with TAO, TSAdjoint and TS.
9   The nonlinear problem is written in a DAE equivalent form.
10   The objective is to minimize the difference between observation and model prediction by finding an optimal value for parameter \mu.
11   The gradient is computed with the discrete adjoint of an implicit theta method, see ex20adj.c for details.
12   ------------------------------------------------------------------------- */
13 #include <petsctao.h>
14 #include <petscts.h>
15 
16 typedef struct _n_User *User;
17 struct _n_User {
18   TS        ts;
19   PetscReal mu;
20   PetscReal next_output;
21 
22   /* Sensitivity analysis support */
23   PetscReal ftime;
24   Mat       A;                    /* Jacobian matrix */
25   Mat       Jacp;                 /* JacobianP matrix */
26   Mat       H;                    /* Hessian matrix for optimization */
27   Vec       U, Lambda[1], Mup[1]; /* adjoint variables */
28   Vec       Lambda2[1], Mup2[1];  /* second-order adjoint variables */
29   Vec       Ihp1[1];              /* working space for Hessian evaluations */
30   Vec       Ihp2[1];              /* working space for Hessian evaluations */
31   Vec       Ihp3[1];              /* working space for Hessian evaluations */
32   Vec       Ihp4[1];              /* working space for Hessian evaluations */
33   Vec       Dir;                  /* direction vector */
34   PetscReal ob[2];                /* observation used by the cost function */
35   PetscBool implicitform;         /* implicit ODE? */
36 };
37 
38 PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
39 PetscErrorCode FormHessian(Tao, Vec, Mat, Mat, void *);
40 PetscErrorCode Adjoint2(Vec, PetscScalar[], User);
41 
42 /* ----------------------- Explicit form of the ODE  -------------------- */
43 
44 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, void *ctx) {
45   User               user = (User)ctx;
46   PetscScalar       *f;
47   const PetscScalar *u;
48 
49   PetscFunctionBeginUser;
50   PetscCall(VecGetArrayRead(U, &u));
51   PetscCall(VecGetArray(F, &f));
52   f[0] = u[1];
53   f[1] = user->mu * ((1. - u[0] * u[0]) * u[1] - u[0]);
54   PetscCall(VecRestoreArrayRead(U, &u));
55   PetscCall(VecRestoreArray(F, &f));
56   PetscFunctionReturn(0);
57 }
58 
59 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat A, Mat B, void *ctx) {
60   User               user     = (User)ctx;
61   PetscReal          mu       = user->mu;
62   PetscInt           rowcol[] = {0, 1};
63   PetscScalar        J[2][2];
64   const PetscScalar *u;
65 
66   PetscFunctionBeginUser;
67   PetscCall(VecGetArrayRead(U, &u));
68 
69   J[0][0] = 0;
70   J[1][0] = -mu * (2.0 * u[1] * u[0] + 1.);
71   J[0][1] = 1.0;
72   J[1][1] = mu * (1.0 - u[0] * u[0]);
73   PetscCall(MatSetValues(A, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
74   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
75   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
76   if (B && A != B) {
77     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
78     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
79   }
80 
81   PetscCall(VecRestoreArrayRead(U, &u));
82   PetscFunctionReturn(0);
83 }
84 
85 static PetscErrorCode RHSHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
86   const PetscScalar *vl, *vr, *u;
87   PetscScalar       *vhv;
88   PetscScalar        dJdU[2][2][2] = {{{0}}};
89   PetscInt           i, j, k;
90   User               user = (User)ctx;
91 
92   PetscFunctionBeginUser;
93   PetscCall(VecGetArrayRead(U, &u));
94   PetscCall(VecGetArrayRead(Vl[0], &vl));
95   PetscCall(VecGetArrayRead(Vr, &vr));
96   PetscCall(VecGetArray(VHV[0], &vhv));
97 
98   dJdU[1][0][0] = -2. * user->mu * u[1];
99   dJdU[1][1][0] = -2. * user->mu * u[0];
100   dJdU[1][0][1] = -2. * user->mu * u[0];
101   for (j = 0; j < 2; j++) {
102     vhv[j] = 0;
103     for (k = 0; k < 2; k++)
104       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
105   }
106 
107   PetscCall(VecRestoreArrayRead(U, &u));
108   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
109   PetscCall(VecRestoreArrayRead(Vr, &vr));
110   PetscCall(VecRestoreArray(VHV[0], &vhv));
111   PetscFunctionReturn(0);
112 }
113 
114 static PetscErrorCode RHSHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
115   const PetscScalar *vl, *vr, *u;
116   PetscScalar       *vhv;
117   PetscScalar        dJdP[2][2][1] = {{{0}}};
118   PetscInt           i, j, k;
119 
120   PetscFunctionBeginUser;
121   PetscCall(VecGetArrayRead(U, &u));
122   PetscCall(VecGetArrayRead(Vl[0], &vl));
123   PetscCall(VecGetArrayRead(Vr, &vr));
124   PetscCall(VecGetArray(VHV[0], &vhv));
125 
126   dJdP[1][0][0] = -(1. + 2. * u[0] * u[1]);
127   dJdP[1][1][0] = 1. - u[0] * u[0];
128   for (j = 0; j < 2; j++) {
129     vhv[j] = 0;
130     for (k = 0; k < 1; k++)
131       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdP[i][j][k] * vr[k];
132   }
133 
134   PetscCall(VecRestoreArrayRead(U, &u));
135   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
136   PetscCall(VecRestoreArrayRead(Vr, &vr));
137   PetscCall(VecRestoreArray(VHV[0], &vhv));
138   PetscFunctionReturn(0);
139 }
140 
141 static PetscErrorCode RHSHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
142   const PetscScalar *vl, *vr, *u;
143   PetscScalar       *vhv;
144   PetscScalar        dJdU[2][1][2] = {{{0}}};
145   PetscInt           i, j, k;
146 
147   PetscFunctionBeginUser;
148   PetscCall(VecGetArrayRead(U, &u));
149   PetscCall(VecGetArrayRead(Vl[0], &vl));
150   PetscCall(VecGetArrayRead(Vr, &vr));
151   PetscCall(VecGetArray(VHV[0], &vhv));
152 
153   dJdU[1][0][0] = -1. - 2. * u[1] * u[0];
154   dJdU[1][0][1] = 1. - u[0] * u[0];
155   for (j = 0; j < 1; j++) {
156     vhv[j] = 0;
157     for (k = 0; k < 2; k++)
158       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
159   }
160 
161   PetscCall(VecRestoreArrayRead(U, &u));
162   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
163   PetscCall(VecRestoreArrayRead(Vr, &vr));
164   PetscCall(VecRestoreArray(VHV[0], &vhv));
165   PetscFunctionReturn(0);
166 }
167 
168 static PetscErrorCode RHSHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
169   PetscFunctionBeginUser;
170   PetscFunctionReturn(0);
171 }
172 
173 /* ----------------------- Implicit form of the ODE  -------------------- */
174 
175 static PetscErrorCode IFunction(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, void *ctx) {
176   User               user = (User)ctx;
177   PetscScalar       *f;
178   const PetscScalar *u, *udot;
179 
180   PetscFunctionBeginUser;
181   PetscCall(VecGetArrayRead(U, &u));
182   PetscCall(VecGetArrayRead(Udot, &udot));
183   PetscCall(VecGetArray(F, &f));
184 
185   f[0] = udot[0] - u[1];
186   f[1] = udot[1] - user->mu * ((1.0 - u[0] * u[0]) * u[1] - u[0]);
187 
188   PetscCall(VecRestoreArrayRead(U, &u));
189   PetscCall(VecRestoreArrayRead(Udot, &udot));
190   PetscCall(VecRestoreArray(F, &f));
191   PetscFunctionReturn(0);
192 }
193 
194 static PetscErrorCode IJacobian(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal a, Mat A, Mat B, void *ctx) {
195   User               user     = (User)ctx;
196   PetscInt           rowcol[] = {0, 1};
197   PetscScalar        J[2][2];
198   const PetscScalar *u;
199 
200   PetscFunctionBeginUser;
201   PetscCall(VecGetArrayRead(U, &u));
202 
203   J[0][0] = a;
204   J[0][1] = -1.0;
205   J[1][0] = user->mu * (1.0 + 2.0 * u[0] * u[1]);
206   J[1][1] = a - user->mu * (1.0 - u[0] * u[0]);
207   PetscCall(MatSetValues(B, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
208   PetscCall(VecRestoreArrayRead(U, &u));
209   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
210   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
211   if (A != B) {
212     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
213     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
214   }
215 
216   PetscCall(VecRestoreArrayRead(U, &u));
217   PetscFunctionReturn(0);
218 }
219 
220 static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec U, Mat A, void *ctx) {
221   PetscInt           row[] = {0, 1}, col[] = {0};
222   PetscScalar        J[2][1];
223   const PetscScalar *u;
224 
225   PetscFunctionBeginUser;
226   PetscCall(VecGetArrayRead(U, &u));
227 
228   J[0][0] = 0;
229   J[1][0] = (1. - u[0] * u[0]) * u[1] - u[0];
230   PetscCall(MatSetValues(A, 2, row, 1, col, &J[0][0], INSERT_VALUES));
231   PetscCall(VecRestoreArrayRead(U, &u));
232   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
233   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
234 
235   PetscCall(VecRestoreArrayRead(U, &u));
236   PetscFunctionReturn(0);
237 }
238 
239 static PetscErrorCode IHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
240   const PetscScalar *vl, *vr, *u;
241   PetscScalar       *vhv;
242   PetscScalar        dJdU[2][2][2] = {{{0}}};
243   PetscInt           i, j, k;
244   User               user = (User)ctx;
245 
246   PetscFunctionBeginUser;
247   PetscCall(VecGetArrayRead(U, &u));
248   PetscCall(VecGetArrayRead(Vl[0], &vl));
249   PetscCall(VecGetArrayRead(Vr, &vr));
250   PetscCall(VecGetArray(VHV[0], &vhv));
251 
252   dJdU[1][0][0] = 2. * user->mu * u[1];
253   dJdU[1][1][0] = 2. * user->mu * u[0];
254   dJdU[1][0][1] = 2. * user->mu * u[0];
255   for (j = 0; j < 2; j++) {
256     vhv[j] = 0;
257     for (k = 0; k < 2; k++)
258       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
259   }
260 
261   PetscCall(VecRestoreArrayRead(U, &u));
262   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
263   PetscCall(VecRestoreArrayRead(Vr, &vr));
264   PetscCall(VecRestoreArray(VHV[0], &vhv));
265   PetscFunctionReturn(0);
266 }
267 
268 static PetscErrorCode IHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
269   const PetscScalar *vl, *vr, *u;
270   PetscScalar       *vhv;
271   PetscScalar        dJdP[2][2][1] = {{{0}}};
272   PetscInt           i, j, k;
273 
274   PetscFunctionBeginUser;
275   PetscCall(VecGetArrayRead(U, &u));
276   PetscCall(VecGetArrayRead(Vl[0], &vl));
277   PetscCall(VecGetArrayRead(Vr, &vr));
278   PetscCall(VecGetArray(VHV[0], &vhv));
279 
280   dJdP[1][0][0] = 1. + 2. * u[0] * u[1];
281   dJdP[1][1][0] = u[0] * u[0] - 1.;
282   for (j = 0; j < 2; j++) {
283     vhv[j] = 0;
284     for (k = 0; k < 1; k++)
285       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdP[i][j][k] * vr[k];
286   }
287 
288   PetscCall(VecRestoreArrayRead(U, &u));
289   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
290   PetscCall(VecRestoreArrayRead(Vr, &vr));
291   PetscCall(VecRestoreArray(VHV[0], &vhv));
292   PetscFunctionReturn(0);
293 }
294 
295 static PetscErrorCode IHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
296   const PetscScalar *vl, *vr, *u;
297   PetscScalar       *vhv;
298   PetscScalar        dJdU[2][1][2] = {{{0}}};
299   PetscInt           i, j, k;
300 
301   PetscFunctionBeginUser;
302   PetscCall(VecGetArrayRead(U, &u));
303   PetscCall(VecGetArrayRead(Vl[0], &vl));
304   PetscCall(VecGetArrayRead(Vr, &vr));
305   PetscCall(VecGetArray(VHV[0], &vhv));
306 
307   dJdU[1][0][0] = 1. + 2. * u[1] * u[0];
308   dJdU[1][0][1] = u[0] * u[0] - 1.;
309   for (j = 0; j < 1; j++) {
310     vhv[j] = 0;
311     for (k = 0; k < 2; k++)
312       for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
313   }
314 
315   PetscCall(VecRestoreArrayRead(U, &u));
316   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
317   PetscCall(VecRestoreArrayRead(Vr, &vr));
318   PetscCall(VecRestoreArray(VHV[0], &vhv));
319   PetscFunctionReturn(0);
320 }
321 
322 static PetscErrorCode IHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
323   PetscFunctionBeginUser;
324   PetscFunctionReturn(0);
325 }
326 
327 /* Monitor timesteps and use interpolation to output at integer multiples of 0.1 */
328 static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, void *ctx) {
329   const PetscScalar *x;
330   PetscReal          tfinal, dt;
331   User               user = (User)ctx;
332   Vec                interpolatedX;
333 
334   PetscFunctionBeginUser;
335   PetscCall(TSGetTimeStep(ts, &dt));
336   PetscCall(TSGetMaxTime(ts, &tfinal));
337 
338   while (user->next_output <= t && user->next_output <= tfinal) {
339     PetscCall(VecDuplicate(X, &interpolatedX));
340     PetscCall(TSInterpolate(ts, user->next_output, interpolatedX));
341     PetscCall(VecGetArrayRead(interpolatedX, &x));
342     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%g] %" PetscInt_FMT " TS %g (dt = %g) X %g %g\n", (double)user->next_output, step, (double)t, (double)dt, (double)PetscRealPart(x[0]), (double)PetscRealPart(x[1])));
343     PetscCall(VecRestoreArrayRead(interpolatedX, &x));
344     PetscCall(VecDestroy(&interpolatedX));
345     user->next_output += PetscRealConstant(0.1);
346   }
347   PetscFunctionReturn(0);
348 }
349 
350 int main(int argc, char **argv) {
351   Vec                P;
352   PetscBool          monitor = PETSC_FALSE;
353   PetscScalar       *x_ptr;
354   const PetscScalar *y_ptr;
355   PetscMPIInt        size;
356   struct _n_User     user;
357   Tao                tao;
358   KSP                ksp;
359   PC                 pc;
360 
361   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
362      Initialize program
363      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
364   PetscFunctionBeginUser;
365   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
366   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
367   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
368 
369   /* Create TAO solver and set desired solution method */
370   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
371   PetscCall(TaoSetType(tao, TAOBQNLS));
372 
373   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
374     Set runtime options
375     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
376   user.next_output  = 0.0;
377   user.mu           = PetscRealConstant(1.0e3);
378   user.ftime        = PetscRealConstant(0.5);
379   user.implicitform = PETSC_TRUE;
380 
381   PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL));
382   PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &user.mu, NULL));
383   PetscCall(PetscOptionsGetBool(NULL, NULL, "-implicitform", &user.implicitform, NULL));
384 
385   /* Create necessary matrix and vectors, solve same ODE on every process */
386   PetscCall(MatCreate(PETSC_COMM_WORLD, &user.A));
387   PetscCall(MatSetSizes(user.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
388   PetscCall(MatSetFromOptions(user.A));
389   PetscCall(MatSetUp(user.A));
390   PetscCall(MatCreateVecs(user.A, &user.U, NULL));
391   PetscCall(MatCreateVecs(user.A, &user.Lambda[0], NULL));
392   PetscCall(MatCreateVecs(user.A, &user.Lambda2[0], NULL));
393   PetscCall(MatCreateVecs(user.A, &user.Ihp1[0], NULL));
394   PetscCall(MatCreateVecs(user.A, &user.Ihp2[0], NULL));
395 
396   PetscCall(MatCreate(PETSC_COMM_WORLD, &user.Jacp));
397   PetscCall(MatSetSizes(user.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 1));
398   PetscCall(MatSetFromOptions(user.Jacp));
399   PetscCall(MatSetUp(user.Jacp));
400   PetscCall(MatCreateVecs(user.Jacp, &user.Dir, NULL));
401   PetscCall(MatCreateVecs(user.Jacp, &user.Mup[0], NULL));
402   PetscCall(MatCreateVecs(user.Jacp, &user.Mup2[0], NULL));
403   PetscCall(MatCreateVecs(user.Jacp, &user.Ihp3[0], NULL));
404   PetscCall(MatCreateVecs(user.Jacp, &user.Ihp4[0], NULL));
405 
406   /* Create timestepping solver context */
407   PetscCall(TSCreate(PETSC_COMM_WORLD, &user.ts));
408   PetscCall(TSSetEquationType(user.ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
409   if (user.implicitform) {
410     PetscCall(TSSetIFunction(user.ts, NULL, IFunction, &user));
411     PetscCall(TSSetIJacobian(user.ts, user.A, user.A, IJacobian, &user));
412     PetscCall(TSSetType(user.ts, TSCN));
413   } else {
414     PetscCall(TSSetRHSFunction(user.ts, NULL, RHSFunction, &user));
415     PetscCall(TSSetRHSJacobian(user.ts, user.A, user.A, RHSJacobian, &user));
416     PetscCall(TSSetType(user.ts, TSRK));
417   }
418   PetscCall(TSSetRHSJacobianP(user.ts, user.Jacp, RHSJacobianP, &user));
419   PetscCall(TSSetMaxTime(user.ts, user.ftime));
420   PetscCall(TSSetExactFinalTime(user.ts, TS_EXACTFINALTIME_MATCHSTEP));
421 
422   if (monitor) { PetscCall(TSMonitorSet(user.ts, Monitor, &user, NULL)); }
423 
424   /* Set ODE initial conditions */
425   PetscCall(VecGetArray(user.U, &x_ptr));
426   x_ptr[0] = 2.0;
427   x_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * user.mu) - 292.0 / (2187.0 * user.mu * user.mu);
428   PetscCall(VecRestoreArray(user.U, &x_ptr));
429   PetscCall(TSSetTimeStep(user.ts, PetscRealConstant(0.001)));
430 
431   /* Set runtime options */
432   PetscCall(TSSetFromOptions(user.ts));
433 
434   PetscCall(TSSolve(user.ts, user.U));
435   PetscCall(VecGetArrayRead(user.U, &y_ptr));
436   user.ob[0] = y_ptr[0];
437   user.ob[1] = y_ptr[1];
438   PetscCall(VecRestoreArrayRead(user.U, &y_ptr));
439 
440   /* Save trajectory of solution so that TSAdjointSolve() may be used.
441      Skip checkpointing for the first TSSolve since no adjoint run follows it.
442    */
443   PetscCall(TSSetSaveTrajectory(user.ts));
444 
445   /* Optimization starts */
446   PetscCall(MatCreate(PETSC_COMM_WORLD, &user.H));
447   PetscCall(MatSetSizes(user.H, PETSC_DECIDE, PETSC_DECIDE, 1, 1));
448   PetscCall(MatSetUp(user.H)); /* Hessian should be symmetric. Do we need to do MatSetOption(user.H,MAT_SYMMETRIC,PETSC_TRUE) ? */
449 
450   /* Set initial solution guess */
451   PetscCall(MatCreateVecs(user.Jacp, &P, NULL));
452   PetscCall(VecGetArray(P, &x_ptr));
453   x_ptr[0] = PetscRealConstant(1.2);
454   PetscCall(VecRestoreArray(P, &x_ptr));
455   PetscCall(TaoSetSolution(tao, P));
456 
457   /* Set routine for function and gradient evaluation */
458   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&user));
459   PetscCall(TaoSetHessian(tao, user.H, user.H, FormHessian, (void *)&user));
460 
461   /* Check for any TAO command line options */
462   PetscCall(TaoGetKSP(tao, &ksp));
463   if (ksp) {
464     PetscCall(KSPGetPC(ksp, &pc));
465     PetscCall(PCSetType(pc, PCNONE));
466   }
467   PetscCall(TaoSetFromOptions(tao));
468 
469   PetscCall(TaoSolve(tao));
470 
471   PetscCall(VecView(P, PETSC_VIEWER_STDOUT_WORLD));
472   /* Free TAO data structures */
473   PetscCall(TaoDestroy(&tao));
474 
475   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
476      Free work space.  All PETSc objects should be destroyed when they
477      are no longer needed.
478    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
479   PetscCall(MatDestroy(&user.H));
480   PetscCall(MatDestroy(&user.A));
481   PetscCall(VecDestroy(&user.U));
482   PetscCall(MatDestroy(&user.Jacp));
483   PetscCall(VecDestroy(&user.Lambda[0]));
484   PetscCall(VecDestroy(&user.Mup[0]));
485   PetscCall(VecDestroy(&user.Lambda2[0]));
486   PetscCall(VecDestroy(&user.Mup2[0]));
487   PetscCall(VecDestroy(&user.Ihp1[0]));
488   PetscCall(VecDestroy(&user.Ihp2[0]));
489   PetscCall(VecDestroy(&user.Ihp3[0]));
490   PetscCall(VecDestroy(&user.Ihp4[0]));
491   PetscCall(VecDestroy(&user.Dir));
492   PetscCall(TSDestroy(&user.ts));
493   PetscCall(VecDestroy(&P));
494   PetscCall(PetscFinalize());
495   return 0;
496 }
497 
498 /* ------------------------------------------------------------------ */
499 /*
500    FormFunctionGradient - Evaluates the function and corresponding gradient.
501 
502    Input Parameters:
503    tao - the Tao context
504    X   - the input vector
505    ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()
506 
507    Output Parameters:
508    f   - the newly evaluated function
509    G   - the newly evaluated gradient
510 */
511 PetscErrorCode FormFunctionGradient(Tao tao, Vec P, PetscReal *f, Vec G, void *ctx) {
512   User               user_ptr = (User)ctx;
513   TS                 ts       = user_ptr->ts;
514   PetscScalar       *x_ptr, *g;
515   const PetscScalar *y_ptr;
516 
517   PetscFunctionBeginUser;
518   PetscCall(VecGetArrayRead(P, &y_ptr));
519   user_ptr->mu = y_ptr[0];
520   PetscCall(VecRestoreArrayRead(P, &y_ptr));
521 
522   PetscCall(TSSetTime(ts, 0.0));
523   PetscCall(TSSetStepNumber(ts, 0));
524   PetscCall(TSSetTimeStep(ts, PetscRealConstant(0.001))); /* can be overwritten by command line options */
525   PetscCall(TSSetFromOptions(ts));
526   PetscCall(VecGetArray(user_ptr->U, &x_ptr));
527   x_ptr[0] = 2.0;
528   x_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * user_ptr->mu) - 292.0 / (2187.0 * user_ptr->mu * user_ptr->mu);
529   PetscCall(VecRestoreArray(user_ptr->U, &x_ptr));
530 
531   PetscCall(TSSolve(ts, user_ptr->U));
532 
533   PetscCall(VecGetArrayRead(user_ptr->U, &y_ptr));
534   *f = (y_ptr[0] - user_ptr->ob[0]) * (y_ptr[0] - user_ptr->ob[0]) + (y_ptr[1] - user_ptr->ob[1]) * (y_ptr[1] - user_ptr->ob[1]);
535 
536   /*   Reset initial conditions for the adjoint integration */
537   PetscCall(VecGetArray(user_ptr->Lambda[0], &x_ptr));
538   x_ptr[0] = 2. * (y_ptr[0] - user_ptr->ob[0]);
539   x_ptr[1] = 2. * (y_ptr[1] - user_ptr->ob[1]);
540   PetscCall(VecRestoreArrayRead(user_ptr->U, &y_ptr));
541   PetscCall(VecRestoreArray(user_ptr->Lambda[0], &x_ptr));
542 
543   PetscCall(VecGetArray(user_ptr->Mup[0], &x_ptr));
544   x_ptr[0] = 0.0;
545   PetscCall(VecRestoreArray(user_ptr->Mup[0], &x_ptr));
546   PetscCall(TSSetCostGradients(ts, 1, user_ptr->Lambda, user_ptr->Mup));
547 
548   PetscCall(TSAdjointSolve(ts));
549 
550   PetscCall(VecGetArray(user_ptr->Mup[0], &x_ptr));
551   PetscCall(VecGetArrayRead(user_ptr->Lambda[0], &y_ptr));
552   PetscCall(VecGetArray(G, &g));
553   g[0] = y_ptr[1] * (-10.0 / (81.0 * user_ptr->mu * user_ptr->mu) + 2.0 * 292.0 / (2187.0 * user_ptr->mu * user_ptr->mu * user_ptr->mu)) + x_ptr[0];
554   PetscCall(VecRestoreArray(user_ptr->Mup[0], &x_ptr));
555   PetscCall(VecRestoreArrayRead(user_ptr->Lambda[0], &y_ptr));
556   PetscCall(VecRestoreArray(G, &g));
557   PetscFunctionReturn(0);
558 }
559 
560 PetscErrorCode FormHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx) {
561   User           user_ptr = (User)ctx;
562   PetscScalar    harr[1];
563   const PetscInt rows[1] = {0};
564   PetscInt       col     = 0;
565 
566   PetscFunctionBeginUser;
567   PetscCall(Adjoint2(P, harr, user_ptr));
568   PetscCall(MatSetValues(H, 1, rows, 1, &col, harr, INSERT_VALUES));
569 
570   PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
571   PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
572   if (H != Hpre) {
573     PetscCall(MatAssemblyBegin(Hpre, MAT_FINAL_ASSEMBLY));
574     PetscCall(MatAssemblyEnd(Hpre, MAT_FINAL_ASSEMBLY));
575   }
576   PetscFunctionReturn(0);
577 }
578 
579 PetscErrorCode Adjoint2(Vec P, PetscScalar arr[], User ctx) {
580   TS                 ts = ctx->ts;
581   const PetscScalar *z_ptr;
582   PetscScalar       *x_ptr, *y_ptr, dzdp, dzdp2;
583   Mat                tlmsen;
584 
585   PetscFunctionBeginUser;
586   /* Reset TSAdjoint so that AdjointSetUp will be called again */
587   PetscCall(TSAdjointReset(ts));
588 
589   /* The directional vector should be 1 since it is one-dimensional */
590   PetscCall(VecGetArray(ctx->Dir, &x_ptr));
591   x_ptr[0] = 1.;
592   PetscCall(VecRestoreArray(ctx->Dir, &x_ptr));
593 
594   PetscCall(VecGetArrayRead(P, &z_ptr));
595   ctx->mu = z_ptr[0];
596   PetscCall(VecRestoreArrayRead(P, &z_ptr));
597 
598   dzdp  = -10.0 / (81.0 * ctx->mu * ctx->mu) + 2.0 * 292.0 / (2187.0 * ctx->mu * ctx->mu * ctx->mu);
599   dzdp2 = 2. * 10.0 / (81.0 * ctx->mu * ctx->mu * ctx->mu) - 3.0 * 2.0 * 292.0 / (2187.0 * ctx->mu * ctx->mu * ctx->mu * ctx->mu);
600 
601   PetscCall(TSSetTime(ts, 0.0));
602   PetscCall(TSSetStepNumber(ts, 0));
603   PetscCall(TSSetTimeStep(ts, PetscRealConstant(0.001))); /* can be overwritten by command line options */
604   PetscCall(TSSetFromOptions(ts));
605   PetscCall(TSSetCostHessianProducts(ts, 1, ctx->Lambda2, ctx->Mup2, ctx->Dir));
606 
607   PetscCall(MatZeroEntries(ctx->Jacp));
608   PetscCall(MatSetValue(ctx->Jacp, 1, 0, dzdp, INSERT_VALUES));
609   PetscCall(MatAssemblyBegin(ctx->Jacp, MAT_FINAL_ASSEMBLY));
610   PetscCall(MatAssemblyEnd(ctx->Jacp, MAT_FINAL_ASSEMBLY));
611 
612   PetscCall(TSAdjointSetForward(ts, ctx->Jacp));
613   PetscCall(VecGetArray(ctx->U, &y_ptr));
614   y_ptr[0] = 2.0;
615   y_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * ctx->mu) - 292.0 / (2187.0 * ctx->mu * ctx->mu);
616   PetscCall(VecRestoreArray(ctx->U, &y_ptr));
617   PetscCall(TSSolve(ts, ctx->U));
618 
619   /* Set terminal conditions for first- and second-order adjonts */
620   PetscCall(VecGetArrayRead(ctx->U, &z_ptr));
621   PetscCall(VecGetArray(ctx->Lambda[0], &y_ptr));
622   y_ptr[0] = 2. * (z_ptr[0] - ctx->ob[0]);
623   y_ptr[1] = 2. * (z_ptr[1] - ctx->ob[1]);
624   PetscCall(VecRestoreArray(ctx->Lambda[0], &y_ptr));
625   PetscCall(VecRestoreArrayRead(ctx->U, &z_ptr));
626   PetscCall(VecGetArray(ctx->Mup[0], &y_ptr));
627   y_ptr[0] = 0.0;
628   PetscCall(VecRestoreArray(ctx->Mup[0], &y_ptr));
629   PetscCall(TSForwardGetSensitivities(ts, NULL, &tlmsen));
630   PetscCall(MatDenseGetColumn(tlmsen, 0, &x_ptr));
631   PetscCall(VecGetArray(ctx->Lambda2[0], &y_ptr));
632   y_ptr[0] = 2. * x_ptr[0];
633   y_ptr[1] = 2. * x_ptr[1];
634   PetscCall(VecRestoreArray(ctx->Lambda2[0], &y_ptr));
635   PetscCall(VecGetArray(ctx->Mup2[0], &y_ptr));
636   y_ptr[0] = 0.0;
637   PetscCall(VecRestoreArray(ctx->Mup2[0], &y_ptr));
638   PetscCall(MatDenseRestoreColumn(tlmsen, &x_ptr));
639   PetscCall(TSSetCostGradients(ts, 1, ctx->Lambda, ctx->Mup));
640   if (ctx->implicitform) {
641     PetscCall(TSSetIHessianProduct(ts, ctx->Ihp1, IHessianProductUU, ctx->Ihp2, IHessianProductUP, ctx->Ihp3, IHessianProductPU, ctx->Ihp4, IHessianProductPP, ctx));
642   } else {
643     PetscCall(TSSetRHSHessianProduct(ts, ctx->Ihp1, RHSHessianProductUU, ctx->Ihp2, RHSHessianProductUP, ctx->Ihp3, RHSHessianProductPU, ctx->Ihp4, RHSHessianProductPP, ctx));
644   }
645   PetscCall(TSAdjointSolve(ts));
646 
647   PetscCall(VecGetArray(ctx->Lambda[0], &x_ptr));
648   PetscCall(VecGetArray(ctx->Lambda2[0], &y_ptr));
649   PetscCall(VecGetArrayRead(ctx->Mup2[0], &z_ptr));
650 
651   arr[0] = x_ptr[1] * dzdp2 + y_ptr[1] * dzdp2 + z_ptr[0];
652 
653   PetscCall(VecRestoreArray(ctx->Lambda2[0], &x_ptr));
654   PetscCall(VecRestoreArray(ctx->Lambda2[0], &y_ptr));
655   PetscCall(VecRestoreArrayRead(ctx->Mup2[0], &z_ptr));
656 
657   /* Disable second-order adjoint mode */
658   PetscCall(TSAdjointReset(ts));
659   PetscCall(TSAdjointResetForward(ts));
660   PetscFunctionReturn(0);
661 }
662 
663 /*TEST
664     build:
665       requires: !complex !single
666     test:
667       args:  -implicitform 0 -ts_type rk -ts_adapt_type none -mu 10 -ts_dt 0.1 -viewer_binary_skip_info -tao_monitor -tao_view
668       output_file: output/ex20opt_p_1.out
669 
670     test:
671       suffix: 2
672       args:  -implicitform 0 -ts_type rk -ts_adapt_type none -mu 10 -ts_dt 0.01 -viewer_binary_skip_info -tao_monitor -tao_type bntr -tao_bnk_pc_type none
673       output_file: output/ex20opt_p_2.out
674 
675     test:
676       suffix: 3
677       args:  -ts_type cn -ts_adapt_type none -mu 100 -ts_dt 0.01 -viewer_binary_skip_info -tao_monitor -tao_view
678       output_file: output/ex20opt_p_3.out
679 
680     test:
681       suffix: 4
682       args:  -ts_type cn -ts_adapt_type none -mu 100 -ts_dt 0.01 -viewer_binary_skip_info -tao_monitor -tao_type bntr -tao_bnk_pc_type none
683       output_file: output/ex20opt_p_4.out
684 
685 TEST*/
686