xref: /petsc/src/ts/tutorials/autodiff/adr_ex1.cxx (revision d5b43468fb8780a8feea140ccd6fa3e6a50411cc)
1 static char help[] = "Demonstrates automatic Jacobian generation using ADOL-C for a nonlinear reaction problem from chemistry.\n";
2 
3 /*
4    REQUIRES configuration of PETSc with option --download-adolc.
5 
6    For documentation on ADOL-C, see
7      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
8 */
9 /* ------------------------------------------------------------------------
10   See ../advection-diffusion-reaction/ex1 for a description of the problem
11   ------------------------------------------------------------------------- */
12 #include <petscts.h>
13 #include "adolc-utils/drivers.cxx"
14 #include <adolc/adolc.h>
15 
16 typedef struct {
17   PetscScalar k;
18   Vec         initialsolution;
19   AdolcCtx   *adctx; /* Automatic differentiation support */
20 } AppCtx;
21 
22 PetscErrorCode IFunctionView(AppCtx *ctx, PetscViewer v)
23 {
24   PetscFunctionBegin;
25   PetscCall(PetscViewerBinaryWrite(v, &ctx->k, 1, PETSC_SCALAR));
26   PetscFunctionReturn(0);
27 }
28 
29 PetscErrorCode IFunctionLoad(AppCtx **ctx, PetscViewer v)
30 {
31   PetscFunctionBegin;
32   PetscCall(PetscNew(ctx));
33   PetscCall(PetscViewerBinaryRead(v, &(*ctx)->k, 1, NULL, PETSC_SCALAR));
34   PetscFunctionReturn(0);
35 }
36 
37 /*
38   Defines the ODE passed to the ODE solver
39 */
40 PetscErrorCode IFunctionPassive(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, AppCtx *ctx)
41 {
42   PetscScalar       *f;
43   const PetscScalar *u, *udot;
44 
45   PetscFunctionBegin;
46   /*  The next three lines allow us to access the entries of the vectors directly */
47   PetscCall(VecGetArrayRead(U, &u));
48   PetscCall(VecGetArrayRead(Udot, &udot));
49   PetscCall(VecGetArray(F, &f));
50   f[0] = udot[0] + ctx->k * u[0] * u[1];
51   f[1] = udot[1] + ctx->k * u[0] * u[1];
52   f[2] = udot[2] - ctx->k * u[0] * u[1];
53   PetscCall(VecRestoreArray(F, &f));
54   PetscCall(VecRestoreArrayRead(Udot, &udot));
55   PetscCall(VecRestoreArrayRead(U, &u));
56   PetscFunctionReturn(0);
57 }
58 
59 /*
60   'Active' ADOL-C annotated version, marking dependence upon u.
61 */
62 PetscErrorCode IFunctionActive1(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, AppCtx *ctx)
63 {
64   PetscScalar       *f;
65   const PetscScalar *u, *udot;
66 
67   adouble f_a[3]; /* 'active' double for dependent variables */
68   adouble u_a[3]; /* 'active' double for independent variables */
69 
70   PetscFunctionBegin;
71   /*  The next three lines allow us to access the entries of the vectors directly */
72   PetscCall(VecGetArrayRead(U, &u));
73   PetscCall(VecGetArrayRead(Udot, &udot));
74   PetscCall(VecGetArray(F, &f));
75 
76   /* Start of active section */
77   trace_on(1);
78   u_a[0] <<= u[0];
79   u_a[1] <<= u[1];
80   u_a[2] <<= u[2]; /* Mark independence */
81   f_a[0] = udot[0] + ctx->k * u_a[0] * u_a[1];
82   f_a[1] = udot[1] + ctx->k * u_a[0] * u_a[1];
83   f_a[2] = udot[2] - ctx->k * u_a[0] * u_a[1];
84   f_a[0] >>= f[0];
85   f_a[1] >>= f[1];
86   f_a[2] >>= f[2]; /* Mark dependence */
87   trace_off();
88   /* End of active section */
89 
90   PetscCall(VecRestoreArray(F, &f));
91   PetscCall(VecRestoreArrayRead(Udot, &udot));
92   PetscCall(VecRestoreArrayRead(U, &u));
93   PetscFunctionReturn(0);
94 }
95 
96 /*
97   'Active' ADOL-C annotated version, marking dependence upon udot.
98 */
99 PetscErrorCode IFunctionActive2(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, AppCtx *ctx)
100 {
101   PetscScalar       *f;
102   const PetscScalar *u, *udot;
103 
104   adouble f_a[3];    /* 'active' double for dependent variables */
105   adouble udot_a[3]; /* 'active' double for independent variables */
106 
107   PetscFunctionBegin;
108   /*  The next three lines allow us to access the entries of the vectors directly */
109   PetscCall(VecGetArrayRead(U, &u));
110   PetscCall(VecGetArrayRead(Udot, &udot));
111   PetscCall(VecGetArray(F, &f));
112 
113   /* Start of active section */
114   trace_on(2);
115   udot_a[0] <<= udot[0];
116   udot_a[1] <<= udot[1];
117   udot_a[2] <<= udot[2]; /* Mark independence */
118   f_a[0] = udot_a[0] + ctx->k * u[0] * u[1];
119   f_a[1] = udot_a[1] + ctx->k * u[0] * u[1];
120   f_a[2] = udot_a[2] - ctx->k * u[0] * u[1];
121   f_a[0] >>= f[0];
122   f_a[1] >>= f[1];
123   f_a[2] >>= f[2]; /* Mark dependence */
124   trace_off();
125   /* End of active section */
126 
127   PetscCall(VecRestoreArray(F, &f));
128   PetscCall(VecRestoreArrayRead(Udot, &udot));
129   PetscCall(VecRestoreArrayRead(U, &u));
130   PetscFunctionReturn(0);
131 }
132 
133 /*
134  Defines the Jacobian of the ODE passed to the ODE solver, using the PETSc-ADOL-C driver for
135  implicit TS.
136 */
137 PetscErrorCode IJacobian(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal a, Mat A, Mat B, AppCtx *ctx)
138 {
139   AppCtx            *appctx = (AppCtx *)ctx;
140   const PetscScalar *u;
141 
142   PetscFunctionBegin;
143   PetscCall(VecGetArrayRead(U, &u));
144   PetscCall(PetscAdolcComputeIJacobian(1, 2, A, u, a, appctx->adctx));
145   PetscCall(VecRestoreArrayRead(U, &u));
146   PetscFunctionReturn(0);
147 }
148 
149 /*
150      Defines the exact (analytic) solution to the ODE
151 */
152 static PetscErrorCode Solution(TS ts, PetscReal t, Vec U, AppCtx *ctx)
153 {
154   const PetscScalar *uinit;
155   PetscScalar       *u, d0, q;
156 
157   PetscFunctionBegin;
158   PetscCall(VecGetArrayRead(ctx->initialsolution, &uinit));
159   PetscCall(VecGetArray(U, &u));
160   d0 = uinit[0] - uinit[1];
161   if (d0 == 0.0) q = ctx->k * t;
162   else q = (1.0 - PetscExpScalar(-ctx->k * t * d0)) / d0;
163   u[0] = uinit[0] / (1.0 + uinit[1] * q);
164   u[1] = u[0] - d0;
165   u[2] = uinit[1] + uinit[2] - u[1];
166   PetscCall(VecRestoreArray(U, &u));
167   PetscCall(VecRestoreArrayRead(ctx->initialsolution, &uinit));
168   PetscFunctionReturn(0);
169 }
170 
171 int main(int argc, char **argv)
172 {
173   TS                ts;         /* ODE integrator */
174   Vec               U, Udot, R; /* solution, derivative, residual */
175   Mat               A;          /* Jacobian matrix */
176   PetscMPIInt       size;
177   PetscInt          n = 3;
178   AppCtx            ctx;
179   AdolcCtx         *adctx;
180   PetscScalar      *u;
181   const char *const names[] = {"U1", "U2", "U3", NULL};
182 
183   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
184      Initialize program
185      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
186   PetscFunctionBeginUser;
187   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
188   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
189   PetscCheck(size <= 1, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Only for sequential runs");
190   PetscCall(PetscNew(&adctx));
191   adctx->m  = n;
192   adctx->n  = n;
193   adctx->p  = n;
194   ctx.adctx = adctx;
195 
196   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
197     Create necessary matrix and vectors
198     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
199   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
200   PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
201   PetscCall(MatSetFromOptions(A));
202   PetscCall(MatSetUp(A));
203 
204   PetscCall(MatCreateVecs(A, &U, NULL));
205 
206   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
207     Set runtime options
208     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
209   ctx.k = .9;
210   PetscCall(PetscOptionsGetScalar(NULL, NULL, "-k", &ctx.k, NULL));
211   PetscCall(VecDuplicate(U, &ctx.initialsolution));
212   PetscCall(VecGetArray(ctx.initialsolution, &u));
213   u[0] = 1;
214   u[1] = .7;
215   u[2] = 0;
216   PetscCall(VecRestoreArray(ctx.initialsolution, &u));
217   PetscCall(PetscOptionsGetVec(NULL, NULL, "-initial", ctx.initialsolution, NULL));
218 
219   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
220      Create timestepping solver context
221      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
222   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
223   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
224   PetscCall(TSSetType(ts, TSROSW));
225   PetscCall(TSSetIFunction(ts, NULL, (TSIFunction)IFunctionPassive, &ctx));
226 
227   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
228      Set initial conditions
229    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
230   PetscCall(Solution(ts, 0, U, &ctx));
231   PetscCall(TSSetSolution(ts, U));
232 
233   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
234      Trace just once for each tape
235      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
236   PetscCall(VecDuplicate(U, &Udot));
237   PetscCall(VecDuplicate(U, &R));
238   PetscCall(IFunctionActive1(ts, 0., U, Udot, R, &ctx));
239   PetscCall(IFunctionActive2(ts, 0., U, Udot, R, &ctx));
240   PetscCall(VecDestroy(&R));
241   PetscCall(VecDestroy(&Udot));
242 
243   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
244      Set Jacobian
245      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
246   PetscCall(TSSetIJacobian(ts, A, A, (TSIJacobian)IJacobian, &ctx));
247   PetscCall(TSSetSolutionFunction(ts, (TSSolutionFunction)Solution, &ctx));
248 
249   {
250     DM    dm;
251     void *ptr;
252     PetscCall(TSGetDM(ts, &dm));
253     PetscCall(PetscDLSym(NULL, "IFunctionView", &ptr));
254     PetscCall(PetscDLSym(NULL, "IFunctionLoad", &ptr));
255     PetscCall(DMTSSetIFunctionSerialize(dm, (PetscErrorCode(*)(void *, PetscViewer))IFunctionView, (PetscErrorCode(*)(void **, PetscViewer))IFunctionLoad));
256     PetscCall(DMTSSetIJacobianSerialize(dm, (PetscErrorCode(*)(void *, PetscViewer))IFunctionView, (PetscErrorCode(*)(void **, PetscViewer))IFunctionLoad));
257   }
258 
259   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
260      Set solver options
261    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
262   PetscCall(TSSetTimeStep(ts, .001));
263   PetscCall(TSSetMaxSteps(ts, 1000));
264   PetscCall(TSSetMaxTime(ts, 20.0));
265   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
266   PetscCall(TSSetFromOptions(ts));
267   PetscCall(TSMonitorLGSetVariableNames(ts, names));
268 
269   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
270      Solve nonlinear system
271      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
272   PetscCall(TSSolve(ts, U));
273 
274   PetscCall(TSView(ts, PETSC_VIEWER_BINARY_WORLD));
275 
276   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
277      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
278    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
279   PetscCall(VecDestroy(&ctx.initialsolution));
280   PetscCall(MatDestroy(&A));
281   PetscCall(VecDestroy(&U));
282   PetscCall(TSDestroy(&ts));
283   PetscCall(PetscFree(adctx));
284   PetscCall(PetscFinalize());
285   return 0;
286 }
287 
288 /*TEST
289 
290   build:
291     requires: double !complex adolc
292 
293   test:
294     suffix: 1
295     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
296     output_file: output/adr_ex1_1.out
297 
298   test:
299     suffix: 2
300     args: -ts_max_steps 1 -snes_test_jacobian
301     output_file: output/adr_ex1_2.out
302 
303 TEST*/
304