User-callable routine to retrieve the current stage index

CHANGELOG.md

@@ -8,6 +8,10 @@
 
 Multiple minor updates were made to the ARKODE package.  We removed an extraneous
 copy of the output vector when using ARKODE in `ARK_ONE_STEP` mode.  We added the
+function `ARKodeGetStageIndex` that returns the index of the stage currently
+being processed, and the total number of stages in the method, for users who must
+compute auxiliary quantities in their IVP right-hand side functions during some
+stages and not others (e.g., in all but the first or last stage). We added the
 function `ARKodeAllocateInternalData` to ARKODE to enable stage-related
 data allocation before the first call to `ARKodeEvolve` (but after all other
 optional input routines have been called), to support users who measure memory

bindings/sundials4py/arkode/arkode_generated.hpp

@@ -500,6 +500,26 @@ m.def("ARKodeGetReturnFlagName", ARKodeGetReturnFlagName, nb::arg("flag"));
 m.def("ARKodeWriteParameters", ARKodeWriteParameters, nb::arg("arkode_mem"),
       nb::arg("fp"));
 
+m.def(
+  "ARKodeGetStageIndex",
+  [](void* arkode_mem) -> std::tuple<int, int, int>
+  {
+    auto ARKodeGetStageIndex_adapt_modifiable_immutable_to_return =
+      [](void* arkode_mem) -> std::tuple<int, int, int>
+    {
+      int stage_adapt_modifiable;
+      int max_stages_adapt_modifiable;
+
+      int r = ARKodeGetStageIndex(arkode_mem, &stage_adapt_modifiable,
+                                  &max_stages_adapt_modifiable);
+      return std::make_tuple(r, stage_adapt_modifiable,
+                             max_stages_adapt_modifiable);
+    };
+
+    return ARKodeGetStageIndex_adapt_modifiable_immutable_to_return(arkode_mem);
+  },
+  nb::arg("arkode_mem"));
+
 m.def(
   "ARKodeGetNumExpSteps",
   [](void* arkode_mem) -> std::tuple<int, long>
@@ -674,6 +694,40 @@ m.def(
   },
   nb::arg("arkode_mem"));
 
