xref: /petsc/src/ts/tutorials/advection-diffusion-reaction/ex5opt_ic.c (revision a69119a591a03a9d906b29c0a4e9802e4d7c9795)
1 static char help[] = "Demonstrates adjoint sensitivity analysis for Reaction-Diffusion Equations.\n";
2 
3 /*
4    See ex5.c for details on the equation.
5    This code demonestrates the TSAdjoint/TAO interface to solve an inverse initial value problem built on a system of
6    time-dependent partial differential equations.
7    In this problem, the initial value for the PDE is unknown, but the output (the final solution of the PDE) is known.
8    We want to determine the initial value that can produce the given output.
9    We formulate the problem as a nonlinear optimization problem that minimizes the discrepency between the simulated
10    result and given reference solution, calculate the gradient of the objective function with the discrete adjoint
11    solver, and solve the optimization problem with TAO.
12 
13    Runtime options:
14      -forwardonly  - run only the forward simulation
15      -implicitform - provide IFunction and IJacobian to TS, if not set, RHSFunction and RHSJacobian will be used
16  */
17 
18 #include "reaction_diffusion.h"
19 #include <petscdm.h>
20 #include <petscdmda.h>
21 #include <petsctao.h>
22 
23 /* User-defined routines */
24 extern PetscErrorCode FormFunctionAndGradient(Tao, Vec, PetscReal *, Vec, void *);
25 
26 /*
27    Set terminal condition for the adjoint variable
28  */
29 PetscErrorCode InitializeLambda(DM da, Vec lambda, Vec U, AppCtx *appctx) {
30   char        filename[PETSC_MAX_PATH_LEN] = "";
31   PetscViewer viewer;
32   Vec         Uob;
33 
34   PetscFunctionBegin;
35   PetscCall(VecDuplicate(U, &Uob));
36   PetscCall(PetscSNPrintf(filename, sizeof filename, "ex5opt.ob"));
37   PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, filename, FILE_MODE_READ, &viewer));
38   PetscCall(VecLoad(Uob, viewer));
39   PetscCall(PetscViewerDestroy(&viewer));
40   PetscCall(VecAYPX(Uob, -1., U));
41   PetscCall(VecScale(Uob, 2.0));
42   PetscCall(VecAXPY(lambda, 1., Uob));
43   PetscCall(VecDestroy(&Uob));
44   PetscFunctionReturn(0);
45 }
46 
47 /*
48    Set up a viewer that dumps data into a binary file
49  */
50 PetscErrorCode OutputBIN(DM da, const char *filename, PetscViewer *viewer) {
51   PetscFunctionBegin;
52   PetscCall(PetscViewerCreate(PetscObjectComm((PetscObject)da), viewer));
53   PetscCall(PetscViewerSetType(*viewer, PETSCVIEWERBINARY));
54   PetscCall(PetscViewerFileSetMode(*viewer, FILE_MODE_WRITE));
55   PetscCall(PetscViewerFileSetName(*viewer, filename));
56   PetscFunctionReturn(0);
57 }
58 
59 /*
60    Generate a reference solution and save it to a binary file
61  */
62 PetscErrorCode GenerateOBs(TS ts, Vec U, AppCtx *appctx) {
63   char        filename[PETSC_MAX_PATH_LEN] = "";
64   PetscViewer viewer;
65   DM          da;
66 
67   PetscFunctionBegin;
68   PetscCall(TSGetDM(ts, &da));
69   PetscCall(TSSolve(ts, U));
70   PetscCall(PetscSNPrintf(filename, sizeof filename, "ex5opt.ob"));
71   PetscCall(OutputBIN(da, filename, &viewer));
72   PetscCall(VecView(U, viewer));
73   PetscCall(PetscViewerDestroy(&viewer));
74   PetscFunctionReturn(0);
75 }
76 
77 PetscErrorCode InitialConditions(DM da, Vec U) {
78   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
79   Field   **u;
80   PetscReal hx, hy, x, y;
81 
82   PetscFunctionBegin;
83   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));
84 
85   hx = 2.5 / (PetscReal)Mx;
86   hy = 2.5 / (PetscReal)My;
87 
88   PetscCall(DMDAVecGetArray(da, U, &u));
89   /* Get local grid boundaries */
90   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
91 
92   /* Compute function over the locally owned part of the grid */
93   for (j = ys; j < ys + ym; j++) {
94     y = j * hy;
95     for (i = xs; i < xs + xm; i++) {
96       x = i * hx;
97       if ((1.0 <= x) && (x <= 1.5) && (1.0 <= y) && (y <= 1.5)) u[j][i].v = .25 * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0);
98       else u[j][i].v = 0.0;
99 
100       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
101     }
102   }
103 
104   PetscCall(DMDAVecRestoreArray(da, U, &u));
105   PetscFunctionReturn(0);
106 }
107 
108 PetscErrorCode PerturbedInitialConditions(DM da, Vec U) {
109   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
110   Field   **u;
111   PetscReal hx, hy, x, y;
112 
113   PetscFunctionBegin;
114   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));
115 
116   hx = 2.5 / (PetscReal)Mx;
117   hy = 2.5 / (PetscReal)My;
118 
119   PetscCall(DMDAVecGetArray(da, U, &u));
120   /* Get local grid boundaries */
121   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
122 
123   /* Compute function over the locally owned part of the grid */
124   for (j = ys; j < ys + ym; j++) {
125     y = j * hy;
126     for (i = xs; i < xs + xm; i++) {
127       x = i * hx;
128       if ((1.0 <= x) && (x <= 1.5) && (1.0 <= y) && (y <= 1.5)) u[j][i].v = .25 * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * 0.5 * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * 0.5 * y), 2.0);
129       else u[j][i].v = 0.0;
130 
131       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
132     }
133   }
134 
135   PetscCall(DMDAVecRestoreArray(da, U, &u));
136   PetscFunctionReturn(0);
137 }
138 
139 PetscErrorCode PerturbedInitialConditions2(DM da, Vec U) {
140   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
141   Field   **u;
142   PetscReal hx, hy, x, y;
143 
144   PetscFunctionBegin;
145   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));
146 
147   hx = 2.5 / (PetscReal)Mx;
148   hy = 2.5 / (PetscReal)My;
149 
150   PetscCall(DMDAVecGetArray(da, U, &u));
151   /* Get local grid boundaries */
152   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
153 
154   /* Compute function over the locally owned part of the grid */
155   for (j = ys; j < ys + ym; j++) {
156     y = j * hy;
157     for (i = xs; i < xs + xm; i++) {
158       x = i * hx;
159       if (PetscApproximateGTE(x, 1.0) && PetscApproximateLTE(x, 1.5) && PetscApproximateGTE(y, 1.0) && PetscApproximateLTE(y, 1.5))
160         u[j][i].v = PetscPowReal(PetscSinReal(4.0 * PETSC_PI * (x - 0.5)), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0) / 4.0;
161       else u[j][i].v = 0.0;
162 
163       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
164     }
165   }
166 
167   PetscCall(DMDAVecRestoreArray(da, U, &u));
168   PetscFunctionReturn(0);
169 }
170 
171 PetscErrorCode PerturbedInitialConditions3(DM da, Vec U) {
172   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
173   Field   **u;
174   PetscReal hx, hy, x, y;
175 
176   PetscFunctionBegin;
177   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));
178 
179   hx = 2.5 / (PetscReal)Mx;
180   hy = 2.5 / (PetscReal)My;
181 
182   PetscCall(DMDAVecGetArray(da, U, &u));
183   /* Get local grid boundaries */
184   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));
185 
186   /* Compute function over the locally owned part of the grid */
187   for (j = ys; j < ys + ym; j++) {
188     y = j * hy;
189     for (i = xs; i < xs + xm; i++) {
190       x = i * hx;
191       if ((0.5 <= x) && (x <= 2.0) && (0.5 <= y) && (y <= 2.0)) u[j][i].v = .25 * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0);
192       else u[j][i].v = 0.0;
193 
194       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
195     }
196   }
197 
198   PetscCall(DMDAVecRestoreArray(da, U, &u));
199   PetscFunctionReturn(0);
200 }
201 
202 int main(int argc, char **argv) {
203   DM        da;
204   AppCtx    appctx;
205   PetscBool forwardonly = PETSC_FALSE, implicitform = PETSC_FALSE;
206   PetscInt  perturbic = 1;
207 
208   PetscFunctionBeginUser;
209   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
210   PetscCall(PetscOptionsGetBool(NULL, NULL, "-forwardonly", &forwardonly, NULL));
211   PetscCall(PetscOptionsGetBool(NULL, NULL, "-implicitform", &implicitform, NULL));
212   PetscCall(PetscOptionsGetInt(NULL, NULL, "-perturbic", &perturbic, NULL));
213 
214   appctx.D1    = 8.0e-5;
215   appctx.D2    = 4.0e-5;
216   appctx.gamma = .024;
217   appctx.kappa = .06;
218   appctx.aijpc = PETSC_FALSE;
219 
220   /* Create distributed array (DMDA) to manage parallel grid and vectors */
221   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR, 64, 64, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, NULL, &da));
222   PetscCall(DMSetFromOptions(da));
223   PetscCall(DMSetUp(da));
224   PetscCall(DMDASetFieldName(da, 0, "u"));
225   PetscCall(DMDASetFieldName(da, 1, "v"));
226 
227   /* Extract global vectors from DMDA; then duplicate for remaining
228      vectors that are the same types */
229   PetscCall(DMCreateGlobalVector(da, &appctx.U));
230 
231   /* Create timestepping solver context */
232   PetscCall(TSCreate(PETSC_COMM_WORLD, &appctx.ts));
233   PetscCall(TSSetType(appctx.ts, TSCN));
234   PetscCall(TSSetDM(appctx.ts, da));
235   PetscCall(TSSetProblemType(appctx.ts, TS_NONLINEAR));
236   PetscCall(TSSetEquationType(appctx.ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
237   if (!implicitform) {
238     PetscCall(TSSetRHSFunction(appctx.ts, NULL, RHSFunction, &appctx));
239     PetscCall(TSSetRHSJacobian(appctx.ts, NULL, NULL, RHSJacobian, &appctx));
240   } else {
241     PetscCall(TSSetIFunction(appctx.ts, NULL, IFunction, &appctx));
242     PetscCall(TSSetIJacobian(appctx.ts, NULL, NULL, IJacobian, &appctx));
243   }
244 
245   /* Set initial conditions */
246   PetscCall(InitialConditions(da, appctx.U));
247   PetscCall(TSSetSolution(appctx.ts, appctx.U));
248 
249   /* Set solver options */
250   PetscCall(TSSetMaxTime(appctx.ts, 2000.0));
251   PetscCall(TSSetTimeStep(appctx.ts, 0.5));
252   PetscCall(TSSetExactFinalTime(appctx.ts, TS_EXACTFINALTIME_MATCHSTEP));
253   PetscCall(TSSetFromOptions(appctx.ts));
254 
255   PetscCall(GenerateOBs(appctx.ts, appctx.U, &appctx));
256 
257   if (!forwardonly) {
258     Tao tao;
259     Vec P;
260     Vec lambda[1];
261 #if defined(PETSC_USE_LOG)
262     PetscLogStage opt_stage;
263 #endif
264 
265     PetscCall(PetscLogStageRegister("Optimization", &opt_stage));
266     PetscCall(PetscLogStagePush(opt_stage));
267     if (perturbic == 1) {
268       PetscCall(PerturbedInitialConditions(da, appctx.U));
269     } else if (perturbic == 2) {
270       PetscCall(PerturbedInitialConditions2(da, appctx.U));
271     } else if (perturbic == 3) {
272       PetscCall(PerturbedInitialConditions3(da, appctx.U));
273     }
274 
275     PetscCall(VecDuplicate(appctx.U, &lambda[0]));
276     PetscCall(TSSetCostGradients(appctx.ts, 1, lambda, NULL));
277 
278     /* Have the TS save its trajectory needed by TSAdjointSolve() */
279     PetscCall(TSSetSaveTrajectory(appctx.ts));
280 
281     /* Create TAO solver and set desired solution method */
282     PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
283     PetscCall(TaoSetType(tao, TAOBLMVM));
284 
285     /* Set initial guess for TAO */
286     PetscCall(VecDuplicate(appctx.U, &P));
287     PetscCall(VecCopy(appctx.U, P));
288     PetscCall(TaoSetSolution(tao, P));
289 
290     /* Set routine for function and gradient evaluation */
291     PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionAndGradient, &appctx));
292 
293     /* Check for any TAO command line options */
294     PetscCall(TaoSetFromOptions(tao));
295 
296     PetscCall(TaoSolve(tao));
297     PetscCall(TaoDestroy(&tao));
298     PetscCall(VecDestroy(&lambda[0]));
299     PetscCall(VecDestroy(&P));
300     PetscCall(PetscLogStagePop());
301   }
302 
303   /* Free work space.  All PETSc objects should be destroyed when they
304      are no longer needed. */
305   PetscCall(VecDestroy(&appctx.U));
306   PetscCall(TSDestroy(&appctx.ts));
307   PetscCall(DMDestroy(&da));
308   PetscCall(PetscFinalize());
309   return 0;
310 }
311 
312 /* ------------------------ TAO callbacks ---------------------------- */
313 
314 /*
315    FormFunctionAndGradient - Evaluates the function and corresponding gradient.
316 
317    Input Parameters:
318    tao - the Tao context
319    P   - the input vector
320    ctx - optional user-defined context, as set by TaoSetObjectiveAndGradient()
321 
322    Output Parameters:
323    f   - the newly evaluated function
324    G   - the newly evaluated gradient
325 */
326 PetscErrorCode FormFunctionAndGradient(Tao tao, Vec P, PetscReal *f, Vec G, void *ctx) {
327   AppCtx     *appctx = (AppCtx *)ctx;
328   PetscReal   soberr, timestep;
329   Vec        *lambda;
330   Vec         SDiff;
331   DM          da;
332   char        filename[PETSC_MAX_PATH_LEN] = "";
333   PetscViewer viewer;
334 
335   PetscFunctionBeginUser;
336   PetscCall(TSSetTime(appctx->ts, 0.0));
337   PetscCall(TSGetTimeStep(appctx->ts, &timestep));
338   if (timestep < 0) PetscCall(TSSetTimeStep(appctx->ts, -timestep));
339   PetscCall(TSSetStepNumber(appctx->ts, 0));
340   PetscCall(TSSetFromOptions(appctx->ts));
341 
342   PetscCall(VecDuplicate(P, &SDiff));
343   PetscCall(VecCopy(P, appctx->U));
344   PetscCall(TSGetDM(appctx->ts, &da));
345   *f = 0;
346 
347   PetscCall(TSSolve(appctx->ts, appctx->U));
348   PetscCall(PetscSNPrintf(filename, sizeof filename, "ex5opt.ob"));
349   PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, filename, FILE_MODE_READ, &viewer));
350   PetscCall(VecLoad(SDiff, viewer));
351   PetscCall(PetscViewerDestroy(&viewer));
352   PetscCall(VecAYPX(SDiff, -1., appctx->U));
353   PetscCall(VecDot(SDiff, SDiff, &soberr));
354   *f += soberr;
355 
356   PetscCall(TSGetCostGradients(appctx->ts, NULL, &lambda, NULL));
357   PetscCall(VecSet(lambda[0], 0.0));
358   PetscCall(InitializeLambda(da, lambda[0], appctx->U, appctx));
359   PetscCall(TSAdjointSolve(appctx->ts));
360 
361   PetscCall(VecCopy(lambda[0], G));
362 
363   PetscCall(VecDestroy(&SDiff));
364   PetscFunctionReturn(0);
365 }
366 
367 /*TEST
368 
369    build:
370       depends: reaction_diffusion.c
371       requires: !complex !single
372 
373    test:
374       args: -ts_max_steps 5 -ts_type rk -ts_rk_type 3 -ts_trajectory_type memory -tao_monitor -tao_view -tao_gatol 1e-6
375       output_file: output/ex5opt_ic_1.out
376 
377 TEST*/
378