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