Feature: Various ARKODE updates (pre/post functions, remove extra copies, LSRKStep updates, Preallocate, GetStageIndex) (#897)

Combined updates from several PRs:
* #813: step pre/post functions, stage postprocessing, and pre RHS
functions
* #814: removed an extraneous copy of `yn` into `yout` when running in
ONE_STEP mode for ARKODE
* #815: various cleanup to the LSRKStep module
* #819: user-callable routine to retrieve the current stage index
* #886: user-callable `ARKodeInit` routine to finalize ARKODE and
allocate remaining memory before `ARKodeEvolve` is called

---------

Co-authored-by: David Gardner <gardner48@llnl.gov>

Author: drreynolds

.github/actions/test-driver/action.yml

@@ -11,6 +11,10 @@ inputs:
   tpls:
     description: "enable/disable TPLs"
     required: true
+  phase:
+    description: "Optional phase passed to test_driver.sh as --phase <PHASE> (e.g., BUILD)"
+    required: false
+    default: ""
 
 runs:
   using: composite
@@ -29,10 +33,15 @@ runs:
       run: |
         git config --global --add safe.directory $GITHUB_WORKSPACE
         cd test
+        extra_args=""
+        if [ -n "${{ inputs.phase }}" ]; then
+          extra_args=" --phase ${{ inputs.phase }}"
+        fi
         ./test_driver.sh \
         --testtype CUSTOM \
         --env env/docker.sh \
         --tpls ${{ inputs.tpls }} \
         --sunrealtype ${{ inputs.precision }} \
-        --indexsize ${{ inputs.indexsize }}
+        --indexsize ${{ inputs.indexsize }} \
+        ${extra_args}
       shell: bash

.github/workflows/ubuntu-latest.yml

@@ -52,6 +52,7 @@ jobs:
           indexsize: ${{ matrix.indexsize }}
           precision: ${{ matrix.precision }}
           tpls: ${{ matrix.tpls }}
+          phase: ${{ matrix.precision == 'extended' && 'BUILD' || '' }}
         env:
           CMAKE_BUILD_TYPE: ${{ matrix.buildtype }}
       - name: Archive build files from failed build

CHANGELOG.md

@@ -11,6 +11,44 @@ Updated the Kokkos N_Vector to support Kokkos 5.x versions.
 Added `SUNLogger_Set{Error,Warning,Info,Debug}File` functions to allow setting
 logger output streams with a `FILE*`.
 
+ARKODE now allows users to supply functions that will be called before each
+internal time step attempt (`ARKodeSetPreStepFn`), after each successful time
+step (`ARKodeSetPostStepFn`), before right-hand side routines are called on an
+updated state (`ARKodeSetPreRhsFn`), and/or once each internal step/stage is
+computed (`ARKodeSetPostprocessStepFn`/ `ARKodeSetPostprocessStageFn`). These
+are considered **advanced** functions, as they should treat the state vector as
+read-only, otherwise all theoretical guarantees of solution accuracy and
+stability will be lost.  As a result of these new functions, the values of
+multiple ARKODE return codes (e.g., ``ARK_INTERP_FAIL``) have been updated;
+users who key off of the named constants will not be affected, but users who
+rely on the values themselves should update their codes accordingly.
+
+Note to users utilizing the previously undocumented `ARKodeSetPostprocessStepFn`
+function, the supplied function is now called on the newly computed state vector
+for all step attempts not just successful steps. To obtain the previous behavior
+of only calling a function on successful steps, switch to using
+`ARKodeSetPostStepFn`.
+
+Removed extraneous copy of output vector when using ARKODE in ``ARK_ONE_STEP`` mode.
+
+The default number of stages for the SSP Runge-Kutta methods `ARKODE_LSRK_SSP_S_2`
+and `ARKODE_LSRK_SSP_S_3` in LSRKStep were changed from 10 and 9, respectively, to
+their minimum allowable values of 2 and 4. Users may revert to the previous values
+by calling `LSRKStepSetNumSSPStages`.
+
+Added the optional function `ARKodeInit` to ARKODE to enable
+data allocation before the first call to `ARKodeEvolve`
+(but after all other optional input routines have been called), to support
+users who measure memory usage before beginning a simulation.
+
+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 wish to 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).
+
+Added the functions `ARKodeGetLastTime` and `ARKodeGetLastState` to return the last
+successful time and state achieved by ARKODE, respectively.
+
 ### Bug Fixes
 
 Fixed a CMake bug where the SuperLU_MT interface would not be built and
