Actual source code: adr_ex5adj_mf.cxx

      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";

      3: /*
      4:    REQUIRES configuration of PETSc with option --download-adolc.

      6:    For documentation on ADOL-C, see
      7:      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
      8: */
      9: /* ------------------------------------------------------------------------
     10:   See ../advection-diffusion-reaction/ex5 for a description of the problem
     11:   ------------------------------------------------------------------------- */

     13: #include <petscdmda.h>
     14: #include <petscts.h>
     15: #include "adolc-utils/init.cxx"
     16: #include "adolc-utils/matfree.cxx"
     17: #include <adolc/adolc.h>

     19: /* (Passive) field for the two variables */
     20: typedef struct {
     21:   PetscScalar u, v;
     22: } Field;

     24: /* Active field for the two variables */
     25: typedef struct {
     26:   adouble u, v;
     27: } AField;

     29: /* Application context */
     30: typedef struct {
     31:   PetscReal D1, D2, gamma, kappa;
     32:   AField  **u_a, **f_a;
     33:   AdolcCtx *adctx; /* Automatic differentation support */
     34: } AppCtx;

     36: extern PetscErrorCode InitialConditions(DM da, Vec U);
     37: extern PetscErrorCode InitializeLambda(DM da, Vec lambda, PetscReal x, PetscReal y);
     38: extern PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info, PetscReal t, Field **u, Field **udot, Field **f, void *ptr);
     39: extern PetscErrorCode IFunctionActive(TS ts, PetscReal ftime, Vec U, Vec Udot, Vec F, void *ptr);
     40: extern PetscErrorCode IJacobianMatFree(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A_shell, Mat B, void *ctx);

     42: int main(int argc, char **argv)
     43: {
     44:   TS          ts;   /* ODE integrator */
     45:   Vec         x, r; /* solution, residual */
     46:   DM          da;
     47:   AppCtx      appctx; /* Application context */
     48:   AdolcMatCtx matctx; /* Matrix (free) context */
     49:   Vec         lambda[1];
     50:   PetscBool   forwardonly = PETSC_FALSE;
     51:   Mat         A; /* (Matrix free) Jacobian matrix */
     52:   PetscInt    gxm, gym;

     54:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     55:      Initialize program
     56:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
     57:   PetscFunctionBeginUser;
     58:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
     59:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-forwardonly", &forwardonly, NULL));
     60:   appctx.D1    = 8.0e-5;
     61:   appctx.D2    = 4.0e-5;
     62:   appctx.gamma = .024;
     63:   appctx.kappa = .06;
     64:   PetscCall(PetscLogEventRegister("df/dx forward", MAT_CLASSID, &matctx.event1));
     65:   PetscCall(PetscLogEventRegister("df/d(xdot) forward", MAT_CLASSID, &matctx.event2));
     66:   PetscCall(PetscLogEventRegister("df/dx reverse", MAT_CLASSID, &matctx.event3));
     67:   PetscCall(PetscLogEventRegister("df/d(xdot) reverse", MAT_CLASSID, &matctx.event4));

     69:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     70:      Create distributed array (DMDA) to manage parallel grid and vectors
     71:   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
     72:   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DMDA_STENCIL_STAR, 65, 65, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, NULL, &da));
     73:   PetscCall(DMSetFromOptions(da));
     74:   PetscCall(DMSetUp(da));
     75:   PetscCall(DMDASetFieldName(da, 0, "u"));
     76:   PetscCall(DMDASetFieldName(da, 1, "v"));

     78:   /*  - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     79:      Extract global vectors from DMDA; then duplicate for remaining
     80:      vectors that are the same types
     81:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
     82:   PetscCall(DMCreateGlobalVector(da, &x));
     83:   PetscCall(VecDuplicate(x, &r));

     85:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     86:     Create matrix-free context and specify usage of PETSc-ADOL-C drivers
     87:     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
     88:   PetscCall(DMSetMatType(da, MATSHELL));
     89:   PetscCall(DMCreateMatrix(da, &A));
     90:   PetscCall(MatShellSetContext(A, &matctx));
     91:   PetscCall(MatShellSetOperation(A, MATOP_MULT, (void (*)(void))PetscAdolcIJacobianVectorProductIDMass));
     92:   PetscCall(MatShellSetOperation(A, MATOP_MULT_TRANSPOSE, (void (*)(void))PetscAdolcIJacobianTransposeVectorProductIDMass));
     93:   PetscCall(VecDuplicate(x, &matctx.X));
     94:   PetscCall(VecDuplicate(x, &matctx.Xdot));
     95:   PetscCall(DMGetLocalVector(da, &matctx.localX0));

     97:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     98:      Create timestepping solver context
     99:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    100:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
    101:   PetscCall(TSSetType(ts, TSCN));
    102:   PetscCall(TSSetDM(ts, da));
    103:   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
    104:   PetscCall(DMDATSSetIFunctionLocal(da, INSERT_VALUES, (DMDATSIFunctionLocalFn *)IFunctionLocalPassive, &appctx));

    106:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    107:     Some data required for matrix-free context
    108:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    109:   PetscCall(DMDAGetGhostCorners(da, NULL, NULL, NULL, &gxm, &gym, NULL));
    110:   matctx.m    = 2 * gxm * gym;
    111:   matctx.n    = 2 * gxm * gym; /* Number of dependent and independent variables */
    112:   matctx.flg  = PETSC_FALSE;   /* Flag for reverse mode */
    113:   matctx.tag1 = 1;             /* Tape identifier */

    115:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    116:      Trace function just once
    117:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    118:   PetscCall(PetscNew(&appctx.adctx));
    119:   PetscCall(IFunctionActive(ts, 1., x, matctx.Xdot, r, &appctx));
    120:   PetscCall(PetscFree(appctx.adctx));

    122:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    123:      Set Jacobian. In this case, IJacobian simply acts to pass context
    124:      information to the matrix-free Jacobian vector product.
    125:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    126:   PetscCall(TSSetIJacobian(ts, A, A, IJacobianMatFree, &appctx));

    128:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    129:      Set initial conditions
    130:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    131:   PetscCall(InitialConditions(da, x));
    132:   PetscCall(TSSetSolution(ts, x));

    134:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    135:     Have the TS save its trajectory so that TSAdjointSolve() may be used
    136:     and set solver options
    137:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    138:   if (!forwardonly) {
    139:     PetscCall(TSSetSaveTrajectory(ts));
    140:     PetscCall(TSSetMaxTime(ts, 200.0));
    141:     PetscCall(TSSetTimeStep(ts, 0.5));
    142:   } else {
    143:     PetscCall(TSSetMaxTime(ts, 2000.0));
    144:     PetscCall(TSSetTimeStep(ts, 10));
    145:   }
    146:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
    147:   PetscCall(TSSetFromOptions(ts));

    149:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    150:      Solve ODE system
    151:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    152:   PetscCall(TSSolve(ts, x));
    153:   if (!forwardonly) {
    154:     /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    155:        Start the Adjoint model
    156:        - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    157:     PetscCall(VecDuplicate(x, &lambda[0]));
    158:     /*   Reset initial conditions for the adjoint integration */
    159:     PetscCall(InitializeLambda(da, lambda[0], 0.5, 0.5));
    160:     PetscCall(TSSetCostGradients(ts, 1, lambda, NULL));
    161:     PetscCall(TSAdjointSolve(ts));
    162:     PetscCall(VecDestroy(&lambda[0]));
    163:   }

    165:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    166:      Free work space.  All PETSc objects should be destroyed when they
    167:      are no longer needed.
    168:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
    169:   PetscCall(DMRestoreLocalVector(da, &matctx.localX0));
    170:   PetscCall(VecDestroy(&r));
    171:   PetscCall(VecDestroy(&matctx.X));
    172:   PetscCall(VecDestroy(&matctx.Xdot));
    173:   PetscCall(MatDestroy(&A));
    174:   PetscCall(VecDestroy(&x));
    175:   PetscCall(TSDestroy(&ts));
    176:   PetscCall(DMDestroy(&da));

    178:   PetscCall(PetscFinalize());
    179:   return 0;
    180: }

    182: PetscErrorCode InitialConditions(DM da, Vec U)
    183: {
    184:   PetscInt  i, j, xs, ys, xm, ym, Mx, My;
    185:   Field   **u;
    186:   PetscReal hx, hy, x, y;

    188:   PetscFunctionBegin;
    189:   PetscCall(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));

    191:   hx = 2.5 / (PetscReal)Mx;
    192:   hy = 2.5 / (PetscReal)My;

    194:   /*
    195:      Get pointers to vector data
    196:   */
    197:   PetscCall(DMDAVecGetArray(da, U, &u));

    199:   /*
    200:      Get local grid boundaries
    201:   */
    202:   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));

    204:   /*
    205:      Compute function over the locally owned part of the grid
    206:   */
    207:   for (j = ys; j < ys + ym; j++) {
    208:     y = j * hy;
    209:     for (i = xs; i < xs + xm; i++) {
    210:       x = i * hx;
    211:       if (PetscApproximateGTE(x, 1.0) && PetscApproximateLTE(x, 1.5) && PetscApproximateGTE(y, 1.0) && PetscApproximateLTE(y, 1.5))
    212:         u[j][i].v = PetscPowReal(PetscSinReal(4.0 * PETSC_PI * x), 2.0) * PetscPowReal(PetscSinReal(4.0 * PETSC_PI * y), 2.0) / 4.0;
    213:       else u[j][i].v = 0.0;

    215:       u[j][i].u = 1.0 - 2.0 * u[j][i].v;
    216:     }
    217:   }

    219:   /*
    220:      Restore vectors
    221:   */
    222:   PetscCall(DMDAVecRestoreArray(da, U, &u));
    223:   PetscFunctionReturn(PETSC_SUCCESS);
    224: }

    226: PetscErrorCode InitializeLambda(DM da, Vec lambda, PetscReal x, PetscReal y)
    227: {
    228:   PetscInt i, j, Mx, My, xs, ys, xm, ym;
    229:   Field  **l;

    231:   PetscFunctionBegin;
    232:   PetscCall(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));
    233:   /* locate the global i index for x and j index for y */
    234:   i = (PetscInt)(x * (Mx - 1));
    235:   j = (PetscInt)(y * (My - 1));
    236:   PetscCall(DMDAGetCorners(da, &xs, &ys, NULL, &xm, &ym, NULL));

    238:   if (xs <= i && i < xs + xm && ys <= j && j < ys + ym) {
    239:     /* the i,j vertex is on this process */
    240:     PetscCall(DMDAVecGetArray(da, lambda, &l));
    241:     l[j][i].u = 1.0;
    242:     l[j][i].v = 1.0;
    243:     PetscCall(DMDAVecRestoreArray(da, lambda, &l));
    244:   }
    245:   PetscFunctionReturn(PETSC_SUCCESS);
    246: }

    248: PetscErrorCode IFunctionLocalPassive(DMDALocalInfo *info, PetscReal t, Field **u, Field **udot, Field **f, void *ptr)
    249: {
    250:   AppCtx     *appctx = (AppCtx *)ptr;
    251:   PetscInt    i, j, xs, ys, xm, ym;
    252:   PetscReal   hx, hy, sx, sy;
    253:   PetscScalar uc, uxx, uyy, vc, vxx, vyy;

    255:   PetscFunctionBegin;
    256:   hx = 2.50 / (PetscReal)info->mx;
    257:   sx = 1.0 / (hx * hx);
    258:   hy = 2.50 / (PetscReal)info->my;
    259:   sy = 1.0 / (hy * hy);

    261:   /* Get local grid boundaries */
    262:   xs = info->xs;
    263:   xm = info->xm;
    264:   ys = info->ys;
    265:   ym = info->ym;

    267:   /* Compute function over the locally owned part of the grid */
    268:   for (j = ys; j < ys + ym; j++) {
    269:     for (i = xs; i < xs + xm; i++) {
    270:       uc        = u[j][i].u;
    271:       uxx       = (-2.0 * uc + u[j][i - 1].u + u[j][i + 1].u) * sx;
    272:       uyy       = (-2.0 * uc + u[j - 1][i].u + u[j + 1][i].u) * sy;
    273:       vc        = u[j][i].v;
    274:       vxx       = (-2.0 * vc + u[j][i - 1].v + u[j][i + 1].v) * sx;
    275:       vyy       = (-2.0 * vc + u[j - 1][i].v + u[j + 1][i].v) * sy;
    276:       f[j][i].u = udot[j][i].u - appctx->D1 * (uxx + uyy) + uc * vc * vc - appctx->gamma * (1.0 - uc);
    277:       f[j][i].v = udot[j][i].v - appctx->D2 * (vxx + vyy) - uc * vc * vc + (appctx->gamma + appctx->kappa) * vc;
    278:     }
    279:   }
    280:   PetscCall(PetscLogFlops(16.0 * xm * ym));
    281:   PetscFunctionReturn(PETSC_SUCCESS);
    282: }

    284: PetscErrorCode IFunctionActive(TS ts, PetscReal ftime, Vec U, Vec Udot, Vec F, void *ptr)
    285: {
    286:   AppCtx       *appctx = (AppCtx *)ptr;
    287:   DM            da;
    288:   DMDALocalInfo info;
    289:   Field       **u, **f, **udot;
    290:   Vec           localU;
    291:   PetscInt      i, j, xs, ys, xm, ym, gxs, gys, gxm, gym;
    292:   PetscReal     hx, hy, sx, sy;
    293:   adouble       uc, uxx, uyy, vc, vxx, vyy;
    294:   AField      **f_a, *f_c, **u_a, *u_c;
    295:   PetscScalar   dummy;

    297:   PetscFunctionBegin;
    298:   PetscCall(TSGetDM(ts, &da));
    299:   PetscCall(DMDAGetLocalInfo(da, &info));
    300:   PetscCall(DMGetLocalVector(da, &localU));
    301:   hx  = 2.50 / (PetscReal)info.mx;
    302:   sx  = 1.0 / (hx * hx);
    303:   hy  = 2.50 / (PetscReal)info.my;
    304:   sy  = 1.0 / (hy * hy);
    305:   xs  = info.xs;
    306:   xm  = info.xm;
    307:   gxs = info.gxs;
    308:   gxm = info.gxm;
    309:   ys  = info.ys;
    310:   ym  = info.ym;
    311:   gys = info.gys;
    312:   gym = info.gym;

    314:   /*
    315:      Scatter ghost points to local vector,using the 2-step process
    316:         DMGlobalToLocalBegin(),DMGlobalToLocalEnd().
    317:      By placing code between these two statements, computations can be
    318:      done while messages are in transition.
    319:   */
    320:   PetscCall(DMGlobalToLocalBegin(da, U, INSERT_VALUES, localU));
    321:   PetscCall(DMGlobalToLocalEnd(da, U, INSERT_VALUES, localU));

    323:   /*
    324:      Get pointers to vector data
    325:   */
    326:   PetscCall(DMDAVecGetArrayRead(da, localU, &u));
    327:   PetscCall(DMDAVecGetArray(da, F, &f));
    328:   PetscCall(DMDAVecGetArrayRead(da, Udot, &udot));

    330:   /*
    331:     Create contiguous 1-arrays of AFields

    333:     NOTE: Memory for ADOL-C active variables (such as adouble and AField)
    334:           cannot be allocated using PetscMalloc, as this does not call the
    335:           relevant class constructor. Instead, we use the C++ keyword `new`.
    336:   */
    337:   u_c = new AField[info.gxm * info.gym];
    338:   f_c = new AField[info.gxm * info.gym];

    340:   /* Create corresponding 2-arrays of AFields */
    341:   u_a = new AField *[info.gym];
    342:   f_a = new AField *[info.gym];

    344:   /* Align indices between array types to endow 2d array with ghost points */
    345:   PetscCall(GiveGhostPoints(da, u_c, &u_a));
    346:   PetscCall(GiveGhostPoints(da, f_c, &f_a));

    348:   trace_on(1); /* Start of active section on tape 1 */

    350:   /*
    351:     Mark independence

    353:     NOTE: Ghost points are marked as independent, in place of the points they represent on
    354:           other processors / on other boundaries.
    355:   */
    356:   for (j = gys; j < gys + gym; j++) {
    357:     for (i = gxs; i < gxs + gxm; i++) {
    358:       u_a[j][i].u <<= u[j][i].u;
    359:       u_a[j][i].v <<= u[j][i].v;
    360:     }
    361:   }

    363:   /* Compute function over the locally owned part of the grid */
    364:   for (j = ys; j < ys + ym; j++) {
    365:     for (i = xs; i < xs + xm; i++) {
    366:       uc          = u_a[j][i].u;
    367:       uxx         = (-2.0 * uc + u_a[j][i - 1].u + u_a[j][i + 1].u) * sx;
    368:       uyy         = (-2.0 * uc + u_a[j - 1][i].u + u_a[j + 1][i].u) * sy;
    369:       vc          = u_a[j][i].v;
    370:       vxx         = (-2.0 * vc + u_a[j][i - 1].v + u_a[j][i + 1].v) * sx;
    371:       vyy         = (-2.0 * vc + u_a[j - 1][i].v + u_a[j + 1][i].v) * sy;
    372:       f_a[j][i].u = udot[j][i].u - appctx->D1 * (uxx + uyy) + uc * vc * vc - appctx->gamma * (1.0 - uc);
    373:       f_a[j][i].v = udot[j][i].v - appctx->D2 * (vxx + vyy) - uc * vc * vc + (appctx->gamma + appctx->kappa) * vc;
    374:     }
    375:   }

    377:   /*
    378:     Mark dependence

    380:     NOTE: Marking dependence of dummy variables makes the index notation much simpler when forming
    381:           the Jacobian later.
    382:   */
    383:   for (j = gys; j < gys + gym; j++) {
    384:     for (i = gxs; i < gxs + gxm; i++) {
    385:       if ((i < xs) || (i >= xs + xm) || (j < ys) || (j >= ys + ym)) {
    386:         f_a[j][i].u >>= dummy;
    387:         f_a[j][i].v >>= dummy;
    388:       } else {
    389:         f_a[j][i].u >>= f[j][i].u;
    390:         f_a[j][i].v >>= f[j][i].v;
    391:       }
    392:     }
    393:   }
    394:   trace_off(); /* End of active section */
    395:   PetscCall(PetscLogFlops(16.0 * xm * ym));

    397:   /* Restore vectors */
    398:   PetscCall(DMDAVecRestoreArray(da, F, &f));
    399:   PetscCall(DMDAVecRestoreArrayRead(da, localU, &u));
    400:   PetscCall(DMDAVecRestoreArrayRead(da, Udot, &udot));

    402:   PetscCall(DMRestoreLocalVector(da, &localU));

    404:   /* Destroy AFields appropriately */
    405:   f_a += info.gys;
    406:   u_a += info.gys;
    407:   delete[] f_a;
    408:   delete[] u_a;
    409:   delete[] f_c;
    410:   delete[] u_c;
    411:   PetscFunctionReturn(PETSC_SUCCESS);
    412: }

    414: /*
    415:   Simply acts to pass TS information to the AdolcMatCtx
    416: */
    417: PetscErrorCode IJacobianMatFree(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A_shell, Mat B, void *ctx)
    418: {
    419:   AdolcMatCtx *mctx;
    420:   DM           da;

    422:   PetscFunctionBeginUser;
    423:   PetscCall(MatShellGetContext(A_shell, &mctx));

    425:   mctx->time  = t;
    426:   mctx->shift = a;
    427:   if (mctx->ts != ts) mctx->ts = ts;
    428:   PetscCall(VecCopy(X, mctx->X));
    429:   PetscCall(VecCopy(Xdot, mctx->Xdot));
    430:   PetscCall(TSGetDM(ts, &da));
    431:   PetscCall(DMGlobalToLocalBegin(da, mctx->X, INSERT_VALUES, mctx->localX0));
    432:   PetscCall(DMGlobalToLocalEnd(da, mctx->X, INSERT_VALUES, mctx->localX0));
    433:   PetscFunctionReturn(PETSC_SUCCESS);
    434: }

    436: /*TEST

    438:   build:
    439:     requires: double !complex adolc

    441:   test:
    442:     suffix: 1
    443:     args: -ts_max_steps 1 -da_grid_x 12 -da_grid_y 12 -snes_test_jacobian
    444:     output_file: output/adr_ex5adj_mf_1.out

    446:   test:
    447:     suffix: 2
    448:     nsize: 4
    449:     args: -ts_max_steps 10 -da_grid_x 12 -da_grid_y 12 -ts_monitor -ts_adjoint_monitor
    450:     output_file: output/adr_ex5adj_mf_2.out

    452: TEST*/