xref: /petsc/src/ts/tutorials/optimal_control/ex1.c (revision 40badf4fbc550ac1f60bd080eaff6de6d55b946d)
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   CHKERRQ(TSGetStepNumber(ts,&step));
33   CHKERRQ(VecGetArrayRead(U,&u));
34   CHKERRQ(VecGetArrayRead(actx->V,&v));
35   CHKERRQ(VecGetArrayRead(actx->W,&w));
36   CHKERRQ(VecGetArray(F,&f));
37   f[0] = v[step]*PetscCosReal(w[step]);
38   f[1] = v[step]*PetscSinReal(w[step]);
39   CHKERRQ(VecRestoreArrayRead(U,&u));
40   CHKERRQ(VecRestoreArrayRead(actx->V,&v));
41   CHKERRQ(VecRestoreArrayRead(actx->W,&w));
42   CHKERRQ(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   CHKERRQ(MatZeroEntries(A));
55   CHKERRQ(TSGetStepNumber(ts,&step));
56   CHKERRQ(VecGetArrayRead(U,&u));
57   CHKERRQ(VecGetArrayRead(actx->V,&v));
58   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(U,&u));
66   CHKERRQ(VecRestoreArrayRead(actx->V,&v));
67   CHKERRQ(VecRestoreArrayRead(actx->W,&w));
68 
69   rowcol[0] = 2*step;
70   rowcol[1] = 2*step+1;
71   CHKERRQ(MatSetValues(A,2,rows,2,rowcol,&Jp[0][0],INSERT_VALUES));
72 
73   CHKERRQ(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
74   CHKERRQ(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   CHKERRQ(TSGetStepNumber(ts,&step));
106   CHKERRQ(VecGetArrayRead(U,&u));
107   CHKERRQ(VecGetArrayRead(actx->V,&v));
108   CHKERRQ(VecGetArrayRead(actx->W,&w));
109   CHKERRQ(VecGetArrayRead(Vl[0],&vl));
110   CHKERRQ(VecGetArrayRead(Vr,&vr));
111   CHKERRQ(VecSet(VHV[0],0.0));
112   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(U,&u));
128   CHKERRQ(VecRestoreArrayRead(Vl[0],&vl));
129   CHKERRQ(VecRestoreArrayRead(Vr,&vr));
130   CHKERRQ(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   CHKERRQ(TSGetStepNumber(ts,&step));
145   CHKERRQ(VecGetArrayRead(U,&u));
146   CHKERRQ(VecGetArrayRead(actx->V,&v));
147   CHKERRQ(VecGetArrayRead(actx->W,&w));
148   CHKERRQ(VecGetArrayRead(Vr,&vr));
149   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(U,&u));
160   CHKERRQ(VecRestoreArrayRead(Vr,&vr));
161   CHKERRQ(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   CHKERRQ(VecGetArrayRead(U,&u));
192   CHKERRQ(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   CHKERRQ(VecRestoreArray(R,&r));
197   CHKERRQ(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   CHKERRQ(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   CHKERRQ(MatSetValues(DRDU,2,row,1,col,&drdu[0][0],INSERT_VALUES));
216   CHKERRQ(VecRestoreArrayRead(U,&u));
217   CHKERRQ(MatAssemblyBegin(DRDU,MAT_FINAL_ASSEMBLY));
218   CHKERRQ(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   CHKERRQ(MatZeroEntries(DRDP));
226   CHKERRQ(MatAssemblyBegin(DRDP,MAT_FINAL_ASSEMBLY));
227   CHKERRQ(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   PetscErrorCode     ierr;
242 
243   /* Initialize program */
244   ierr = PetscInitialize(&argc,&argv,NULL,NULL);if (ierr) return ierr;
245   CHKERRMPI(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   CHKERRQ(PetscOptionsGetReal(NULL,NULL,"-ftime",&aircraft.ftime,NULL));
255   CHKERRQ(PetscOptionsGetInt(NULL,NULL,"-nsteps",&aircraft.nsteps,NULL));
256   CHKERRQ(PetscOptionsHasName(NULL,NULL,"-matrixfree",&aircraft.mf));
257   CHKERRQ(PetscOptionsHasName(NULL,NULL,"-exacthessian",&aircraft.eh));
258 
259   /* Create TAO solver and set desired solution method */
260   CHKERRQ(TaoCreate(PETSC_COMM_WORLD,&tao));
261   CHKERRQ(TaoSetType(tao,TAOBQNLS));
262 
263   /* Create necessary matrix and vectors, solve same ODE on every process */
264   CHKERRQ(MatCreate(PETSC_COMM_WORLD,&aircraft.A));
265   CHKERRQ(MatSetSizes(aircraft.A,PETSC_DECIDE,PETSC_DECIDE,2,2));
266   CHKERRQ(MatSetFromOptions(aircraft.A));
267   CHKERRQ(MatSetUp(aircraft.A));
268   CHKERRQ(MatAssemblyBegin(aircraft.A,MAT_FINAL_ASSEMBLY));
269   CHKERRQ(MatAssemblyEnd(aircraft.A,MAT_FINAL_ASSEMBLY));
270   CHKERRQ(MatShift(aircraft.A,1));
271   CHKERRQ(MatShift(aircraft.A,-1));
272 
273   CHKERRQ(MatCreate(PETSC_COMM_WORLD,&aircraft.Jacp));
274   CHKERRQ(MatSetSizes(aircraft.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2*aircraft.nsteps));
275   CHKERRQ(MatSetFromOptions(aircraft.Jacp));
276   CHKERRQ(MatSetUp(aircraft.Jacp));
277   CHKERRQ(MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2*aircraft.nsteps,1,NULL,&aircraft.DRDP));
278   CHKERRQ(MatSetUp(aircraft.DRDP));
279   CHKERRQ(MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2,1,NULL,&aircraft.DRDU));
280   CHKERRQ(MatSetUp(aircraft.DRDU));
281 
282   /* Create timestepping solver context */
283   CHKERRQ(TSCreate(PETSC_COMM_WORLD,&aircraft.ts));
284   CHKERRQ(TSSetType(aircraft.ts,TSRK));
285   CHKERRQ(TSSetRHSFunction(aircraft.ts,NULL,RHSFunction,&aircraft));
286   CHKERRQ(TSSetRHSJacobian(aircraft.ts,aircraft.A,aircraft.A,TSComputeRHSJacobianConstant,&aircraft));
287   CHKERRQ(TSSetRHSJacobianP(aircraft.ts,aircraft.Jacp,RHSJacobianP,&aircraft));
288   CHKERRQ(TSSetExactFinalTime(aircraft.ts,TS_EXACTFINALTIME_MATCHSTEP));
289   CHKERRQ(TSSetEquationType(aircraft.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
290 
291   /* Set initial conditions */
292   CHKERRQ(MatCreateVecs(aircraft.A,&aircraft.U,NULL));
293   CHKERRQ(TSSetSolution(aircraft.ts,aircraft.U));
294   CHKERRQ(VecGetArray(aircraft.U,&u));
295   u[0] = 1.5;
296   u[1] = 0;
297   CHKERRQ(VecRestoreArray(aircraft.U,&u));
298   CHKERRQ(VecCreate(PETSC_COMM_WORLD,&aircraft.V));
299   CHKERRQ(VecSetSizes(aircraft.V,PETSC_DECIDE,aircraft.nsteps));
300   CHKERRQ(VecSetUp(aircraft.V));
301   CHKERRQ(VecDuplicate(aircraft.V,&aircraft.W));
302   CHKERRQ(VecSet(aircraft.V,1.));
303   CHKERRQ(VecSet(aircraft.W,PETSC_PI/4.));
304 
305   /* Save trajectory of solution so that TSAdjointSolve() may be used */
306   CHKERRQ(TSSetSaveTrajectory(aircraft.ts));
307 
308   /* Set sensitivity context */
309   CHKERRQ(TSCreateQuadratureTS(aircraft.ts,PETSC_FALSE,&aircraft.quadts));
310   CHKERRQ(TSSetRHSFunction(aircraft.quadts,NULL,(TSRHSFunction)CostIntegrand,&aircraft));
311   CHKERRQ(TSSetRHSJacobian(aircraft.quadts,aircraft.DRDU,aircraft.DRDU,(TSRHSJacobian)DRDUJacobianTranspose,&aircraft));
312   CHKERRQ(TSSetRHSJacobianP(aircraft.quadts,aircraft.DRDP,(TSRHSJacobianP)DRDPJacobianTranspose,&aircraft));
313   CHKERRQ(MatCreateVecs(aircraft.A,&aircraft.Lambda[0],NULL));
314   CHKERRQ(MatCreateVecs(aircraft.Jacp,&aircraft.Mup[0],NULL));
315   if (aircraft.eh) {
316     CHKERRQ(MatCreateVecs(aircraft.A,&aircraft.rhshp1[0],NULL));
317     CHKERRQ(MatCreateVecs(aircraft.A,&aircraft.rhshp2[0],NULL));
318     CHKERRQ(MatCreateVecs(aircraft.Jacp,&aircraft.rhshp3[0],NULL));
319     CHKERRQ(MatCreateVecs(aircraft.Jacp,&aircraft.rhshp4[0],NULL));
320     CHKERRQ(MatCreateVecs(aircraft.DRDU,&aircraft.inthp1[0],NULL));
321     CHKERRQ(MatCreateVecs(aircraft.DRDU,&aircraft.inthp2[0],NULL));
322     CHKERRQ(MatCreateVecs(aircraft.DRDP,&aircraft.inthp3[0],NULL));
323     CHKERRQ(MatCreateVecs(aircraft.DRDP,&aircraft.inthp4[0],NULL));
324     CHKERRQ(MatCreateVecs(aircraft.Jacp,&aircraft.Dir,NULL));
325     CHKERRQ(TSSetRHSHessianProduct(aircraft.ts,aircraft.rhshp1,RHSHessianProductUU,aircraft.rhshp2,RHSHessianProductUP,aircraft.rhshp3,RHSHessianProductPU,aircraft.rhshp4,RHSHessianProductPP,&aircraft));
326     CHKERRQ(TSSetRHSHessianProduct(aircraft.quadts,aircraft.inthp1,IntegrandHessianProductUU,aircraft.inthp2,IntegrandHessianProductUP,aircraft.inthp3,IntegrandHessianProductPU,aircraft.inthp4,IntegrandHessianProductPP,&aircraft));
327     CHKERRQ(MatCreateVecs(aircraft.A,&aircraft.Lambda2[0],NULL));
328     CHKERRQ(MatCreateVecs(aircraft.Jacp,&aircraft.Mup2[0],NULL));
329   }
330   CHKERRQ(TSSetFromOptions(aircraft.ts));
331   CHKERRQ(TSSetMaxTime(aircraft.ts,aircraft.ftime));
332   CHKERRQ(TSSetTimeStep(aircraft.ts,aircraft.ftime/aircraft.nsteps));
333 
334   /* Set initial solution guess */
335   CHKERRQ(MatCreateVecs(aircraft.Jacp,&P,NULL));
336   CHKERRQ(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   CHKERRQ(VecRestoreArray(P,&p));
342   CHKERRQ(VecDuplicate(P,&PU));
343   CHKERRQ(VecDuplicate(P,&PL));
344   CHKERRQ(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   CHKERRQ(VecRestoreArray(PU,&p));
350   CHKERRQ(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   CHKERRQ(VecRestoreArray(PL,&p));
356 
357   CHKERRQ(TaoSetSolution(tao,P));
358   CHKERRQ(TaoSetVariableBounds(tao,PL,PU));
359   /* Set routine for function and gradient evaluation */
360   CHKERRQ(TaoSetObjectiveAndGradient(tao,NULL,FormObjFunctionGradient,(void *)&aircraft));
361 
362   if (aircraft.eh) {
363     if (aircraft.mf) {
364       CHKERRQ(MatCreateShell(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,2*aircraft.nsteps,2*aircraft.nsteps,(void*)&aircraft,&aircraft.H));
365       CHKERRQ(MatShellSetOperation(aircraft.H,MATOP_MULT,(void(*)(void))MyMatMult));
366       CHKERRQ(MatSetOption(aircraft.H,MAT_SYMMETRIC,PETSC_TRUE));
367       CHKERRQ(TaoSetHessian(tao,aircraft.H,aircraft.H,MatrixFreeObjHessian,(void*)&aircraft));
368     } else {
369       CHKERRQ(MatCreateDense(MPI_COMM_WORLD,PETSC_DETERMINE,PETSC_DETERMINE,2*aircraft.nsteps,2*aircraft.nsteps,NULL,&(aircraft.H)));
370       CHKERRQ(MatSetOption(aircraft.H,MAT_SYMMETRIC,PETSC_TRUE));
371       CHKERRQ(TaoSetHessian(tao,aircraft.H,aircraft.H,FormObjHessian,(void *)&aircraft));
372     }
373   }
374 
375   /* Check for any TAO command line options */
376   CHKERRQ(TaoGetKSP(tao,&ksp));
377   if (ksp) {
378     CHKERRQ(KSPGetPC(ksp,&pc));
379     CHKERRQ(PCSetType(pc,PCNONE));
380   }
381   CHKERRQ(TaoSetFromOptions(tao));
382 
383   CHKERRQ(TaoSolve(tao));
384   CHKERRQ(VecView(P,PETSC_VIEWER_STDOUT_WORLD));
385 
386   /* Free TAO data structures */
387   CHKERRQ(TaoDestroy(&tao));
388 
389   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
390      Free work space.  All PETSc objects should be destroyed when they
391      are no longer needed.
392    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
393   CHKERRQ(TSDestroy(&aircraft.ts));
394   CHKERRQ(MatDestroy(&aircraft.A));
395   CHKERRQ(VecDestroy(&aircraft.U));
396   CHKERRQ(VecDestroy(&aircraft.V));
397   CHKERRQ(VecDestroy(&aircraft.W));
398   CHKERRQ(VecDestroy(&P));
399   CHKERRQ(VecDestroy(&PU));
400   CHKERRQ(VecDestroy(&PL));
401   CHKERRQ(MatDestroy(&aircraft.Jacp));
402   CHKERRQ(MatDestroy(&aircraft.DRDU));
403   CHKERRQ(MatDestroy(&aircraft.DRDP));
404   CHKERRQ(VecDestroy(&aircraft.Lambda[0]));
405   CHKERRQ(VecDestroy(&aircraft.Mup[0]));
406   CHKERRQ(VecDestroy(&P));
407   if (aircraft.eh) {
408     CHKERRQ(VecDestroy(&aircraft.Lambda2[0]));
409     CHKERRQ(VecDestroy(&aircraft.Mup2[0]));
410     CHKERRQ(VecDestroy(&aircraft.Dir));
411     CHKERRQ(VecDestroy(&aircraft.rhshp1[0]));
412     CHKERRQ(VecDestroy(&aircraft.rhshp2[0]));
413     CHKERRQ(VecDestroy(&aircraft.rhshp3[0]));
414     CHKERRQ(VecDestroy(&aircraft.rhshp4[0]));
415     CHKERRQ(VecDestroy(&aircraft.inthp1[0]));
416     CHKERRQ(VecDestroy(&aircraft.inthp2[0]));
417     CHKERRQ(VecDestroy(&aircraft.inthp3[0]));
418     CHKERRQ(VecDestroy(&aircraft.inthp4[0]));
419     CHKERRQ(MatDestroy(&aircraft.H));
420   }
421   ierr = PetscFinalize();
422   return ierr;
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   CHKERRQ(VecGetArrayRead(P,&p));
448   CHKERRQ(VecGetArray(actx->V,&v));
449   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(P,&p));
455   CHKERRQ(VecRestoreArray(actx->V,&v));
456   CHKERRQ(VecRestoreArray(actx->W,&w));
457 
458   CHKERRQ(TSSetTime(ts,0.0));
459   CHKERRQ(TSSetStepNumber(ts,0));
460   CHKERRQ(TSSetFromOptions(ts));
461   CHKERRQ(TSSetTimeStep(ts,actx->ftime/actx->nsteps));
462 
463   /* reinitialize system state */
464   CHKERRQ(VecGetArray(actx->U,&u));
465   u[0] = 2.0;
466   u[1] = 0;
467   CHKERRQ(VecRestoreArray(actx->U,&u));
468 
469   /* reinitialize the integral value */
470   CHKERRQ(TSGetCostIntegral(ts,&Q));
471   CHKERRQ(VecSet(Q,0.0));
472 
473   CHKERRQ(TSSolve(ts,actx->U));
474 
475   /* Reset initial conditions for the adjoint integration */
476   CHKERRQ(VecSet(actx->Lambda[0],0.0));
477   CHKERRQ(VecSet(actx->Mup[0],0.0));
478   CHKERRQ(TSSetCostGradients(ts,1,actx->Lambda,actx->Mup));
479 
480   CHKERRQ(TSAdjointSolve(ts));
481   CHKERRQ(VecCopy(actx->Mup[0],G));
482   CHKERRQ(TSGetCostIntegral(ts,&Q));
483   CHKERRQ(VecGetArrayRead(Q,&q));
484   *f   = q[0];
485   CHKERRQ(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   CHKERRQ(VecGetArrayRead(P,&p));
501   CHKERRQ(VecGetArray(actx->V,&v));
502   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(P,&p));
508   CHKERRQ(VecRestoreArray(actx->V,&v));
509   CHKERRQ(VecRestoreArray(actx->W,&w));
510 
511   CHKERRQ(PetscMalloc1(2*actx->nsteps,&harr));
512   CHKERRQ(PetscMalloc1(2*actx->nsteps,&cols));
513   for (i=0; i<2*actx->nsteps; i++) cols[i] = i;
514   CHKERRQ(VecDuplicate(P,&Dir));
515   for (i=0; i<2*actx->nsteps; i++) {
516     ind[0] = i;
517     CHKERRQ(VecSet(Dir,0.0));
518     CHKERRQ(VecSetValues(Dir,1,ind,&one,INSERT_VALUES));
519     CHKERRQ(VecAssemblyBegin(Dir));
520     CHKERRQ(VecAssemblyEnd(Dir));
521     CHKERRQ(ComputeObjHessianWithSOA(Dir,harr,actx));
522     CHKERRQ(MatSetValues(H,1,ind,2*actx->nsteps,cols,harr,INSERT_VALUES));
523     CHKERRQ(MatAssemblyBegin(H,MAT_FINAL_ASSEMBLY));
524     CHKERRQ(MatAssemblyEnd(H,MAT_FINAL_ASSEMBLY));
525     if (H != Hpre) {
526       CHKERRQ(MatAssemblyBegin(Hpre,MAT_FINAL_ASSEMBLY));
527       CHKERRQ(MatAssemblyEnd(Hpre,MAT_FINAL_ASSEMBLY));
528     }
529   }
530   CHKERRQ(PetscFree(cols));
531   CHKERRQ(PetscFree(harr));
532   CHKERRQ(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   CHKERRQ(VecGetArrayRead(P,&p));
545   CHKERRQ(VecGetArray(actx->V,&v));
546   CHKERRQ(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   CHKERRQ(VecRestoreArrayRead(P,&p));
552   CHKERRQ(VecRestoreArray(actx->V,&v));
553   CHKERRQ(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   CHKERRQ(MatShellGetContext(H_shell,&ptr));
564   CHKERRQ(VecGetArray(Y,&y));
565   CHKERRQ(ComputeObjHessianWithSOA(X,y,(Aircraft)ptr));
566   CHKERRQ(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   CHKERRQ(TSAdjointReset(ts));
581 
582   CHKERRQ(TSSetTime(ts,0.0));
583   CHKERRQ(TSSetStepNumber(ts,0));
584   CHKERRQ(TSSetFromOptions(ts));
585   CHKERRQ(TSSetTimeStep(ts,actx->ftime/actx->nsteps));
586   CHKERRQ(TSSetCostHessianProducts(actx->ts,1,actx->Lambda2,actx->Mup2,Dir));
587 
588   /* reinitialize system state */
589   CHKERRQ(VecGetArray(actx->U,&u));
590   u[0] = 2.0;
591   u[1] = 0;
592   CHKERRQ(VecRestoreArray(actx->U,&u));
593 
594   /* reinitialize the integral value */
595   CHKERRQ(TSGetCostIntegral(ts,&Q));
596   CHKERRQ(VecSet(Q,0.0));
597 
598   /* initialize tlm variable */
599   CHKERRQ(MatZeroEntries(actx->Jacp));
600   CHKERRQ(MatAssemblyBegin(actx->Jacp,MAT_FINAL_ASSEMBLY));
601   CHKERRQ(MatAssemblyEnd(actx->Jacp,MAT_FINAL_ASSEMBLY));
602   CHKERRQ(TSAdjointSetForward(ts,actx->Jacp));
603 
604   CHKERRQ(TSSolve(ts,actx->U));
605 
606   /* Set terminal conditions for first- and second-order adjonts */
607   CHKERRQ(VecSet(actx->Lambda[0],0.0));
608   CHKERRQ(VecSet(actx->Mup[0],0.0));
609   CHKERRQ(VecSet(actx->Lambda2[0],0.0));
610   CHKERRQ(VecSet(actx->Mup2[0],0.0));
611   CHKERRQ(TSSetCostGradients(ts,1,actx->Lambda,actx->Mup));
612 
613   CHKERRQ(TSGetCostIntegral(ts,&Q));
614 
615   /* Reset initial conditions for the adjoint integration */
616   CHKERRQ(TSAdjointSolve(ts));
617 
618   /* initial condition does not depend on p, so that lambda is not needed to assemble G */
619   CHKERRQ(VecGetArrayRead(actx->Mup2[0],&z_ptr));
620   for (i=0; i<2*actx->nsteps; i++) arr[i] = z_ptr[i];
621   CHKERRQ(VecRestoreArrayRead(actx->Mup2[0],&z_ptr));
622 
623   /* Disable second-order adjoint mode */
624   CHKERRQ(TSAdjointReset(ts));
625   CHKERRQ(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