@@ -22,6 +60,37 @@ Fixed a bug where passing an empty string to `SUNLogger_Set{Error,Warning,Info,D
 did not disable the corresponding logging stream ([Issue
 #844](https://github.com/llnl/sundials/issues/844)).
 
+Fixed a bug in logging output from ARKODE, where for some time stepping modules,
+the current "time" output in the logger was incorrect.
+
+Fixed a bug in the ARKODE discrete adjoint checkpointing where an incorrect
+state would be stored on the first step if the output vector passed to
+`ARKodeEvolve` did not contain the initial condition on the first call.
+
+Fixed a bug in MRIStep when using a custom inner integrator that relies on the
+input state being the initial condition for the fast integration rather than
+retaining the result from the last inner integration or most recent reset call
+and the output vector passed to `ARKodeEvolve` does not contain the initial
+condition on the first call or the last returned solution on subsequent calls.
+
+Removed an extraneous copy of the output vector in each step with SplittingStep.
+
+Added a missing call to `SUNNonlinSolSetup` in MRIStep when using an
+IMEX-MRI-SR method.
+
+Fixed a potential bug in LSRKStep's `ARKODE_LSRK_SSP_S_3` method, where a real
+number was used instead of an integer, potentially resulting in a rounding error.
+
+Fixed a bug in LSRKStep where an incorrect state vector could be passed to a
+user-supplied dominant eigenvalue function on the first step unless the output
+vector passed to `ARKodeEvolve` contained the initial condition and when an
+eigenvalue estimate is requested on the first step in a subsequent call to
+`ARKodeEvolve` unless the output vector passed contained the most recently returned
+solution.
+
+Fixed a bug in MRIStep for estimating the first "slow" time step in an adaptive
+multirate calculation.
+
 ### Deprecation Notices
 
 Several CMake options have been deprecated in favor of namespaced versions

bindings/sundials4py/arkode/arkode.cpp

@@ -137,6 +137,18 @@ void bind_arkode(nb::module_& m)
                         arkode_relaxjacfn_wrapper, nb::arg("arkode_mem"),
                         nb::arg("rfn").none(), nb::arg("rjacfn").none());
 
+  BIND_ARKODE_CALLBACK(ARKodeSetPreStepFn, ARKPreStepFn, prestepfn,
+                       arkode_prestepfn_wrapper, nb::arg("arkode_mem"),
+                       nb::arg("prestep").none());
+
+  BIND_ARKODE_CALLBACK(ARKodeSetPostStepFn, ARKPostStepFn, poststepfn,
+                       arkode_poststepfn_wrapper, nb::arg("arkode_mem"),
+                       nb::arg("poststep").none());
+
+  BIND_ARKODE_CALLBACK(ARKodeSetPreRhsFn, ARKPreRhsFn, prerhsfn,
+                       arkode_prerhsfn_wrapper, nb::arg("arkode_mem"),
+                       nb::arg("prerhs").none());
+
   BIND_ARKODE_CALLBACK(ARKodeSetPostprocessStepFn, ARKPostProcessFn,
                        postprocessstepfn, arkode_postprocessstepfn_wrapper,
                        nb::arg("arkode_mem"), nb::arg("postprocessstep").none());

bindings/sundials4py/arkode/arkode_generated.hpp

@@ -62,24 +62,28 @@ m.attr("ARK_INNERTOOUTER_FAIL")      = -36;
 m.attr("ARK_POSTPROCESS_FAIL")       = -37;
 m.attr("ARK_POSTPROCESS_STEP_FAIL")  = -37;
 m.attr("ARK_POSTPROCESS_STAGE_FAIL") = -38;
-m.attr("ARK_USER_PREDICT_FAIL")      = -39;
-m.attr("ARK_INTERP_FAIL")            = -40;
-m.attr("ARK_INVALID_TABLE")          = -41;
-m.attr("ARK_CONTEXT_ERR")            = -42;
-m.attr("ARK_RELAX_FAIL")             = -43;
-m.attr("ARK_RELAX_MEM_NULL")         = -44;
-m.attr("ARK_RELAX_FUNC_FAIL")        = -45;
-m.attr("ARK_RELAX_JAC_FAIL")         = -46;
-m.attr("ARK_CONTROLLER_ERR")         = -47;
-m.attr("ARK_STEPPER_UNSUPPORTED")    = -48;
-m.attr("ARK_DOMEIG_FAIL")            = -49;
-m.attr("ARK_MAX_STAGE_LIMIT_FAIL")   = -50;
-m.attr("ARK_SUNSTEPPER_ERR")         = -51;
-m.attr("ARK_STEP_DIRECTION_ERR")     = -52;
-m.attr("ARK_ADJ_CHECKPOINT_FAIL")    = -53;
-m.attr("ARK_ADJ_RECOMPUTE_FAIL")     = -54;
-m.attr("ARK_SUNADJSTEPPER_ERR")      = -55;
-m.attr("ARK_DEE_FAIL")               = -56;
+m.attr("ARK_PRESTEPFN_FAIL")         = -39;
+m.attr("ARK_POSTSTEPFN_FAIL")        = -40;
+m.attr("ARK_PRERHSFN_FAIL")          = -41;
+m.attr("ARK_USER_PREDICT_FAIL")      = -42;
+m.attr("ARK_INTERP_FAIL")            = -43;
+m.attr("ARK_INVALID_TABLE")          = -44;
+m.attr("ARK_CONTEXT_ERR")            = -45;
+m.attr("ARK_RELAX_FAIL")             = -46;
+m.attr("ARK_RELAX_MEM_NULL")         = -47;
+m.attr("ARK_RELAX_FUNC_FAIL")        = -48;
+m.attr("ARK_RELAX_JAC_FAIL")         = -49;
+m.attr("ARK_CONTROLLER_ERR")         = -50;
+m.attr("ARK_STEPPER_UNSUPPORTED")    = -51;
+m.attr("ARK_DOMEIG_FAIL")            = -52;
+m.attr("ARK_MAX_STAGE_LIMIT_FAIL")   = -53;
+m.attr("ARK_SUNSTEPPER_ERR")         = -54;
+m.attr("ARK_STEP_DIRECTION_ERR")     = -55;
+m.attr("ARK_ADJ_CHECKPOINT_FAIL")    = -56;
+m.attr("ARK_ADJ_RECOMPUTE_FAIL")     = -57;
+m.attr("ARK_SUNADJSTEPPER_ERR")      = -58;
+m.attr("ARK_DEE_FAIL")               = -59;
+m.attr("ARK_STEP_H0_FAIL")           = -60;
 m.attr("ARK_UNRECOGNIZED_ERROR")     = -99;
 
 auto pyEnumARKRelaxSolver = nb::enum_<ARKRelaxSolver>(m, "ARKRelaxSolver",
@@ -107,6 +111,9 @@ auto pyEnumARKAccumError =
 m.def("ARKodeReset", ARKodeReset, nb::arg("arkode_mem"), nb::arg("tR"),
       nb::arg("yR"));
 
+m.def("ARKodeInit", ARKodeInit, nb::arg("arkode_mem"),
+      "Optional data allocation function");
+
 m.def(
   "ARKodeCreateMRIStepInnerStepper",
   [](void* arkode_mem)
@@ -497,6 +504,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>
@@ -671,6 +698,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>

bindings/sundials4py/arkode/arkode_usersupplied.hpp

@@ -45,6 +45,9 @@ struct arkode_user_supplied_fn_table
   nb::object ewtn;
   nb::object rwtn;
   nb::object vecresizefn;
+  nb::object prerhsfn;
+  nb::object prestepfn;
+  nb::object poststepfn;
   nb::object postprocessstepfn;
   nb::object postprocessstagefn;
   nb::object stagepredictfn;
@@ -159,6 +162,30 @@ inline int arkode_vecresizefn_wrapper(Args... args)
     ARKodeMem, 1>(&arkode_user_supplied_fn_table::vecresizefn, args...);
 }
 
+template<typename... Args>
+inline int arkode_prerhsfn_wrapper(Args... args)
+{
+  return sundials4py::user_supplied_fn_caller<
+    std::remove_pointer_t<ARKPreRhsFn>, arkode_user_supplied_fn_table,
+    ARKodeMem, 1>(&arkode_user_supplied_fn_table::prerhsfn, args...);
+}
+
+template<typename... Args>
+inline int arkode_prestepfn_wrapper(Args... args)
+{
+  return sundials4py::user_supplied_fn_caller<
+    std::remove_pointer_t<ARKPreStepFn>, arkode_user_supplied_fn_table,
+    ARKodeMem, 1>(&arkode_user_supplied_fn_table::prestepfn, args...);
+}
+
+template<typename... Args>
+inline int arkode_poststepfn_wrapper(Args... args)
+{
+  return sundials4py::user_supplied_fn_caller<
+    std::remove_pointer_t<ARKPostStepFn>, arkode_user_supplied_fn_table,
+    ARKodeMem, 1>(&arkode_user_supplied_fn_table::poststepfn, args...);
+}
+
 template<typename... Args>
 inline int arkode_postprocessstepfn_wrapper(Args... args)
 {