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