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

 31: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
 32: {
 33:   PetscFunctionBegin;
 34:   if (X0)
 35:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0));
 36:   if (Xdot)
 37:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot));
 38:   PetscFunctionReturn(PETSC_SUCCESS);
 39: }

 41: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
 42: {
 43:   PetscFunctionBegin;
 44:   PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar));
 45:   PetscCall(DMGetNamedGlobalVector(dm, "TSMimex_G", G));
 46:   PetscFunctionReturn(PETSC_SUCCESS);
 47: }

 49: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
 50: {
 51:   PetscFunctionBegin;
 52:   PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar));
 53:   PetscCall(DMRestoreNamedGlobalVector(dm, "TSMimex_G", G));
 54:   PetscFunctionReturn(PETSC_SUCCESS);
 55: }

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

 68:   PetscFunctionBegin;
 69:   PetscCall(SNESGetDM(snes, &dm));
 70:   PetscCall(TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot));
 71:   PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));

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

 86:     PetscCall(TSGetDM(ts, &dm));
 87:     PetscCall(DMGetDS(dm, &prob));
 88:     PetscCall(DMGetLocalSection(dm, &s));
 89:     PetscCall(PetscDSGetNumFields(prob, &Nf));
 90:     PetscCall(PetscSectionGetChart(s, &pStart, &pEnd));
 91:     PetscCall(TSMimexGetXstarAndG(ts, dm, &Xstar, &G));
 92:     PetscCall(VecCopy(X0, Xstar));
 93:     PetscCall(VecGetArrayRead(x, &ax));
 94:     PetscCall(VecGetArray(Xstar, &axstar));
 95:     for (f = 0; f < Nf; ++f) {
 96:       PetscBool implicit;

 98:       PetscCall(PetscDSGetImplicit(prob, f, &implicit));
 99:       if (!implicit) continue;
100:       for (p = pStart; p < pEnd; ++p) {
101:         PetscScalar *a, *axs;
102:         PetscInt     fdof, fcdof, d;

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

122: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
123: {
124:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
125:   DM        dm, dmsave;
126:   Vec       Xdot;
127:   PetscReal shift = 1. / ts->time_step;

129:   PetscFunctionBegin;
130:   /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
131:   PetscCall(SNESGetDM(snes, &dm));
132:   PetscCall(TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot));

134:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
135:   dmsave = ts->dm;
136:   ts->dm = dm;
137:   PetscCall(TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE));
138:   ts->dm = dmsave;
139:   PetscCall(TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot));
140:   PetscFunctionReturn(PETSC_SUCCESS);
141: }

143: static PetscErrorCode TSStep_Mimex_Split(TS ts)
144: {
145:   TS_Mimex          *mimex = (TS_Mimex *)ts->data;
146:   DM                 dm;
147:   PetscDS            prob;
148:   PetscSection       s;
149:   Vec                sol = ts->vec_sol, update = mimex->update;
150:   const PetscScalar *aupdate;
151:   PetscScalar       *asol, dt = ts->time_step;
152:   PetscInt           Nf, f, pStart, pEnd, p;

154:   PetscFunctionBegin;
155:   PetscCall(TSGetDM(ts, &dm));
156:   PetscCall(DMGetDS(dm, &prob));
157:   PetscCall(DMGetLocalSection(dm, &s));
158:   PetscCall(PetscDSGetNumFields(prob, &Nf));
159:   PetscCall(PetscSectionGetChart(s, &pStart, &pEnd));
160:   PetscCall(TSPreStage(ts, ts->ptime));
161:   /* Compute implicit update */
162:   mimex->stage_time = ts->ptime + ts->time_step;
163:   PetscCall(VecCopy(sol, update));
164:   PetscCall(SNESSolve(ts->snes, NULL, update));
165:   PetscCall(VecGetArrayRead(update, &aupdate));
166:   PetscCall(VecGetArray(sol, &asol));
167:   for (f = 0; f < Nf; ++f) {
168:     PetscBool implicit;

170:     PetscCall(PetscDSGetImplicit(prob, f, &implicit));
171:     if (!implicit) continue;
172:     for (p = pStart; p < pEnd; ++p) {
173:       PetscScalar *au, *as;
174:       PetscInt     fdof, fcdof, d;

176:       PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof));
177:       PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof));
178:       PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au));
179:       PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as));
180:       for (d = 0; d < fdof - fcdof; ++d) as[d] = au[d];
181:     }
182:   }
183:   PetscCall(VecRestoreArrayRead(update, &aupdate));
184:   PetscCall(VecRestoreArray(sol, &asol));
185:   /* Compute explicit update */
186:   PetscCall(TSComputeRHSFunction(ts, ts->ptime, sol, update));
187:   PetscCall(VecGetArrayRead(update, &aupdate));
188:   PetscCall(VecGetArray(sol, &asol));
189:   for (f = 0; f < Nf; ++f) {
190:     PetscBool implicit;

192:     PetscCall(PetscDSGetImplicit(prob, f, &implicit));
193:     if (implicit) continue;
194:     for (p = pStart; p < pEnd; ++p) {
195:       PetscScalar *au, *as;
196:       PetscInt     fdof, fcdof, d;

198:       PetscCall(PetscSectionGetFieldDof(s, p, f, &fdof));
199:       PetscCall(PetscSectionGetFieldConstraintDof(s, p, f, &fcdof));
200:       PetscCall(DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au));
201:       PetscCall(DMPlexPointGlobalFieldRef(dm, p, f, asol, &as));
202:       for (d = 0; d < fdof - fcdof; ++d) as[d] += dt * au[d];
203:     }
204:   }
205:   PetscCall(VecRestoreArrayRead(update, &aupdate));
206:   PetscCall(VecRestoreArray(sol, &asol));
207:   PetscCall(TSPostStage(ts, ts->ptime, 0, &sol));
208:   ts->ptime += ts->time_step;
209:   PetscFunctionReturn(PETSC_SUCCESS);
210: }

