Actual source code: mimex.c
petsc-3.13.6 2020-09-29
Report Typos and Errors

      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;

     21:   if (X0) {
     22:     if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);}
     23:     else                    {*X0  = ts->vec_sol;}
     24:   }
     25:   if (Xdot) {
     26:     if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
     27:     else                    {*Xdot = mimex->Xdot;}
     28:   }
     29:   return(0);
     30: }

     32: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     33: {

     37:   if (X0)   if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);}
     38:   if (Xdot) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
     39:   return(0);
     40: }

     42: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     43: {

     47:   DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
     48:   DMGetNamedGlobalVector(dm, "TSMimex_G", G);
     49:   return(0);
     50: }

     52: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
     53: {

     57:   DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
     58:   DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
     59:   return(0);
     60: }

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

     75:   SNESGetDM(snes, &dm);
     76:   TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
     77:   VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);

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

     92:     TSGetDM(ts, &dm);
     93:     DMGetDS(dm, &prob);
     94:     DMGetLocalSection(dm, &s);
     95:     PetscDSGetNumFields(prob, &Nf);
     96:     PetscSectionGetChart(s, &pStart, &pEnd);
     97:     TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
     98:     VecCopy(X0, Xstar);
     99:     VecGetArrayRead(x, &ax);
    100:     VecGetArray(Xstar, &axstar);
    101:     for (f = 0; f < Nf; ++f) {
    102:       PetscBool implicit;

    104:       PetscDSGetImplicit(prob, f, &implicit);
    105:       if (!implicit) continue;
    106:       for (p = pStart; p < pEnd; ++p) {
    107:         PetscScalar *a, *axs;
    108:         PetscInt     fdof, fcdof, d;

    110:         PetscSectionGetFieldDof(s, p, f, &fdof);
    111:         PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
    112:         DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
    113:         DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
    114:         for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
    115:       }
    116:     }
    117:     VecRestoreArrayRead(x, &ax);
    118:     VecRestoreArray(Xstar, &axstar);
    119:     TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
    120:     VecAXPY(y, -1.0, G);
    121:     TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
    122:   }
    123:   ts->dm = dmsave;
    124:   TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
    125:   return(0);
    126: }

    128: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
    129: {
    130:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
    131:   DM             dm, dmsave;
    132:   Vec            Xdot;
    133:   PetscReal      shift = 1./ts->time_step;

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

    141:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    142:   dmsave = ts->dm;
    143:   ts->dm = dm;
    144:   TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
    145:   ts->dm = dmsave;
    146:   TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
    147:   return(0);
    148: }

    150: static PetscErrorCode TSStep_Mimex_Split(TS ts)
    151: {
    152:   TS_Mimex          *mimex = (TS_Mimex *) ts->data;
    153:   DM                 dm;
    154:   PetscDS            prob;
    155:   PetscSection       s;
    156:   Vec                sol = ts->vec_sol, update = mimex->update;
    157:   const PetscScalar *aupdate;
    158:   PetscScalar       *asol, dt = ts->time_step;
    159:   PetscInt           Nf, f, pStart, pEnd, p;
    160:   PetscErrorCode     ierr;

    163:   TSGetDM(ts, &dm);
    164:   DMGetDS(dm, &prob);
    165:   DMGetLocalSection(dm, &s);
    166:   PetscDSGetNumFields(prob, &Nf);
    167:   PetscSectionGetChart(s, &pStart, &pEnd);
    168:   TSPreStage(ts, ts->ptime);
    169:   /* Compute implicit update */
    170:   mimex->stage_time = ts->ptime + ts->time_step;
    171:   VecCopy(sol, update);
    172:   SNESSolve(ts->snes, NULL, update);
    173:   VecGetArrayRead(update, &aupdate);
    174:   VecGetArray(sol, &asol);
    175:   for (f = 0; f < Nf; ++f) {
    176:     PetscBool implicit;

    178:     PetscDSGetImplicit(prob, f, &implicit);
    179:     if (!implicit) continue;
    180:     for (p = pStart; p < pEnd; ++p) {
    181:       PetscScalar *au, *as;
    182:       PetscInt     fdof, fcdof, d;

    184:       PetscSectionGetFieldDof(s, p, f, &fdof);
    185:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
    186:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
    187:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
    188:       for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
    189:     }
    190:   }
    191:   VecRestoreArrayRead(update, &aupdate);
    192:   VecRestoreArray(sol, &asol);
    193:   /* Compute explicit update */
    194:   TSComputeRHSFunction(ts, ts->ptime, sol, update);
    195:   VecGetArrayRead(update, &aupdate);
    196:   VecGetArray(sol, &asol);
    197:   for (f = 0; f < Nf; ++f) {
    198:     PetscBool implicit;

    200:     PetscDSGetImplicit(prob, f, &implicit);
    201:     if (implicit) continue;
    202:     for (p = pStart; p < pEnd; ++p) {
    203:       PetscScalar *au, *as;
    204:       PetscInt     fdof, fcdof, d;

    206:       PetscSectionGetFieldDof(s, p, f, &fdof);
    207:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
    208:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
    209:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
    210:       for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
    211:     }
    212:   }
    213:   VecRestoreArrayRead(update, &aupdate);
    214:   VecRestoreArray(sol, &asol);
    215:   TSPostStage(ts, ts->ptime, 0, &sol);
    216:   ts->ptime += ts->time_step;
    217:   return(0);
    218: }

    221: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
    222: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
    223: {
    224:   TS_Mimex      *mimex  = (TS_Mimex *) ts->data;
    225:   Vec            sol    = ts->vec_sol;
    226:   Vec            update = mimex->update;

    230:   TSPreStage(ts, ts->ptime);
    231:   /* Compute implicit update */
    232:   mimex->stage_time = ts->ptime + ts->time_step;
    233:   ts->ptime += ts->time_step;
    234:   VecCopy(sol, update);
    235:   SNESSolve(ts->snes, NULL, update);
    236:   VecCopy(update, sol);
    237:   TSPostStage(ts, ts->ptime, 0, &sol);
    238:   return(0);
    239: }

    241: static PetscErrorCode TSStep_Mimex(TS ts)
    242: {
    243:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
    244:   PetscErrorCode  ierr;

    247:   switch(mimex->version) {
    248:   case 0:
    249:     TSStep_Mimex_Split(ts); break;
    250:   case 1:
    251:     TSStep_Mimex_Implicit(ts); break;
    252:   default:
    253:     SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
    254:   }
    255:   return(0);
    256: }

    258: /*------------------------------------------------------------*/

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

    266:   VecDuplicate(ts->vec_sol, &mimex->update);
    267:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
    268:   return(0);
    269: }

    271: static PetscErrorCode TSReset_Mimex(TS ts)
    272: {
    273:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

    277:   VecDestroy(&mimex->update);
    278:   VecDestroy(&mimex->Xdot);
    279:   return(0);
    280: }

    282: static PetscErrorCode TSDestroy_Mimex(TS ts)
    283: {

    287:   TSReset_Mimex(ts);
    288:   PetscFree(ts->data);
    289:   return(0);
    290: }
    291: /*------------------------------------------------------------*/

    293: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
    294: {
    295:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

    299:   PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
    300:   {
    301:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
    302:   }
    303:   PetscOptionsTail();
    304:   return(0);
    305: }

    307: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
    308: {
    309:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
    310:   PetscBool      iascii;

    314:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
    315:   if (iascii) {
    316:     PetscViewerASCIIPrintf(viewer, "  Version = %D\n", mimex->version);
    317:   }
    318:   return(0);
    319: }

    321: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
    322: {
    323:   PetscReal      alpha = (ts->ptime - t)/ts->time_step;

    327:   VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
    328:   return(0);
    329: }

    331: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
    332: {
    334:   *yr = 1.0 + xr;
    335:   *yi = xi;
    336:   return(0);
    337: }
    338: /* ------------------------------------------------------------ */

    340: /*MC
    341:       TSMIMEX - ODE solver using the explicit forward Mimex method

    343:   Level: beginner

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

    347: M*/
    348: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
    349: {
    350:   TS_Mimex       *mimex;

    354:   ts->ops->setup           = TSSetUp_Mimex;
    355:   ts->ops->step            = TSStep_Mimex;
    356:   ts->ops->reset           = TSReset_Mimex;
    357:   ts->ops->destroy         = TSDestroy_Mimex;
    358:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
    359:   ts->ops->view            = TSView_Mimex;
    360:   ts->ops->interpolate     = TSInterpolate_Mimex;
    361:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
    362:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
    363:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
    364:   ts->default_adapt_type   = TSADAPTNONE;

    366:   PetscNewLog(ts,&mimex);
    367:   ts->data = (void*)mimex;

    369:   mimex->version = 1;
    370:   return(0);
    371: }