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