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