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 User user = (User)ctx; 31 const PetscScalar *x, *xdot; 32 PetscScalar *f; 33 34 PetscFunctionBeginUser; 35 PetscCall(VecGetArrayRead(X, &x)); 36 PetscCall(VecGetArrayRead(Xdot, &xdot)); 37 PetscCall(VecGetArrayWrite(F, &f)); 38 f[0] = user->c * xdot[0] - user->b * x[0]; 39 PetscCall(VecRestoreArrayRead(X, &x)); 40 PetscCall(VecRestoreArrayRead(Xdot, &xdot)); 41 PetscCall(VecRestoreArrayWrite(F, &f)); 42 PetscFunctionReturn(0); 43 } 44 45 static PetscErrorCode IJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A, Mat B, void *ctx) { 46 User user = (User)ctx; 47 PetscInt rowcol[] = {0}; 48 PetscScalar J[1][1]; 49 const PetscScalar *x; 50 51 PetscFunctionBeginUser; 52 PetscCall(VecGetArrayRead(X, &x)); 53 J[0][0] = user->c * a - user->b * 1.0; 54 PetscCall(MatSetValues(B, 1, rowcol, 1, rowcol, &J[0][0], INSERT_VALUES)); 55 PetscCall(VecRestoreArrayRead(X, &x)); 56 57 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 58 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 59 if (A != B) { 60 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY)); 61 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY)); 62 } 63 PetscFunctionReturn(0); 64 } 65 66 static PetscErrorCode IJacobianP(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal shift, Mat A, void *ctx) { 67 User user = (User)ctx; 68 PetscInt row[] = {0}, col[] = {0}; 69 PetscScalar J[1][1]; 70 const PetscScalar *x, *xdot; 71 PetscReal dt; 72 73 PetscFunctionBeginUser; 74 PetscCall(VecGetArrayRead(X, &x)); 75 PetscCall(VecGetArrayRead(Xdot, &xdot)); 76 PetscCall(TSGetTimeStep(ts, &dt)); 77 if (user->der == 1) J[0][0] = xdot[0]; 78 if (user->der == 2) J[0][0] = -x[0]; 79 PetscCall(MatSetValues(A, 1, row, 1, col, &J[0][0], INSERT_VALUES)); 80 PetscCall(VecRestoreArrayRead(X, &x)); 81 82 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 83 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 84 PetscFunctionReturn(0); 85 } 86 87 int main(int argc, char **argv) { 88 TS ts; 89 PetscScalar *x_ptr; 90 PetscMPIInt size; 91 struct _n_User user; 92 PetscInt rows, cols; 93 94 PetscFunctionBeginUser; 95 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 96 97 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 98 PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!"); 99 100 user.a = 2.0; 101 user.b = 4.0; 102 user.c = 3.0; 103 user.steps = 0; 104 user.ftime = 1.0; 105 user.der = 1; 106 PetscCall(PetscOptionsGetInt(NULL, NULL, "-der", &user.der, NULL)); 107 108 rows = 1; 109 cols = 1; 110 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.Jac)); 111 PetscCall(MatSetSizes(user.Jac, PETSC_DECIDE, PETSC_DECIDE, 1, 1)); 112 PetscCall(MatSetFromOptions(user.Jac)); 113 PetscCall(MatSetUp(user.Jac)); 114 PetscCall(MatCreateVecs(user.Jac, &user.x, NULL)); 115 116 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts)); 117 PetscCall(TSSetType(ts, TSBEULER)); 118 PetscCall(TSSetIFunction(ts, NULL, IFunction, &user)); 119 PetscCall(TSSetIJacobian(ts, user.Jac, user.Jac, IJacobian, &user)); 120 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 121 PetscCall(TSSetMaxTime(ts, user.ftime)); 122 123 PetscCall(VecGetArrayWrite(user.x, &x_ptr)); 124 x_ptr[0] = user.a; 125 PetscCall(VecRestoreArrayWrite(user.x, &x_ptr)); 126 PetscCall(TSSetTimeStep(ts, 0.001)); 127 128 /* Set up forward sensitivity */ 129 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.Jacp)); 130 PetscCall(MatSetSizes(user.Jacp, PETSC_DECIDE, PETSC_DECIDE, rows, cols)); 131 PetscCall(MatSetFromOptions(user.Jacp)); 132 PetscCall(MatSetUp(user.Jacp)); 133 PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, rows, cols, NULL, &user.sp)); 134 PetscCall(MatZeroEntries(user.sp)); 135 PetscCall(TSForwardSetSensitivities(ts, cols, user.sp)); 136 PetscCall(TSSetIJacobianP(ts, user.Jacp, IJacobianP, &user)); 137 138 PetscCall(TSSetSaveTrajectory(ts)); 139 PetscCall(TSSetFromOptions(ts)); 140 141 PetscCall(TSSolve(ts, user.x)); 142 PetscCall(TSGetSolveTime(ts, &user.ftime)); 143 PetscCall(TSGetStepNumber(ts, &user.steps)); 144 PetscCall(VecGetArray(user.x, &x_ptr)); 145 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n ode solution %g\n", (double)PetscRealPart(x_ptr[0]))); 146 PetscCall(VecRestoreArray(user.x, &x_ptr)); 147 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n analytical solution %g\n", (double)(user.a * PetscExpReal(user.b / user.c * user.ftime)))); 148 149 if (user.der == 1) { PetscCall(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)))); } 150 if (user.der == 2) { PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n analytical derivative w.r.t. b %g\n", (double)(user.a * user.ftime / user.c * PetscExpReal(user.b / user.c * user.ftime)))); } 151 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n forward sensitivity:\n")); 152 PetscCall(MatView(user.sp, PETSC_VIEWER_STDOUT_WORLD)); 153 154 PetscCall(MatCreateVecs(user.Jac, &user.lambda[0], NULL)); 155 /* Set initial conditions for the adjoint integration */ 156 PetscCall(VecGetArrayWrite(user.lambda[0], &x_ptr)); 157 x_ptr[0] = 1.0; 158 PetscCall(VecRestoreArrayWrite(user.lambda[0], &x_ptr)); 159 PetscCall(MatCreateVecs(user.Jacp, &user.mup[0], NULL)); 160 PetscCall(VecGetArrayWrite(user.mup[0], &x_ptr)); 161 x_ptr[0] = 0.0; 162 PetscCall(VecRestoreArrayWrite(user.mup[0], &x_ptr)); 163 164 PetscCall(TSSetCostGradients(ts, 1, user.lambda, user.mup)); 165 PetscCall(TSAdjointSolve(ts)); 166 167 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n adjoint sensitivity:\n")); 168 PetscCall(VecView(user.mup[0], PETSC_VIEWER_STDOUT_WORLD)); 169 170 PetscCall(MatDestroy(&user.Jac)); 171 PetscCall(MatDestroy(&user.sp)); 172 PetscCall(MatDestroy(&user.Jacp)); 173 PetscCall(VecDestroy(&user.x)); 174 PetscCall(VecDestroy(&user.lambda[0])); 175 PetscCall(VecDestroy(&user.mup[0])); 176 PetscCall(TSDestroy(&ts)); 177 178 PetscCall(PetscFinalize()); 179 return 0; 180 } 181 182 /*TEST 183 184 test: 185 args: -ts_type beuler 186 187 test: 188 suffix: 2 189 args: -ts_type cn 190 output_file: output/ex23fwdadj_1.out 191 192 TEST*/ 193