xref: /petsc/src/ts/tutorials/ex20opt_ic.c (revision 63a3b9bc7a1f24f247904ccba9383635fe6abade)
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));
166*63a3b9bcSJacob 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 */
4109566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc,&argv,NULL,help));
4119566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
4123c633725SBarry Smith   PetscCheck(size == 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
413c4762a1bSJed Brown 
414c4762a1bSJed Brown   /* Set runtime options */
415c4762a1bSJed Brown   user.next_output  = 0.0;
416c4762a1bSJed Brown   user.mu           = 1.0e3;
417c4762a1bSJed Brown   user.steps        = 0;
418c4762a1bSJed Brown   user.ftime        = 0.5;
419c4762a1bSJed Brown   user.implicitform = PETSC_TRUE;
420c4762a1bSJed Brown 
4219566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL));
4229566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-mu",&user.mu,NULL));
4239566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetInt(NULL,NULL,"-mode",&mode,NULL));
4249566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-ic1",&ic1,NULL));
4259566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetReal(NULL,NULL,"-ic2",&ic2,NULL));
4269566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-my_tao_mf",&mf,NULL)); /* matrix-free hessian for optimization */
4279566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetBool(NULL,NULL,"-implicitform",&user.implicitform,NULL));
428c4762a1bSJed Brown 
429c4762a1bSJed Brown   /* Create necessary matrix and vectors, solve same ODE on every process */
4309566063dSJacob Faibussowitsch   PetscCall(MatCreate(PETSC_COMM_WORLD,&user.A));
4319566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2));
4329566063dSJacob Faibussowitsch   PetscCall(MatSetFromOptions(user.A));
4339566063dSJacob Faibussowitsch   PetscCall(MatSetUp(user.A));
4349566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.U,NULL));
4359566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Dir,NULL));
4369566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Lambda[0],NULL));
4379566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Lambda2[0],NULL));
4389566063dSJacob Faibussowitsch   PetscCall(MatCreateVecs(user.A,&user.Ihp1[0],NULL));
439c4762a1bSJed Brown 
440c4762a1bSJed Brown   /* Create timestepping solver context */
4419566063dSJacob Faibussowitsch   PetscCall(TSCreate(PETSC_COMM_WORLD,&user.ts));
4429566063dSJacob Faibussowitsch   PetscCall(TSSetEquationType(user.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
443c4762a1bSJed Brown   if (user.implicitform) {
4449566063dSJacob Faibussowitsch     PetscCall(TSSetIFunction(user.ts,NULL,IFunction,&user));
4459566063dSJacob Faibussowitsch     PetscCall(TSSetIJacobian(user.ts,user.A,user.A,IJacobian,&user));
4469566063dSJacob Faibussowitsch     PetscCall(TSSetType(user.ts,TSCN));
447c4762a1bSJed Brown   } else {
4489566063dSJacob Faibussowitsch     PetscCall(TSSetRHSFunction(user.ts,NULL,RHSFunction,&user));
4499566063dSJacob Faibussowitsch     PetscCall(TSSetRHSJacobian(user.ts,user.A,user.A,RHSJacobian,&user));
4509566063dSJacob Faibussowitsch     PetscCall(TSSetType(user.ts,TSRK));
451c4762a1bSJed Brown   }
4529566063dSJacob Faibussowitsch   PetscCall(TSSetMaxTime(user.ts,user.ftime));
4539566063dSJacob Faibussowitsch   PetscCall(TSSetExactFinalTime(user.ts,TS_EXACTFINALTIME_MATCHSTEP));
454c4762a1bSJed Brown 
455c4762a1bSJed Brown   if (monitor) {
4569566063dSJacob Faibussowitsch     PetscCall(TSMonitorSet(user.ts,Monitor,&user,NULL));
457c4762a1bSJed Brown   }
458c4762a1bSJed Brown 
459c4762a1bSJed Brown   /* Set ODE initial conditions */
4609566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user.U,&x_ptr));
461c4762a1bSJed Brown   x_ptr[0] = 2.0;
462c4762a1bSJed Brown   x_ptr[1] = -2.0/3.0 + 10.0/(81.0*user.mu) - 292.0/(2187.0*user.mu*user.mu);
4639566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user.U,&x_ptr));
464c4762a1bSJed Brown 
465c4762a1bSJed Brown   /* Set runtime options */
4669566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(user.ts));
467c4762a1bSJed Brown 
468c4762a1bSJed Brown   /* Obtain the observation */
4699566063dSJacob Faibussowitsch   PetscCall(TSSolve(user.ts,user.U));
4709566063dSJacob Faibussowitsch   PetscCall(VecGetArray(user.U,&x_ptr));
471c4762a1bSJed Brown   user.ob[0] = x_ptr[0];
472c4762a1bSJed Brown   user.ob[1] = x_ptr[1];
4739566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(user.U,&x_ptr));
474c4762a1bSJed Brown 
4759566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(user.U,&x));
4769566063dSJacob Faibussowitsch   PetscCall(VecGetArray(x,&x_ptr));
477c4762a1bSJed Brown   x_ptr[0] = ic1;
478c4762a1bSJed Brown   x_ptr[1] = ic2;
4799566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(x,&x_ptr));
480c4762a1bSJed Brown 
481c4762a1bSJed Brown   /* Save trajectory of solution so that TSAdjointSolve() may be used */
4829566063dSJacob Faibussowitsch   PetscCall(TSSetSaveTrajectory(user.ts));
483c4762a1bSJed Brown 
484c4762a1bSJed Brown   /* Compare finite difference and second-order adjoint. */
485c4762a1bSJed Brown   switch (mode) {
486c4762a1bSJed Brown     case 2 :
4879566063dSJacob Faibussowitsch       PetscCall(FiniteDiff(x,arr,&user));
4889566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\n Finite difference approximation of the Hessian\n"));
4899566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g %g\n%g %g\n",(double)arr[0],(double)arr[1],(double)arr[2],(double)arr[3]));
490c4762a1bSJed Brown       break;
491c4762a1bSJed Brown     case 1 : /* Compute the Hessian column by column */
4929566063dSJacob Faibussowitsch       PetscCall(VecCopy(x,user.U));
4939566063dSJacob Faibussowitsch       PetscCall(VecGetArray(user.Dir,&x_ptr));
494c4762a1bSJed Brown       x_ptr[0] = 1.;
495c4762a1bSJed Brown       x_ptr[1] = 0.;
4969566063dSJacob Faibussowitsch       PetscCall(VecRestoreArray(user.Dir,&x_ptr));
4979566063dSJacob Faibussowitsch       PetscCall(Adjoint2(user.U,arr,&user));
4989566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\nFirst column of the Hessian\n"));
4999566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g\n%g\n",(double)arr[0],(double)arr[1]));
5009566063dSJacob Faibussowitsch       PetscCall(VecCopy(x,user.U));
5019566063dSJacob Faibussowitsch       PetscCall(VecGetArray(user.Dir,&x_ptr));
502c4762a1bSJed Brown       x_ptr[0] = 0.;
503c4762a1bSJed Brown       x_ptr[1] = 1.;
5049566063dSJacob Faibussowitsch       PetscCall(VecRestoreArray(user.Dir,&x_ptr));
5059566063dSJacob Faibussowitsch       PetscCall(Adjoint2(user.U,arr,&user));
5069566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"\nSecond column of the Hessian\n"));
5079566063dSJacob Faibussowitsch       PetscCall(PetscPrintf(PETSC_COMM_WORLD,"%g\n%g\n",(double)arr[0],(double)arr[1]));
508c4762a1bSJed Brown       break;
509c4762a1bSJed Brown     case 0 :
510c4762a1bSJed Brown       /* Create optimization context and set up */
5119566063dSJacob Faibussowitsch       PetscCall(TaoCreate(PETSC_COMM_WORLD,&tao));
5129566063dSJacob Faibussowitsch       PetscCall(TaoSetType(tao,TAOBLMVM));
5139566063dSJacob Faibussowitsch       PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionGradient,(void *)&user));
514c4762a1bSJed Brown 
515c4762a1bSJed Brown       if (mf) {
5169566063dSJacob Faibussowitsch         PetscCall(MatCreateShell(PETSC_COMM_SELF,2,2,2,2,(void*)&user,&user.H));
5179566063dSJacob Faibussowitsch         PetscCall(MatShellSetOperation(user.H,MATOP_MULT,(void(*)(void))HessianProductMat));
5189566063dSJacob Faibussowitsch         PetscCall(MatSetOption(user.H,MAT_SYMMETRIC,PETSC_TRUE));
5199566063dSJacob Faibussowitsch         PetscCall(TaoSetHessian(tao,user.H,user.H,MatrixFreeHessian,(void *)&user));
520c4762a1bSJed Brown       } else { /* Create Hessian matrix */
5219566063dSJacob Faibussowitsch         PetscCall(MatCreate(PETSC_COMM_WORLD,&user.H));
5229566063dSJacob Faibussowitsch         PetscCall(MatSetSizes(user.H,PETSC_DECIDE,PETSC_DECIDE,2,2));
5239566063dSJacob Faibussowitsch         PetscCall(MatSetUp(user.H));
5249566063dSJacob Faibussowitsch         PetscCall(TaoSetHessian(tao,user.H,user.H,FormHessian,(void *)&user));
525c4762a1bSJed Brown       }
526c4762a1bSJed Brown 
527c4762a1bSJed Brown       /* Not use any preconditioner */
5289566063dSJacob Faibussowitsch       PetscCall(TaoGetKSP(tao,&ksp));
529c4762a1bSJed Brown       if (ksp) {
5309566063dSJacob Faibussowitsch         PetscCall(KSPGetPC(ksp,&pc));
5319566063dSJacob Faibussowitsch         PetscCall(PCSetType(pc,PCNONE));
532c4762a1bSJed Brown       }
533c4762a1bSJed Brown 
534c4762a1bSJed Brown       /* Set initial solution guess */
5359566063dSJacob Faibussowitsch       PetscCall(TaoSetSolution(tao,x));
5369566063dSJacob Faibussowitsch       PetscCall(TaoSetFromOptions(tao));
5379566063dSJacob Faibussowitsch       PetscCall(TaoSolve(tao));
5389566063dSJacob Faibussowitsch       PetscCall(TaoDestroy(&tao));
5399566063dSJacob Faibussowitsch       PetscCall(MatDestroy(&user.H));
540c4762a1bSJed Brown       break;
541c4762a1bSJed Brown     default:
542c4762a1bSJed Brown       break;
543c4762a1bSJed Brown   }
544c4762a1bSJed Brown 
545c4762a1bSJed Brown   /* Free work space.  All PETSc objects should be destroyed when they are no longer needed. */
5469566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&user.A));
5479566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.U));
5489566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Lambda[0]));
5499566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Lambda2[0]));
5509566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Ihp1[0]));
5519566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&user.Dir));
5529566063dSJacob Faibussowitsch   PetscCall(TSDestroy(&user.ts));
5539566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&x));
5549566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
555b122ec5aSJacob Faibussowitsch   return 0;
556c4762a1bSJed Brown }
557c4762a1bSJed Brown 
558c4762a1bSJed Brown /*TEST
559c4762a1bSJed Brown     build:
560c4762a1bSJed Brown       requires: !complex !single
561c4762a1bSJed Brown 
562c4762a1bSJed Brown     test:
563c4762a1bSJed Brown       args:  -ts_type cn -viewer_binary_skip_info -tao_monitor -tao_view -mu 1000 -ts_dt 0.03125
564c4762a1bSJed Brown       output_file: output/ex20opt_ic_1.out
565c4762a1bSJed Brown 
566c4762a1bSJed Brown     test:
567c4762a1bSJed Brown       suffix: 2
568c4762a1bSJed 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
569c4762a1bSJed Brown       output_file: output/ex20opt_ic_2.out
570c4762a1bSJed Brown 
571c4762a1bSJed Brown     test:
572c4762a1bSJed Brown       suffix: 3
573c4762a1bSJed 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
574c4762a1bSJed Brown       output_file: output/ex20opt_ic_3.out
575c4762a1bSJed Brown 
576c4762a1bSJed Brown     test:
577c4762a1bSJed Brown       suffix: 4
578c4762a1bSJed 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
579c4762a1bSJed Brown TEST*/
580