xref: /petsc/src/ts/tutorials/autodiff/adr_ex1.cxx (revision daa037dfd3c3bec8dc8659548d2b20b07c1dc6de)
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]; u_a[1] <<= u[1]; u_a[2] <<= u[2]; /* Mark independence */
79   f_a[0] = udot[0] + ctx->k*u_a[0]*u_a[1];
80   f_a[1] = udot[1] + ctx->k*u_a[0]*u_a[1];
81   f_a[2] = udot[2] - ctx->k*u_a[0]*u_a[1];
82   f_a[0] >>= f[0]; f_a[1] >>= f[1]; 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 {
97   PetscScalar       *f;
98   const PetscScalar *u,*udot;
99 
100   adouble           f_a[3];    /* 'active' double for dependent variables */
101   adouble           udot_a[3]; /* 'active' double for independent variables */
102 
103   PetscFunctionBegin;
104   /*  The next three lines allow us to access the entries of the vectors directly */
105   PetscCall(VecGetArrayRead(U,&u));
106   PetscCall(VecGetArrayRead(Udot,&udot));
107   PetscCall(VecGetArray(F,&f));
108 
109   /* Start of active section */
110   trace_on(2);
111   udot_a[0] <<= udot[0]; udot_a[1] <<= udot[1]; udot_a[2] <<= udot[2]; /* Mark independence */
112   f_a[0] = udot_a[0] + ctx->k*u[0]*u[1];
113   f_a[1] = udot_a[1] + ctx->k*u[0]*u[1];
114   f_a[2] = udot_a[2] - ctx->k*u[0]*u[1];
115   f_a[0] >>= f[0]; f_a[1] >>= f[1]; f_a[2] >>= f[2];                   /* Mark dependence */
116   trace_off();
117   /* End of active section */
118 
119   PetscCall(VecRestoreArray(F,&f));
120   PetscCall(VecRestoreArrayRead(Udot,&udot));
121   PetscCall(VecRestoreArrayRead(U,&u));
122   PetscFunctionReturn(0);
123 }
124 
125 /*
126  Defines the Jacobian of the ODE passed to the ODE solver, using the PETSc-ADOL-C driver for
127  implicit TS.
128 */
129 PetscErrorCode IJacobian(TS ts,PetscReal t,Vec U,Vec Udot,PetscReal a,Mat A,Mat B,AppCtx *ctx)
130 {
131   AppCtx            *appctx = (AppCtx*)ctx;
132   const PetscScalar *u;
133 
134   PetscFunctionBegin;
135   PetscCall(VecGetArrayRead(U,&u));
136   PetscCall(PetscAdolcComputeIJacobian(1,2,A,u,a,appctx->adctx));
137   PetscCall(VecRestoreArrayRead(U,&u));
138   PetscFunctionReturn(0);
139 }
140 
141 /*
142      Defines the exact (analytic) solution to the ODE
143 */
144 static PetscErrorCode Solution(TS ts,PetscReal t,Vec U,AppCtx *ctx)
145 {
146   const PetscScalar *uinit;
147   PetscScalar       *u,d0,q;
148 
149   PetscFunctionBegin;
150   PetscCall(VecGetArrayRead(ctx->initialsolution,&uinit));
151   PetscCall(VecGetArray(U,&u));
152   d0   = uinit[0] - uinit[1];
153   if (d0 == 0.0) q = ctx->k*t;
154   else q = (1.0 - PetscExpScalar(-ctx->k*t*d0))/d0;
155   u[0] = uinit[0]/(1.0 + uinit[1]*q);
156   u[1] = u[0] - d0;
157   u[2] = uinit[1] + uinit[2] - u[1];
158   PetscCall(VecRestoreArray(U,&u));
159   PetscCall(VecRestoreArrayRead(ctx->initialsolution,&uinit));
160   PetscFunctionReturn(0);
161 }
162 
163 int main(int argc,char **argv)
164 {
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   PetscCall(PetscInitialize(&argc,&argv,NULL,help));
179   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size));
180   PetscCheckFalse(size > 1,PETSC_COMM_WORLD,PETSC_ERR_SUP,"Only for sequential runs");
181   PetscCall(PetscNew(&adctx));
182   adctx->m = n;adctx->n = n;adctx->p = n;
183   ctx.adctx = adctx;
184 
185   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
186     Create necessary matrix and vectors
187     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
188   PetscCall(MatCreate(PETSC_COMM_WORLD,&A));
189   PetscCall(MatSetSizes(A,n,n,PETSC_DETERMINE,PETSC_DETERMINE));
190   PetscCall(MatSetFromOptions(A));
191   PetscCall(MatSetUp(A));
192 
193   PetscCall(MatCreateVecs(A,&U,NULL));
194 
195   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
196     Set runtime options
197     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
198   ctx.k = .9;
199   PetscCall(PetscOptionsGetScalar(NULL,NULL,"-k",&ctx.k,NULL));
200   PetscCall(VecDuplicate(U,&ctx.initialsolution));
201   PetscCall(VecGetArray(ctx.initialsolution,&u));
202   u[0]  = 1;
203   u[1]  = .7;
204   u[2]  = 0;
205   PetscCall(VecRestoreArray(ctx.initialsolution,&u));
206   PetscCall(PetscOptionsGetVec(NULL,NULL,"-initial",ctx.initialsolution,NULL));
207 
208   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
209      Create timestepping solver context
210      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
211   PetscCall(TSCreate(PETSC_COMM_WORLD,&ts));
212   PetscCall(TSSetProblemType(ts,TS_NONLINEAR));
213   PetscCall(TSSetType(ts,TSROSW));
214   PetscCall(TSSetIFunction(ts,NULL,(TSIFunction) IFunctionPassive,&ctx));
215 
216   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
217      Set initial conditions
218    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
219   PetscCall(Solution(ts,0,U,&ctx));
220   PetscCall(TSSetSolution(ts,U));
221 
222   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
223      Trace just once for each tape
224      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
225   PetscCall(VecDuplicate(U,&Udot));
226   PetscCall(VecDuplicate(U,&R));
227   PetscCall(IFunctionActive1(ts,0.,U,Udot,R,&ctx));
228   PetscCall(IFunctionActive2(ts,0.,U,Udot,R,&ctx));
229   PetscCall(VecDestroy(&R));
230   PetscCall(VecDestroy(&Udot));
231 
232   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
233      Set Jacobian
234      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
235   PetscCall(TSSetIJacobian(ts,A,A,(TSIJacobian)IJacobian,&ctx));
236   PetscCall(TSSetSolutionFunction(ts,(TSSolutionFunction)Solution,&ctx));
237 
238   {
239     DM   dm;
240     void *ptr;
241     PetscCall(TSGetDM(ts,&dm));
242     PetscCall(PetscDLSym(NULL,"IFunctionView",&ptr));
243     PetscCall(PetscDLSym(NULL,"IFunctionLoad",&ptr));
244     PetscCall(DMTSSetIFunctionSerialize(dm,(PetscErrorCode (*)(void*,PetscViewer))IFunctionView,(PetscErrorCode (*)(void**,PetscViewer))IFunctionLoad));
245     PetscCall(DMTSSetIJacobianSerialize(dm,(PetscErrorCode (*)(void*,PetscViewer))IFunctionView,(PetscErrorCode (*)(void**,PetscViewer))IFunctionLoad));
246   }
247 
248   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
249      Set solver options
250    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
251   PetscCall(TSSetTimeStep(ts,.001));
252   PetscCall(TSSetMaxSteps(ts,1000));
253   PetscCall(TSSetMaxTime(ts,20.0));
254   PetscCall(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_STEPOVER));
255   PetscCall(TSSetFromOptions(ts));
256   PetscCall(TSMonitorLGSetVariableNames(ts,names));
257 
258   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
259      Solve nonlinear system
260      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
261   PetscCall(TSSolve(ts,U));
262 
263   PetscCall(TSView(ts,PETSC_VIEWER_BINARY_WORLD));
264 
265   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
266      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
267    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
268   PetscCall(VecDestroy(&ctx.initialsolution));
269   PetscCall(MatDestroy(&A));
270   PetscCall(VecDestroy(&U));
271   PetscCall(TSDestroy(&ts));
272   PetscCall(PetscFree(adctx));
273   PetscCall(PetscFinalize());
274   return 0;
275 }
276 
277 /*TEST
278 
279   build:
280     requires: double !complex adolc
281 
282   test:
283     suffix: 1
284     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
285     output_file: output/adr_ex1_1.out
286 
287   test:
288     suffix: 2
289     args: -ts_max_steps 1 -snes_test_jacobian
290     output_file: output/adr_ex1_2.out
291 
292 TEST*/
293