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;
21: if (X0) {
22: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);}
23: else {*X0 = ts->vec_sol;}
24: }
25: if (Xdot) {
26: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
27: else {*Xdot = mimex->Xdot;}
28: }
29: return(0);
30: }
32: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
33: {
37: if (X0) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);}
38: if (Xdot) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
39: return(0);
40: }
42: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
43: {
47: DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
48: DMGetNamedGlobalVector(dm, "TSMimex_G", G);
49: return(0);
50: }
52: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
53: {
57: DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
58: DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
59: return(0);
60: }
62: /*
63: This defines the nonlinear equation that is to be solved with SNES
64: G(U) = F[t0+dt, U, (U-U0)*shift] = 0
65: */
66: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
67: {
68: TS_Mimex *mimex = (TS_Mimex *) ts->data;
69: DM dm, dmsave;
70: Vec X0, Xdot;
71: PetscReal shift = 1./ts->time_step;
75: SNESGetDM(snes, &dm);
76: TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
77: VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);
79: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
80: dmsave = ts->dm;
81: ts->dm = dm;
82: TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
83: if (mimex->version == 1) {
84: DM dm;
85: PetscDS prob;
86: PetscSection s;
87: Vec Xstar = NULL, G = NULL;
88: const PetscScalar *ax;
89: PetscScalar *axstar;
90: PetscInt Nf, f, pStart, pEnd, p;
92: TSGetDM(ts, &dm);
93: DMGetDS(dm, &prob);
94: DMGetLocalSection(dm, &s);
95: PetscDSGetNumFields(prob, &Nf);
96: PetscSectionGetChart(s, &pStart, &pEnd);
97: TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
98: VecCopy(X0, Xstar);
99: VecGetArrayRead(x, &ax);
100: VecGetArray(Xstar, &axstar);
101: for (f = 0; f < Nf; ++f) {
102: PetscBool implicit;
104: PetscDSGetImplicit(prob, f, &implicit);
105: if (!implicit) continue;
106: for (p = pStart; p < pEnd; ++p) {
107: PetscScalar *a, *axs;
108: PetscInt fdof, fcdof, d;
110: PetscSectionGetFieldDof(s, p, f, &fdof);
111: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
112: DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
113: DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
114: for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
115: }
116: }
117: VecRestoreArrayRead(x, &ax);
118: VecRestoreArray(Xstar, &axstar);
119: TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
120: VecAXPY(y, -1.0, G);
121: TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
122: }
123: ts->dm = dmsave;
124: TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
125: return(0);
126: }
128: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
129: {
130: TS_Mimex *mimex = (TS_Mimex *) ts->data;
131: DM dm, dmsave;
132: Vec Xdot;
133: PetscReal shift = 1./ts->time_step;
137: /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
138: SNESGetDM(snes, &dm);
139: TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);
141: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
142: dmsave = ts->dm;
143: ts->dm = dm;
144: TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
145: ts->dm = dmsave;
146: TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
147: return(0);
148: }
150: static PetscErrorCode TSStep_Mimex_Split(TS ts)
151: {
152: TS_Mimex *mimex = (TS_Mimex *) ts->data;
153: DM dm;
154: PetscDS prob;
155: PetscSection s;
156: Vec sol = ts->vec_sol, update = mimex->update;
157: const PetscScalar *aupdate;
158: PetscScalar *asol, dt = ts->time_step;
159: PetscInt Nf, f, pStart, pEnd, p;
160: PetscErrorCode ierr;
163: TSGetDM(ts, &dm);
164: DMGetDS(dm, &prob);
165: DMGetLocalSection(dm, &s);
166: PetscDSGetNumFields(prob, &Nf);
167: PetscSectionGetChart(s, &pStart, &pEnd);
168: TSPreStage(ts, ts->ptime);
169: /* Compute implicit update */
170: mimex->stage_time = ts->ptime + ts->time_step;
171: VecCopy(sol, update);
172: SNESSolve(ts->snes, NULL, update);
173: VecGetArrayRead(update, &aupdate);
174: VecGetArray(sol, &asol);
175: for (f = 0; f < Nf; ++f) {
176: PetscBool implicit;
178: PetscDSGetImplicit(prob, f, &implicit);
179: if (!implicit) continue;
180: for (p = pStart; p < pEnd; ++p) {
181: PetscScalar *au, *as;
182: PetscInt fdof, fcdof, d;
184: PetscSectionGetFieldDof(s, p, f, &fdof);
185: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
186: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
187: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
188: for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
189: }
190: }
191: VecRestoreArrayRead(update, &aupdate);
192: VecRestoreArray(sol, &asol);
193: /* Compute explicit update */
194: TSComputeRHSFunction(ts, ts->ptime, sol, update);
195: VecGetArrayRead(update, &aupdate);
196: VecGetArray(sol, &asol);
197: for (f = 0; f < Nf; ++f) {
198: PetscBool implicit;
200: PetscDSGetImplicit(prob, f, &implicit);
201: if (implicit) continue;
202: for (p = pStart; p < pEnd; ++p) {
203: PetscScalar *au, *as;
204: PetscInt fdof, fcdof, d;
206: PetscSectionGetFieldDof(s, p, f, &fdof);
207: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
208: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
209: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
210: for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
211: }
212: }
213: VecRestoreArrayRead(update, &aupdate);
214: VecRestoreArray(sol, &asol);
215: TSPostStage(ts, ts->ptime, 0, &sol);
216: ts->ptime += ts->time_step;
217: return(0);
218: }
221: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
222: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
223: {
224: TS_Mimex *mimex = (TS_Mimex *) ts->data;
225: Vec sol = ts->vec_sol;
226: Vec update = mimex->update;
230: TSPreStage(ts, ts->ptime);
231: /* Compute implicit update */
232: mimex->stage_time = ts->ptime + ts->time_step;
233: ts->ptime += ts->time_step;
234: VecCopy(sol, update);
235: SNESSolve(ts->snes, NULL, update);
236: VecCopy(update, sol);
237: TSPostStage(ts, ts->ptime, 0, &sol);
238: return(0);
239: }
241: static PetscErrorCode TSStep_Mimex(TS ts)
242: {
243: TS_Mimex *mimex = (TS_Mimex*)ts->data;
244: PetscErrorCode ierr;
247: switch(mimex->version) {
248: case 0:
249: TSStep_Mimex_Split(ts); break;
250: case 1:
251: TSStep_Mimex_Implicit(ts); break;
252: default:
253: SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
254: }
255: return(0);
256: }
258: /*------------------------------------------------------------*/
260: static PetscErrorCode TSSetUp_Mimex(TS ts)
261: {
262: TS_Mimex *mimex = (TS_Mimex*)ts->data;
266: VecDuplicate(ts->vec_sol, &mimex->update);
267: VecDuplicate(ts->vec_sol, &mimex->Xdot);
268: return(0);
269: }
271: static PetscErrorCode TSReset_Mimex(TS ts)
272: {
273: TS_Mimex *mimex = (TS_Mimex*)ts->data;
277: VecDestroy(&mimex->update);
278: VecDestroy(&mimex->Xdot);
279: return(0);
280: }
282: static PetscErrorCode TSDestroy_Mimex(TS ts)
283: {
287: TSReset_Mimex(ts);
288: PetscFree(ts->data);
289: return(0);
290: }
291: /*------------------------------------------------------------*/
293: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
294: {
295: TS_Mimex *mimex = (TS_Mimex *) ts->data;
299: PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
300: {
301: PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
302: }
303: PetscOptionsTail();
304: return(0);
305: }
307: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
308: {
309: TS_Mimex *mimex = (TS_Mimex *) ts->data;
310: PetscBool iascii;
314: PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
315: if (iascii) {
316: PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);
317: }
318: return(0);
319: }
321: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
322: {
323: PetscReal alpha = (ts->ptime - t)/ts->time_step;
327: VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
328: return(0);
329: }
331: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
332: {
334: *yr = 1.0 + xr;
335: *yi = xi;
336: return(0);
337: }
338: /* ------------------------------------------------------------ */
340: /*MC
341: TSMIMEX - ODE solver using the explicit forward Mimex method
343: Level: beginner
345: .seealso: TSCreate(), TS, TSSetType(), TSBEULER
347: M*/
348: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
349: {
350: TS_Mimex *mimex;
354: ts->ops->setup = TSSetUp_Mimex;
355: ts->ops->step = TSStep_Mimex;
356: ts->ops->reset = TSReset_Mimex;
357: ts->ops->destroy = TSDestroy_Mimex;
358: ts->ops->setfromoptions = TSSetFromOptions_Mimex;
359: ts->ops->view = TSView_Mimex;
360: ts->ops->interpolate = TSInterpolate_Mimex;
361: ts->ops->linearstability = TSComputeLinearStability_Mimex;
362: ts->ops->snesfunction = SNESTSFormFunction_Mimex;
363: ts->ops->snesjacobian = SNESTSFormJacobian_Mimex;
364: ts->default_adapt_type = TSADAPTNONE;
366: PetscNewLog(ts,&mimex);
367: ts->data = (void*)mimex;
369: mimex->version = 1;
370: return(0);
371: }