Actual source code: dmdats.c

      1: #include <petscdmda.h>
      2: #include <petsc/private/dmimpl.h>
      3: #include <petsc/private/tsimpl.h>
      4: #include <petscdraw.h>

      6: /* This structure holds the user-provided DMDA callbacks */
      7: typedef struct {
      8:   PetscErrorCode (*ifunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *, void *);
      9:   PetscErrorCode (*rhsfunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *);
     10:   PetscErrorCode (*ijacobianlocal)(DMDALocalInfo *, PetscReal, void *, void *, PetscReal, Mat, Mat, void *);
     11:   PetscErrorCode (*rhsjacobianlocal)(DMDALocalInfo *, PetscReal, void *, Mat, Mat, void *);
     12:   void      *ifunctionlocalctx;
     13:   void      *ijacobianlocalctx;
     14:   void      *rhsfunctionlocalctx;
     15:   void      *rhsjacobianlocalctx;
     16:   InsertMode ifunctionlocalimode;
     17:   InsertMode rhsfunctionlocalimode;
     18: } DMTS_DA;

     20: static PetscErrorCode DMTSDestroy_DMDA(DMTS sdm)
     21: {
     22:   PetscFunctionBegin;
     23:   PetscCall(PetscFree(sdm->data));
     24:   PetscFunctionReturn(PETSC_SUCCESS);
     25: }

     27: static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm, DMTS sdm)
     28: {
     29:   PetscFunctionBegin;
     30:   PetscCall(PetscNew((DMTS_DA **)&sdm->data));
     31:   if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMTS_DA)));
     32:   PetscFunctionReturn(PETSC_SUCCESS);
     33: }

     35: static PetscErrorCode DMDATSGetContext(DM dm, DMTS sdm, DMTS_DA **dmdats)
     36: {
     37:   PetscFunctionBegin;
     38:   *dmdats = NULL;
     39:   if (!sdm->data) {
     40:     PetscCall(PetscNew((DMTS_DA **)&sdm->data));
     41:     sdm->ops->destroy   = DMTSDestroy_DMDA;
     42:     sdm->ops->duplicate = DMTSDuplicate_DMDA;
     43:   }
     44:   *dmdats = (DMTS_DA *)sdm->data;
     45:   PetscFunctionReturn(PETSC_SUCCESS);
     46: }

     48: static PetscErrorCode TSComputeIFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, Vec F, PetscCtx ctx)
     49: {
     50:   DM            dm;
     51:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
     52:   DMDALocalInfo info;
     53:   Vec           Xloc, Xdotloc;
     54:   void         *x, *f, *xdot;

     56:   PetscFunctionBegin;
     60:   PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
     61:   PetscCall(TSGetDM(ts, &dm));
     62:   PetscCall(DMGetLocalVector(dm, &Xdotloc));
     63:   PetscCall(DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc));
     64:   PetscCall(DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc));
     65:   PetscCall(DMGetLocalVector(dm, &Xloc));
     66:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
     67:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
     68:   PetscCall(DMDAGetLocalInfo(dm, &info));
     69:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
     70:   PetscCall(DMDAVecGetArray(dm, Xdotloc, &xdot));
     71:   switch (dmdats->ifunctionlocalimode) {
     72:   case INSERT_VALUES: {
     73:     PetscCall(DMDAVecGetArray(dm, F, &f));
     74:     CHKMEMQ;
     75:     PetscCall((*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx));
     76:     CHKMEMQ;
     77:     PetscCall(DMDAVecRestoreArray(dm, F, &f));
     78:   } break;
     79:   case ADD_VALUES: {
     80:     Vec Floc;
     81:     PetscCall(DMGetLocalVector(dm, &Floc));
     82:     PetscCall(VecZeroEntries(Floc));
     83:     PetscCall(DMDAVecGetArray(dm, Floc, &f));
     84:     CHKMEMQ;
     85:     PetscCall((*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx));
     86:     CHKMEMQ;
     87:     PetscCall(DMDAVecRestoreArray(dm, Floc, &f));
     88:     PetscCall(VecZeroEntries(F));
     89:     PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
     90:     PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
     91:     PetscCall(DMRestoreLocalVector(dm, &Floc));
     92:   } break;
     93:   default:
     94:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->ifunctionlocalimode);
     95:   }
     96:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
     97:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
     98:   PetscCall(DMDAVecRestoreArray(dm, Xdotloc, &xdot));
     99:   PetscCall(DMRestoreLocalVector(dm, &Xdotloc));
    100:   PetscFunctionReturn(PETSC_SUCCESS);
    101: }

    103: static PetscErrorCode TSComputeIJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, PetscReal shift, Mat A, Mat B, PetscCtx ctx)
    104: {
    105:   DM            dm;
    106:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
    107:   DMDALocalInfo info;
    108:   Vec           Xloc, Xdotloc;
    109:   void         *x, *xdot;

    111:   PetscFunctionBegin;
    112:   PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
    113:   PetscCall(TSGetDM(ts, &dm));

    115:   PetscCheck(dmdats->ijacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
    116:   PetscCall(DMGetLocalVector(dm, &Xloc));
    117:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
    118:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
    119:   PetscCall(DMGetLocalVector(dm, &Xdotloc));
    120:   PetscCall(DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc));
    121:   PetscCall(DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc));
    122:   PetscCall(DMDAGetLocalInfo(dm, &info));
    123:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
    124:   PetscCall(DMDAVecGetArray(dm, Xdotloc, &xdot));
    125:   CHKMEMQ;
    126:   PetscCall((*dmdats->ijacobianlocal)(&info, ptime, x, xdot, shift, A, B, dmdats->ijacobianlocalctx));
    127:   CHKMEMQ;
    128:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
    129:   PetscCall(DMDAVecRestoreArray(dm, Xdotloc, &xdot));
    130:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
    131:   PetscCall(DMRestoreLocalVector(dm, &Xdotloc));
    132:   /* This will be redundant if the user called both, but it's too common to forget. */
    133:   if (A != B) {
    134:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    135:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    136:   }
    137:   PetscFunctionReturn(PETSC_SUCCESS);
    138: }

    140: static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, PetscCtx ctx)
    141: {
    142:   DM            dm;
    143:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
    144:   DMDALocalInfo info;
    145:   Vec           Xloc;
    146:   void         *x, *f;

    148:   PetscFunctionBegin;
    152:   PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
    153:   PetscCall(TSGetDM(ts, &dm));
    154:   PetscCall(DMGetLocalVector(dm, &Xloc));
    155:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
    156:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
    157:   PetscCall(DMDAGetLocalInfo(dm, &info));
    158:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
    159:   switch (dmdats->rhsfunctionlocalimode) {
    160:   case INSERT_VALUES: {
    161:     PetscCall(DMDAVecGetArray(dm, F, &f));
    162:     CHKMEMQ;
    163:     PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx));
    164:     CHKMEMQ;
    165:     PetscCall(DMDAVecRestoreArray(dm, F, &f));
    166:   } break;
    167:   case ADD_VALUES: {
    168:     Vec Floc;
    169:     PetscCall(DMGetLocalVector(dm, &Floc));
    170:     PetscCall(VecZeroEntries(Floc));
    171:     PetscCall(DMDAVecGetArray(dm, Floc, &f));
    172:     CHKMEMQ;
    173:     PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx));
    174:     CHKMEMQ;
    175:     PetscCall(DMDAVecRestoreArray(dm, Floc, &f));
    176:     PetscCall(VecZeroEntries(F));
    177:     PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
    178:     PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
    179:     PetscCall(DMRestoreLocalVector(dm, &Floc));
    180:   } break;
    181:   default:
    182:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->rhsfunctionlocalimode);
    183:   }
    184:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
    185:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
    186:   PetscFunctionReturn(PETSC_SUCCESS);
    187: }

    189: static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, PetscCtx ctx)
    190: {
    191:   DM            dm;
    192:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
    193:   DMDALocalInfo info;
    194:   Vec           Xloc;
    195:   void         *x;

    197:   PetscFunctionBegin;
    198:   PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
    199:   PetscCall(TSGetDM(ts, &dm));

    201:   PetscCheck(dmdats->rhsjacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()");
    202:   PetscCall(DMGetLocalVector(dm, &Xloc));
    203:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
    204:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
    205:   PetscCall(DMDAGetLocalInfo(dm, &info));
    206:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
    207:   CHKMEMQ;
    208:   PetscCall((*dmdats->rhsjacobianlocal)(&info, ptime, x, A, B, dmdats->rhsjacobianlocalctx));
    209:   CHKMEMQ;
    210:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
    211:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
    212:   /* This will be redundant if the user called both, but it's too common to forget. */
    213:   if (A != B) {
    214:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    215:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    216:   }
    217:   PetscFunctionReturn(PETSC_SUCCESS);
    218: }

    220: /*@C
    221:   DMDATSSetRHSFunctionLocal - set a local residual evaluation function for use with `DMDA`

    223:   Logically Collective

    225:   Input Parameters:
    226: + dm    - `DM` to associate callback with
    227: . imode - insert mode for the residual
    228: . func  - local residual evaluation, see `DMDATSRHSFunctionLocalFn` for the calling sequence
    229: - ctx   - optional context for local residual evaluation

    231:   Level: beginner

    233: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSFunctionLocalFn`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
    234: @*/
    235: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocalFn *func, PetscCtx ctx)
    236: {
    237:   DMTS     sdm;
    238:   DMTS_DA *dmdats;

    240:   PetscFunctionBegin;
    242:   PetscCall(DMGetDMTSWrite(dm, &sdm));
    243:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
    244:   dmdats->rhsfunctionlocalimode = imode;
    245:   dmdats->rhsfunctionlocal      = func;
    246:   dmdats->rhsfunctionlocalctx   = ctx;
    247:   PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats));
    248:   PetscFunctionReturn(PETSC_SUCCESS);
    249: }

    251: /*@C
    252:   DMDATSSetRHSJacobianLocal - set a local residual evaluation function for use with `DMDA`

    254:   Logically Collective

    256:   Input Parameters:
    257: + dm   - `DM` to associate callback with
    258: . func - local RHS Jacobian evaluation routine, see `DMDATSRHSJacobianLocalFn` for the calling sequence
    259: - ctx  - optional context for local jacobian evaluation

    261:   Level: beginner

    263: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSJacobianLocalFn`, `DMTSSetRHSJacobian()`,
    264: `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
    265: @*/
    266: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocalFn *func, PetscCtx ctx)
    267: {
    268:   DMTS     sdm;
    269:   DMTS_DA *dmdats;

    271:   PetscFunctionBegin;
    273:   PetscCall(DMGetDMTSWrite(dm, &sdm));
    274:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
    275:   dmdats->rhsjacobianlocal    = func;
    276:   dmdats->rhsjacobianlocalctx = ctx;
    277:   PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats));
    278:   PetscFunctionReturn(PETSC_SUCCESS);
    279: }

    281: /*@C
    282:   DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA`

    284:   Logically Collective

    286:   Input Parameters:
    287: + dm    - `DM` to associate callback with
    288: . imode - the insert mode of the function
    289: . func  - local residual evaluation, see `DMDATSIFunctionLocalFn` for the calling sequence
    290: - ctx   - optional context for local residual evaluation

    292:   Level: beginner

    294: .seealso: [](ch_ts), `DMDA`, `DMDATSIFunctionLocalFn`, `DMTSSetIFunction()`,
    295: `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
    296: @*/
    297: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocalFn *func, PetscCtx ctx)
    298: {
    299:   DMTS     sdm;
    300:   DMTS_DA *dmdats;

    302:   PetscFunctionBegin;
    304:   PetscCall(DMGetDMTSWrite(dm, &sdm));
    305:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
    306:   dmdats->ifunctionlocalimode = imode;
    307:   dmdats->ifunctionlocal      = func;
    308:   dmdats->ifunctionlocalctx   = ctx;
    309:   PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats));
    310:   PetscFunctionReturn(PETSC_SUCCESS);
    311: }

    313: /*@C
    314:   DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA`

    316:   Logically Collective

    318:   Input Parameters:
    319: + dm   - `DM` to associate callback with
    320: . func - local residual evaluation, see `DMDATSIJacobianLocalFn` for the calling sequence
    321: - ctx  - optional context for local residual evaluation

    323:   Level: beginner

    325: .seealso: [](ch_ts), `DMDA`, `DMDATSIJacobianLocalFn`, `DMTSSetIJacobian()`,
    326: `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
    327: @*/
    328: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocalFn *func, PetscCtx ctx)
    329: {
    330:   DMTS     sdm;
    331:   DMTS_DA *dmdats;

    333:   PetscFunctionBegin;
    335:   PetscCall(DMGetDMTSWrite(dm, &sdm));
    336:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
    337:   dmdats->ijacobianlocal    = func;
    338:   dmdats->ijacobianlocalctx = ctx;
    339:   PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats));
    340:   PetscFunctionReturn(PETSC_SUCCESS);
    341: }

    343: PetscErrorCode TSMonitorDMDARayDestroy(PetscCtxRt mctx)
    344: {
    345:   TSMonitorDMDARayCtx *rayctx = *(TSMonitorDMDARayCtx **)mctx;

    347:   PetscFunctionBegin;
    348:   PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
    349:   PetscCall(VecDestroy(&rayctx->ray));
    350:   PetscCall(VecScatterDestroy(&rayctx->scatter));
    351:   PetscCall(PetscViewerDestroy(&rayctx->viewer));
    352:   PetscCall(PetscFree(rayctx));
    353:   PetscFunctionReturn(PETSC_SUCCESS);
    354: }

    356: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
    357: {
    358:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
    359:   Vec                  solution;

    361:   PetscFunctionBegin;
    362:   PetscCall(TSGetSolution(ts, &solution));
    363:   PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
    364:   PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
    365:   if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer));
    366:   PetscFunctionReturn(PETSC_SUCCESS);
    367: }

    369: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, PetscCtx ctx)
    370: {
    371:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
    372:   TSMonitorLGCtx       lgctx  = rayctx->lgctx;
    373:   Vec                  v      = rayctx->ray;
    374:   const PetscScalar   *a;
    375:   PetscInt             dim;

    377:   PetscFunctionBegin;
    378:   PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
    379:   PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
    380:   if (!step) {
    381:     PetscDrawAxis axis;

    383:     PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis));
    384:     PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution"));
    385:     PetscCall(VecGetLocalSize(rayctx->ray, &dim));
    386:     PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim));
    387:     PetscCall(PetscDrawLGReset(lgctx->lg));
    388:   }
    389:   PetscCall(VecGetArrayRead(v, &a));
    390: #if defined(PETSC_USE_COMPLEX)
    391:   {
    392:     PetscReal *areal;
    393:     PetscInt   i, n;
    394:     PetscCall(VecGetLocalSize(v, &n));
    395:     PetscCall(PetscMalloc1(n, &areal));
    396:     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
    397:     PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal));
    398:     PetscCall(PetscFree(areal));
    399:   }
    400: #else
    401:   PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a));
    402: #endif
    403:   PetscCall(VecRestoreArrayRead(v, &a));
    404:   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
    405:     PetscCall(PetscDrawLGDraw(lgctx->lg));
    406:     PetscCall(PetscDrawLGSave(lgctx->lg));
    407:   }
    408:   PetscFunctionReturn(PETSC_SUCCESS);
    409: }