xref: /petsc/src/ts/tutorials/optimal_control/ex1.c (revision f97672e55eacc8688507b9471cd7ec2664d7f203)
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   PetscCall(PetscInitialize(&argc,&argv,NULL,NULL));
244   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
245   PetscCheck(size == 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
246 
247   /* Parameter settings */
248   aircraft.ftime = 1.;   /* time interval in hour */
249   aircraft.nsteps = 10; /* number of steps */
250   aircraft.lv = 2.0; /* leader speed in kmph */
251   aircraft.lw = PETSC_PI/4.; /* leader heading angle */
252 
253   PetscCall(PetscOptionsGetReal(NULL,NULL,"-ftime",&aircraft.ftime,NULL));
254   PetscCall(PetscOptionsGetInt(NULL,NULL,"-nsteps",&aircraft.nsteps,NULL));
255   PetscCall(PetscOptionsHasName(NULL,NULL,"-matrixfree",&aircraft.mf));
256   PetscCall(PetscOptionsHasName(NULL,NULL,"-exacthessian",&aircraft.eh));
257 
258   /* Create TAO solver and set desired solution method */
259   PetscCall(TaoCreate(PETSC_COMM_WORLD,&tao));
260   PetscCall(TaoSetType(tao,TAOBQNLS));
261 
262   /* Create necessary matrix and vectors, solve same ODE on every process */
263   PetscCall(MatCreate(PETSC_COMM_WORLD,&aircraft.A));
264   PetscCall(MatSetSizes(aircraft.A,PETSC_DECIDE,PETSC_DECIDE,2,2));
265   PetscCall(MatSetFromOptions(aircraft.A));
266   PetscCall(MatSetUp(aircraft.A));
267   PetscCall(MatAssemblyBegin(aircraft.A,MAT_FINAL_ASSEMBLY));
268   PetscCall(MatAssemblyEnd(aircraft.A,MAT_FINAL_ASSEMBLY));
269   PetscCall(MatShift(aircraft.A,1));
270   PetscCall(MatShift(aircraft.A,-1));
271 
272   PetscCall(MatCreate(PETSC_COMM_WORLD,&aircraft.Jacp));
273   PetscCall(MatSetSizes(aircraft.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2*aircraft.nsteps));
274   PetscCall(MatSetFromOptions(aircraft.Jacp));
275   PetscCall(MatSetUp(aircraft.Jacp));
276   PetscCall(MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2*aircraft.nsteps,1,NULL,&aircraft.DRDP));
277   PetscCall(MatSetUp(aircraft.DRDP));
278   PetscCall(MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2,1,NULL,&aircraft.DRDU));
279   PetscCall(MatSetUp(aircraft.DRDU));
280 
281   /* Create timestepping solver context */
282   PetscCall(TSCreate(PETSC_COMM_WORLD,&aircraft.ts));
283   PetscCall(TSSetType(aircraft.ts,TSRK));
284   PetscCall(TSSetRHSFunction(aircraft.ts,NULL,RHSFunction,&aircraft));
285   PetscCall(TSSetRHSJacobian(aircraft.ts,aircraft.A,aircraft.A,TSComputeRHSJacobianConstant,&aircraft));
286   PetscCall(TSSetRHSJacobianP(aircraft.ts,aircraft.Jacp,RHSJacobianP,&aircraft));
287   PetscCall(TSSetExactFinalTime(aircraft.ts,TS_EXACTFINALTIME_MATCHSTEP));
288   PetscCall(TSSetEquationType(aircraft.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
289 
290   /* Set initial conditions */
291   PetscCall(MatCreateVecs(aircraft.A,&aircraft.U,NULL));
292   PetscCall(TSSetSolution(aircraft.ts,aircraft.U));
293   PetscCall(VecGetArray(aircraft.U,&u));
294   u[0] = 1.5;
295   u[1] = 0;
296   PetscCall(VecRestoreArray(aircraft.U,&u));
297   PetscCall(VecCreate(PETSC_COMM_WORLD,&aircraft.V));
298   PetscCall(VecSetSizes(aircraft.V,PETSC_DECIDE,aircraft.nsteps));
299   PetscCall(VecSetUp(aircraft.V));
300   PetscCall(VecDuplicate(aircraft.V,&aircraft.W));
301   PetscCall(VecSet(aircraft.V,1.));
302   PetscCall(VecSet(aircraft.W,PETSC_PI/4.));
303 
304   /* Save trajectory of solution so that TSAdjointSolve() may be used */
305   PetscCall(TSSetSaveTrajectory(aircraft.ts));
306 
307   /* Set sensitivity context */
308   PetscCall(TSCreateQuadratureTS(aircraft.ts,PETSC_FALSE,&aircraft.quadts));
309   PetscCall(TSSetRHSFunction(aircraft.quadts,NULL,(TSRHSFunction)CostIntegrand,&aircraft));
310   PetscCall(TSSetRHSJacobian(aircraft.quadts,aircraft.DRDU,aircraft.DRDU,(TSRHSJacobian)DRDUJacobianTranspose,&aircraft));
311   PetscCall(TSSetRHSJacobianP(aircraft.quadts,aircraft.DRDP,(TSRHSJacobianP)DRDPJacobianTranspose,&aircraft));
312   PetscCall(MatCreateVecs(aircraft.A,&aircraft.Lambda[0],NULL));
313   PetscCall(MatCreateVecs(aircraft.Jacp,&aircraft.Mup[0],NULL));
314   if (aircraft.eh) {
315     PetscCall(MatCreateVecs(aircraft.A,&aircraft.rhshp1[0],NULL));
316     PetscCall(MatCreateVecs(aircraft.A,&aircraft.rhshp2[0],NULL));
317     PetscCall(MatCreateVecs(aircraft.Jacp,&aircraft.rhshp3[0],NULL));
318     PetscCall(MatCreateVecs(aircraft.Jacp,&aircraft.rhshp4[0],NULL));
319     PetscCall(MatCreateVecs(aircraft.DRDU,&aircraft.inthp1[0],NULL));
320     PetscCall(MatCreateVecs(aircraft.DRDU,&aircraft.inthp2[0],NULL));
321     PetscCall(MatCreateVecs(aircraft.DRDP,&aircraft.inthp3[0],NULL));
322     PetscCall(MatCreateVecs(aircraft.DRDP,&aircraft.inthp4[0],NULL));
323     PetscCall(MatCreateVecs(aircraft.Jacp,&aircraft.Dir,NULL));
324     PetscCall(TSSetRHSHessianProduct(aircraft.ts,aircraft.rhshp1,RHSHessianProductUU,aircraft.rhshp2,RHSHessianProductUP,aircraft.rhshp3,RHSHessianProductPU,aircraft.rhshp4,RHSHessianProductPP,&aircraft));
325     PetscCall(TSSetRHSHessianProduct(aircraft.quadts,aircraft.inthp1,IntegrandHessianProductUU,aircraft.inthp2,IntegrandHessianProductUP,aircraft.inthp3,IntegrandHessianProductPU,aircraft.inthp4,IntegrandHessianProductPP,&aircraft));
326     PetscCall(MatCreateVecs(aircraft.A,&aircraft.Lambda2[0],NULL));
327     PetscCall(MatCreateVecs(aircraft.Jacp,&aircraft.Mup2[0],NULL));
328   }
329   PetscCall(TSSetFromOptions(aircraft.ts));
330   PetscCall(TSSetMaxTime(aircraft.ts,aircraft.ftime));
331   PetscCall(TSSetTimeStep(aircraft.ts,aircraft.ftime/aircraft.nsteps));
332 
333   /* Set initial solution guess */
334   PetscCall(MatCreateVecs(aircraft.Jacp,&P,NULL));
335   PetscCall(VecGetArray(P,&p));
336   for (i=0; i<aircraft.nsteps; i++) {
337     p[2*i] = 2.0;
338     p[2*i+1] = PETSC_PI/2.0;
339   }
340   PetscCall(VecRestoreArray(P,&p));
341   PetscCall(VecDuplicate(P,&PU));
342   PetscCall(VecDuplicate(P,&PL));
343   PetscCall(VecGetArray(PU,&p));
344   for (i=0; i<aircraft.nsteps; i++) {
345     p[2*i] = 2.0;
346     p[2*i+1] = PETSC_PI;
347   }
348   PetscCall(VecRestoreArray(PU,&p));
349   PetscCall(VecGetArray(PL,&p));
350   for (i=0; i<aircraft.nsteps; i++) {
351     p[2*i] = 0.0;
352     p[2*i+1] = -PETSC_PI;
353   }
354   PetscCall(VecRestoreArray(PL,&p));
355 
356   PetscCall(TaoSetSolution(tao,P));
357   PetscCall(TaoSetVariableBounds(tao,PL,PU));
358   /* Set routine for function and gradient evaluation */
359   PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormObjFunctionGradient,(void *)&aircraft));
360 
361   if (aircraft.eh) {
362     if (aircraft.mf) {
363       PetscCall(MatCreateShell(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2*aircraft.nsteps,2*aircraft.nsteps,(void*)&aircraft,&aircraft.H));
364       PetscCall(MatShellSetOperation(aircraft.H,MATOP_MULT,(void(*)(void))MyMatMult));
365       PetscCall(MatSetOption(aircraft.H,MAT_SYMMETRIC,PETSC_TRUE));
366       PetscCall(TaoSetHessian(tao,aircraft.H,aircraft.H,MatrixFreeObjHessian,(void*)&aircraft));
367     } else {
368       PetscCall(MatCreateDense(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE,2*aircraft.nsteps,2*aircraft.nsteps,NULL,&(aircraft.H)));
369       PetscCall(MatSetOption(aircraft.H,MAT_SYMMETRIC,PETSC_TRUE));
370       PetscCall(TaoSetHessian(tao,aircraft.H,aircraft.H,FormObjHessian,(void *)&aircraft));
371     }
372   }
373 
374   /* Check for any TAO command line options */
375   PetscCall(TaoGetKSP(tao,&ksp));
376   if (ksp) {
377     PetscCall(KSPGetPC(ksp,&pc));
378     PetscCall(PCSetType(pc,PCNONE));
379   }
380   PetscCall(TaoSetFromOptions(tao));
381 
382   PetscCall(TaoSolve(tao));
383   PetscCall(VecView(P,PETSC_VIEWER_STDOUT_WORLD));
384 
385   /* Free TAO data structures */
386   PetscCall(TaoDestroy(&tao));
387 
388   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
389      Free work space.  All PETSc objects should be destroyed when they
390      are no longer needed.
391    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
392   PetscCall(TSDestroy(&aircraft.ts));
393   PetscCall(MatDestroy(&aircraft.A));
394   PetscCall(VecDestroy(&aircraft.U));
395   PetscCall(VecDestroy(&aircraft.V));
396   PetscCall(VecDestroy(&aircraft.W));
397   PetscCall(VecDestroy(&P));
398   PetscCall(VecDestroy(&PU));
399   PetscCall(VecDestroy(&PL));
400   PetscCall(MatDestroy(&aircraft.Jacp));
401   PetscCall(MatDestroy(&aircraft.DRDU));
402   PetscCall(MatDestroy(&aircraft.DRDP));
403   PetscCall(VecDestroy(&aircraft.Lambda[0]));
404   PetscCall(VecDestroy(&aircraft.Mup[0]));
405   PetscCall(VecDestroy(&P));
406   if (aircraft.eh) {
407     PetscCall(VecDestroy(&aircraft.Lambda2[0]));
408     PetscCall(VecDestroy(&aircraft.Mup2[0]));
409     PetscCall(VecDestroy(&aircraft.Dir));
410     PetscCall(VecDestroy(&aircraft.rhshp1[0]));
411     PetscCall(VecDestroy(&aircraft.rhshp2[0]));
412     PetscCall(VecDestroy(&aircraft.rhshp3[0]));
413     PetscCall(VecDestroy(&aircraft.rhshp4[0]));
414     PetscCall(VecDestroy(&aircraft.inthp1[0]));
415     PetscCall(VecDestroy(&aircraft.inthp2[0]));
416     PetscCall(VecDestroy(&aircraft.inthp3[0]));
417     PetscCall(VecDestroy(&aircraft.inthp4[0]));
418     PetscCall(MatDestroy(&aircraft.H));
419   }
420   PetscCall(PetscFinalize());
421   return 0;
422 }
423 
424 /*
425    FormObjFunctionGradient - Evaluates the function and corresponding gradient.
426 
427    Input Parameters:
428    tao - the Tao context
429    P   - the input vector
430    ctx - optional aircraft-defined context, as set by TaoSetObjectiveAndGradient()
431 
432    Output Parameters:
433    f   - the newly evaluated function
434    G   - the newly evaluated gradient
435 */
436 PetscErrorCode FormObjFunctionGradient(Tao tao,Vec P,PetscReal *f,Vec G,void *ctx)
437 {
438   Aircraft          actx = (Aircraft)ctx;
439   TS                ts = actx->ts;
440   Vec               Q;
441   const PetscScalar *p,*q;
442   PetscScalar       *u,*v,*w;
443   PetscInt          i;
444 
445   PetscFunctionBeginUser;
446   PetscCall(VecGetArrayRead(P,&p));
447   PetscCall(VecGetArray(actx->V,&v));
448   PetscCall(VecGetArray(actx->W,&w));
449   for (i=0; i<actx->nsteps; i++) {
450     v[i] = p[2*i];
451     w[i] = p[2*i+1];
452   }
453   PetscCall(VecRestoreArrayRead(P,&p));
454   PetscCall(VecRestoreArray(actx->V,&v));
455   PetscCall(VecRestoreArray(actx->W,&w));
456 
457   PetscCall(TSSetTime(ts,0.0));
458   PetscCall(TSSetStepNumber(ts,0));
459   PetscCall(TSSetFromOptions(ts));
460   PetscCall(TSSetTimeStep(ts,actx->ftime/actx->nsteps));
461 
462   /* reinitialize system state */
463   PetscCall(VecGetArray(actx->U,&u));
464   u[0] = 2.0;
465   u[1] = 0;
466   PetscCall(VecRestoreArray(actx->U,&u));
467 
468   /* reinitialize the integral value */
469   PetscCall(TSGetCostIntegral(ts,&Q));
470   PetscCall(VecSet(Q,0.0));
471 
472   PetscCall(TSSolve(ts,actx->U));
473 
474   /* Reset initial conditions for the adjoint integration */
475   PetscCall(VecSet(actx->Lambda[0],0.0));
476   PetscCall(VecSet(actx->Mup[0],0.0));
477   PetscCall(TSSetCostGradients(ts,1,actx->Lambda,actx->Mup));
478 
479   PetscCall(TSAdjointSolve(ts));
480   PetscCall(VecCopy(actx->Mup[0],G));
481   PetscCall(TSGetCostIntegral(ts,&Q));
482   PetscCall(VecGetArrayRead(Q,&q));
483   *f   = q[0];
484   PetscCall(VecRestoreArrayRead(Q,&q));
485   PetscFunctionReturn(0);
486 }
487 
488 PetscErrorCode FormObjHessian(Tao tao,Vec P,Mat H,Mat Hpre,void *ctx)
489 {
490   Aircraft          actx = (Aircraft)ctx;
491   const PetscScalar *p;
492   PetscScalar       *harr,*v,*w,one = 1.0;
493   PetscInt          ind[1];
494   PetscInt          *cols,i;
495   Vec               Dir;
496 
497   PetscFunctionBeginUser;
498   /* set up control parameters */
499   PetscCall(VecGetArrayRead(P,&p));
500   PetscCall(VecGetArray(actx->V,&v));
501   PetscCall(VecGetArray(actx->W,&w));
502   for (i=0; i<actx->nsteps; i++) {
503     v[i] = p[2*i];
504     w[i] = p[2*i+1];
505   }
506   PetscCall(VecRestoreArrayRead(P,&p));
507   PetscCall(VecRestoreArray(actx->V,&v));
508   PetscCall(VecRestoreArray(actx->W,&w));
509 
510   PetscCall(PetscMalloc1(2*actx->nsteps,&harr));
511   PetscCall(PetscMalloc1(2*actx->nsteps,&cols));
512   for (i=0; i<2*actx->nsteps; i++) cols[i] = i;
513   PetscCall(VecDuplicate(P,&Dir));
514   for (i=0; i<2*actx->nsteps; i++) {
515     ind[0] = i;
516     PetscCall(VecSet(Dir,0.0));
517     PetscCall(VecSetValues(Dir,1,ind,&one,INSERT_VALUES));
518     PetscCall(VecAssemblyBegin(Dir));
519     PetscCall(VecAssemblyEnd(Dir));
520     PetscCall(ComputeObjHessianWithSOA(Dir,harr,actx));
521     PetscCall(MatSetValues(H,1,ind,2*actx->nsteps,cols,harr,INSERT_VALUES));
522     PetscCall(MatAssemblyBegin(H,MAT_FINAL_ASSEMBLY));
523     PetscCall(MatAssemblyEnd(H,MAT_FINAL_ASSEMBLY));
524     if (H != Hpre) {
525       PetscCall(MatAssemblyBegin(Hpre,MAT_FINAL_ASSEMBLY));
526       PetscCall(MatAssemblyEnd(Hpre,MAT_FINAL_ASSEMBLY));
527     }
528   }
529   PetscCall(PetscFree(cols));
530   PetscCall(PetscFree(harr));
531   PetscCall(VecDestroy(&Dir));
532   PetscFunctionReturn(0);
533 }
534 
535 PetscErrorCode MatrixFreeObjHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx)
536 {
537   Aircraft          actx = (Aircraft)ctx;
538   PetscScalar       *v,*w;
539   const PetscScalar *p;
540   PetscInt          i;
541 
542   PetscFunctionBegin;
543   PetscCall(VecGetArrayRead(P,&p));
544   PetscCall(VecGetArray(actx->V,&v));
545   PetscCall(VecGetArray(actx->W,&w));
546   for (i=0; i<actx->nsteps; i++) {
547     v[i] = p[2*i];
548     w[i] = p[2*i+1];
549   }
550   PetscCall(VecRestoreArrayRead(P,&p));
551   PetscCall(VecRestoreArray(actx->V,&v));
552   PetscCall(VecRestoreArray(actx->W,&w));
553   PetscFunctionReturn(0);
554 }
555 
556 PetscErrorCode MyMatMult(Mat H_shell, Vec X, Vec Y)
557 {
558   PetscScalar    *y;
559   void           *ptr;
560 
561   PetscFunctionBegin;
562   PetscCall(MatShellGetContext(H_shell,&ptr));
563   PetscCall(VecGetArray(Y,&y));
564   PetscCall(ComputeObjHessianWithSOA(X,y,(Aircraft)ptr));
565   PetscCall(VecRestoreArray(Y,&y));
566   PetscFunctionReturn(0);
567 }
568 
569 PetscErrorCode ComputeObjHessianWithSOA(Vec Dir,PetscScalar arr[],Aircraft actx)
570 {
571   TS                ts = actx->ts;
572   const PetscScalar *z_ptr;
573   PetscScalar       *u;
574   Vec               Q;
575   PetscInt          i;
576 
577   PetscFunctionBeginUser;
578   /* Reset TSAdjoint so that AdjointSetUp will be called again */
579   PetscCall(TSAdjointReset(ts));
580 
581   PetscCall(TSSetTime(ts,0.0));
582   PetscCall(TSSetStepNumber(ts,0));
583   PetscCall(TSSetFromOptions(ts));
584   PetscCall(TSSetTimeStep(ts,actx->ftime/actx->nsteps));
585   PetscCall(TSSetCostHessianProducts(actx->ts,1,actx->Lambda2,actx->Mup2,Dir));
586 
587   /* reinitialize system state */
588   PetscCall(VecGetArray(actx->U,&u));
589   u[0] = 2.0;
590   u[1] = 0;
591   PetscCall(VecRestoreArray(actx->U,&u));
592 
593   /* reinitialize the integral value */
594   PetscCall(TSGetCostIntegral(ts,&Q));
595   PetscCall(VecSet(Q,0.0));
596 
597   /* initialize tlm variable */
598   PetscCall(MatZeroEntries(actx->Jacp));
599   PetscCall(MatAssemblyBegin(actx->Jacp,MAT_FINAL_ASSEMBLY));
600   PetscCall(MatAssemblyEnd(actx->Jacp,MAT_FINAL_ASSEMBLY));
601   PetscCall(TSAdjointSetForward(ts,actx->Jacp));
602 
603   PetscCall(TSSolve(ts,actx->U));
604 
605   /* Set terminal conditions for first- and second-order adjonts */
606   PetscCall(VecSet(actx->Lambda[0],0.0));
607   PetscCall(VecSet(actx->Mup[0],0.0));
608   PetscCall(VecSet(actx->Lambda2[0],0.0));
609   PetscCall(VecSet(actx->Mup2[0],0.0));
610   PetscCall(TSSetCostGradients(ts,1,actx->Lambda,actx->Mup));
611 
612   PetscCall(TSGetCostIntegral(ts,&Q));
613 
614   /* Reset initial conditions for the adjoint integration */
615   PetscCall(TSAdjointSolve(ts));
616 
617   /* initial condition does not depend on p, so that lambda is not needed to assemble G */
618   PetscCall(VecGetArrayRead(actx->Mup2[0],&z_ptr));
619   for (i=0; i<2*actx->nsteps; i++) arr[i] = z_ptr[i];
620   PetscCall(VecRestoreArrayRead(actx->Mup2[0],&z_ptr));
621 
622   /* Disable second-order adjoint mode */
623   PetscCall(TSAdjointReset(ts));
624   PetscCall(TSAdjointResetForward(ts));
625   PetscFunctionReturn(0);
626 }
627 
628 /*TEST
629     build:
630       requires: !complex !single
631 
632     test:
633       args:  -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_gatol 1e-7
634 
635     test:
636       suffix: 2
637       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
638 
639     test:
640       suffix: 3
641       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
642 TEST*/
643