1 static char help[] = "A toy example for testing forward and adjoint sensitivity analysis of an implicit ODE with a paramerized mass matrice.\n"; 2 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). 7 8 */ 9 10 #include <petscts.h> 11 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 }; 28 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; 34 35 PetscFunctionBeginUser; 36 CHKERRQ(VecGetArrayRead(X,&x)); 37 CHKERRQ(VecGetArrayRead(Xdot,&xdot)); 38 CHKERRQ(VecGetArrayWrite(F,&f)); 39 f[0] = user->c*xdot[0] - user->b*x[0]; 40 CHKERRQ(VecRestoreArrayRead(X,&x)); 41 CHKERRQ(VecRestoreArrayRead(Xdot,&xdot)); 42 CHKERRQ(VecRestoreArrayWrite(F,&f)); 43 PetscFunctionReturn(0); 44 } 45 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; 52 53 PetscFunctionBeginUser; 54 CHKERRQ(VecGetArrayRead(X,&x)); 55 J[0][0] = user->c*a - user->b*1.0; 56 CHKERRQ(MatSetValues(B,1,rowcol,1,rowcol,&J[0][0],INSERT_VALUES)); 57 CHKERRQ(VecRestoreArrayRead(X,&x)); 58 59 CHKERRQ(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY)); 60 CHKERRQ(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY)); 61 if (A != B) { 62 CHKERRQ(MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY)); 63 CHKERRQ(MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY)); 64 } 65 PetscFunctionReturn(0); 66 } 67 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; 75 76 PetscFunctionBeginUser; 77 CHKERRQ(VecGetArrayRead(X,&x)); 78 CHKERRQ(VecGetArrayRead(Xdot,&xdot)); 79 CHKERRQ(TSGetTimeStep(ts,&dt)); 80 if (user->der == 1) J[0][0] = xdot[0]; 81 if (user->der == 2) J[0][0] = -x[0]; 82 CHKERRQ(MatSetValues(A,1,row,1,col,&J[0][0],INSERT_VALUES)); 83 CHKERRQ(VecRestoreArrayRead(X,&x)); 84 85 CHKERRQ(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY)); 86 CHKERRQ(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY)); 87 PetscFunctionReturn(0); 88 } 89 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; 97 PetscErrorCode ierr; 98 99 ierr = PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr; 100 101 CHKERRMPI(MPI_Comm_size(PETSC_COMM_WORLD,&size)); 102 PetscCheck(size == 1,PETSC_COMM_WORLD,PETSC_ERR_WRONG_MPI_SIZE,"This is a uniprocessor example only!"); 103 104 user.a = 2.0; 105 user.b = 4.0; 106 user.c = 3.0; 107 user.steps = 0; 108 user.ftime = 1.0; 109 user.der = 1; 110 CHKERRQ(PetscOptionsGetInt(NULL,NULL,"-der",&user.der,NULL)); 111 112 rows = 1; 113 cols = 1; 114 CHKERRQ(MatCreate(PETSC_COMM_WORLD,&user.Jac)); 115 CHKERRQ(MatSetSizes(user.Jac,PETSC_DECIDE,PETSC_DECIDE,1,1)); 116 CHKERRQ(MatSetFromOptions(user.Jac)); 117 CHKERRQ(MatSetUp(user.Jac)); 118 CHKERRQ(MatCreateVecs(user.Jac,&user.x,NULL)); 119 120 CHKERRQ(TSCreate(PETSC_COMM_WORLD,&ts)); 121 CHKERRQ(TSSetType(ts,TSBEULER)); 122 CHKERRQ(TSSetIFunction(ts,NULL,IFunction,&user)); 123 CHKERRQ(TSSetIJacobian(ts,user.Jac,user.Jac,IJacobian,&user)); 124 CHKERRQ(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP)); 125 CHKERRQ(TSSetMaxTime(ts,user.ftime)); 126 127 CHKERRQ(VecGetArrayWrite(user.x,&x_ptr)); 128 x_ptr[0] = user.a; 129 CHKERRQ(VecRestoreArrayWrite(user.x,&x_ptr)); 130 CHKERRQ(TSSetTimeStep(ts,0.001)); 131 132 /* Set up forward sensitivity */ 133 CHKERRQ(MatCreate(PETSC_COMM_WORLD,&user.Jacp)); 134 CHKERRQ(MatSetSizes(user.Jacp,PETSC_DECIDE,PETSC_DECIDE,rows,cols)); 135 CHKERRQ(MatSetFromOptions(user.Jacp)); 136 CHKERRQ(MatSetUp(user.Jacp)); 137 CHKERRQ(MatCreateDense(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,rows,cols,NULL,&user.sp)); 138 CHKERRQ(MatZeroEntries(user.sp)); 139 CHKERRQ(TSForwardSetSensitivities(ts,cols,user.sp)); 140 CHKERRQ(TSSetIJacobianP(ts,user.Jacp,IJacobianP,&user)); 141 142 CHKERRQ(TSSetSaveTrajectory(ts)); 143 CHKERRQ(TSSetFromOptions(ts)); 144 145 CHKERRQ(TSSolve(ts,user.x)); 146 CHKERRQ(TSGetSolveTime(ts,&user.ftime)); 147 CHKERRQ(TSGetStepNumber(ts,&user.steps)); 148 CHKERRQ(VecGetArray(user.x,&x_ptr)); 149 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\n ode solution %g\n",(double)PetscRealPart(x_ptr[0]))); 150 CHKERRQ(VecRestoreArray(user.x,&x_ptr)); 151 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\n analytical solution %g\n",(double)user.a*PetscExpReal(user.b/user.c*user.ftime))); 152 153 if (user.der == 1) { 154 CHKERRQ(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))); 155 } 156 if (user.der == 2) { 157 CHKERRQ(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))); 158 } 159 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\n forward sensitivity:\n")); 160 CHKERRQ(MatView(user.sp,PETSC_VIEWER_STDOUT_WORLD)); 161 162 CHKERRQ(MatCreateVecs(user.Jac,&user.lambda[0],NULL)); 163 /* Set initial conditions for the adjoint integration */ 164 CHKERRQ(VecGetArrayWrite(user.lambda[0],&x_ptr)); 165 x_ptr[0] = 1.0; 166 CHKERRQ(VecRestoreArrayWrite(user.lambda[0],&x_ptr)); 167 CHKERRQ(MatCreateVecs(user.Jacp,&user.mup[0],NULL)); 168 CHKERRQ(VecGetArrayWrite(user.mup[0],&x_ptr)); 169 x_ptr[0] = 0.0; 170 CHKERRQ(VecRestoreArrayWrite(user.mup[0],&x_ptr)); 171 172 CHKERRQ(TSSetCostGradients(ts,1,user.lambda,user.mup)); 173 CHKERRQ(TSAdjointSolve(ts)); 174 175 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"\n adjoint sensitivity:\n")); 176 CHKERRQ(VecView(user.mup[0],PETSC_VIEWER_STDOUT_WORLD)); 177 178 CHKERRQ(MatDestroy(&user.Jac)); 179 CHKERRQ(MatDestroy(&user.sp)); 180 CHKERRQ(MatDestroy(&user.Jacp)); 181 CHKERRQ(VecDestroy(&user.x)); 182 CHKERRQ(VecDestroy(&user.lambda[0])); 183 CHKERRQ(VecDestroy(&user.mup[0])); 184 CHKERRQ(TSDestroy(&ts)); 185 186 ierr = PetscFinalize(); 187 return(ierr); 188 } 189 190 /*TEST 191 192 test: 193 args: -ts_type beuler 194 195 test: 196 suffix: 2 197 args: -ts_type cn 198 output_file: output/ex23fwdadj_1.out 199 200 TEST*/ 201