xref: /petsc/src/ts/tutorials/advection-diffusion-reaction/ex5opt_ic.c (revision 40badf4fbc550ac1f60bd080eaff6de6d55b946d)
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   PetscErrorCode ierr;
211   DM             da;
212   AppCtx         appctx;
213   PetscBool      forwardonly = PETSC_FALSE,implicitform = PETSC_FALSE;
214   PetscInt       perturbic = 1;
215 
216   ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;
217   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-forwardonly",&forwardonly,NULL));
218   CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-implicitform",&implicitform,NULL));
219   CHKERRQ(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   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));
229   CHKERRQ(DMSetFromOptions(da));
230   CHKERRQ(DMSetUp(da));
231   CHKERRQ(DMDASetFieldName(da,0,"u"));
232   CHKERRQ(DMDASetFieldName(da,1,"v"));
233 
234   /* Extract global vectors from DMDA; then duplicate for remaining
235      vectors that are the same types */
236   CHKERRQ(DMCreateGlobalVector(da,&appctx.U));
237 
238   /* Create timestepping solver context */
239   CHKERRQ(TSCreate(PETSC_COMM_WORLD,&appctx.ts));
240   CHKERRQ(TSSetType(appctx.ts,TSCN));
241   CHKERRQ(TSSetDM(appctx.ts,da));
242   CHKERRQ(TSSetProblemType(appctx.ts,TS_NONLINEAR));
243   CHKERRQ(TSSetEquationType(appctx.ts,TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
244   if (!implicitform) {
245     CHKERRQ(TSSetRHSFunction(appctx.ts,NULL,RHSFunction,&appctx));
246     CHKERRQ(TSSetRHSJacobian(appctx.ts,NULL,NULL,RHSJacobian,&appctx));
247   } else {
248     CHKERRQ(TSSetIFunction(appctx.ts,NULL,IFunction,&appctx));
249     CHKERRQ(TSSetIJacobian(appctx.ts,NULL,NULL,IJacobian,&appctx));
250   }
251 
252   /* Set initial conditions */
253   CHKERRQ(InitialConditions(da,appctx.U));
254   CHKERRQ(TSSetSolution(appctx.ts,appctx.U));
255 
256   /* Set solver options */
257   CHKERRQ(TSSetMaxTime(appctx.ts,2000.0));
258   CHKERRQ(TSSetTimeStep(appctx.ts,0.5));
259   CHKERRQ(TSSetExactFinalTime(appctx.ts,TS_EXACTFINALTIME_MATCHSTEP));
260   CHKERRQ(TSSetFromOptions(appctx.ts));
261 
262   CHKERRQ(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     CHKERRQ(PetscLogStageRegister("Optimization",&opt_stage));
273     CHKERRQ(PetscLogStagePush(opt_stage));
274     if (perturbic == 1) {
275       CHKERRQ(PerturbedInitialConditions(da,appctx.U));
276     } else if (perturbic == 2) {
277       CHKERRQ(PerturbedInitialConditions2(da,appctx.U));
278     } else if (perturbic == 3) {
279       CHKERRQ(PerturbedInitialConditions3(da,appctx.U));
280     }
281 
282     CHKERRQ(VecDuplicate(appctx.U,&lambda[0]));
283     CHKERRQ(TSSetCostGradients(appctx.ts,1,lambda,NULL));
284 
285     /* Have the TS save its trajectory needed by TSAdjointSolve() */
286     CHKERRQ(TSSetSaveTrajectory(appctx.ts));
287 
288     /* Create TAO solver and set desired solution method */
289     CHKERRQ(TaoCreate(PETSC_COMM_WORLD,&tao));
290     CHKERRQ(TaoSetType(tao,TAOBLMVM));
291 
292     /* Set initial guess for TAO */
293     CHKERRQ(VecDuplicate(appctx.U,&P));
294     CHKERRQ(VecCopy(appctx.U,P));
295     CHKERRQ(TaoSetSolution(tao,P));
296 
297     /* Set routine for function and gradient evaluation */
298     CHKERRQ(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionAndGradient,&appctx));
299 
300     /* Check for any TAO command line options */
301     CHKERRQ(TaoSetFromOptions(tao));
302 
303     CHKERRQ(TaoSolve(tao));
304     CHKERRQ(TaoDestroy(&tao));
305     CHKERRQ(VecDestroy(&lambda[0]));
306     CHKERRQ(VecDestroy(&P));
307     CHKERRQ(PetscLogStagePop());
308   }
309 
310   /* Free work space.  All PETSc objects should be destroyed when they
311      are no longer needed. */
312   CHKERRQ(VecDestroy(&appctx.U));
313   CHKERRQ(TSDestroy(&appctx.ts));
314   CHKERRQ(DMDestroy(&da));
315   ierr = PetscFinalize();
316   return ierr;
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   CHKERRQ(TSSetTime(appctx->ts,0.0));
345   CHKERRQ(TSGetTimeStep(appctx->ts,&timestep));
346   if (timestep<0) {
347     CHKERRQ(TSSetTimeStep(appctx->ts,-timestep));
348   }
349   CHKERRQ(TSSetStepNumber(appctx->ts,0));
350   CHKERRQ(TSSetFromOptions(appctx->ts));
351 
352   CHKERRQ(VecDuplicate(P,&SDiff));
353   CHKERRQ(VecCopy(P,appctx->U));
354   CHKERRQ(TSGetDM(appctx->ts,&da));
355   *f = 0;
356 
357   CHKERRQ(TSSolve(appctx->ts,appctx->U));
358   CHKERRQ(PetscSNPrintf(filename,sizeof filename,"ex5opt.ob"));
359   CHKERRQ(PetscViewerBinaryOpen(PETSC_COMM_WORLD,filename,FILE_MODE_READ,&viewer));
360   CHKERRQ(VecLoad(SDiff,viewer));
361   CHKERRQ(PetscViewerDestroy(&viewer));
362   CHKERRQ(VecAYPX(SDiff,-1.,appctx->U));
363   CHKERRQ(VecDot(SDiff,SDiff,&soberr));
364   *f += soberr;
365 
366   CHKERRQ(TSGetCostGradients(appctx->ts,NULL,&lambda,NULL));
367   CHKERRQ(VecSet(lambda[0],0.0));
368   CHKERRQ(InitializeLambda(da,lambda[0],appctx->U,appctx));
369   CHKERRQ(TSAdjointSolve(appctx->ts));
370 
371   CHKERRQ(VecCopy(lambda[0],G));
372 
373   CHKERRQ(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