Python Interfaces

This chapter covers details developers need to know about the SUNDIALS Python interfaces, distributed as the Python package sundials4py.

We use nanobind for the Python bindings. nanobind is a sleeker, faster, alternative to pybind11. It is a C++ library, i.e. you write your binding code in C++. Nanobind does have some restrictions:

  • Cannot bind to functions which take double, or more pointer arguments. I.e., it cannot bind to ** or *** and so on. These have to be flattened somehow.

  • Cannot implicitly convert between a “View” container class and the underlying C type. That is, it cannot implicitly convert ARKodeView to void*. This means that users must explicitly convert from the “View” class by calling the get member function.

We use litgen to generate a large portion of the nanobind code.

  • We have generate.yaml files that designate headers to generate bindings from and functions to exclude.

  • A separate package, sundials4py-generate, includes a generate.py script which uses litgen to generate the bindings as a C++ header according to the generate.yaml.

  • For each generated file, there is at least one hand-coded file that includes the generated header.

Note

Litgen itself is licensed under GPLv3. This means that sundials4py-generate is effectively governed by GPLv3, but the binding code generated by the script/litgen falls only under our SUNDIALS license.

Structure

sundials4py code lives in bindings/sundials4py. The main python module and all of its submodules are defined in sundials4py.cpp. sundials4py consists of the following submodules:

  • sundials4py.arkode

    • Implements bindings for all of ARKODE.

    • Source directory: arkode/

  • sundials4py.cvodes

    • Provides bindings for all of CVODES.

    • Source directory: cvodes/

  • sundials4py.idas

    • Contains bindings for all of IDAS.

    • Source directory: idas/

  • sundials4py.kinsol

    • Facilitates bindings for all of KINSOL.

    • Source directory: kinsol/.

  • sundials4py.core

    • All SUNDIALS shared classes/modules and implementations.

    • Source directories:

      • nvector/

      • sunadaptcontroller/

      • sunadjointcheckpointscheme/

      • sunlinsol/

      • sunmatrix/

      • sunmemory/

      • sunnonlinsol/

      • sundials/

      • sundomeigest/

Development

sundials4py requires Python 3.12+ and the interpreter/development components e.g., if you were installing Python on a RedHat Linux system, you could install Python 3.12 with:

yum install python3.12 python3.12-devel

The recommended method for development is to use a typical Python development workflow with pip rather than invoking CMake directly.

cd sundials_root_directory
python -m venv .venv  # create python virtual environment
. .venv/bin/activate  # activate the python virtual environment
pip install scikit-build-core[pyproject] hatchling nanobind # this is a prerequisite for the next step
MAKEFLAGS="-j$(nproc)" pip install --no-build-isolation -Ceditable.rebuild=true -ve .[dev] # install sundials4py into the virtual environment

The last pip install command will allow automatic incremental builds. It will invoke the SUNDIALS CMake build system with the -DSUNDIALS_ENABLE_PYTHON=ON option through scikit-build-core. After the initial build, if you make any changes within SUNDIALS a rebuild will be triggered when you import the sundials4py module within a Python script.

Different CMake options can be controlled by passing them through the --config-settings (or -C for short) option of pip install. For example, the following builds the interface with 32-bit index types

MAKEFLAGS="-j$(nproc)" pip install --no-build-isolation -Ceditable.rebuild=true -ve .[dev] \
   -C cmake.define.SUNDIALS_INDEX_SIZE=32

Alternatively, you can set the CMAKE_ARGS environment variable:

export CMAKE_ARGS="-DSUNDIALS_INDEX_SIZE=32"
MAKEFLAGS="-j$(nproc)" pip install --no-build-isolation -Ceditable.rebuild=true -ve .[dev]

If you are developing new features inside of sundials, you may also need to install the Python code generator package, sundials4py-generate. The steps to install this package are found in its repo.

Tests

We use pytest for setting up unit/smoke tests of the interfaces. All tests are in bindings/sundials4py/test. The goal is to test the interfacing, not the correctness of SUNDIALS itself.

User-supplied functions

All user-supplied Python functions have to be wrapped with functions that convert between a std::function and a raw C function pointer. This is done by smuggling in a “function table” – a struct of std::function members – in a python member inside each integrator memory structure, and then storing the integrator memory structure in the user_data pointer. For the objects which are not the integrator, we still stuff the function table in the python member of the struct so it will be available in all of the module/class methods. As a result, every time we add a user-supplied function, we need to add a new member to the function table struct, and add a wrapper for it. We also have to add a wrapper for the “Set” function that takes the user-supplied function.

Consider the following example for ARKODE. In bindings/sundials4py/arkode/arkode_usersupplied.hpp, the function table struct is defined as:

struct arkode_user_supplied_fn_table
{
   // common user-supplied function pointers
   nb::object rootfn;
   nb::object ewtn;
   nb::object rwtn;
   nb::object adaptfn;
   nb::object expstabfn;
   nb::object vecresizefn;
   nb::object postprocessstepfn;
   nb::object postprocessstagefn;
   nb::object stagepredictfn;
   nb::object relaxfn;
   nb::object relaxjacfn;
   nb::object nlsfi;

   // truncated ...
};

Then each one of the functions in the table has a wrapper function defined below this struct definition, e.g.,