+m.def(
+  "ARKodeGetLastTime",
+  [](void* arkode_mem) -> std::tuple<int, sunrealtype>
+  {
+    auto ARKodeGetLastTime_adapt_modifiable_immutable_to_return =
+      [](void* arkode_mem) -> std::tuple<int, sunrealtype>
+    {
+      sunrealtype tn_adapt_modifiable;
+
+      int r = ARKodeGetLastTime(arkode_mem, &tn_adapt_modifiable);
+      return std::make_tuple(r, tn_adapt_modifiable);
+    };
+
+    return ARKodeGetLastTime_adapt_modifiable_immutable_to_return(arkode_mem);
+  },
+  nb::arg("arkode_mem"));
+
+m.def(
+  "ARKodeGetLastState",
+  [](void* arkode_mem) -> std::tuple<int, N_Vector>
+  {
+    auto ARKodeGetLastState_adapt_modifiable_immutable_to_return =
+      [](void* arkode_mem) -> std::tuple<int, N_Vector>
+    {
+      N_Vector state_adapt_modifiable;
+
+      int r = ARKodeGetLastState(arkode_mem, &state_adapt_modifiable);
+      return std::make_tuple(r, state_adapt_modifiable);
+    };
+
+    return ARKodeGetLastState_adapt_modifiable_immutable_to_return(arkode_mem);
+  },
+  nb::arg("arkode_mem"), "nb::rv_policy::reference", nb::rv_policy::reference);
+
 m.def(
   "ARKodeGetCurrentTime",
   [](void* arkode_mem) -> std::tuple<int, sunrealtype>

doc/arkode/guide/source/Usage/User_callable.rst

@@ -3622,6 +3622,8 @@ Actual initial time step size used                     :c:func:`ARKodeGetActualI
 Step size used for the last successful step            :c:func:`ARKodeGetLastStep`
 Step size to be attempted on the next step             :c:func:`ARKodeGetCurrentStep`
 Integration direction, e.g., forward or backward       :c:func:`ARKodeGetStepDirection`
+Last saved time reached by the solver                  :c:func:`ARKodeGetLastTime`
+Last saved solution reached by the solver              :c:func:`ARKodeGetLastState`
 Current internal time reached by the solver            :c:func:`ARKodeGetCurrentTime`
 Current internal solution reached by the solver        :c:func:`ARKodeGetCurrentState`
 Current :math:`\gamma` value used by the solver        :c:func:`ARKodeGetCurrentGamma`
@@ -3641,6 +3643,7 @@ Estimated local truncation error vector                :c:func:`ARKodeGetEstLoca
 Number of constraint test failures                     :c:func:`ARKodeGetNumConstrFails`
 Retrieve a pointer for user data                       :c:func:`ARKodeGetUserData`
 Retrieve the accumulated temporal error estimate       :c:func:`ARKodeGetAccumulatedError`
+Current stage index, and total number of stages        :c:func:`ARKodeGetStageIndex`
 =====================================================  ============================================
 
 
@@ -3744,6 +3747,38 @@ Retrieve the accumulated temporal error estimate       :c:func:`ARKodeGetAccumul
    .. versionadded:: 6.2.0
 
 
+.. c:function:: int ARKodeGetLastTime(void* arkode_mem, sunrealtype* tn)
+
+   Returns the last saved time reached by the solver.
+
+   :param arkode_mem: pointer to the ARKODE memory block.
+   :param tn: last saved time reached.
+
+   :retval ARK_SUCCESS: the function exited successfully.
+   :retval ARK_MEM_NULL: ``arkode_mem`` was ``NULL``.
+
+   .. versionadded:: x.y.z
+
+
+.. c:function:: int ARKodeGetLastState(void *arkode_mem, N_Vector *yn)
+
+   Returns the last saved solution reached by the solver.
+
+   :param arkode_mem: pointer to the ARKODE memory block.
+   :param yn: last saved solution.
+
+   :retval ARK_SUCCESS: the function exited successfully.
+   :retval ARK_MEM_NULL: ``arkode_mem`` was ``NULL``.
+
+   .. note::
+
+      Users should exercise extreme caution when using this function,
+      as altering values of *yn* may lead to undesirable behavior, depending
+      on the particular use case and on when this routine is called.
+
+   .. versionadded:: x.y.z
+
+
 .. c:function:: int ARKodeGetCurrentTime(void* arkode_mem, sunrealtype* tcur)
 
    Returns the current internal time reached by the solver.
@@ -4117,6 +4152,21 @@ Retrieve the accumulated temporal error estimate       :c:func:`ARKodeGetAccumul
    .. versionadded:: 6.2.0
 
 
+.. c:function:: int ARKodeGetStageIndex(void* arkode_mem, int* stage, int *max_stages)
+
+   Returns the index of the current stage (0-based) and the total number of
+   stages in the method.
+
+   :param arkode_mem: pointer to the ARKODE memory block.
+   :param stage: pointer to storage for the current stage index.
+   :param max_stages: pointer to storage for the number of stages in the method.
+
+   :retval ARK_SUCCESS: the function exited successfully.
+   :retval ARK_MEM_NULL: ``arkode_mem`` was ``NULL``.
+
+   .. versionadded:: x.y.z
+
+
 
 
 

doc/shared/RecentChanges.rst

@@ -7,6 +7,10 @@
 
 Multiple minor updates were made to the ARKODE package.  We removed an extraneous
 copy of the output vector when using ARKODE in ``ARK_ONE_STEP`` mode.  We added the
+function :c:func:`ARKodeGetStageIndex` that returns the index of the stage currently
+being processed, and the total number of stages in the method, for users who must
+compute auxiliary quantities in their IVP right-hand side functions during some
+stages and not others (e.g., in all but the first or last stage). We added the
 function :c:func:`ARKodeAllocateInternalData` to ARKODE to enable stage-related
 data allocation before the first call to :c:func:`ARKodeEvolve` (but after all other
 optional input routines have been called), to support users who measure memory

include/arkode/arkode.h

@@ -382,6 +382,8 @@ SUNDIALS_EXPORT int ARKodePrintAllStats(void* arkode_mem, FILE* outfile,
                                         SUNOutputFormat fmt);
 SUNDIALS_EXPORT char* ARKodeGetReturnFlagName(long int flag);
 SUNDIALS_EXPORT int ARKodeWriteParameters(void* arkode_mem, FILE* fp);
+SUNDIALS_EXPORT int ARKodeGetStageIndex(void* arkode_mem, int* stage,
+                                        int* max_stages);
 
 /* Optional output functions (temporal adaptivity) */
 SUNDIALS_EXPORT int ARKodeGetNumExpSteps(void* arkode_mem, long int* expsteps);
@@ -404,6 +406,9 @@ SUNDIALS_EXPORT int ARKodeGetAccumulatedError(void* arkode_mem,
 /* Optional output functions (implicit solver) */
 SUNDIALS_EXPORT int ARKodeGetNumLinSolvSetups(void* arkode_mem,
                                               long int* nlinsetups);
+SUNDIALS_EXPORT int ARKodeGetLastTime(void* arkode_mem, sunrealtype* tn);
+SUNDIALS_EXPORT int ARKodeGetLastState(void* arkode_mem,
+                                       N_Vector* state); // nb::rv_policy::reference
 SUNDIALS_EXPORT int ARKodeGetCurrentTime(void* arkode_mem, sunrealtype* tcur);
 SUNDIALS_EXPORT int ARKodeGetCurrentState(void* arkode_mem,
                                           N_Vector* state); // nb::rv_policy::reference

src/arkode/arkode.c

@@ -1580,6 +1580,7 @@ ARKodeMem arkCreate(SUNContext sunctx)
   ark_mem->step_getnumnonlinsolviters     = NULL;
   ark_mem->step_getnumnonlinsolvconvfails = NULL;
   ark_mem->step_getnonlinsolvstats        = NULL;
+  ark_mem->step_getstageindex             = NULL;
   ark_mem->step_setforcing                = NULL;
   ark_mem->step_mem                       = NULL;
   ark_mem->step_supports_adaptive         = SUNFALSE;

src/arkode/arkode_arkstep.c

@@ -143,6 +143,7 @@ void* ARKStepCreate(ARKRhsFn fe, ARKRhsFn fi, sunrealtype t0, N_Vector y0,
   ark_mem->step_getnumnonlinsolvconvfails = arkStep_GetNumNonlinSolvConvFails;
   ark_mem->step_getnonlinsolvstats        = arkStep_GetNonlinSolvStats;
   ark_mem->step_setforcing                = arkStep_SetInnerForcing;
+  ark_mem->step_getstageindex             = arkStep_GetStageIndex;
   ark_mem->step_supports_adaptive         = SUNTRUE;
   ark_mem->step_supports_implicit         = SUNTRUE;
   ark_mem->step_supports_massmatrix       = SUNTRUE;
@@ -1711,6 +1712,9 @@ int arkStep_TakeStep_Z(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
     *nflagPtr = ARK_SUCCESS;
   }
 
+  /* initialize the current stage index */
+  step_mem->istage = 0;
+
   /* call nonlinear solver setup if it exists */
   if (step_mem->NLS)
   {

src/arkode/arkode_arkstep_impl.h

@@ -243,6 +243,7 @@ void arkStep_Free(ARKodeMem ark_mem);
 void arkStep_PrintMem(ARKodeMem ark_mem, FILE* outfile);
 int arkStep_SetInnerForcing(ARKodeMem ark_mem, sunrealtype tshift,
                             sunrealtype tscale, N_Vector* f, int nvecs);
+int arkStep_GetStageIndex(ARKodeMem ark_mem, int* stage, int* max_stages);
 
 /* Internal utility routines */
 int arkStep_AccessARKODEStepMem(void* arkode_mem, const char* fname,

src/arkode/arkode_arkstep_io.c

@@ -1298,6 +1298,26 @@ int arkStep_GetNonlinSolvStats(ARKodeMem ark_mem, long int* nniters,
   return (ARK_SUCCESS);
 }
 
+/*---------------------------------------------------------------
+  arkStep_GetStageIndex:
+
+  Returns the current stage index and number of stages
+  ---------------------------------------------------------------*/
+int arkStep_GetStageIndex(ARKodeMem ark_mem, int* stage, int* max_stages)
+{
+  ARKodeARKStepMem step_mem;
+  int retval;
+
+  /* access ARKodeARKStepMem structure */
+  retval = arkStep_AccessStepMem(ark_mem, __func__, &step_mem);
+  if (retval != ARK_SUCCESS) { return (retval); }
+
+  *stage      = step_mem->istage;
+  *max_stages = step_mem->stages;
+
+  return (ARK_SUCCESS);
+}
+
 /*---------------------------------------------------------------
   arkStep_PrintAllStats:
 

src/arkode/arkode_erkstep.c

@@ -109,6 +109,7 @@ void* ERKStepCreate(ARKRhsFn f, sunrealtype t0, N_Vector y0, SUNContext sunctx)
   ark_mem->step_getnumrhsevals        = erkStep_GetNumRhsEvals;
   ark_mem->step_getestlocalerrors     = erkStep_GetEstLocalErrors;
   ark_mem->step_setforcing            = erkStep_SetInnerForcing;
+  ark_mem->step_getstageindex         = erkStep_GetStageIndex;
   ark_mem->step_supports_adaptive     = SUNTRUE;
   ark_mem->step_supports_relaxation   = SUNTRUE;
   ark_mem->step_mem                   = (void*)step_mem;
@@ -364,6 +365,7 @@ void erkStep_PrintMem(ARKodeMem ark_mem, FILE* outfile)
   /* output integer quantities */
   fprintf(outfile, "ERKStep: q = %i\n", step_mem->q);
   fprintf(outfile, "ERKStep: p = %i\n", step_mem->p);
+  fprintf(outfile, "ERKStep: istage = %i\n", step_mem->istage);
   fprintf(outfile, "ERKStep: stages = %i\n", step_mem->stages);
 
   /* output long integer quantities */
@@ -770,6 +772,9 @@ int erkStep_TakeStep(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
   cvals = step_mem->cvals;
   Xvecs = step_mem->Xvecs;
 
+  /* initialize the current stage index */
+  step_mem->istage = 0;
+
   SUNLogInfo(ARK_LOGGER, "begin-stages-list", "stage = 0, tcur = " SUN_FORMAT_G,
              ark_mem->tcur);
   SUNLogExtraDebugVec(ARK_LOGGER, "stage", ark_mem->yn, "z_0(:) =");
@@ -836,6 +841,9 @@ int erkStep_TakeStep(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
      the first stage RHS is just the full RHS from the start of the step */
   for (is = 1; is < step_mem->stages; is++)
   {
+    /* store current stage index */
+    step_mem->istage = is;
+
     /* Set current stage time(s) */
     ark_mem->tcur = ark_mem->tn + step_mem->B->c[is] * ark_mem->h;
 

src/arkode/arkode_erkstep_impl.h

@@ -58,6 +58,7 @@ typedef struct ARKodeERKStepMemRec
   N_Vector* F;          /* explicit RHS at each stage */
   int q;                /* method order               */
   int p;                /* embedding order            */
+  int istage;           /* current stage              */
   int stages;           /* number of stages           */
   ARKodeButcherTable B; /* ERK Butcher table          */
 
@@ -106,6 +107,7 @@ int erkStep_GetNumRhsEvals(ARKodeMem ark_mem, int partition_index,
 int erkStep_GetEstLocalErrors(ARKodeMem ark_mem, N_Vector ele);
 int erkStep_SetInnerForcing(ARKodeMem ark_mem, sunrealtype tshift,
                             sunrealtype tscale, N_Vector* f, int nvecs);
+int erkStep_GetStageIndex(ARKodeMem ark_mem, int* stage, int* max_stages);
 
 /* Internal utility routines */
 int erkStep_AccessARKODEStepMem(void* arkode_mem, const char* fname,

src/arkode/arkode_erkstep_io.c

@@ -398,6 +398,26 @@ int erkStep_GetEstLocalErrors(ARKodeMem ark_mem, N_Vector ele)
   return (ARK_SUCCESS);
 }
 
+/*---------------------------------------------------------------
+  erkStep_GetStageIndex:
+
+  Returns the current stage index and number of stages
+  ---------------------------------------------------------------*/
+int erkStep_GetStageIndex(ARKodeMem ark_mem, int* stage, int* max_stages)
+{
+  ARKodeERKStepMem step_mem;
+  int retval;
+
+  /* access ARKodeARKStepMem structure */
+  retval = erkStep_AccessStepMem(ark_mem, __func__, &step_mem);
+  if (retval != ARK_SUCCESS) { return (retval); }
+
+  *stage      = step_mem->istage;
+  *max_stages = step_mem->stages;
+
+  return (ARK_SUCCESS);
+}
+
 /*---------------------------------------------------------------
   erkStep_PrintAllStats:
 

src/arkode/arkode_impl.h

@@ -230,6 +230,8 @@ typedef int (*ARKTimestepSetUseCompensatedSums)(ARKodeMem ark_mem,
                                                 sunbooleantype onoff);
 typedef int (*ARKTimestepSetOptions)(ARKodeMem ark_mem, int* argidx, char* argv[],
                                      size_t offset, sunbooleantype* arg_used);
+typedef int (*ARKTimestepGetStageIndex)(ARKodeMem ark_mem, int* stage,
+                                        int* max_stages);
 
 /* time stepper interface functions -- temporal adaptivity */
 typedef int (*ARKTimestepGetEstLocalErrors)(ARKodeMem ark_mem, N_Vector ele);
@@ -423,6 +425,7 @@ struct ARKodeMemRec
   ARKTimestepSetStepDirection step_setstepdirection;
   ARKTimestepSetUseCompensatedSums step_setusecompensatedsums;
   ARKTimestepSetOptions step_setoptions;
+  ARKTimestepGetStageIndex step_getstageindex;
 
   /* Time stepper module -- temporal adaptivity */
   sunbooleantype step_supports_adaptive;