Date: 2022-03-31T04:48:46+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:   AppCtx            *user = (AppCtx*) ctx;
    105:   PetscScalar       *f;
    106:   PetscInt          curr_step;
    107:   const PetscScalar *u;
    108:   const PetscScalar *mu1;
    109:   const PetscScalar *mu2;

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

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

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

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

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

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

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

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

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

    200: /* Dump solution to console if called */
    201: PetscErrorCode Monitor(TS ts,PetscInt step,PetscReal t,Vec U,void *ctx)
    202: {
    204:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution at time %e is \n", t);
    205:   VecView(U,PETSC_VIEWER_STDOUT_WORLD);
    206:   return 0;
    207: }

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

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

    244: int main(int argc,char **argv)
    245: {
    246:   TS             ts;
    247:   AppCtx         user;
    248:   PetscScalar    *x_ptr,*y_ptr,*u_ptr;
    249:   PetscMPIInt    size;
    250:   PetscBool      monitor = PETSC_FALSE;
    251:   SAMethod       sa = SA_GLOBAL;

    254:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    255:      Initialize program
    256:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    257:   PetscInitialize(&argc,&argv,NULL,help);
    258:   MPI_Comm_size(PETSC_COMM_WORLD,&size);

    261:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    262:      Set runtime options
    263:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    264:   PetscOptionsBegin(PETSC_COMM_WORLD,NULL,"SA analysis options.","");{
    265:   PetscOptionsGetBool(NULL,NULL,"-monitor",&monitor,NULL);
    266:   PetscOptionsEnum("-sa_method","Sensitivity analysis method (track or global)","",SAMethods,(PetscEnum)sa,(PetscEnum*)&sa,NULL);
    267:   }
    268:   PetscOptionsEnd();

    270:   user.final_time = 0.1;
    271:   user.max_steps  = 5;
    272:   user.time_step  = user.final_time/user.max_steps;

    274:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    275:      Create necessary matrix and vectors for forward solve.
    276:      Create Jacp matrix for adjoint solve.
    277:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    278:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu1);
    279:   VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.mu2);
    280:   VecSet(user.mu1,1.25);
    281:   VecSet(user.mu2,1.0e2);

    283:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    284:       For tracking method : create the global sensitivity array to
    285:       accumulate sensitivity with respect to parameters at each step.
    286:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    287:   if (sa == SA_TRACK) {
    288:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu1);
    289:     VecCreateSeq(PETSC_COMM_WORLD,user.max_steps,&user.sens_mu2);
    290:   }

    292:   MatCreate(PETSC_COMM_WORLD,&user.A);
    293:   MatSetSizes(user.A,PETSC_DECIDE,PETSC_DECIDE,2,2);
    294:   MatSetFromOptions(user.A);
    295:   MatSetUp(user.A);
    296:   MatCreateVecs(user.A,&user.U,NULL);

    298:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    299:       Note that the dimensions of the Jacp matrix depend upon the
    300:       sensitivity analysis method being used !
    301:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    302:   MatCreate(PETSC_COMM_WORLD,&user.Jacp);
    303:   if (sa == SA_TRACK) {
    304:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,2);
    305:   }
    306:   if (sa == SA_GLOBAL) {
    307:     MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,2,user.max_steps*2);
    308:   }
    309:   MatSetFromOptions(user.Jacp);
    310:   MatSetUp(user.Jacp);

    312:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    313:      Create timestepping solver context
    314:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    315:   TSCreate(PETSC_COMM_WORLD,&ts);
    316:   TSSetEquationType(ts,TS_EQ_ODE_EXPLICIT);
    317:   TSSetType(ts,TSCN);

    319:   TSSetRHSFunction(ts,NULL,RHSFunction,&user);
    320:   TSSetRHSJacobian(ts,user.A,user.A,RHSJacobian,&user);
    321:   if (sa == SA_TRACK) {
    322:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_track,&user);
    323:   }
    324:   if (sa == SA_GLOBAL) {
    325:     TSSetRHSJacobianP(ts,user.Jacp,RHSJacobianP_global,&user);
    326:   }

    328:   TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
    329:   TSSetMaxTime(ts,user.final_time);
    330:   TSSetTimeStep(ts,user.final_time/user.max_steps);

    332:   if (monitor) {
    333:     TSMonitorSet(ts,Monitor,&user,NULL);
    334:   }
    335:   if (sa == SA_TRACK) {
    336:     TSAdjointMonitorSet(ts,AdjointMonitor,&user,NULL);
    337:   }

    339:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    340:      Set initial conditions
    341:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    342:   VecGetArray(user.U,&x_ptr);
    343:   x_ptr[0] = 2.0;
    344:   x_ptr[1] = -2.0/3.0;
    345:   VecRestoreArray(user.U,&x_ptr);

    347:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    348:      Save trajectory of solution so that TSAdjointSolve() may be used
    349:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    350:   TSSetSaveTrajectory(ts);

    352:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    353:      Set runtime options
    354:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    355:   TSSetFromOptions(ts);

    357:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    358:      Execute forward model and print solution.
    359:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    360:   TSSolve(ts,user.U);
    361:   PetscPrintf(PETSC_COMM_WORLD,"\n Solution of forward TS :\n");
    362:   VecView(user.U,PETSC_VIEWER_STDOUT_WORLD);
    363:   PetscPrintf(PETSC_COMM_WORLD,"\n Forward TS solve successfull! Adjoint run begins!\n");

    365:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    366:      Adjoint model starts here! Create adjoint vectors.
    367:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    368:   MatCreateVecs(user.A,&user.lambda,NULL);
    369:   MatCreateVecs(user.Jacp,&user.mup,NULL);

    371:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    372:      Set initial conditions for the adjoint vector
    373:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    374:   VecGetArray(user.U,&u_ptr);
    375:   VecGetArray(user.lambda,&y_ptr);
    376:   y_ptr[0] = 2*(u_ptr[0] - 1.5967);
    377:   y_ptr[1] = 2*(u_ptr[1] - -(1.02969));
    378:   VecRestoreArray(user.lambda,&y_ptr);
    379:   VecRestoreArray(user.U,&y_ptr);
    380:   VecSet(user.mup,0);

    382:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    383:      Set number of cost functions.
    384:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    385:   TSSetCostGradients(ts,1,&user.lambda,&user.mup);

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

    406:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    407:      Dispaly adjoint sensitivities wrt parameters and initial conditions
    408:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    409:   if (sa == SA_TRACK) {
    410:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu1: d[cost]/d[mu1]\n");
    411:     VecView(user.sens_mu1,PETSC_VIEWER_STDOUT_WORLD);
    412:     PetscPrintf(PETSC_COMM_WORLD,"\n sensitivity wrt  mu2: d[cost]/d[mu2]\n");
    413:     VecView(user.sens_mu2,PETSC_VIEWER_STDOUT_WORLD);
    414:   }

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

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

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

    445: /*TEST

    447:   test:
    448:     requires: !complex
    449:     suffix : track
    450:     args : -sa_method track

    452:   test:
    453:     requires: !complex
    454:     suffix : global
    455:     args : -sa_method global

    457: TEST*/