xref: /petsc/src/ts/tutorials/autodiff/adr_ex5adj_mf.cxx (revision 4e8208cbcbc709572b8abe32f33c78b69c819375)
1 static char help[] = "Demonstrates automatic, matrix-free Jacobian generation using ADOL-C for a time-dependent PDE in 2d, solved using implicit timestepping.\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/ex5 for a description of the problem
11   ------------------------------------------------------------------------- */
12 
13 #include <petscdmda.h>
14 #include <petscts.h>
15 #include "adolc-utils/init.cxx"
16 #include "adolc-utils/matfree.cxx"
17 #include <adolc/adolc.h>
18 
19 /* (Passive) field for the two variables */
20 typedef struct {
21   PetscScalar u, v;
22 } Field;
23 
24 /* Active field for the two variables */
25 typedef struct {
26   adouble u, v;
27 } AField;
28 
29 /* Application context */
30 typedef struct {
31   PetscReal D1, D2, gamma, kappa;
32   AField  **u_a, **f_a;
33   AdolcCtx *adctx; /* Automatic differentation support */
34 } AppCtx;
35 
36 extern PetscErrorCode InitialConditions(DM da, Vec U);
37 extern PetscErrorCode InitializeLambda(DM da, Vec lambda, PetscReal x, PetscReal y);
38 extern PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info, PetscReal t, Field **u, Field **udot, Field **f, void *ptr);
39 extern PetscErrorCode IFunctionActive(TS ts, PetscReal ftime, Vec U, Vec Udot, Vec F, void *ptr);
40 extern PetscErrorCode IJacobianMatFree(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A_shell, Mat B, PetscCtx ctx);
41 
main(int argc,char ** argv)42 int main(int argc, char **argv)
43 {
44   TS          ts;   /* ODE integrator */
45   Vec         x, r; /* solution, residual */
46   DM          da;
47   AppCtx      appctx; /* Application context */
48   AdolcMatCtx matctx; /* Matrix (free) context */
49   Vec         lambda[1];
50   PetscBool   forwardonly = PETSC_FALSE;
51   Mat         A; /* (Matrix free) Jacobian matrix */
52   PetscInt    gxm, gym;
53 
54   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
55      Initialize program
56      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
57   PetscFunctionBeginUser;
58   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
59   PetscCall(PetscOptionsGetBool(NULL, NULL, "-forwardonly", &forwardonly, NULL));
60   appctx.D1    = 8.0e-5;
61   appctx.D2    = 4.0e-5;
62   appctx.gamma = .024;
63   appctx.kappa = .06;
64   PetscCall(PetscLogEventRegister("df/dx forward", MAT_CLASSID, &matctx.event1));
65   PetscCall(PetscLogEventRegister("df/d(xdot) forward", MAT_CLASSID, &matctx.event2));
66   PetscCall(PetscLogEventRegister("df/dx reverse", MAT_CLASSID, &matctx.event3));
67   PetscCall(PetscLogEventRegister("df/d(xdot) reverse", MAT_CLASSID, &matctx.event4));
68 
69   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
70      Create distributed array (DMDA) to manage parallel grid and vectors
71   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
72   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR, 65, 65, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, NULL, &da));
73   PetscCall(DMSetFromOptions(da));
74   PetscCall(DMSetUp(da));
75   PetscCall(DMDASetFieldName(da, 0, "u"));
76   PetscCall(DMDASetFieldName(da, 1, "v"));
77 
78   /*  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
79      Extract global vectors from DMDA; then duplicate for remaining
80      vectors that are the same types
81    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
82   PetscCall(DMCreateGlobalVector(da, &x));
83   PetscCall(VecDuplicate(x, &r));
84 
85   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
86     Create matrix-free context and specify usage of PETSc-ADOL-C drivers
87     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
88   PetscCall(DMSetMatType(da, MATSHELL));
89   PetscCall(DMCreateMatrix(da, &A));
90   PetscCall(MatShellSetContext(A, &matctx));
91   PetscCall(MatShellSetOperation(A, MATOP_MULT, (PetscErrorCodeFn *)PetscAdolcIJacobianVectorProductIDMass));
92   PetscCall(MatShellSetOperation(A, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)PetscAdolcIJacobianTransposeVectorProductIDMass));
93   PetscCall(VecDuplicate(x, &matctx.X));
94   PetscCall(VecDuplicate(x, &matctx.Xdot));
95   PetscCall(DMGetLocalVector(da, &matctx.localX0));
96 
97   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
98      Create timestepping solver context
99      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
100   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
101   PetscCall(TSSetType(ts, TSCN));
102   PetscCall(TSSetDM(ts, da));
103   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
104   PetscCall(DMDATSSetIFunctionLocal(da, INSERT_VALUES, (DMDATSIFunctionLocalFn *)IFunctionLocalPassive, &appctx));
105 
106   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
107     Some data required for matrix-free context
108      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
109   PetscCall(DMDAGetGhostCorners(da, NULL, NULL, NULL, &gxm, &gym, NULL));
110   matctx.m    = 2 * gxm * gym;
111   matctx.n    = 2 * gxm * gym; /* Number of dependent and independent variables */
112   matctx.flg  = PETSC_FALSE;   /* Flag for reverse mode */
113   matctx.tag1 = 1;             /* Tape identifier */
114 
115   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
116      Trace function just once
117    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
118   PetscCall(PetscNew(&appctx.adctx));
119   PetscCall(IFunctionActive(ts, 1., x, matctx.Xdot, r, &appctx));
120   PetscCall(PetscFree(appctx.adctx));
121 
122   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
123      Set Jacobian. In this case, IJacobian simply acts to pass context
124      information to the matrix-free Jacobian vector product.
125    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
126   PetscCall(TSSetIJacobian(ts, A, A, IJacobianMatFree, &appctx));
127 
128   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
129      Set initial conditions
130    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
131   PetscCall(InitialConditions(da, x));
132   PetscCall(TSSetSolution(ts, x));
133 
134   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
135     Have the TS save its trajectory so that TSAdjointSolve() may be used
136     and set solver options
137    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
138   if (!forwardonly) {
139     PetscCall(TSSetSaveTrajectory(ts));
140     PetscCall(TSSetMaxTime(ts, 200.0));
141     PetscCall(TSSetTimeStep(ts, 0.5));
142   } else {
143     PetscCall(TSSetMaxTime(ts, 2000.0));
144     PetscCall(TSSetTimeStep(ts, 10));
145   }
146   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
147   PetscCall(TSSetFromOptions(ts));
148 
149   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
150      Solve ODE system
151      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
152   PetscCall(TSSolve(ts, x));
153   if (!forwardonly) {
154     /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
155        Start the Adjoint model
156        - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
157     PetscCall(VecDuplicate(x, &lambda[0]));
158     /*   Reset initial conditions for the adjoint integration */
159     PetscCall(InitializeLambda(da, lambda[0], 0.5, 0.5));
160     PetscCall(TSSetCostGradients(ts, 1, lambda, NULL));
161     PetscCall(TSAdjointSolve(ts));
162     PetscCall(VecDestroy(&lambda[0]));
163   }
164 
165   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
166      Free work space.  All PETSc objects should be destroyed when they
167      are no longer needed.
168    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
169   PetscCall(DMRestoreLocalVector(da, &matctx.localX0));
170   PetscCall(VecDestroy(&r));
171   PetscCall(VecDestroy(&matctx.X));
172   PetscCall(VecDestroy(&matctx.Xdot));
173   PetscCall(MatDestroy(&A));
174   PetscCall(VecDestroy(&x));
175   PetscCall(TSDestroy(&ts));
176   PetscCall(DMDestroy(&da));
177 
178   PetscCall(PetscFinalize());
179   return 0;
180 }
181 
InitialConditions(DM da,Vec U)182 PetscErrorCode InitialConditions(DM da, Vec U)
183 {
184   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
185   Field   **u;
186   PetscReal hx, hy, x, y;
187 
188   PetscFunctionBegin;
189   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
190 
191   hx = 2.5 / (PetscReal)Mx;
192   hy = 2.5 / (PetscReal)My;
193 
194   /*
195      Get pointers to vector data
196   */
197   PetscCall(DMDAVecGetArray(da, U, &u));
198 
199   /*
200      Get local grid boundaries
201   */
202   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
203 
204   /*
205      Compute function over the locally owned part of the grid
206   */
207   for (j = ys; j < ys + ym; j++) {
208     y = j * hy;
209     for (i = xs; i < xs + xm; i++) {
210       x = i * hx;
211       if (PetscApproximateGTE(x, 1.0) && PetscApproximateLTE(x, 1.5) && PetscApproximateGTE(y, 1.0) && PetscApproximateLTE(y, 1.5))
212         u[j][i].v = PetscPowReal(PetscSinReal(4.0 * PETSC_PI * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0) / 4.0;
213       else u[j][i].v = 0.0;
214 
215       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
216     }
217   }
218 
219   /*
220      Restore vectors
221   */
222   PetscCall(DMDAVecRestoreArray(da, U, &u));
223   PetscFunctionReturn(PETSC_SUCCESS);
224 }
225 
InitializeLambda(DM da,Vec lambda,PetscReal x,PetscReal y)226 PetscErrorCode InitializeLambda(DM da, Vec lambda, PetscReal x, PetscReal y)
227 {
228   PetscInt i, j, Mx, My, xs, ys, xm, ym;
229   Field  **l;
230 
231   PetscFunctionBegin;
232   PetscCall(DMDAGetInfo(da, PETSC_IGNORE, &Mx, &My, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE, PETSC_IGNORE));
233   /* locate the global i index for x and j index for y */
234   i = (PetscInt)(x * (Mx - 1));
235   j = (PetscInt)(y * (My - 1));
236   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
237 
238   if (xs <= i && i < xs + xm && ys <= j && j < ys + ym) {
239     /* the i,j vertex is on this process */
240     PetscCall(DMDAVecGetArray(da, lambda, &l));
241     l[j][i].u = 1.0;
242     l[j][i].v = 1.0;
243     PetscCall(DMDAVecRestoreArray(da, lambda, &l));
244   }
245   PetscFunctionReturn(PETSC_SUCCESS);
246 }
247 
IFunctionLocalPassive(DMDALocalInfo * info,PetscReal t,Field ** u,Field ** udot,Field ** f,void * ptr)248 PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info, PetscReal t, Field **u, Field **udot, Field **f, void *ptr)
249 {
250   AppCtx     *appctx = (AppCtx *)ptr;
251   PetscInt    i, j, xs, ys, xm, ym;
252   PetscReal   hx, hy, sx, sy;
253   PetscScalar uc, uxx, uyy, vc, vxx, vyy;
254 
255   PetscFunctionBegin;
256   hx = 2.50 / (PetscReal)info->mx;
257   sx = 1.0 / (hx * hx);
258   hy = 2.50 / (PetscReal)info->my;
259   sy = 1.0 / (hy * hy);
260 
261   /* Get local grid boundaries */
262   xs = info->xs;
263   xm = info->xm;
264   ys = info->ys;
265   ym = info->ym;
266 
267   /* Compute function over the locally owned part of the grid */
268   for (j = ys; j < ys + ym; j++) {
269     for (i = xs; i < xs + xm; i++) {
270       uc        = u[j][i].u;
271       uxx       = (-2.0 * uc + u[j][i - 1].u + u[j][i + 1].u) * sx;
272       uyy       = (-2.0 * uc + u[j - 1][i].u + u[j + 1][i].u) * sy;
273       vc        = u[j][i].v;
274       vxx       = (-2.0 * vc + u[j][i - 1].v + u[j][i + 1].v) * sx;
275       vyy       = (-2.0 * vc + u[j - 1][i].v + u[j + 1][i].v) * sy;
276       f[j][i].u = udot[j][i].u - appctx->D1 * (uxx + uyy) + uc * vc * vc - appctx->gamma * (1.0 - uc);
277       f[j][i].v = udot[j][i].v - appctx->D2 * (vxx + vyy) - uc * vc * vc + (appctx->gamma + appctx->kappa) * vc;
278     }
279   }
280   PetscCall(PetscLogFlops(16.0 * xm * ym));
281   PetscFunctionReturn(PETSC_SUCCESS);
282 }
283 
IFunctionActive(TS ts,PetscReal ftime,Vec U,Vec Udot,Vec F,void * ptr)284 PetscErrorCode IFunctionActive(TS ts, PetscReal ftime, Vec U, Vec Udot, Vec F, void *ptr)
285 {
286   AppCtx       *appctx = (AppCtx *)ptr;
287   DM            da;
288   DMDALocalInfo info;
289   Field       **u, **f, **udot;
290   Vec           localU;
291   PetscInt      i, j, xs, ys, xm, ym, gxs, gys, gxm, gym;
292   PetscReal     hx, hy, sx, sy;
293   adouble       uc, uxx, uyy, vc, vxx, vyy;
294   AField      **f_a, *f_c, **u_a, *u_c;
295   PetscScalar   dummy;
296 
297   PetscFunctionBegin;
298   PetscCall(TSGetDM(ts, &da));
299   PetscCall(DMDAGetLocalInfo(da, &info));
300   PetscCall(DMGetLocalVector(da, &localU));
301   hx  = 2.50 / (PetscReal)info.mx;
302   sx  = 1.0 / (hx * hx);
303   hy  = 2.50 / (PetscReal)info.my;
304   sy  = 1.0 / (hy * hy);
305   xs  = info.xs;
306   xm  = info.xm;
307   gxs = info.gxs;
308   gxm = info.gxm;
309   ys  = info.ys;
310   ym  = info.ym;
311   gys = info.gys;
312   gym = info.gym;
313 
314   /*
315      Scatter ghost points to local vector,using the 2-step process
316         DMGlobalToLocalBegin(),DMGlobalToLocalEnd().
317      By placing code between these two statements, computations can be
318      done while messages are in transition.
319   */
320   PetscCall(DMGlobalToLocalBegin(da, U, INSERT_VALUES, localU));
321   PetscCall(DMGlobalToLocalEnd(da, U, INSERT_VALUES, localU));
322 
323   /*
324      Get pointers to vector data
325   */
326   PetscCall(DMDAVecGetArrayRead(da, localU, &u));
327   PetscCall(DMDAVecGetArray(da, F, &f));
328   PetscCall(DMDAVecGetArrayRead(da, Udot, &udot));
329 
330   /*
331     Create contiguous 1-arrays of AFields
332 
333     NOTE: Memory for ADOL-C active variables (such as adouble and AField)
334           cannot be allocated using PetscMalloc, as this does not call the
335           relevant class constructor. Instead, we use the C++ keyword `new`.
336   */
337   u_c = new AField[info.gxm * info.gym];
338   f_c = new AField[info.gxm * info.gym];
339 
340   /* Create corresponding 2-arrays of AFields */
341   u_a = new AField *[info.gym];
342   f_a = new AField *[info.gym];
343 
344   /* Align indices between array types to endow 2d array with ghost points */
345   PetscCall(GiveGhostPoints(da, u_c, &u_a));
346   PetscCall(GiveGhostPoints(da, f_c, &f_a));
347 
348   trace_on(1); /* Start of active section on tape 1 */
349 
350   /*
351     Mark independence
352 
353     NOTE: Ghost points are marked as independent, in place of the points they represent on
354           other processors / on other boundaries.
355   */
356   for (j = gys; j < gys + gym; j++) {
357     for (i = gxs; i < gxs + gxm; i++) {
358       u_a[j][i].u <<= u[j][i].u;
359       u_a[j][i].v <<= u[j][i].v;
360     }
361   }
362 
363   /* Compute function over the locally owned part of the grid */
364   for (j = ys; j < ys + ym; j++) {
365     for (i = xs; i < xs + xm; i++) {
366       uc          = u_a[j][i].u;
367       uxx         = (-2.0 * uc + u_a[j][i - 1].u + u_a[j][i + 1].u) * sx;
368       uyy         = (-2.0 * uc + u_a[j - 1][i].u + u_a[j + 1][i].u) * sy;
369       vc          = u_a[j][i].v;
370       vxx         = (-2.0 * vc + u_a[j][i - 1].v + u_a[j][i + 1].v) * sx;
371       vyy         = (-2.0 * vc + u_a[j - 1][i].v + u_a[j + 1][i].v) * sy;
372       f_a[j][i].u = udot[j][i].u - appctx->D1 * (uxx + uyy) + uc * vc * vc - appctx->gamma * (1.0 - uc);
373       f_a[j][i].v = udot[j][i].v - appctx->D2 * (vxx + vyy) - uc * vc * vc + (appctx->gamma + appctx->kappa) * vc;
374     }
375   }
376 
377   /*
378     Mark dependence
379 
380     NOTE: Marking dependence of dummy variables makes the index notation much simpler when forming
381           the Jacobian later.
382   */
383   for (j = gys; j < gys + gym; j++) {
384     for (i = gxs; i < gxs + gxm; i++) {
385       if ((i < xs) || (i >= xs + xm) || (j < ys) || (j >= ys + ym)) {
386         f_a[j][i].u >>= dummy;
387         f_a[j][i].v >>= dummy;
388       } else {
389         f_a[j][i].u >>= f[j][i].u;
390         f_a[j][i].v >>= f[j][i].v;
391       }
392     }
393   }
394   trace_off(); /* End of active section */
395   PetscCall(PetscLogFlops(16.0 * xm * ym));
396 
397   /* Restore vectors */
398   PetscCall(DMDAVecRestoreArray(da, F, &f));
399   PetscCall(DMDAVecRestoreArrayRead(da, localU, &u));
400   PetscCall(DMDAVecRestoreArrayRead(da, Udot, &udot));
401 
402   PetscCall(DMRestoreLocalVector(da, &localU));
403 
404   /* Destroy AFields appropriately */
405   f_a += info.gys;
406   u_a += info.gys;
407   delete[] f_a;
408   delete[] u_a;
409   delete[] f_c;
410   delete[] u_c;
411   PetscFunctionReturn(PETSC_SUCCESS);
412 }
413 
414 /*
415   Simply acts to pass TS information to the AdolcMatCtx
416 */
IJacobianMatFree(TS ts,PetscReal t,Vec X,Vec Xdot,PetscReal a,Mat A_shell,Mat B,PetscCtx ctx)417 PetscErrorCode IJacobianMatFree(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A_shell, Mat B, PetscCtx ctx)
418 {
419   AdolcMatCtx *mctx;
420   DM           da;
421 
422   PetscFunctionBeginUser;
423   PetscCall(MatShellGetContext(A_shell, &mctx));
424 
425   mctx->time  = t;
426   mctx->shift = a;
427   if (mctx->ts != ts) mctx->ts = ts;
428   PetscCall(VecCopy(X, mctx->X));
429   PetscCall(VecCopy(Xdot, mctx->Xdot));
430   PetscCall(TSGetDM(ts, &da));
431   PetscCall(DMGlobalToLocalBegin(da, mctx->X, INSERT_VALUES, mctx->localX0));
432   PetscCall(DMGlobalToLocalEnd(da, mctx->X, INSERT_VALUES, mctx->localX0));
433   PetscFunctionReturn(PETSC_SUCCESS);
434 }
435 
436 /*TEST
437 
438   build:
439     requires: double !complex adolc
440 
441   test:
442     suffix: 1
443     args: -ts_max_steps 1 -da_grid_x 12 -da_grid_y 12 -snes_test_jacobian
444     output_file: output/adr_ex5adj_mf_1.out
445 
446   test:
447     suffix: 2
448     nsize: 4
449     args: -ts_max_steps 10 -da_grid_x 12 -da_grid_y 12 -ts_monitor -ts_adjoint_monitor
450     output_file: output/adr_ex5adj_mf_2.out
451 
452 TEST*/
453