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