Actual source code: ex20td.c

petsc-master 2020-11-30
Report Typos and Errors
  1: static char help[] = "Performs adjoint sensitivity analysis for a van der Pol like \n\
  2: equation with time dependent parameters using two approaches :  \n\
  3: track  : track only local sensitivities at each adjoint step \n\
  4:          and accumulate them in a global array \n\
  5: global : track parameters at all timesteps together \n\
  6: Choose one of the two at runtime by -sa_method {track,global}. \n";

  8: /*
  9:   Concepts: TS^adjoint for time dependent parameters
 10:   Concepts: Customized adjoint monitor based sensitivity tracking
 11:   Concepts: All at once approach to sensitivity tracking
 12:   Processors: 1
 13: */


 16: /*
 17:    Simple example to demonstrate TSAdjoint capabilities for time dependent params
 18:    without integral cost terms using either a tracking or global method.

 20:    Modify the Van Der Pol Eq to :
 21:    [u1'] = [mu1(t)*u1]
 22:    [u2'] = [mu2(t)*((1-u1^2)*u2-u1)]
 23:    (with initial conditions & params independent)

 25:    Define uref to be solution with initail conditions (2,-2/3), mu=(1,1e3)
 26:    - u_ref : (1.5967,-1.02969)

 28:    Define const function as cost = 2-norm(u - u_ref);

 30:    Initialization for the adjoint TS :
 31:    - dcost/dy|final_time = 2*(u-u_ref)|final_time
 32:    - dcost/dp|final_time = 0

 34:    The tracking method only tracks local sensitivity at each time step
 35:    and accumulates these sensitivities in a global array. Since the structure
 36:    of the equations being solved at each time step does not change, the jacobian
 37:    wrt parameters is defined analogous to constant RHSJacobian for a liner
 38:    TSSolve and the size of the jacP is independent of the number of time
 39:    steps. Enable this mode of adjoint analysis by -sa_method track.

 41:    The global method combines the parameters at all timesteps and tracks them
 42:    together. Thus, the columns of the jacP matrix are filled dependent upon the
 43:    time step. Also, the dimensions of the jacP matrix now depend upon the number
 44:    of time steps. Enable this mode of adjoint analysis by -sa_method global.

 46:    Since the equations here have parameters at predefined time steps, this
 47:    example should be run with non adaptive time stepping solvers only. This
 48:    can be ensured by -ts_adapt_type none (which is the default behavior only
 49:    for certain TS solvers like TSCN. If using an explicit TS like TSRK,
 50:    please be sure to add the aforementioned option to disable adaptive
 51:    timestepping.)
 52: */

 54: /*
 55:    Include "petscts.h" so that we can use TS solvers.  Note that this file
 56:    automatically includes:
 57:      petscsys.h    - base PETSc routines   petscvec.h  - vectors
 58:      petscmat.h    - matrices
 59:      petscis.h     - index sets            petscksp.h  - Krylov subspace methods
 60:      petscviewer.h - viewers               petscpc.h   - preconditioners
 61:      petscksp.h    - linear solvers        petscsnes.h - nonlinear solvers
 62: */
 63: #include <petscts.h>

 65: extern PetscErrorCode RHSFunction(TS ,PetscReal ,Vec ,Vec ,void *);
 66: extern PetscErrorCode RHSJacobian(TS ,PetscReal ,Vec ,Mat ,Mat ,void *);
 67: extern PetscErrorCode RHSJacobianP_track(TS ,PetscReal ,Vec ,Mat ,void *);
 68: extern PetscErrorCode RHSJacobianP_global(TS ,PetscReal ,Vec ,Mat ,void *);
 69: extern PetscErrorCode Monitor(TS ,PetscInt ,PetscReal ,Vec ,void *);
 70: extern PetscErrorCode AdjointMonitor(TS ,PetscInt ,PetscReal ,Vec ,PetscInt ,Vec *, Vec *,void *);

 72: /*
 73:    User-defined application context - contains data needed by the
 74:    application-provided call-back routines.
 75: */

 77: typedef struct {
 78:  /*------------- Forward solve data structures --------------*/
 79:   PetscInt  max_steps;     /* number of steps to run ts for */
 80:   PetscReal final_time;    /* final time to integrate to*/
 81:   PetscReal time_step;     /* ts integration time step */
 82:   Vec       mu1;           /* time dependent params */
 83:   Vec       mu2;           /* time dependent params */
 84:   Vec       U;             /* solution vector */
 85:   Mat       A;             /* Jacobian matrix */

 87:   /*------------- Adjoint solve data structures --------------*/
 88:   Mat       Jacp;          /* JacobianP matrix */
 89:   Vec       lambda;        /* adjoint variable */
 90:   Vec       mup;           /* adjoint variable */

 92:   /*------------- Global accumation vecs for monitor based tracking --------------*/
 93:   Vec       sens_mu1;      /* global sensitivity accumulation */
 94:   Vec       sens_mu2;      /* global sensitivity accumulation */
 95:   PetscInt  adj_idx;       /* to keep track of adjoint solve index */
 96: } AppCtx;

 98: typedef enum {SA_TRACK, SA_GLOBAL} SAMethod;
 99: static const char *const SAMethods[] = {"TRACK","GLOBAL","SAMethod","SA_",0};

