Date: 2025-03-28T21:11:19+00:00

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, void *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, void *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:   if (dmdats->ijacobianlocal) {
    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:   } else SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
    133:   /* This will be redundant if the user called both, but it's too common to forget. */
    134:   if (A != B) {
    135:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    136:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    137:   }
    138:   PetscFunctionReturn(PETSC_SUCCESS);
    139: }

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

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

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

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

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

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

    225:   Logically Collective

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

    233:   Level: beginner

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

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

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

    256:   Logically Collective

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

    263:   Level: beginner

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

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

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

    286:   Logically Collective

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

    294:   Level: beginner

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

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

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

    318:   Logically Collective

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

    325:   Level: beginner

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

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

    345: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
    346: {
    347:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;

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

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

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

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

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

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