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: }