xref: /petsc/src/ts/tutorials/autodiff/ex16opt_ic.cxx (revision ebead697dbf761eb322f829370bbe90b3bd93fa3)
1 static char help[] = "Demonstrates automatic Jacobian generation using ADOL-C for an ODE-constrained optimization problem.\n\
2 Input parameters include:\n\
3       -mu : stiffness parameter\n\n";
4 
5 /*
6    REQUIRES configuration of PETSc with option --download-adolc.
7 
8    For documentation on ADOL-C, see
9      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
10 */
11 /* ------------------------------------------------------------------------
12   See ex16opt_ic for a description of the problem being solved.
13   ------------------------------------------------------------------------- */
14 #include <petsctao.h>
15 #include <petscts.h>
16 #include <petscmat.h>
17 #include "adolc-utils/drivers.cxx"
18 #include <adolc/adolc.h>
19 
20 typedef struct _n_User *User;
21 struct _n_User {
22   PetscReal mu;
23   PetscReal next_output;
24   PetscInt  steps;
25 
26   /* Sensitivity analysis support */
27   PetscReal ftime,x_ob[2];
28   Mat       A;             /* Jacobian matrix */
29   Vec       x,lambda[2];   /* adjoint variables */
30 
31   /* Automatic differentiation support */
32   AdolcCtx  *adctx;
33 };
34 
35 PetscErrorCode FormFunctionGradient(Tao,Vec,PetscReal*,Vec,void*);
36 
37 /*
38   'Passive' RHS function, used in residual evaluations during the time integration.
39 */
40 static PetscErrorCode RHSFunctionPassive(TS ts,PetscReal t,Vec X,Vec F,void *ctx)
41 {
42   User              user = (User)ctx;
43   PetscScalar       *f;
44   const PetscScalar *x;
45 
46   PetscFunctionBeginUser;
47   PetscCall(VecGetArrayRead(X,&x));
48   PetscCall(VecGetArray(F,&f));
49   f[0] = x[1];
50   f[1] = user->mu*(1.-x[0]*x[0])*x[1]-x[0];
51   PetscCall(VecRestoreArrayRead(X,&x));
52   PetscCall(VecRestoreArray(F,&f));
53   PetscFunctionReturn(0);
54 }
55 
56 /*
57   Trace RHS to mark on tape 1 the dependence of f upon x. This tape is used in generating the
58   Jacobian transform.
59 */
60 static PetscErrorCode RHSFunctionActive(TS ts,PetscReal t,Vec X,Vec F,void *ctx)
61 {
62   User              user = (User)ctx;
63   PetscReal         mu   = user->mu;
64   PetscScalar       *f;
65   const PetscScalar *x;
66 
67   adouble           f_a[2];                     /* adouble for dependent variables */
68   adouble           x_a[2];                     /* adouble for independent variables */
69 
70   PetscFunctionBeginUser;
71   PetscCall(VecGetArrayRead(X,&x));
72   PetscCall(VecGetArray(F,&f));
73 
74   trace_on(1);                                  /* Start of active section */
75   x_a[0] <<= x[0]; x_a[1] <<= x[1];             /* Mark as independent */
76   f_a[0] = x_a[1];
77   f_a[1] = mu*(1.-x_a[0]*x_a[0])*x_a[1]-x_a[0];
78   f_a[0] >>= f[0]; f_a[1] >>= f[1];             /* Mark as dependent */
79   trace_off(1);                                 /* End of active section */
80 
81   PetscCall(VecRestoreArrayRead(X,&x));
82   PetscCall(VecRestoreArray(F,&f));
83   PetscFunctionReturn(0);
84 }
85 
86 /*
87   Compute the Jacobian w.r.t. x using PETSc-ADOL-C driver.
88 */
89 static PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec X,Mat A,Mat B,void *ctx)
90 {
91   User              user=(User)ctx;
92   const PetscScalar *x;
93 
94   PetscFunctionBeginUser;
95   PetscCall(VecGetArrayRead(X,&x));
96   PetscCall(PetscAdolcComputeRHSJacobian(1,A,x,user->adctx));
97   PetscCall(VecRestoreArrayRead(X,&x));
98   PetscFunctionReturn(0);
99 }
100 
101 /*
102   Monitor timesteps and use interpolation to output at integer multiples of 0.1
103 */
104 static PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec X,void *ctx)
105 {
106   const PetscScalar *x;
107   PetscReal         tfinal, dt, tprev;
108   User              user = (User)ctx;
109 
110   PetscFunctionBeginUser;
111   PetscCall(TSGetTimeStep(ts,&dt));
112   PetscCall(TSGetMaxTime(ts,&tfinal));
113   PetscCall(TSGetPrevTime(ts,&tprev));
114   PetscCall(VecGetArrayRead(X,&x));
115   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"[%.1f] %" PetscInt_FMT " TS %.6f (dt = %.6f) X % 12.6e % 12.6e\n",(double)user->next_output,step,(double)t,(double)dt,(double)PetscRealPart(x[0]),(double)PetscRealPart(x[1])));
116   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"t %.6f (tprev = %.6f) \n",(double)t,(double)tprev));
117   PetscCall(VecGetArrayRead(X,&x));
118   PetscFunctionReturn(0);
119 }
120 
121 int main(int argc,char **argv)
122 {
123   TS                 ts = NULL;          /* nonlinear solver */
124   Vec                ic,r;
125   PetscBool          monitor = PETSC_FALSE;
126   PetscScalar        *x_ptr;
127   PetscMPIInt        size;
128   struct _n_User     user;
129   AdolcCtx           *adctx;
130   Tao                tao;
131   KSP                ksp;
132   PC                 pc;
133 
134   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
135      Initialize program
136      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
137   PetscFunctionBeginUser;
138   PetscCall(PetscInitialize(&argc,&argv,NULL,help));
139   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
140   PetscCheck(size == 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
141 
142   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
143     Set runtime options and create AdolcCtx
144     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
145   PetscCall(PetscNew(&adctx));
146   user.mu          = 1.0;
147   user.next_output = 0.0;
148   user.steps       = 0;
149   user.ftime       = 0.5;
150   adctx->m = 2;adctx->n = 2;adctx->p = 2;
151   user.adctx = adctx;
152 
153   PetscCall(PetscOptionsGetReal(NULL,NULL,"-mu",&user.mu,NULL));
154   PetscCall(PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL));
155 
156   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
157     Create necessary matrix and vectors, solve same ODE on every process
158     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
159   PetscCall(MatCreate(PETSC_COMM_WORLD,&user.A));
160   PetscCall(MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2));
161   PetscCall(MatSetFromOptions(user.A));
162   PetscCall(MatSetUp(user.A));
163   PetscCall(MatCreateVecs(user.A,&user.x,NULL));
164 
165   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
166      Set initial conditions
167    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
168   PetscCall(VecGetArray(user.x,&x_ptr));
169   x_ptr[0] = 2.0;   x_ptr[1] = 0.66666654321;
170   PetscCall(VecRestoreArray(user.x,&x_ptr));
171 
172   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
173      Trace just once on each tape and put zeros on Jacobian diagonal
174      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
175   PetscCall(VecDuplicate(user.x,&r));
176   PetscCall(RHSFunctionActive(ts,0.,user.x,r,&user));
177   PetscCall(VecSet(r,0));
178   PetscCall(MatDiagonalSet(user.A,r,INSERT_VALUES));
179   PetscCall(VecDestroy(&r));
180 
181   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
182      Create timestepping solver context
183      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
184   PetscCall(TSCreate(PETSC_COMM_WORLD,&ts));
185   PetscCall(TSSetType(ts,TSRK));
186   PetscCall(TSSetRHSFunction(ts,NULL,RHSFunctionPassive,&user));
187   PetscCall(TSSetRHSJacobian(ts,user.A,user.A,RHSJacobian,&user));
188   PetscCall(TSSetMaxTime(ts,user.ftime));
189   PetscCall(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP));
190   if (monitor) {
191     PetscCall(TSMonitorSet(ts,Monitor,&user,NULL));
192   }
193 
194   PetscCall(TSSetTime(ts,0.0));
195   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"mu %g, steps %" PetscInt_FMT ", ftime %g\n",(double)user.mu,user.steps,(double)(user.ftime)));
196 
197   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
198     Save trajectory of solution so that TSAdjointSolve() may be used
199    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
200   PetscCall(TSSetSaveTrajectory(ts));
201 
202   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
203      Set runtime options
204    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
205   PetscCall(TSSetFromOptions(ts));
206 
207   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
208      Solve nonlinear system
209      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
210   PetscCall(TSSolve(ts,user.x));
211   PetscCall(TSGetSolveTime(ts,&(user.ftime)));
212   PetscCall(TSGetStepNumber(ts,&user.steps));
213   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"mu %g, steps %" PetscInt_FMT ", ftime %g\n",(double)user.mu,user.steps,(double)user.ftime));
214 
215   PetscCall(VecGetArray(user.x,&x_ptr));
216   user.x_ob[0] = x_ptr[0];
217   user.x_ob[1] = x_ptr[1];
218   PetscCall(VecRestoreArray(user.x,&x_ptr));
219 
220   PetscCall(MatCreateVecs(user.A,&user.lambda[0],NULL));
221 
222   /* Create TAO solver and set desired solution method */
223   PetscCall(TaoCreate(PETSC_COMM_WORLD,&tao));
224   PetscCall(TaoSetType(tao,TAOCG));
225 
226   /* Set initial solution guess */
227   PetscCall(MatCreateVecs(user.A,&ic,NULL));
228   PetscCall(VecGetArray(ic,&x_ptr));
229   x_ptr[0]  = 2.1;
230   x_ptr[1]  = 0.7;
231   PetscCall(VecRestoreArray(ic,&x_ptr));
232 
233   PetscCall(TaoSetSolution(tao,ic));
234 
235   /* Set routine for function and gradient evaluation */
236   PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionGradient,(void *)&user));
237 
238   /* Check for any TAO command line options */
239   PetscCall(TaoSetFromOptions(tao));
240   PetscCall(TaoGetKSP(tao,&ksp));
241   if (ksp) {
242     PetscCall(KSPGetPC(ksp,&pc));
243     PetscCall(PCSetType(pc,PCNONE));
244   }
245 
246   PetscCall(TaoSetTolerances(tao,1e-10,PETSC_DEFAULT,PETSC_DEFAULT));
247 
248   /* SOLVE THE APPLICATION */
249   PetscCall(TaoSolve(tao));
250 
251   /* Free TAO data structures */
252   PetscCall(TaoDestroy(&tao));
253 
254   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
255      Free work space.  All PETSc objects should be destroyed when they
256      are no longer needed.
257    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
258   PetscCall(MatDestroy(&user.A));
259   PetscCall(VecDestroy(&user.x));
260   PetscCall(VecDestroy(&user.lambda[0]));
261   PetscCall(TSDestroy(&ts));
262   PetscCall(VecDestroy(&ic));
263   PetscCall(PetscFree(adctx));
264   PetscCall(PetscFinalize());
265   return 0;
266 }
267 
268 /* ------------------------------------------------------------------ */
269 /*
270    FormFunctionGradient - Evaluates the function and corresponding gradient.
271 
272    Input Parameters:
273    tao - the Tao context
274    X   - the input vector
275    ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()
276 
277    Output Parameters:
278    f   - the newly evaluated function
279    G   - the newly evaluated gradient
280 */
281 PetscErrorCode FormFunctionGradient(Tao tao,Vec IC,PetscReal *f,Vec G,void *ctx)
282 {
283   User              user = (User)ctx;
284   TS                ts;
285   PetscScalar       *x_ptr,*y_ptr;
286 
287   PetscFunctionBeginUser;
288   PetscCall(VecCopy(IC,user->x));
289 
290   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
291      Create timestepping solver context
292      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
293   PetscCall(TSCreate(PETSC_COMM_WORLD,&ts));
294   PetscCall(TSSetType(ts,TSRK));
295   PetscCall(TSSetRHSFunction(ts,NULL,RHSFunctionPassive,user));
296   /*   Set RHS Jacobian  for the adjoint integration */
297   PetscCall(TSSetRHSJacobian(ts,user->A,user->A,RHSJacobian,user));
298 
299   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
300      Set time
301    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
302   PetscCall(TSSetTime(ts,0.0));
303   PetscCall(TSSetTimeStep(ts,.001));
304   PetscCall(TSSetMaxTime(ts,0.5));
305   PetscCall(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP));
306 
307   PetscCall(TSSetTolerances(ts,1e-7,NULL,1e-7,NULL));
308 
309   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
310     Save trajectory of solution so that TSAdjointSolve() may be used
311    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
312   PetscCall(TSSetSaveTrajectory(ts));
313 
314   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
315      Set runtime options
316    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
317   PetscCall(TSSetFromOptions(ts));
318 
319   PetscCall(TSSolve(ts,user->x));
320   PetscCall(TSGetSolveTime(ts,&user->ftime));
321   PetscCall(TSGetStepNumber(ts,&user->steps));
322   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"mu %.6f, steps %" PetscInt_FMT ", ftime %g\n",(double)user->mu,user->steps,(double)user->ftime));
323 
324   PetscCall(VecGetArray(user->x,&x_ptr));
325   *f   = (x_ptr[0]-user->x_ob[0])*(x_ptr[0]-user->x_ob[0])+(x_ptr[1]-user->x_ob[1])*(x_ptr[1]-user->x_ob[1]);
326   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"Observed value y_ob=[%f; %f], ODE solution y=[%f;%f], Cost function f=%f\n",(double)user->x_ob[0],(double)user->x_ob[1],(double)x_ptr[0],(double)x_ptr[1],(double)(*f)));
327 
328   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
329      Adjoint model starts here
330      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
331   /*   Redet initial conditions for the adjoint integration */
332   PetscCall(VecGetArray(user->lambda[0],&y_ptr));
333   y_ptr[0] = 2.*(x_ptr[0]-user->x_ob[0]);
334   y_ptr[1] = 2.*(x_ptr[1]-user->x_ob[1]);
335   PetscCall(VecRestoreArray(user->lambda[0],&y_ptr));
336   PetscCall(VecRestoreArray(user->x,&x_ptr));
337   PetscCall(TSSetCostGradients(ts,1,user->lambda,NULL));
338 
339   PetscCall(TSAdjointSolve(ts));
340 
341   PetscCall(VecCopy(user->lambda[0],G));
342 
343   PetscCall(TSDestroy(&ts));
344   PetscFunctionReturn(0);
345 }
346 
347 /*TEST
348 
349   build:
350     requires: double !complex adolc
351 
352   test:
353     suffix: 1
354     args: -ts_rhs_jacobian_test_mult_transpose FALSE -tao_max_it 2 -ts_rhs_jacobian_test_mult FALSE
355     output_file: output/ex16opt_ic_1.out
356 
357 TEST*/
358