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)
 33:     if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);
 34:   if (Xdot)
 35:     if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);
 36:   return 0;
 37: }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

240: /*------------------------------------------------------------*/

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

246:   VecDuplicate(ts->vec_sol, &mimex->update);
247:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
248:   return 0;
249: }

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

255:   VecDestroy(&mimex->update);
256:   VecDestroy(&mimex->Xdot);
257:   return 0;
258: }

260: static PetscErrorCode TSDestroy_Mimex(TS ts)
261: {
262:   TSReset_Mimex(ts);
263:   PetscFree(ts->data);
264:   return 0;
265: }
266: /*------------------------------------------------------------*/

268: static PetscErrorCode TSSetFromOptions_Mimex(TS ts, PetscOptionItems *PetscOptionsObject)
269: {
270:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

272:   PetscOptionsHeadBegin(PetscOptionsObject, "MIMEX ODE solver options");
273:   {
274:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
275:   }
276:   PetscOptionsHeadEnd();
277:   return 0;
278: }

280: static PetscErrorCode TSView_Mimex(TS ts, PetscViewer viewer)
281: {
282:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
283:   PetscBool iascii;

285:   PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii);
286:   if (iascii) PetscViewerASCIIPrintf(viewer, "  Version = %" PetscInt_FMT "\n", mimex->version);
287:   return 0;
288: }

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

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

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

306: /*MC
307:       TSMIMEX - ODE solver using the explicit forward Mimex method

309:   Level: beginner

311: .seealso: `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`

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

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

330:   PetscNew(&mimex);
331:   ts->data = (void *)mimex;

333:   mimex->version = 1;
334:   return 0;
335: }