212: /* Evaluate F at U and G at U0 for explicit fields and U for implicit fields */
213: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
214: {
215:   TS_Mimex *mimex  = (TS_Mimex *)ts->data;
216:   Vec       sol    = ts->vec_sol;
217:   Vec       update = mimex->update;

219:   PetscFunctionBegin;
220:   PetscCall(TSPreStage(ts, ts->ptime));
221:   /* Compute implicit update */
222:   mimex->stage_time = ts->ptime + ts->time_step;
223:   ts->ptime += ts->time_step;
224:   PetscCall(VecCopy(sol, update));
225:   PetscCall(SNESSolve(ts->snes, NULL, update));
226:   PetscCall(VecCopy(update, sol));
227:   PetscCall(TSPostStage(ts, ts->ptime, 0, &sol));
228:   PetscFunctionReturn(PETSC_SUCCESS);
229: }

231: static PetscErrorCode TSStep_Mimex(TS ts)
232: {
233:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

235:   PetscFunctionBegin;
236:   switch (mimex->version) {
237:   case 0:
238:     PetscCall(TSStep_Mimex_Split(ts));
239:     break;
240:   case 1:
241:     PetscCall(TSStep_Mimex_Implicit(ts));
242:     break;
243:   default:
244:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %" PetscInt_FMT, mimex->version);
245:   }
246:   PetscFunctionReturn(PETSC_SUCCESS);
247: }

249: /*------------------------------------------------------------*/

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

255:   PetscFunctionBegin;
256:   PetscCall(VecDuplicate(ts->vec_sol, &mimex->update));
257:   PetscCall(VecDuplicate(ts->vec_sol, &mimex->Xdot));
258:   PetscFunctionReturn(PETSC_SUCCESS);
259: }

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

265:   PetscFunctionBegin;
266:   PetscCall(VecDestroy(&mimex->update));
267:   PetscCall(VecDestroy(&mimex->Xdot));
268:   PetscFunctionReturn(PETSC_SUCCESS);
269: }

271: static PetscErrorCode TSDestroy_Mimex(TS ts)
272: {
273:   PetscFunctionBegin;
274:   PetscCall(TSReset_Mimex(ts));
275:   PetscCall(PetscFree(ts->data));
276:   PetscFunctionReturn(PETSC_SUCCESS);
277: }
278: /*------------------------------------------------------------*/

280: static PetscErrorCode TSSetFromOptions_Mimex(TS ts, PetscOptionItems *PetscOptionsObject)
281: {
282:   TS_Mimex *mimex = (TS_Mimex *)ts->data;

284:   PetscFunctionBegin;
285:   PetscOptionsHeadBegin(PetscOptionsObject, "MIMEX ODE solver options");
286:   {
287:     PetscCall(PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL));
288:   }
289:   PetscOptionsHeadEnd();
290:   PetscFunctionReturn(PETSC_SUCCESS);
291: }

293: static PetscErrorCode TSView_Mimex(TS ts, PetscViewer viewer)
294: {
295:   TS_Mimex *mimex = (TS_Mimex *)ts->data;
296:   PetscBool iascii;

298:   PetscFunctionBegin;
299:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
300:   if (iascii) PetscCall(PetscViewerASCIIPrintf(viewer, "  Version = %" PetscInt_FMT "\n", mimex->version));
301:   PetscFunctionReturn(PETSC_SUCCESS);
302: }

304: static PetscErrorCode TSInterpolate_Mimex(TS ts, PetscReal t, Vec X)
305: {
306:   PetscReal alpha = (ts->ptime - t) / ts->time_step;

308:   PetscFunctionBegin;
309:   PetscCall(VecAXPBY(ts->vec_sol, 1.0 - alpha, alpha, X));
310:   PetscFunctionReturn(PETSC_SUCCESS);
311: }

313: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
314: {
315:   PetscFunctionBegin;
316:   *yr = 1.0 + xr;
317:   *yi = xi;
318:   PetscFunctionReturn(PETSC_SUCCESS);
319: }
320: /* ------------------------------------------------------------ */

322: /*MC
323:       TSMIMEX - ODE solver using the explicit forward Mimex method

325:   Level: beginner

327: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`
328: M*/
329: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
330: {
331:   TS_Mimex *mimex;

333:   PetscFunctionBegin;
334:   ts->ops->setup           = TSSetUp_Mimex;
335:   ts->ops->step            = TSStep_Mimex;
336:   ts->ops->reset           = TSReset_Mimex;
337:   ts->ops->destroy         = TSDestroy_Mimex;
338:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
339:   ts->ops->view            = TSView_Mimex;
340:   ts->ops->interpolate     = TSInterpolate_Mimex;
341:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
342:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
343:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
344:   ts->default_adapt_type   = TSADAPTNONE;

346:   PetscCall(PetscNew(&mimex));
347:   ts->data = (void *)mimex;

349:   mimex->version = 1;
350:   PetscFunctionReturn(PETSC_SUCCESS);
351: }