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