xref: /petsc/src/ts/tutorials/autodiff/ex16adj_tl.cxx (revision 503c0ea9b45bcfbcebbb1ea5341243bbc69f0bea)
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   PetscCall(VecGetArrayRead(X,&x));
67   PetscCall(VecGetArray(F,&f));
68   PetscCall(EvaluateResidual(x,user->mu,f));
69   PetscCall(VecRestoreArrayRead(X,&x));
70   PetscCall(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   PetscCall(VecGetArrayRead(X,&x));
89   x_a[0].setValue(x[0]);
90   x_a[1].setValue(x[1]);
91   mu_a.setValue(user->mu);
92   PetscCall(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   PetscCall(EvaluateResidual(x_a,mu_a,f_a));
101 
102   /* Extract derivatives */
103   PetscCall(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       PetscCall(MatSetValues(A,1,&i,1,&j,&J[i][j],INSERT_VALUES));
111     }
112   }
113   PetscCall(PetscFree(J));
114   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
115   PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
116   if (A != B) {
117     PetscCall(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY));
118     PetscCall(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   PetscCall(VecGetArray(X,&x));
139   x_a[0].setValue(x[0]);
140   x_a[1].setValue(x[1]);
141   mu_a.setValue(user->mu);
142   PetscCall(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   PetscCall(EvaluateResidual(x_a,mu_a,f_a));
151 
152   /* Extract derivatives */
153   PetscCall(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     PetscCall(MatSetValues(A,1,&i,1,&j,&J[i][user->adctx->n],INSERT_VALUES));
160   }
161   PetscCall(PetscFree(J));
162   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
163   PetscCall(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   PetscCall(TSGetTimeStep(ts,&dt));
178   PetscCall(TSGetMaxTime(ts,&tfinal));
179   PetscCall(TSGetPrevTime(ts,&tprev));
180   PetscCall(VecGetArrayRead(X,&x));
181   PetscCall(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   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"t %.6f (tprev = %.6f) \n",(double)t,(double)tprev));
183   PetscCall(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   Vec            lambda[2],mu[2];
201 
202   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
203      Initialize program
204      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
205   PetscCall(PetscInitialize(&argc,&argv,NULL,help));
206   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
207   PetscCheckFalse(size != 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");
208 
209   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
210     Set runtime options and create AdolcCtx
211     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
212   PetscCall(PetscNew(&adctx));
213   user.mu          = 1;
214   user.next_output = 0.0;
215   adctx->m = 2;adctx->n = 2;adctx->p = 2;
216   user.adctx = adctx;
217   adtl::setNumDir(adctx->n+1); /* #indep. variables, plus parameters */
218 
219   PetscCall(PetscOptionsGetReal(NULL,NULL,"-mu",&user.mu,NULL));
220   PetscCall(PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL));
221 
222   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
223     Create necessary matrix and vectors, solve same ODE on every process
224     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
225   PetscCall(MatCreate(PETSC_COMM_WORLD,&A));
226   PetscCall(MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,2,2));
227   PetscCall(MatSetFromOptions(A));
228   PetscCall(MatSetUp(A));
229   PetscCall(MatCreateVecs(A,&x,NULL));
230 
231   PetscCall(MatCreate(PETSC_COMM_WORLD,&Jacp));
232   PetscCall(MatSetSizes(Jacp,PETSC_DECIDE,PETSC_DECIDE,2,1));
233   PetscCall(MatSetFromOptions(Jacp));
234   PetscCall(MatSetUp(Jacp));
235 
236   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
237      Create timestepping solver context
238      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
239   PetscCall(TSCreate(PETSC_COMM_WORLD,&ts));
240   PetscCall(TSSetType(ts,TSRK));
241   PetscCall(TSSetRHSFunction(ts,NULL,RHSFunctionPassive,&user));
242 
243   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
244      Set initial conditions
245    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
246   PetscCall(VecGetArray(x,&x_ptr));
247   x_ptr[0] = 2;   x_ptr[1] = 0.66666654321;
248   PetscCall(VecRestoreArray(x,&x_ptr));
249 
250   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
251      Set RHS Jacobian for the adjoint integration
252      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
253   PetscCall(TSSetRHSJacobian(ts,A,A,RHSJacobian,&user));
254   PetscCall(TSSetMaxTime(ts,ftime));
255   PetscCall(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP));
256   if (monitor) {
257     PetscCall(TSMonitorSet(ts,Monitor,&user,NULL));
258   }
259   PetscCall(TSSetTimeStep(ts,.001));
260 
261   /*
262     Have the TS save its trajectory so that TSAdjointSolve() may be used
263   */
264   PetscCall(TSSetSaveTrajectory(ts));
265 
266   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
267      Set runtime options
268    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
269   PetscCall(TSSetFromOptions(ts));
270 
271   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
272      Solve nonlinear system
273      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
274   PetscCall(TSSolve(ts,x));
275   PetscCall(TSGetSolveTime(ts,&ftime));
276   PetscCall(TSGetStepNumber(ts,&steps));
277   PetscCall(PetscPrintf(PETSC_COMM_WORLD,"mu %g, steps %D, ftime %g\n",(double)user.mu,steps,(double)ftime));
278   PetscCall(VecView(x,PETSC_VIEWER_STDOUT_WORLD));
279 
280   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
281      Start the Adjoint model
282      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
283   PetscCall(MatCreateVecs(A,&lambda[0],NULL));
284   PetscCall(MatCreateVecs(A,&lambda[1],NULL));
285   /*   Reset initial conditions for the adjoint integration */
286   PetscCall(VecGetArray(lambda[0],&x_ptr));
287   x_ptr[0] = 1.0;   x_ptr[1] = 0.0;
288   PetscCall(VecRestoreArray(lambda[0],&x_ptr));
289   PetscCall(VecGetArray(lambda[1],&x_ptr));
290   x_ptr[0] = 0.0;   x_ptr[1] = 1.0;
291   PetscCall(VecRestoreArray(lambda[1],&x_ptr));
292 
293   PetscCall(MatCreateVecs(Jacp,&mu[0],NULL));
294   PetscCall(MatCreateVecs(Jacp,&mu[1],NULL));
295   PetscCall(VecGetArray(mu[0],&x_ptr));
296   x_ptr[0] = 0.0;
297   PetscCall(VecRestoreArray(mu[0],&x_ptr));
298   PetscCall(VecGetArray(mu[1],&x_ptr));
299   x_ptr[0] = 0.0;
300   PetscCall(VecRestoreArray(mu[1],&x_ptr));
301   PetscCall(TSSetCostGradients(ts,2,lambda,mu));
302 
303   /*   Set RHS JacobianP */
304   PetscCall(TSSetRHSJacobianP(ts,Jacp,RHSJacobianP,&user));
305 
306   PetscCall(TSAdjointSolve(ts));
307 
308   PetscCall(VecView(lambda[0],PETSC_VIEWER_STDOUT_WORLD));
309   PetscCall(VecView(lambda[1],PETSC_VIEWER_STDOUT_WORLD));
310   PetscCall(VecView(mu[0],PETSC_VIEWER_STDOUT_WORLD));
311   PetscCall(VecView(mu[1],PETSC_VIEWER_STDOUT_WORLD));
312 
313   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
314      Free work space.  All PETSc objects should be destroyed when they
315      are no longer needed.
316    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
317   PetscCall(MatDestroy(&A));
318   PetscCall(MatDestroy(&Jacp));
319   PetscCall(VecDestroy(&x));
320   PetscCall(VecDestroy(&lambda[0]));
321   PetscCall(VecDestroy(&lambda[1]));
322   PetscCall(VecDestroy(&mu[0]));
323   PetscCall(VecDestroy(&mu[1]));
324   PetscCall(TSDestroy(&ts));
325   PetscCall(PetscFree(adctx));
326   PetscCall(PetscFinalize());
327   return 0;
328 }
329 
330 /*TEST
331 
332   build:
333     requires: double !complex adolc
334 
335   test:
336     suffix: 1
337     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
338     output_file: output/ex16adj_tl_1.out
339 
340   test:
341     suffix: 2
342     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor -mu 5
343     output_file: output/ex16adj_tl_2.out
344 
345   test:
346     suffix: 3
347     args: -ts_max_steps 10 -monitor
348     output_file: output/ex16adj_tl_3.out
349 
350   test:
351     suffix: 4
352     args: -ts_max_steps 10 -monitor -mu 5
353     output_file: output/ex16adj_tl_4.out
354 
355 TEST*/
356