xref: /petsc/src/ts/tutorials/optimal_control/ex1.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
1 #include <petsctao.h>
2 #include <petscts.h>
3 
4 typedef struct _n_aircraft *Aircraft;
5 struct _n_aircraft {
6   TS        ts, quadts;
7   Vec       V, W;   /* control variables V and W */
8   PetscInt  nsteps; /* number of time steps */
9   PetscReal ftime;
10   Mat       A, H;
11   Mat       Jacp, DRDU, DRDP;
12   Vec       U, Lambda[1], Mup[1], Lambda2[1], Mup2[1], Dir;
13   Vec       rhshp1[1], rhshp2[1], rhshp3[1], rhshp4[1], inthp1[1], inthp2[1], inthp3[1], inthp4[1];
14   PetscReal lv, lw;
15   PetscBool mf, eh;
16 };
17 
18 PetscErrorCode FormObjFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
19 PetscErrorCode FormObjHessian(Tao, Vec, Mat, Mat, void *);
20 PetscErrorCode ComputeObjHessianWithSOA(Vec, PetscScalar[], Aircraft);
21 PetscErrorCode MatrixFreeObjHessian(Tao, Vec, Mat, Mat, void *);
22 PetscErrorCode MyMatMult(Mat, Vec, Vec);
23 
24 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, void *ctx) {
25   Aircraft           actx = (Aircraft)ctx;
26   const PetscScalar *u, *v, *w;
27   PetscScalar       *f;
28   PetscInt           step;
29 
30   PetscFunctionBeginUser;
31   PetscCall(TSGetStepNumber(ts, &step));
32   PetscCall(VecGetArrayRead(U, &u));
33   PetscCall(VecGetArrayRead(actx->V, &v));
34   PetscCall(VecGetArrayRead(actx->W, &w));
35   PetscCall(VecGetArray(F, &f));
36   f[0] = v[step] * PetscCosReal(w[step]);
37   f[1] = v[step] * PetscSinReal(w[step]);
38   PetscCall(VecRestoreArrayRead(U, &u));
39   PetscCall(VecRestoreArrayRead(actx->V, &v));
40   PetscCall(VecRestoreArrayRead(actx->W, &w));
41   PetscCall(VecRestoreArray(F, &f));
42   PetscFunctionReturn(0);
43 }
44 
45 static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec U, Mat A, void *ctx) {
46   Aircraft           actx = (Aircraft)ctx;
47   const PetscScalar *u, *v, *w;
48   PetscInt           step, rows[2] = {0, 1}, rowcol[2];
49   PetscScalar        Jp[2][2];
50 
51   PetscFunctionBeginUser;
52   PetscCall(MatZeroEntries(A));
53   PetscCall(TSGetStepNumber(ts, &step));
54   PetscCall(VecGetArrayRead(U, &u));
55   PetscCall(VecGetArrayRead(actx->V, &v));
56   PetscCall(VecGetArrayRead(actx->W, &w));
57 
58   Jp[0][0] = PetscCosReal(w[step]);
59   Jp[0][1] = -v[step] * PetscSinReal(w[step]);
60   Jp[1][0] = PetscSinReal(w[step]);
61   Jp[1][1] = v[step] * PetscCosReal(w[step]);
62 
63   PetscCall(VecRestoreArrayRead(U, &u));
64   PetscCall(VecRestoreArrayRead(actx->V, &v));
65   PetscCall(VecRestoreArrayRead(actx->W, &w));
66 
67   rowcol[0] = 2 * step;
68   rowcol[1] = 2 * step + 1;
69   PetscCall(MatSetValues(A, 2, rows, 2, rowcol, &Jp[0][0], INSERT_VALUES));
70 
71   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
72   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
73   PetscFunctionReturn(0);
74 }
75 
76 static PetscErrorCode RHSHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
77   PetscFunctionBeginUser;
78   PetscFunctionReturn(0);
79 }
80 
81 static PetscErrorCode RHSHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
82   PetscFunctionBeginUser;
83   PetscFunctionReturn(0);
84 }
85 
86 static PetscErrorCode RHSHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
87   PetscFunctionBeginUser;
88   PetscFunctionReturn(0);
89 }
90 
91 static PetscErrorCode RHSHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
92   Aircraft           actx = (Aircraft)ctx;
93   const PetscScalar *v, *w, *vl, *vr, *u;
94   PetscScalar       *vhv;
95   PetscScalar        dJpdP[2][2][2] = {{{0}}};
96   PetscInt           step, i, j, k;
97 
98   PetscFunctionBeginUser;
99   PetscCall(TSGetStepNumber(ts, &step));
100   PetscCall(VecGetArrayRead(U, &u));
101   PetscCall(VecGetArrayRead(actx->V, &v));
102   PetscCall(VecGetArrayRead(actx->W, &w));
103   PetscCall(VecGetArrayRead(Vl[0], &vl));
104   PetscCall(VecGetArrayRead(Vr, &vr));
105   PetscCall(VecSet(VHV[0], 0.0));
106   PetscCall(VecGetArray(VHV[0], &vhv));
107 
108   dJpdP[0][0][1] = -PetscSinReal(w[step]);
109   dJpdP[0][1][0] = -PetscSinReal(w[step]);
110   dJpdP[0][1][1] = -v[step] * PetscCosReal(w[step]);
111   dJpdP[1][0][1] = PetscCosReal(w[step]);
112   dJpdP[1][1][0] = PetscCosReal(w[step]);
113   dJpdP[1][1][1] = -v[step] * PetscSinReal(w[step]);
114 
115   for (j = 0; j < 2; j++) {
116     vhv[2 * step + j] = 0;
117     for (k = 0; k < 2; k++)
118       for (i = 0; i < 2; i++) vhv[2 * step + j] += vl[i] * dJpdP[i][j][k] * vr[2 * step + k];
119   }
120   PetscCall(VecRestoreArrayRead(U, &u));
121   PetscCall(VecRestoreArrayRead(Vl[0], &vl));
122   PetscCall(VecRestoreArrayRead(Vr, &vr));
123   PetscCall(VecRestoreArray(VHV[0], &vhv));
124   PetscFunctionReturn(0);
125 }
126 
127 /* Vl in NULL,updates to VHV must be added */
128 static PetscErrorCode IntegrandHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
129   Aircraft           actx = (Aircraft)ctx;
130   const PetscScalar *v, *w, *vr, *u;
131   PetscScalar       *vhv;
132   PetscScalar        dRudU[2][2] = {{0}};
133   PetscInt           step, j, k;
134 
135   PetscFunctionBeginUser;
136   PetscCall(TSGetStepNumber(ts, &step));
137   PetscCall(VecGetArrayRead(U, &u));
138   PetscCall(VecGetArrayRead(actx->V, &v));
139   PetscCall(VecGetArrayRead(actx->W, &w));
140   PetscCall(VecGetArrayRead(Vr, &vr));
141   PetscCall(VecGetArray(VHV[0], &vhv));
142 
143   dRudU[0][0] = 2.0;
144   dRudU[1][1] = 2.0;
145 
146   for (j = 0; j < 2; j++) {
147     vhv[j] = 0;
148     for (k = 0; k < 2; k++) vhv[j] += dRudU[j][k] * vr[k];
149   }
150   PetscCall(VecRestoreArrayRead(U, &u));
151   PetscCall(VecRestoreArrayRead(Vr, &vr));
152   PetscCall(VecRestoreArray(VHV[0], &vhv));
153   PetscFunctionReturn(0);
154 }
155 
156 static PetscErrorCode IntegrandHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
157   PetscFunctionBeginUser;
158   PetscFunctionReturn(0);
159 }
160 
161 static PetscErrorCode IntegrandHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
162   PetscFunctionBeginUser;
163   PetscFunctionReturn(0);
164 }
165 
166 static PetscErrorCode IntegrandHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) {
167   PetscFunctionBeginUser;
168   PetscFunctionReturn(0);
169 }
170 
171 static PetscErrorCode CostIntegrand(TS ts, PetscReal t, Vec U, Vec R, void *ctx) {
172   Aircraft           actx = (Aircraft)ctx;
173   PetscScalar       *r;
174   PetscReal          dx, dy;
175   const PetscScalar *u;
176 
177   PetscFunctionBegin;
178   PetscCall(VecGetArrayRead(U, &u));
179   PetscCall(VecGetArray(R, &r));
180   dx   = u[0] - actx->lv * t * PetscCosReal(actx->lw);
181   dy   = u[1] - actx->lv * t * PetscSinReal(actx->lw);
182   r[0] = dx * dx + dy * dy;
183   PetscCall(VecRestoreArray(R, &r));
184   PetscCall(VecRestoreArrayRead(U, &u));
185   PetscFunctionReturn(0);
186 }
187 
188 static PetscErrorCode DRDUJacobianTranspose(TS ts, PetscReal t, Vec U, Mat DRDU, Mat B, void *ctx) {
189   Aircraft           actx = (Aircraft)ctx;
190   PetscScalar        drdu[2][1];
191   const PetscScalar *u;
192   PetscReal          dx, dy;
193   PetscInt           row[] = {0, 1}, col[] = {0};
194 
195   PetscFunctionBegin;
196   PetscCall(VecGetArrayRead(U, &u));
197   dx         = u[0] - actx->lv * t * PetscCosReal(actx->lw);
198   dy         = u[1] - actx->lv * t * PetscSinReal(actx->lw);
199   drdu[0][0] = 2. * dx;
200   drdu[1][0] = 2. * dy;
201   PetscCall(MatSetValues(DRDU, 2, row, 1, col, &drdu[0][0], INSERT_VALUES));
202   PetscCall(VecRestoreArrayRead(U, &u));
203   PetscCall(MatAssemblyBegin(DRDU, MAT_FINAL_ASSEMBLY));
204   PetscCall(MatAssemblyEnd(DRDU, MAT_FINAL_ASSEMBLY));
205   PetscFunctionReturn(0);
206 }
207 
208 static PetscErrorCode DRDPJacobianTranspose(TS ts, PetscReal t, Vec U, Mat DRDP, void *ctx) {
209   PetscFunctionBegin;
210   PetscCall(MatZeroEntries(DRDP));
211   PetscCall(MatAssemblyBegin(DRDP, MAT_FINAL_ASSEMBLY));
212   PetscCall(MatAssemblyEnd(DRDP, MAT_FINAL_ASSEMBLY));
213   PetscFunctionReturn(0);
214 }
215 
216 int main(int argc, char **argv) {
217   Vec                P, PL, PU;
218   struct _n_aircraft aircraft;
219   PetscMPIInt        size;
220   Tao                tao;
221   KSP                ksp;
222   PC                 pc;
223   PetscScalar       *u, *p;
224   PetscInt           i;
225 
226   /* Initialize program */
227   PetscFunctionBeginUser;
228   PetscCall(PetscInitialize(&argc, &argv, NULL, NULL));
229   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
230   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
231 
232   /* Parameter settings */
233   aircraft.ftime  = 1.;            /* time interval in hour */
234   aircraft.nsteps = 10;            /* number of steps */
235   aircraft.lv     = 2.0;           /* leader speed in kmph */
236   aircraft.lw     = PETSC_PI / 4.; /* leader heading angle */
237 
238   PetscCall(PetscOptionsGetReal(NULL, NULL, "-ftime", &aircraft.ftime, NULL));
239   PetscCall(PetscOptionsGetInt(NULL, NULL, "-nsteps", &aircraft.nsteps, NULL));
240   PetscCall(PetscOptionsHasName(NULL, NULL, "-matrixfree", &aircraft.mf));
241   PetscCall(PetscOptionsHasName(NULL, NULL, "-exacthessian", &aircraft.eh));
242 
243   /* Create TAO solver and set desired solution method */
244   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
245   PetscCall(TaoSetType(tao, TAOBQNLS));
246 
247   /* Create necessary matrix and vectors, solve same ODE on every process */
248   PetscCall(MatCreate(PETSC_COMM_WORLD, &aircraft.A));
249   PetscCall(MatSetSizes(aircraft.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
250   PetscCall(MatSetFromOptions(aircraft.A));
251   PetscCall(MatSetUp(aircraft.A));
252   PetscCall(MatAssemblyBegin(aircraft.A, MAT_FINAL_ASSEMBLY));
253   PetscCall(MatAssemblyEnd(aircraft.A, MAT_FINAL_ASSEMBLY));
254   PetscCall(MatShift(aircraft.A, 1));
255   PetscCall(MatShift(aircraft.A, -1));
256 
257   PetscCall(MatCreate(PETSC_COMM_WORLD, &aircraft.Jacp));
258   PetscCall(MatSetSizes(aircraft.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 2 * aircraft.nsteps));
259   PetscCall(MatSetFromOptions(aircraft.Jacp));
260   PetscCall(MatSetUp(aircraft.Jacp));
261   PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2 * aircraft.nsteps, 1, NULL, &aircraft.DRDP));
262   PetscCall(MatSetUp(aircraft.DRDP));
263   PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, &aircraft.DRDU));
264   PetscCall(MatSetUp(aircraft.DRDU));
265 
266   /* Create timestepping solver context */
267   PetscCall(TSCreate(PETSC_COMM_WORLD, &aircraft.ts));
268   PetscCall(TSSetType(aircraft.ts, TSRK));
269   PetscCall(TSSetRHSFunction(aircraft.ts, NULL, RHSFunction, &aircraft));
270   PetscCall(TSSetRHSJacobian(aircraft.ts, aircraft.A, aircraft.A, TSComputeRHSJacobianConstant, &aircraft));
271   PetscCall(TSSetRHSJacobianP(aircraft.ts, aircraft.Jacp, RHSJacobianP, &aircraft));
272   PetscCall(TSSetExactFinalTime(aircraft.ts, TS_EXACTFINALTIME_MATCHSTEP));
273   PetscCall(TSSetEquationType(aircraft.ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
274 
275   /* Set initial conditions */
276   PetscCall(MatCreateVecs(aircraft.A, &aircraft.U, NULL));
277   PetscCall(TSSetSolution(aircraft.ts, aircraft.U));
278   PetscCall(VecGetArray(aircraft.U, &u));
279   u[0] = 1.5;
280   u[1] = 0;
281   PetscCall(VecRestoreArray(aircraft.U, &u));
282   PetscCall(VecCreate(PETSC_COMM_WORLD, &aircraft.V));
283   PetscCall(VecSetSizes(aircraft.V, PETSC_DECIDE, aircraft.nsteps));
284   PetscCall(VecSetUp(aircraft.V));
285   PetscCall(VecDuplicate(aircraft.V, &aircraft.W));
286   PetscCall(VecSet(aircraft.V, 1.));
287   PetscCall(VecSet(aircraft.W, PETSC_PI / 4.));
288 
289   /* Save trajectory of solution so that TSAdjointSolve() may be used */
290   PetscCall(TSSetSaveTrajectory(aircraft.ts));
291 
292   /* Set sensitivity context */
293   PetscCall(TSCreateQuadratureTS(aircraft.ts, PETSC_FALSE, &aircraft.quadts));
294   PetscCall(TSSetRHSFunction(aircraft.quadts, NULL, (TSRHSFunction)CostIntegrand, &aircraft));
295   PetscCall(TSSetRHSJacobian(aircraft.quadts, aircraft.DRDU, aircraft.DRDU, (TSRHSJacobian)DRDUJacobianTranspose, &aircraft));
296   PetscCall(TSSetRHSJacobianP(aircraft.quadts, aircraft.DRDP, (TSRHSJacobianP)DRDPJacobianTranspose, &aircraft));
297   PetscCall(MatCreateVecs(aircraft.A, &aircraft.Lambda[0], NULL));
298   PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Mup[0], NULL));
299   if (aircraft.eh) {
300     PetscCall(MatCreateVecs(aircraft.A, &aircraft.rhshp1[0], NULL));
301     PetscCall(MatCreateVecs(aircraft.A, &aircraft.rhshp2[0], NULL));
302     PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.rhshp3[0], NULL));
303     PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.rhshp4[0], NULL));
304     PetscCall(MatCreateVecs(aircraft.DRDU, &aircraft.inthp1[0], NULL));
305     PetscCall(MatCreateVecs(aircraft.DRDU, &aircraft.inthp2[0], NULL));
306     PetscCall(MatCreateVecs(aircraft.DRDP, &aircraft.inthp3[0], NULL));
307     PetscCall(MatCreateVecs(aircraft.DRDP, &aircraft.inthp4[0], NULL));
308     PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Dir, NULL));
309     PetscCall(TSSetRHSHessianProduct(aircraft.ts, aircraft.rhshp1, RHSHessianProductUU, aircraft.rhshp2, RHSHessianProductUP, aircraft.rhshp3, RHSHessianProductPU, aircraft.rhshp4, RHSHessianProductPP, &aircraft));
310     PetscCall(TSSetRHSHessianProduct(aircraft.quadts, aircraft.inthp1, IntegrandHessianProductUU, aircraft.inthp2, IntegrandHessianProductUP, aircraft.inthp3, IntegrandHessianProductPU, aircraft.inthp4, IntegrandHessianProductPP, &aircraft));
311     PetscCall(MatCreateVecs(aircraft.A, &aircraft.Lambda2[0], NULL));
312     PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Mup2[0], NULL));
313   }
314   PetscCall(TSSetFromOptions(aircraft.ts));
315   PetscCall(TSSetMaxTime(aircraft.ts, aircraft.ftime));
316   PetscCall(TSSetTimeStep(aircraft.ts, aircraft.ftime / aircraft.nsteps));
317 
318   /* Set initial solution guess */
319   PetscCall(MatCreateVecs(aircraft.Jacp, &P, NULL));
320   PetscCall(VecGetArray(P, &p));
321   for (i = 0; i < aircraft.nsteps; i++) {
322     p[2 * i]     = 2.0;
323     p[2 * i + 1] = PETSC_PI / 2.0;
324   }
325   PetscCall(VecRestoreArray(P, &p));
326   PetscCall(VecDuplicate(P, &PU));
327   PetscCall(VecDuplicate(P, &PL));
328   PetscCall(VecGetArray(PU, &p));
329   for (i = 0; i < aircraft.nsteps; i++) {
330     p[2 * i]     = 2.0;
331     p[2 * i + 1] = PETSC_PI;
332   }
333   PetscCall(VecRestoreArray(PU, &p));
334   PetscCall(VecGetArray(PL, &p));
335   for (i = 0; i < aircraft.nsteps; i++) {
336     p[2 * i]     = 0.0;
337     p[2 * i + 1] = -PETSC_PI;
338   }
339   PetscCall(VecRestoreArray(PL, &p));
340 
341   PetscCall(TaoSetSolution(tao, P));
342   PetscCall(TaoSetVariableBounds(tao, PL, PU));
343   /* Set routine for function and gradient evaluation */
344   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormObjFunctionGradient, (void *)&aircraft));
345 
346   if (aircraft.eh) {
347     if (aircraft.mf) {
348       PetscCall(MatCreateShell(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2 * aircraft.nsteps, 2 * aircraft.nsteps, (void *)&aircraft, &aircraft.H));
349       PetscCall(MatShellSetOperation(aircraft.H, MATOP_MULT, (void (*)(void))MyMatMult));
350       PetscCall(MatSetOption(aircraft.H, MAT_SYMMETRIC, PETSC_TRUE));
351       PetscCall(TaoSetHessian(tao, aircraft.H, aircraft.H, MatrixFreeObjHessian, (void *)&aircraft));
352     } else {
353       PetscCall(MatCreateDense(MPI_COMM_WORLD, PETSC_DETERMINE, PETSC_DETERMINE, 2 * aircraft.nsteps, 2 * aircraft.nsteps, NULL, &(aircraft.H)));
354       PetscCall(MatSetOption(aircraft.H, MAT_SYMMETRIC, PETSC_TRUE));
355       PetscCall(TaoSetHessian(tao, aircraft.H, aircraft.H, FormObjHessian, (void *)&aircraft));
356     }
357   }
358 
359   /* Check for any TAO command line options */
360   PetscCall(TaoGetKSP(tao, &ksp));
361   if (ksp) {
362     PetscCall(KSPGetPC(ksp, &pc));
363     PetscCall(PCSetType(pc, PCNONE));
364   }
365   PetscCall(TaoSetFromOptions(tao));
366 
367   PetscCall(TaoSolve(tao));
368   PetscCall(VecView(P, PETSC_VIEWER_STDOUT_WORLD));
369 
370   /* Free TAO data structures */
371   PetscCall(TaoDestroy(&tao));
372 
373   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
374      Free work space.  All PETSc objects should be destroyed when they
375      are no longer needed.
376    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
377   PetscCall(TSDestroy(&aircraft.ts));
378   PetscCall(MatDestroy(&aircraft.A));
379   PetscCall(VecDestroy(&aircraft.U));
380   PetscCall(VecDestroy(&aircraft.V));
381   PetscCall(VecDestroy(&aircraft.W));
382   PetscCall(VecDestroy(&P));
383   PetscCall(VecDestroy(&PU));
384   PetscCall(VecDestroy(&PL));
385   PetscCall(MatDestroy(&aircraft.Jacp));
386   PetscCall(MatDestroy(&aircraft.DRDU));
387   PetscCall(MatDestroy(&aircraft.DRDP));
388   PetscCall(VecDestroy(&aircraft.Lambda[0]));
389   PetscCall(VecDestroy(&aircraft.Mup[0]));
390   PetscCall(VecDestroy(&P));
391   if (aircraft.eh) {
392     PetscCall(VecDestroy(&aircraft.Lambda2[0]));
393     PetscCall(VecDestroy(&aircraft.Mup2[0]));
394     PetscCall(VecDestroy(&aircraft.Dir));
395     PetscCall(VecDestroy(&aircraft.rhshp1[0]));
396     PetscCall(VecDestroy(&aircraft.rhshp2[0]));
397     PetscCall(VecDestroy(&aircraft.rhshp3[0]));
398     PetscCall(VecDestroy(&aircraft.rhshp4[0]));
399     PetscCall(VecDestroy(&aircraft.inthp1[0]));
400     PetscCall(VecDestroy(&aircraft.inthp2[0]));
401     PetscCall(VecDestroy(&aircraft.inthp3[0]));
402     PetscCall(VecDestroy(&aircraft.inthp4[0]));
403     PetscCall(MatDestroy(&aircraft.H));
404   }
405   PetscCall(PetscFinalize());
406   return 0;
407 }
408 
409 /*
410    FormObjFunctionGradient - Evaluates the function and corresponding gradient.
411 
412    Input Parameters:
413    tao - the Tao context
414    P   - the input vector
415    ctx - optional aircraft-defined context, as set by TaoSetObjectiveAndGradient()
416 
417    Output Parameters:
418    f   - the newly evaluated function
419    G   - the newly evaluated gradient
420 */
421 PetscErrorCode FormObjFunctionGradient(Tao tao, Vec P, PetscReal *f, Vec G, void *ctx) {
422   Aircraft           actx = (Aircraft)ctx;
423   TS                 ts   = actx->ts;
424   Vec                Q;
425   const PetscScalar *p, *q;
426   PetscScalar       *u, *v, *w;
427   PetscInt           i;
428 
429   PetscFunctionBeginUser;
430   PetscCall(VecGetArrayRead(P, &p));
431   PetscCall(VecGetArray(actx->V, &v));
432   PetscCall(VecGetArray(actx->W, &w));
433   for (i = 0; i < actx->nsteps; i++) {
434     v[i] = p[2 * i];
435     w[i] = p[2 * i + 1];
436   }
437   PetscCall(VecRestoreArrayRead(P, &p));
438   PetscCall(VecRestoreArray(actx->V, &v));
439   PetscCall(VecRestoreArray(actx->W, &w));
440 
441   PetscCall(TSSetTime(ts, 0.0));
442   PetscCall(TSSetStepNumber(ts, 0));
443   PetscCall(TSSetFromOptions(ts));
444   PetscCall(TSSetTimeStep(ts, actx->ftime / actx->nsteps));
445 
446   /* reinitialize system state */
447   PetscCall(VecGetArray(actx->U, &u));
448   u[0] = 2.0;
449   u[1] = 0;
450   PetscCall(VecRestoreArray(actx->U, &u));
451 
452   /* reinitialize the integral value */
453   PetscCall(TSGetCostIntegral(ts, &Q));
454   PetscCall(VecSet(Q, 0.0));
455 
456   PetscCall(TSSolve(ts, actx->U));
457 
458   /* Reset initial conditions for the adjoint integration */
459   PetscCall(VecSet(actx->Lambda[0], 0.0));
460   PetscCall(VecSet(actx->Mup[0], 0.0));
461   PetscCall(TSSetCostGradients(ts, 1, actx->Lambda, actx->Mup));
462 
463   PetscCall(TSAdjointSolve(ts));
464   PetscCall(VecCopy(actx->Mup[0], G));
465   PetscCall(TSGetCostIntegral(ts, &Q));
466   PetscCall(VecGetArrayRead(Q, &q));
467   *f = q[0];
468   PetscCall(VecRestoreArrayRead(Q, &q));
469   PetscFunctionReturn(0);
470 }
471 
472 PetscErrorCode FormObjHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx) {
473   Aircraft           actx = (Aircraft)ctx;
474   const PetscScalar *p;
475   PetscScalar       *harr, *v, *w, one = 1.0;
476   PetscInt           ind[1];
477   PetscInt          *cols, i;
478   Vec                Dir;
479 
480   PetscFunctionBeginUser;
481   /* set up control parameters */
482   PetscCall(VecGetArrayRead(P, &p));
483   PetscCall(VecGetArray(actx->V, &v));
484   PetscCall(VecGetArray(actx->W, &w));
485   for (i = 0; i < actx->nsteps; i++) {
486     v[i] = p[2 * i];
487     w[i] = p[2 * i + 1];
488   }
489   PetscCall(VecRestoreArrayRead(P, &p));
490   PetscCall(VecRestoreArray(actx->V, &v));
491   PetscCall(VecRestoreArray(actx->W, &w));
492 
493   PetscCall(PetscMalloc1(2 * actx->nsteps, &harr));
494   PetscCall(PetscMalloc1(2 * actx->nsteps, &cols));
495   for (i = 0; i < 2 * actx->nsteps; i++) cols[i] = i;
496   PetscCall(VecDuplicate(P, &Dir));
497   for (i = 0; i < 2 * actx->nsteps; i++) {
498     ind[0] = i;
499     PetscCall(VecSet(Dir, 0.0));
500     PetscCall(VecSetValues(Dir, 1, ind, &one, INSERT_VALUES));
501     PetscCall(VecAssemblyBegin(Dir));
502     PetscCall(VecAssemblyEnd(Dir));
503     PetscCall(ComputeObjHessianWithSOA(Dir, harr, actx));
504     PetscCall(MatSetValues(H, 1, ind, 2 * actx->nsteps, cols, harr, INSERT_VALUES));
505     PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
506     PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
507     if (H != Hpre) {
508       PetscCall(MatAssemblyBegin(Hpre, MAT_FINAL_ASSEMBLY));
509       PetscCall(MatAssemblyEnd(Hpre, MAT_FINAL_ASSEMBLY));
510     }
511   }
512   PetscCall(PetscFree(cols));
513   PetscCall(PetscFree(harr));
514   PetscCall(VecDestroy(&Dir));
515   PetscFunctionReturn(0);
516 }
517 
518 PetscErrorCode MatrixFreeObjHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx) {
519   Aircraft           actx = (Aircraft)ctx;
520   PetscScalar       *v, *w;
521   const PetscScalar *p;
522   PetscInt           i;
523 
524   PetscFunctionBegin;
525   PetscCall(VecGetArrayRead(P, &p));
526   PetscCall(VecGetArray(actx->V, &v));
527   PetscCall(VecGetArray(actx->W, &w));
528   for (i = 0; i < actx->nsteps; i++) {
529     v[i] = p[2 * i];
530     w[i] = p[2 * i + 1];
531   }
532   PetscCall(VecRestoreArrayRead(P, &p));
533   PetscCall(VecRestoreArray(actx->V, &v));
534   PetscCall(VecRestoreArray(actx->W, &w));
535   PetscFunctionReturn(0);
536 }
537 
538 PetscErrorCode MyMatMult(Mat H_shell, Vec X, Vec Y) {
539   PetscScalar *y;
540   void        *ptr;
541 
542   PetscFunctionBegin;
543   PetscCall(MatShellGetContext(H_shell, &ptr));
544   PetscCall(VecGetArray(Y, &y));
545   PetscCall(ComputeObjHessianWithSOA(X, y, (Aircraft)ptr));
546   PetscCall(VecRestoreArray(Y, &y));
547   PetscFunctionReturn(0);
548 }
549 
550 PetscErrorCode ComputeObjHessianWithSOA(Vec Dir, PetscScalar arr[], Aircraft actx) {
551   TS                 ts = actx->ts;
552   const PetscScalar *z_ptr;
553   PetscScalar       *u;
554   Vec                Q;
555   PetscInt           i;
556 
557   PetscFunctionBeginUser;
558   /* Reset TSAdjoint so that AdjointSetUp will be called again */
559   PetscCall(TSAdjointReset(ts));
560 
561   PetscCall(TSSetTime(ts, 0.0));
562   PetscCall(TSSetStepNumber(ts, 0));
563   PetscCall(TSSetFromOptions(ts));
564   PetscCall(TSSetTimeStep(ts, actx->ftime / actx->nsteps));
565   PetscCall(TSSetCostHessianProducts(actx->ts, 1, actx->Lambda2, actx->Mup2, Dir));
566 
567   /* reinitialize system state */
568   PetscCall(VecGetArray(actx->U, &u));
569   u[0] = 2.0;
570   u[1] = 0;
571   PetscCall(VecRestoreArray(actx->U, &u));
572 
573   /* reinitialize the integral value */
574   PetscCall(TSGetCostIntegral(ts, &Q));
575   PetscCall(VecSet(Q, 0.0));
576 
577   /* initialize tlm variable */
578   PetscCall(MatZeroEntries(actx->Jacp));
579   PetscCall(MatAssemblyBegin(actx->Jacp, MAT_FINAL_ASSEMBLY));
580   PetscCall(MatAssemblyEnd(actx->Jacp, MAT_FINAL_ASSEMBLY));
581   PetscCall(TSAdjointSetForward(ts, actx->Jacp));
582 
583   PetscCall(TSSolve(ts, actx->U));
584 
585   /* Set terminal conditions for first- and second-order adjonts */
586   PetscCall(VecSet(actx->Lambda[0], 0.0));
587   PetscCall(VecSet(actx->Mup[0], 0.0));
588   PetscCall(VecSet(actx->Lambda2[0], 0.0));
589   PetscCall(VecSet(actx->Mup2[0], 0.0));
590   PetscCall(TSSetCostGradients(ts, 1, actx->Lambda, actx->Mup));
591 
592   PetscCall(TSGetCostIntegral(ts, &Q));
593 
594   /* Reset initial conditions for the adjoint integration */
595   PetscCall(TSAdjointSolve(ts));
596 
597   /* initial condition does not depend on p, so that lambda is not needed to assemble G */
598   PetscCall(VecGetArrayRead(actx->Mup2[0], &z_ptr));
599   for (i = 0; i < 2 * actx->nsteps; i++) arr[i] = z_ptr[i];
600   PetscCall(VecRestoreArrayRead(actx->Mup2[0], &z_ptr));
601 
602   /* Disable second-order adjoint mode */
603   PetscCall(TSAdjointReset(ts));
604   PetscCall(TSAdjointResetForward(ts));
605   PetscFunctionReturn(0);
606 }
607 
608 /*TEST
609     build:
610       requires: !complex !single
611 
612     test:
613       args:  -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_gatol 1e-7
614 
615     test:
616       suffix: 2
617       args:  -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_view -tao_type bntr -tao_bnk_pc_type none -exacthessian
618 
619     test:
620       suffix: 3
621       args:  -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_view -tao_type bntr -tao_bnk_pc_type none -exacthessian -matrixfree
622 TEST*/
623