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 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. I.e., 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 designate headers to generate bindings from and functions to exclude.

  • A generate.py script 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 the generate.py script is effectively governed by GPLv3, but the binding code generated by the script/litgen falls only under our SUNDIALS license. Because of this, the generate.py script and litgen extensions are kept in a separate git repository as package, sundials4py-generator.

Structure

sundials4py code lives in bindings/sundials4py. The main python module and all of its submodules are defined in sundials4py.cpp. sundials4py consists of 5 modules. Below we list how each one maps to the directory layout:

  • sundials.arkode: - Implements bindings for all of ARKODE. - Source directory: arkode/

  • sundials.cvodes: - Provides bindings for all of CVODES. - Source directory: cvodes/

  • sundials.idas: - Contains bindings for all of IDAS. - Source directory: idas/

  • sundials.kinsol: - Facilitates bindings for all of KINSOL. - Source directory: kinsol/.

  • sundials.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 these modules like this:

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. E.g.,

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]

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.

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. The upshot is that 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.

Here is an example for ARKODE:

In bindings/sundials4py/arkode/arkode_usersupplied.hpp, the function table struct is defined:

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,
                  std::forward<Args>(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 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 sundials_stepper.cpp,

m.def(
   "SUNStepper_SetEvolveFn",
   [](SUNStepper stepper,
      std::function<std::remove_pointer_t<SUNStepperEvolveFn>> fn) -> SUNErrCode
   {
      if (!stepper->python)
      {
         stepper->python = SUNStepperFunctionTable_Alloc();
      }
      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());

We are again creating a nanobind wrapper for SUNStepper_SetEvolveFn(), but this time, the function table is smuggled inside of the SUNStepper structure’s python member.