Actual source code: ex23fwdadj.c

petsc-main 2021-04-20
Report Typos and Errors
  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:   PetscErrorCode    ierr;
 32:   User              user = (User)ctx;
 33:   const PetscScalar *x,*xdot;
 34:   PetscScalar       *f;

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

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

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

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

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

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

 88:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
 89:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
 90:   return(0);
 91: }

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

102:   PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr;

104:   MPI_Comm_size(PETSC_COMM_WORLD,&size);
105:   if (size != 1) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!");

107:   user.a           = 2.0;
108:   user.b           = 4.0;
109:   user.c           = 3.0;
110:   user.steps       = 0;
111:   user.ftime       = 1.0;
112:   user.der         = 1;
113:   PetscOptionsGetInt(NULL,NULL,"-der",&user.der,NULL);

115:   rows = 1;
116:   cols = 1;
117:   MatCreate(PETSC_COMM_WORLD,&user.Jac);
118:   MatSetSizes(user.Jac,PETSC_DECIDE,PETSC_DECIDE,1,1);
119:   MatSetFromOptions(user.Jac);
120:   MatSetUp(user.Jac);
121:   MatCreateVecs(user.Jac,&user.x,NULL);

123:   TSCreate(PETSC_COMM_WORLD,&ts);
124:   TSSetType(ts,TSBEULER);
125:   TSSetIFunction(ts,NULL,IFunction,&user);
126:   TSSetIJacobian(ts,user.Jac,user.Jac,IJacobian,&user);
127:   TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);
128:   TSSetMaxTime(ts,user.ftime);

130:   VecGetArrayWrite(user.x,&x_ptr);
131:   x_ptr[0] = user.a;
132:   VecRestoreArrayWrite(user.x,&x_ptr);
133:   TSSetTimeStep(ts,0.001);

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

145:   TSSetSaveTrajectory(ts);
146:   TSSetFromOptions(ts);

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

156:   if (user.der == 1) {
157:     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));
158:   }
159:   if (user.der == 2) {
160:     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));
161:   }
162:   PetscPrintf(PETSC_COMM_WORLD,"\n forward sensitivity:\n");
163:   MatView(user.sp,PETSC_VIEWER_STDOUT_WORLD);

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

175:   TSSetCostGradients(ts,1,user.lambda,user.mup);
176:   TSAdjointSolve(ts);

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

181:   MatDestroy(&user.Jac);
182:   MatDestroy(&user.sp);
183:   MatDestroy(&user.Jacp);
184:   VecDestroy(&user.x);
185:   VecDestroy(&user.lambda[0]);
186:   VecDestroy(&user.mup[0]);
187:   TSDestroy(&ts);

189:   PetscFinalize();
190:   return(ierr);
191: }

193: /*TEST

195:     test:
196:       args: -ts_type beuler

198:     test:
199:       suffix: 2
200:       args: -ts_type cn
201:       output_file: output/ex23fwdadj_1.out

203: TEST*/