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