xref: /petsc/src/ts/tutorials/autodiff/ex16adj_tl.cxx (revision c87ba875e4007ad659b117ea274f03d5f4cd5ea7)
1 static char help[] = "Demonstrates tapeless automatic Jacobian generation using ADOL-C for an adjoint sensitivity analysis of the van der Pol equation.\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: TS^adjoint sensitivity analysis
9    Concepts: Automatic differentation using ADOL-C
10    Concepts: Tapeless automatic differentiation using ADOL-C
11    Concepts: Automatic differentation w.r.t. a parameter using ADOL-C
12    Processors: 1
13 */
14 /*
15    REQUIRES configuration of PETSc with option --download-adolc.
16 
17    For documentation on ADOL-C, see
18      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
19 */
20 /* ------------------------------------------------------------------------
21    See ex16adj for a description of the problem being solved.
22   ------------------------------------------------------------------------- */
23 
24 #include <petscts.h>
25 #include <petscmat.h>
26 
27 #define ADOLC_TAPELESS
28 #define NUMBER_DIRECTIONS 3
29 #include "adolc-utils/drivers.cxx"
30 #include <adolc/adtl.h>
31 using namespace adtl;
32 
33 typedef struct _n_User *User;
34 struct _n_User {
35   PetscReal mu;
36   PetscReal next_output;
37   PetscReal tprev;
38 
39   /* Automatic differentiation support */
40   AdolcCtx  *adctx;
41   Vec       F;
42 };
43 
44 /*
45   Residual evaluation templated, so as to allow for PetscScalar or adouble
46   arguments.
47 */
48 template <class T> PetscErrorCode EvaluateResidual(T *x,T mu,T *f)
49 {
50   PetscFunctionBegin;
51   f[0] = x[1];
52   f[1] = mu*(1.-x[0]*x[0])*x[1]-x[0];
53   PetscFunctionReturn(0);
54 }
55 
56 /*
57   'Passive' RHS function, used in residual evaluations during the time integration.
58 */
59 static PetscErrorCode RHSFunctionPassive(TS ts,PetscReal t,Vec X,Vec F,void *ctx)
60 {
61   PetscErrorCode    ierr;
62   User              user = (User)ctx;
63   PetscScalar       *f,*x;
64 
65   PetscFunctionBeginUser;
66   ierr = VecGetArray(X,&x);CHKERRQ(ierr);
67   ierr = VecGetArray(F,&f);CHKERRQ(ierr);
68   ierr = EvaluateResidual(x,user->mu,f);CHKERRQ(ierr);
69   ierr = VecRestoreArray(X,&x);CHKERRQ(ierr);
70   ierr = VecRestoreArray(F,&f);CHKERRQ(ierr);
71   PetscFunctionReturn(0);
72 }
73 
74 /*
75   Compute the Jacobian w.r.t. x using tapeless mode of ADOL-C.
76 */
77 static PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec X,Mat A,Mat B,void *ctx)
78 {
79   PetscErrorCode ierr;
80   User           user = (User)ctx;
81   PetscScalar    *x,**J;
82   adouble        f_a[2];      /* 'active' double for dependent variables */
83   adouble        x_a[2],mu_a; /* 'active' doubles for independent variables */
84   PetscInt       i,j;
85 
86   PetscFunctionBeginUser;
87 
88   /* Set values for independent variables and parameters */
89   ierr = VecGetArray(X,&x);CHKERRQ(ierr);
90   x_a[0].setValue(x[0]);
91   x_a[1].setValue(x[1]);
92   mu_a.setValue(user->mu);
93   ierr = VecRestoreArray(X,&x);CHKERRQ(ierr);
94 
95   /* Set seed matrix as 3x3 identity matrix */
96   x_a[0].setADValue(0,1.);x_a[0].setADValue(1,0.);x_a[0].setADValue(2,0.);
97   x_a[1].setADValue(0,0.);x_a[1].setADValue(1,1.);x_a[1].setADValue(2,0.);
98   mu_a.setADValue(0,0.);mu_a.setADValue(1,0.);mu_a.setADValue(2,1.);
99 
100   /* Evaluate residual (on active variables) */
101   ierr = EvaluateResidual(x_a,mu_a,f_a);CHKERRQ(ierr);
102 
103   /* Extract derivatives */
104   ierr = PetscMalloc1(user->adctx->n,&J);
105   J[0] = (PetscScalar*) f_a[0].getADValue();
106   J[1] = (PetscScalar*) f_a[1].getADValue();
107 
108   /* Set matrix values */
109   for (i=0; i<user->adctx->m; i++) {
110     for (j=0; j<user->adctx->n; j++) {
111       ierr = MatSetValues(A,1,&i,1,&j,&J[i][j],INSERT_VALUES);CHKERRQ(ierr);
112     }
113   }
114   ierr = PetscFree(J);CHKERRQ(ierr);
115   ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
116   ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
117   if (A != B) {
118     ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
119     ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
120   }
121   PetscFunctionReturn(0);
122 }
123 
124 /*
125   Compute the Jacobian w.r.t. mu using tapeless mode of ADOL-C.
126 */
127 static PetscErrorCode RHSJacobianP(TS ts,PetscReal t,Vec X,Mat A,void *ctx)
128 {
129   PetscErrorCode ierr;
130   User           user = (User)ctx;
131   PetscScalar    **J;
132   PetscScalar    *x;
133   adouble        f_a[2];      /* 'active' double for dependent variables */
134   adouble        x_a[2],mu_a; /* 'active' doubles for independent variables */
135   PetscInt       i,j = 0;
136 
137   PetscFunctionBeginUser;
138 
139   /* Set values for independent variables and parameters */
140   ierr = VecGetArray(X,&x);CHKERRQ(ierr);
141   x_a[0].setValue(x[0]);
142   x_a[1].setValue(x[1]);
143   mu_a.setValue(user->mu);
144   ierr = VecRestoreArray(X,&x);CHKERRQ(ierr);
145 
146   /* Set seed matrix as 3x3 identity matrix */
147   x_a[0].setADValue(0,1.);x_a[0].setADValue(1,0.);x_a[0].setADValue(2,0.);
148   x_a[1].setADValue(0,0.);x_a[1].setADValue(1,1.);x_a[1].setADValue(2,0.);
149   mu_a.setADValue(0,0.);mu_a.setADValue(1,0.);mu_a.setADValue(2,1.);
150 
151   /* Evaluate residual (on active variables) */
152   ierr = EvaluateResidual(x_a,mu_a,f_a);CHKERRQ(ierr);
153 
154   /* Extract derivatives */
155   ierr = PetscMalloc1(2,&J);CHKERRQ(ierr);
156   J[0] = (PetscScalar*) f_a[0].getADValue();
157   J[1] = (PetscScalar*) f_a[1].getADValue();
158 
159   /* Set matrix values */
160   for (i=0; i<user->adctx->m; i++) {
161     ierr = MatSetValues(A,1,&i,1,&j,&J[i][user->adctx->n],INSERT_VALUES);CHKERRQ(ierr);
162   }
163   ierr = PetscFree(J);CHKERRQ(ierr);
164   ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
165   ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
166   PetscFunctionReturn(0);
167 }
168 
169 /*
170   Monitor timesteps and use interpolation to output at integer multiples of 0.1
171 */
172 static PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec X,void *ctx)
173 {
174   PetscErrorCode    ierr;
175   const PetscScalar *x;
176   PetscReal         tfinal, dt, tprev;
177   User              user = (User)ctx;
178 
179   PetscFunctionBeginUser;
180   ierr = TSGetTimeStep(ts,&dt);CHKERRQ(ierr);
181   ierr = TSGetMaxTime(ts,&tfinal);CHKERRQ(ierr);
182   ierr = TSGetPrevTime(ts,&tprev);CHKERRQ(ierr);
183   ierr = VecGetArrayRead(X,&x);CHKERRQ(ierr);
184   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);
185   ierr = PetscPrintf(PETSC_COMM_WORLD,"t %.6f (tprev = %.6f) \n",(double)t,(double)tprev);CHKERRQ(ierr);
186   ierr = VecRestoreArrayRead(X,&x);CHKERRQ(ierr);
187   PetscFunctionReturn(0);
188 }
189 
190 int main(int argc,char **argv)
191 {
192   TS             ts;            /* nonlinear solver */
193   Vec            x;             /* solution, residual vectors */
194   Mat            A;             /* Jacobian matrix */
195   Mat            Jacp;          /* JacobianP matrix */
196   PetscInt       steps;
197   PetscReal      ftime   = 0.5;
198   PetscBool      monitor = PETSC_FALSE;
199   PetscScalar    *x_ptr;
200   PetscMPIInt    size;
201   struct _n_User user;
202   AdolcCtx       *adctx;
203   PetscErrorCode ierr;
204   Vec            lambda[2],mu[2];
205 
206   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
207      Initialize program
208      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
209   ierr = PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;
210   ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRQ(ierr);
211   if (size != 1) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_MPI_WRONG_SIZE,"This is a uniprocessor example only!");
212 
213   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
214     Set runtime options and create AdolcCtx
215     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
216   ierr = PetscNew(&adctx);CHKERRQ(ierr);
217   user.mu          = 1;
218   user.next_output = 0.0;
219   adctx->m = 2;adctx->n = 2;adctx->p = 2;
220   user.adctx = adctx;
221   adtl::setNumDir(adctx->n+1); /* #indep. variables, plus parameters */
222 
223   ierr = PetscOptionsGetReal(NULL,NULL,"-mu",&user.mu,NULL);CHKERRQ(ierr);
224   ierr = PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL);CHKERRQ(ierr);
225 
226   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
227     Create necessary matrix and vectors, solve same ODE on every process
228     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
229   ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
230   ierr = MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,2,2);CHKERRQ(ierr);
231   ierr = MatSetFromOptions(A);CHKERRQ(ierr);
232   ierr = MatSetUp(A);CHKERRQ(ierr);
233   ierr = MatCreateVecs(A,&x,NULL);CHKERRQ(ierr);
234 
235   ierr = MatCreate(PETSC_COMM_WORLD,&Jacp);CHKERRQ(ierr);
236   ierr = MatSetSizes(Jacp,PETSC_DECIDE,PETSC_DECIDE,2,1);CHKERRQ(ierr);
237   ierr = MatSetFromOptions(Jacp);CHKERRQ(ierr);
238   ierr = MatSetUp(Jacp);CHKERRQ(ierr);
239 
240   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
241      Create timestepping solver context
242      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
243   ierr = TSCreate(PETSC_COMM_WORLD,&ts);CHKERRQ(ierr);
244   ierr = TSSetType(ts,TSRK);CHKERRQ(ierr);
245   ierr = TSSetRHSFunction(ts,NULL,RHSFunctionPassive,&user);CHKERRQ(ierr);
246 
247   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
248      Set initial conditions
249    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
250   ierr = VecGetArray(x,&x_ptr);CHKERRQ(ierr);
251   x_ptr[0] = 2;   x_ptr[1] = 0.66666654321;
252   ierr = VecRestoreArray(x,&x_ptr);CHKERRQ(ierr);
253 
254   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
255      Set RHS Jacobian for the adjoint integration
256      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
257   ierr = TSSetRHSJacobian(ts,A,A,RHSJacobian,&user);CHKERRQ(ierr);
258   ierr = TSSetMaxTime(ts,ftime);CHKERRQ(ierr);
259   ierr = TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
260   if (monitor) {
261     ierr = TSMonitorSet(ts,Monitor,&user,NULL);CHKERRQ(ierr);
262   }
263   ierr = TSSetTimeStep(ts,.001);CHKERRQ(ierr);
264 
265   /*
266     Have the TS save its trajectory so that TSAdjointSolve() may be used
267   */
268   ierr = TSSetSaveTrajectory(ts);CHKERRQ(ierr);
269 
270   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
271      Set runtime options
272    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
273   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
274 
275   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
276      Solve nonlinear system
277      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
278   ierr = TSSolve(ts,x);CHKERRQ(ierr);
279   ierr = TSGetSolveTime(ts,&ftime);CHKERRQ(ierr);
280   ierr = TSGetStepNumber(ts,&steps);CHKERRQ(ierr);
281   ierr = PetscPrintf(PETSC_COMM_WORLD,"mu %g, steps %D, ftime %g\n",(double)user.mu,steps,(double)ftime);CHKERRQ(ierr);
282   ierr = VecView(x,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
283 
284   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
285      Start the Adjoint model
286      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
287   ierr = MatCreateVecs(A,&lambda[0],NULL);CHKERRQ(ierr);
288   ierr = MatCreateVecs(A,&lambda[1],NULL);CHKERRQ(ierr);
289   /*   Reset initial conditions for the adjoint integration */
290   ierr = VecGetArray(lambda[0],&x_ptr);CHKERRQ(ierr);
291   x_ptr[0] = 1.0;   x_ptr[1] = 0.0;
292   ierr = VecRestoreArray(lambda[0],&x_ptr);CHKERRQ(ierr);
293   ierr = VecGetArray(lambda[1],&x_ptr);CHKERRQ(ierr);
294   x_ptr[0] = 0.0;   x_ptr[1] = 1.0;
295   ierr = VecRestoreArray(lambda[1],&x_ptr);CHKERRQ(ierr);
296 
297   ierr = MatCreateVecs(Jacp,&mu[0],NULL);CHKERRQ(ierr);
298   ierr = MatCreateVecs(Jacp,&mu[1],NULL);CHKERRQ(ierr);
299   ierr = VecGetArray(mu[0],&x_ptr);CHKERRQ(ierr);
300   x_ptr[0] = 0.0;
301   ierr = VecRestoreArray(mu[0],&x_ptr);CHKERRQ(ierr);
302   ierr = VecGetArray(mu[1],&x_ptr);CHKERRQ(ierr);
303   x_ptr[0] = 0.0;
304   ierr = VecRestoreArray(mu[1],&x_ptr);CHKERRQ(ierr);
305   ierr = TSSetCostGradients(ts,2,lambda,mu);CHKERRQ(ierr);
306 
307 
308   /*   Set RHS JacobianP */
309   ierr = TSSetRHSJacobianP(ts,Jacp,RHSJacobianP,&user);CHKERRQ(ierr);
310 
311   ierr = TSAdjointSolve(ts);CHKERRQ(ierr);
312 
313   ierr = VecView(lambda[0],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
314   ierr = VecView(lambda[1],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
315   ierr = VecView(mu[0],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
316   ierr = VecView(mu[1],PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
317 
318   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
319      Free work space.  All PETSc objects should be destroyed when they
320      are no longer needed.
321    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
322   ierr = MatDestroy(&A);CHKERRQ(ierr);
323   ierr = MatDestroy(&Jacp);CHKERRQ(ierr);
324   ierr = VecDestroy(&x);CHKERRQ(ierr);
325   ierr = VecDestroy(&lambda[0]);CHKERRQ(ierr);
326   ierr = VecDestroy(&lambda[1]);CHKERRQ(ierr);
327   ierr = VecDestroy(&mu[0]);CHKERRQ(ierr);
328   ierr = VecDestroy(&mu[1]);CHKERRQ(ierr);
329   ierr = TSDestroy(&ts);CHKERRQ(ierr);
330   ierr = PetscFree(adctx);CHKERRQ(ierr);
331   ierr = PetscFinalize();
332   return ierr;
333 }
334 
335 /*TEST
336 
337   build:
338     requires: double !complex adolc
339 
340   test:
341     suffix: 1
342     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
343     output_file: output/ex16adj_tl_1.out
344 
345   test:
346     suffix: 2
347     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor -mu 5
348     output_file: output/ex16adj_tl_2.out
349 
350   test:
351     suffix: 3
352     args: -ts_max_steps 10 -monitor
353     output_file: output/ex16adj_tl_3.out
354 
355   test:
356     suffix: 4
357     args: -ts_max_steps 10 -monitor -mu 5
358     output_file: output/ex16adj_tl_4.out
359 
360 TEST*/
361