Date: 2022-03-31T04:48:43+00:00

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:   if (X0) {
     20:     if (dm && dm != ts->dm) DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);
     21:     else                    {*X0  = ts->vec_sol;}
     22:   }
     23:   if (Xdot) {
     24:     if (dm && dm != ts->dm) DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);
     25:     else                    {*Xdot = mimex->Xdot;}
     26:   }
     27:   return 0;
     28: }

     30: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     31: {
     32:   if (X0)   if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);
     33:   if (Xdot) if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);
     34:   return 0;
     35: }

     37: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     38: {
     39:   DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
     40:   DMGetNamedGlobalVector(dm, "TSMimex_G", G);
     41:   return 0;
     42: }

     44: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     45: {
     46:   DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
     47:   DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
     48:   return 0;
     49: }

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

     62:   SNESGetDM(snes, &dm);
     63:   TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
     64:   VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);

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

     79:     TSGetDM(ts, &dm);
     80:     DMGetDS(dm, &prob);
     81:     DMGetLocalSection(dm, &s);
     82:     PetscDSGetNumFields(prob, &Nf);
     83:     PetscSectionGetChart(s, &pStart, &pEnd);
     84:     TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
     85:     VecCopy(X0, Xstar);
     86:     VecGetArrayRead(x, &ax);
     87:     VecGetArray(Xstar, &axstar);
     88:     for (f = 0; f < Nf; ++f) {
     89:       PetscBool implicit;

     91:       PetscDSGetImplicit(prob, f, &implicit);
     92:       if (!implicit) continue;
     93:       for (p = pStart; p < pEnd; ++p) {
     94:         PetscScalar *a, *axs;
     95:         PetscInt     fdof, fcdof, d;

     97:         PetscSectionGetFieldDof(s, p, f, &fdof);
     98:         PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
     99:         DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
    100:         DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
    101:         for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
    102:       }
    103:     }
    104:     VecRestoreArrayRead(x, &ax);
    105:     VecRestoreArray(Xstar, &axstar);
    106:     TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
    107:     VecAXPY(y, -1.0, G);
    108:     TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
    109:   }
    110:   ts->dm = dmsave;
    111:   TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
    112:   return 0;
    113: }

    115: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
    116: {
    117:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
    118:   DM             dm, dmsave;
    119:   Vec            Xdot;
    120:   PetscReal      shift = 1./ts->time_step;

    122:   /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
    123:   SNESGetDM(snes, &dm);
    124:   TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);

    126:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    127:   dmsave = ts->dm;
    128:   ts->dm = dm;
    129:   TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
    130:   ts->dm = dmsave;
    131:   TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
    132:   return 0;
    133: }

    135: static PetscErrorCode TSStep_Mimex_Split(TS ts)
    136: {
    137:   TS_Mimex          *mimex = (TS_Mimex *) ts->data;
    138:   DM                 dm;
    139:   PetscDS            prob;
    140:   PetscSection       s;
    141:   Vec                sol = ts->vec_sol, update = mimex->update;
    142:   const PetscScalar *aupdate;
    143:   PetscScalar       *asol, dt = ts->time_step;
    144:   PetscInt           Nf, f, pStart, pEnd, p;

    146:   TSGetDM(ts, &dm);
    147:   DMGetDS(dm, &prob);
    148:   DMGetLocalSection(dm, &s);
    149:   PetscDSGetNumFields(prob, &Nf);
    150:   PetscSectionGetChart(s, &pStart, &pEnd);
    151:   TSPreStage(ts, ts->ptime);
    152:   /* Compute implicit update */
    153:   mimex->stage_time = ts->ptime + ts->time_step;
    154:   VecCopy(sol, update);
    155:   SNESSolve(ts->snes, NULL, update);
    156:   VecGetArrayRead(update, &aupdate);
    157:   VecGetArray(sol, &asol);
    158:   for (f = 0; f < Nf; ++f) {
    159:     PetscBool implicit;

    161:     PetscDSGetImplicit(prob, f, &implicit);
    162:     if (!implicit) continue;
    163:     for (p = pStart; p < pEnd; ++p) {
    164:       PetscScalar *au, *as;
    165:       PetscInt     fdof, fcdof, d;

    167:       PetscSectionGetFieldDof(s, p, f, &fdof);
    168:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
    169:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
    170:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
    171:       for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
    172:     }
    173:   }
    174:   VecRestoreArrayRead(update, &aupdate);
    175:   VecRestoreArray(sol, &asol);
    176:   /* Compute explicit update */
    177:   TSComputeRHSFunction(ts, ts->ptime, sol, update);
    178:   VecGetArrayRead(update, &aupdate);
    179:   VecGetArray(sol, &asol);
    180:   for (f = 0; f < Nf; ++f) {
    181:     PetscBool implicit;

    183:     PetscDSGetImplicit(prob, f, &implicit);
    184:     if (implicit) continue;
    185:     for (p = pStart; p < pEnd; ++p) {
    186:       PetscScalar *au, *as;
    187:       PetscInt     fdof, fcdof, d;

    189:       PetscSectionGetFieldDof(s, p, f, &fdof);
    190:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
    191:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
    192:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
    193:       for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
    194:     }
    195:   }
    196:   VecRestoreArrayRead(update, &aupdate);
    197:   VecRestoreArray(sol, &asol);
    198:   TSPostStage(ts, ts->ptime, 0, &sol);
    199:   ts->ptime += ts->time_step;
    200:   return 0;
    201: }

    203: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
    204: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
    205: {
    206:   TS_Mimex      *mimex  = (TS_Mimex *) ts->data;
    207:   Vec            sol    = ts->vec_sol;
    208:   Vec            update = mimex->update;

    210:   TSPreStage(ts, ts->ptime);
    211:   /* Compute implicit update */
    212:   mimex->stage_time = ts->ptime + ts->time_step;
    213:   ts->ptime += ts->time_step;
    214:   VecCopy(sol, update);
    215:   SNESSolve(ts->snes, NULL, update);
    216:   VecCopy(update, sol);
    217:   TSPostStage(ts, ts->ptime, 0, &sol);
    218:   return 0;
    219: }

    221: static PetscErrorCode TSStep_Mimex(TS ts)
    222: {
    223:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

    225:   switch(mimex->version) {
    226:   case 0:
    227:     TSStep_Mimex_Split(ts); break;
    228:   case 1:
    229:     TSStep_Mimex_Implicit(ts); break;
    230:   default:
    231:     SETERRQ(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
    232:   }
    233:   return 0;
    234: }

    236: /*------------------------------------------------------------*/

    238: static PetscErrorCode TSSetUp_Mimex(TS ts)
    239: {
    240:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

    242:   VecDuplicate(ts->vec_sol, &mimex->update);
    243:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
    244:   return 0;
    245: }

    247: static PetscErrorCode TSReset_Mimex(TS ts)
    248: {
    249:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

    251:   VecDestroy(&mimex->update);
    252:   VecDestroy(&mimex->Xdot);
    253:   return 0;
    254: }

    256: static PetscErrorCode TSDestroy_Mimex(TS ts)
    257: {
    258:   TSReset_Mimex(ts);
    259:   PetscFree(ts->data);
    260:   return 0;
    261: }
    262: /*------------------------------------------------------------*/

    264: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
    265: {
    266:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

    268:   PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
    269:   {
    270:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
    271:   }
    272:   PetscOptionsTail();
    273:   return 0;
    274: }

    276: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
    277: {
    278:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
    279:   PetscBool      iascii;

    281:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
    282:   if (iascii) {
    283:     PetscViewerASCIIPrintf(viewer, "  Version = %D\n", mimex->version);
    284:   }
    285:   return 0;
    286: }

    288: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
    289: {
    290:   PetscReal      alpha = (ts->ptime - t)/ts->time_step;

    292:   VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
    293:   return 0;
    294: }

    296: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
    297: {
    298:   *yr = 1.0 + xr;
    299:   *yi = xi;
    300:   return 0;
    301: }
    302: /* ------------------------------------------------------------ */

    304: /*MC
    305:       TSMIMEX - ODE solver using the explicit forward Mimex method

    307:   Level: beginner

    309: .seealso:  TSCreate(), TS, TSSetType(), TSBEULER

    311: M*/
    312: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
    313: {
    314:   TS_Mimex       *mimex;

    316:   ts->ops->setup           = TSSetUp_Mimex;
    317:   ts->ops->step            = TSStep_Mimex;
    318:   ts->ops->reset           = TSReset_Mimex;
    319:   ts->ops->destroy         = TSDestroy_Mimex;
    320:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
    321:   ts->ops->view            = TSView_Mimex;
    322:   ts->ops->interpolate     = TSInterpolate_Mimex;
    323:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
    324:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
    325:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
    326:   ts->default_adapt_type   = TSADAPTNONE;

    328:   PetscNewLog(ts,&mimex);
    329:   ts->data = (void*)mimex;

    331:   mimex->version = 1;
    332:   return 0;
    333: }