101: /* ----------------------- Explicit form of the ODE  -------------------- */

103: PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec U,Vec F,void *ctx)
104: {
105:   PetscErrorCode    ierr;
106:   AppCtx            *user = (AppCtx*) ctx;
107:   PetscScalar       *f;
108:   PetscInt          curr_step;
109:   const PetscScalar *u;
110:   const PetscScalar *mu1;
111:   const PetscScalar *mu2;

114:   TSGetStepNumber(ts,&curr_step);
115:   VecGetArrayRead(U,&u);
116:   VecGetArrayRead(user->mu1,&mu1);
117:   VecGetArrayRead(user->mu2,&mu2);
118:   VecGetArray(F,&f);
119:   f[0] = mu1[curr_step]*u[1];
120:   f[1] = mu2[curr_step]*((1.-u[0]*u[0])*u[1]-u[0]);
121:   VecRestoreArrayRead(U,&u);
122:   VecRestoreArrayRead(user->mu1,&mu1);
123:   VecRestoreArrayRead(user->mu2,&mu2);
124:   VecRestoreArray(F,&f);
125:   return(0);
126: }

128: PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,void *ctx)
129: {
130:   PetscErrorCode    ierr;
131:   AppCtx            *user = (AppCtx*) ctx;
132:   PetscInt          rowcol[] = {0,1};
133:   PetscScalar       J[2][2];
134:   PetscInt          curr_step;
135:   const PetscScalar *u;
136:   const PetscScalar *mu1;
137:   const PetscScalar *mu2;

140:   TSGetStepNumber(ts,&curr_step);
141:   VecGetArrayRead(user->mu1,&mu1);
142:   VecGetArrayRead(user->mu2,&mu2);
143:   VecGetArrayRead(U,&u);
144:   J[0][0] = 0;
145:   J[1][0] = -mu2[curr_step]*(2.0*u[1]*u[0]+1.);
146:   J[0][1] = mu1[curr_step];
147:   J[1][1] = mu2[curr_step]*(1.0-u[0]*u[0]);
148:   MatSetValues(A,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);
149:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
150:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
151:   VecRestoreArrayRead(U,&u);
152:   VecRestoreArrayRead(user->mu1,&mu1);
153:   VecRestoreArrayRead(user->mu2,&mu2);
154:   return(0);
155: }

157: /* ------------------ Jacobian wrt parameters for tracking method ------------------ */

