Actual source code: ex23fwdadj.c

      1: static char help[] = "A toy example for testing forward and adjoint sensitivity analysis of an implicit ODE with a paramerized mass matrice.\n";

      3: /*
      4:   This example solves the simple ODE
      5:     c x' = b x, x(0) = a,
      6:   whose analytical solution is x(T)=a*exp(b/c*T), and calculates the derivative of x(T) w.r.t. c (by default) or w.r.t. b (can be enabled with command line option -der 2).

      8: */

     10: #include <petscts.h>

     12: typedef struct _n_User *User;
     13: struct _n_User {
     14:   PetscReal a;
     15:   PetscReal b;
     16:   PetscReal c;
     17:   /* Sensitivity analysis support */
     18:   PetscInt  steps;
     19:   PetscReal ftime;
     20:   Mat       Jac;                    /* Jacobian matrix */
     21:   Mat       Jacp;                   /* JacobianP matrix */
     22:   Vec       x;
     23:   Mat       sp;                     /* forward sensitivity variables */
     24:   Vec       lambda[1];              /* adjoint sensitivity variables */
     25:   Vec       mup[1];                 /* adjoint sensitivity variables */
     26:   PetscInt  der;
     27: };

     29: static PetscErrorCode IFunction(TS ts,PetscReal t,Vec X,Vec Xdot,Vec F,void *ctx)
     30: {
     31:   User              user = (User)ctx;
     32:   const PetscScalar *x,*xdot;
     33:   PetscScalar       *f;

     36:   VecGetArrayRead(X,&x);
     37:   VecGetArrayRead(Xdot,&xdot);
     38:   VecGetArrayWrite(F,&f);
     39:   f[0] = user->c*xdot[0] - user->b*x[0];
     40:   VecRestoreArrayRead(X,&x);
     41:   VecRestoreArrayRead(Xdot,&xdot);
     42:   VecRestoreArrayWrite(F,&f);
     43:   return 0;
     44: }

     46: static PetscErrorCode IJacobian(TS ts,PetscReal t,Vec X,Vec Xdot,PetscReal a,Mat A,Mat B,void *ctx)
     47: {
     48:   User              user     = (User)ctx;
     49:   PetscInt          rowcol[] = {0};
     50:   PetscScalar       J[1][1];
     51:   const PetscScalar *x;

     54:   VecGetArrayRead(X,&x);
     55:   J[0][0] = user->c*a - user->b*1.0;
     56:   MatSetValues(B,1,rowcol,1,rowcol,&J[0][0],INSERT_VALUES);
     57:   VecRestoreArrayRead(X,&x);

     59:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
     60:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
     61:   if (A != B) {
     62:     MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);
     63:     MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);
     64:   }
     65:   return 0;
     66: }

     68: static PetscErrorCode IJacobianP(TS ts,PetscReal t,Vec X,Vec Xdot,PetscReal shift,Mat A,void *ctx)
     69: {
     70:   User              user = (User)ctx;
     71:   PetscInt          row[] = {0},col[]={0};
     72:   PetscScalar       J[1][1];
     73:   const PetscScalar *x,*xdot;
     74:   PetscReal         dt;

     77:   VecGetArrayRead(X,&x);
     78:   VecGetArrayRead(Xdot,&xdot);
     79:   TSGetTimeStep(ts,&dt);
     80:   if (user->der == 1) J[0][0] = xdot[0];
     81:   if (user->der == 2) J[0][0] = -x[0];
     82:   MatSetValues(A,1,row,1,col,&J[0][0],INSERT_VALUES);
     83:   VecRestoreArrayRead(X,&x);

     85:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
     86:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
     87:   return 0;
     88: }

     90: int main(int argc,char **argv)
     91: {
     92:   TS             ts;
     93:   PetscScalar    *x_ptr;
     94:   PetscMPIInt    size;
     95:   struct _n_User user;
     96:   PetscInt       rows,cols;

     98:   PetscInitialize(&argc,&argv,NULL,help);

    100:   MPI_Comm_size(PETSC_COMM_WORLD,&size);

    103:   user.a           = 2.0;
    104:   user.b           = 4.0;
    105:   user.c           = 3.0;
    106:   user.steps       = 0;
    107:   user.ftime       = 1.0;
    108:   user.der         = 1;
    109:   PetscOptionsGetInt(NULL,NULL,"-der",&user.der,NULL);

    111:   rows = 1;
    112:   cols = 1;
    113:   MatCreate(PETSC_COMM_WORLD,&user.Jac);
    114:   MatSetSizes(user.Jac,PETSC_DECIDE,PETSC_DECIDE,1,1);
    115:   MatSetFromOptions(user.Jac);
    116:   MatSetUp(user.Jac);
    117:   MatCreateVecs(user.Jac,&user.x,NULL);

    119:   TSCreate(PETSC_COMM_WORLD,&ts);
    120:   TSSetType(ts,TSBEULER);
    121:   TSSetIFunction(ts,NULL,IFunction,&user);
    122:   TSSetIJacobian(ts,user.Jac,user.Jac,IJacobian,&user);
    123:   TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
    124:   TSSetMaxTime(ts,user.ftime);

    126:   VecGetArrayWrite(user.x,&x_ptr);
    127:   x_ptr[0] = user.a;
    128:   VecRestoreArrayWrite(user.x,&x_ptr);
    129:   TSSetTimeStep(ts,0.001);

    131:   /* Set up forward sensitivity */
    132:   MatCreate(PETSC_COMM_WORLD,&user.Jacp);
    133:   MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,rows,cols);
    134:   MatSetFromOptions(user.Jacp);
    135:   MatSetUp(user.Jacp);
    136:   MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,rows,cols,NULL,&user.sp);
    137:   MatZeroEntries(user.sp);
    138:   TSForwardSetSensitivities(ts,cols,user.sp);
    139:   TSSetIJacobianP(ts,user.Jacp,IJacobianP,&user);

    141:   TSSetSaveTrajectory(ts);
    142:   TSSetFromOptions(ts);

    144:   TSSolve(ts,user.x);
    145:   TSGetSolveTime(ts,&user.ftime);
    146:   TSGetStepNumber(ts,&user.steps);
    147:   VecGetArray(user.x,&x_ptr);
    148:   PetscPrintf(PETSC_COMM_WORLD,"\n ode solution %g\n",(double)PetscRealPart(x_ptr[0]));
    149:   VecRestoreArray(user.x,&x_ptr);
    150:   PetscPrintf(PETSC_COMM_WORLD,"\n analytical solution %g\n",(double)user.a*PetscExpReal(user.b/user.c*user.ftime));

    152:   if (user.der == 1) {
    153:     PetscPrintf(PETSC_COMM_WORLD,"\n analytical derivative w.r.t. c %g\n",(double)-user.a*user.ftime*user.b/(user.c*user.c)*PetscExpReal(user.b/user.c*user.ftime));
    154:   }
    155:   if (user.der == 2) {
    156:     PetscPrintf(PETSC_COMM_WORLD,"\n analytical derivative w.r.t. b %g\n",user.a*user.ftime/user.c*PetscExpReal(user.b/user.c*user.ftime));
    157:   }
    158:   PetscPrintf(PETSC_COMM_WORLD,"\n forward sensitivity:\n");
    159:   MatView(user.sp,PETSC_VIEWER_STDOUT_WORLD);

    161:   MatCreateVecs(user.Jac,&user.lambda[0],NULL);
    162:   /* Set initial conditions for the adjoint integration */
    163:   VecGetArrayWrite(user.lambda[0],&x_ptr);
    164:   x_ptr[0] = 1.0;
    165:   VecRestoreArrayWrite(user.lambda[0],&x_ptr);
    166:   MatCreateVecs(user.Jacp,&user.mup[0],NULL);
    167:   VecGetArrayWrite(user.mup[0],&x_ptr);
    168:   x_ptr[0] = 0.0;
    169:   VecRestoreArrayWrite(user.mup[0],&x_ptr);

    171:   TSSetCostGradients(ts,1,user.lambda,user.mup);
    172:   TSAdjointSolve(ts);

    174:   PetscPrintf(PETSC_COMM_WORLD,"\n adjoint sensitivity:\n");
    175:   VecView(user.mup[0],PETSC_VIEWER_STDOUT_WORLD);

    177:   MatDestroy(&user.Jac);
    178:   MatDestroy(&user.sp);
    179:   MatDestroy(&user.Jacp);
    180:   VecDestroy(&user.x);
    181:   VecDestroy(&user.lambda[0]);
    182:   VecDestroy(&user.mup[0]);
    183:   TSDestroy(&ts);

    185:   PetscFinalize();
    186:   return 0;
    187: }

    189: /*TEST

    191:     test:
    192:       args: -ts_type beuler

    194:     test:
    195:       suffix: 2
    196:       args: -ts_type cn
    197:       output_file: output/ex23fwdadj_1.out

    199: TEST*/