template<typename... Args>
inline int arkode_postprocessstepfn_wrapper(Args... args)
{
  return sundials4py::user_supplied_fn_caller<
    std::remove_pointer_t<ARKPostProcessFn>, arkode_user_supplied_fn_table,
    ARKodeMem, 1>(&arkode_user_supplied_fn_table::postprocessstepfn, args...);
}

Finally, in bindings/sundials4py/arkode/arkode.cpp, the Set function is registered with nanobind:

BIND_ARKODE_CALLBACK(ARKodeSetPostprocessStepFn, ARKPostProcessFn,
                     postprocessstepfn, arkode_postprocessstepfn_wrapper,
                     nb::arg("arkode_mem"), nb::arg("postprocessstep").none());

BIND_ARKODE_CALLBACK is a macro which expands to

m.def(
   ARKodeSetPostprocessStepFn,
   [](void* ark_mem, std::function<std::remove_pointer_t<ARKPostProcessFn>> fn)
   {
      auto fn_table    = get_arkode_fn_table(ark_mem);
      fn_table->MEMBER = nb::cast(fn);
      fntable->postprocessstepfn = nb::cast(fn);
      if (fn) { return NAME(ark_mem, arkode_postprocessstepfn_wrapper); }
      else { return NAME(ark_mem, nullptr); }
   },
   nb::arg("arkode_mem"), nb::arg("postprocessstep").none())

What we are doing is creating a custom nanobind wrapper of ARKodeSetPostprocessStepFn which takes the user-supplied Python side function as a std::function and stores it in the function table (which is stored in user data). The nb::arg arguments are needed so that we can make postprocessstep nullable (or None from Python).

Here is another example, but this time for the SUNStepper and with the python member instead of user_data. From bindings/sundials4py/sundials/sundials_stepper_usersupplied.hpp:

struct SUNStepperFunctionTable
{
   nb::object evolve;
   nb::object one_step;
   nb::object full_rhs;
   nb::object reinit;
   nb::object reset;
   nb::object reset_ckpt_idx;
   nb::object set_stop_time;
   nb::object set_step_direction;
   nb::object set_forcing;
   nb::object get_num_steps;
};

template<typename... Args>
inline SUNErrCode sunstepper_evolve_wrapper(Args... args)
{
  return sundials4py::user_supplied_fn_caller<
    std::remove_pointer_t<SUNStepperEvolveFn>, SUNStepperFunctionTable,
    SUNStepper>(&SUNStepperFunctionTable::evolve, std::forward<Args>(args)...);
}

From bindings/sundials4py/sundials/sundials_stepper.cpp,

m.def(
  "SUNStepper_SetEvolveFn",
  [](SUNStepper stepper,
     std::function<std::remove_pointer_t<SUNStepperEvolveFn>> fn) -> SUNErrCode
  {
    if (!stepper->python) { stepper->python = new SUNStepperFunctionTable; }
    auto fntable    = static_cast<SUNStepperFunctionTable*>(stepper->python);
    fntable->evolve = nb::cast(fn);
    if (fn)
    {
      return SUNStepper_SetEvolveFn(stepper, sunstepper_evolve_wrapper);
    }
    else { return SUNStepper_SetEvolveFn(stepper, nullptr); }
  },
  nb::arg("stepper"), nb::arg("fn").none());

Similar to the ARKODE case we are creating a nanobind wrapper for a callback function, but this time, the function table is smuggled inside of the SUNStepper structure’s python member.

Special user-supplied functions

Some user-supplied functions take parameters that need special treatment. These parameters are generally arrays or array-like pointers and output parameters (pointers used for the purpose of returning a value).

In this case, we need to write a custom std::function type which mirrors the C function typedef, but uses C++ analogs for the special parameters. For example, consider the wrapper below from bindings/sundials4py/cvodes/cvode_usersupplied.hpp.

// This std::function type mirrors the C typedef with C++ analogs for
// the array of N_Vectors and the jcurPtrB output parameter moved to a return
// value packed into a tuple.
using CVLsPrecSetupStdFnBS = std::tuple<int, sunbooleantype>(
  sunrealtype t, N_Vector y, std::vector<N_Vector> yS_1d, N_Vector yB,
  N_Vector fyB, sunbooleantype jokB, sunrealtype gammaB, void* user_dataB);

inline int cvode_lsprecsetupfnBS_wrapper(sunrealtype t, N_Vector y,
                                         N_Vector* yS_1d, N_Vector yB,
                                         N_Vector fyB, sunbooleantype jokB,
                                         sunbooleantype* jcurPtrB,
                                         sunrealtype gammaB, void* user_dataB)
{
  auto cv_mem   = static_cast<CVodeMem>(user_dataB);
  auto fn_table = get_cvode_fn_table(user_dataB);
  auto fn =
    nb::cast<std::function<CVLsPrecSetupStdFnBS>>(fn_table->lsprecsetupfnBS);
  auto Ns = cv_mem->cv_Ns;

  std::vector<N_Vector> yS(yS_1d, yS_1d + Ns);

  auto result = fn(t, y, yS, yB, fyB, jokB, gammaB, nullptr);

  *jcurPtrB = std::get<1>(result);

  return std::get<0>(result);
}