xref: /petsc/src/ts/tutorials/ex20opt_ic.c (revision 327415f76d85372a4417cf1aaa14db707d4d6c04)
1c4762a1bSJed Brown static char help[] = "Solves a ODE-constrained optimization problem -- finding the optimal initial conditions for the van der Pol equation.\n";
2c4762a1bSJed Brown 
30e3d61c9SBarry Smith /*
4c4762a1bSJed Brown   Notes:
5c4762a1bSJed Brown   This code demonstrates how to solve an ODE-constrained optimization problem with TAO, TSAdjoint and TS.
6c4762a1bSJed Brown   The nonlinear problem is written in an ODE equivalent form.
7c4762a1bSJed Brown   The objective is to minimize the difference between observation and model prediction by finding optimal values for initial conditions.
8c4762a1bSJed Brown   The gradient is computed with the discrete adjoint of an implicit method or an explicit method, see ex20adj.c for details.
9c4762a1bSJed Brown */
10c4762a1bSJed Brown 
11c4762a1bSJed Brown #include <petsctao.h>
12c4762a1bSJed Brown #include <petscts.h>
13c4762a1bSJed Brown 
14c4762a1bSJed Brown typedef struct _n_User *User;
15c4762a1bSJed Brown struct _n_User {
16c4762a1bSJed Brown   TS        ts;
17c4762a1bSJed Brown   PetscReal mu;
18c4762a1bSJed Brown   PetscReal next_output;
19c4762a1bSJed Brown 
20c4762a1bSJed Brown   /* Sensitivity analysis support */
21c4762a1bSJed Brown   PetscInt  steps;
22c4762a1bSJed Brown   PetscReal ftime;
23c4762a1bSJed Brown   Mat       A;                       /* Jacobian matrix for ODE */
24c4762a1bSJed Brown   Mat       Jacp;                    /* JacobianP matrix for ODE*/
25c4762a1bSJed Brown   Mat       H;                       /* Hessian matrix for optimization */
26c4762a1bSJed Brown   Vec       U,Lambda[1],Mup[1];      /* first-order adjoint variables */
27c4762a1bSJed Brown   Vec       Lambda2[2];              /* second-order adjoint variables */
28c4762a1bSJed Brown   Vec       Ihp1[1];                 /* working space for Hessian evaluations */
29c4762a1bSJed Brown   Vec       Dir;                     /* direction vector */
30c4762a1bSJed Brown   PetscReal ob[2];                   /* observation used by the cost function */
31c4762a1bSJed Brown   PetscBool implicitform;            /* implicit ODE? */
32c4762a1bSJed Brown };
33c4762a1bSJed Brown PetscErrorCode Adjoint2(Vec,PetscScalar[],User);
34c4762a1bSJed Brown 
35c4762a1bSJed Brown /* ----------------------- Explicit form of the ODE  -------------------- */
36c4762a1bSJed Brown 
37c4762a1bSJed Brown static PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec U,Vec F,void *ctx)
38c4762a1bSJed Brown {
39c4762a1bSJed Brown   User              user = (User)ctx;
40c4762a1bSJed Brown   PetscScalar       *f;
41c4762a1bSJed Brown   const PetscScalar *u;
42c4762a1bSJed Brown 
43c4762a1bSJed Brown   PetscFunctionBeginUser;
449566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
459566063dSJacob Faibussowitsch   PetscCall(VecGetArray(F,&f));
46c4762a1bSJed Brown   f[0] = u[1];
47c4762a1bSJed Brown   f[1] = user->mu*((1.-u[0]*u[0])*u[1]-u[0]);
489566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
499566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(F,&f));
50c4762a1bSJed Brown   PetscFunctionReturn(0);
51c4762a1bSJed Brown }
52c4762a1bSJed Brown 
53c4762a1bSJed Brown static PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,void *ctx)
54c4762a1bSJed Brown {
55c4762a1bSJed Brown   User              user = (User)ctx;
56c4762a1bSJed Brown   PetscReal         mu   = user->mu;
57c4762a1bSJed Brown   PetscInt          rowcol[] = {0,1};
58c4762a1bSJed Brown   PetscScalar       J[2][2];
59c4762a1bSJed Brown   const PetscScalar *u;
60c4762a1bSJed Brown 
61c4762a1bSJed Brown   PetscFunctionBeginUser;
629566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
63c4762a1bSJed Brown   J[0][0] = 0;
64c4762a1bSJed Brown   J[1][0] = -mu*(2.0*u[1]*u[0]+1.);
65c4762a1bSJed Brown   J[0][1] = 1.0;
66c4762a1bSJed Brown   J[1][1] = mu*(1.0-u[0]*u[0]);
679566063dSJacob Faibussowitsch   PetscCall(MatSetValues(A,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES));
689566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
699566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
70c4762a1bSJed Brown   if (A != B) {
719566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
729566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY));
73c4762a1bSJed Brown   }
749566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
75c4762a1bSJed Brown   PetscFunctionReturn(0);
76c4762a1bSJed Brown }
77c4762a1bSJed Brown 
78c4762a1bSJed Brown static PetscErrorCode RHSHessianProductUU(TS ts,PetscReal t,Vec U,Vec *Vl,Vec Vr,Vec *VHV,void *ctx)
79c4762a1bSJed Brown {
80c4762a1bSJed Brown   const PetscScalar *vl,*vr,*u;
81c4762a1bSJed Brown   PetscScalar       *vhv;
82c4762a1bSJed Brown   PetscScalar       dJdU[2][2][2]={{{0}}};
83c4762a1bSJed Brown   PetscInt          i,j,k;
84c4762a1bSJed Brown   User              user = (User)ctx;
85c4762a1bSJed Brown 
86c4762a1bSJed Brown   PetscFunctionBeginUser;
879566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
889566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Vl[0],&vl));
899566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Vr,&vr));
909566063dSJacob Faibussowitsch   PetscCall(VecGetArray(VHV[0],&vhv));
91c4762a1bSJed Brown 
92c4762a1bSJed Brown   dJdU[1][0][0] = -2.*user->mu*u[1];
93c4762a1bSJed Brown   dJdU[1][1][0] = -2.*user->mu*u[0];
94c4762a1bSJed Brown   dJdU[1][0][1] = -2.*user->mu*u[0];
95c4762a1bSJed Brown   for (j=0;j<2;j++) {
96c4762a1bSJed Brown     vhv[j] = 0;
97c4762a1bSJed Brown     for (k=0;k<2;k++)
98c4762a1bSJed Brown       for (i=0;i<2;i++)
99c4762a1bSJed Brown         vhv[j] += vl[i]*dJdU[i][j][k]*vr[k];
100c4762a1bSJed Brown   }
1019566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
1029566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Vl[0],&vl));
1039566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Vr,&vr));
1049566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(VHV[0],&vhv));
105c4762a1bSJed Brown   PetscFunctionReturn(0);
106c4762a1bSJed Brown }
107c4762a1bSJed Brown 
108c4762a1bSJed Brown /* ----------------------- Implicit form of the ODE  -------------------- */
109c4762a1bSJed Brown 
110c4762a1bSJed Brown static PetscErrorCode IFunction(TS ts,PetscReal t,Vec U,Vec Udot,Vec F,void *ctx)
111c4762a1bSJed Brown {
112c4762a1bSJed Brown   User              user = (User)ctx;
113c4762a1bSJed Brown   const PetscScalar *u,*udot;
114c4762a1bSJed Brown   PetscScalar       *f;
115c4762a1bSJed Brown 
116c4762a1bSJed Brown   PetscFunctionBeginUser;
1179566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
1189566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Udot,&udot));
1199566063dSJacob Faibussowitsch   PetscCall(VecGetArray(F,&f));
120c4762a1bSJed Brown   f[0] = udot[0] - u[1];
121c4762a1bSJed Brown   f[1] = udot[1] - user->mu*((1.0-u[0]*u[0])*u[1] - u[0]) ;
1229566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
1239566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Udot,&udot));
1249566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(F,&f));
125c4762a1bSJed Brown   PetscFunctionReturn(0);
126c4762a1bSJed Brown }
127c4762a1bSJed Brown 
128c4762a1bSJed Brown static PetscErrorCode IJacobian(TS ts,PetscReal t,Vec U,Vec Udot,PetscReal a,Mat A,Mat B,void *ctx)
129c4762a1bSJed Brown {
130c4762a1bSJed Brown   User              user = (User)ctx;
131c4762a1bSJed Brown   PetscInt          rowcol[] = {0,1};
132c4762a1bSJed Brown   PetscScalar       J[2][2];
133c4762a1bSJed Brown   const PetscScalar *u;
134c4762a1bSJed Brown 
135c4762a1bSJed Brown   PetscFunctionBeginUser;
1369566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
137c4762a1bSJed Brown   J[0][0] = a;     J[0][1] =  -1.0;
138c4762a1bSJed Brown   J[1][0] = user->mu*(1.0 + 2.0*u[0]*u[1]);   J[1][1] = a - user->mu*(1.0-u[0]*u[0]);
1399566063dSJacob Faibussowitsch   PetscCall(MatSetValues(B,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES));
1409566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
1419566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
1429566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
143c4762a1bSJed Brown   if (A != B) {
1449566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
1459566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY));
146c4762a1bSJed Brown   }
147c4762a1bSJed Brown   PetscFunctionReturn(0);
148c4762a1bSJed Brown }
149c4762a1bSJed Brown 
150c4762a1bSJed Brown /* Monitor timesteps and use interpolation to output at integer multiples of 0.1 */
151c4762a1bSJed Brown static PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec U,void *ctx)
152c4762a1bSJed Brown {
153c4762a1bSJed Brown   const PetscScalar *u;
154c4762a1bSJed Brown   PetscReal         tfinal, dt;
155c4762a1bSJed Brown   User              user = (User)ctx;
156c4762a1bSJed Brown   Vec               interpolatedU;
157c4762a1bSJed Brown 
158c4762a1bSJed Brown   PetscFunctionBeginUser;
1599566063dSJacob Faibussowitsch   PetscCall(TSGetTimeStep(ts,&dt));
1609566063dSJacob Faibussowitsch   PetscCall(TSGetMaxTime(ts,&tfinal));
161c4762a1bSJed Brown 
162c4762a1bSJed Brown   while (user->next_output <= t && user->next_output <= tfinal) {
1639566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(U,&interpolatedU));
1649566063dSJacob Faibussowitsch     PetscCall(TSInterpolate(ts,user->next_output,interpolatedU));
1659566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(interpolatedU,&u));
16663a3b9bcSJacob Faibussowitsch     PetscCall(PetscPrintf(PETSC_COMM_WORLD,"[%g] %" PetscInt_FMT " TS %g (dt = %g) X %g %g\n",
167c4762a1bSJed Brown                           (double)user->next_output,step,(double)t,(double)dt,(double)PetscRealPart(u[0]),
168b122ec5aSJacob Faibussowitsch                           (double)PetscRealPart(u[1])));
1699566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(interpolatedU,&u));
1709566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&interpolatedU));
171c4762a1bSJed Brown     user->next_output += 0.1;
172c4762a1bSJed Brown   }
173c4762a1bSJed Brown   PetscFunctionReturn(0);
174c4762a1bSJed Brown }
175c4762a1bSJed Brown 
176c4762a1bSJed Brown static PetscErrorCode IHessianProductUU(TS ts,PetscReal t,Vec U,Vec *Vl,Vec Vr,Vec *VHV,void *ctx)
177c4762a1bSJed Brown {
178c4762a1bSJed Brown   const PetscScalar *vl,*vr,*u;
179c4762a1bSJed Brown   PetscScalar       *vhv;
180c4762a1bSJed Brown   PetscScalar       dJdU[2][2][2]={{{0}}};
181c4762a1bSJed Brown   PetscInt          i,j,k;
182c4762a1bSJed Brown   User              user = (User)ctx;
183c4762a1bSJed Brown 
184c4762a1bSJed Brown   PetscFunctionBeginUser;
1859566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(U,&u));
1869566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Vl[0],&vl));
1879566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(Vr,&vr));
1889566063dSJacob Faibussowitsch   PetscCall(VecGetArray(VHV[0],&vhv));
189c4762a1bSJed Brown   dJdU[1][0][0] = 2.*user->mu*u[1];
190c4762a1bSJed Brown   dJdU[1][1][0] = 2.*user->mu*u[0];
191c4762a1bSJed Brown   dJdU[1][0][1] = 2.*user->mu*u[0];
192c4762a1bSJed Brown   for (j=0;j<2;j++) {
193c4762a1bSJed Brown     vhv[j] = 0;
194c4762a1bSJed Brown     for (k=0;k<2;k++)
195c4762a1bSJed Brown       for (i=0;i<2;i++)
196c4762a1bSJed Brown         vhv[j] += vl[i]*dJdU[i][j][k]*vr[k];
197c4762a1bSJed Brown   }
1989566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(U,&u));
1999566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Vl[0],&vl));
2009566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(Vr,&vr));
2019566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(VHV[0],&vhv));
202c4762a1bSJed Brown   PetscFunctionReturn(0);
203c4762a1bSJed Brown }
204c4762a1bSJed Brown 
205c4762a1bSJed Brown /* ------------------ User-defined routines for TAO -------------------------- */
206c4762a1bSJed Brown 
207c4762a1bSJed Brown static PetscErrorCode FormFunctionGradient(Tao tao,Vec IC,PetscReal *f,Vec G,void *ctx)
208c4762a1bSJed Brown {
209c4762a1bSJed Brown   User              user_ptr = (User)ctx;
210c4762a1bSJed Brown   TS                ts = user_ptr->ts;
211c4762a1bSJed Brown   const PetscScalar *x_ptr;
212c4762a1bSJed Brown   PetscScalar       *y_ptr;
213c4762a1bSJed Brown 
214c4762a1bSJed Brown   PetscFunctionBeginUser;
2159566063dSJacob Faibussowitsch   PetscCall(VecCopy(IC,user_ptr->U)); /* set up the initial condition */
216c4762a1bSJed Brown 
2179566063dSJacob Faibussowitsch   PetscCall(TSSetTime(ts,0.0));
2189566063dSJacob Faibussowitsch   PetscCall(TSSetStepNumber(ts,0));
2199566063dSJacob Faibussowitsch   PetscCall(TSSetTimeStep(ts,0.001)); /* can be overwritten by command line options */
2209566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(ts));
2219566063dSJacob Faibussowitsch   PetscCall(TSSolve(ts,user_ptr->U));
222c4762a1bSJed Brown 
2239566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(user_ptr->U,&x_ptr));
2249566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user_ptr->Lambda[0],&y_ptr));
225c4762a1bSJed Brown   *f       = (x_ptr[0]-user_ptr->ob[0])*(x_ptr[0]-user_ptr->ob[0])+(x_ptr[1]-user_ptr->ob[1])*(x_ptr[1]-user_ptr->ob[1]);
226c4762a1bSJed Brown   y_ptr[0] = 2.*(x_ptr[0]-user_ptr->ob[0]);
227c4762a1bSJed Brown   y_ptr[1] = 2.*(x_ptr[1]-user_ptr->ob[1]);
2289566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user_ptr->Lambda[0],&y_ptr));
2299566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(user_ptr->U,&x_ptr));
230c4762a1bSJed Brown 
2319566063dSJacob Faibussowitsch   PetscCall(TSSetCostGradients(ts,1,user_ptr->Lambda,NULL));
2329566063dSJacob Faibussowitsch   PetscCall(TSAdjointSolve(ts));
2339566063dSJacob Faibussowitsch   PetscCall(VecCopy(user_ptr->Lambda[0],G));
234c4762a1bSJed Brown   PetscFunctionReturn(0);
235c4762a1bSJed Brown }
236c4762a1bSJed Brown 
237c4762a1bSJed Brown static PetscErrorCode FormHessian(Tao tao,Vec U,Mat H,Mat Hpre,void *ctx)
238c4762a1bSJed Brown {
239c4762a1bSJed Brown   User           user_ptr = (User)ctx;
240c4762a1bSJed Brown   PetscScalar    harr[2];
241c4762a1bSJed Brown   PetscScalar    *x_ptr;
242c4762a1bSJed Brown   const PetscInt rows[2] = {0,1};
243c4762a1bSJed Brown   PetscInt       col;
244c4762a1bSJed Brown 
245c4762a1bSJed Brown   PetscFunctionBeginUser;
2469566063dSJacob Faibussowitsch   PetscCall(VecCopy(U,user_ptr->U));
2479566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user_ptr->Dir,&x_ptr));
248c4762a1bSJed Brown   x_ptr[0] = 1.;
249c4762a1bSJed Brown   x_ptr[1] = 0.;
2509566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user_ptr->Dir,&x_ptr));
2519566063dSJacob Faibussowitsch   PetscCall(Adjoint2(user_ptr->U,harr,user_ptr));
252c4762a1bSJed Brown   col      = 0;
2539566063dSJacob Faibussowitsch   PetscCall(MatSetValues(H,2,rows,1,&col,harr,INSERT_VALUES));
254c4762a1bSJed Brown 
2559566063dSJacob Faibussowitsch   PetscCall(VecCopy(U,user_ptr->U));
2569566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user_ptr->Dir,&x_ptr));
257c4762a1bSJed Brown   x_ptr[0] = 0.;
258c4762a1bSJed Brown   x_ptr[1] = 1.;
2599566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user_ptr->Dir,&x_ptr));
2609566063dSJacob Faibussowitsch   PetscCall(Adjoint2(user_ptr->U,harr,user_ptr));
261c4762a1bSJed Brown   col      = 1;
2629566063dSJacob Faibussowitsch   PetscCall(MatSetValues(H,2,rows,1,&col,harr,INSERT_VALUES));
263c4762a1bSJed Brown 
2649566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(H,MAT_FINAL_ASSEMBLY));
2659566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(H,MAT_FINAL_ASSEMBLY));
266c4762a1bSJed Brown   if (H != Hpre) {
2679566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(Hpre,MAT_FINAL_ASSEMBLY));
2689566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(Hpre,MAT_FINAL_ASSEMBLY));
269c4762a1bSJed Brown   }
270c4762a1bSJed Brown   PetscFunctionReturn(0);
271c4762a1bSJed Brown }
272c4762a1bSJed Brown 
273c4762a1bSJed Brown static PetscErrorCode MatrixFreeHessian(Tao tao,Vec U,Mat H,Mat Hpre,void *ctx)
274c4762a1bSJed Brown {
275c4762a1bSJed Brown   User           user_ptr = (User)ctx;
276c4762a1bSJed Brown 
277c4762a1bSJed Brown   PetscFunctionBeginUser;
2789566063dSJacob Faibussowitsch   PetscCall(VecCopy(U,user_ptr->U));
279c4762a1bSJed Brown   PetscFunctionReturn(0);
280c4762a1bSJed Brown }
281c4762a1bSJed Brown 
282c4762a1bSJed Brown /* ------------ Routines calculating second-order derivatives -------------- */
283c4762a1bSJed Brown 
2840e3d61c9SBarry Smith /*
285c4762a1bSJed Brown   Compute the Hessian-vector product for the cost function using Second-order adjoint
286c4762a1bSJed Brown */
287c4762a1bSJed Brown PetscErrorCode Adjoint2(Vec U,PetscScalar arr[],User ctx)
288c4762a1bSJed Brown {
289c4762a1bSJed Brown   TS             ts = ctx->ts;
290c4762a1bSJed Brown   PetscScalar    *x_ptr,*y_ptr;
291c4762a1bSJed Brown   Mat            tlmsen;
292c4762a1bSJed Brown 
293c4762a1bSJed Brown   PetscFunctionBeginUser;
2949566063dSJacob Faibussowitsch   PetscCall(TSAdjointReset(ts));
295c4762a1bSJed Brown 
2969566063dSJacob Faibussowitsch   PetscCall(TSSetTime(ts,0.0));
2979566063dSJacob Faibussowitsch   PetscCall(TSSetStepNumber(ts,0));
2989566063dSJacob Faibussowitsch   PetscCall(TSSetTimeStep(ts,0.001));
2999566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(ts));
3009566063dSJacob Faibussowitsch   PetscCall(TSSetCostHessianProducts(ts,1,ctx->Lambda2,NULL,ctx->Dir));
3019566063dSJacob Faibussowitsch   PetscCall(TSAdjointSetForward(ts,NULL));
3029566063dSJacob Faibussowitsch   PetscCall(TSSolve(ts,U));
303c4762a1bSJed Brown 
304c4762a1bSJed Brown   /* Set terminal conditions for first- and second-order adjonts */
3059566063dSJacob Faibussowitsch   PetscCall(VecGetArray(U,&x_ptr));
3069566063dSJacob Faibussowitsch   PetscCall(VecGetArray(ctx->Lambda[0],&y_ptr));
307c4762a1bSJed Brown   y_ptr[0] = 2.*(x_ptr[0]-ctx->ob[0]);
308c4762a1bSJed Brown   y_ptr[1] = 2.*(x_ptr[1]-ctx->ob[1]);
3099566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(ctx->Lambda[0],&y_ptr));
3109566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(U,&x_ptr));
3119566063dSJacob Faibussowitsch   PetscCall(TSForwardGetSensitivities(ts,NULL,&tlmsen));
3129566063dSJacob Faibussowitsch   PetscCall(MatDenseGetColumn(tlmsen,0,&x_ptr));
3139566063dSJacob Faibussowitsch   PetscCall(VecGetArray(ctx->Lambda2[0],&y_ptr));
314c4762a1bSJed Brown   y_ptr[0] = 2.*x_ptr[0];
315c4762a1bSJed Brown   y_ptr[1] = 2.*x_ptr[1];
3169566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(ctx->Lambda2[0],&y_ptr));
3179566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreColumn(tlmsen,&x_ptr));
318c4762a1bSJed Brown 
3199566063dSJacob Faibussowitsch   PetscCall(TSSetCostGradients(ts,1,ctx->Lambda,NULL));
320c4762a1bSJed Brown   if (ctx->implicitform) {
3219566063dSJacob Faibussowitsch     PetscCall(TSSetIHessianProduct(ts,ctx->Ihp1,IHessianProductUU,NULL,NULL,NULL,NULL,NULL,NULL,ctx));
322c4762a1bSJed Brown   } else {
3239566063dSJacob Faibussowitsch     PetscCall(TSSetRHSHessianProduct(ts,ctx->Ihp1,RHSHessianProductUU,NULL,NULL,NULL,NULL,NULL,NULL,ctx));
324c4762a1bSJed Brown   }
3259566063dSJacob Faibussowitsch   PetscCall(TSAdjointSolve(ts));
326c4762a1bSJed Brown 
3279566063dSJacob Faibussowitsch   PetscCall(VecGetArray(ctx->Lambda2[0],&x_ptr));
328c4762a1bSJed Brown   arr[0] = x_ptr[0];
329c4762a1bSJed Brown   arr[1] = x_ptr[1];
3309566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(ctx->Lambda2[0],&x_ptr));
331c4762a1bSJed Brown 
3329566063dSJacob Faibussowitsch   PetscCall(TSAdjointReset(ts));
3339566063dSJacob Faibussowitsch   PetscCall(TSAdjointResetForward(ts));
334c4762a1bSJed Brown   PetscFunctionReturn(0);
335c4762a1bSJed Brown }
336c4762a1bSJed Brown 
337c4762a1bSJed Brown PetscErrorCode FiniteDiff(Vec U,PetscScalar arr[],User ctx)
338c4762a1bSJed Brown {
339c4762a1bSJed Brown   Vec               Up,G,Gp;
340c4762a1bSJed Brown   const PetscScalar eps = PetscRealConstant(1e-7);
341c4762a1bSJed Brown   PetscScalar       *u;
342c4762a1bSJed Brown   Tao               tao = NULL;
343c4762a1bSJed Brown   PetscReal         f;
344c4762a1bSJed Brown 
345c4762a1bSJed Brown   PetscFunctionBeginUser;
3469566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(U,&Up));
3479566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(U,&G));
3489566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(U,&Gp));
349c4762a1bSJed Brown 
3509566063dSJacob Faibussowitsch   PetscCall(FormFunctionGradient(tao,U,&f,G,ctx));
351c4762a1bSJed Brown 
3529566063dSJacob Faibussowitsch   PetscCall(VecCopy(U,Up));
3539566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Up,&u));
354c4762a1bSJed Brown   u[0] += eps;
3559566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Up,&u));
3569566063dSJacob Faibussowitsch   PetscCall(FormFunctionGradient(tao,Up,&f,Gp,ctx));
3579566063dSJacob Faibussowitsch   PetscCall(VecAXPY(Gp,-1,G));
3589566063dSJacob Faibussowitsch   PetscCall(VecScale(Gp,1./eps));
3599566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Gp,&u));
360c4762a1bSJed Brown   arr[0] = u[0];
361c4762a1bSJed Brown   arr[1] = u[1];
3629566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Gp,&u));
363c4762a1bSJed Brown 
3649566063dSJacob Faibussowitsch   PetscCall(VecCopy(U,Up));
3659566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Up,&u));
366c4762a1bSJed Brown   u[1] += eps;
3679566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Up,&u));
3689566063dSJacob Faibussowitsch   PetscCall(FormFunctionGradient(tao,Up,&f,Gp,ctx));
3699566063dSJacob Faibussowitsch   PetscCall(VecAXPY(Gp,-1,G));
3709566063dSJacob Faibussowitsch   PetscCall(VecScale(Gp,1./eps));
3719566063dSJacob Faibussowitsch   PetscCall(VecGetArray(Gp,&u));
372c4762a1bSJed Brown   arr[2] = u[0];
373c4762a1bSJed Brown   arr[3] = u[1];
3749566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(Gp,&u));
375c4762a1bSJed Brown 
3769566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&G));
3779566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&Gp));
3789566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&Up));
379c4762a1bSJed Brown   PetscFunctionReturn(0);
380c4762a1bSJed Brown }
381c4762a1bSJed Brown 
382c4762a1bSJed Brown static PetscErrorCode HessianProductMat(Mat mat,Vec svec,Vec y)
383c4762a1bSJed Brown {
384c4762a1bSJed Brown   User           user_ptr;
385c4762a1bSJed Brown   PetscScalar    *y_ptr;
386c4762a1bSJed Brown 
387c4762a1bSJed Brown   PetscFunctionBeginUser;
3889566063dSJacob Faibussowitsch   PetscCall(MatShellGetContext(mat,&user_ptr));
3899566063dSJacob Faibussowitsch   PetscCall(VecCopy(svec,user_ptr->Dir));
3909566063dSJacob Faibussowitsch   PetscCall(VecGetArray(y,&y_ptr));
3919566063dSJacob Faibussowitsch   PetscCall(Adjoint2(user_ptr->U,y_ptr,user_ptr));
3929566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(y,&y_ptr));
393c4762a1bSJed Brown   PetscFunctionReturn(0);
394c4762a1bSJed Brown }
395c4762a1bSJed Brown 
396c4762a1bSJed Brown int main(int argc,char **argv)
397c4762a1bSJed Brown {
398c4762a1bSJed Brown   PetscBool      monitor = PETSC_FALSE,mf = PETSC_TRUE;
399c4762a1bSJed Brown   PetscInt       mode = 0;
400c4762a1bSJed Brown   PetscMPIInt    size;
401c4762a1bSJed Brown   struct _n_User user;
402c4762a1bSJed Brown   Vec            x; /* working vector for TAO */
403c4762a1bSJed Brown   PetscScalar    *x_ptr,arr[4];
404c4762a1bSJed Brown   PetscScalar    ic1 = 2.2,ic2 = -0.7; /* initial guess for TAO */
405c4762a1bSJed Brown   Tao            tao;
406c4762a1bSJed Brown   KSP            ksp;
407c4762a1bSJed Brown   PC             pc;
408c4762a1bSJed Brown 
409c4762a1bSJed Brown   /* Initialize program */
410*327415f7SBarry Smith   PetscFunctionBeginUser;
4119566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc,&argv,NULL,help));
4129566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
4133c633725SBarry Smith   PetscCheck(size == 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
414c4762a1bSJed Brown 
415c4762a1bSJed Brown   /* Set runtime options */
416c4762a1bSJed Brown   user.next_output  = 0.0;
417c4762a1bSJed Brown   user.mu           = 1.0e3;
418c4762a1bSJed Brown   user.steps        = 0;
419c4762a1bSJed Brown   user.ftime        = 0.5;
420c4762a1bSJed Brown   user.implicitform = PETSC_TRUE;
421c4762a1bSJed Brown 
4229566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL));
4239566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-mu",&user.mu,NULL));
4249566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetInt(NULL,NULL,"-mode",&mode,NULL));
4259566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-ic1",&ic1,NULL));
4269566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-ic2",&ic2,NULL));
4279566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-my_tao_mf",&mf,NULL)); /* matrix-free hessian for optimization */
4289566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-implicitform",&user.implicitform,NULL));
429c4762a1bSJed Brown 
430c4762a1bSJed Brown   /* Create necessary matrix and vectors, solve same ODE on every process */
4319566063dSJacob Faibussowitsch   PetscCall(MatCreate(PETSC_COMM_WORLD,&user.A));
4329566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2));
4339566063dSJacob Faibussowitsch   PetscCall(MatSetFromOptions(user.A));
4349566063dSJacob Faibussowitsch   PetscCall(MatSetUp(user.A));
4359566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.U,NULL));
4369566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Dir,NULL));
4379566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Lambda[0],NULL));
4389566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Lambda2[0],NULL));
4399566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Ihp1[0],NULL));
440c4762a1bSJed Brown 
441c4762a1bSJed Brown   /* Create timestepping solver context */
4429566063dSJacob Faibussowitsch   PetscCall(TSCreate(PETSC_COMM_WORLD,&user.ts));
4439566063dSJacob Faibussowitsch   PetscCall(TSSetEquationType(user.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
444c4762a1bSJed Brown   if (user.implicitform) {
4459566063dSJacob Faibussowitsch     PetscCall(TSSetIFunction(user.ts,NULL,IFunction,&user));
4469566063dSJacob Faibussowitsch     PetscCall(TSSetIJacobian(user.ts,user.A,user.A,IJacobian,&user));
4479566063dSJacob Faibussowitsch     PetscCall(TSSetType(user.ts,TSCN));
448c4762a1bSJed Brown   } else {
4499566063dSJacob Faibussowitsch     PetscCall(TSSetRHSFunction(user.ts,NULL,RHSFunction,&user));
4509566063dSJacob Faibussowitsch     PetscCall(TSSetRHSJacobian(user.ts,user.A,user.A,RHSJacobian,&user));
4519566063dSJacob Faibussowitsch     PetscCall(TSSetType(user.ts,TSRK));
452c4762a1bSJed Brown   }
4539566063dSJacob Faibussowitsch   PetscCall(TSSetMaxTime(user.ts,user.ftime));
4549566063dSJacob Faibussowitsch   PetscCall(TSSetExactFinalTime(user.ts,TS_EXACTFINALTIME_MATCHSTEP));
455c4762a1bSJed Brown 
456c4762a1bSJed Brown   if (monitor) {
4579566063dSJacob Faibussowitsch     PetscCall(TSMonitorSet(user.ts,Monitor,&user,NULL));
458c4762a1bSJed Brown   }
459c4762a1bSJed Brown 
460c4762a1bSJed Brown   /* Set ODE initial conditions */
4619566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user.U,&x_ptr));
462c4762a1bSJed Brown   x_ptr[0] = 2.0;
463c4762a1bSJed Brown   x_ptr[1] = -2.0/3.0 + 10.0/(81.0*user.mu) - 292.0/(2187.0*user.mu*user.mu);
4649566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user.U,&x_ptr));
465c4762a1bSJed Brown 
466c4762a1bSJed Brown   /* Set runtime options */
4679566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(user.ts));
468c4762a1bSJed Brown 
469c4762a1bSJed Brown   /* Obtain the observation */
4709566063dSJacob Faibussowitsch   PetscCall(TSSolve(user.ts,user.U));
4719566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user.U,&x_ptr));
472c4762a1bSJed Brown   user.ob[0] = x_ptr[0];
473c4762a1bSJed Brown   user.ob[1] = x_ptr[1];
4749566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user.U,&x_ptr));
475c4762a1bSJed Brown 
4769566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(user.U,&x));
4779566063dSJacob Faibussowitsch   PetscCall(VecGetArray(x,&x_ptr));
478c4762a1bSJed Brown   x_ptr[0] = ic1;
479c4762a1bSJed Brown   x_ptr[1] = ic2;
4809566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(x,&x_ptr));
481c4762a1bSJed Brown 
482c4762a1bSJed Brown   /* Save trajectory of solution so that TSAdjointSolve() may be used */
4839566063dSJacob Faibussowitsch   PetscCall(TSSetSaveTrajectory(user.ts));
484c4762a1bSJed Brown 
485c4762a1bSJed Brown   /* Compare finite difference and second-order adjoint. */
486c4762a1bSJed Brown   switch (mode) {
487c4762a1bSJed Brown     case 2 :
4889566063dSJacob Faibussowitsch       PetscCall(FiniteDiff(x,arr,&user));
4899566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\n Finite difference approximation of the Hessian\n"));
4909566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g %g\n%g %g\n",(double)arr[0],(double)arr[1],(double)arr[2],(double)arr[3]));
491c4762a1bSJed Brown       break;
492c4762a1bSJed Brown     case 1 : /* Compute the Hessian column by column */
4939566063dSJacob Faibussowitsch       PetscCall(VecCopy(x,user.U));
4949566063dSJacob Faibussowitsch       PetscCall(VecGetArray(user.Dir,&x_ptr));
495c4762a1bSJed Brown       x_ptr[0] = 1.;
496c4762a1bSJed Brown       x_ptr[1] = 0.;
4979566063dSJacob Faibussowitsch       PetscCall(VecRestoreArray(user.Dir,&x_ptr));
4989566063dSJacob Faibussowitsch       PetscCall(Adjoint2(user.U,arr,&user));
4999566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\nFirst column of the Hessian\n"));
5009566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g\n%g\n",(double)arr[0],(double)arr[1]));
5019566063dSJacob Faibussowitsch       PetscCall(VecCopy(x,user.U));
5029566063dSJacob Faibussowitsch       PetscCall(VecGetArray(user.Dir,&x_ptr));
503c4762a1bSJed Brown       x_ptr[0] = 0.;
504c4762a1bSJed Brown       x_ptr[1] = 1.;
5059566063dSJacob Faibussowitsch       PetscCall(VecRestoreArray(user.Dir,&x_ptr));
5069566063dSJacob Faibussowitsch       PetscCall(Adjoint2(user.U,arr,&user));
5079566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\nSecond column of the Hessian\n"));
5089566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g\n%g\n",(double)arr[0],(double)arr[1]));
509c4762a1bSJed Brown       break;
510c4762a1bSJed Brown     case 0 :
511c4762a1bSJed Brown       /* Create optimization context and set up */
5129566063dSJacob Faibussowitsch       PetscCall(TaoCreate(PETSC_COMM_WORLD,&tao));
5139566063dSJacob Faibussowitsch       PetscCall(TaoSetType(tao,TAOBLMVM));
5149566063dSJacob Faibussowitsch       PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionGradient,(void *)&user));
515c4762a1bSJed Brown 
516c4762a1bSJed Brown       if (mf) {
5179566063dSJacob Faibussowitsch         PetscCall(MatCreateShell(PETSC_COMM_SELF,2,2,2,2,(void*)&user,&user.H));
5189566063dSJacob Faibussowitsch         PetscCall(MatShellSetOperation(user.H,MATOP_MULT,(void(*)(void))HessianProductMat));
5199566063dSJacob Faibussowitsch         PetscCall(MatSetOption(user.H,MAT_SYMMETRIC,PETSC_TRUE));
5209566063dSJacob Faibussowitsch         PetscCall(TaoSetHessian(tao,user.H,user.H,MatrixFreeHessian,(void *)&user));
521c4762a1bSJed Brown       } else { /* Create Hessian matrix */
5229566063dSJacob Faibussowitsch         PetscCall(MatCreate(PETSC_COMM_WORLD,&user.H));
5239566063dSJacob Faibussowitsch         PetscCall(MatSetSizes(user.H,PETSC_DECIDE,PETSC_DECIDE,2,2));
5249566063dSJacob Faibussowitsch         PetscCall(MatSetUp(user.H));
5259566063dSJacob Faibussowitsch         PetscCall(TaoSetHessian(tao,user.H,user.H,FormHessian,(void *)&user));
526c4762a1bSJed Brown       }
527c4762a1bSJed Brown 
528c4762a1bSJed Brown       /* Not use any preconditioner */
5299566063dSJacob Faibussowitsch       PetscCall(TaoGetKSP(tao,&ksp));
530c4762a1bSJed Brown       if (ksp) {
5319566063dSJacob Faibussowitsch         PetscCall(KSPGetPC(ksp,&pc));
5329566063dSJacob Faibussowitsch         PetscCall(PCSetType(pc,PCNONE));
533c4762a1bSJed Brown       }
534c4762a1bSJed Brown 
535c4762a1bSJed Brown       /* Set initial solution guess */
5369566063dSJacob Faibussowitsch       PetscCall(TaoSetSolution(tao,x));
5379566063dSJacob Faibussowitsch       PetscCall(TaoSetFromOptions(tao));
5389566063dSJacob Faibussowitsch       PetscCall(TaoSolve(tao));
5399566063dSJacob Faibussowitsch       PetscCall(TaoDestroy(&tao));
5409566063dSJacob Faibussowitsch       PetscCall(MatDestroy(&user.H));
541c4762a1bSJed Brown       break;
542c4762a1bSJed Brown     default:
543c4762a1bSJed Brown       break;
544c4762a1bSJed Brown   }
545c4762a1bSJed Brown 
546c4762a1bSJed Brown   /* Free work space.  All PETSc objects should be destroyed when they are no longer needed. */
5479566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&user.A));
5489566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.U));
5499566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Lambda[0]));
5509566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Lambda2[0]));
5519566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Ihp1[0]));
5529566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Dir));
5539566063dSJacob Faibussowitsch   PetscCall(TSDestroy(&user.ts));
5549566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&x));
5559566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
556b122ec5aSJacob Faibussowitsch   return 0;
557c4762a1bSJed Brown }
558c4762a1bSJed Brown 
559c4762a1bSJed Brown /*TEST
560c4762a1bSJed Brown     build:
561c4762a1bSJed Brown       requires: !complex !single
562c4762a1bSJed Brown 
563c4762a1bSJed Brown     test:
564c4762a1bSJed Brown       args:  -ts_type cn -viewer_binary_skip_info -tao_monitor -tao_view -mu 1000 -ts_dt 0.03125
565c4762a1bSJed Brown       output_file: output/ex20opt_ic_1.out
566c4762a1bSJed Brown 
567c4762a1bSJed Brown     test:
568c4762a1bSJed Brown       suffix: 2
569c4762a1bSJed Brown       args:  -ts_type beuler -viewer_binary_skip_info -tao_monitor -tao_view -mu 100 -ts_dt 0.01 -tao_type bntr -tao_bnk_pc_type none
570c4762a1bSJed Brown       output_file: output/ex20opt_ic_2.out
571c4762a1bSJed Brown 
572c4762a1bSJed Brown     test:
573c4762a1bSJed Brown       suffix: 3
574c4762a1bSJed Brown       args:  -ts_type cn -viewer_binary_skip_info -tao_monitor -tao_view -mu 100 -ts_dt 0.01 -tao_type bntr -tao_bnk_pc_type none
575c4762a1bSJed Brown       output_file: output/ex20opt_ic_3.out
576c4762a1bSJed Brown 
577c4762a1bSJed Brown     test:
578c4762a1bSJed Brown       suffix: 4
579c4762a1bSJed Brown       args: -implicitform 0 -ts_dt 0.01 -ts_max_steps 2 -ts_rhs_jacobian_test_mult_transpose -mat_shell_test_mult_transpose_view -ts_rhs_jacobian_test_mult -mat_shell_test_mult_view -mode 1 -my_tao_mf
580c4762a1bSJed Brown TEST*/
581