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: if (done) *done = PETSC_FALSE;
200: else
201: SETERRQ(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.", tab->name,
202: tab->order, order);
203: PetscFunctionReturn(PETSC_SUCCESS);
204: }
206: /*
207: Additive Runge-Kutta methods for a fast-slow (component-wise partitioned) system in the form
208: Ufdot = Ff(t,Uf,Us)
209: Usdot = Fs(t,Uf,Us)
211: Ys[i] = Us_n + dt \sum_{j=1}^{i-1} a[i][j] Fs(t+c_j*dt,Yf[j],Ys[j])
212: Ys[i] = Us_n + dt \sum_{j=1}^{i-1} a[i][j] Fs(t+c_j*dt,Yf[j],Ys[j])
214: */
215: static PetscErrorCode TSStep_ARKIMEX_FastSlowSplit(TS ts)
216: {
217: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
218: ARKTableau tab = ark->tableau;
219: const PetscInt s = tab->s;
220: const PetscReal *At = tab->At, *A = tab->A, *ct = tab->ct;
221: PetscScalar *w = ark->work;
222: 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;
223: PetscBool extrapolate = ark->extrapolate;
224: TSAdapt adapt;
225: SNES snes;
226: PetscInt i, j, its, lits;
227: PetscInt rejections = 0;
228: PetscBool fasthasE = PETSC_FALSE, stageok, accept = PETSC_TRUE;
229: PetscReal next_time_step = ts->time_step;
231: PetscFunctionBegin;
232: if (ark->is_fast) PetscCall(TSHasRHSFunction(ark->subts_fast, &fasthasE));
233: if (ark->extrapolate && !ark->Y_prev) {
234: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_fast, &Xfast));
235: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->Y_prev));
236: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_prev));
237: if (fasthasE) PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_prev));
238: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xfast));
239: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_slow, &Xslow));
240: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xslow));
241: }
243: if (!ts->steprollback) {
244: if (ts->equation_type >= TS_EQ_IMPLICIT) { /* Save the initial slope for the next step */
245: PetscCall(VecCopy(YdotI_fast[s - 1], Ydot0_fast));
246: }
247: if (ark->extrapolate && !ts->steprestart) { /* Save the Y, YdotI, YdotRHS for extrapolation initial guess */
248: for (i = 0; i < s; i++) {
249: PetscCall(VecISCopy(Y[i], ark->is_fast, SCATTER_REVERSE, ark->Y_prev[i]));
250: PetscCall(VecCopy(YdotI_fast[i], ark->YdotI_prev[i]));
251: if (fasthasE) PetscCall(VecCopy(YdotRHS_fast[i], ark->YdotRHS_prev[i]));
252: }
253: }
254: }
256: /* For IMEX we compute a step */
257: if (ts->equation_type >= TS_EQ_IMPLICIT && tab->explicit_first_stage && ts->steprestart) {
258: TS ts_start;
259: PetscCall(TSClone(ts, &ts_start));
260: PetscCall(TSSetSolution(ts_start, ts->vec_sol));
261: PetscCall(TSSetTime(ts_start, ts->ptime));
262: PetscCall(TSSetMaxSteps(ts_start, ts->steps + 1));
263: PetscCall(TSSetMaxTime(ts_start, ts->ptime + ts->time_step));
264: PetscCall(TSSetExactFinalTime(ts_start, TS_EXACTFINALTIME_STEPOVER));
265: PetscCall(TSSetTimeStep(ts_start, ts->time_step));
266: PetscCall(TSSetType(ts_start, TSARKIMEX));
267: PetscCall(TSARKIMEXSetFullyImplicit(ts_start, PETSC_TRUE));
268: PetscCall(TSARKIMEXSetType(ts_start, TSARKIMEX1BEE));
270: PetscCall(TSRestartStep(ts_start));
271: PetscCall(TSSolve(ts_start, ts->vec_sol));
272: PetscCall(TSGetTime(ts_start, &ts->ptime));
273: PetscCall(TSGetTimeStep(ts_start, &ts->time_step));
275: { /* Save the initial slope for the next step */
276: TS_ARKIMEX *ark_start = (TS_ARKIMEX *)ts_start->data;
277: PetscCall(VecCopy(ark_start->YdotI[ark_start->tableau->s - 1], Ydot0_fast));
278: }
279: ts->steps++;
280: /* Set the correct TS in SNES */
281: /* We'll try to bypass this by changing the method on the fly */
282: {
283: PetscCall(TSRHSSplitGetSNES(ts, &snes));
284: PetscCall(TSRHSSplitSetSNES(ts, snes));
285: }
286: PetscCall(TSDestroy(&ts_start));
287: }
289: ark->status = TS_STEP_INCOMPLETE;
290: while (!ts->reason && ark->status != TS_STEP_COMPLETE) {
291: PetscReal t = ts->ptime;
292: PetscReal h = ts->time_step;
293: for (i = 0; i < s; i++) {
294: ark->stage_time = t + h * ct[i];
295: PetscCall(TSPreStage(ts, ark->stage_time));
296: PetscCall(VecCopy(ts->vec_sol, Y[i]));
297: /* fast components */
298: if (ark->is_fast) {
299: if (At[i * s + i] == 0) { /* This stage is explicit */
300: 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");
301: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
302: for (j = 0; j < i; j++) w[j] = h * At[i * s + j];
303: PetscCall(VecMAXPY(Yfast, i, w, YdotI_fast));
304: if (fasthasE) {
305: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
306: PetscCall(VecMAXPY(Yfast, i, w, YdotRHS_fast));
307: }
308: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
309: } else {
310: ark->scoeff = 1. / At[i * s + i];
311: /* Ydot = shift*(Y-Z) */
312: PetscCall(VecISCopy(ts->vec_sol, ark->is_fast, SCATTER_REVERSE, Z));
313: for (j = 0; j < i; j++) w[j] = h * At[i * s + j];
314: PetscCall(VecMAXPY(Z, i, w, YdotI_fast));
315: if (fasthasE) {
316: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
317: PetscCall(VecMAXPY(Z, i, w, YdotRHS_fast));
318: }
319: PetscCall(TSRHSSplitGetSNES(ts, &snes));
320: if (ark->is_slow) PetscCall(VecCopy(i > 0 ? Y[i - 1] : ts->vec_sol, ark->Y_snes));
321: else ark->Y_snes = Y[i];
322: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
323: if (extrapolate && !ts->steprestart) {
324: /* Initial guess extrapolated from previous time step stage values */
325: PetscCall(TSExtrapolate_ARKIMEX_FastSlowSplit(ts, ct[i], Yfast));
326: } else {
327: /* Initial guess taken from last stage */
328: PetscCall(VecISCopy(i > 0 ? Y[i - 1] : ts->vec_sol, ark->is_fast, SCATTER_REVERSE, Yfast));
329: }
330: PetscCall(SNESSolve(snes, NULL, Yfast));
331: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
332: PetscCall(SNESGetIterationNumber(snes, &its));
333: PetscCall(SNESGetLinearSolveIterations(snes, &lits));
334: ts->snes_its += its;
335: ts->ksp_its += lits;
336: PetscCall(TSGetAdapt(ts, &adapt));
337: PetscCall(TSAdaptCheckStage(adapt, ts, ark->stage_time, Y[i], &stageok));
338: if (!stageok) {
339: /* We are likely rejecting the step because of solver or function domain problems so we should not attempt to
340: * use extrapolation to initialize the solves on the next attempt. */
341: extrapolate = PETSC_FALSE;
342: goto reject_step;
343: }
344: }
346: if (ts->equation_type >= TS_EQ_IMPLICIT) {
347: if (i == 0 && tab->explicit_first_stage) {
348: 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",
349: ((PetscObject)ts)->type_name, ark->tableau->name);
350: PetscCall(VecCopy(Ydot0_fast, YdotI_fast[0])); /* YdotI_fast = YdotI_fast(tn-1) */
351: } else {
352: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
353: PetscCall(VecAXPBYPCZ(YdotI_fast[i], -ark->scoeff / h, ark->scoeff / h, 0, Z, Yfast)); /* YdotI = shift*(X-Z) */
354: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
355: }
356: } else {
357: if (i == 0 && tab->explicit_first_stage) {
358: PetscCall(VecZeroEntries(Ydot_fast));
359: PetscCall(TSComputeIFunction(ark->subts_fast, ark->stage_time, Y[i], Ydot_fast, YdotI_fast[i], ark->imex)); /* YdotI = -G(t,Y,0) */
360: PetscCall(VecScale(YdotI_fast[i], -1.0));
361: } else {
362: PetscCall(VecGetSubVector(Y[i], ark->is_fast, &Yfast));
363: PetscCall(VecAXPBYPCZ(YdotI_fast[i], -ark->scoeff / h, ark->scoeff / h, 0, Z, Yfast)); /* YdotI = shift*(X-Z) */
364: PetscCall(VecRestoreSubVector(Y[i], ark->is_fast, &Yfast));
365: }
366: if (fasthasE) PetscCall(TSComputeRHSFunction(ark->subts_fast, ark->stage_time, Y[i], YdotRHS_fast[i]));
367: }
368: }
369: /* slow components */
370: if (ark->is_slow) {
371: for (j = 0; j < i; j++) w[j] = h * A[i * s + j];
372: PetscCall(VecGetSubVector(Y[i], ark->is_slow, &Yslow));
373: PetscCall(VecMAXPY(Yslow, i, w, YdotRHS_slow));
374: PetscCall(VecRestoreSubVector(Y[i], ark->is_slow, &Yslow));
375: PetscCall(TSComputeRHSFunction(ark->subts_slow, ark->stage_time, Y[i], YdotRHS_slow[i]));
376: }
377: PetscCall(TSPostStage(ts, ark->stage_time, i, Y));
378: }
379: ark->status = TS_STEP_INCOMPLETE;
380: PetscCall(TSEvaluateStep_ARKIMEX_FastSlowSplit(ts, tab->order, ts->vec_sol, NULL));
381: ark->status = TS_STEP_PENDING;
382: PetscCall(TSGetAdapt(ts, &adapt));
383: PetscCall(TSAdaptCandidatesClear(adapt));
384: PetscCall(TSAdaptCandidateAdd(adapt, tab->name, tab->order, 1, tab->ccfl, (PetscReal)tab->s, PETSC_TRUE));
385: PetscCall(TSAdaptChoose(adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
386: ark->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
387: if (!accept) { /* Roll back the current step */
388: PetscCall(VecCopy(ts->vec_sol0, ts->vec_sol));
389: ts->time_step = next_time_step;
390: goto reject_step;
391: }
393: ts->ptime += ts->time_step;
394: ts->time_step = next_time_step;
395: break;
397: reject_step:
398: ts->reject++;
399: accept = PETSC_FALSE;
400: if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
401: ts->reason = TS_DIVERGED_STEP_REJECTED;
402: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
403: }
404: }
405: PetscFunctionReturn(PETSC_SUCCESS);
406: }
408: static PetscErrorCode TSSetUp_ARKIMEX_FastSlowSplit(TS ts)
409: {
410: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
411: ARKTableau tab = ark->tableau;
412: Vec Xfast, Xslow;
414: PetscFunctionBegin;
415: PetscCall(PetscMalloc1(2 * tab->s, &ark->work));
416: PetscCall(VecDuplicateVecs(ts->vec_sol, tab->s, &ark->Y));
417: PetscCall(TSRHSSplitGetIS(ts, "slow", &ark->is_slow));
418: PetscCall(TSRHSSplitGetIS(ts, "fast", &ark->is_fast));
419: 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");
420: /* The following vectors need to be resized */
421: PetscCall(VecDestroy(&ark->Ydot));
422: PetscCall(VecDestroy(&ark->Ydot0));
423: PetscCall(VecDestroy(&ark->Z));
424: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_fast));
425: if (ark->extrapolate && ark->is_slow) { // need to resize these vectors if the fast subvectors is smaller than their original counterparts (which means IS)
426: PetscCall(VecDestroyVecs(tab->s, &ark->Y_prev));
427: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_prev));
428: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_prev));
429: }
430: /* Allocate working vectors */
431: if (ark->is_fast && ark->is_slow) PetscCall(VecDuplicate(ts->vec_sol, &ark->Y_snes)); // need an additional work vector for SNES
432: if (ark->is_fast) {
433: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_fast, &Xfast));
434: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_fast));
435: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_fast));
436: PetscCall(VecDuplicate(Xfast, &ark->Ydot));
437: PetscCall(VecDuplicate(Xfast, &ark->Ydot0));
438: PetscCall(VecDuplicate(Xfast, &ark->Z));
439: if (ark->extrapolate) {
440: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->Y_prev));
441: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotI_prev));
442: PetscCall(VecDuplicateVecs(Xfast, tab->s, &ark->YdotRHS_prev));
443: }
444: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_fast, &Xfast));
445: }
446: if (ark->is_slow) {
447: PetscCall(VecGetSubVector(ts->vec_sol, ark->is_slow, &Xslow));
448: PetscCall(VecDuplicateVecs(Xslow, tab->s, &ark->YdotRHS_slow));
449: PetscCall(VecRestoreSubVector(ts->vec_sol, ark->is_slow, &Xslow));
450: }
451: ts->ops->step = TSStep_ARKIMEX_FastSlowSplit;
452: ts->ops->evaluatestep = TSEvaluateStep_ARKIMEX_FastSlowSplit;
453: PetscCall(TSARKIMEXSetSplits(ts));
454: if (ark->subts_fast) { // subts SNESJacobian is set when users set the subts Jacobian, but the main ts SNESJacobian needs to be set too
455: SNES snes, snes_fast;
456: PetscErrorCode (*func)(SNES, Vec, Mat, Mat, void *);
457: PetscCall(TSRHSSplitGetSNES(ts, &snes));
458: PetscCall(TSGetSNES(ark->subts_fast, &snes_fast));
459: PetscCall(SNESGetJacobian(snes_fast, NULL, NULL, &func, NULL));
460: if (func == SNESTSFormJacobian) PetscCall(SNESSetJacobian(snes, NULL, NULL, SNESTSFormJacobian, ts));
461: ts->ops->snesfunction = SNESTSFormFunction_ARKIMEX_FastSlowSplit;
462: ts->ops->snesjacobian = SNESTSFormJacobian_ARKIMEX_FastSlowSplit;
463: }
464: PetscFunctionReturn(PETSC_SUCCESS);
465: }
467: static PetscErrorCode TSReset_ARKIMEX_FastSlowSplit(TS ts)
468: {
469: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
470: ARKTableau tab = ark->tableau;
472: PetscFunctionBegin;
473: if (tab) {
474: PetscCall(PetscFree(ark->work));
475: PetscCall(VecDestroyVecs(tab->s, &ark->Y));
476: if (ark->is_fast && ark->is_slow) PetscCall(VecDestroy(&ark->Y_snes));
477: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_slow));
478: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_fast));
479: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_fast));
480: PetscCall(VecDestroy(&ark->Ydot));
481: PetscCall(VecDestroy(&ark->Ydot0));
482: PetscCall(VecDestroy(&ark->Z));
483: if (ark->extrapolate) {
484: PetscCall(VecDestroyVecs(tab->s, &ark->Y_prev));
485: PetscCall(VecDestroyVecs(tab->s, &ark->YdotI_prev));
486: PetscCall(VecDestroyVecs(tab->s, &ark->YdotRHS_prev));
487: }
488: }
489: PetscFunctionReturn(PETSC_SUCCESS);
490: }
492: PetscErrorCode TSARKIMEXSetFastSlowSplit_ARKIMEX(TS ts, PetscBool fastslowsplit)
493: {
494: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
496: PetscFunctionBegin;
498: ark->fastslowsplit = fastslowsplit;
499: if (fastslowsplit) {
500: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSSetUp_ARKIMEX_FastSlowSplit_C", TSSetUp_ARKIMEX_FastSlowSplit));
501: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSReset_ARKIMEX_FastSlowSplit_C", TSReset_ARKIMEX_FastSlowSplit));
502: }
503: PetscFunctionReturn(PETSC_SUCCESS);
504: }
506: PetscErrorCode TSARKIMEXGetFastSlowSplit_ARKIMEX(TS ts, PetscBool *fastslowsplit)
507: {
508: TS_ARKIMEX *ark = (TS_ARKIMEX *)ts->data;
510: PetscFunctionBegin;
512: *fastslowsplit = ark->fastslowsplit;
513: PetscFunctionReturn(PETSC_SUCCESS);
514: }