Actual source code: mimex.c
petsc-3.11.4 2019-09-28
1: /*
2: Code for Timestepping with my makeshift IMEX.
3: */
4: #include <petsc/private/tsimpl.h>
5: #include <petscds.h>
6: #include <petscdmplex.h>
8: typedef struct {
9: Vec Xdot, update;
10: PetscReal stage_time;
11: PetscInt version;
12: } TS_Mimex;
14: static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
15: {
16: TS_Mimex *mimex = (TS_Mimex *) ts->data;
20: if (X0) {
21: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);}
22: else {*X0 = ts->vec_sol;}
23: }
24: if (Xdot) {
25: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
26: else {*Xdot = mimex->Xdot;}
27: }
28: return(0);
29: }
31: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
32: {
36: if (X0) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);}
37: if (Xdot) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
38: return(0);
39: }
41: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
42: {
46: DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
47: DMGetNamedGlobalVector(dm, "TSMimex_G", G);
48: return(0);
49: }
51: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
52: {
56: DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
57: DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
58: return(0);
59: }
61: /*
62: This defines the nonlinear equation that is to be solved with SNES
63: G(U) = F[t0+dt, U, (U-U0)*shift] = 0
64: */
65: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
66: {
67: TS_Mimex *mimex = (TS_Mimex *) ts->data;
68: DM dm, dmsave;
69: Vec X0, Xdot;
70: PetscReal shift = 1./ts->time_step;
74: SNESGetDM(snes, &dm);
75: TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
76: VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);
78: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
79: dmsave = ts->dm;
80: ts->dm = dm;
81: TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
82: if (mimex->version == 1) {
83: DM dm;
84: PetscDS prob;
85: PetscSection s;
86: Vec Xstar = NULL, G = NULL;
87: const PetscScalar *ax;
88: PetscScalar *axstar;
89: PetscInt Nf, f, pStart, pEnd, p;
91: TSGetDM(ts, &dm);
92: DMGetDS(dm, &prob);
93: DMGetSection(dm, &s);
94: PetscDSGetNumFields(prob, &Nf);
95: PetscSectionGetChart(s, &pStart, &pEnd);
96: TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
97: VecCopy(X0, Xstar);
98: VecGetArrayRead(x, &ax);
99: VecGetArray(Xstar, &axstar);
100: for (f = 0; f < Nf; ++f) {
101: PetscBool implicit;
103: PetscDSGetImplicit(prob, f, &implicit);
104: if (!implicit) continue;
105: for (p = pStart; p < pEnd; ++p) {
106: PetscScalar *a, *axs;
107: PetscInt fdof, fcdof, d;
109: PetscSectionGetFieldDof(s, p, f, &fdof);
110: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
111: DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
112: DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
113: for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
114: }
115: }
116: VecRestoreArrayRead(x, &ax);
117: VecRestoreArray(Xstar, &axstar);
118: TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
119: VecAXPY(y, -1.0, G);
120: TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
121: }
122: ts->dm = dmsave;
123: TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
124: return(0);
125: }
127: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
128: {
129: TS_Mimex *mimex = (TS_Mimex *) ts->data;
130: DM dm, dmsave;
131: Vec Xdot;
132: PetscReal shift = 1./ts->time_step;
136: /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
137: SNESGetDM(snes, &dm);
138: TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);
140: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
141: dmsave = ts->dm;
142: ts->dm = dm;
143: TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
144: ts->dm = dmsave;
145: TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
146: return(0);
147: }
149: static PetscErrorCode TSStep_Mimex_Split(TS ts)
150: {
151: TS_Mimex *mimex = (TS_Mimex *) ts->data;
152: DM dm;
153: PetscDS prob;
154: PetscSection s;
155: Vec sol = ts->vec_sol, update = mimex->update;
156: const PetscScalar *aupdate;
157: PetscScalar *asol, dt = ts->time_step;
158: PetscInt Nf, f, pStart, pEnd, p;
159: PetscErrorCode ierr;
162: TSGetDM(ts, &dm);
163: DMGetDS(dm, &prob);
164: DMGetSection(dm, &s);
165: PetscDSGetNumFields(prob, &Nf);
166: PetscSectionGetChart(s, &pStart, &pEnd);
167: TSPreStage(ts, ts->ptime);
168: /* Compute implicit update */
169: mimex->stage_time = ts->ptime + ts->time_step;
170: VecCopy(sol, update);
171: SNESSolve(ts->snes, NULL, update);
172: VecGetArrayRead(update, &aupdate);
173: VecGetArray(sol, &asol);
174: for (f = 0; f < Nf; ++f) {
175: PetscBool implicit;
177: PetscDSGetImplicit(prob, f, &implicit);
178: if (!implicit) continue;
179: for (p = pStart; p < pEnd; ++p) {
180: PetscScalar *au, *as;
181: PetscInt fdof, fcdof, d;
183: PetscSectionGetFieldDof(s, p, f, &fdof);
184: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
185: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
186: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
187: for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
188: }
189: }
190: VecRestoreArrayRead(update, &aupdate);
191: VecRestoreArray(sol, &asol);
192: /* Compute explicit update */
193: TSComputeRHSFunction(ts, ts->ptime, sol, update);
194: VecGetArrayRead(update, &aupdate);
195: VecGetArray(sol, &asol);
196: for (f = 0; f < Nf; ++f) {
197: PetscBool implicit;
199: PetscDSGetImplicit(prob, f, &implicit);
200: if (implicit) continue;
201: for (p = pStart; p < pEnd; ++p) {
202: PetscScalar *au, *as;
203: PetscInt fdof, fcdof, d;
205: PetscSectionGetFieldDof(s, p, f, &fdof);
206: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
207: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
208: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
209: for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
210: }
211: }
212: VecRestoreArrayRead(update, &aupdate);
213: VecRestoreArray(sol, &asol);
214: TSPostStage(ts, ts->ptime, 0, &sol);
215: ts->ptime += ts->time_step;
216: return(0);
217: }
220: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
221: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
222: {
223: TS_Mimex *mimex = (TS_Mimex *) ts->data;
224: Vec sol = ts->vec_sol;
225: Vec update = mimex->update;
229: TSPreStage(ts, ts->ptime);
230: /* Compute implicit update */
231: mimex->stage_time = ts->ptime + ts->time_step;
232: ts->ptime += ts->time_step;
233: VecCopy(sol, update);
234: SNESSolve(ts->snes, NULL, update);
235: VecCopy(update, sol);
236: TSPostStage(ts, ts->ptime, 0, &sol);
237: return(0);
238: }
240: static PetscErrorCode TSStep_Mimex(TS ts)
241: {
242: TS_Mimex *mimex = (TS_Mimex*)ts->data;
243: PetscErrorCode ierr;
246: switch(mimex->version) {
247: case 0:
248: TSStep_Mimex_Split(ts); break;
249: case 1:
250: TSStep_Mimex_Implicit(ts); break;
251: default:
252: SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
253: }
254: return(0);
255: }
257: /*------------------------------------------------------------*/
259: static PetscErrorCode TSSetUp_Mimex(TS ts)
260: {
261: TS_Mimex *mimex = (TS_Mimex*)ts->data;
265: VecDuplicate(ts->vec_sol, &mimex->update);
266: VecDuplicate(ts->vec_sol, &mimex->Xdot);
267: return(0);
268: }
270: static PetscErrorCode TSReset_Mimex(TS ts)
271: {
272: TS_Mimex *mimex = (TS_Mimex*)ts->data;
276: VecDestroy(&mimex->update);
277: VecDestroy(&mimex->Xdot);
278: return(0);
279: }
281: static PetscErrorCode TSDestroy_Mimex(TS ts)
282: {
286: TSReset_Mimex(ts);
287: PetscFree(ts->data);
288: return(0);
289: }
290: /*------------------------------------------------------------*/
292: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
293: {
294: TS_Mimex *mimex = (TS_Mimex *) ts->data;
298: PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
299: {
300: PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
301: }
302: PetscOptionsTail();
303: return(0);
304: }
306: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
307: {
308: TS_Mimex *mimex = (TS_Mimex *) ts->data;
309: PetscBool iascii;
313: PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
314: if (iascii) {
315: PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);
316: }
317: return(0);
318: }
320: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
321: {
322: PetscReal alpha = (ts->ptime - t)/ts->time_step;
326: VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
327: return(0);
328: }
330: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
331: {
333: *yr = 1.0 + xr;
334: *yi = xi;
335: return(0);
336: }
337: /* ------------------------------------------------------------ */
339: /*MC
340: TSMIMEX - ODE solver using the explicit forward Mimex method
342: Level: beginner
344: .seealso: TSCreate(), TS, TSSetType(), TSBEULER
346: M*/
347: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
348: {
349: TS_Mimex *mimex;
353: ts->ops->setup = TSSetUp_Mimex;
354: ts->ops->step = TSStep_Mimex;
355: ts->ops->reset = TSReset_Mimex;
356: ts->ops->destroy = TSDestroy_Mimex;
357: ts->ops->setfromoptions = TSSetFromOptions_Mimex;
358: ts->ops->view = TSView_Mimex;
359: ts->ops->interpolate = TSInterpolate_Mimex;
360: ts->ops->linearstability = TSComputeLinearStability_Mimex;
361: ts->ops->snesfunction = SNESTSFormFunction_Mimex;
362: ts->ops->snesjacobian = SNESTSFormJacobian_Mimex;
363: ts->default_adapt_type = TSADAPTNONE;
365: PetscNewLog(ts,&mimex);
366: ts->data = (void*)mimex;
368: mimex->version = 1;
369: return(0);
370: }