Actual source code: ex20td.c

  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:    Simple example to demonstrate TSAdjoint capabilities for time dependent params
 10:    without integral cost terms using either a tracking or global method.

 12:    Modify the Van Der Pol Eq to :
 13:    [u1'] = [mu1(t)*u1]
 14:    [u2'] = [mu2(t)*((1-u1^2)*u2-u1)]
 15:    (with initial conditions & params independent)

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

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

 22:    Initialization for the adjoint TS :
 23:    - dcost/dy|final_time = 2*(u-u_ref)|final_time
 24:    - dcost/dp|final_time = 0

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

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

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

 46: /*
 47:    Include "petscts.h" so that we can use TS solvers.  Note that this file
 48:    automatically includes:
 49:      petscsys.h    - base PETSc routines   petscvec.h  - vectors
 50:      petscmat.h    - matrices
 51:      petscis.h     - index sets            petscksp.h  - Krylov subspace methods
 52:      petscviewer.h - viewers               petscpc.h   - preconditioners
 53:      petscksp.h    - linear solvers        petscsnes.h - nonlinear solvers
 54: */
 55: #include <petscts.h>

 57: extern PetscErrorCode RHSFunction(TS ,PetscReal ,Vec ,Vec ,void *);
 58: extern PetscErrorCode RHSJacobian(TS ,PetscReal ,Vec ,Mat ,Mat ,void *);
 59: extern PetscErrorCode RHSJacobianP_track(TS ,PetscReal ,Vec ,Mat ,void *);
 60: extern PetscErrorCode RHSJacobianP_global(TS ,PetscReal ,Vec ,Mat ,void *);
 61: extern PetscErrorCode Monitor(TS ,PetscInt ,PetscReal ,Vec ,void *);
 62: extern PetscErrorCode AdjointMonitor(TS ,PetscInt ,PetscReal ,Vec ,PetscInt ,Vec *, Vec *,void *);

 64: /*
 65:    User-defined application context - contains data needed by the
 66:    application-provided call-back routines.
 67: */

 69: typedef struct {
 70:  /*------------- Forward solve data structures --------------*/
 71:   PetscInt  max_steps;     /* number of steps to run ts for */
 72:   PetscReal final_time;    /* final time to integrate to*/
 73:   PetscReal time_step;     /* ts integration time step */
 74:   Vec       mu1;           /* time dependent params */
 75:   Vec       mu2;           /* time dependent params */
 76:   Vec       U;             /* solution vector */
 77:   Mat       A;             /* Jacobian matrix */

 79:   /*------------- Adjoint solve data structures --------------*/
 80:   Mat       Jacp;          /* JacobianP matrix */
 81:   Vec       lambda;        /* adjoint variable */
 82:   Vec       mup;           /* adjoint variable */

 84:   /*------------- Global accumation vecs for monitor based tracking --------------*/
 85:   Vec       sens_mu1;      /* global sensitivity accumulation */
 86:   Vec       sens_mu2;      /* global sensitivity accumulation */
 87:   PetscInt  adj_idx;       /* to keep track of adjoint solve index */
 88: } AppCtx;

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

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

 95: PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec U,Vec F,void *ctx)
 96: {
 97:   AppCtx            *user = (AppCtx*) ctx;
 98:   PetscScalar       *f;
 99:   PetscInt          curr_step;
100:   const PetscScalar *u;
101:   const PetscScalar *mu1;
102:   const PetscScalar *mu2;

105:   TSGetStepNumber(ts,&curr_step);
106:   VecGetArrayRead(U,&u);
107:   VecGetArrayRead(user->mu1,&mu1);
108:   VecGetArrayRead(user->mu2,&mu2);
109:   VecGetArray(F,&f);
110:   f[0] = mu1[curr_step]*u[1];
111:   f[1] = mu2[curr_step]*((1.-u[0]*u[0])*u[1]-u[0]);
112:   VecRestoreArrayRead(U,&u);
113:   VecRestoreArrayRead(user->mu1,&mu1);
114:   VecRestoreArrayRead(user->mu2,&mu2);
115:   VecRestoreArray(F,&f);
116:   return 0;
117: }

119: PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,void *ctx)
120: {
121:   AppCtx            *user = (AppCtx*) ctx;
122:   PetscInt          rowcol[] = {0,1};
123:   PetscScalar       J[2][2];
124:   PetscInt          curr_step;
125:   const PetscScalar *u;
126:   const PetscScalar *mu1;
127:   const PetscScalar *mu2;

130:   TSGetStepNumber(ts,&curr_step);
131:   VecGetArrayRead(user->mu1,&mu1);
132:   VecGetArrayRead(user->mu2,&mu2);
133:   VecGetArrayRead(U,&u);
134:   J[0][0] = 0;
135:   J[1][0] = -mu2[curr_step]*(2.0*u[1]*u[0]+1.);
136:   J[0][1] = mu1[curr_step];
137:   J[1][1] = mu2[curr_step]*(1.0-u[0]*u[0]);
138:   MatSetValues(A,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);
139:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
140:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
141:   VecRestoreArrayRead(U,&u);
142:   VecRestoreArrayRead(user->mu1,&mu1);
143:   VecRestoreArrayRead(user->mu2,&mu2);
144:   return 0;
145: }

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

