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