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) if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);
33: if (Xdot) if (dm && dm != ts->dm) DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);
34: return 0;
35: }
37: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
38: {
39: DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
40: DMGetNamedGlobalVector(dm, "TSMimex_G", G);
41: return 0;
42: }
44: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
45: {
46: DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
47: DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
48: return 0;
49: }
51: /*
52: This defines the nonlinear equation that is to be solved with SNES
53: G(U) = F[t0+dt, U, (U-U0)*shift] = 0
54: */
55: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
56: {
57: TS_Mimex *mimex = (TS_Mimex *) ts->data;
58: DM dm, dmsave;
59: Vec X0, Xdot;
60: PetscReal shift = 1./ts->time_step;
62: SNESGetDM(snes, &dm);
63: TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
64: VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);
66: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
67: dmsave = ts->dm;
68: ts->dm = dm;
69: TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
70: if (mimex->version == 1) {
71: DM dm;
72: PetscDS prob;
73: PetscSection s;
74: Vec Xstar = NULL, G = NULL;
75: const PetscScalar *ax;
76: PetscScalar *axstar;
77: PetscInt Nf, f, pStart, pEnd, p;
79: TSGetDM(ts, &dm);
80: DMGetDS(dm, &prob);
81: DMGetLocalSection(dm, &s);
82: PetscDSGetNumFields(prob, &Nf);
83: PetscSectionGetChart(s, &pStart, &pEnd);
84: TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
85: VecCopy(X0, Xstar);
86: VecGetArrayRead(x, &ax);
87: VecGetArray(Xstar, &axstar);
88: for (f = 0; f < Nf; ++f) {
89: PetscBool implicit;
91: PetscDSGetImplicit(prob, f, &implicit);
92: if (!implicit) continue;
93: for (p = pStart; p < pEnd; ++p) {
94: PetscScalar *a, *axs;
95: PetscInt fdof, fcdof, d;
97: PetscSectionGetFieldDof(s, p, f, &fdof);
98: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
99: DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
100: DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
101: for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
102: }
103: }
104: VecRestoreArrayRead(x, &ax);
105: VecRestoreArray(Xstar, &axstar);
106: TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
107: VecAXPY(y, -1.0, G);
108: TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
109: }
110: ts->dm = dmsave;
111: TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
112: return 0;
113: }
115: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
116: {
117: TS_Mimex *mimex = (TS_Mimex *) ts->data;
118: DM dm, dmsave;
119: Vec Xdot;
120: PetscReal shift = 1./ts->time_step;
122: /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
123: SNESGetDM(snes, &dm);
124: TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);
126: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
127: dmsave = ts->dm;
128: ts->dm = dm;
129: TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
130: ts->dm = dmsave;
131: TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
132: return 0;
133: }
135: static PetscErrorCode TSStep_Mimex_Split(TS ts)
136: {
137: TS_Mimex *mimex = (TS_Mimex *) ts->data;
138: DM dm;
139: PetscDS prob;
140: PetscSection s;
141: Vec sol = ts->vec_sol, update = mimex->update;
142: const PetscScalar *aupdate;
143: PetscScalar *asol, dt = ts->time_step;
144: PetscInt Nf, f, pStart, pEnd, p;
146: TSGetDM(ts, &dm);
147: DMGetDS(dm, &prob);
148: DMGetLocalSection(dm, &s);
149: PetscDSGetNumFields(prob, &Nf);
150: PetscSectionGetChart(s, &pStart, &pEnd);
151: TSPreStage(ts, ts->ptime);
152: /* Compute implicit update */
153: mimex->stage_time = ts->ptime + ts->time_step;
154: VecCopy(sol, update);
155: SNESSolve(ts->snes, NULL, update);
156: VecGetArrayRead(update, &aupdate);
157: VecGetArray(sol, &asol);
158: for (f = 0; f < Nf; ++f) {
159: PetscBool implicit;
161: PetscDSGetImplicit(prob, f, &implicit);
162: if (!implicit) continue;
163: for (p = pStart; p < pEnd; ++p) {
164: PetscScalar *au, *as;
165: PetscInt fdof, fcdof, d;
167: PetscSectionGetFieldDof(s, p, f, &fdof);
168: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
169: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
170: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
171: for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
172: }
173: }
174: VecRestoreArrayRead(update, &aupdate);
175: VecRestoreArray(sol, &asol);
176: /* Compute explicit update */
177: TSComputeRHSFunction(ts, ts->ptime, sol, update);
178: VecGetArrayRead(update, &aupdate);
179: VecGetArray(sol, &asol);
180: for (f = 0; f < Nf; ++f) {
181: PetscBool implicit;
183: PetscDSGetImplicit(prob, f, &implicit);
184: if (implicit) continue;
185: for (p = pStart; p < pEnd; ++p) {
186: PetscScalar *au, *as;
187: PetscInt fdof, fcdof, d;
189: PetscSectionGetFieldDof(s, p, f, &fdof);
190: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
191: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
192: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
193: for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
194: }
195: }
196: VecRestoreArrayRead(update, &aupdate);
197: VecRestoreArray(sol, &asol);
198: TSPostStage(ts, ts->ptime, 0, &sol);
199: ts->ptime += ts->time_step;
200: return 0;
201: }
203: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
204: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
205: {
206: TS_Mimex *mimex = (TS_Mimex *) ts->data;
207: Vec sol = ts->vec_sol;
208: Vec update = mimex->update;
210: TSPreStage(ts, ts->ptime);
211: /* Compute implicit update */
212: mimex->stage_time = ts->ptime + ts->time_step;
213: ts->ptime += ts->time_step;
214: VecCopy(sol, update);
215: SNESSolve(ts->snes, NULL, update);
216: VecCopy(update, sol);
217: TSPostStage(ts, ts->ptime, 0, &sol);
218: return 0;
219: }
221: static PetscErrorCode TSStep_Mimex(TS ts)
222: {
223: TS_Mimex *mimex = (TS_Mimex*)ts->data;
225: switch(mimex->version) {
226: case 0:
227: TSStep_Mimex_Split(ts); break;
228: case 1:
229: TSStep_Mimex_Implicit(ts); break;
230: default:
231: SETERRQ(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
232: }
233: return 0;
234: }
236: /*------------------------------------------------------------*/
238: static PetscErrorCode TSSetUp_Mimex(TS ts)
239: {
240: TS_Mimex *mimex = (TS_Mimex*)ts->data;
242: VecDuplicate(ts->vec_sol, &mimex->update);
243: VecDuplicate(ts->vec_sol, &mimex->Xdot);
244: return 0;
245: }
247: static PetscErrorCode TSReset_Mimex(TS ts)
248: {
249: TS_Mimex *mimex = (TS_Mimex*)ts->data;
251: VecDestroy(&mimex->update);
252: VecDestroy(&mimex->Xdot);
253: return 0;
254: }
256: static PetscErrorCode TSDestroy_Mimex(TS ts)
257: {
258: TSReset_Mimex(ts);
259: PetscFree(ts->data);
260: return 0;
261: }
262: /*------------------------------------------------------------*/
264: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
265: {
266: TS_Mimex *mimex = (TS_Mimex *) ts->data;
268: PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
269: {
270: PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
271: }
272: PetscOptionsTail();
273: return 0;
274: }
276: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
277: {
278: TS_Mimex *mimex = (TS_Mimex *) ts->data;
279: PetscBool iascii;
281: PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
282: if (iascii) {
283: PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);
284: }
285: return 0;
286: }
288: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
289: {
290: PetscReal alpha = (ts->ptime - t)/ts->time_step;
292: VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
293: return 0;
294: }
296: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
297: {
298: *yr = 1.0 + xr;
299: *yi = xi;
300: return 0;
301: }
302: /* ------------------------------------------------------------ */
304: /*MC
305: TSMIMEX - ODE solver using the explicit forward Mimex method
307: Level: beginner
309: .seealso: TSCreate(), TS, TSSetType(), TSBEULER
311: M*/
312: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
313: {
314: TS_Mimex *mimex;
316: ts->ops->setup = TSSetUp_Mimex;
317: ts->ops->step = TSStep_Mimex;
318: ts->ops->reset = TSReset_Mimex;
319: ts->ops->destroy = TSDestroy_Mimex;
320: ts->ops->setfromoptions = TSSetFromOptions_Mimex;
321: ts->ops->view = TSView_Mimex;
322: ts->ops->interpolate = TSInterpolate_Mimex;
323: ts->ops->linearstability = TSComputeLinearStability_Mimex;
324: ts->ops->snesfunction = SNESTSFormFunction_Mimex;
325: ts->ops->snesjacobian = SNESTSFormJacobian_Mimex;
326: ts->default_adapt_type = TSADAPTNONE;
328: PetscNewLog(ts,&mimex);
329: ts->data = (void*)mimex;
331: mimex->version = 1;
332: return 0;
333: }