Actual source code: mimex.c
petsc-3.10.5 2019-03-28
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 <petscdmplex.h>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    257: /*------------------------------------------------------------*/

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    342:   Level: beginner

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

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

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

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

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