xref: /petsc/src/ts/tutorials/autodiff/ex16adj.cxx (revision 4e8208cbcbc709572b8abe32f33c78b69c819375)
1 static char help[] = "Demonstrates 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 #include "adolc-utils/drivers.cxx"
18 #include <adolc/adolc.h>
19 
20 typedef struct _n_User *User;
21 struct _n_User {
22   PetscReal mu;
23   PetscReal next_output;
24   PetscReal tprev;
25 
26   /* Automatic differentiation support */
27   AdolcCtx *adctx;
28 };
29 
30 /*
31   'Passive' RHS function, used in residual evaluations during the time integration.
32 */
RHSFunctionPassive(TS ts,PetscReal t,Vec X,Vec F,PetscCtx ctx)33 static PetscErrorCode RHSFunctionPassive(TS ts, PetscReal t, Vec X, Vec F, PetscCtx ctx)
34 {
35   User               user = (User)ctx;
36   PetscScalar       *f;
37   const PetscScalar *x;
38 
39   PetscFunctionBeginUser;
40   PetscCall(VecGetArrayRead(X, &x));
41   PetscCall(VecGetArray(F, &f));
42   f[0] = x[1];
43   f[1] = user->mu * (1. - x[0] * x[0]) * x[1] - x[0];
44   PetscCall(VecRestoreArrayRead(X, &x));
45   PetscCall(VecRestoreArray(F, &f));
46   PetscFunctionReturn(PETSC_SUCCESS);
47 }
48 
49 /*
50   Trace RHS to mark on tape 1 the dependence of f upon x. This tape is used in generating the
51   Jacobian transform.
52 */
RHSFunctionActive(TS ts,PetscReal t,Vec X,Vec F,PetscCtx ctx)53 static PetscErrorCode RHSFunctionActive(TS ts, PetscReal t, Vec X, Vec F, PetscCtx ctx)
54 {
55   User               user = (User)ctx;
56   PetscScalar       *f;
57   const PetscScalar *x;
58 
59   adouble f_a[2]; /* 'active' double for dependent variables */
60   adouble x_a[2]; /* 'active' double for independent variables */
61 
62   PetscFunctionBeginUser;
63   PetscCall(VecGetArrayRead(X, &x));
64   PetscCall(VecGetArray(F, &f));
65 
66   /* Start of active section */
67   trace_on(1);
68   x_a[0] <<= x[0];
69   x_a[1] <<= x[1]; /* Mark independence */
70   f_a[0] = x_a[1];
71   f_a[1] = user->mu * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0];
72   f_a[0] >>= f[0];
73   f_a[1] >>= f[1]; /* Mark dependence */
74   trace_off();
75   /* End of active section */
76 
77   PetscCall(VecRestoreArrayRead(X, &x));
78   PetscCall(VecRestoreArray(F, &f));
79   PetscFunctionReturn(PETSC_SUCCESS);
80 }
81 
82 /*
83   Trace RHS again to mark on tape 2 the dependence of f upon the parameter mu. This tape is used in
84   generating JacobianP.
85 */
RHSFunctionActiveP(TS ts,PetscReal t,Vec X,Vec F,PetscCtx ctx)86 static PetscErrorCode RHSFunctionActiveP(TS ts, PetscReal t, Vec X, Vec F, PetscCtx ctx)
87 {
88   User               user = (User)ctx;
89   PetscScalar       *f;
90   const PetscScalar *x;
91 
92   adouble f_a[2];       /* 'active' double for dependent variables */
93   adouble x_a[2], mu_a; /* 'active' double for independent variables */
94 
95   PetscFunctionBeginUser;
96   PetscCall(VecGetArrayRead(X, &x));
97   PetscCall(VecGetArray(F, &f));
98 
99   /* Start of active section */
100   trace_on(3);
101   x_a[0] <<= x[0];
102   x_a[1] <<= x[1];
103   mu_a <<= user->mu; /* Mark independence */
104   f_a[0] = x_a[1];
105   f_a[1] = mu_a * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0];
106   f_a[0] >>= f[0];
107   f_a[1] >>= f[1]; /* Mark dependence */
108   trace_off();
109   /* End of active section */
110 
111   PetscCall(VecRestoreArrayRead(X, &x));
112   PetscCall(VecRestoreArray(F, &f));
113   PetscFunctionReturn(PETSC_SUCCESS);
114 }
115 
116 /*
117   Compute the Jacobian w.r.t. x using PETSc-ADOL-C driver for explicit TS.
118 */
RHSJacobian(TS ts,PetscReal t,Vec X,Mat A,Mat B,PetscCtx ctx)119 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec X, Mat A, Mat B, PetscCtx ctx)
120 {
121   User               user = (User)ctx;
122   const PetscScalar *x;
123 
124   PetscFunctionBeginUser;
125   PetscCall(VecGetArrayRead(X, &x));
126   PetscCall(PetscAdolcComputeRHSJacobian(1, A, x, user->adctx));
127   PetscCall(VecRestoreArrayRead(X, &x));
128   PetscFunctionReturn(PETSC_SUCCESS);
129 }
130 
131 /*
132   Compute the Jacobian w.r.t. mu using PETSc-ADOL-C driver for explicit TS.
133 */
RHSJacobianP(TS ts,PetscReal t,Vec X,Mat A,PetscCtx ctx)134 static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec X, Mat A, PetscCtx ctx)
135 {
136   User               user = (User)ctx;
137   const PetscScalar *x;
138 
139   PetscFunctionBeginUser;
140   PetscCall(VecGetArrayRead(X, &x));
141   PetscCall(PetscAdolcComputeRHSJacobianP(3, A, x, &user->mu, user->adctx));
142   PetscCall(VecRestoreArrayRead(X, &x));
143   PetscFunctionReturn(PETSC_SUCCESS);
144 }
145 
146 /*
147   Monitor timesteps and use interpolation to output at integer multiples of 0.1
148 */
Monitor(TS ts,PetscInt step,PetscReal t,Vec X,PetscCtx ctx)149 static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, PetscCtx ctx)
150 {
151   const PetscScalar *x;
152   PetscReal          tfinal, dt, tprev;
153   User               user = (User)ctx;
154 
155   PetscFunctionBeginUser;
156   PetscCall(TSGetTimeStep(ts, &dt));
157   PetscCall(TSGetMaxTime(ts, &tfinal));
158   PetscCall(TSGetPrevTime(ts, &tprev));
159   PetscCall(VecGetArrayRead(X, &x));
160   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%.1f] %" PetscInt_FMT " 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])));
161   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "t %.6f (tprev = %.6f) \n", (double)t, (double)tprev));
162   PetscCall(VecRestoreArrayRead(X, &x));
163   PetscFunctionReturn(PETSC_SUCCESS);
164 }
165 
main(int argc,char ** argv)166 int main(int argc, char **argv)
167 {
168   TS             ts;   /* nonlinear solver */
169   Vec            x;    /* solution, residual vectors */
170   Mat            A;    /* Jacobian matrix */
171   Mat            Jacp; /* JacobianP matrix */
172   PetscInt       steps;
173   PetscReal      ftime   = 0.5;
174   PetscBool      monitor = PETSC_FALSE;
175   PetscScalar   *x_ptr;
176   PetscMPIInt    size;
177   struct _n_User user;
178   AdolcCtx      *adctx;
179   Vec            lambda[2], mu[2], r;
180 
181   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
182      Initialize program
183      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
184   PetscFunctionBeginUser;
185   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
186   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
187   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
188 
189   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
190     Set runtime options and create AdolcCtx
191     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
192   PetscCall(PetscNew(&adctx));
193   user.mu           = 1;
194   user.next_output  = 0.0;
195   adctx->m          = 2;
196   adctx->n          = 2;
197   adctx->p          = 2;
198   adctx->num_params = 1;
199   user.adctx        = adctx;
200 
201   PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &user.mu, NULL));
202   PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL));
203 
204   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
205     Create necessary matrix and vectors, solve same ODE on every process
206     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
207   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
208   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
209   PetscCall(MatSetFromOptions(A));
210   PetscCall(MatSetUp(A));
211   PetscCall(MatCreateVecs(A, &x, NULL));
212 
213   PetscCall(MatCreate(PETSC_COMM_WORLD, &Jacp));
214   PetscCall(MatSetSizes(Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 1));
215   PetscCall(MatSetFromOptions(Jacp));
216   PetscCall(MatSetUp(Jacp));
217 
218   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
219      Create timestepping solver context
220      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
221   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
222   PetscCall(TSSetType(ts, TSRK));
223   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, &user));
224 
225   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
226      Set initial conditions
227    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
228   PetscCall(VecGetArray(x, &x_ptr));
229   x_ptr[0] = 2;
230   x_ptr[1] = 0.66666654321;
231   PetscCall(VecRestoreArray(x, &x_ptr));
232 
233   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
234      Trace just once on each tape and put zeros on Jacobian diagonal
235      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
236   PetscCall(VecDuplicate(x, &r));
237   PetscCall(RHSFunctionActive(ts, 0., x, r, &user));
238   PetscCall(RHSFunctionActiveP(ts, 0., x, r, &user));
239   PetscCall(VecSet(r, 0));
240   PetscCall(MatDiagonalSet(A, r, INSERT_VALUES));
241   PetscCall(VecDestroy(&r));
242 
243   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
244      Set RHS Jacobian for the adjoint integration
245      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
246   PetscCall(TSSetRHSJacobian(ts, A, A, RHSJacobian, &user));
247   PetscCall(TSSetMaxTime(ts, ftime));
248   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
249   if (monitor) PetscCall(TSMonitorSet(ts, Monitor, &user, NULL));
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 %" PetscInt_FMT ", 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;
279   x_ptr[1] = 0.0;
280   PetscCall(VecRestoreArray(lambda[0], &x_ptr));
281   PetscCall(VecGetArray(lambda[1], &x_ptr));
282   x_ptr[0] = 0.0;
283   x_ptr[1] = 1.0;
284   PetscCall(VecRestoreArray(lambda[1], &x_ptr));
285 
286   PetscCall(MatCreateVecs(Jacp, &mu[0], NULL));
287   PetscCall(MatCreateVecs(Jacp, &mu[1], NULL));
288   PetscCall(VecGetArray(mu[0], &x_ptr));
289   x_ptr[0] = 0.0;
290   PetscCall(VecRestoreArray(mu[0], &x_ptr));
291   PetscCall(VecGetArray(mu[1], &x_ptr));
292   x_ptr[0] = 0.0;
293   PetscCall(VecRestoreArray(mu[1], &x_ptr));
294   PetscCall(TSSetCostGradients(ts, 2, lambda, mu));
295 
296   /*   Set RHS JacobianP */
297   PetscCall(TSSetRHSJacobianP(ts, Jacp, RHSJacobianP, &user));
298 
299   PetscCall(TSAdjointSolve(ts));
300 
301   PetscCall(VecView(lambda[0], PETSC_VIEWER_STDOUT_WORLD));
302   PetscCall(VecView(lambda[1], PETSC_VIEWER_STDOUT_WORLD));
303   PetscCall(VecView(mu[0], PETSC_VIEWER_STDOUT_WORLD));
304   PetscCall(VecView(mu[1], PETSC_VIEWER_STDOUT_WORLD));
305 
306   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
307      Free work space.  All PETSc objects should be destroyed when they
308      are no longer needed.
309    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
310   PetscCall(MatDestroy(&A));
311   PetscCall(MatDestroy(&Jacp));
312   PetscCall(VecDestroy(&x));
313   PetscCall(VecDestroy(&lambda[0]));
314   PetscCall(VecDestroy(&lambda[1]));
315   PetscCall(VecDestroy(&mu[0]));
316   PetscCall(VecDestroy(&mu[1]));
317   PetscCall(TSDestroy(&ts));
318   PetscCall(PetscFree(adctx));
319   PetscCall(PetscFinalize());
320   return 0;
321 }
322 
323 /*TEST
324 
325   build:
326     requires: double !complex adolc
327 
328   test:
329     suffix: 1
330     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
331     output_file: output/ex16adj_1.out
332 
333   test:
334     suffix: 2
335     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor -mu 5
336     output_file: output/ex16adj_2.out
337 
338   test:
339     suffix: 3
340     args: -ts_max_steps 10 -monitor
341     output_file: output/ex16adj_3.out
342 
343   test:
344     suffix: 4
345     args: -ts_max_steps 10 -monitor -mu 5
346     output_file: output/ex16adj_4.out
347 
348 TEST*/
349