159: PetscErrorCode RHSJacobianP_track(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
160: {
161:   PetscErrorCode    ierr;
162:   PetscInt          row[] = {0,1},col[] = {0,1};
163:   PetscScalar       J[2][2];
164:   const PetscScalar *u;

167:   VecGetArrayRead(U,&u);
168:   J[0][0] = u[1];
169:   J[1][0] = 0;
170:   J[0][1] = 0;
171:   J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
172:   MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
173:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
174:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
175:   VecRestoreArrayRead(U,&u);
176:   return(0);
177: }

179: /* ------------------ Jacobian wrt parameters for global method ------------------ */

181: PetscErrorCode RHSJacobianP_global(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
182: {
183:   PetscErrorCode    ierr;
184:   PetscInt          row[] = {0,1},col[] = {0,1};
185:   PetscScalar       J[2][2];
186:   const PetscScalar *u;
187:   PetscInt          curr_step;

190:   TSGetStepNumber(ts,&curr_step);
191:   VecGetArrayRead(U,&u);
192:   J[0][0] = u[1];
193:   J[1][0] = 0;
194:   J[0][1] = 0;
195:   J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
196:   col[0] = (curr_step)*2;
197:   col[1] = (curr_step)*2+1;
198:   MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
199:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
200:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
201:   VecRestoreArrayRead(U,&u);
202:   return(0);
203: }

205: /* Dump solution to console if called */
206: PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec U,void *ctx)
207: {
208:   PetscErrorCode    ierr;

211:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution at time %e is \n", t);
212:   VecView(U,PETSC_VIEWER_STDOUT_WORLD);
213:   return(0);
214: }

216: /* Customized adjoint monitor to keep track of local
217:    sensitivities by storing them in a global sensitivity array.
218:    Note : This routine is only used for the tracking method. */
219: PetscErrorCode AdjointMonitor(TS ts,PetscInt steps,PetscReal time,Vec u,PetscInt numcost,Vec *lambda, Vec *mu,void *ctx)
220: {
221:   PetscErrorCode    ierr;
222:   AppCtx            *user = (AppCtx*) ctx;
223:   PetscInt          curr_step;
224:   PetscScalar       *sensmu1_glob;
225:   PetscScalar       *sensmu2_glob;
226:   const PetscScalar *sensmu_loc;

229:   TSGetStepNumber(ts,&curr_step);
230:   /* Note that we skip the first call to the monitor in the adjoint
231:      solve since the sensitivities are already set (during
232:      initialization of adjoint vectors).
233:      We also note that each indvidial TSAdjointSolve calls the monitor
234:      twice, once at the step it is integrating from and once at the step
235:      it integrates to. Only the second call is useful for transferring
236:      local sensitivities to the global array. */
237:   if (curr_step == user->adj_idx) {
238:     return(0);
239:   } else {
240:     VecGetArrayRead(*mu,&sensmu_loc);
241:     VecGetArray(user->sens_mu1,&sensmu1_glob);
242:     VecGetArray(user->sens_mu2,&sensmu2_glob);
243:     sensmu1_glob[curr_step] = sensmu_loc[0];
244:     sensmu2_glob[curr_step] = sensmu_loc[1];
245:     VecRestoreArray(user->sens_mu1,&sensmu1_glob);
246:     VecRestoreArray(user->sens_mu2,&sensmu2_glob);
247:     VecRestoreArrayRead(*mu,&sensmu_loc);
248:     return(0);
249:   }
250: }


253: int main(int argc,char **argv)
254: {
255:   TS             ts;
256:   AppCtx         user;
257:   PetscScalar    *x_ptr,*y_ptr,*u_ptr;
258:   PetscMPIInt    size;
259:   PetscBool      monitor = PETSC_FALSE;
260:   SAMethod       sa = SA_GLOBAL;

263:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
264:      Initialize program
265:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
266:   PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;
267:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
268:   if (size != 1) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");

270:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
271:      Set runtime options
272:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
273:   PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"SA analysis options.","");{
274:   PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL);
275:   PetscOptionsEnum("-sa_method","Sensitivity analysis method (track or global)","",SAMethods,(PetscEnum)sa,(PetscEnum*)&sa,NULL);
276:   }
277:   PetscOptionsEnd();

279:   user.final_time = 0.1;
280:   user.max_steps  = 5;
281:   user.time_step  = user.final_time/user.max_steps;

283:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
284:      Create necessary matrix and vectors for forward solve.
285:      Create Jacp matrix for adjoint solve.
286:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
287:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu1);
288:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu2);
289:   VecSet(user.mu1,1.25);
290:   VecSet(user.mu2,1.0e2);

292:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
293:       For tracking method : create the global sensitivity array to
294:       accumulate sensitivity with respect to parameters at each step.
295:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
296:   if (sa == SA_TRACK) {
297:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu1);
298:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu2);
299:   }

301:   MatCreate(PETSC_COMM_WORLD,&user.A);
302:   MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2);
303:   MatSetFromOptions(user.A);
304:   MatSetUp(user.A);
305:   MatCreateVecs(user.A,&user.U,NULL);

307:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
308:       Note that the dimensions of the Jacp matrix depend upon the
309:       sensitivity analysis method being used !
310:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
311:   MatCreate(PETSC_COMM_WORLD,&user.Jacp);
312:   if (sa == SA_TRACK) {
313:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2);
314:   }
315:   if (sa == SA_GLOBAL) {
316:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,user.max_steps*2);
317:   }
318:   MatSetFromOptions(user.Jacp);
319:   MatSetUp(user.Jacp);

321:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
322:      Create timestepping solver context
323:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
324:   TSCreate(PETSC_COMM_WORLD,&ts);
325:   TSSetEquationType(ts,TS_EQ_ODE_EXPLICIT);
326:   TSSetType(ts,TSCN);

328:   TSSetRHSFunction(ts,NULL,RHSFunction,&user);
329:   TSSetRHSJacobian(ts,user.A,user.A,RHSJacobian,&user);
330:   if (sa == SA_TRACK) {
331:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_track,&user);
332:   }
333:   if (sa == SA_GLOBAL) {
334:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_global,&user);
335:   }

