Actual source code: matfree.cxx

      1: #include <petscdmda.h>
      2: #include <petscts.h>
      3: #include <adolc/interfaces.h>
      4: #include "contexts.cxx"

      6: /*
      7:    REQUIRES configuration of PETSc with option --download-adolc.

      9:    For documentation on ADOL-C, see
     10:      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
     11: */

     13: /*
     14:   ADOL-C implementation for Jacobian vector product, using the forward mode of AD.
     15:   Intended to overload MatMult in matrix-free methods where implicit timestepping
     16:   has been used.

     18:   For an implicit Jacobian we may use the rule that
     19:      G = M*xdot + f(x)    ==>     dG/dx = a*M + df/dx,
     20:   where a = d(xdot)/dx is a constant. Evaluated at x0 and acting upon a vector x1:
     21:      (dG/dx)(x0) * x1 = (a*df/d(xdot)(x0) + df/dx)(x0) * x1.

     23:   Input parameters:
     24:   A_shell - Jacobian matrix of MatShell type
     25:   X       - vector to be multiplied by A_shell

     27:   Output parameters:
     28:   Y       - product of A_shell and X
     29: */
     30: PetscErrorCode PetscAdolcIJacobianVectorProduct(Mat A_shell, Vec X, Vec Y)
     31: {
     32:   AdolcMatCtx       *mctx;
     33:   PetscInt           m, n, i, j, k = 0, d;
     34:   const PetscScalar *x0;
     35:   PetscScalar       *action, *x1;
     36:   Vec                localX1;
     37:   DM                 da;
     38:   DMDALocalInfo      info;

     40:   PetscFunctionBegin;
     41:   /* Get matrix-free context info */
     42:   PetscCall(MatShellGetContext(A_shell, &mctx));
     43:   m = mctx->m;
     44:   n = mctx->n;

     46:   /* Get local input vectors and extract data, x0 and x1*/
     47:   PetscCall(TSGetDM(mctx->ts, &da));
     48:   PetscCall(DMDAGetLocalInfo(da, &info));
     49:   PetscCall(DMGetLocalVector(da, &localX1));
     50:   PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX1));
     51:   PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX1));

     53:   PetscCall(VecGetArrayRead(mctx->localX0, &x0));
     54:   PetscCall(VecGetArray(localX1, &x1));

     56:   /* dF/dx part */
     57:   PetscCall(PetscMalloc1(m, &action));
     58:   PetscCall(PetscLogEventBegin(mctx->event1, 0, 0, 0, 0));
     59:   fos_forward(mctx->tag1, m, n, 0, x0, x1, NULL, action);
     60:   for (j = info.gys; j < info.gys + info.gym; j++) {
     61:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
     62:       for (d = 0; d < 2; d++) {
     63:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], INSERT_VALUES));
     64:         k++;
     65:       }
     66:     }
     67:   }
     68:   PetscCall(PetscLogEventEnd(mctx->event1, 0, 0, 0, 0));
     69:   k = 0;
     70:   PetscCall(VecAssemblyBegin(Y)); /* Note: Need to assemble between separate calls */
     71:   PetscCall(VecAssemblyEnd(Y));   /*       to INSERT_VALUES and ADD_VALUES         */

     73:   /* a * dF/d(xdot) part */
     74:   PetscCall(PetscLogEventBegin(mctx->event2, 0, 0, 0, 0));
     75:   fos_forward(mctx->tag2, m, n, 0, x0, x1, NULL, action);
     76:   for (j = info.gys; j < info.gys + info.gym; j++) {
     77:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
     78:       for (d = 0; d < 2; d++) {
     79:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) {
     80:           action[k] *= mctx->shift;
     81:           PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], ADD_VALUES));
     82:         }
     83:         k++;
     84:       }
     85:     }
     86:   }
     87:   PetscCall(PetscLogEventEnd(mctx->event2, 0, 0, 0, 0));
     88:   PetscCall(VecAssemblyBegin(Y));
     89:   PetscCall(VecAssemblyEnd(Y));
     90:   PetscCall(PetscFree(action));

     92:   /* Restore local vector */
     93:   PetscCall(VecRestoreArray(localX1, &x1));
     94:   PetscCall(VecRestoreArrayRead(mctx->localX0, &x0));
     95:   PetscCall(DMRestoreLocalVector(da, &localX1));
     96:   PetscFunctionReturn(PETSC_SUCCESS);
     97: }

     99: /*
    100:   ADOL-C implementation for Jacobian vector product, using the forward mode of AD.
    101:   Intended to overload MatMult in matrix-free methods where implicit timestepping
    102:   has been applied to a problem of the form
    103:                              du/dt = F(u).

    105:   Input parameters:
    106:   A_shell - Jacobian matrix of MatShell type
    107:   X       - vector to be multiplied by A_shell

    109:   Output parameters:
    110:   Y       - product of A_shell and X
    111: */
    112: PetscErrorCode PetscAdolcIJacobianVectorProductIDMass(Mat A_shell, Vec X, Vec Y)
    113: {
    114:   AdolcMatCtx       *mctx;
    115:   PetscInt           m, n, i, j, k = 0, d;
    116:   const PetscScalar *x0;
    117:   PetscScalar       *action, *x1;
    118:   Vec                localX1;
    119:   DM                 da;
    120:   DMDALocalInfo      info;

    122:   PetscFunctionBegin;
    123:   /* Get matrix-free context info */
    124:   PetscCall(MatShellGetContext(A_shell, &mctx));
    125:   m = mctx->m;
    126:   n = mctx->n;

    128:   /* Get local input vectors and extract data, x0 and x1*/
    129:   PetscCall(TSGetDM(mctx->ts, &da));
    130:   PetscCall(DMDAGetLocalInfo(da, &info));
    131:   PetscCall(DMGetLocalVector(da, &localX1));
    132:   PetscCall(DMGlobalToLocalBegin(da, X, INSERT_VALUES, localX1));
    133:   PetscCall(DMGlobalToLocalEnd(da, X, INSERT_VALUES, localX1));

    135:   PetscCall(VecGetArrayRead(mctx->localX0, &x0));
    136:   PetscCall(VecGetArray(localX1, &x1));

    138:   /* dF/dx part */
    139:   PetscCall(PetscMalloc1(m, &action));
    140:   PetscCall(PetscLogEventBegin(mctx->event1, 0, 0, 0, 0));
    141:   fos_forward(mctx->tag1, m, n, 0, x0, x1, NULL, action);
    142:   for (j = info.gys; j < info.gys + info.gym; j++) {
    143:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
    144:       for (d = 0; d < 2; d++) {
    145:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(Y, 1, &k, &action[k], INSERT_VALUES));
    146:         k++;
    147:       }
    148:     }
    149:   }
    150:   PetscCall(PetscLogEventEnd(mctx->event1, 0, 0, 0, 0));
    151:   k = 0;
    152:   PetscCall(VecAssemblyBegin(Y)); /* Note: Need to assemble between separate calls */
    153:   PetscCall(VecAssemblyEnd(Y));   /*       to INSERT_VALUES and ADD_VALUES         */
    154:   PetscCall(PetscFree(action));

    156:   /* Restore local vector */
    157:   PetscCall(VecRestoreArray(localX1, &x1));
    158:   PetscCall(VecRestoreArrayRead(mctx->localX0, &x0));
    159:   PetscCall(DMRestoreLocalVector(da, &localX1));

    161:   /* a * dF/d(xdot) part */
    162:   PetscCall(PetscLogEventBegin(mctx->event2, 0, 0, 0, 0));
    163:   PetscCall(VecAXPY(Y, mctx->shift, X));
    164:   PetscCall(PetscLogEventEnd(mctx->event2, 0, 0, 0, 0));
    165:   PetscFunctionReturn(PETSC_SUCCESS);
    166: }

    168: /*
    169:   ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD.
    170:   Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping
    171:   has been used.

    173:   Input parameters:
    174:   A_shell - Jacobian matrix of MatShell type
    175:   Y       - vector to be multiplied by A_shell transpose

    177:   Output parameters:
    178:   X       - product of A_shell transpose and X
    179: */
    180: PetscErrorCode PetscAdolcIJacobianTransposeVectorProduct(Mat A_shell, Vec Y, Vec X)
    181: {
    182:   AdolcMatCtx       *mctx;
    183:   PetscInt           m, n, i, j, k = 0, d;
    184:   const PetscScalar *x;
    185:   PetscScalar       *action, *y;
    186:   Vec                localY;
    187:   DM                 da;
    188:   DMDALocalInfo      info;

    190:   PetscFunctionBegin;
    191:   /* Get matrix-free context info */
    192:   PetscCall(MatShellGetContext(A_shell, &mctx));
    193:   m = mctx->m;
    194:   n = mctx->n;

    196:   /* Get local input vectors and extract data, x0 and x1*/
    197:   PetscCall(TSGetDM(mctx->ts, &da));
    198:   PetscCall(DMDAGetLocalInfo(da, &info));
    199:   PetscCall(DMGetLocalVector(da, &localY));
    200:   PetscCall(DMGlobalToLocalBegin(da, Y, INSERT_VALUES, localY));
    201:   PetscCall(DMGlobalToLocalEnd(da, Y, INSERT_VALUES, localY));
    202:   PetscCall(VecGetArrayRead(mctx->localX0, &x));
    203:   PetscCall(VecGetArray(localY, &y));

    205:   /* dF/dx part */
    206:   PetscCall(PetscMalloc1(n, &action));
    207:   PetscCall(PetscLogEventBegin(mctx->event3, 0, 0, 0, 0));
    208:   if (!mctx->flg) zos_forward(mctx->tag1, m, n, 1, x, NULL);
    209:   fos_reverse(mctx->tag1, m, n, y, action);
    210:   for (j = info.gys; j < info.gys + info.gym; j++) {
    211:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
    212:       for (d = 0; d < 2; d++) {
    213:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], INSERT_VALUES));
    214:         k++;
    215:       }
    216:     }
    217:   }
    218:   PetscCall(PetscLogEventEnd(mctx->event3, 0, 0, 0, 0));
    219:   k = 0;
    220:   PetscCall(VecAssemblyBegin(X)); /* Note: Need to assemble between separate calls */
    221:   PetscCall(VecAssemblyEnd(X));   /*       to INSERT_VALUES and ADD_VALUES         */

    223:   /* a * dF/d(xdot) part */
    224:   PetscCall(PetscLogEventBegin(mctx->event4, 0, 0, 0, 0));
    225:   if (!mctx->flg) {
    226:     zos_forward(mctx->tag2, m, n, 1, x, NULL);
    227:     mctx->flg = PETSC_TRUE;
    228:   }
    229:   fos_reverse(mctx->tag2, m, n, y, action);
    230:   for (j = info.gys; j < info.gys + info.gym; j++) {
    231:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
    232:       for (d = 0; d < 2; d++) {
    233:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) {
    234:           action[k] *= mctx->shift;
    235:           PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], ADD_VALUES));
    236:         }
    237:         k++;
    238:       }
    239:     }
    240:   }
    241:   PetscCall(PetscLogEventEnd(mctx->event4, 0, 0, 0, 0));
    242:   PetscCall(VecAssemblyBegin(X));
    243:   PetscCall(VecAssemblyEnd(X));
    244:   PetscCall(PetscFree(action));

    246:   /* Restore local vector */
    247:   PetscCall(VecRestoreArray(localY, &y));
    248:   PetscCall(VecRestoreArrayRead(mctx->localX0, &x));
    249:   PetscCall(DMRestoreLocalVector(da, &localY));
    250:   PetscFunctionReturn(PETSC_SUCCESS);
    251: }

    253: /*
    254:   ADOL-C implementation for Jacobian transpose vector product, using the reverse mode of AD.
    255:   Intended to overload MatMultTranspose in matrix-free methods where implicit timestepping
    256:   has been applied to a problem of the form
    257:                           du/dt = F(u).

    259:   Input parameters:
    260:   A_shell - Jacobian matrix of MatShell type
    261:   Y       - vector to be multiplied by A_shell transpose

    263:   Output parameters:
    264:   X       - product of A_shell transpose and X
    265: */
    266: PetscErrorCode PetscAdolcIJacobianTransposeVectorProductIDMass(Mat A_shell, Vec Y, Vec X)
    267: {
    268:   AdolcMatCtx       *mctx;
    269:   PetscInt           m, n, i, j, k = 0, d;
    270:   const PetscScalar *x;
    271:   PetscScalar       *action, *y;
    272:   Vec                localY;
    273:   DM                 da;
    274:   DMDALocalInfo      info;

    276:   PetscFunctionBegin;
    277:   /* Get matrix-free context info */
    278:   PetscCall(MatShellGetContext(A_shell, &mctx));
    279:   m = mctx->m;
    280:   n = mctx->n;

    282:   /* Get local input vectors and extract data, x0 and x1*/
    283:   PetscCall(TSGetDM(mctx->ts, &da));
    284:   PetscCall(DMDAGetLocalInfo(da, &info));
    285:   PetscCall(DMGetLocalVector(da, &localY));
    286:   PetscCall(DMGlobalToLocalBegin(da, Y, INSERT_VALUES, localY));
    287:   PetscCall(DMGlobalToLocalEnd(da, Y, INSERT_VALUES, localY));
    288:   PetscCall(VecGetArrayRead(mctx->localX0, &x));
    289:   PetscCall(VecGetArray(localY, &y));

    291:   /* dF/dx part */
    292:   PetscCall(PetscMalloc1(n, &action));
    293:   PetscCall(PetscLogEventBegin(mctx->event3, 0, 0, 0, 0));
    294:   if (!mctx->flg) zos_forward(mctx->tag1, m, n, 1, x, NULL);
    295:   fos_reverse(mctx->tag1, m, n, y, action);
    296:   for (j = info.gys; j < info.gys + info.gym; j++) {
    297:     for (i = info.gxs; i < info.gxs + info.gxm; i++) {
    298:       for (d = 0; d < 2; d++) {
    299:         if ((i >= info.xs) && (i < info.xs + info.xm) && (j >= info.ys) && (j < info.ys + info.ym)) PetscCall(VecSetValuesLocal(X, 1, &k, &action[k], INSERT_VALUES));
    300:         k++;
    301:       }
    302:     }
    303:   }
    304:   PetscCall(PetscLogEventEnd(mctx->event3, 0, 0, 0, 0));
    305:   k = 0;
    306:   PetscCall(VecAssemblyBegin(X)); /* Note: Need to assemble between separate calls */
    307:   PetscCall(VecAssemblyEnd(X));   /*       to INSERT_VALUES and ADD_VALUES         */
    308:   PetscCall(PetscFree(action));

    310:   /* Restore local vector */
    311:   PetscCall(VecRestoreArray(localY, &y));
    312:   PetscCall(VecRestoreArrayRead(mctx->localX0, &x));
    313:   PetscCall(DMRestoreLocalVector(da, &localY));

    315:   /* a * dF/d(xdot) part */
    316:   PetscCall(PetscLogEventBegin(mctx->event4, 0, 0, 0, 0));
    317:   PetscCall(VecAXPY(X, mctx->shift, Y));
    318:   PetscCall(PetscLogEventEnd(mctx->event4, 0, 0, 0, 0));
    319:   PetscFunctionReturn(PETSC_SUCCESS);
    320: }