Actual source code: fsarkimex.c
1: #include <petsc/private/tsimpl.h>
2: #include <petscdm.h>
3: #include <../src/ts/impls/arkimex/arkimex.h>
4: #include <../src/ts/impls/arkimex/fsarkimex.h>
6: static PetscErrorCode TSARKIMEXSetSplits(TS ts)
7: {
8: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
9: DM dm, subdm, newdm;
11: PetscFunctionBegin;
12: PetscCall(TSRHSSplitGetSubTS(ts, "slow", &ark->subts_slow));
13: PetscCall(TSRHSSplitGetSubTS(ts, "fast", &ark->subts_fast));
14: /* Only copy the DM */
15: PetscCall(TSGetDM(ts, &dm));
16: if (ark->subts_slow) {
17: PetscCall(DMClone(dm, &newdm));
18: PetscCall(TSGetDM(ark->subts_slow, &subdm));
19: PetscCall(DMCopyDMTS(subdm, newdm));
20: PetscCall(TSSetDM(ark->subts_slow, newdm));
21: PetscCall(DMDestroy(&newdm));
22: }
23: if (ark->subts_fast) {
24: PetscCall(DMClone(dm, &newdm));
25: PetscCall(TSGetDM(ark->subts_fast, &subdm));
26: PetscCall(DMCopyDMTS(subdm, newdm));
27: PetscCall(TSSetDM(ark->subts_fast, newdm));
28: PetscCall(DMDestroy(&newdm));
29: }
30: PetscFunctionReturn(PETSC_SUCCESS);
31: }
33: static PetscErrorCode SNESTSFormFunction_ARKIMEX_FastSlowSplit(SNES snes, Vec X, Vec F, TS ts)
34: {
35: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
36: DM dm, dmsave;
37: Vec Z = ark->Z, Ydot = ark->Ydot, Y = ark->Y_snes;
39: PetscFunctionBegin;
40: PetscCall(SNESGetDM(snes, &dm));
41: dmsave = ts->dm;
42: ts->dm = dm; // Use the SNES DM to compute IFunction
44: PetscReal shift = ark->scoeff / ts->time_step;
45: PetscCall(VecAXPBYPCZ(Ydot, -shift, shift, 0, Z, X)); /* Ydot = shift*(X-Z) */
46: if (ark->is_slow) PetscCall(VecISCopy(Y, ark->is_fast, SCATTER_FORWARD, X));
47: else Y = Z;
48: PetscCall(TSComputeIFunction(ark->subts_fast, ark->stage_time, Y, Ydot, F, ark->imex));
50: ts->dm = dmsave;
51: PetscFunctionReturn(PETSC_SUCCESS);
52: }
54: static PetscErrorCode SNESTSFormJacobian_ARKIMEX_FastSlowSplit(SNES snes, Vec X, Mat A, Mat B, TS ts)
55: {
56: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
57: DM dm, dmsave;
58: Vec Z = ark->Z, Ydot = ark->Ydot, Y = ark->Y_snes;
59: PetscReal shift;
61: PetscFunctionBegin;
62: PetscCall(SNESGetDM(snes, &dm));
63: dmsave = ts->dm;
64: ts->dm = dm;
66: shift = ark->scoeff / ts->time_step;
67: if (ark->is_slow) PetscCall(VecISCopy(Y, ark->is_fast, SCATTER_FORWARD, X));
68: else Y = Z;
69: PetscCall(TSComputeIJacobian(ark->subts_fast, ark->stage_time, Y, Ydot, shift, A, B, ark->imex));
71: ts->dm = dmsave;
72: PetscFunctionReturn(PETSC_SUCCESS);
73: }
75: static PetscErrorCode TSExtrapolate_ARKIMEX_FastSlowSplit(TS ts, PetscReal c, Vec X)
76: {
77: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
78: ARKTableau tab = ark->tableau;
79: PetscInt s = tab->s, pinterp = tab->pinterp, i, j;
80: PetscReal h, h_prev, t, tt;
81: PetscScalar *bt = ark->work, *b = ark->work + s;
82: const PetscReal *Bt = tab->binterpt, *B = tab->binterp;
83: PetscBool fasthasE;
85: PetscFunctionBegin;
86: PetscCheck(Bt && B, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "TSARKIMEX %s does not have an interpolation formula", ark->tableau->name);
87: h = ts->time_step;
88: h_prev = ts->ptime - ts->ptime_prev;
89: t = 1 + h / h_prev * c;
90: for (i = 0; i < s; i++) bt[i] = b[i] = 0;
91: for (j = 0, tt = t; j < pinterp; j++, tt *= t) {
92: for (i = 0; i < s; i++) {
93: bt[i] += h * Bt[i * pinterp + j] * tt;
94: b[i] += h * B[i * pinterp + j] * tt;
95: }
96: }
97: PetscCheck(ark->Y_prev, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Stages from previous step have not been stored");
98: PetscCall(VecCopy(ark->Y_prev[0], X));
99: PetscCall(VecMAXPY(X, s, bt, ark->YdotI_prev));
100: PetscCall(TSHasRHSFunction(ark->subts_fast, &fasthasE));
101: if (fasthasE) PetscCall(VecMAXPY(X, s, b, ark->YdotRHS_prev));
102: PetscFunctionReturn(PETSC_SUCCESS);
103: }
105: /*
106: The step completion formula is
108: x1 = x0 - h bt^T YdotI + h b^T YdotRHS
110: This function can be called before or after ts->vec_sol has been updated.
111: Suppose we have a completion formula (bt,b) and an embedded formula (bet,be) of different order.
112: We can write
114: x1e = x0 - h bet^T YdotI + h be^T YdotRHS
115: = x1 + h bt^T YdotI - h b^T YdotRHS - h bet^T YdotI + h be^T YdotRHS
116: = x1 - h (bet - bt)^T YdotI + h (be - b)^T YdotRHS
118: so we can evaluate the method with different order even after the step has been optimistically completed.
119: */
120: static PetscErrorCode TSEvaluateStep_ARKIMEX_FastSlowSplit(TS ts, PetscInt order, Vec X, PetscBool *done)
121: {
122: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
123: ARKTableau tab = ark->tableau;
124: Vec Xfast, Xslow;
125: PetscScalar *w = ark->work;
126: PetscReal h;
127: PetscInt s = tab->s, j;
128: PetscBool fasthasE;
130: PetscFunctionBegin;
131: switch (ark->status) {
132: case TS_STEP_INCOMPLETE:
133: case TS_STEP_PENDING:
134: h = ts->time_step;
135: break;
136: case TS_STEP_COMPLETE:
137: h = ts->ptime - ts->ptime_prev;
138: break;
139: default:
140: SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Invalid TSStepStatus");
141: }
142: if (ark->is_fast) PetscCall(TSHasRHSFunction(ark->subts_fast, &fasthasE));
143: if (order == tab->order) {
144: if (ark->status == TS_STEP_INCOMPLETE) {
145: PetscCall(VecCopy(ts->vec_sol, X));
146: for (j = 0; j < s; j++) w[j] = h * tab->b[j];
147: if (ark->is_slow) {
148: PetscCall(VecGetSubVector(X, ark->is_slow, &Xslow));
149: PetscCall(VecMAXPY(Xslow, s, w, ark->YdotRHS_slow));
150: PetscCall(VecRestoreSubVector(X, ark->is_slow, &Xslow));
151: }
152: if (ark->is_fast) {
153: PetscCall(VecGetSubVector(X, ark->is_fast, &Xfast));
154: if (fasthasE) PetscCall(VecMAXPY(Xfast, s, w, ark->YdotRHS_fast));
155: for (j = 0; j < s; j++) w[j] = h * tab->bt[j];
156: PetscCall(VecMAXPY(Xfast, s, w, ark->YdotI_fast));
157: PetscCall(VecRestoreSubVector(X, ark->is_fast, &Xfast));
158: }
159: } else PetscCall(VecCopy(ts->vec_sol, X));
160: if (done) *done = PETSC_TRUE;
161: PetscFunctionReturn(PETSC_SUCCESS);
162: } else if (order == tab->order - 1) {
163: if (!tab->bembedt) goto unavailable;
164: if (ark->status == TS_STEP_INCOMPLETE) { /* Complete with the embedded method (bet,be) */
165: PetscCall(VecCopy(ts->vec_sol, X));
166: for (j = 0; j < s; j++) w[j] = h * tab->bembed[j];
167: if (ark->is_slow) {
168: PetscCall(VecGetSubVector(X, ark->is_slow, &Xslow));
169: PetscCall(VecMAXPY(Xslow, s, w, ark->YdotRHS_slow));
170: PetscCall(VecRestoreSubVector(X, ark->is_slow, &Xslow));
171: }
172: if (ark->is_fast) {
173: PetscCall(VecGetSubVector(X, ark->is_fast, &Xfast));
174: if (fasthasE) PetscCall(VecMAXPY(Xfast, s, w, ark->YdotRHS_fast));
175: for (j = 0; j < s; j++) w[j] = h * tab->bembedt[j];
176: PetscCall(VecMAXPY(Xfast, s, w, ark->YdotI_fast));
177: PetscCall(VecRestoreSubVector(X, ark->is_fast, &Xfast));
178: }
179: } else { /* Rollback and re-complete using (bet-be,be-b) */
180: PetscCall(VecCopy(ts->vec_sol, X));
181: for (j = 0; j < s; j++) w[j] = h * (tab->bembed[j] - tab->b[j]);
182: if (ark->is_slow) {
183: PetscCall(VecGetSubVector(X, ark->is_slow, &Xslow));
184: PetscCall(VecMAXPY(Xslow, s, w, ark->YdotRHS_slow));
185: PetscCall(VecRestoreSubVector(X, ark->is_slow, &Xslow));
186: }
187: if (ark->is_fast) {
188: PetscCall(VecGetSubVector(X, ark->is_fast, &Xfast));
189: if (fasthasE) PetscCall(VecMAXPY(Xfast, tab->s, w, ark->YdotRHS_fast));
190: for (j = 0; j < s; j++) w[j] = h * (tab->bembedt[j] - tab->bt[j]);
191: PetscCall(VecMAXPY(Xfast, tab->s, w, ark->YdotI_fast));
192: PetscCall(VecRestoreSubVector(X, ark->is_fast, &Xfast));
193: }
194: }
195: if (done) *done = PETSC_TRUE;
196: PetscFunctionReturn(PETSC_SUCCESS);
197: }
198: unavailable:
199: PetscCheck(done, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "ARKIMEX '%s' of order %" PetscInt_FMT " cannot evaluate step at order %" PetscInt_FMT ". Consider using -ts_adapt_type none or a different method that has an embedded estimate.",
200: tab->name, tab->order, order);
201: *done = PETSC_FALSE;
202: PetscFunctionReturn(PETSC_SUCCESS);
203: }
205: /*
206: Additive Runge-Kutta methods for a fast-slow (component-wise partitioned) system in the form
207: Ufdot = Ff(t,Uf,Us)
208: Usdot = Fs(t,Uf,Us)
210: Ys[i] = Us_n + dt \sum_{j=1}^{i-1} a[i][j] Fs(t+c_j*dt,Yf[j],Ys[j])
211: Ys[i] = Us_n + dt \sum_{j=1}^{i-1} a[i][j] Fs(t+c_j*dt,Yf[j],Ys[j])
213: */
214: static PetscErrorCode TSStep_ARKIMEX_FastSlowSplit(TS ts)
215: {
216: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
217: ARKTableau tab = ark->tableau;
218: const PetscInt s = tab->s;
219: const PetscReal *At = tab->At, *A = tab->A, *ct = tab->ct;
220: PetscScalar *w = ark->work;
221: Vec *Y = ark->Y, Ydot_fast = ark->Ydot, Ydot0_fast = ark->Ydot0, Z = ark->Z, *YdotRHS_fast = ark->YdotRHS_fast, *YdotRHS_slow = ark->YdotRHS_slow, *YdotI_fast = ark->YdotI_fast, Yfast, Yslow, Xfast, Xslow;
222: PetscBool extrapolate = ark->extrapolate;
223: TSAdapt adapt;
224: SNES snes;
225: PetscInt i, j, its, lits;
226: PetscInt rejections = 0;
227: PetscBool fasthasE = PETSC_FALSE, stageok, accept = PETSC_TRUE;
228: PetscReal next_time_step = ts->time_step;
230: PetscFunctionBegin;
231: if (ark->is_fast) PetscCall(TSHasRHSFunction(ark->subts_fast, &fasthasE));
232: if (ark->extrapolate && !ark->Y_prev) {
233: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_fast, &Xfast));
234: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->Y_prev));
235: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_prev));
236: if (fasthasE) PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_prev));
237: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xfast));
238: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_slow, &Xslow));
239: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xslow));
240: }
242: if (!ts->steprollback) {
243: if (ts->equation_type >= TS_EQ_IMPLICIT) { /* Save the initial slope for the next step */
244: PetscCall(VecCopy(YdotI_fast[s - 1], Ydot0_fast));
245: }
246: if (ark->extrapolate && !ts->steprestart) { /* Save the Y, YdotI, YdotRHS for extrapolation initial guess */
247: for (i = 0; i < s; i++) {
248: PetscCall(VecISCopy(Y[i], ark->is_fast, SCATTER_REVERSE, ark->Y_prev[i]));
249: PetscCall(VecCopy(YdotI_fast[i], ark->YdotI_prev[i]));
250: if (fasthasE) PetscCall(VecCopy(YdotRHS_fast[i], ark->YdotRHS_prev[i]));
251: }
252: }
253: }
255: /* For IMEX we compute a step */
256: if (ts->equation_type >= TS_EQ_IMPLICIT && tab->explicit_first_stage && ts->steprestart) {
257: TS ts_start;
258: PetscCall(TSClone(ts, &ts_start));
259: PetscCall(TSSetSolution(ts_start, ts->vec_sol));
260: PetscCall(TSSetTime(ts_start, ts->ptime));
261: PetscCall(TSSetMaxSteps(ts_start, ts->steps + 1));
262: PetscCall(TSSetMaxTime(ts_start, ts->ptime + ts->time_step));
263: PetscCall(TSSetExactFinalTime(ts_start, TS_EXACTFINALTIME_STEPOVER));
264: PetscCall(TSSetTimeStep(ts_start, ts->time_step));
265: PetscCall(TSSetType(ts_start, TSARKIMEX));
266: PetscCall(TSARKIMEXSetFullyImplicit(ts_start, PETSC_TRUE));
267: PetscCall(TSARKIMEXSetType(ts_start, TSARKIMEX1BEE));
269: PetscCall(TSRestartStep(ts_start));
270: PetscCall(TSSolve(ts_start, ts->vec_sol));
271: PetscCall(TSGetTime(ts_start, &ts->ptime));
272: PetscCall(TSGetTimeStep(ts_start, &ts->time_step));
274: { /* Save the initial slope for the next step */
275: TS_ARKIMEX *ark_start = (TS_ARKIMEX *)ts_start->data;
276: PetscCall(VecCopy(ark_start->YdotI[ark_start->tableau->s - 1], Ydot0_fast));
277: }
278: ts->steps++;
279: /* Set the correct TS in SNES */
280: /* We'll try to bypass this by changing the method on the fly */
281: {
282: PetscCall(TSRHSSplitGetSNES(ts, &snes));
283: PetscCall(TSRHSSplitSetSNES(ts, snes));
284: }
285: PetscCall(TSDestroy(&ts_start));
286: }
288: ark->status = TS_STEP_INCOMPLETE;
289: while (!ts->reason && ark->status != TS_STEP_COMPLETE) {
290: PetscReal t = ts->ptime;
291: PetscReal h = ts->time_step;
292: for (i = 0; i < s; i++) {
293: ark->stage_time = t + h * ct[i];
294: PetscCall(TSPreStage(ts, ark->stage_time));
295: PetscCall(VecCopy(ts->vec_sol, Y[i]));
296: /* fast components */
297: if (ark->is_fast) {
298: if (At[i * s + i] == 0) { /* This stage is explicit */
299: PetscCheck(i == 0 || ts->equation_type < TS_EQ_IMPLICIT, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Explicit stages other than the first one are not supported for implicit problems");
300: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
301: for (j = 0; j < i; j++) w[j] = h * At[i * s + j];
302: PetscCall(VecMAXPY(Yfast, i, w, YdotI_fast));
303: if (fasthasE) {
304: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
305: PetscCall(VecMAXPY(Yfast, i, w, YdotRHS_fast));
306: }
307: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
308: } else {
309: ark->scoeff = 1. / At[i * s + i];
310: /* Ydot = shift*(Y-Z) */
311: PetscCall(VecISCopy(ts->vec_sol, ark->is_fast, SCATTER_REVERSE, Z));
312: for (j = 0; j < i; j++) w[j] = h * At[i * s + j];
313: PetscCall(VecMAXPY(Z, i, w, YdotI_fast));
314: if (fasthasE) {
315: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
316: PetscCall(VecMAXPY(Z, i, w, YdotRHS_fast));
317: }
318: PetscCall(TSRHSSplitGetSNES(ts, &snes));
319: if (ark->is_slow) PetscCall(VecCopy(i > 0 ? Y[i - 1] : ts->vec_sol, ark->Y_snes));
320: else ark->Y_snes = Y[i];
321: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
322: if (extrapolate && !ts->steprestart) {
323: /* Initial guess extrapolated from previous time step stage values */
324: PetscCall(TSExtrapolate_ARKIMEX_FastSlowSplit(ts, ct[i], Yfast));
325: } else {
326: /* Initial guess taken from last stage */
327: PetscCall(VecISCopy(i > 0 ? Y[i - 1] : ts->vec_sol, ark->is_fast, SCATTER_REVERSE, Yfast));
328: }
329: PetscCall(SNESSolve(snes, NULL, Yfast));
330: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
331: PetscCall(SNESGetIterationNumber(snes, &its));
332: PetscCall(SNESGetLinearSolveIterations(snes, &lits));
333: ts->snes_its += its;
334: ts->ksp_its += lits;
335: PetscCall(TSGetAdapt(ts, &adapt));
336: PetscCall(TSAdaptCheckStage(adapt, ts, ark->stage_time, Y[i], &stageok));
337: if (!stageok) {
338: /* We are likely rejecting the step because of solver or function domain problems so we should not attempt to
339: * use extrapolation to initialize the solves on the next attempt. */
340: extrapolate = PETSC_FALSE;
341: goto reject_step;
342: }
343: }
345: if (ts->equation_type >= TS_EQ_IMPLICIT) {
346: if (i == 0 && tab->explicit_first_stage) {
347: PetscCheck(tab->stiffly_accurate, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "%s %s is not stiffly accurate and therefore explicit-first stage methods cannot be used if the equation is implicit because the slope cannot be evaluated",
348: ((PetscObject)ts)->type_name, ark->tableau->name);
349: PetscCall(VecCopy(Ydot0_fast, YdotI_fast[0])); /* YdotI_fast = YdotI_fast(tn-1) */
350: } else {
351: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
352: PetscCall(VecAXPBYPCZ(YdotI_fast[i], -ark->scoeff / h, ark->scoeff / h, 0, Z, Yfast)); /* YdotI = shift*(X-Z) */
353: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
354: }
355: } else {
356: if (i == 0 && tab->explicit_first_stage) {
357: PetscCall(VecZeroEntries(Ydot_fast));
358: PetscCall(TSComputeIFunction(ark->subts_fast, ark->stage_time, Y[i], Ydot_fast, YdotI_fast[i], ark->imex)); /* YdotI = -G(t,Y,0) */
359: PetscCall(VecScale(YdotI_fast[i], -1.0));
360: } else {
361: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
362: PetscCall(VecAXPBYPCZ(YdotI_fast[i], -ark->scoeff / h, ark->scoeff / h, 0, Z, Yfast)); /* YdotI = shift*(X-Z) */
363: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
364: }
365: if (fasthasE) PetscCall(TSComputeRHSFunction(ark->subts_fast, ark->stage_time, Y[i], YdotRHS_fast[i]));
366: }
367: }
368: /* slow components */
369: if (ark->is_slow) {
370: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
371: PetscCall(VecGetSubVector(Y[i], ark->is_slow, &Yslow));
372: PetscCall(VecMAXPY(Yslow, i, w, YdotRHS_slow));
373: PetscCall(VecRestoreSubVector(Y[i], ark->is_slow, &Yslow));
374: PetscCall(TSComputeRHSFunction(ark->subts_slow, ark->stage_time, Y[i], YdotRHS_slow[i]));
375: }
376: PetscCall(TSPostStage(ts, ark->stage_time, i, Y));
377: }
378: ark->status = TS_STEP_INCOMPLETE;
379: PetscCall(TSEvaluateStep_ARKIMEX_FastSlowSplit(ts, tab->order, ts->vec_sol, NULL));
380: ark->status = TS_STEP_PENDING;
381: PetscCall(TSGetAdapt(ts, &adapt));
382: PetscCall(TSAdaptCandidatesClear(adapt));
383: PetscCall(TSAdaptCandidateAdd(adapt, tab->name, tab->order, 1, tab->ccfl, (PetscReal)tab->s, PETSC_TRUE));
384: PetscCall(TSAdaptChoose(adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
385: ark->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
386: if (!accept) { /* Roll back the current step */
387: PetscCall(VecCopy(ts->vec_sol0, ts->vec_sol));
388: ts->time_step = next_time_step;
389: goto reject_step;
390: }
392: ts->ptime += ts->time_step;
393: ts->time_step = next_time_step;
394: break;
396: reject_step:
397: ts->reject++;
398: accept = PETSC_FALSE;
399: if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
400: ts->reason = TS_DIVERGED_STEP_REJECTED;
401: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
402: }
403: }
404: PetscFunctionReturn(PETSC_SUCCESS);
405: }
407: static PetscErrorCode TSSetUp_ARKIMEX_FastSlowSplit(TS ts)
408: {
409: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
410: ARKTableau tab = ark->tableau;
411: Vec Xfast, Xslow;
413: PetscFunctionBegin;
414: PetscCall(PetscMalloc1(2 * tab->s, &ark->work));
415: PetscCall(VecDuplicateVecs(ts->vec_sol, tab->s, &ark->Y));
416: PetscCall(TSRHSSplitGetIS(ts, "slow", &ark->is_slow));
417: PetscCall(TSRHSSplitGetIS(ts, "fast", &ark->is_fast));
418: PetscCheck(ark->is_slow || ark->is_fast, PetscObjectComm((PetscObject)ts), PETSC_ERR_USER, "Must set up RHSSplits with TSRHSSplitSetIS() using split names 'slow' or 'fast' or both in order to use -ts_arkimex_fastslow true");
419: /* The following vectors need to be resized */
420: PetscCall(VecDestroy(&ark->Ydot));
421: PetscCall(VecDestroy(&ark->Ydot0));
422: PetscCall(VecDestroy(&ark->Z));
423: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_fast));
424: if (ark->extrapolate && ark->is_slow) { // need to resize these vectors if the fast subvectors is smaller than their original counterparts (which means IS)
425: PetscCall(VecDestroyVecs(tab->s, &ark->Y_prev));
426: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_prev));
427: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_prev));
428: }
429: /* Allocate working vectors */
430: if (ark->is_fast && ark->is_slow) PetscCall(VecDuplicate(ts->vec_sol, &ark->Y_snes)); // need an additional work vector for SNES
431: if (ark->is_fast) {
432: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_fast, &Xfast));
433: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_fast));
434: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_fast));
435: PetscCall(VecDuplicate(Xfast, &ark->Ydot));
436: PetscCall(VecDuplicate(Xfast, &ark->Ydot0));
437: PetscCall(VecDuplicate(Xfast, &ark->Z));
438: if (ark->extrapolate) {
439: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->Y_prev));
440: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_prev));
441: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_prev));
442: }
443: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xfast));
444: }
445: if (ark->is_slow) {
446: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_slow, &Xslow));
447: PetscCall(VecDuplicateVecs(Xslow, tab->s, &ark->YdotRHS_slow));
448: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_slow, &Xslow));
449: }
450: ts->ops->step = TSStep_ARKIMEX_FastSlowSplit;
451: ts->ops->evaluatestep = TSEvaluateStep_ARKIMEX_FastSlowSplit;
452: PetscCall(TSARKIMEXSetSplits(ts));
453: if (ark->subts_fast) { // subts SNESJacobian is set when users set the subts Jacobian, but the main ts SNESJacobian needs to be set too
454: SNES snes, snes_fast;
455: PetscErrorCode (*func)(SNES, Vec, Mat, Mat, void *);
456: PetscCall(TSRHSSplitGetSNES(ts, &snes));
457: PetscCall(TSGetSNES(ark->subts_fast, &snes_fast));
458: PetscCall(SNESGetJacobian(snes_fast, NULL, NULL, &func, NULL));
459: if (func == SNESTSFormJacobian) PetscCall(SNESSetJacobian(snes, NULL, NULL, SNESTSFormJacobian, ts));
460: ts->ops->snesfunction = SNESTSFormFunction_ARKIMEX_FastSlowSplit;
461: ts->ops->snesjacobian = SNESTSFormJacobian_ARKIMEX_FastSlowSplit;
462: }
463: PetscFunctionReturn(PETSC_SUCCESS);
464: }
466: static PetscErrorCode TSReset_ARKIMEX_FastSlowSplit(TS ts)
467: {
468: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
469: ARKTableau tab = ark->tableau;
471: PetscFunctionBegin;
472: if (tab) {
473: PetscCall(PetscFree(ark->work));
474: PetscCall(VecDestroyVecs(tab->s, &ark->Y));
475: if (ark->is_fast && ark->is_slow) PetscCall(VecDestroy(&ark->Y_snes));
476: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_slow));
477: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_fast));
478: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_fast));
479: PetscCall(VecDestroy(&ark->Ydot));
480: PetscCall(VecDestroy(&ark->Ydot0));
481: PetscCall(VecDestroy(&ark->Z));
482: if (ark->extrapolate) {
483: PetscCall(VecDestroyVecs(tab->s, &ark->Y_prev));
484: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_prev));
485: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_prev));
486: }
487: }
488: PetscFunctionReturn(PETSC_SUCCESS);
489: }
491: PetscErrorCode TSARKIMEXSetFastSlowSplit_ARKIMEX(TS ts, PetscBool fastslowsplit)
492: {
493: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
495: PetscFunctionBegin;
497: ark->fastslowsplit = fastslowsplit;
498: if (fastslowsplit) {
499: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSSetUp_ARKIMEX_FastSlowSplit_C", TSSetUp_ARKIMEX_FastSlowSplit));
500: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSReset_ARKIMEX_FastSlowSplit_C", TSReset_ARKIMEX_FastSlowSplit));
501: }
502: PetscFunctionReturn(PETSC_SUCCESS);
503: }
505: PetscErrorCode TSARKIMEXGetFastSlowSplit_ARKIMEX(TS ts, PetscBool *fastslowsplit)
506: {
507: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
509: PetscFunctionBegin;
511: *fastslowsplit = ark->fastslowsplit;
512: PetscFunctionReturn(PETSC_SUCCESS);
513: }