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