1 static char help[] = "Demonstrates automatic Jacobian generation using ADOL-C for an ODE-constrained optimization problem.\n\ 2 Input parameters include:\n\ 3 -mu : stiffness parameter\n\n"; 4 5 /* 6 REQUIRES configuration of PETSc with option --download-adolc. 7 8 For documentation on ADOL-C, see 9 $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf 10 */ 11 /* ------------------------------------------------------------------------ 12 See ex16opt_ic for a description of the problem being solved. 13 ------------------------------------------------------------------------- */ 14 #include <petsctao.h> 15 #include <petscts.h> 16 #include <petscmat.h> 17 #include "adolc-utils/drivers.cxx" 18 #include <adolc/adolc.h> 19 20 typedef struct _n_User *User; 21 struct _n_User { 22 PetscReal mu; 23 PetscReal next_output; 24 PetscInt steps; 25 26 /* Sensitivity analysis support */ 27 PetscReal ftime, x_ob[2]; 28 Mat A; /* Jacobian matrix */ 29 Vec x, lambda[2]; /* adjoint variables */ 30 31 /* Automatic differentiation support */ 32 AdolcCtx *adctx; 33 }; 34 35 PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *); 36 37 /* 38 'Passive' RHS function, used in residual evaluations during the time integration. 39 */ 40 static PetscErrorCode RHSFunctionPassive(TS ts, PetscReal t, Vec X, Vec F, void *ctx) { 41 User user = (User)ctx; 42 PetscScalar *f; 43 const PetscScalar *x; 44 45 PetscFunctionBeginUser; 46 PetscCall(VecGetArrayRead(X, &x)); 47 PetscCall(VecGetArray(F, &f)); 48 f[0] = x[1]; 49 f[1] = user->mu * (1. - x[0] * x[0]) * x[1] - x[0]; 50 PetscCall(VecRestoreArrayRead(X, &x)); 51 PetscCall(VecRestoreArray(F, &f)); 52 PetscFunctionReturn(0); 53 } 54 55 /* 56 Trace RHS to mark on tape 1 the dependence of f upon x. This tape is used in generating the 57 Jacobian transform. 58 */ 59 static PetscErrorCode RHSFunctionActive(TS ts, PetscReal t, Vec X, Vec F, void *ctx) { 60 User user = (User)ctx; 61 PetscReal mu = user->mu; 62 PetscScalar *f; 63 const PetscScalar *x; 64 65 adouble f_a[2]; /* adouble for dependent variables */ 66 adouble x_a[2]; /* adouble for independent variables */ 67 68 PetscFunctionBeginUser; 69 PetscCall(VecGetArrayRead(X, &x)); 70 PetscCall(VecGetArray(F, &f)); 71 72 trace_on(1); /* Start of active section */ 73 x_a[0] <<= x[0]; 74 x_a[1] <<= x[1]; /* Mark as independent */ 75 f_a[0] = x_a[1]; 76 f_a[1] = mu * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0]; 77 f_a[0] >>= f[0]; 78 f_a[1] >>= f[1]; /* Mark as dependent */ 79 trace_off(1); /* End of active section */ 80 81 PetscCall(VecRestoreArrayRead(X, &x)); 82 PetscCall(VecRestoreArray(F, &f)); 83 PetscFunctionReturn(0); 84 } 85 86 /* 87 Compute the Jacobian w.r.t. x using PETSc-ADOL-C driver. 88 */ 89 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec X, Mat A, Mat B, void *ctx) { 90 User user = (User)ctx; 91 const PetscScalar *x; 92 93 PetscFunctionBeginUser; 94 PetscCall(VecGetArrayRead(X, &x)); 95 PetscCall(PetscAdolcComputeRHSJacobian(1, A, x, user->adctx)); 96 PetscCall(VecRestoreArrayRead(X, &x)); 97 PetscFunctionReturn(0); 98 } 99 100 /* 101 Monitor timesteps and use interpolation to output at integer multiples of 0.1 102 */ 103 static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, void *ctx) { 104 const PetscScalar *x; 105 PetscReal tfinal, dt, tprev; 106 User user = (User)ctx; 107 108 PetscFunctionBeginUser; 109 PetscCall(TSGetTimeStep(ts, &dt)); 110 PetscCall(TSGetMaxTime(ts, &tfinal)); 111 PetscCall(TSGetPrevTime(ts, &tprev)); 112 PetscCall(VecGetArrayRead(X, &x)); 113 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%.1f] %" PetscInt_FMT " TS %.6f (dt = %.6f) X % 12.6e % 12.6e\n", (double)user->next_output, step, (double)t, (double)dt, (double)PetscRealPart(x[0]), (double)PetscRealPart(x[1]))); 114 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "t %.6f (tprev = %.6f) \n", (double)t, (double)tprev)); 115 PetscCall(VecGetArrayRead(X, &x)); 116 PetscFunctionReturn(0); 117 } 118 119 int main(int argc, char **argv) { 120 TS ts = NULL; /* nonlinear solver */ 121 Vec ic, r; 122 PetscBool monitor = PETSC_FALSE; 123 PetscScalar *x_ptr; 124 PetscMPIInt size; 125 struct _n_User user; 126 AdolcCtx *adctx; 127 Tao tao; 128 KSP ksp; 129 PC pc; 130 131 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 132 Initialize program 133 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 134 PetscFunctionBeginUser; 135 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 136 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 137 PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!"); 138 139 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 140 Set runtime options and create AdolcCtx 141 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 142 PetscCall(PetscNew(&adctx)); 143 user.mu = 1.0; 144 user.next_output = 0.0; 145 user.steps = 0; 146 user.ftime = 0.5; 147 adctx->m = 2; 148 adctx->n = 2; 149 adctx->p = 2; 150 user.adctx = adctx; 151 152 PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &user.mu, NULL)); 153 PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL)); 154 155 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 156 Create necessary matrix and vectors, solve same ODE on every process 157 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 158 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.A)); 159 PetscCall(MatSetSizes(user.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2)); 160 PetscCall(MatSetFromOptions(user.A)); 161 PetscCall(MatSetUp(user.A)); 162 PetscCall(MatCreateVecs(user.A, &user.x, NULL)); 163 164 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 165 Set initial conditions 166 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 167 PetscCall(VecGetArray(user.x, &x_ptr)); 168 x_ptr[0] = 2.0; 169 x_ptr[1] = 0.66666654321; 170 PetscCall(VecRestoreArray(user.x, &x_ptr)); 171 172 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 173 Trace just once on each tape and put zeros on Jacobian diagonal 174 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 175 PetscCall(VecDuplicate(user.x, &r)); 176 PetscCall(RHSFunctionActive(ts, 0., user.x, r, &user)); 177 PetscCall(VecSet(r, 0)); 178 PetscCall(MatDiagonalSet(user.A, r, INSERT_VALUES)); 179 PetscCall(VecDestroy(&r)); 180 181 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 182 Create timestepping solver context 183 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 184 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts)); 185 PetscCall(TSSetType(ts, TSRK)); 186 PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, &user)); 187 PetscCall(TSSetRHSJacobian(ts, user.A, user.A, RHSJacobian, &user)); 188 PetscCall(TSSetMaxTime(ts, user.ftime)); 189 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 190 if (monitor) { PetscCall(TSMonitorSet(ts, Monitor, &user, NULL)); } 191 192 PetscCall(TSSetTime(ts, 0.0)); 193 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %g, steps %" PetscInt_FMT ", ftime %g\n", (double)user.mu, user.steps, (double)(user.ftime))); 194 195 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 196 Save trajectory of solution so that TSAdjointSolve() may be used 197 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 198 PetscCall(TSSetSaveTrajectory(ts)); 199 200 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 201 Set runtime options 202 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 203 PetscCall(TSSetFromOptions(ts)); 204 205 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 206 Solve nonlinear system 207 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 208 PetscCall(TSSolve(ts, user.x)); 209 PetscCall(TSGetSolveTime(ts, &(user.ftime))); 210 PetscCall(TSGetStepNumber(ts, &user.steps)); 211 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %g, steps %" PetscInt_FMT ", ftime %g\n", (double)user.mu, user.steps, (double)user.ftime)); 212 213 PetscCall(VecGetArray(user.x, &x_ptr)); 214 user.x_ob[0] = x_ptr[0]; 215 user.x_ob[1] = x_ptr[1]; 216 PetscCall(VecRestoreArray(user.x, &x_ptr)); 217 218 PetscCall(MatCreateVecs(user.A, &user.lambda[0], NULL)); 219 220 /* Create TAO solver and set desired solution method */ 221 PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao)); 222 PetscCall(TaoSetType(tao, TAOCG)); 223 224 /* Set initial solution guess */ 225 PetscCall(MatCreateVecs(user.A, &ic, NULL)); 226 PetscCall(VecGetArray(ic, &x_ptr)); 227 x_ptr[0] = 2.1; 228 x_ptr[1] = 0.7; 229 PetscCall(VecRestoreArray(ic, &x_ptr)); 230 231 PetscCall(TaoSetSolution(tao, ic)); 232 233 /* Set routine for function and gradient evaluation */ 234 PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&user)); 235 236 /* Check for any TAO command line options */ 237 PetscCall(TaoSetFromOptions(tao)); 238 PetscCall(TaoGetKSP(tao, &ksp)); 239 if (ksp) { 240 PetscCall(KSPGetPC(ksp, &pc)); 241 PetscCall(PCSetType(pc, PCNONE)); 242 } 243 244 PetscCall(TaoSetTolerances(tao, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT)); 245 246 /* SOLVE THE APPLICATION */ 247 PetscCall(TaoSolve(tao)); 248 249 /* Free TAO data structures */ 250 PetscCall(TaoDestroy(&tao)); 251 252 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 253 Free work space. All PETSc objects should be destroyed when they 254 are no longer needed. 255 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 256 PetscCall(MatDestroy(&user.A)); 257 PetscCall(VecDestroy(&user.x)); 258 PetscCall(VecDestroy(&user.lambda[0])); 259 PetscCall(TSDestroy(&ts)); 260 PetscCall(VecDestroy(&ic)); 261 PetscCall(PetscFree(adctx)); 262 PetscCall(PetscFinalize()); 263 return 0; 264 } 265 266 /* ------------------------------------------------------------------ */ 267 /* 268 FormFunctionGradient - Evaluates the function and corresponding gradient. 269 270 Input Parameters: 271 tao - the Tao context 272 X - the input vector 273 ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient() 274 275 Output Parameters: 276 f - the newly evaluated function 277 G - the newly evaluated gradient 278 */ 279 PetscErrorCode FormFunctionGradient(Tao tao, Vec IC, PetscReal *f, Vec G, void *ctx) { 280 User user = (User)ctx; 281 TS ts; 282 PetscScalar *x_ptr, *y_ptr; 283 284 PetscFunctionBeginUser; 285 PetscCall(VecCopy(IC, user->x)); 286 287 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 288 Create timestepping solver context 289 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 290 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts)); 291 PetscCall(TSSetType(ts, TSRK)); 292 PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, user)); 293 /* Set RHS Jacobian for the adjoint integration */ 294 PetscCall(TSSetRHSJacobian(ts, user->A, user->A, RHSJacobian, user)); 295 296 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 297 Set time 298 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 299 PetscCall(TSSetTime(ts, 0.0)); 300 PetscCall(TSSetTimeStep(ts, .001)); 301 PetscCall(TSSetMaxTime(ts, 0.5)); 302 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 303 304 PetscCall(TSSetTolerances(ts, 1e-7, NULL, 1e-7, NULL)); 305 306 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 307 Save trajectory of solution so that TSAdjointSolve() may be used 308 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 309 PetscCall(TSSetSaveTrajectory(ts)); 310 311 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 312 Set runtime options 313 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 314 PetscCall(TSSetFromOptions(ts)); 315 316 PetscCall(TSSolve(ts, user->x)); 317 PetscCall(TSGetSolveTime(ts, &user->ftime)); 318 PetscCall(TSGetStepNumber(ts, &user->steps)); 319 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %.6f, steps %" PetscInt_FMT ", ftime %g\n", (double)user->mu, user->steps, (double)user->ftime)); 320 321 PetscCall(VecGetArray(user->x, &x_ptr)); 322 *f = (x_ptr[0] - user->x_ob[0]) * (x_ptr[0] - user->x_ob[0]) + (x_ptr[1] - user->x_ob[1]) * (x_ptr[1] - user->x_ob[1]); 323 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Observed value y_ob=[%f; %f], ODE solution y=[%f;%f], Cost function f=%f\n", (double)user->x_ob[0], (double)user->x_ob[1], (double)x_ptr[0], (double)x_ptr[1], (double)(*f))); 324 325 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 326 Adjoint model starts here 327 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 328 /* Redet initial conditions for the adjoint integration */ 329 PetscCall(VecGetArray(user->lambda[0], &y_ptr)); 330 y_ptr[0] = 2. * (x_ptr[0] - user->x_ob[0]); 331 y_ptr[1] = 2. * (x_ptr[1] - user->x_ob[1]); 332 PetscCall(VecRestoreArray(user->lambda[0], &y_ptr)); 333 PetscCall(VecRestoreArray(user->x, &x_ptr)); 334 PetscCall(TSSetCostGradients(ts, 1, user->lambda, NULL)); 335 336 PetscCall(TSAdjointSolve(ts)); 337 338 PetscCall(VecCopy(user->lambda[0], G)); 339 340 PetscCall(TSDestroy(&ts)); 341 PetscFunctionReturn(0); 342 } 343 344 /*TEST 345 346 build: 347 requires: double !complex adolc 348 349 test: 350 suffix: 1 351 args: -ts_rhs_jacobian_test_mult_transpose FALSE -tao_max_it 2 -ts_rhs_jacobian_test_mult FALSE 352 output_file: output/ex16opt_ic_1.out 353 354 TEST*/ 355