Date: 2022-03-04T17:39:38+00:00

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:   Concepts: TS^adjoint for time dependent parameters
     10:   Concepts: TS^Customized adjoint monitor based sensitivity tracking
     11:   Concepts: TS^All at once approach to sensitivity tracking
     12:   Processors: 1
     13: */

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    277:   user.final_time = 0.1;
    278:   user.max_steps  = 5;
    279:   user.time_step  = user.final_time/user.max_steps;

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

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

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

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

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

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

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

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

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

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

    359:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    360:      Set runtime options
    361:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    362:   TSSetFromOptions(ts);

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

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

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

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

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

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

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

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

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

    452: /*TEST

    454:   test:
    455:     requires: !complex
    456:     suffix : track
    457:     args : -sa_method track

    459:   test:
    460:     requires: !complex
    461:     suffix : global
    462:     args : -sa_method global

    464: TEST*/