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