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