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/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

@@ -385,6 +385,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);
@@ -407,6 +409,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

@@ -1585,6 +1585,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(:) =");
@@ -888,6 +893,9 @@ int erkStep_TakeStep(ARKodeMem ark_mem, sunrealtype* dsmPtr, int* nflagPtr)
       }
     }
 
+    /* update current stage index */
+    step_mem->istage = is;
+
     /* apply user-supplied stage preprocessing function (if supplied) */
     if (ark_mem->PreProcessRHS != NULL)
     {

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;

src/arkode/arkode_io.c

@@ -2569,6 +2569,46 @@ int ARKodeGetEstLocalErrors(void* arkode_mem, N_Vector ele)
   }
 }
 
+/*---------------------------------------------------------------
+  ARKodeGetLastTime:
+
+  Returns the last saved value of the independent variable
+  ---------------------------------------------------------------*/
+int ARKodeGetLastTime(void* arkode_mem, sunrealtype* tn)
+{
+  ARKodeMem ark_mem;
+  if (arkode_mem == NULL)
+  {
+    arkProcessError(NULL, ARK_MEM_NULL, __LINE__, __func__, __FILE__,
+                    MSG_ARK_NO_MEM);
+    return (ARK_MEM_NULL);
+  }
+  ark_mem = (ARKodeMem)arkode_mem;
+
+  *tn = ark_mem->tn;
+  return (ARK_SUCCESS);
+}
+
+/*---------------------------------------------------------------
+  ARKodeGetLastState:
+
+  Returns the last saved time step solution.
+  ---------------------------------------------------------------*/
+int ARKodeGetLastState(void* arkode_mem, N_Vector* yn)
+{
+  ARKodeMem ark_mem;
+  if (arkode_mem == NULL)
+  {
+    arkProcessError(NULL, ARK_MEM_NULL, __LINE__, __func__, __FILE__,
+                    MSG_ARK_NO_MEM);
+    return (ARK_MEM_NULL);
+  }
+  ark_mem = (ARKodeMem)arkode_mem;
+
+  *yn = ark_mem->yn;
+  return (ARK_SUCCESS);
+}
+
 /*---------------------------------------------------------------
   ARKodeGetCurrentTime:
 
@@ -3184,6 +3224,38 @@ int ARKodeGetUserData(void* arkode_mem, void** user_data)
   return (ARK_SUCCESS);
 }
 
+/*---------------------------------------------------------------
+  ARKodeGetStageIndex:
+
+  Returns the index of the current stage and the total number of
+  stages.  If this is not supplied by the time-stepping module
+  then it returns (0,1), indicating that it is currently in the
+  first of only a single stage.
+  ---------------------------------------------------------------*/
+int ARKodeGetStageIndex(void* arkode_mem, int* stage, int* max_stages)
+{
+  ARKodeMem ark_mem;
+  if (arkode_mem == NULL)
+  {
+    arkProcessError(NULL, ARK_MEM_NULL, __LINE__, __func__, __FILE__,
+                    MSG_ARK_NO_MEM);
+    return (ARK_MEM_NULL);
+  }
+  ark_mem = (ARKodeMem)arkode_mem;
+
+  /* Call stepper routine to compute the state (if provided) */
+  if (ark_mem->step_getstageindex)
+  {
+    return (ark_mem->step_getstageindex(ark_mem, stage, max_stages));
+  }
+  else
+  {
+    *stage      = 0;
+    *max_stages = 1;
+    return (ARK_SUCCESS);
+  }
+}
+
 /*-----------------------------------------------------------------
   ARKodePrintAllStats