1 static char help[] = "Demonstrates automatic, matrix-free Jacobian generation using ADOL-C for a time-dependent PDE in 2d, solved using implicit timestepping.\n"; 2 3 /* 4 Concepts: TS^time-dependent nonlinear problems 5 Concepts: Automatic differentiation using ADOL-C 6 Concepts: Matrix-free methods 7 */ 8 /* 9 REQUIRES configuration of PETSc with option --download-adolc. 10 11 For documentation on ADOL-C, see 12 $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf 13 */ 14 /* ------------------------------------------------------------------------ 15 See ../advection-diffusion-reaction/ex5 for a description of the problem 16 ------------------------------------------------------------------------- */ 17 18 #include <petscdmda.h> 19 #include <petscts.h> 20 #include "adolc-utils/init.cxx" 21 #include "adolc-utils/matfree.cxx" 22 #include <adolc/adolc.h> 23 24 /* (Passive) field for the two variables */ 25 typedef struct { 26 PetscScalar u,v; 27 } Field; 28 29 /* Active field for the two variables */ 30 typedef struct { 31 adouble u,v; 32 } AField; 33 34 /* Application context */ 35 typedef struct { 36 PetscReal D1,D2,gamma,kappa; 37 AField **u_a,**f_a; 38 AdolcCtx *adctx; /* Automatic differentation support */ 39 } AppCtx; 40 41 extern PetscErrorCode InitialConditions(DM da,Vec U); 42 extern PetscErrorCode InitializeLambda(DM da,Vec lambda,PetscReal x,PetscReal y); 43 extern PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info,PetscReal t,Field**u,Field**udot,Field**f,void *ptr); 44 extern PetscErrorCode IFunctionActive(TS ts,PetscReal ftime,Vec U,Vec Udot,Vec F,void *ptr); 45 extern PetscErrorCode IJacobianMatFree(TS ts,PetscReal t,Vec X,Vec Xdot,PetscReal a,Mat A_shell,Mat B,void *ctx); 46 47 int main(int argc,char **argv) 48 { 49 TS ts; /* ODE integrator */ 50 Vec x,r; /* solution, residual */ 51 PetscErrorCode ierr; 52 DM da; 53 AppCtx appctx; /* Application context */ 54 AdolcMatCtx matctx; /* Matrix (free) context */ 55 Vec lambda[1]; 56 PetscBool forwardonly=PETSC_FALSE; 57 Mat A; /* (Matrix free) Jacobian matrix */ 58 PetscInt gxm,gym; 59 60 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 61 Initialize program 62 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 63 ierr = PetscInitialize(&argc,&argv,NULL,help);if (ierr) return ierr; 64 CHKERRQ(PetscOptionsGetBool(NULL,NULL,"-forwardonly",&forwardonly,NULL)); 65 PetscFunctionBeginUser; 66 appctx.D1 = 8.0e-5; 67 appctx.D2 = 4.0e-5; 68 appctx.gamma = .024; 69 appctx.kappa = .06; 70 CHKERRQ(PetscLogEventRegister("df/dx forward",MAT_CLASSID,&matctx.event1)); 71 CHKERRQ(PetscLogEventRegister("df/d(xdot) forward",MAT_CLASSID,&matctx.event2)); 72 CHKERRQ(PetscLogEventRegister("df/dx reverse",MAT_CLASSID,&matctx.event3)); 73 CHKERRQ(PetscLogEventRegister("df/d(xdot) reverse",MAT_CLASSID,&matctx.event4)); 74 75 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 76 Create distributed array (DMDA) to manage parallel grid and vectors 77 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 78 CHKERRQ(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_PERIODIC,DM_BOUNDARY_PERIODIC,DMDA_STENCIL_STAR,65,65,PETSC_DECIDE,PETSC_DECIDE,2,1,NULL,NULL,&da)); 79 CHKERRQ(DMSetFromOptions(da)); 80 CHKERRQ(DMSetUp(da)); 81 CHKERRQ(DMDASetFieldName(da,0,"u")); 82 CHKERRQ(DMDASetFieldName(da,1,"v")); 83 84 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 85 Extract global vectors from DMDA; then duplicate for remaining 86 vectors that are the same types 87 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 88 CHKERRQ(DMCreateGlobalVector(da,&x)); 89 CHKERRQ(VecDuplicate(x,&r)); 90 91 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 92 Create matrix free context and specify usage of PETSc-ADOL-C drivers 93 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 94 CHKERRQ(DMSetMatType(da,MATSHELL)); 95 CHKERRQ(DMCreateMatrix(da,&A)); 96 CHKERRQ(MatShellSetContext(A,&matctx)); 97 CHKERRQ(MatShellSetOperation(A,MATOP_MULT,(void (*)(void))PetscAdolcIJacobianVectorProductIDMass)); 98 CHKERRQ(MatShellSetOperation(A,MATOP_MULT_TRANSPOSE,(void (*)(void))PetscAdolcIJacobianTransposeVectorProductIDMass)); 99 CHKERRQ(VecDuplicate(x,&matctx.X)); 100 CHKERRQ(VecDuplicate(x,&matctx.Xdot)); 101 CHKERRQ(DMGetLocalVector(da,&matctx.localX0)); 102 103 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 104 Create timestepping solver context 105 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 106 CHKERRQ(TSCreate(PETSC_COMM_WORLD,&ts)); 107 CHKERRQ(TSSetType(ts,TSCN)); 108 CHKERRQ(TSSetDM(ts,da)); 109 CHKERRQ(TSSetProblemType(ts,TS_NONLINEAR)); 110 CHKERRQ(DMDATSSetIFunctionLocal(da,INSERT_VALUES,(DMDATSIFunctionLocal)IFunctionLocalPassive,&appctx)); 111 112 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 113 Some data required for matrix-free context 114 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 115 CHKERRQ(DMDAGetGhostCorners(da,NULL,NULL,NULL,&gxm,&gym,NULL)); 116 matctx.m = 2*gxm*gym;matctx.n = 2*gxm*gym; /* Number of dependent and independent variables */ 117 matctx.flg = PETSC_FALSE; /* Flag for reverse mode */ 118 matctx.tag1 = 1; /* Tape identifier */ 119 120 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 121 Trace function just once 122 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 123 CHKERRQ(PetscNew(&appctx.adctx)); 124 CHKERRQ(IFunctionActive(ts,1.,x,matctx.Xdot,r,&appctx)); 125 CHKERRQ(PetscFree(appctx.adctx)); 126 127 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 128 Set Jacobian. In this case, IJacobian simply acts to pass context 129 information to the matrix-free Jacobian vector product. 130 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 131 CHKERRQ(TSSetIJacobian(ts,A,A,IJacobianMatFree,&appctx)); 132 133 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 134 Set initial conditions 135 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 136 CHKERRQ(InitialConditions(da,x)); 137 CHKERRQ(TSSetSolution(ts,x)); 138 139 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 140 Have the TS save its trajectory so that TSAdjointSolve() may be used 141 and set solver options 142 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 143 if (!forwardonly) { 144 CHKERRQ(TSSetSaveTrajectory(ts)); 145 CHKERRQ(TSSetMaxTime(ts,200.0)); 146 CHKERRQ(TSSetTimeStep(ts,0.5)); 147 } else { 148 CHKERRQ(TSSetMaxTime(ts,2000.0)); 149 CHKERRQ(TSSetTimeStep(ts,10)); 150 } 151 CHKERRQ(TSSetExactFinalTime(ts,TS_EXACTFINALTIME_STEPOVER)); 152 CHKERRQ(TSSetFromOptions(ts)); 153 154 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 155 Solve ODE system 156 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 157 CHKERRQ(TSSolve(ts,x)); 158 if (!forwardonly) { 159 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 160 Start the Adjoint model 161 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 162 CHKERRQ(VecDuplicate(x,&lambda[0])); 163 /* Reset initial conditions for the adjoint integration */ 164 CHKERRQ(InitializeLambda(da,lambda[0],0.5,0.5)); 165 CHKERRQ(TSSetCostGradients(ts,1,lambda,NULL)); 166 CHKERRQ(TSAdjointSolve(ts)); 167 CHKERRQ(VecDestroy(&lambda[0])); 168 } 169 170 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 171 Free work space. All PETSc objects should be destroyed when they 172 are no longer needed. 173 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 174 CHKERRQ(DMRestoreLocalVector(da,&matctx.localX0)); 175 CHKERRQ(VecDestroy(&r)); 176 CHKERRQ(VecDestroy(&matctx.X)); 177 CHKERRQ(VecDestroy(&matctx.Xdot)); 178 CHKERRQ(MatDestroy(&A)); 179 CHKERRQ(VecDestroy(&x)); 180 CHKERRQ(TSDestroy(&ts)); 181 CHKERRQ(DMDestroy(&da)); 182 183 ierr = PetscFinalize(); 184 return ierr; 185 } 186 187 PetscErrorCode InitialConditions(DM da,Vec U) 188 { 189 PetscInt i,j,xs,ys,xm,ym,Mx,My; 190 Field **u; 191 PetscReal hx,hy,x,y; 192 193 PetscFunctionBegin; 194 CHKERRQ(DMDAGetInfo(da,PETSC_IGNORE,&Mx,&My,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE)); 195 196 hx = 2.5/(PetscReal)Mx; 197 hy = 2.5/(PetscReal)My; 198 199 /* 200 Get pointers to vector data 201 */ 202 CHKERRQ(DMDAVecGetArray(da,U,&u)); 203 204 /* 205 Get local grid boundaries 206 */ 207 CHKERRQ(DMDAGetCorners(da,&xs,&ys,NULL,&xm,&ym,NULL)); 208 209 /* 210 Compute function over the locally owned part of the grid 211 */ 212 for (j=ys; j<ys+ym; j++) { 213 y = j*hy; 214 for (i=xs; i<xs+xm; i++) { 215 x = i*hx; 216 if (PetscApproximateGTE(x,1.0) && PetscApproximateLTE(x,1.5) && PetscApproximateGTE(y,1.0) && PetscApproximateLTE(y,1.5)) u[j][i].v = PetscPowReal(PetscSinReal(4.0*PETSC_PI*x),2.0)*PetscPowReal(PetscSinReal(4.0*PETSC_PI*y),2.0)/4.0; 217 else u[j][i].v = 0.0; 218 219 u[j][i].u = 1.0 - 2.0*u[j][i].v; 220 } 221 } 222 223 /* 224 Restore vectors 225 */ 226 CHKERRQ(DMDAVecRestoreArray(da,U,&u)); 227 PetscFunctionReturn(0); 228 } 229 230 PetscErrorCode InitializeLambda(DM da,Vec lambda,PetscReal x,PetscReal y) 231 { 232 PetscInt i,j,Mx,My,xs,ys,xm,ym; 233 Field **l; 234 235 PetscFunctionBegin; 236 CHKERRQ(DMDAGetInfo(da,PETSC_IGNORE,&Mx,&My,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE,PETSC_IGNORE)); 237 /* locate the global i index for x and j index for y */ 238 i = (PetscInt)(x*(Mx-1)); 239 j = (PetscInt)(y*(My-1)); 240 CHKERRQ(DMDAGetCorners(da,&xs,&ys,NULL,&xm,&ym,NULL)); 241 242 if (xs <= i && i < xs+xm && ys <= j && j < ys+ym) { 243 /* the i,j vertex is on this process */ 244 CHKERRQ(DMDAVecGetArray(da,lambda,&l)); 245 l[j][i].u = 1.0; 246 l[j][i].v = 1.0; 247 CHKERRQ(DMDAVecRestoreArray(da,lambda,&l)); 248 } 249 PetscFunctionReturn(0); 250 } 251 252 PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info,PetscReal t,Field**u,Field**udot,Field**f,void *ptr) 253 { 254 AppCtx *appctx = (AppCtx*)ptr; 255 PetscInt i,j,xs,ys,xm,ym; 256 PetscReal hx,hy,sx,sy; 257 PetscScalar uc,uxx,uyy,vc,vxx,vyy; 258 259 PetscFunctionBegin; 260 hx = 2.50/(PetscReal)(info->mx); sx = 1.0/(hx*hx); 261 hy = 2.50/(PetscReal)(info->my); sy = 1.0/(hy*hy); 262 263 /* Get local grid boundaries */ 264 xs = info->xs; xm = info->xm; ys = info->ys; ym = info->ym; 265 266 /* Compute function over the locally owned part of the grid */ 267 for (j=ys; j<ys+ym; j++) { 268 for (i=xs; i<xs+xm; i++) { 269 uc = u[j][i].u; 270 uxx = (-2.0*uc + u[j][i-1].u + u[j][i+1].u)*sx; 271 uyy = (-2.0*uc + u[j-1][i].u + u[j+1][i].u)*sy; 272 vc = u[j][i].v; 273 vxx = (-2.0*vc + u[j][i-1].v + u[j][i+1].v)*sx; 274 vyy = (-2.0*vc + u[j-1][i].v + u[j+1][i].v)*sy; 275 f[j][i].u = udot[j][i].u - appctx->D1*(uxx + uyy) + uc*vc*vc - appctx->gamma*(1.0 - uc); 276 f[j][i].v = udot[j][i].v - appctx->D2*(vxx + vyy) - uc*vc*vc + (appctx->gamma + appctx->kappa)*vc; 277 } 278 } 279 CHKERRQ(PetscLogFlops(16.0*xm*ym)); 280 PetscFunctionReturn(0); 281 } 282 283 PetscErrorCode IFunctionActive(TS ts,PetscReal ftime,Vec U,Vec Udot,Vec F,void *ptr) 284 { 285 AppCtx *appctx = (AppCtx*)ptr; 286 DM da; 287 DMDALocalInfo info; 288 Field **u,**f,**udot; 289 Vec localU; 290 PetscInt i,j,xs,ys,xm,ym,gxs,gys,gxm,gym; 291 PetscReal hx,hy,sx,sy; 292 adouble uc,uxx,uyy,vc,vxx,vyy; 293 AField **f_a,*f_c,**u_a,*u_c; 294 PetscScalar dummy; 295 296 PetscFunctionBegin; 297 CHKERRQ(TSGetDM(ts,&da)); 298 CHKERRQ(DMDAGetLocalInfo(da,&info)); 299 CHKERRQ(DMGetLocalVector(da,&localU)); 300 hx = 2.50/(PetscReal)(info.mx); sx = 1.0/(hx*hx); 301 hy = 2.50/(PetscReal)(info.my); sy = 1.0/(hy*hy); 302 xs = info.xs; xm = info.xm; gxs = info.gxs; gxm = info.gxm; 303 ys = info.ys; ym = info.ym; gys = info.gys; gym = info.gym; 304 305 /* 306 Scatter ghost points to local vector,using the 2-step process 307 DMGlobalToLocalBegin(),DMGlobalToLocalEnd(). 308 By placing code between these two statements, computations can be 309 done while messages are in transition. 310 */ 311 CHKERRQ(DMGlobalToLocalBegin(da,U,INSERT_VALUES,localU)); 312 CHKERRQ(DMGlobalToLocalEnd(da,U,INSERT_VALUES,localU)); 313 314 /* 315 Get pointers to vector data 316 */ 317 CHKERRQ(DMDAVecGetArrayRead(da,localU,&u)); 318 CHKERRQ(DMDAVecGetArray(da,F,&f)); 319 CHKERRQ(DMDAVecGetArrayRead(da,Udot,&udot)); 320 321 /* 322 Create contiguous 1-arrays of AFields 323 324 NOTE: Memory for ADOL-C active variables (such as adouble and AField) 325 cannot be allocated using PetscMalloc, as this does not call the 326 relevant class constructor. Instead, we use the C++ keyword `new`. 327 */ 328 u_c = new AField[info.gxm*info.gym]; 329 f_c = new AField[info.gxm*info.gym]; 330 331 /* Create corresponding 2-arrays of AFields */ 332 u_a = new AField*[info.gym]; 333 f_a = new AField*[info.gym]; 334 335 /* Align indices between array types to endow 2d array with ghost points */ 336 CHKERRQ(GiveGhostPoints(da,u_c,&u_a)); 337 CHKERRQ(GiveGhostPoints(da,f_c,&f_a)); 338 339 trace_on(1); /* Start of active section on tape 1 */ 340 341 /* 342 Mark independence 343 344 NOTE: Ghost points are marked as independent, in place of the points they represent on 345 other processors / on other boundaries. 346 */ 347 for (j=gys; j<gys+gym; j++) { 348 for (i=gxs; i<gxs+gxm; i++) { 349 u_a[j][i].u <<= u[j][i].u; 350 u_a[j][i].v <<= u[j][i].v; 351 } 352 } 353 354 /* Compute function over the locally owned part of the grid */ 355 for (j=ys; j<ys+ym; j++) { 356 for (i=xs; i<xs+xm; i++) { 357 uc = u_a[j][i].u; 358 uxx = (-2.0*uc + u_a[j][i-1].u + u_a[j][i+1].u)*sx; 359 uyy = (-2.0*uc + u_a[j-1][i].u + u_a[j+1][i].u)*sy; 360 vc = u_a[j][i].v; 361 vxx = (-2.0*vc + u_a[j][i-1].v + u_a[j][i+1].v)*sx; 362 vyy = (-2.0*vc + u_a[j-1][i].v + u_a[j+1][i].v)*sy; 363 f_a[j][i].u = udot[j][i].u - appctx->D1*(uxx + uyy) + uc*vc*vc - appctx->gamma*(1.0 - uc); 364 f_a[j][i].v = udot[j][i].v - appctx->D2*(vxx + vyy) - uc*vc*vc + (appctx->gamma + appctx->kappa)*vc; 365 } 366 } 367 368 /* 369 Mark dependence 370 371 NOTE: Marking dependence of dummy variables makes the index notation much simpler when forming 372 the Jacobian later. 373 */ 374 for (j=gys; j<gys+gym; j++) { 375 for (i=gxs; i<gxs+gxm; i++) { 376 if ((i < xs) || (i >= xs+xm) || (j < ys) || (j >= ys+ym)) { 377 f_a[j][i].u >>= dummy; 378 f_a[j][i].v >>= dummy; 379 } else { 380 f_a[j][i].u >>= f[j][i].u; 381 f_a[j][i].v >>= f[j][i].v; 382 } 383 } 384 } 385 trace_off(); /* End of active section */ 386 CHKERRQ(PetscLogFlops(16.0*xm*ym)); 387 388 /* Restore vectors */ 389 CHKERRQ(DMDAVecRestoreArray(da,F,&f)); 390 CHKERRQ(DMDAVecRestoreArrayRead(da,localU,&u)); 391 CHKERRQ(DMDAVecRestoreArrayRead(da,Udot,&udot)); 392 393 CHKERRQ(DMRestoreLocalVector(da,&localU)); 394 395 /* Destroy AFields appropriately */ 396 f_a += info.gys; 397 u_a += info.gys; 398 delete[] f_a; 399 delete[] u_a; 400 delete[] f_c; 401 delete[] u_c; 402 PetscFunctionReturn(0); 403 } 404 405 /* 406 Simply acts to pass TS information to the AdolcMatCtx 407 */ 408 PetscErrorCode IJacobianMatFree(TS ts,PetscReal t,Vec X,Vec Xdot,PetscReal a,Mat A_shell,Mat B,void *ctx) 409 { 410 AdolcMatCtx *mctx; 411 DM da; 412 413 PetscFunctionBeginUser; 414 CHKERRQ(MatShellGetContext(A_shell,&mctx)); 415 416 mctx->time = t; 417 mctx->shift = a; 418 if (mctx->ts != ts) mctx->ts = ts; 419 CHKERRQ(VecCopy(X,mctx->X)); 420 CHKERRQ(VecCopy(Xdot,mctx->Xdot)); 421 CHKERRQ(TSGetDM(ts,&da)); 422 CHKERRQ(DMGlobalToLocalBegin(da,mctx->X,INSERT_VALUES,mctx->localX0)); 423 CHKERRQ(DMGlobalToLocalEnd(da,mctx->X,INSERT_VALUES,mctx->localX0)); 424 PetscFunctionReturn(0); 425 } 426 427 /*TEST 428 429 build: 430 requires: double !complex adolc 431 432 test: 433 suffix: 1 434 args: -ts_max_steps 1 -da_grid_x 12 -da_grid_y 12 -snes_test_jacobian 435 output_file: output/adr_ex5adj_mf_1.out 436 437 test: 438 suffix: 2 439 nsize: 4 440 args: -ts_max_steps 10 -da_grid_x 12 -da_grid_y 12 -ts_monitor -ts_adjoint_monitor 441 output_file: output/adr_ex5adj_mf_2.out 442 443 TEST*/ 444