337:   TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
338:   TSSetMaxTime(ts,user.final_time);
339:   TSSetTimeStep(ts,user.final_time/user.max_steps);

341:   if (monitor) {
342:     TSMonitorSet(ts,Monitor,&user,NULL);
343:   }
344:   if (sa == SA_TRACK) {
345:     TSAdjointMonitorSet(ts,AdjointMonitor,&user,NULL);
346:   }

348:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
349:      Set initial conditions
350:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
351:   VecGetArray(user.U,&x_ptr);
352:   x_ptr[0] = 2.0;
353:   x_ptr[1] = -2.0/3.0;
354:   VecRestoreArray(user.U,&x_ptr);

356:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
357:      Save trajectory of solution so that TSAdjointSolve() may be used
358:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
359:   TSSetSaveTrajectory(ts);

361:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
362:      Set runtime options
363:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
364:   TSSetFromOptions(ts);

366:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
367:      Execute forward model and print solution.
368:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
369:   TSSolve(ts,user.U);
370:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution of forward TS :\n");
371:   VecView(user.U,PETSC_VIEWER_STDOUT_WORLD);
372:   PetscPrintf(PETSC_COMM_WORLD,"\n Forward TS solve successfull! Adjoint run begins!\n");

374:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
375:      Adjoint model starts here! Create adjoint vectors.
376:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
377:   MatCreateVecs(user.A,&user.lambda,NULL);
378:   MatCreateVecs(user.Jacp,&user.mup,NULL);

380:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
381:      Set initial conditions for the adjoint vector
382:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
383:   VecGetArray(user.U,&u_ptr);
384:   VecGetArray(user.lambda,&y_ptr);
385:   y_ptr[0] = 2*(u_ptr[0] - 1.5967);
386:   y_ptr[1] = 2*(u_ptr[1] - -(1.02969));
387:   VecRestoreArray(user.lambda,&y_ptr);
388:   VecRestoreArray(user.U,&y_ptr);
389:   VecSet(user.mup,0);

391:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
392:      Set number of cost functions.
393:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
394:   TSSetCostGradients(ts,1,&user.lambda,&user.mup);

396:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
397:      The adjoint vector mup has to be reset for each adjoint step when
398:      using the tracking method as we want to treat the parameters at each
399:      time step one at a time and prevent accumulation of the sensitivities
400:      from parameters at previous time steps.
401:      This is not necessary for the global method as each time dependent
402:      parameter is treated as an independent parameter.
403:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
404:   if (sa == SA_TRACK) {
405:     for (user.adj_idx=user.max_steps; user.adj_idx>0; user.adj_idx--) {
406:       VecSet(user.mup,0);
407:       TSAdjointSetSteps(ts, 1);
408:       TSAdjointSolve(ts);
409:     }
410:   }
411:   if (sa == SA_GLOBAL) {
412:     TSAdjointSolve(ts);
413:   }

415:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
416:      Dispaly adjoint sensitivities wrt parameters and initial conditions
417:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
418:   if (sa == SA_TRACK) {
419:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu1: d[cost]/d[mu1]\n");
420:     VecView(user.sens_mu1,PETSC_VIEWER_STDOUT_WORLD);
421:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu2: d[cost]/d[mu2]\n");
422:     VecView(user.sens_mu2,PETSC_VIEWER_STDOUT_WORLD);
423:   }

425:   if (sa == SA_GLOBAL) {
426:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  params: d[cost]/d[p], where p refers to \n\
427:                     the interlaced vector made by combining mu1,mu2\n");
428:     VecView(user.mup,PETSC_VIEWER_STDOUT_WORLD);
429:   }

431:   PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt initial conditions: d[cost]/d[u(t=0)]\n");
432:   VecView(user.lambda,PETSC_VIEWER_STDOUT_WORLD);

434:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
435:      Free work space!
436:      All PETSc objects should be destroyed when they are no longer needed.
437:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
438:   MatDestroy(&user.A);
439:   MatDestroy(&user.Jacp);
440:   VecDestroy(&user.U);
441:   VecDestroy(&user.lambda);
442:   VecDestroy(&user.mup);
443:   VecDestroy(&user.mu1);
444:   VecDestroy(&user.mu2);
445:   if (sa == SA_TRACK) {
446:     VecDestroy(&user.sens_mu1);
447:     VecDestroy(&user.sens_mu2);
448:   }
449:   TSDestroy(&ts);
450:   PetscFinalize();
451:   return(ierr);
452: }


455: /*TEST

457:   test:
458:     requires: !complex
459:     suffix : track
460:     args : -sa_method track

462:   test:
463:     requires: !complex
464:     suffix : global
465:     args : -sa_method global

467: TEST*/