149: PetscErrorCode RHSJacobianP_track(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
150: {
151:   PetscInt          row[] = {0,1},col[] = {0,1};
152:   PetscScalar       J[2][2];
153:   const PetscScalar *u;

156:   VecGetArrayRead(U,&u);
157:   J[0][0] = u[1];
158:   J[1][0] = 0;
159:   J[0][1] = 0;
160:   J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
161:   MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
162:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
163:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
164:   VecRestoreArrayRead(U,&u);
165:   return 0;
166: }

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

170: PetscErrorCode RHSJacobianP_global(TS ts,PetscReal t,Vec U,Mat A,void *ctx)
171: {
172:   PetscInt          row[] = {0,1},col[] = {0,1};
173:   PetscScalar       J[2][2];
174:   const PetscScalar *u;
175:   PetscInt          curr_step;

178:   TSGetStepNumber(ts,&curr_step);
179:   VecGetArrayRead(U,&u);
180:   J[0][0] = u[1];
181:   J[1][0] = 0;
182:   J[0][1] = 0;
183:   J[1][1] = (1.-u[0]*u[0])*u[1]-u[0];
184:   col[0] = (curr_step)*2;
185:   col[1] = (curr_step)*2+1;
186:   MatSetValues(A,2,row,2,col,&J[0][0],INSERT_VALUES);
187:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
188:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
189:   VecRestoreArrayRead(U,&u);
190:   return 0;
191: }

193: /* Dump solution to console if called */
194: PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec U,void *ctx)
195: {
197:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution at time %e is \n", t);
198:   VecView(U,PETSC_VIEWER_STDOUT_WORLD);
199:   return 0;
200: }

202: /* Customized adjoint monitor to keep track of local
203:    sensitivities by storing them in a global sensitivity array.
204:    Note : This routine is only used for the tracking method. */
205: PetscErrorCode AdjointMonitor(TS ts,PetscInt steps,PetscReal time,Vec u,PetscInt numcost,Vec *lambda, Vec *mu,void *ctx)
206: {
207:   AppCtx            *user = (AppCtx*) ctx;
208:   PetscInt          curr_step;
209:   PetscScalar       *sensmu1_glob;
210:   PetscScalar       *sensmu2_glob;
211:   const PetscScalar *sensmu_loc;

214:   TSGetStepNumber(ts,&curr_step);
215:   /* Note that we skip the first call to the monitor in the adjoint
216:      solve since the sensitivities are already set (during
217:      initialization of adjoint vectors).
218:      We also note that each indvidial TSAdjointSolve calls the monitor
219:      twice, once at the step it is integrating from and once at the step
220:      it integrates to. Only the second call is useful for transferring
221:      local sensitivities to the global array. */
222:   if (curr_step == user->adj_idx) {
223:     return 0;
224:   } else {
225:     VecGetArrayRead(*mu,&sensmu_loc);
226:     VecGetArray(user->sens_mu1,&sensmu1_glob);
227:     VecGetArray(user->sens_mu2,&sensmu2_glob);
228:     sensmu1_glob[curr_step] = sensmu_loc[0];
229:     sensmu2_glob[curr_step] = sensmu_loc[1];
230:     VecRestoreArray(user->sens_mu1,&sensmu1_glob);
231:     VecRestoreArray(user->sens_mu2,&sensmu2_glob);
232:     VecRestoreArrayRead(*mu,&sensmu_loc);
233:     return 0;
234:   }
235: }

