Date: 2025-03-28T21:11:16+00:00

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: }