xref: /petsc/src/ts/impls/eimex/eimex.c (revision 4dfa11a44d5adf2389f1d3acbc8f3c1116dc6c3a)
17d5697caSHong Zhang 
2af0996ceSBarry Smith #include <petsc/private/tsimpl.h> /*I   "petscts.h"   I*/
34042b796SJed Brown #include <petscdm.h>
47d5697caSHong Zhang 
57d5697caSHong Zhang static const PetscInt TSEIMEXDefault = 3;
67d5697caSHong Zhang 
77d5697caSHong Zhang typedef struct {
87d5697caSHong Zhang   PetscInt     row_ind;    /* Return the term T[row_ind][col_ind] */
97d5697caSHong Zhang   PetscInt     col_ind;    /* Return the term T[row_ind][col_ind] */
107d5697caSHong Zhang   PetscInt     nstages;    /* Numbers of stages in current scheme */
117d5697caSHong Zhang   PetscInt     max_rows;   /* Maximum number of rows */
127d5697caSHong Zhang   PetscInt    *N;          /* Harmonic sequence N[max_rows] */
137d5697caSHong Zhang   Vec          Y;          /* States computed during the step, used to complete the step */
147d5697caSHong Zhang   Vec          Z;          /* For shift*(Y-Z) */
157d5697caSHong Zhang   Vec         *T;          /* Working table, size determined by nstages */
167d5697caSHong Zhang   Vec          YdotRHS;    /* f(x) Work vector holding YdotRHS during residual evaluation */
177d5697caSHong Zhang   Vec          YdotI;      /* xdot-g(x) Work vector holding YdotI = G(t,x,xdot) when xdot =0 */
187d5697caSHong Zhang   Vec          Ydot;       /* f(x)+g(x) Work vector */
197d5697caSHong Zhang   Vec          VecSolPrev; /* Work vector holding the solution from the previous step (used for interpolation) */
207d5697caSHong Zhang   PetscReal    shift;
217d5697caSHong Zhang   PetscReal    ctime;
227d5697caSHong Zhang   PetscBool    recompute_jacobian; /* Recompute the Jacobian at each stage, default is to freeze the Jacobian at the start of each step */
237d5697caSHong Zhang   PetscBool    ord_adapt;          /* order adapativity */
247d5697caSHong Zhang   TSStepStatus status;
257d5697caSHong Zhang } TS_EIMEX;
267d5697caSHong Zhang 
277d5697caSHong Zhang /* This function is pure */
289371c9d4SSatish Balay static PetscInt Map(PetscInt i, PetscInt j, PetscInt s) {
297d5697caSHong Zhang   return ((2 * s - j + 1) * j / 2 + i - j);
307d5697caSHong Zhang }
317d5697caSHong Zhang 
329371c9d4SSatish Balay static PetscErrorCode TSEvaluateStep_EIMEX(TS ts, PetscInt order, Vec X, PetscBool *done) {
337d5697caSHong Zhang   TS_EIMEX      *ext = (TS_EIMEX *)ts->data;
347d5697caSHong Zhang   const PetscInt ns  = ext->nstages;
357d5697caSHong Zhang   PetscFunctionBegin;
369566063dSJacob Faibussowitsch   PetscCall(VecCopy(ext->T[Map(ext->row_ind, ext->col_ind, ns)], X));
377d5697caSHong Zhang   PetscFunctionReturn(0);
387d5697caSHong Zhang }
397d5697caSHong Zhang 
409371c9d4SSatish Balay static PetscErrorCode TSStage_EIMEX(TS ts, PetscInt istage) {
417d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
427d5697caSHong Zhang   PetscReal h;
437d5697caSHong Zhang   Vec       Y = ext->Y, Z = ext->Z;
447d5697caSHong Zhang   SNES      snes;
457d5697caSHong Zhang   TSAdapt   adapt;
467d5697caSHong Zhang   PetscInt  i, its, lits;
477d5697caSHong Zhang   PetscBool accept;
487d5697caSHong Zhang 
497d5697caSHong Zhang   PetscFunctionBegin;
509566063dSJacob Faibussowitsch   PetscCall(TSGetSNES(ts, &snes));
517d5697caSHong Zhang   h          = ts->time_step / ext->N[istage]; /* step size for the istage-th stage */
527d5697caSHong Zhang   ext->shift = 1. / h;
539566063dSJacob Faibussowitsch   PetscCall(SNESSetLagJacobian(snes, -2)); /* Recompute the Jacobian on this solve, but not again */
549566063dSJacob Faibussowitsch   PetscCall(VecCopy(ext->VecSolPrev, Y));  /* Take the previous solution as initial step */
557d5697caSHong Zhang 
567d5697caSHong Zhang   for (i = 0; i < ext->N[istage]; i++) {
577d5697caSHong Zhang     ext->ctime = ts->ptime + h * i;
589566063dSJacob Faibussowitsch     PetscCall(VecCopy(Y, Z)); /* Save the solution of the previous substep */
599566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Y));
609566063dSJacob Faibussowitsch     PetscCall(SNESGetIterationNumber(snes, &its));
619566063dSJacob Faibussowitsch     PetscCall(SNESGetLinearSolveIterations(snes, &lits));
629371c9d4SSatish Balay     ts->snes_its += its;
639371c9d4SSatish Balay     ts->ksp_its += lits;
649566063dSJacob Faibussowitsch     PetscCall(TSGetAdapt(ts, &adapt));
659566063dSJacob Faibussowitsch     PetscCall(TSAdaptCheckStage(adapt, ts, ext->ctime, Y, &accept));
667d5697caSHong Zhang   }
677d5697caSHong Zhang   PetscFunctionReturn(0);
687d5697caSHong Zhang }
697d5697caSHong Zhang 
709371c9d4SSatish Balay static PetscErrorCode TSStep_EIMEX(TS ts) {
717d5697caSHong Zhang   TS_EIMEX      *ext = (TS_EIMEX *)ts->data;
727d5697caSHong Zhang   const PetscInt ns  = ext->nstages;
737d5697caSHong Zhang   Vec           *T = ext->T, Y = ext->Y;
747d5697caSHong Zhang   SNES           snes;
757d5697caSHong Zhang   PetscInt       i, j;
767d5697caSHong Zhang   PetscBool      accept = PETSC_FALSE;
777453f775SEmil Constantinescu   PetscReal      alpha, local_error, local_error_a, local_error_r;
787d5697caSHong Zhang 
79d0609cedSBarry Smith   PetscFunctionBegin;
809566063dSJacob Faibussowitsch   PetscCall(TSGetSNES(ts, &snes));
819566063dSJacob Faibussowitsch   PetscCall(SNESSetType(snes, "ksponly"));
827d5697caSHong Zhang   ext->status = TS_STEP_INCOMPLETE;
837d5697caSHong Zhang 
849566063dSJacob Faibussowitsch   PetscCall(VecCopy(ts->vec_sol, ext->VecSolPrev));
857d5697caSHong Zhang 
867d5697caSHong Zhang   /* Apply n_j steps of the base method to obtain solutions of T(j,1),1<=j<=s */
877d5697caSHong Zhang   for (j = 0; j < ns; j++) {
889566063dSJacob Faibussowitsch     PetscCall(TSStage_EIMEX(ts, j));
899566063dSJacob Faibussowitsch     PetscCall(VecCopy(Y, T[j]));
907d5697caSHong Zhang   }
917d5697caSHong Zhang 
927d5697caSHong Zhang   for (i = 1; i < ns; i++) {
937d5697caSHong Zhang     for (j = i; j < ns; j++) {
947d5697caSHong Zhang       alpha = -(PetscReal)ext->N[j] / ext->N[j - i];
95d0609cedSBarry Smith       PetscCall(VecAXPBYPCZ(T[Map(j, i, ns)], alpha, 1.0, 0, T[Map(j, i - 1, ns)], T[Map(j - 1, i - 1, ns)])); /* T[j][i]=alpha*T[j][i-1]+T[j-1][i-1] */
967d5697caSHong Zhang       alpha = 1.0 / (1.0 + alpha);
979566063dSJacob Faibussowitsch       PetscCall(VecScale(T[Map(j, i, ns)], alpha));
987d5697caSHong Zhang     }
997d5697caSHong Zhang   }
1007d5697caSHong Zhang 
1019566063dSJacob Faibussowitsch   PetscCall(TSEvaluateStep(ts, ns, ts->vec_sol, NULL)); /*update ts solution */
1027d5697caSHong Zhang 
1037d5697caSHong Zhang   if (ext->ord_adapt && ext->nstages < ext->max_rows) {
1047d5697caSHong Zhang     accept = PETSC_FALSE;
1057d5697caSHong Zhang     while (!accept && ext->nstages < ext->max_rows) {
1069566063dSJacob Faibussowitsch       PetscCall(TSErrorWeightedNorm(ts, ts->vec_sol, T[Map(ext->nstages - 1, ext->nstages - 2, ext->nstages)], ts->adapt->wnormtype, &local_error, &local_error_a, &local_error_r));
1077d5697caSHong Zhang       accept = (local_error < 1.0) ? PETSC_TRUE : PETSC_FALSE;
1087d5697caSHong Zhang 
1097d5697caSHong Zhang       if (!accept) { /* add one more stage*/
1109566063dSJacob Faibussowitsch         PetscCall(TSStage_EIMEX(ts, ext->nstages));
1119371c9d4SSatish Balay         ext->nstages++;
1129371c9d4SSatish Balay         ext->row_ind++;
1139371c9d4SSatish Balay         ext->col_ind++;
1147d5697caSHong Zhang         /*T table need to be recycled*/
1159566063dSJacob Faibussowitsch         PetscCall(VecDuplicateVecs(ts->vec_sol, (1 + ext->nstages) * ext->nstages / 2, &ext->T));
1167d5697caSHong Zhang         for (i = 0; i < ext->nstages - 1; i++) {
11748a46eb9SPierre Jolivet           for (j = 0; j <= i; j++) PetscCall(VecCopy(T[Map(i, j, ext->nstages - 1)], ext->T[Map(i, j, ext->nstages)]));
1187d5697caSHong Zhang         }
1199566063dSJacob Faibussowitsch         PetscCall(VecDestroyVecs(ext->nstages * (ext->nstages - 1) / 2, &T));
1207d5697caSHong Zhang         T = ext->T; /*reset the pointer*/
1217d5697caSHong Zhang         /*recycling finished, store the new solution*/
1229566063dSJacob Faibussowitsch         PetscCall(VecCopy(Y, T[ext->nstages - 1]));
1237d5697caSHong Zhang         /*extrapolation for the newly added stage*/
1247d5697caSHong Zhang         for (i = 1; i < ext->nstages; i++) {
1257d5697caSHong Zhang           alpha = -(PetscReal)ext->N[ext->nstages - 1] / ext->N[ext->nstages - 1 - i];
126d0609cedSBarry Smith           PetscCall(VecAXPBYPCZ(T[Map(ext->nstages - 1, i, ext->nstages)], alpha, 1.0, 0, T[Map(ext->nstages - 1, i - 1, ext->nstages)], T[Map(ext->nstages - 1 - 1, i - 1, ext->nstages)])); /*T[ext->nstages-1][i]=alpha*T[ext->nstages-1][i-1]+T[ext->nstages-1-1][i-1]*/
1277d5697caSHong Zhang           alpha = 1.0 / (1.0 + alpha);
1289566063dSJacob Faibussowitsch           PetscCall(VecScale(T[Map(ext->nstages - 1, i, ext->nstages)], alpha));
1297d5697caSHong Zhang         }
1307d5697caSHong Zhang         /*update ts solution */
1319566063dSJacob Faibussowitsch         PetscCall(TSEvaluateStep(ts, ext->nstages, ts->vec_sol, NULL));
1327d5697caSHong Zhang       } /*end if !accept*/
1337d5697caSHong Zhang     }   /*end while*/
1347d5697caSHong Zhang 
13548a46eb9SPierre Jolivet     if (ext->nstages == ext->max_rows) PetscCall(PetscInfo(ts, "Max number of rows has been used\n"));
1367d5697caSHong Zhang   } /*end if ext->ord_adapt*/
1377d5697caSHong Zhang   ts->ptime += ts->time_step;
1387d5697caSHong Zhang   ext->status = TS_STEP_COMPLETE;
1397d5697caSHong Zhang 
1407d5697caSHong Zhang   if (ext->status != TS_STEP_COMPLETE && !ts->reason) ts->reason = TS_DIVERGED_STEP_REJECTED;
1417d5697caSHong Zhang   PetscFunctionReturn(0);
1427d5697caSHong Zhang }
1437d5697caSHong Zhang 
1447d5697caSHong Zhang /* cubic Hermit spline */
1459371c9d4SSatish Balay static PetscErrorCode TSInterpolate_EIMEX(TS ts, PetscReal itime, Vec X) {
1467d5697caSHong Zhang   TS_EIMEX       *ext = (TS_EIMEX *)ts->data;
1477d5697caSHong Zhang   PetscReal       t, a, b;
148be5899b3SLisandro Dalcin   Vec             Y0 = ext->VecSolPrev, Y1 = ext->Y, Ydot = ext->Ydot, YdotI = ext->YdotI;
149be5899b3SLisandro Dalcin   const PetscReal h = ts->ptime - ts->ptime_prev;
1507d5697caSHong Zhang   PetscFunctionBegin;
1517d5697caSHong Zhang   t = (itime - ts->ptime + h) / h;
1527d5697caSHong Zhang   /* YdotI = -f(x)-g(x) */
1537d5697caSHong Zhang 
1549566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(Ydot));
1559566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ts->ptime - h, Y0, Ydot, YdotI, PETSC_FALSE));
1567d5697caSHong Zhang 
15716dc5bacSEmil Constantinescu   a = 2.0 * t * t * t - 3.0 * t * t + 1.0;
15816dc5bacSEmil Constantinescu   b = -(t * t * t - 2.0 * t * t + t) * h;
1599566063dSJacob Faibussowitsch   PetscCall(VecAXPBYPCZ(X, a, b, 0.0, Y0, YdotI));
1607d5697caSHong Zhang 
1619566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ts->ptime, Y1, Ydot, YdotI, PETSC_FALSE));
16216dc5bacSEmil Constantinescu   a = -2.0 * t * t * t + 3.0 * t * t;
16316dc5bacSEmil Constantinescu   b = -(t * t * t - t * t) * h;
1649566063dSJacob Faibussowitsch   PetscCall(VecAXPBYPCZ(X, a, b, 1.0, Y1, YdotI));
1657d5697caSHong Zhang 
1667d5697caSHong Zhang   PetscFunctionReturn(0);
1677d5697caSHong Zhang }
1687d5697caSHong Zhang 
1699371c9d4SSatish Balay static PetscErrorCode TSReset_EIMEX(TS ts) {
1707d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
1717d5697caSHong Zhang   PetscInt  ns;
1727d5697caSHong Zhang 
1737d5697caSHong Zhang   PetscFunctionBegin;
1747d5697caSHong Zhang   ns = ext->nstages;
1759566063dSJacob Faibussowitsch   PetscCall(VecDestroyVecs((1 + ns) * ns / 2, &ext->T));
1769566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Y));
1779566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Z));
1789566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->YdotRHS));
1799566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->YdotI));
1809566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Ydot));
1819566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->VecSolPrev));
1829566063dSJacob Faibussowitsch   PetscCall(PetscFree(ext->N));
1837d5697caSHong Zhang   PetscFunctionReturn(0);
1847d5697caSHong Zhang }
1857d5697caSHong Zhang 
1869371c9d4SSatish Balay static PetscErrorCode TSDestroy_EIMEX(TS ts) {
1877d5697caSHong Zhang   PetscFunctionBegin;
1889566063dSJacob Faibussowitsch   PetscCall(TSReset_EIMEX(ts));
1899566063dSJacob Faibussowitsch   PetscCall(PetscFree(ts->data));
1909566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetMaxRows_C", NULL));
1919566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetRowCol_C", NULL));
1929566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetOrdAdapt_C", NULL));
1937d5697caSHong Zhang   PetscFunctionReturn(0);
1947d5697caSHong Zhang }
1957d5697caSHong Zhang 
1969371c9d4SSatish Balay static PetscErrorCode TSEIMEXGetVecs(TS ts, DM dm, Vec *Z, Vec *Ydot, Vec *YdotI, Vec *YdotRHS) {
1977d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
1987d5697caSHong Zhang 
1997d5697caSHong Zhang   PetscFunctionBegin;
2007d5697caSHong Zhang   if (Z) {
2017d5697caSHong Zhang     if (dm && dm != ts->dm) {
2029566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_Z", Z));
2037d5697caSHong Zhang     } else *Z = ext->Z;
2047d5697caSHong Zhang   }
2057d5697caSHong Zhang   if (Ydot) {
2067d5697caSHong Zhang     if (dm && dm != ts->dm) {
2079566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_Ydot", Ydot));
2087d5697caSHong Zhang     } else *Ydot = ext->Ydot;
2097d5697caSHong Zhang   }
2107d5697caSHong Zhang   if (YdotI) {
2117d5697caSHong Zhang     if (dm && dm != ts->dm) {
2129566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_YdotI", YdotI));
2137d5697caSHong Zhang     } else *YdotI = ext->YdotI;
2147d5697caSHong Zhang   }
2157d5697caSHong Zhang   if (YdotRHS) {
2167d5697caSHong Zhang     if (dm && dm != ts->dm) {
2179566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_YdotRHS", YdotRHS));
2187d5697caSHong Zhang     } else *YdotRHS = ext->YdotRHS;
2197d5697caSHong Zhang   }
2207d5697caSHong Zhang   PetscFunctionReturn(0);
2217d5697caSHong Zhang }
2227d5697caSHong Zhang 
2239371c9d4SSatish Balay static PetscErrorCode TSEIMEXRestoreVecs(TS ts, DM dm, Vec *Z, Vec *Ydot, Vec *YdotI, Vec *YdotRHS) {
2247d5697caSHong Zhang   PetscFunctionBegin;
2257d5697caSHong Zhang   if (Z) {
22648a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_Z", Z));
2277d5697caSHong Zhang   }
2287d5697caSHong Zhang   if (Ydot) {
22948a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_Ydot", Ydot));
2307d5697caSHong Zhang   }
2317d5697caSHong Zhang   if (YdotI) {
23248a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_YdotI", YdotI));
2337d5697caSHong Zhang   }
2347d5697caSHong Zhang   if (YdotRHS) {
23548a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_YdotRHS", YdotRHS));
2367d5697caSHong Zhang   }
2377d5697caSHong Zhang   PetscFunctionReturn(0);
2387d5697caSHong Zhang }
2397d5697caSHong Zhang 
2407d5697caSHong Zhang /*
2417d5697caSHong Zhang   This defines the nonlinear equation that is to be solved with SNES
2427d5697caSHong Zhang   Fn[t0+Theta*dt, U, (U-U0)*shift] = 0
2437d5697caSHong Zhang   In the case of Backward Euler, Fn = (U-U0)/h-g(t1,U))
2447d5697caSHong Zhang   Since FormIFunction calculates G = ydot - g(t,y), ydot will be set to (U-U0)/h
2457d5697caSHong Zhang */
2469371c9d4SSatish Balay static PetscErrorCode SNESTSFormFunction_EIMEX(SNES snes, Vec X, Vec G, TS ts) {
2477d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
2487d5697caSHong Zhang   Vec       Ydot, Z;
2497d5697caSHong Zhang   DM        dm, dmsave;
2507d5697caSHong Zhang 
2517d5697caSHong Zhang   PetscFunctionBegin;
2529566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(G));
2537d5697caSHong Zhang 
2549566063dSJacob Faibussowitsch   PetscCall(SNESGetDM(snes, &dm));
2559566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, dm, &Z, &Ydot, NULL, NULL));
2569566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(Ydot));
2577d5697caSHong Zhang   dmsave = ts->dm;
2587d5697caSHong Zhang   ts->dm = dm;
2599566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ext->ctime, X, Ydot, G, PETSC_FALSE));
2607d5697caSHong Zhang   /* PETSC_FALSE indicates non-imex, adding explicit RHS to the implicit I function.  */
2619566063dSJacob Faibussowitsch   PetscCall(VecCopy(G, Ydot));
2627d5697caSHong Zhang   ts->dm = dmsave;
2639566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, dm, &Z, &Ydot, NULL, NULL));
2647d5697caSHong Zhang 
2657d5697caSHong Zhang   PetscFunctionReturn(0);
2667d5697caSHong Zhang }
2677d5697caSHong Zhang 
2687d5697caSHong Zhang /*
2697d5697caSHong Zhang  This defined the Jacobian matrix for SNES. Jn = (I/h-g'(t,y))
2707d5697caSHong Zhang  */
2719371c9d4SSatish Balay static PetscErrorCode SNESTSFormJacobian_EIMEX(SNES snes, Vec X, Mat A, Mat B, TS ts) {
2727d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
2737d5697caSHong Zhang   Vec       Ydot;
2747d5697caSHong Zhang   DM        dm, dmsave;
2757d5697caSHong Zhang   PetscFunctionBegin;
2769566063dSJacob Faibussowitsch   PetscCall(SNESGetDM(snes, &dm));
2779566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, dm, NULL, &Ydot, NULL, NULL));
2789566063dSJacob Faibussowitsch   /*  PetscCall(VecZeroEntries(Ydot)); */
2797d5697caSHong Zhang   /* ext->Ydot have already been computed in SNESTSFormFunction_EIMEX (SNES guarantees this) */
2807d5697caSHong Zhang   dmsave = ts->dm;
2817d5697caSHong Zhang   ts->dm = dm;
2829566063dSJacob Faibussowitsch   PetscCall(TSComputeIJacobian(ts, ts->ptime, X, Ydot, ext->shift, A, B, PETSC_TRUE));
2837d5697caSHong Zhang   ts->dm = dmsave;
2849566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, dm, NULL, &Ydot, NULL, NULL));
2857d5697caSHong Zhang   PetscFunctionReturn(0);
2867d5697caSHong Zhang }
2877d5697caSHong Zhang 
2889371c9d4SSatish Balay static PetscErrorCode DMCoarsenHook_TSEIMEX(DM fine, DM coarse, void *ctx) {
2897d5697caSHong Zhang   PetscFunctionBegin;
2907d5697caSHong Zhang   PetscFunctionReturn(0);
2917d5697caSHong Zhang }
2927d5697caSHong Zhang 
2939371c9d4SSatish Balay static PetscErrorCode DMRestrictHook_TSEIMEX(DM fine, Mat restrct, Vec rscale, Mat inject, DM coarse, void *ctx) {
2947d5697caSHong Zhang   TS  ts = (TS)ctx;
2957d5697caSHong Zhang   Vec Z, Z_c;
2967d5697caSHong Zhang 
2977d5697caSHong Zhang   PetscFunctionBegin;
2989566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, fine, &Z, NULL, NULL, NULL));
2999566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, coarse, &Z_c, NULL, NULL, NULL));
3009566063dSJacob Faibussowitsch   PetscCall(MatRestrict(restrct, Z, Z_c));
3019566063dSJacob Faibussowitsch   PetscCall(VecPointwiseMult(Z_c, rscale, Z_c));
3029566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, fine, &Z, NULL, NULL, NULL));
3039566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, coarse, &Z_c, NULL, NULL, NULL));
3047d5697caSHong Zhang   PetscFunctionReturn(0);
3057d5697caSHong Zhang }
3067d5697caSHong Zhang 
3079371c9d4SSatish Balay static PetscErrorCode TSSetUp_EIMEX(TS ts) {
3087d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
3097d5697caSHong Zhang   DM        dm;
3107d5697caSHong Zhang 
311955c1f14SBarry Smith   PetscFunctionBegin;
3127d5697caSHong Zhang   if (!ext->N) { /* ext->max_rows not set */
3139566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetMaxRows(ts, TSEIMEXDefault));
3147d5697caSHong Zhang   }
3157d5697caSHong Zhang   if (-1 == ext->row_ind && -1 == ext->col_ind) {
3169566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetRowCol(ts, ext->max_rows, ext->max_rows));
3177d5697caSHong Zhang   } else { /* ext->row_ind and col_ind already set */
31848a46eb9SPierre Jolivet     if (ext->ord_adapt) PetscCall(PetscInfo(ts, "Order adaptivity is enabled and TSEIMEXSetRowCol or -ts_eimex_row_col option will take no effect\n"));
3197d5697caSHong Zhang   }
3207d5697caSHong Zhang 
3217d5697caSHong Zhang   if (ext->ord_adapt) {
3227d5697caSHong Zhang     ext->nstages = 2; /* Start with the 2-stage scheme */
3239566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetRowCol(ts, ext->nstages, ext->nstages));
3247d5697caSHong Zhang   } else {
3257d5697caSHong Zhang     ext->nstages = ext->max_rows; /* by default nstages is the same as max_rows, this can be changed by setting order adaptivity */
3267d5697caSHong Zhang   }
3277d5697caSHong Zhang 
3289566063dSJacob Faibussowitsch   PetscCall(TSGetAdapt(ts, &ts->adapt));
3292ffb9264SLisandro Dalcin 
3309566063dSJacob Faibussowitsch   PetscCall(VecDuplicateVecs(ts->vec_sol, (1 + ext->nstages) * ext->nstages / 2, &ext->T)); /* full T table */
3319566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->YdotI));
3329566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->YdotRHS));
3339566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Ydot));
3349566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->VecSolPrev));
3359566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Y));
3369566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Z));
3379566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
3381baa6e33SBarry Smith   if (dm) PetscCall(DMCoarsenHookAdd(dm, DMCoarsenHook_TSEIMEX, DMRestrictHook_TSEIMEX, ts));
3397d5697caSHong Zhang   PetscFunctionReturn(0);
3407d5697caSHong Zhang }
3417d5697caSHong Zhang 
3429371c9d4SSatish Balay static PetscErrorCode TSSetFromOptions_EIMEX(TS ts, PetscOptionItems *PetscOptionsObject) {
3437d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
3449200755eSBarry Smith   PetscInt  tindex[2];
3457d5697caSHong Zhang   PetscInt  np = 2, nrows = TSEIMEXDefault;
3469200755eSBarry Smith 
3477d5697caSHong Zhang   PetscFunctionBegin;
3489200755eSBarry Smith   tindex[0] = TSEIMEXDefault;
3499200755eSBarry Smith   tindex[1] = TSEIMEXDefault;
350d0609cedSBarry Smith   PetscOptionsHeadBegin(PetscOptionsObject, "EIMEX ODE solver options");
3517d5697caSHong Zhang   {
3527d5697caSHong Zhang     PetscBool flg;
3539566063dSJacob Faibussowitsch     PetscCall(PetscOptionsInt("-ts_eimex_max_rows", "Define the maximum number of rows used", "TSEIMEXSetMaxRows", nrows, &nrows, &flg)); /* default value 3 */
3541baa6e33SBarry Smith     if (flg) PetscCall(TSEIMEXSetMaxRows(ts, nrows));
3559566063dSJacob Faibussowitsch     PetscCall(PetscOptionsIntArray("-ts_eimex_row_col", "Return the specific term in the T table", "TSEIMEXSetRowCol", tindex, &np, &flg));
35648a46eb9SPierre Jolivet     if (flg) PetscCall(TSEIMEXSetRowCol(ts, tindex[0], tindex[1]));
3579566063dSJacob Faibussowitsch     PetscCall(PetscOptionsBool("-ts_eimex_order_adapt", "Solve the problem with adaptive order", "TSEIMEXSetOrdAdapt", ext->ord_adapt, &ext->ord_adapt, NULL));
3587d5697caSHong Zhang   }
359d0609cedSBarry Smith   PetscOptionsHeadEnd();
3607d5697caSHong Zhang   PetscFunctionReturn(0);
3617d5697caSHong Zhang }
3627d5697caSHong Zhang 
3639371c9d4SSatish Balay static PetscErrorCode TSView_EIMEX(TS ts, PetscViewer viewer) {
3647d5697caSHong Zhang   PetscFunctionBegin;
3657d5697caSHong Zhang   PetscFunctionReturn(0);
3667d5697caSHong Zhang }
3677d5697caSHong Zhang 
3687d5697caSHong Zhang /*@C
3697d5697caSHong Zhang   TSEIMEXSetMaxRows - Set the maximum number of rows for EIMEX schemes
3707d5697caSHong Zhang 
3717d5697caSHong Zhang   Logically collective
3727d5697caSHong Zhang 
373d8d19677SJose E. Roman   Input Parameters:
3747d5697caSHong Zhang +  ts - timestepping context
3757d5697caSHong Zhang -  nrows - maximum number of rows
3767d5697caSHong Zhang 
3777d5697caSHong Zhang   Level: intermediate
3787d5697caSHong Zhang 
379db781477SPatrick Sanan .seealso: `TSEIMEXSetRowCol()`, `TSEIMEXSetOrdAdapt()`, `TSEIMEX`
3807d5697caSHong Zhang @*/
3819371c9d4SSatish Balay PetscErrorCode TSEIMEXSetMaxRows(TS ts, PetscInt nrows) {
3827d5697caSHong Zhang   PetscFunctionBegin;
3837d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
384cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetMaxRows_C", (TS, PetscInt), (ts, nrows));
3857d5697caSHong Zhang   PetscFunctionReturn(0);
3867d5697caSHong Zhang }
3877d5697caSHong Zhang 
3887d5697caSHong Zhang /*@C
3897d5697caSHong Zhang   TSEIMEXSetRowCol - Set the type index in the T table for the return value
3907d5697caSHong Zhang 
3917d5697caSHong Zhang   Logically collective
3927d5697caSHong Zhang 
393d8d19677SJose E. Roman   Input Parameters:
3947d5697caSHong Zhang +  ts - timestepping context
3957d5697caSHong Zhang -  tindex - index in the T table
3967d5697caSHong Zhang 
3977d5697caSHong Zhang   Level: intermediate
3987d5697caSHong Zhang 
399db781477SPatrick Sanan .seealso: `TSEIMEXSetMaxRows()`, `TSEIMEXSetOrdAdapt()`, `TSEIMEX`
4007d5697caSHong Zhang @*/
4019371c9d4SSatish Balay PetscErrorCode TSEIMEXSetRowCol(TS ts, PetscInt row, PetscInt col) {
4027d5697caSHong Zhang   PetscFunctionBegin;
4037d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
404cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetRowCol_C", (TS, PetscInt, PetscInt), (ts, row, col));
4057d5697caSHong Zhang   PetscFunctionReturn(0);
4067d5697caSHong Zhang }
4077d5697caSHong Zhang 
4087d5697caSHong Zhang /*@C
4097d5697caSHong Zhang   TSEIMEXSetOrdAdapt - Set the order adaptativity
4107d5697caSHong Zhang 
4117d5697caSHong Zhang   Logically collective
4127d5697caSHong Zhang 
413d8d19677SJose E. Roman   Input Parameters:
4147d5697caSHong Zhang +  ts - timestepping context
4157d5697caSHong Zhang -  tindex - index in the T table
4167d5697caSHong Zhang 
4177d5697caSHong Zhang   Level: intermediate
4187d5697caSHong Zhang 
419db781477SPatrick Sanan .seealso: `TSEIMEXSetRowCol()`, `TSEIMEXSetOrdAdapt()`, `TSEIMEX`
4207d5697caSHong Zhang @*/
4219371c9d4SSatish Balay PetscErrorCode TSEIMEXSetOrdAdapt(TS ts, PetscBool flg) {
4227d5697caSHong Zhang   PetscFunctionBegin;
4237d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
424cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetOrdAdapt_C", (TS, PetscBool), (ts, flg));
4257d5697caSHong Zhang   PetscFunctionReturn(0);
4267d5697caSHong Zhang }
4277d5697caSHong Zhang 
4289371c9d4SSatish Balay static PetscErrorCode TSEIMEXSetMaxRows_EIMEX(TS ts, PetscInt nrows) {
4297d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
4307d5697caSHong Zhang   PetscInt  i;
4317d5697caSHong Zhang 
4327d5697caSHong Zhang   PetscFunctionBegin;
43363a3b9bcSJacob Faibussowitsch   PetscCheck(nrows >= 0 && nrows <= 100, ((PetscObject)ts)->comm, PETSC_ERR_ARG_OUTOFRANGE, "Max number of rows (current value %" PetscInt_FMT ") should be an integer number between 1 and 100", nrows);
4349566063dSJacob Faibussowitsch   PetscCall(PetscFree(ext->N));
4357d5697caSHong Zhang   ext->max_rows = nrows;
4369566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nrows, &ext->N));
4377d5697caSHong Zhang   for (i = 0; i < nrows; i++) ext->N[i] = i + 1;
4387d5697caSHong Zhang   PetscFunctionReturn(0);
4397d5697caSHong Zhang }
4407d5697caSHong Zhang 
4419371c9d4SSatish Balay static PetscErrorCode TSEIMEXSetRowCol_EIMEX(TS ts, PetscInt row, PetscInt col) {
4427d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
4437d5697caSHong Zhang 
4447d5697caSHong Zhang   PetscFunctionBegin;
44563a3b9bcSJacob Faibussowitsch   PetscCheck(row >= 1 && col >= 1, ((PetscObject)ts)->comm, PETSC_ERR_ARG_OUTOFRANGE, "The row or column index (current value %" PetscInt_FMT ",%" PetscInt_FMT ") should not be less than 1 ", row, col);
4469371c9d4SSatish Balay   PetscCheck(row <= ext->max_rows && col <= ext->max_rows, ((PetscObject)ts)->comm, PETSC_ERR_ARG_OUTOFRANGE, "The row or column index (current value %" PetscInt_FMT ",%" PetscInt_FMT ") exceeds the maximum number of rows %" PetscInt_FMT, row, col,
4479371c9d4SSatish Balay              ext->max_rows);
44863a3b9bcSJacob Faibussowitsch   PetscCheck(col <= row, ((PetscObject)ts)->comm, PETSC_ERR_ARG_OUTOFRANGE, "The column index (%" PetscInt_FMT ") exceeds the row index (%" PetscInt_FMT ")", col, row);
4497d5697caSHong Zhang 
4507d5697caSHong Zhang   ext->row_ind = row - 1;
4517d5697caSHong Zhang   ext->col_ind = col - 1; /* Array index in C starts from 0 */
4527d5697caSHong Zhang   PetscFunctionReturn(0);
4537d5697caSHong Zhang }
4547d5697caSHong Zhang 
4559371c9d4SSatish Balay static PetscErrorCode TSEIMEXSetOrdAdapt_EIMEX(TS ts, PetscBool flg) {
4567d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
4577d5697caSHong Zhang   PetscFunctionBegin;
4587d5697caSHong Zhang   ext->ord_adapt = flg;
4597d5697caSHong Zhang   PetscFunctionReturn(0);
4607d5697caSHong Zhang }
4617d5697caSHong Zhang 
4627d5697caSHong Zhang /*MC
463679f3ef1SBarry Smith       TSEIMEX - Time stepping with Extrapolated IMEX methods.
464679f3ef1SBarry Smith 
465679f3ef1SBarry Smith    These methods are intended for problems with well-separated time scales, especially when a slow scale is strongly nonlinear such that it
466679f3ef1SBarry Smith    is expensive to solve with a fully implicit method. The user should provide the stiff part of the equation using TSSetIFunction() and the
467679f3ef1SBarry Smith    non-stiff part with TSSetRHSFunction().
4687d5697caSHong Zhang 
4697d5697caSHong Zhang    Notes:
4707d5697caSHong Zhang   The default is a 3-stage scheme, it can be changed with TSEIMEXSetMaxRows() or -ts_eimex_max_rows
4717d5697caSHong Zhang 
4727d5697caSHong Zhang   This method currently only works with ODE, for which the stiff part G(t,X,Xdot) has the form Xdot + Ghat(t,X).
4737d5697caSHong Zhang 
474679f3ef1SBarry Smith   The general system is written as
475679f3ef1SBarry Smith 
476679f3ef1SBarry Smith   G(t,X,Xdot) = F(t,X)
477679f3ef1SBarry Smith 
478679f3ef1SBarry Smith   where G represents the stiff part and F represents the non-stiff part. The user should provide the stiff part
479679f3ef1SBarry Smith   of the equation using TSSetIFunction() and the non-stiff part with TSSetRHSFunction().
480679f3ef1SBarry Smith   This method is designed to be linearly implicit on G and can use an approximate and lagged Jacobian.
481679f3ef1SBarry Smith 
482679f3ef1SBarry Smith   Another common form for the system is
483679f3ef1SBarry Smith 
484679f3ef1SBarry Smith   y'=f(x)+g(x)
485679f3ef1SBarry Smith 
486679f3ef1SBarry Smith   The relationship between F,G and f,g is
487679f3ef1SBarry Smith 
488679f3ef1SBarry Smith   G = y'-g(x), F = f(x)
489679f3ef1SBarry Smith 
490679f3ef1SBarry Smith  References
491679f3ef1SBarry Smith   E. Constantinescu and A. Sandu, Extrapolated implicit-explicit time stepping, SIAM Journal on Scientific
492679f3ef1SBarry Smith Computing, 31 (2010), pp. 4452-4477.
493679f3ef1SBarry Smith 
4947d5697caSHong Zhang       Level: beginner
4957d5697caSHong Zhang 
496db781477SPatrick Sanan .seealso: `TSCreate()`, `TS`, `TSSetType()`, `TSEIMEXSetMaxRows()`, `TSEIMEXSetRowCol()`, `TSEIMEXSetOrdAdapt()`
497679f3ef1SBarry Smith 
4987d5697caSHong Zhang  M*/
4999371c9d4SSatish Balay PETSC_EXTERN PetscErrorCode TSCreate_EIMEX(TS ts) {
5007d5697caSHong Zhang   TS_EIMEX *ext;
5017d5697caSHong Zhang 
5027d5697caSHong Zhang   PetscFunctionBegin;
5037d5697caSHong Zhang 
5047d5697caSHong Zhang   ts->ops->reset          = TSReset_EIMEX;
5057d5697caSHong Zhang   ts->ops->destroy        = TSDestroy_EIMEX;
5067d5697caSHong Zhang   ts->ops->view           = TSView_EIMEX;
5077d5697caSHong Zhang   ts->ops->setup          = TSSetUp_EIMEX;
5087d5697caSHong Zhang   ts->ops->step           = TSStep_EIMEX;
5097d5697caSHong Zhang   ts->ops->interpolate    = TSInterpolate_EIMEX;
5107d5697caSHong Zhang   ts->ops->evaluatestep   = TSEvaluateStep_EIMEX;
5117d5697caSHong Zhang   ts->ops->setfromoptions = TSSetFromOptions_EIMEX;
5127d5697caSHong Zhang   ts->ops->snesfunction   = SNESTSFormFunction_EIMEX;
5137d5697caSHong Zhang   ts->ops->snesjacobian   = SNESTSFormJacobian_EIMEX;
5142ffb9264SLisandro Dalcin   ts->default_adapt_type  = TSADAPTNONE;
5157d5697caSHong Zhang 
516efd4aadfSBarry Smith   ts->usessnes = PETSC_TRUE;
517efd4aadfSBarry Smith 
518*4dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&ext));
5197d5697caSHong Zhang   ts->data = (void *)ext;
5207d5697caSHong Zhang 
5217d5697caSHong Zhang   ext->ord_adapt = PETSC_FALSE; /* By default, no order adapativity */
5227d5697caSHong Zhang   ext->row_ind   = -1;
5237d5697caSHong Zhang   ext->col_ind   = -1;
5247d5697caSHong Zhang   ext->max_rows  = TSEIMEXDefault;
5257d5697caSHong Zhang   ext->nstages   = TSEIMEXDefault;
5267d5697caSHong Zhang 
5279566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetMaxRows_C", TSEIMEXSetMaxRows_EIMEX));
5289566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetRowCol_C", TSEIMEXSetRowCol_EIMEX));
5299566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetOrdAdapt_C", TSEIMEXSetOrdAdapt_EIMEX));
5307d5697caSHong Zhang   PetscFunctionReturn(0);
5317d5697caSHong Zhang }
532