237: int main(int argc,char **argv)
238: {
239:   TS             ts;
240:   AppCtx         user;
241:   PetscScalar    *x_ptr,*y_ptr,*u_ptr;
242:   PetscMPIInt    size;
243:   PetscBool      monitor = PETSC_FALSE;
244:   SAMethod       sa = SA_GLOBAL;

247:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
248:      Initialize program
249:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
250:   PetscInitialize(&argc,&argv,NULL,help);
251:   MPI_Comm_size(PETSC_COMM_WORLD,&size);

254:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
255:      Set runtime options
256:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
257:   PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"SA analysis options.","");{
258:   PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL);
259:   PetscOptionsEnum("-sa_method","Sensitivity analysis method (track or global)","",SAMethods,(PetscEnum)sa,(PetscEnum*)&sa,NULL);
260:   }
261:   PetscOptionsEnd();

263:   user.final_time = 0.1;
264:   user.max_steps  = 5;
265:   user.time_step  = user.final_time/user.max_steps;

267:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
268:      Create necessary matrix and vectors for forward solve.
269:      Create Jacp matrix for adjoint solve.
270:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
271:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu1);
272:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu2);
273:   VecSet(user.mu1,1.25);
274:   VecSet(user.mu2,1.0e2);

276:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
277:       For tracking method : create the global sensitivity array to
278:       accumulate sensitivity with respect to parameters at each step.
279:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
280:   if (sa == SA_TRACK) {
281:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu1);
282:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu2);
283:   }

285:   MatCreate(PETSC_COMM_WORLD,&user.A);
286:   MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2);
287:   MatSetFromOptions(user.A);
288:   MatSetUp(user.A);
289:   MatCreateVecs(user.A,&user.U,NULL);

291:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
292:       Note that the dimensions of the Jacp matrix depend upon the
293:       sensitivity analysis method being used !
294:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
295:   MatCreate(PETSC_COMM_WORLD,&user.Jacp);
296:   if (sa == SA_TRACK) {
297:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2);
298:   }
299:   if (sa == SA_GLOBAL) {
300:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,user.max_steps*2);
301:   }
302:   MatSetFromOptions(user.Jacp);
303:   MatSetUp(user.Jacp);

305:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
306:      Create timestepping solver context
307:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
308:   TSCreate(PETSC_COMM_WORLD,&ts);
309:   TSSetEquationType(ts,TS_EQ_ODE_EXPLICIT);
310:   TSSetType(ts,TSCN);

312:   TSSetRHSFunction(ts,NULL,RHSFunction,&user);
313:   TSSetRHSJacobian(ts,user.A,user.A,RHSJacobian,&user);
314:   if (sa == SA_TRACK) {
315:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_track,&user);
316:   }
317:   if (sa == SA_GLOBAL) {
318:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_global,&user);
319:   }

321:   TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
322:   TSSetMaxTime(ts,user.final_time);
323:   TSSetTimeStep(ts,user.final_time/user.max_steps);

325:   if (monitor) {
326:     TSMonitorSet(ts,Monitor,&user,NULL);
327:   }
328:   if (sa == SA_TRACK) {
329:     TSAdjointMonitorSet(ts,AdjointMonitor,&user,NULL);
330:   }

332:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
333:      Set initial conditions
334:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
335:   VecGetArray(user.U,&x_ptr);
336:   x_ptr[0] = 2.0;
337:   x_ptr[1] = -2.0/3.0;
338:   VecRestoreArray(user.U,&x_ptr);

340:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
341:      Save trajectory of solution so that TSAdjointSolve() may be used
342:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
343:   TSSetSaveTrajectory(ts);

345:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
346:      Set runtime options
347:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
348:   TSSetFromOptions(ts);

350:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
351:      Execute forward model and print solution.
352:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
353:   TSSolve(ts,user.U);
354:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution of forward TS :\n");
355:   VecView(user.U,PETSC_VIEWER_STDOUT_WORLD);
356:   PetscPrintf(PETSC_COMM_WORLD,"\n Forward TS solve successful! Adjoint run begins!\n");

358:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
359:      Adjoint model starts here! Create adjoint vectors.
360:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
361:   MatCreateVecs(user.A,&user.lambda,NULL);
362:   MatCreateVecs(user.Jacp,&user.mup,NULL);

364:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
365:      Set initial conditions for the adjoint vector
366:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
367:   VecGetArray(user.U,&u_ptr);
368:   VecGetArray(user.lambda,&y_ptr);
369:   y_ptr[0] = 2*(u_ptr[0] - 1.5967);
370:   y_ptr[1] = 2*(u_ptr[1] - -(1.02969));
371:   VecRestoreArray(user.lambda,&y_ptr);
372:   VecRestoreArray(user.U,&y_ptr);
373:   VecSet(user.mup,0);

375:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
376:      Set number of cost functions.
377:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
378:   TSSetCostGradients(ts,1,&user.lambda,&user.mup);

380:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
381:      The adjoint vector mup has to be reset for each adjoint step when
382:      using the tracking method as we want to treat the parameters at each
383:      time step one at a time and prevent accumulation of the sensitivities
384:      from parameters at previous time steps.
385:      This is not necessary for the global method as each time dependent
386:      parameter is treated as an independent parameter.
387:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
388:   if (sa == SA_TRACK) {
389:     for (user.adj_idx=user.max_steps; user.adj_idx>0; user.adj_idx--) {
390:       VecSet(user.mup,0);
391:       TSAdjointSetSteps(ts, 1);
392:       TSAdjointSolve(ts);
393:     }
394:   }
395:   if (sa == SA_GLOBAL) {
396:     TSAdjointSolve(ts);
397:   }

399:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
400:      Dispaly adjoint sensitivities wrt parameters and initial conditions
401:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
402:   if (sa == SA_TRACK) {
403:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu1: d[cost]/d[mu1]\n");
404:     VecView(user.sens_mu1,PETSC_VIEWER_STDOUT_WORLD);
405:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu2: d[cost]/d[mu2]\n");
406:     VecView(user.sens_mu2,PETSC_VIEWER_STDOUT_WORLD);
407:   }

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

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

418:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
419:      Free work space!
420:      All PETSc objects should be destroyed when they are no longer needed.
421:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
422:   MatDestroy(&user.A);
423:   MatDestroy(&user.Jacp);
424:   VecDestroy(&user.U);
425:   VecDestroy(&user.lambda);
426:   VecDestroy(&user.mup);
427:   VecDestroy(&user.mu1);
428:   VecDestroy(&user.mu2);
429:   if (sa == SA_TRACK) {
430:     VecDestroy(&user.sens_mu1);
431:     VecDestroy(&user.sens_mu2);
432:   }
433:   TSDestroy(&ts);
434:   PetscFinalize();
435:   return(ierr);
436: }

438: /*TEST

440:   test:
441:     requires: !complex
442:     suffix : track
443:     args : -sa_method track

445:   test:
446:     requires: !complex
447:     suffix : global
448:     args : -sa_method global

450: TEST*/