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