Actual source code: mimex.c
petsc-3.6.1 2015-08-06
1: /*
2: Code for Timestepping with my makeshift IMEX.
3: */
4: #include <petsc/private/tsimpl.h> /*I "petscts.h" I*/
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;
16: static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
17: {
18: TS_Mimex *mimex = (TS_Mimex *) ts->data;
22: if (X0) {
23: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);}
24: else {*X0 = ts->vec_sol;}
25: }
26: if (Xdot) {
27: if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
28: else {*Xdot = mimex->Xdot;}
29: }
30: return(0);
31: }
35: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
36: {
40: if (X0) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);}
41: if (Xdot) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
42: return(0);
43: }
47: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
48: {
52: DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
53: DMGetNamedGlobalVector(dm, "TSMimex_G", G);
54: return(0);
55: }
59: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
60: {
64: DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
65: DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
66: return(0);
67: }
69: /*
70: This defines the nonlinear equation that is to be solved with SNES
71: G(U) = F[t0+dt, U, (U-U0)*shift] = 0
72: */
75: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
76: {
77: TS_Mimex *mimex = (TS_Mimex *) ts->data;
78: DM dm, dmsave;
79: Vec X0, Xdot;
80: PetscReal shift = 1./ts->time_step;
84: SNESGetDM(snes, &dm);
85: TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
86: VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);
88: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
89: dmsave = ts->dm;
90: ts->dm = dm;
91: TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
92: if (mimex->version == 1) {
93: DM dm;
94: PetscDS prob;
95: PetscSection s;
96: Vec Xstar = NULL, G = NULL;
97: const PetscScalar *ax;
98: PetscScalar *axstar;
99: PetscInt Nf, f, pStart, pEnd, p;
101: TSGetDM(ts, &dm);
102: DMGetDS(dm, &prob);
103: DMGetDefaultSection(dm, &s);
104: PetscDSGetNumFields(prob, &Nf);
105: PetscSectionGetChart(s, &pStart, &pEnd);
106: TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
107: VecCopy(X0, Xstar);
108: VecGetArrayRead(x, &ax);
109: VecGetArray(Xstar, &axstar);
110: for (f = 0; f < Nf; ++f) {
111: PetscBool implicit;
113: PetscDSGetImplicit(prob, f, &implicit);
114: if (!implicit) continue;
115: for (p = pStart; p < pEnd; ++p) {
116: PetscScalar *a, *axs;
117: PetscInt fdof, fcdof, d;
119: PetscSectionGetFieldDof(s, p, f, &fdof);
120: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
121: DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
122: DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
123: for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
124: }
125: }
126: VecRestoreArrayRead(x, &ax);
127: VecRestoreArray(Xstar, &axstar);
128: TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
129: VecAXPY(y, -1.0, G);
130: TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
131: }
132: ts->dm = dmsave;
133: TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
134: return(0);
135: }
139: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
140: {
141: TS_Mimex *mimex = (TS_Mimex *) ts->data;
142: DM dm, dmsave;
143: Vec Xdot;
144: PetscReal shift = 1./ts->time_step;
148: /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
149: SNESGetDM(snes, &dm);
150: TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);
152: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
153: dmsave = ts->dm;
154: ts->dm = dm;
155: TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
156: ts->dm = dmsave;
157: TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
158: return(0);
159: }
163: static PetscErrorCode TSStep_Mimex_Split(TS ts)
164: {
165: TS_Mimex *mimex = (TS_Mimex *) ts->data;
166: DM dm;
167: PetscDS prob;
168: PetscSection s;
169: Vec sol = ts->vec_sol, update = mimex->update;
170: const PetscScalar *aupdate;
171: PetscScalar *asol, dt = ts->time_step;
172: PetscInt Nf, f, pStart, pEnd, p;
173: PetscErrorCode ierr;
176: TSGetDM(ts, &dm);
177: DMGetDS(dm, &prob);
178: DMGetDefaultSection(dm, &s);
179: PetscDSGetNumFields(prob, &Nf);
180: PetscSectionGetChart(s, &pStart, &pEnd);
181: TSPreStep(ts);
182: TSPreStage(ts, ts->ptime);
183: /* Compute implicit update */
184: mimex->stage_time = ts->ptime + ts->time_step;
185: VecCopy(sol, update);
186: SNESSolve(ts->snes, NULL, update);
187: VecGetArrayRead(update, &aupdate);
188: VecGetArray(sol, &asol);
189: for (f = 0; f < Nf; ++f) {
190: PetscBool implicit;
192: 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: PetscSectionGetFieldDof(s, p, f, &fdof);
199: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
200: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
201: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
202: for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
203: }
204: }
205: VecRestoreArrayRead(update, &aupdate);
206: VecRestoreArray(sol, &asol);
207: /* Compute explicit update */
208: TSComputeRHSFunction(ts, ts->ptime, sol, update);
209: VecGetArrayRead(update, &aupdate);
210: VecGetArray(sol, &asol);
211: for (f = 0; f < Nf; ++f) {
212: PetscBool implicit;
214: PetscDSGetImplicit(prob, f, &implicit);
215: if (implicit) continue;
216: for (p = pStart; p < pEnd; ++p) {
217: PetscScalar *au, *as;
218: PetscInt fdof, fcdof, d;
220: PetscSectionGetFieldDof(s, p, f, &fdof);
221: PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
222: DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
223: DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
224: for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
225: }
226: }
227: VecRestoreArrayRead(update, &aupdate);
228: VecRestoreArray(sol, &asol);
229: TSPostStage(ts, ts->ptime, 0, &sol);
230: ts->ptime += ts->time_step;
231: ts->steps++;
232: return(0);
233: }
238: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
239: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
240: {
241: TS_Mimex *mimex = (TS_Mimex *) ts->data;
242: Vec sol = ts->vec_sol;
243: Vec update = mimex->update;
247: TSPreStep(ts);
248: TSPreStage(ts, ts->ptime);
249: /* Compute implicit update */
250: mimex->stage_time = ts->ptime + ts->time_step;
251: ts->ptime += ts->time_step;
252: VecCopy(sol, update);
253: SNESSolve(ts->snes, NULL, update);
254: VecCopy(update, sol);
255: TSPostStage(ts, ts->ptime, 0, &sol);
256: ts->steps++;
257: return(0);
258: }
262: static PetscErrorCode TSStep_Mimex(TS ts)
263: {
264: TS_Mimex *mimex = (TS_Mimex*)ts->data;
265: PetscErrorCode ierr;
268: switch(mimex->version) {
269: case 0:
270: TSStep_Mimex_Split(ts); break;
271: case 1:
272: TSStep_Mimex_Implicit(ts); break;
273: default:
274: SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
275: }
276: return(0);
277: }
279: /*------------------------------------------------------------*/
283: static PetscErrorCode TSSetUp_Mimex(TS ts)
284: {
285: TS_Mimex *mimex = (TS_Mimex*)ts->data;
289: VecDuplicate(ts->vec_sol, &mimex->update);
290: VecDuplicate(ts->vec_sol, &mimex->Xdot);
291: return(0);
292: }
296: static PetscErrorCode TSReset_Mimex(TS ts)
297: {
298: TS_Mimex *mimex = (TS_Mimex*)ts->data;
302: VecDestroy(&mimex->update);
303: VecDestroy(&mimex->Xdot);
304: return(0);
305: }
309: static PetscErrorCode TSDestroy_Mimex(TS ts)
310: {
314: TSReset_Mimex(ts);
315: PetscFree(ts->data);
316: return(0);
317: }
318: /*------------------------------------------------------------*/
322: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject, TS ts)
323: {
324: TS_Mimex *mimex = (TS_Mimex *) ts->data;
328: PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
329: {
330: PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
331: }
332: PetscOptionsTail();
333: return(0);
334: }
338: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
339: {
340: TS_Mimex *mimex = (TS_Mimex *) ts->data;
341: PetscBool iascii;
345: PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
346: if (iascii) {
347: PetscViewerASCIIPrintf(viewer, " Version = %D\n", mimex->version);
348: }
349: if (ts->snes) {SNESView(ts->snes, viewer);}
350: return(0);
351: }
355: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
356: {
357: PetscReal alpha = (ts->ptime - t)/ts->time_step;
361: VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
362: return(0);
363: }
367: PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
368: {
370: *yr = 1.0 + xr;
371: *yi = xi;
372: return(0);
373: }
374: /* ------------------------------------------------------------ */
376: /*MC
377: TSMIMEX - ODE solver using the explicit forward Mimex method
379: Level: beginner
381: .seealso: TSCreate(), TS, TSSetType(), TSBEULER
383: M*/
386: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
387: {
388: TS_Mimex *mimex;
392: ts->ops->setup = TSSetUp_Mimex;
393: ts->ops->step = TSStep_Mimex;
394: ts->ops->reset = TSReset_Mimex;
395: ts->ops->destroy = TSDestroy_Mimex;
396: ts->ops->setfromoptions = TSSetFromOptions_Mimex;
397: ts->ops->view = TSView_Mimex;
398: ts->ops->interpolate = TSInterpolate_Mimex;
399: ts->ops->linearstability = TSComputeLinearStability_Mimex;
400: ts->ops->snesfunction = SNESTSFormFunction_Mimex;
401: ts->ops->snesjacobian = SNESTSFormJacobian_Mimex;
403: PetscNewLog(ts,&mimex);
404: ts->data = (void*)mimex;
406: mimex->version = 1;
407: return(0);
408: }