Actual source code: mimex.c

      1: /*
      2:        Code for Timestepping with my makeshift IMEX.
      3: */
      4: #include <petsc/private/tsimpl.h>
      5: #include <petscds.h>
      6: #include <petscsection.h>
      7: #include <petscdmplex.h>

      9: typedef struct {
     10:   Vec       Xdot, update;
     11:   PetscReal stage_time;
     12:   PetscInt  version;
     13: } TS_Mimex;

     15: static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     16: {
     17:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

     19:   PetscFunctionBegin;
     20:   if (X0) {
     21:     if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_X0", X0));
     22:     else *X0 = ts->vec_sol;
     23:   }
     24:   if (Xdot) {
     25:     if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot));
     26:     else *Xdot = mimex->Xdot;
     27:   }
     28:   PetscFunctionReturn(PETSC_SUCCESS);
     29: }

     31: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     32: {
     33:   PetscFunctionBegin;
     34:   if (X0)
     35:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0));
     36:   if (Xdot)
     37:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot));
     38:   PetscFunctionReturn(PETSC_SUCCESS);
     39: }

     41: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     42: {
     43:   PetscFunctionBegin;
     44:   PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar));
     45:   PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_G", G));
     46:   PetscFunctionReturn(PETSC_SUCCESS);
     47: }

     49: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     50: {
     51:   PetscFunctionBegin;
     52:   PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar));
     53:   PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_G", G));
     54:   PetscFunctionReturn(PETSC_SUCCESS);
     55: }

     57: /*
     58:   This defines the nonlinear equation that is to be solved with SNES
     59:   G(U) = F[t0+dt, U, (U-U0)*shift] = 0
     60: */
     61: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
     62: {
     63:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
     64:   DM        dm, dmsave;
     65:   Vec       X0, Xdot;
     66:   PetscReal shift = 1. / ts->time_step;

     68:   PetscFunctionBegin;
     69:   PetscCall(SNESGetDM(snes, &dm));
     70:   PetscCall(TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot));
     71:   PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));

     73:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
     74:   dmsave = ts->dm;
     75:   ts->dm = dm;
     76:   PetscCall(TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE));
     77:   if (mimex->version == 1) {
     78:     DM                 dm;
     79:     PetscDS            prob;
     80:     PetscSection       s;
     81:     Vec                Xstar = NULL, G = NULL;
     82:     const PetscScalar *ax;
     83:     PetscScalar       *axstar;
     84:     PetscInt           Nf, f, pStart, pEnd, p;

     86:     PetscCall(TSGetDM(ts, &dm));
     87:     PetscCall(DMGetDS(dm, &prob));
     88:     PetscCall(DMGetLocalSection(dm, &s));
     89:     PetscCall(PetscDSGetNumFields(prob, &Nf));
     90:     PetscCall(PetscSectionGetChart(s, &pStart, &pEnd));
     91:     PetscCall(TSMimexGetXstarAndG(ts, dm, &Xstar, &G));
     92:     PetscCall(VecCopy(X0, Xstar));
     93:     PetscCall(VecGetArrayRead(x, &ax));
     94:     PetscCall(VecGetArray(Xstar, &axstar));
     95:     for (f = 0; f < Nf; ++f) {
     96:       PetscBool implicit;

     98:       PetscCall(PetscDSGetImplicit(prob, f, &implicit));
     99:       if (!implicit) continue;
    100:       for (p = pStart; p < pEnd; ++p) {
    101:         PetscScalar *a, *axs;
    102:         PetscInt     fdof, fcdof, d;

    104:         PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof));
    105:         PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof));
    106:         PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, ax, &a));
    107:         PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs));
    108:         for (d = 0; d < fdof - fcdof; ++d) axs[d] = a[d];
    109:       }
    110:     }
    111:     PetscCall(VecRestoreArrayRead(x, &ax));
    112:     PetscCall(VecRestoreArray(Xstar, &axstar));
    113:     PetscCall(TSComputeRHSFunction(ts, ts->ptime, Xstar, G));
    114:     PetscCall(VecAXPY(y, -1.0, G));
    115:     PetscCall(TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G));
    116:   }
    117:   ts->dm = dmsave;
    118:   PetscCall(TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot));
    119:   PetscFunctionReturn(PETSC_SUCCESS);
    120: }

    122: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
    123: {
    124:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
    125:   DM        dm, dmsave;
    126:   Vec       Xdot;
    127:   PetscReal shift = 1. / ts->time_step;

    129:   PetscFunctionBegin;
    130:   /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
    131:   PetscCall(SNESGetDM(snes, &dm));
    132:   PetscCall(TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot));

    134:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    135:   dmsave = ts->dm;
    136:   ts->dm = dm;
    137:   PetscCall(TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE));
    138:   ts->dm = dmsave;
    139:   PetscCall(TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot));
    140:   PetscFunctionReturn(PETSC_SUCCESS);
    141: }

    143: static PetscErrorCode TSStep_Mimex_Split(TS ts)
    144: {
    145:   TS_Mimex          *mimex = (TS_Mimex *)ts->data;
    146:   DM                 dm;
    147:   PetscDS            prob;
    148:   PetscSection       s;
    149:   Vec                sol = ts->vec_sol, update = mimex->update;
    150:   const PetscScalar *aupdate;
    151:   PetscScalar       *asol, dt = ts->time_step;
    152:   PetscInt           Nf, f, pStart, pEnd, p;

    154:   PetscFunctionBegin;
    155:   PetscCall(TSGetDM(ts, &dm));
    156:   PetscCall(DMGetDS(dm, &prob));
    157:   PetscCall(DMGetLocalSection(dm, &s));
    158:   PetscCall(PetscDSGetNumFields(prob, &Nf));
    159:   PetscCall(PetscSectionGetChart(s, &pStart, &pEnd));
    160:   PetscCall(TSPreStage(ts, ts->ptime));
    161:   /* Compute implicit update */
    162:   mimex->stage_time = ts->ptime + ts->time_step;
    163:   PetscCall(VecCopy(sol, update));
    164:   PetscCall(SNESSolve(ts->snes, NULL, update));
    165:   PetscCall(VecGetArrayRead(update, &aupdate));
    166:   PetscCall(VecGetArray(sol, &asol));
    167:   for (f = 0; f < Nf; ++f) {
    168:     PetscBool implicit;

    170:     PetscCall(PetscDSGetImplicit(prob, f, &implicit));
    171:     if (!implicit) continue;
    172:     for (p = pStart; p < pEnd; ++p) {
    173:       PetscScalar *au, *as;
    174:       PetscInt     fdof, fcdof, d;

    176:       PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof));
    177:       PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof));
    178:       PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au));
    179:       PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as));
    180:       for (d = 0; d < fdof - fcdof; ++d) as[d] = au[d];
    181:     }
    182:   }
    183:   PetscCall(VecRestoreArrayRead(update, &aupdate));
    184:   PetscCall(VecRestoreArray(sol, &asol));
    185:   /* Compute explicit update */
    186:   PetscCall(TSComputeRHSFunction(ts, ts->ptime, sol, update));
    187:   PetscCall(VecGetArrayRead(update, &aupdate));
    188:   PetscCall(VecGetArray(sol, &asol));
    189:   for (f = 0; f < Nf; ++f) {
    190:     PetscBool implicit;

    192:     PetscCall(PetscDSGetImplicit(prob, f, &implicit));
    193:     if (implicit) continue;
    194:     for (p = pStart; p < pEnd; ++p) {
    195:       PetscScalar *au, *as;
    196:       PetscInt     fdof, fcdof, d;

    198:       PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof));
    199:       PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof));
    200:       PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au));
    201:       PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as));
    202:       for (d = 0; d < fdof - fcdof; ++d) as[d] += dt * au[d];
    203:     }
    204:   }
    205:   PetscCall(VecRestoreArrayRead(update, &aupdate));
    206:   PetscCall(VecRestoreArray(sol, &asol));
    207:   PetscCall(TSPostStage(ts, ts->ptime, 0, &sol));
    208:   ts->ptime += ts->time_step;
    209:   PetscFunctionReturn(PETSC_SUCCESS);
    210: }

    212: /* Evaluate F at U and G at U0 for explicit fields and U for implicit fields */
    213: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
    214: {
    215:   TS_Mimex *mimex  = (TS_Mimex *)ts->data;
    216:   Vec       sol    = ts->vec_sol;
    217:   Vec       update = mimex->update;

    219:   PetscFunctionBegin;
    220:   PetscCall(TSPreStage(ts, ts->ptime));
    221:   /* Compute implicit update */
    222:   mimex->stage_time = ts->ptime + ts->time_step;
    223:   ts->ptime += ts->time_step;
    224:   PetscCall(VecCopy(sol, update));
    225:   PetscCall(SNESSolve(ts->snes, NULL, update));
    226:   PetscCall(VecCopy(update, sol));
    227:   PetscCall(TSPostStage(ts, ts->ptime, 0, &sol));
    228:   PetscFunctionReturn(PETSC_SUCCESS);
    229: }

    231: static PetscErrorCode TSStep_Mimex(TS ts)
    232: {
    233:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

    235:   PetscFunctionBegin;
    236:   switch (mimex->version) {
    237:   case 0:
    238:     PetscCall(TSStep_Mimex_Split(ts));
    239:     break;
    240:   case 1:
    241:     PetscCall(TSStep_Mimex_Implicit(ts));
    242:     break;
    243:   default:
    244:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %" PetscInt_FMT, mimex->version);
    245:   }
    246:   PetscFunctionReturn(PETSC_SUCCESS);
    247: }

    249: /*------------------------------------------------------------*/

    251: static PetscErrorCode TSSetUp_Mimex(TS ts)
    252: {
    253:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

    255:   PetscFunctionBegin;
    256:   PetscCall(VecDuplicate(ts->vec_sol, &mimex->update));
    257:   PetscCall(VecDuplicate(ts->vec_sol, &mimex->Xdot));
    258:   PetscFunctionReturn(PETSC_SUCCESS);
    259: }

    261: static PetscErrorCode TSReset_Mimex(TS ts)
    262: {
    263:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

    265:   PetscFunctionBegin;
    266:   PetscCall(VecDestroy(&mimex->update));
    267:   PetscCall(VecDestroy(&mimex->Xdot));
    268:   PetscFunctionReturn(PETSC_SUCCESS);
    269: }

    271: static PetscErrorCode TSDestroy_Mimex(TS ts)
    272: {
    273:   PetscFunctionBegin;
    274:   PetscCall(TSReset_Mimex(ts));
    275:   PetscCall(PetscFree(ts->data));
    276:   PetscFunctionReturn(PETSC_SUCCESS);
    277: }
    278: /*------------------------------------------------------------*/

    280: static PetscErrorCode TSSetFromOptions_Mimex(TS ts, PetscOptionItems PetscOptionsObject)
    281: {
    282:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

    284:   PetscFunctionBegin;
    285:   PetscOptionsHeadBegin(PetscOptionsObject, "MIMEX ODE solver options");
    286:   {
    287:     PetscCall(PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL));
    288:   }
    289:   PetscOptionsHeadEnd();
    290:   PetscFunctionReturn(PETSC_SUCCESS);
    291: }

    293: static PetscErrorCode TSView_Mimex(TS ts, PetscViewer viewer)
    294: {
    295:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
    296:   PetscBool isascii;

    298:   PetscFunctionBegin;
    299:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    300:   if (isascii) PetscCall(PetscViewerASCIIPrintf(viewer, "  Version = %" PetscInt_FMT "\n", mimex->version));
    301:   PetscFunctionReturn(PETSC_SUCCESS);
    302: }

    304: static PetscErrorCode TSInterpolate_Mimex(TS ts, PetscReal t, Vec X)
    305: {
    306:   PetscReal alpha = (ts->ptime - t) / ts->time_step;

    308:   PetscFunctionBegin;
    309:   PetscCall(VecAXPBY(ts->vec_sol, 1.0 - alpha, alpha, X));
    310:   PetscFunctionReturn(PETSC_SUCCESS);
    311: }

    313: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
    314: {
    315:   PetscFunctionBegin;
    316:   *yr = 1.0 + xr;
    317:   *yi = xi;
    318:   PetscFunctionReturn(PETSC_SUCCESS);
    319: }
    320: /* ------------------------------------------------------------ */

    322: /*MC
    323:       TSMIMEX - ODE solver using the explicit forward Mimex method

    325:   Level: beginner

    327: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`
    328: M*/
    329: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
    330: {
    331:   TS_Mimex *mimex;

    333:   PetscFunctionBegin;
    334:   ts->ops->setup           = TSSetUp_Mimex;
    335:   ts->ops->step            = TSStep_Mimex;
    336:   ts->ops->reset           = TSReset_Mimex;
    337:   ts->ops->destroy         = TSDestroy_Mimex;
    338:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
    339:   ts->ops->view            = TSView_Mimex;
    340:   ts->ops->interpolate     = TSInterpolate_Mimex;
    341:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
    342:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
    343:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
    344:   ts->default_adapt_type   = TSADAPTNONE;

    346:   PetscCall(PetscNew(&mimex));
    347:   ts->data = (void *)mimex;

    349:   mimex->version = 1;
    350:   PetscFunctionReturn(PETSC_SUCCESS);
    351: }