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