xref: /petsc/src/ts/impls/eimex/eimex.c (revision 4d86920da9ee67c835173a5dfffa1b3a52fd24ca)
1af0996ceSBarry Smith #include <petsc/private/tsimpl.h> /*I   "petscts.h"   I*/
24042b796SJed Brown #include <petscdm.h>
37d5697caSHong Zhang 
47d5697caSHong Zhang static const PetscInt TSEIMEXDefault = 3;
57d5697caSHong Zhang 
67d5697caSHong Zhang typedef struct {
77d5697caSHong Zhang   PetscInt     row_ind;    /* Return the term T[row_ind][col_ind] */
87d5697caSHong Zhang   PetscInt     col_ind;    /* Return the term T[row_ind][col_ind] */
97d5697caSHong Zhang   PetscInt     nstages;    /* Numbers of stages in current scheme */
107d5697caSHong Zhang   PetscInt     max_rows;   /* Maximum number of rows */
117d5697caSHong Zhang   PetscInt    *N;          /* Harmonic sequence N[max_rows] */
127d5697caSHong Zhang   Vec          Y;          /* States computed during the step, used to complete the step */
137d5697caSHong Zhang   Vec          Z;          /* For shift*(Y-Z) */
147d5697caSHong Zhang   Vec         *T;          /* Working table, size determined by nstages */
157d5697caSHong Zhang   Vec          YdotRHS;    /* f(x) Work vector holding YdotRHS during residual evaluation */
167d5697caSHong Zhang   Vec          YdotI;      /* xdot-g(x) Work vector holding YdotI = G(t,x,xdot) when xdot =0 */
177d5697caSHong Zhang   Vec          Ydot;       /* f(x)+g(x) Work vector */
187d5697caSHong Zhang   Vec          VecSolPrev; /* Work vector holding the solution from the previous step (used for interpolation) */
197d5697caSHong Zhang   PetscReal    shift;
207d5697caSHong Zhang   PetscReal    ctime;
217d5697caSHong Zhang   PetscBool    recompute_jacobian; /* Recompute the Jacobian at each stage, default is to freeze the Jacobian at the start of each step */
227d5697caSHong Zhang   PetscBool    ord_adapt;          /* order adapativity */
237d5697caSHong Zhang   TSStepStatus status;
247d5697caSHong Zhang } TS_EIMEX;
257d5697caSHong Zhang 
267d5697caSHong Zhang /* This function is pure */
27d71ae5a4SJacob Faibussowitsch static PetscInt Map(PetscInt i, PetscInt j, PetscInt s)
28d71ae5a4SJacob Faibussowitsch {
297d5697caSHong Zhang   return ((2 * s - j + 1) * j / 2 + i - j);
307d5697caSHong Zhang }
317d5697caSHong Zhang 
32d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEvaluateStep_EIMEX(TS ts, PetscInt order, Vec X, PetscBool *done)
33d71ae5a4SJacob Faibussowitsch {
347d5697caSHong Zhang   TS_EIMEX      *ext = (TS_EIMEX *)ts->data;
357d5697caSHong Zhang   const PetscInt ns  = ext->nstages;
36*4d86920dSPierre Jolivet 
377d5697caSHong Zhang   PetscFunctionBegin;
389566063dSJacob Faibussowitsch   PetscCall(VecCopy(ext->T[Map(ext->row_ind, ext->col_ind, ns)], X));
393ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
407d5697caSHong Zhang }
417d5697caSHong Zhang 
42d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSStage_EIMEX(TS ts, PetscInt istage)
43d71ae5a4SJacob Faibussowitsch {
447d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
457d5697caSHong Zhang   PetscReal h;
467d5697caSHong Zhang   Vec       Y = ext->Y, Z = ext->Z;
477d5697caSHong Zhang   SNES      snes;
487d5697caSHong Zhang   TSAdapt   adapt;
497d5697caSHong Zhang   PetscInt  i, its, lits;
507d5697caSHong Zhang   PetscBool accept;
517d5697caSHong Zhang 
527d5697caSHong Zhang   PetscFunctionBegin;
539566063dSJacob Faibussowitsch   PetscCall(TSGetSNES(ts, &snes));
547d5697caSHong Zhang   h          = ts->time_step / ext->N[istage]; /* step size for the istage-th stage */
557d5697caSHong Zhang   ext->shift = 1. / h;
569566063dSJacob Faibussowitsch   PetscCall(SNESSetLagJacobian(snes, -2)); /* Recompute the Jacobian on this solve, but not again */
579566063dSJacob Faibussowitsch   PetscCall(VecCopy(ext->VecSolPrev, Y));  /* Take the previous solution as initial step */
587d5697caSHong Zhang 
597d5697caSHong Zhang   for (i = 0; i < ext->N[istage]; i++) {
607d5697caSHong Zhang     ext->ctime = ts->ptime + h * i;
619566063dSJacob Faibussowitsch     PetscCall(VecCopy(Y, Z)); /* Save the solution of the previous substep */
629566063dSJacob Faibussowitsch     PetscCall(SNESSolve(snes, NULL, Y));
639566063dSJacob Faibussowitsch     PetscCall(SNESGetIterationNumber(snes, &its));
649566063dSJacob Faibussowitsch     PetscCall(SNESGetLinearSolveIterations(snes, &lits));
659371c9d4SSatish Balay     ts->snes_its += its;
669371c9d4SSatish Balay     ts->ksp_its += lits;
679566063dSJacob Faibussowitsch     PetscCall(TSGetAdapt(ts, &adapt));
689566063dSJacob Faibussowitsch     PetscCall(TSAdaptCheckStage(adapt, ts, ext->ctime, Y, &accept));
697d5697caSHong Zhang   }
703ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
717d5697caSHong Zhang }
727d5697caSHong Zhang 
73d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSStep_EIMEX(TS ts)
74d71ae5a4SJacob Faibussowitsch {
757d5697caSHong Zhang   TS_EIMEX      *ext = (TS_EIMEX *)ts->data;
767d5697caSHong Zhang   const PetscInt ns  = ext->nstages;
777d5697caSHong Zhang   Vec           *T = ext->T, Y = ext->Y;
787d5697caSHong Zhang   SNES           snes;
797d5697caSHong Zhang   PetscInt       i, j;
807d5697caSHong Zhang   PetscBool      accept = PETSC_FALSE;
817453f775SEmil Constantinescu   PetscReal      alpha, local_error, local_error_a, local_error_r;
827d5697caSHong Zhang 
83d0609cedSBarry Smith   PetscFunctionBegin;
849566063dSJacob Faibussowitsch   PetscCall(TSGetSNES(ts, &snes));
859566063dSJacob Faibussowitsch   PetscCall(SNESSetType(snes, "ksponly"));
867d5697caSHong Zhang   ext->status = TS_STEP_INCOMPLETE;
877d5697caSHong Zhang 
889566063dSJacob Faibussowitsch   PetscCall(VecCopy(ts->vec_sol, ext->VecSolPrev));
897d5697caSHong Zhang 
907d5697caSHong Zhang   /* Apply n_j steps of the base method to obtain solutions of T(j,1),1<=j<=s */
917d5697caSHong Zhang   for (j = 0; j < ns; j++) {
929566063dSJacob Faibussowitsch     PetscCall(TSStage_EIMEX(ts, j));
939566063dSJacob Faibussowitsch     PetscCall(VecCopy(Y, T[j]));
947d5697caSHong Zhang   }
957d5697caSHong Zhang 
967d5697caSHong Zhang   for (i = 1; i < ns; i++) {
977d5697caSHong Zhang     for (j = i; j < ns; j++) {
987d5697caSHong Zhang       alpha = -(PetscReal)ext->N[j] / ext->N[j - i];
99d0609cedSBarry 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] */
1007d5697caSHong Zhang       alpha = 1.0 / (1.0 + alpha);
1019566063dSJacob Faibussowitsch       PetscCall(VecScale(T[Map(j, i, ns)], alpha));
1027d5697caSHong Zhang     }
1037d5697caSHong Zhang   }
1047d5697caSHong Zhang 
1059566063dSJacob Faibussowitsch   PetscCall(TSEvaluateStep(ts, ns, ts->vec_sol, NULL)); /*update ts solution */
1067d5697caSHong Zhang 
1077d5697caSHong Zhang   if (ext->ord_adapt && ext->nstages < ext->max_rows) {
1087d5697caSHong Zhang     accept = PETSC_FALSE;
1097d5697caSHong Zhang     while (!accept && ext->nstages < ext->max_rows) {
1109566063dSJacob 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));
1117d5697caSHong Zhang       accept = (local_error < 1.0) ? PETSC_TRUE : PETSC_FALSE;
1127d5697caSHong Zhang 
1137d5697caSHong Zhang       if (!accept) { /* add one more stage*/
1149566063dSJacob Faibussowitsch         PetscCall(TSStage_EIMEX(ts, ext->nstages));
1159371c9d4SSatish Balay         ext->nstages++;
1169371c9d4SSatish Balay         ext->row_ind++;
1179371c9d4SSatish Balay         ext->col_ind++;
1187d5697caSHong Zhang         /*T table need to be recycled*/
1199566063dSJacob Faibussowitsch         PetscCall(VecDuplicateVecs(ts->vec_sol, (1 + ext->nstages) * ext->nstages / 2, &ext->T));
1207d5697caSHong Zhang         for (i = 0; i < ext->nstages - 1; i++) {
12148a46eb9SPierre Jolivet           for (j = 0; j <= i; j++) PetscCall(VecCopy(T[Map(i, j, ext->nstages - 1)], ext->T[Map(i, j, ext->nstages)]));
1227d5697caSHong Zhang         }
1239566063dSJacob Faibussowitsch         PetscCall(VecDestroyVecs(ext->nstages * (ext->nstages - 1) / 2, &T));
1247d5697caSHong Zhang         T = ext->T; /*reset the pointer*/
1257d5697caSHong Zhang         /*recycling finished, store the new solution*/
1269566063dSJacob Faibussowitsch         PetscCall(VecCopy(Y, T[ext->nstages - 1]));
1277d5697caSHong Zhang         /*extrapolation for the newly added stage*/
1287d5697caSHong Zhang         for (i = 1; i < ext->nstages; i++) {
1297d5697caSHong Zhang           alpha = -(PetscReal)ext->N[ext->nstages - 1] / ext->N[ext->nstages - 1 - i];
130d0609cedSBarry 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]*/
1317d5697caSHong Zhang           alpha = 1.0 / (1.0 + alpha);
1329566063dSJacob Faibussowitsch           PetscCall(VecScale(T[Map(ext->nstages - 1, i, ext->nstages)], alpha));
1337d5697caSHong Zhang         }
1347d5697caSHong Zhang         /*update ts solution */
1359566063dSJacob Faibussowitsch         PetscCall(TSEvaluateStep(ts, ext->nstages, ts->vec_sol, NULL));
1367d5697caSHong Zhang       } /*end if !accept*/
1377d5697caSHong Zhang     }   /*end while*/
1387d5697caSHong Zhang 
13948a46eb9SPierre Jolivet     if (ext->nstages == ext->max_rows) PetscCall(PetscInfo(ts, "Max number of rows has been used\n"));
1407d5697caSHong Zhang   } /*end if ext->ord_adapt*/
1417d5697caSHong Zhang   ts->ptime += ts->time_step;
1427d5697caSHong Zhang   ext->status = TS_STEP_COMPLETE;
1437d5697caSHong Zhang 
1447d5697caSHong Zhang   if (ext->status != TS_STEP_COMPLETE && !ts->reason) ts->reason = TS_DIVERGED_STEP_REJECTED;
1453ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1467d5697caSHong Zhang }
1477d5697caSHong Zhang 
1487d5697caSHong Zhang /* cubic Hermit spline */
149d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSInterpolate_EIMEX(TS ts, PetscReal itime, Vec X)
150d71ae5a4SJacob Faibussowitsch {
1517d5697caSHong Zhang   TS_EIMEX       *ext = (TS_EIMEX *)ts->data;
1527d5697caSHong Zhang   PetscReal       t, a, b;
153be5899b3SLisandro Dalcin   Vec             Y0 = ext->VecSolPrev, Y1 = ext->Y, Ydot = ext->Ydot, YdotI = ext->YdotI;
154be5899b3SLisandro Dalcin   const PetscReal h = ts->ptime - ts->ptime_prev;
155*4d86920dSPierre Jolivet 
1567d5697caSHong Zhang   PetscFunctionBegin;
1577d5697caSHong Zhang   t = (itime - ts->ptime + h) / h;
1587d5697caSHong Zhang   /* YdotI = -f(x)-g(x) */
1597d5697caSHong Zhang 
1609566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(Ydot));
1619566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ts->ptime - h, Y0, Ydot, YdotI, PETSC_FALSE));
1627d5697caSHong Zhang 
16316dc5bacSEmil Constantinescu   a = 2.0 * t * t * t - 3.0 * t * t + 1.0;
16416dc5bacSEmil Constantinescu   b = -(t * t * t - 2.0 * t * t + t) * h;
1659566063dSJacob Faibussowitsch   PetscCall(VecAXPBYPCZ(X, a, b, 0.0, Y0, YdotI));
1667d5697caSHong Zhang 
1679566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ts->ptime, Y1, Ydot, YdotI, PETSC_FALSE));
16816dc5bacSEmil Constantinescu   a = -2.0 * t * t * t + 3.0 * t * t;
16916dc5bacSEmil Constantinescu   b = -(t * t * t - t * t) * h;
1709566063dSJacob Faibussowitsch   PetscCall(VecAXPBYPCZ(X, a, b, 1.0, Y1, YdotI));
1713ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1727d5697caSHong Zhang }
1737d5697caSHong Zhang 
174d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSReset_EIMEX(TS ts)
175d71ae5a4SJacob Faibussowitsch {
1767d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
1777d5697caSHong Zhang   PetscInt  ns;
1787d5697caSHong Zhang 
1797d5697caSHong Zhang   PetscFunctionBegin;
1807d5697caSHong Zhang   ns = ext->nstages;
1819566063dSJacob Faibussowitsch   PetscCall(VecDestroyVecs((1 + ns) * ns / 2, &ext->T));
1829566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Y));
1839566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Z));
1849566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->YdotRHS));
1859566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->YdotI));
1869566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->Ydot));
1879566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&ext->VecSolPrev));
1889566063dSJacob Faibussowitsch   PetscCall(PetscFree(ext->N));
1893ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1907d5697caSHong Zhang }
1917d5697caSHong Zhang 
192d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSDestroy_EIMEX(TS ts)
193d71ae5a4SJacob Faibussowitsch {
1947d5697caSHong Zhang   PetscFunctionBegin;
1959566063dSJacob Faibussowitsch   PetscCall(TSReset_EIMEX(ts));
1969566063dSJacob Faibussowitsch   PetscCall(PetscFree(ts->data));
1979566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetMaxRows_C", NULL));
1989566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetRowCol_C", NULL));
1999566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetOrdAdapt_C", NULL));
2003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2017d5697caSHong Zhang }
2027d5697caSHong Zhang 
203d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEIMEXGetVecs(TS ts, DM dm, Vec *Z, Vec *Ydot, Vec *YdotI, Vec *YdotRHS)
204d71ae5a4SJacob Faibussowitsch {
2057d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
2067d5697caSHong Zhang 
2077d5697caSHong Zhang   PetscFunctionBegin;
2087d5697caSHong Zhang   if (Z) {
2097d5697caSHong Zhang     if (dm && dm != ts->dm) {
2109566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_Z", Z));
2117d5697caSHong Zhang     } else *Z = ext->Z;
2127d5697caSHong Zhang   }
2137d5697caSHong Zhang   if (Ydot) {
2147d5697caSHong Zhang     if (dm && dm != ts->dm) {
2159566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_Ydot", Ydot));
2167d5697caSHong Zhang     } else *Ydot = ext->Ydot;
2177d5697caSHong Zhang   }
2187d5697caSHong Zhang   if (YdotI) {
2197d5697caSHong Zhang     if (dm && dm != ts->dm) {
2209566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_YdotI", YdotI));
2217d5697caSHong Zhang     } else *YdotI = ext->YdotI;
2227d5697caSHong Zhang   }
2237d5697caSHong Zhang   if (YdotRHS) {
2247d5697caSHong Zhang     if (dm && dm != ts->dm) {
2259566063dSJacob Faibussowitsch       PetscCall(DMGetNamedGlobalVector(dm, "TSEIMEX_YdotRHS", YdotRHS));
2267d5697caSHong Zhang     } else *YdotRHS = ext->YdotRHS;
2277d5697caSHong Zhang   }
2283ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2297d5697caSHong Zhang }
2307d5697caSHong Zhang 
231d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEIMEXRestoreVecs(TS ts, DM dm, Vec *Z, Vec *Ydot, Vec *YdotI, Vec *YdotRHS)
232d71ae5a4SJacob Faibussowitsch {
2337d5697caSHong Zhang   PetscFunctionBegin;
2347d5697caSHong Zhang   if (Z) {
23548a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_Z", Z));
2367d5697caSHong Zhang   }
2377d5697caSHong Zhang   if (Ydot) {
23848a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_Ydot", Ydot));
2397d5697caSHong Zhang   }
2407d5697caSHong Zhang   if (YdotI) {
24148a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_YdotI", YdotI));
2427d5697caSHong Zhang   }
2437d5697caSHong Zhang   if (YdotRHS) {
24448a46eb9SPierre Jolivet     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSEIMEX_YdotRHS", YdotRHS));
2457d5697caSHong Zhang   }
2463ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2477d5697caSHong Zhang }
2487d5697caSHong Zhang 
2497d5697caSHong Zhang /*
2507d5697caSHong Zhang   This defines the nonlinear equation that is to be solved with SNES
2517d5697caSHong Zhang   Fn[t0+Theta*dt, U, (U-U0)*shift] = 0
2527d5697caSHong Zhang   In the case of Backward Euler, Fn = (U-U0)/h-g(t1,U))
2537d5697caSHong Zhang   Since FormIFunction calculates G = ydot - g(t,y), ydot will be set to (U-U0)/h
2547d5697caSHong Zhang */
255d71ae5a4SJacob Faibussowitsch static PetscErrorCode SNESTSFormFunction_EIMEX(SNES snes, Vec X, Vec G, TS ts)
256d71ae5a4SJacob Faibussowitsch {
2577d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
2587d5697caSHong Zhang   Vec       Ydot, Z;
2597d5697caSHong Zhang   DM        dm, dmsave;
2607d5697caSHong Zhang 
2617d5697caSHong Zhang   PetscFunctionBegin;
2629566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(G));
2637d5697caSHong Zhang 
2649566063dSJacob Faibussowitsch   PetscCall(SNESGetDM(snes, &dm));
2659566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, dm, &Z, &Ydot, NULL, NULL));
2669566063dSJacob Faibussowitsch   PetscCall(VecZeroEntries(Ydot));
2677d5697caSHong Zhang   dmsave = ts->dm;
2687d5697caSHong Zhang   ts->dm = dm;
2699566063dSJacob Faibussowitsch   PetscCall(TSComputeIFunction(ts, ext->ctime, X, Ydot, G, PETSC_FALSE));
2707d5697caSHong Zhang   /* PETSC_FALSE indicates non-imex, adding explicit RHS to the implicit I function.  */
2719566063dSJacob Faibussowitsch   PetscCall(VecCopy(G, Ydot));
2727d5697caSHong Zhang   ts->dm = dmsave;
2739566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, dm, &Z, &Ydot, NULL, NULL));
2743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2757d5697caSHong Zhang }
2767d5697caSHong Zhang 
2777d5697caSHong Zhang /*
2787d5697caSHong Zhang  This defined the Jacobian matrix for SNES. Jn = (I/h-g'(t,y))
2797d5697caSHong Zhang  */
280d71ae5a4SJacob Faibussowitsch static PetscErrorCode SNESTSFormJacobian_EIMEX(SNES snes, Vec X, Mat A, Mat B, TS ts)
281d71ae5a4SJacob Faibussowitsch {
2827d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
2837d5697caSHong Zhang   Vec       Ydot;
2847d5697caSHong Zhang   DM        dm, dmsave;
285*4d86920dSPierre Jolivet 
2867d5697caSHong Zhang   PetscFunctionBegin;
2879566063dSJacob Faibussowitsch   PetscCall(SNESGetDM(snes, &dm));
2889566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, dm, NULL, &Ydot, NULL, NULL));
2899566063dSJacob Faibussowitsch   /*  PetscCall(VecZeroEntries(Ydot)); */
2907d5697caSHong Zhang   /* ext->Ydot have already been computed in SNESTSFormFunction_EIMEX (SNES guarantees this) */
2917d5697caSHong Zhang   dmsave = ts->dm;
2927d5697caSHong Zhang   ts->dm = dm;
2939566063dSJacob Faibussowitsch   PetscCall(TSComputeIJacobian(ts, ts->ptime, X, Ydot, ext->shift, A, B, PETSC_TRUE));
2947d5697caSHong Zhang   ts->dm = dmsave;
2959566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, dm, NULL, &Ydot, NULL, NULL));
2963ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2977d5697caSHong Zhang }
2987d5697caSHong Zhang 
299d71ae5a4SJacob Faibussowitsch static PetscErrorCode DMCoarsenHook_TSEIMEX(DM fine, DM coarse, void *ctx)
300d71ae5a4SJacob Faibussowitsch {
3017d5697caSHong Zhang   PetscFunctionBegin;
3023ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3037d5697caSHong Zhang }
3047d5697caSHong Zhang 
305d71ae5a4SJacob Faibussowitsch static PetscErrorCode DMRestrictHook_TSEIMEX(DM fine, Mat restrct, Vec rscale, Mat inject, DM coarse, void *ctx)
306d71ae5a4SJacob Faibussowitsch {
3077d5697caSHong Zhang   TS  ts = (TS)ctx;
3087d5697caSHong Zhang   Vec Z, Z_c;
3097d5697caSHong Zhang 
3107d5697caSHong Zhang   PetscFunctionBegin;
3119566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, fine, &Z, NULL, NULL, NULL));
3129566063dSJacob Faibussowitsch   PetscCall(TSEIMEXGetVecs(ts, coarse, &Z_c, NULL, NULL, NULL));
3139566063dSJacob Faibussowitsch   PetscCall(MatRestrict(restrct, Z, Z_c));
3149566063dSJacob Faibussowitsch   PetscCall(VecPointwiseMult(Z_c, rscale, Z_c));
3159566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, fine, &Z, NULL, NULL, NULL));
3169566063dSJacob Faibussowitsch   PetscCall(TSEIMEXRestoreVecs(ts, coarse, &Z_c, NULL, NULL, NULL));
3173ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3187d5697caSHong Zhang }
3197d5697caSHong Zhang 
320d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSSetUp_EIMEX(TS ts)
321d71ae5a4SJacob Faibussowitsch {
3227d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
3237d5697caSHong Zhang   DM        dm;
3247d5697caSHong Zhang 
325955c1f14SBarry Smith   PetscFunctionBegin;
3267d5697caSHong Zhang   if (!ext->N) { /* ext->max_rows not set */
3279566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetMaxRows(ts, TSEIMEXDefault));
3287d5697caSHong Zhang   }
3297d5697caSHong Zhang   if (-1 == ext->row_ind && -1 == ext->col_ind) {
3309566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetRowCol(ts, ext->max_rows, ext->max_rows));
3317d5697caSHong Zhang   } else { /* ext->row_ind and col_ind already set */
33248a46eb9SPierre 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"));
3337d5697caSHong Zhang   }
3347d5697caSHong Zhang 
3357d5697caSHong Zhang   if (ext->ord_adapt) {
3367d5697caSHong Zhang     ext->nstages = 2; /* Start with the 2-stage scheme */
3379566063dSJacob Faibussowitsch     PetscCall(TSEIMEXSetRowCol(ts, ext->nstages, ext->nstages));
3387d5697caSHong Zhang   } else {
3397d5697caSHong Zhang     ext->nstages = ext->max_rows; /* by default nstages is the same as max_rows, this can be changed by setting order adaptivity */
3407d5697caSHong Zhang   }
3417d5697caSHong Zhang 
3429566063dSJacob Faibussowitsch   PetscCall(TSGetAdapt(ts, &ts->adapt));
3432ffb9264SLisandro Dalcin 
3449566063dSJacob Faibussowitsch   PetscCall(VecDuplicateVecs(ts->vec_sol, (1 + ext->nstages) * ext->nstages / 2, &ext->T)); /* full T table */
3459566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->YdotI));
3469566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->YdotRHS));
3479566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Ydot));
3489566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->VecSolPrev));
3499566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Y));
3509566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol, &ext->Z));
3519566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
3521baa6e33SBarry Smith   if (dm) PetscCall(DMCoarsenHookAdd(dm, DMCoarsenHook_TSEIMEX, DMRestrictHook_TSEIMEX, ts));
3533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3547d5697caSHong Zhang }
3557d5697caSHong Zhang 
356d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSSetFromOptions_EIMEX(TS ts, PetscOptionItems *PetscOptionsObject)
357d71ae5a4SJacob Faibussowitsch {
3587d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
3599200755eSBarry Smith   PetscInt  tindex[2];
3607d5697caSHong Zhang   PetscInt  np = 2, nrows = TSEIMEXDefault;
3619200755eSBarry Smith 
3627d5697caSHong Zhang   PetscFunctionBegin;
3639200755eSBarry Smith   tindex[0] = TSEIMEXDefault;
3649200755eSBarry Smith   tindex[1] = TSEIMEXDefault;
365d0609cedSBarry Smith   PetscOptionsHeadBegin(PetscOptionsObject, "EIMEX ODE solver options");
3667d5697caSHong Zhang   {
3677d5697caSHong Zhang     PetscBool flg;
3689566063dSJacob Faibussowitsch     PetscCall(PetscOptionsInt("-ts_eimex_max_rows", "Define the maximum number of rows used", "TSEIMEXSetMaxRows", nrows, &nrows, &flg)); /* default value 3 */
3691baa6e33SBarry Smith     if (flg) PetscCall(TSEIMEXSetMaxRows(ts, nrows));
3709566063dSJacob Faibussowitsch     PetscCall(PetscOptionsIntArray("-ts_eimex_row_col", "Return the specific term in the T table", "TSEIMEXSetRowCol", tindex, &np, &flg));
37148a46eb9SPierre Jolivet     if (flg) PetscCall(TSEIMEXSetRowCol(ts, tindex[0], tindex[1]));
3729566063dSJacob Faibussowitsch     PetscCall(PetscOptionsBool("-ts_eimex_order_adapt", "Solve the problem with adaptive order", "TSEIMEXSetOrdAdapt", ext->ord_adapt, &ext->ord_adapt, NULL));
3737d5697caSHong Zhang   }
374d0609cedSBarry Smith   PetscOptionsHeadEnd();
3753ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3767d5697caSHong Zhang }
3777d5697caSHong Zhang 
378d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSView_EIMEX(TS ts, PetscViewer viewer)
379d71ae5a4SJacob Faibussowitsch {
3807d5697caSHong Zhang   PetscFunctionBegin;
3813ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3827d5697caSHong Zhang }
3837d5697caSHong Zhang 
3847d5697caSHong Zhang /*@C
385bcf0153eSBarry Smith   TSEIMEXSetMaxRows - Set the maximum number of rows for `TSEIMEX` schemes
3867d5697caSHong Zhang 
38720f4b53cSBarry Smith   Logically Collective
3887d5697caSHong Zhang 
389d8d19677SJose E. Roman   Input Parameters:
3907d5697caSHong Zhang + ts    - timestepping context
3917d5697caSHong Zhang - nrows - maximum number of rows
3927d5697caSHong Zhang 
3937d5697caSHong Zhang   Level: intermediate
3947d5697caSHong Zhang 
3951cc06b55SBarry Smith .seealso: [](ch_ts), `TSEIMEXSetRowCol()`, `TSEIMEXSetOrdAdapt()`, `TSEIMEX`
3967d5697caSHong Zhang @*/
397d71ae5a4SJacob Faibussowitsch PetscErrorCode TSEIMEXSetMaxRows(TS ts, PetscInt nrows)
398d71ae5a4SJacob Faibussowitsch {
3997d5697caSHong Zhang   PetscFunctionBegin;
4007d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
401cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetMaxRows_C", (TS, PetscInt), (ts, nrows));
4023ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4037d5697caSHong Zhang }
4047d5697caSHong Zhang 
4057d5697caSHong Zhang /*@C
4062fe279fdSBarry Smith   TSEIMEXSetRowCol - Set the number of rows and the number of columns for the tableau that represents the T solution in the `TSEIMEX` scheme
4077d5697caSHong Zhang 
40820f4b53cSBarry Smith   Logically Collective
4097d5697caSHong Zhang 
410d8d19677SJose E. Roman   Input Parameters:
4117d5697caSHong Zhang + ts  - timestepping context
4122fe279fdSBarry Smith . row - the row
4132fe279fdSBarry Smith - col - the column
4147d5697caSHong Zhang 
4157d5697caSHong Zhang   Level: intermediate
4167d5697caSHong Zhang 
4171cc06b55SBarry Smith .seealso: [](ch_ts), `TSEIMEXSetMaxRows()`, `TSEIMEXSetOrdAdapt()`, `TSEIMEX`
4187d5697caSHong Zhang @*/
419d71ae5a4SJacob Faibussowitsch PetscErrorCode TSEIMEXSetRowCol(TS ts, PetscInt row, PetscInt col)
420d71ae5a4SJacob Faibussowitsch {
4217d5697caSHong Zhang   PetscFunctionBegin;
4227d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
423cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetRowCol_C", (TS, PetscInt, PetscInt), (ts, row, col));
4243ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4257d5697caSHong Zhang }
4267d5697caSHong Zhang 
4277d5697caSHong Zhang /*@C
428bcf0153eSBarry Smith   TSEIMEXSetOrdAdapt - Set the order adaptativity for the `TSEIMEX` schemes
4297d5697caSHong Zhang 
43020f4b53cSBarry Smith   Logically Collective
4317d5697caSHong Zhang 
432d8d19677SJose E. Roman   Input Parameters:
4337d5697caSHong Zhang + ts  - timestepping context
434b43aa488SJacob Faibussowitsch - flg - index in the T table
4357d5697caSHong Zhang 
4367d5697caSHong Zhang   Level: intermediate
4377d5697caSHong Zhang 
43842747ad1SJacob Faibussowitsch .seealso: [](ch_ts), `TSEIMEXSetRowCol()`, `TSEIMEX`
4397d5697caSHong Zhang @*/
440d71ae5a4SJacob Faibussowitsch PetscErrorCode TSEIMEXSetOrdAdapt(TS ts, PetscBool flg)
441d71ae5a4SJacob Faibussowitsch {
4427d5697caSHong Zhang   PetscFunctionBegin;
4437d5697caSHong Zhang   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
444cac4c232SBarry Smith   PetscTryMethod(ts, "TSEIMEXSetOrdAdapt_C", (TS, PetscBool), (ts, flg));
4453ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4467d5697caSHong Zhang }
4477d5697caSHong Zhang 
448d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEIMEXSetMaxRows_EIMEX(TS ts, PetscInt nrows)
449d71ae5a4SJacob Faibussowitsch {
4507d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
4517d5697caSHong Zhang   PetscInt  i;
4527d5697caSHong Zhang 
4537d5697caSHong Zhang   PetscFunctionBegin;
45463a3b9bcSJacob 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);
4559566063dSJacob Faibussowitsch   PetscCall(PetscFree(ext->N));
4567d5697caSHong Zhang   ext->max_rows = nrows;
4579566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nrows, &ext->N));
4587d5697caSHong Zhang   for (i = 0; i < nrows; i++) ext->N[i] = i + 1;
4593ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4607d5697caSHong Zhang }
4617d5697caSHong Zhang 
462d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEIMEXSetRowCol_EIMEX(TS ts, PetscInt row, PetscInt col)
463d71ae5a4SJacob Faibussowitsch {
4647d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
4657d5697caSHong Zhang 
4667d5697caSHong Zhang   PetscFunctionBegin;
46763a3b9bcSJacob 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);
4689371c9d4SSatish 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,
4699371c9d4SSatish Balay              ext->max_rows);
47063a3b9bcSJacob Faibussowitsch   PetscCheck(col <= row, ((PetscObject)ts)->comm, PETSC_ERR_ARG_OUTOFRANGE, "The column index (%" PetscInt_FMT ") exceeds the row index (%" PetscInt_FMT ")", col, row);
4717d5697caSHong Zhang 
4727d5697caSHong Zhang   ext->row_ind = row - 1;
4737d5697caSHong Zhang   ext->col_ind = col - 1; /* Array index in C starts from 0 */
4743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4757d5697caSHong Zhang }
4767d5697caSHong Zhang 
477d71ae5a4SJacob Faibussowitsch static PetscErrorCode TSEIMEXSetOrdAdapt_EIMEX(TS ts, PetscBool flg)
478d71ae5a4SJacob Faibussowitsch {
4797d5697caSHong Zhang   TS_EIMEX *ext = (TS_EIMEX *)ts->data;
480*4d86920dSPierre Jolivet 
4817d5697caSHong Zhang   PetscFunctionBegin;
4827d5697caSHong Zhang   ext->ord_adapt = flg;
4833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4847d5697caSHong Zhang }
4857d5697caSHong Zhang 
4867d5697caSHong Zhang /*MC
4871d27aa22SBarry Smith    TSEIMEX - Time stepping with Extrapolated IMEX methods {cite}`constantinescu_a2010a`.
488679f3ef1SBarry Smith 
489679f3ef1SBarry Smith    These methods are intended for problems with well-separated time scales, especially when a slow scale is strongly nonlinear such that it
490bcf0153eSBarry Smith    is expensive to solve with a fully implicit method. The user should provide the stiff part of the equation using `TSSetIFunction()` and the
491bcf0153eSBarry Smith    non-stiff part with `TSSetRHSFunction()`.
492bcf0153eSBarry Smith 
493bcf0153eSBarry Smith       Level: beginner
4947d5697caSHong Zhang 
4957d5697caSHong Zhang   Notes:
496bcf0153eSBarry Smith   The default is a 3-stage scheme, it can be changed with `TSEIMEXSetMaxRows()` or -ts_eimex_max_rows
4977d5697caSHong Zhang 
4981d27aa22SBarry Smith   This method currently only works with ODE, for which the stiff part $ G(t,X,Xdot) $  has the form $ Xdot + Ghat(t,X)$.
4997d5697caSHong Zhang 
500679f3ef1SBarry Smith   The general system is written as
501679f3ef1SBarry Smith 
5021d27aa22SBarry Smith   $$
503679f3ef1SBarry Smith   G(t,X,Xdot) = F(t,X)
5041d27aa22SBarry Smith   $$
505679f3ef1SBarry Smith 
506679f3ef1SBarry Smith   where G represents the stiff part and F represents the non-stiff part. The user should provide the stiff part
507bcf0153eSBarry Smith   of the equation using TSSetIFunction() and the non-stiff part with `TSSetRHSFunction()`.
508679f3ef1SBarry Smith   This method is designed to be linearly implicit on G and can use an approximate and lagged Jacobian.
509679f3ef1SBarry Smith 
510679f3ef1SBarry Smith   Another common form for the system is
511679f3ef1SBarry Smith 
5121d27aa22SBarry Smith   $$
513679f3ef1SBarry Smith   y'=f(x)+g(x)
5141d27aa22SBarry Smith   $$
515679f3ef1SBarry Smith 
516679f3ef1SBarry Smith   The relationship between F,G and f,g is
517679f3ef1SBarry Smith 
5181d27aa22SBarry Smith   $$
519679f3ef1SBarry Smith   G = y'-g(x), F = f(x)
5201d27aa22SBarry Smith   $$
521679f3ef1SBarry Smith 
5221cc06b55SBarry Smith .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEIMEXSetMaxRows()`, `TSEIMEXSetRowCol()`, `TSEIMEXSetOrdAdapt()`, `TSType`
5237d5697caSHong Zhang  M*/
524d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode TSCreate_EIMEX(TS ts)
525d71ae5a4SJacob Faibussowitsch {
5267d5697caSHong Zhang   TS_EIMEX *ext;
5277d5697caSHong Zhang 
5287d5697caSHong Zhang   PetscFunctionBegin;
5297d5697caSHong Zhang   ts->ops->reset          = TSReset_EIMEX;
5307d5697caSHong Zhang   ts->ops->destroy        = TSDestroy_EIMEX;
5317d5697caSHong Zhang   ts->ops->view           = TSView_EIMEX;
5327d5697caSHong Zhang   ts->ops->setup          = TSSetUp_EIMEX;
5337d5697caSHong Zhang   ts->ops->step           = TSStep_EIMEX;
5347d5697caSHong Zhang   ts->ops->interpolate    = TSInterpolate_EIMEX;
5357d5697caSHong Zhang   ts->ops->evaluatestep   = TSEvaluateStep_EIMEX;
5367d5697caSHong Zhang   ts->ops->setfromoptions = TSSetFromOptions_EIMEX;
5377d5697caSHong Zhang   ts->ops->snesfunction   = SNESTSFormFunction_EIMEX;
5387d5697caSHong Zhang   ts->ops->snesjacobian   = SNESTSFormJacobian_EIMEX;
5392ffb9264SLisandro Dalcin   ts->default_adapt_type  = TSADAPTNONE;
5407d5697caSHong Zhang 
541efd4aadfSBarry Smith   ts->usessnes = PETSC_TRUE;
542efd4aadfSBarry Smith 
5434dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&ext));
5447d5697caSHong Zhang   ts->data = (void *)ext;
5457d5697caSHong Zhang 
5467d5697caSHong Zhang   ext->ord_adapt = PETSC_FALSE; /* By default, no order adapativity */
5477d5697caSHong Zhang   ext->row_ind   = -1;
5487d5697caSHong Zhang   ext->col_ind   = -1;
5497d5697caSHong Zhang   ext->max_rows  = TSEIMEXDefault;
5507d5697caSHong Zhang   ext->nstages   = TSEIMEXDefault;
5517d5697caSHong Zhang 
5529566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetMaxRows_C", TSEIMEXSetMaxRows_EIMEX));
5539566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetRowCol_C", TSEIMEXSetRowCol_EIMEX));
5549566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSEIMEXSetOrdAdapt_C", TSEIMEXSetOrdAdapt_EIMEX));
5553ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5567d5697caSHong Zhang }
557