xref: /petsc/src/ts/tutorials/advection-diffusion-reaction/ex5opt_ic.c (revision b122ec5aa1bd4469eb4e0673542fb7de3f411254)
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   CHKERRQ(VecDuplicate(U,&Uob));
37   CHKERRQ(PetscSNPrintf(filename,sizeof filename,"ex5opt.ob"));
38   CHKERRQ(PetscViewerBinaryOpen(PETSC_COMM_WORLD,filename,FILE_MODE_READ,&viewer));
39   CHKERRQ(VecLoad(Uob,viewer));
40   CHKERRQ(PetscViewerDestroy(&viewer));
41   CHKERRQ(VecAYPX(Uob,-1.,U));
42   CHKERRQ(VecScale(Uob,2.0));
43   CHKERRQ(VecAXPY(lambda,1.,Uob));
44   CHKERRQ(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   CHKERRQ(PetscViewerCreate(PetscObjectComm((PetscObject)da),viewer));
55   CHKERRQ(PetscViewerSetType(*viewer,PETSCVIEWERBINARY));
56   CHKERRQ(PetscViewerFileSetMode(*viewer,FILE_MODE_WRITE));
57   CHKERRQ(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   CHKERRQ(TSGetDM(ts,&da));
72   CHKERRQ(TSSolve(ts,U));
73   CHKERRQ(PetscSNPrintf(filename,sizeof filename,"ex5opt.ob"));
74   CHKERRQ(OutputBIN(da,filename,&viewer));
75   CHKERRQ(VecView(U,viewer));
76   CHKERRQ(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   CHKERRQ(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   CHKERRQ(DMDAVecGetArray(da,U,&u));
93   /* Get local grid boundaries */
94   CHKERRQ(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   CHKERRQ(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   CHKERRQ(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   CHKERRQ(DMDAVecGetArray(da,U,&u));
125   /* Get local grid boundaries */
126   CHKERRQ(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   CHKERRQ(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   CHKERRQ(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   CHKERRQ(DMDAVecGetArray(da,U,&u));
157   /* Get local grid boundaries */
158   CHKERRQ(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   CHKERRQ(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   CHKERRQ(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   CHKERRQ(DMDAVecGetArray(da,U,&u));
189   /* Get local grid boundaries */
190   CHKERRQ(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   CHKERRQ(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   CHKERRQ(PetscInitialize(&argc,&argv,(char*)0,help));
216   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-forwardonly",&forwardonly,NULL));
217   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-implicitform",&implicitform,NULL));
218   CHKERRQ(PetscOptionsGetInt(NULL,NULL,"-perturbic",&perturbic,NULL));
219 
220   appctx.D1    = 8.0e-5;
221   appctx.D2    = 4.0e-5;
222   appctx.gamma = .024;
223   appctx.kappa = .06;
224   appctx.aijpc = PETSC_FALSE;
225 
226   /* Create distributed array (DMDA) to manage parallel grid and vectors */
227   CHKERRQ(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_PERIODIC,DM_BOUNDARY_PERIODIC,DMDA_STENCIL_STAR,64,64,PETSC_DECIDE,PETSC_DECIDE,2,1,NULL,NULL,&da));
228   CHKERRQ(DMSetFromOptions(da));
229   CHKERRQ(DMSetUp(da));
230   CHKERRQ(DMDASetFieldName(da,0,"u"));
231   CHKERRQ(DMDASetFieldName(da,1,"v"));
232 
233   /* Extract global vectors from DMDA; then duplicate for remaining
234      vectors that are the same types */
235   CHKERRQ(DMCreateGlobalVector(da,&appctx.U));
236 
237   /* Create timestepping solver context */
238   CHKERRQ(TSCreate(PETSC_COMM_WORLD,&appctx.ts));
239   CHKERRQ(TSSetType(appctx.ts,TSCN));
240   CHKERRQ(TSSetDM(appctx.ts,da));
241   CHKERRQ(TSSetProblemType(appctx.ts,TS_NONLINEAR));
242   CHKERRQ(TSSetEquationType(appctx.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
243   if (!implicitform) {
244     CHKERRQ(TSSetRHSFunction(appctx.ts,NULL,RHSFunction,&appctx));
245     CHKERRQ(TSSetRHSJacobian(appctx.ts,NULL,NULL,RHSJacobian,&appctx));
246   } else {
247     CHKERRQ(TSSetIFunction(appctx.ts,NULL,IFunction,&appctx));
248     CHKERRQ(TSSetIJacobian(appctx.ts,NULL,NULL,IJacobian,&appctx));
249   }
250 
251   /* Set initial conditions */
252   CHKERRQ(InitialConditions(da,appctx.U));
253   CHKERRQ(TSSetSolution(appctx.ts,appctx.U));
254 
255   /* Set solver options */
256   CHKERRQ(TSSetMaxTime(appctx.ts,2000.0));
257   CHKERRQ(TSSetTimeStep(appctx.ts,0.5));
258   CHKERRQ(TSSetExactFinalTime(appctx.ts,TS_EXACTFINALTIME_MATCHSTEP));
259   CHKERRQ(TSSetFromOptions(appctx.ts));
260 
261   CHKERRQ(GenerateOBs(appctx.ts,appctx.U,&appctx));
262 
263   if (!forwardonly) {
264     Tao           tao;
265     Vec           P;
266     Vec           lambda[1];
267 #if defined(PETSC_USE_LOG)
268     PetscLogStage opt_stage;
269 #endif
270 
271     CHKERRQ(PetscLogStageRegister("Optimization",&opt_stage));
272     CHKERRQ(PetscLogStagePush(opt_stage));
273     if (perturbic == 1) {
274       CHKERRQ(PerturbedInitialConditions(da,appctx.U));
275     } else if (perturbic == 2) {
276       CHKERRQ(PerturbedInitialConditions2(da,appctx.U));
277     } else if (perturbic == 3) {
278       CHKERRQ(PerturbedInitialConditions3(da,appctx.U));
279     }
280 
281     CHKERRQ(VecDuplicate(appctx.U,&lambda[0]));
282     CHKERRQ(TSSetCostGradients(appctx.ts,1,lambda,NULL));
283 
284     /* Have the TS save its trajectory needed by TSAdjointSolve() */
285     CHKERRQ(TSSetSaveTrajectory(appctx.ts));
286 
287     /* Create TAO solver and set desired solution method */
288     CHKERRQ(TaoCreate(PETSC_COMM_WORLD,&tao));
289     CHKERRQ(TaoSetType(tao,TAOBLMVM));
290 
291     /* Set initial guess for TAO */
292     CHKERRQ(VecDuplicate(appctx.U,&P));
293     CHKERRQ(VecCopy(appctx.U,P));
294     CHKERRQ(TaoSetSolution(tao,P));
295 
296     /* Set routine for function and gradient evaluation */
297     CHKERRQ(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionAndGradient,&appctx));
298 
299     /* Check for any TAO command line options */
300     CHKERRQ(TaoSetFromOptions(tao));
301 
302     CHKERRQ(TaoSolve(tao));
303     CHKERRQ(TaoDestroy(&tao));
304     CHKERRQ(VecDestroy(&lambda[0]));
305     CHKERRQ(VecDestroy(&P));
306     CHKERRQ(PetscLogStagePop());
307   }
308 
309   /* Free work space.  All PETSc objects should be destroyed when they
310      are no longer needed. */
311   CHKERRQ(VecDestroy(&appctx.U));
312   CHKERRQ(TSDestroy(&appctx.ts));
313   CHKERRQ(DMDestroy(&da));
314   CHKERRQ(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 */
332 PetscErrorCode FormFunctionAndGradient(Tao tao,Vec P,PetscReal *f,Vec G,void *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   CHKERRQ(TSSetTime(appctx->ts,0.0));
344   CHKERRQ(TSGetTimeStep(appctx->ts,&timestep));
345   if (timestep<0) {
346     CHKERRQ(TSSetTimeStep(appctx->ts,-timestep));
347   }
348   CHKERRQ(TSSetStepNumber(appctx->ts,0));
349   CHKERRQ(TSSetFromOptions(appctx->ts));
350 
351   CHKERRQ(VecDuplicate(P,&SDiff));
352   CHKERRQ(VecCopy(P,appctx->U));
353   CHKERRQ(TSGetDM(appctx->ts,&da));
354   *f = 0;
355 
356   CHKERRQ(TSSolve(appctx->ts,appctx->U));
357   CHKERRQ(PetscSNPrintf(filename,sizeof filename,"ex5opt.ob"));
358   CHKERRQ(PetscViewerBinaryOpen(PETSC_COMM_WORLD,filename,FILE_MODE_READ,&viewer));
359   CHKERRQ(VecLoad(SDiff,viewer));
360   CHKERRQ(PetscViewerDestroy(&viewer));
361   CHKERRQ(VecAYPX(SDiff,-1.,appctx->U));
362   CHKERRQ(VecDot(SDiff,SDiff,&soberr));
363   *f += soberr;
364 
365   CHKERRQ(TSGetCostGradients(appctx->ts,NULL,&lambda,NULL));
366   CHKERRQ(VecSet(lambda[0],0.0));
367   CHKERRQ(InitializeLambda(da,lambda[0],appctx->U,appctx));
368   CHKERRQ(TSAdjointSolve(appctx->ts));
369 
370   CHKERRQ(VecCopy(lambda[0],G));
371 
372   CHKERRQ(VecDestroy(&SDiff));
373   PetscFunctionReturn(0);
374 }
375 
376 /*TEST
377 
378    build:
379       depends: reaction_diffusion.c
380       requires: !complex !single
381 
382    test:
383       args: -ts_max_steps 5 -ts_type rk -ts_rk_type 3 -ts_trajectory_type memory -tao_monitor -tao_view -tao_gatol 1e-6
384       output_file: output/ex5opt_ic_1.out
385 
386 TEST*/
387