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