xref: /petsc/src/ts/tutorials/advection-diffusion-reaction/ex5opt_ic.c (revision ebead697dbf761eb322f829370bbe90b3bd93fa3)
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 {
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(0);
46 }
47 
48 /*
49    Set up a viewer that dumps data into a binary file
50  */
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(0);
59 }
60 
61 /*
62    Generate a reference solution and save it to a binary file
63  */
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(0);
78 }
79 
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(0);
110 }
111 
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(0);
142 }
143 
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)) 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;
166       else u[j][i].v = 0.0;
167 
168       u[j][i].u = 1.0 - 2.0*u[j][i].v;
169     }
170   }
171 
172   PetscCall(DMDAVecRestoreArray(da,U,&u));
173   PetscFunctionReturn(0);
174 }
175 
176 PetscErrorCode PerturbedInitialConditions3(DM da,Vec U)
177 {
178   PetscInt       i,j,xs,ys,xm,ym,Mx,My;
179   Field          **u;
180   PetscReal      hx,hy,x,y;
181 
182   PetscFunctionBegin;
183   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));
184 
185   hx = 2.5/(PetscReal)Mx;
186   hy = 2.5/(PetscReal)My;
187 
188   PetscCall(DMDAVecGetArray(da,U,&u));
189   /* Get local grid boundaries */
190   PetscCall(DMDAGetCorners(da,&xs,&ys,NULL,&xm,&ym,NULL));
191 
192   /* Compute function over the locally owned part of the grid */
193   for (j=ys; j<ys+ym; j++) {
194     y = j*hy;
195     for (i=xs; i<xs+xm; i++) {
196       x = i*hx;
197       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);
198       else u[j][i].v = 0.0;
199 
200       u[j][i].u = 1.0 - 2.0*u[j][i].v;
201     }
202   }
203 
204   PetscCall(DMDAVecRestoreArray(da,U,&u));
205   PetscFunctionReturn(0);
206 }
207 
208 int main(int argc,char **argv)
209 {
210   DM             da;
211   AppCtx         appctx;
212   PetscBool      forwardonly = PETSC_FALSE,implicitform = PETSC_FALSE;
213   PetscInt       perturbic = 1;
214 
215   PetscFunctionBeginUser;
216   PetscCall(PetscInitialize(&argc,&argv,(char*)0,help));
217   PetscCall(PetscOptionsGetBool(NULL,NULL,"-forwardonly",&forwardonly,NULL));
218   PetscCall(PetscOptionsGetBool(NULL,NULL,"-implicitform",&implicitform,NULL));
219   PetscCall(PetscOptionsGetInt(NULL,NULL,"-perturbic",&perturbic,NULL));
220 
221   appctx.D1    = 8.0e-5;
222   appctx.D2    = 4.0e-5;
223   appctx.gamma = .024;
224   appctx.kappa = .06;
225   appctx.aijpc = PETSC_FALSE;
226 
227   /* Create distributed array (DMDA) to manage parallel grid and vectors */
228   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));
229   PetscCall(DMSetFromOptions(da));
230   PetscCall(DMSetUp(da));
231   PetscCall(DMDASetFieldName(da,0,"u"));
232   PetscCall(DMDASetFieldName(da,1,"v"));
233 
234   /* Extract global vectors from DMDA; then duplicate for remaining
235      vectors that are the same types */
236   PetscCall(DMCreateGlobalVector(da,&appctx.U));
237 
238   /* Create timestepping solver context */
239   PetscCall(TSCreate(PETSC_COMM_WORLD,&appctx.ts));
240   PetscCall(TSSetType(appctx.ts,TSCN));
241   PetscCall(TSSetDM(appctx.ts,da));
242   PetscCall(TSSetProblemType(appctx.ts,TS_NONLINEAR));
243   PetscCall(TSSetEquationType(appctx.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
244   if (!implicitform) {
245     PetscCall(TSSetRHSFunction(appctx.ts,NULL,RHSFunction,&appctx));
246     PetscCall(TSSetRHSJacobian(appctx.ts,NULL,NULL,RHSJacobian,&appctx));
247   } else {
248     PetscCall(TSSetIFunction(appctx.ts,NULL,IFunction,&appctx));
249     PetscCall(TSSetIJacobian(appctx.ts,NULL,NULL,IJacobian,&appctx));
250   }
251 
252   /* Set initial conditions */
253   PetscCall(InitialConditions(da,appctx.U));
254   PetscCall(TSSetSolution(appctx.ts,appctx.U));
255 
256   /* Set solver options */
257   PetscCall(TSSetMaxTime(appctx.ts,2000.0));
258   PetscCall(TSSetTimeStep(appctx.ts,0.5));
259   PetscCall(TSSetExactFinalTime(appctx.ts,TS_EXACTFINALTIME_MATCHSTEP));
260   PetscCall(TSSetFromOptions(appctx.ts));
261 
262   PetscCall(GenerateOBs(appctx.ts,appctx.U,&appctx));
263 
264   if (!forwardonly) {
265     Tao           tao;
266     Vec           P;
267     Vec           lambda[1];
268 #if defined(PETSC_USE_LOG)
269     PetscLogStage opt_stage;
270 #endif
271 
272     PetscCall(PetscLogStageRegister("Optimization",&opt_stage));
273     PetscCall(PetscLogStagePush(opt_stage));
274     if (perturbic == 1) {
275       PetscCall(PerturbedInitialConditions(da,appctx.U));
276     } else if (perturbic == 2) {
277       PetscCall(PerturbedInitialConditions2(da,appctx.U));
278     } else if (perturbic == 3) {
279       PetscCall(PerturbedInitialConditions3(da,appctx.U));
280     }
281 
282     PetscCall(VecDuplicate(appctx.U,&lambda[0]));
283     PetscCall(TSSetCostGradients(appctx.ts,1,lambda,NULL));
284 
285     /* Have the TS save its trajectory needed by TSAdjointSolve() */
286     PetscCall(TSSetSaveTrajectory(appctx.ts));
287 
288     /* Create TAO solver and set desired solution method */
289     PetscCall(TaoCreate(PETSC_COMM_WORLD,&tao));
290     PetscCall(TaoSetType(tao,TAOBLMVM));
291 
292     /* Set initial guess for TAO */
293     PetscCall(VecDuplicate(appctx.U,&P));
294     PetscCall(VecCopy(appctx.U,P));
295     PetscCall(TaoSetSolution(tao,P));
296 
297     /* Set routine for function and gradient evaluation */
298     PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionAndGradient,&appctx));
299 
300     /* Check for any TAO command line options */
301     PetscCall(TaoSetFromOptions(tao));
302 
303     PetscCall(TaoSolve(tao));
304     PetscCall(TaoDestroy(&tao));
305     PetscCall(VecDestroy(&lambda[0]));
306     PetscCall(VecDestroy(&P));
307     PetscCall(PetscLogStagePop());
308   }
309 
310   /* Free work space.  All PETSc objects should be destroyed when they
311      are no longer needed. */
312   PetscCall(VecDestroy(&appctx.U));
313   PetscCall(TSDestroy(&appctx.ts));
314   PetscCall(DMDestroy(&da));
315   PetscCall(PetscFinalize());
316   return 0;
317 }
318 
319 /* ------------------------ TAO callbacks ---------------------------- */
320 
321 /*
322    FormFunctionAndGradient - Evaluates the function and corresponding gradient.
323 
324    Input Parameters:
325    tao - the Tao context
326    P   - the input vector
327    ctx - optional user-defined context, as set by TaoSetObjectiveAndGradient()
328 
329    Output Parameters:
330    f   - the newly evaluated function
331    G   - the newly evaluated gradient
332 */
333 PetscErrorCode FormFunctionAndGradient(Tao tao,Vec P,PetscReal *f,Vec G,void *ctx)
334 {
335   AppCtx         *appctx = (AppCtx*)ctx;
336   PetscReal      soberr,timestep;
337   Vec            *lambda;
338   Vec            SDiff;
339   DM             da;
340   char           filename[PETSC_MAX_PATH_LEN]="";
341   PetscViewer    viewer;
342 
343   PetscFunctionBeginUser;
344   PetscCall(TSSetTime(appctx->ts,0.0));
345   PetscCall(TSGetTimeStep(appctx->ts,&timestep));
346   if (timestep<0) {
347     PetscCall(TSSetTimeStep(appctx->ts,-timestep));
348   }
349   PetscCall(TSSetStepNumber(appctx->ts,0));
350   PetscCall(TSSetFromOptions(appctx->ts));
351 
352   PetscCall(VecDuplicate(P,&SDiff));
353   PetscCall(VecCopy(P,appctx->U));
354   PetscCall(TSGetDM(appctx->ts,&da));
355   *f = 0;
356 
357   PetscCall(TSSolve(appctx->ts,appctx->U));
358   PetscCall(PetscSNPrintf(filename,sizeof filename,"ex5opt.ob"));
359   PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD,filename,FILE_MODE_READ,&viewer));
360   PetscCall(VecLoad(SDiff,viewer));
361   PetscCall(PetscViewerDestroy(&viewer));
362   PetscCall(VecAYPX(SDiff,-1.,appctx->U));
363   PetscCall(VecDot(SDiff,SDiff,&soberr));
364   *f += soberr;
365 
366   PetscCall(TSGetCostGradients(appctx->ts,NULL,&lambda,NULL));
367   PetscCall(VecSet(lambda[0],0.0));
368   PetscCall(InitializeLambda(da,lambda[0],appctx->U,appctx));
369   PetscCall(TSAdjointSolve(appctx->ts));
370 
371   PetscCall(VecCopy(lambda[0],G));
372 
373   PetscCall(VecDestroy(&SDiff));
374   PetscFunctionReturn(0);
375 }
376 
377 /*TEST
378 
379    build:
380       depends: reaction_diffusion.c
381       requires: !complex !single
382 
383    test:
384       args: -ts_max_steps 5 -ts_type rk -ts_rk_type 3 -ts_trajectory_type memory -tao_monitor -tao_view -tao_gatol 1e-6
385       output_file: output/ex5opt_ic_1.out
386 
387 TEST*/
388