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