feat(pybammsolvers)!: an IDAKLU backend over the Rust C ABI

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Breaking changes
 
+- `pybamm` now requires `pybammsolvers>=0.12.0`, whose IDAKLU backend can evaluate a Rust `CompiledModel`. Upgrade `pybammsolvers` together with `pybamm`. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961))
 - `pybamm` now ships platform wheels with its Rust core compiled in as `pybamm.rust._core`: one `cp310-abi3` wheel per platform for Python 3.10-3.15. Building from source or the sdist requires `cargo` 1.89+; `pip install --only-binary pybamm pybamm` forces a wheel. ([#5958](https://github.com/pybamm-team/PyBaMM/pull/5958))
 
 ## Deprecated
@@ -19,6 +20,7 @@
 - `pybamm.rust` exposes `ExprGraph`, `CompiledFunction` (one expression, or a named set via `compile_outputs`, with lane-batched trajectory and Hermite evaluation) and `CompiledJacobian`, with hand-written type stubs that `mypy.stubtest` pins against the built extension. ([#5958](https://github.com/pybamm-team/PyBaMM/pull/5958))
 - `Symbol.to_rust` lowers a discretised expression tree into a `pybamm.rust.ExprGraph`, covering the operators, matrix products, concatenations, `Conditional`, the specific functions, `RegPower` and 1D-3D interpolants. Spatial and tensor operators, such as `Norm` and `Component`, raise `TypeError`. ([#5959](https://github.com/pybamm-team/PyBaMM/pull/5959))
 - `model.convert_to_format = "rust"` evaluates a model's rhs, algebraic, Jacobian and event functions through compiled Rust evaluators, for `ScipySolver`, `AlgebraicSolver` and EIS. ([#5960](https://github.com/pybamm-team/PyBaMM/pull/5960))
+- pybammsolvers gains an IDAKLU backend that evaluates a Rust `CompiledModel` through its C ABI, resolved at runtime so IDAKLU still loads without the Rust extension, with a Rust-backed Newton solver and output sensitivities projected natively in-solver. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961))
 
 ## Bug fixes
 
@@ -41,6 +43,8 @@
 - Fixed copying a function without simplification, as `ParameterValues` does for a parameter set loaded with `Serialise.load_parameters`, returning a generic `pybamm.Function`: `pybamm.exp` and the other specific functions lost their class, and an `Arcsinh2` copy could not be evaluated. ([#5954](https://github.com/pybamm-team/PyBaMM/pull/5954))
 - `RegPower` now evaluates the same operations in `evaluate` and CasADi, taking a literal exponent of 0.5 or 2 as square roots, so they agree bitwise; CasADi used a `hypot` form. With `delta=0`, `evaluate` and the Jacobian use `base ** exponent`, as CasADi did. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955))
 - CasADi linear and `pchip` interpolants on a uniform grid now find a point's segment by binary search. The `"exact"` lookup could pick the neighbouring segment at or one ulp beside a knot, giving the Jacobian that segment's slope. ([#5955](https://github.com/pybamm-team/PyBaMM/pull/5955))
+- `IDAKLUSolver` no longer reads one value past the inputs of a model without input parameters. Inputs with a different number of values than at set-up now raise a `SolverError` instead of being read at the set-up length. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961))
+- `IDAKLUSolver` now finds a DAE's consistent initial sensitivities with one linear solve; at tight tolerances `IDACalcIC` could leave them zero at the first time point. A failed `IDACalcIC` is logged, and raises a `SolverError` when computing sensitivities. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961))
 
 ## Optimizations
 

packages/pybamm/pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "pybammsolvers>=0.11.0,<0.12.0",
+    "pybammsolvers>=0.12.0,<0.13.0",
     "black",
     "numpy>=2.0.0",
     "scipy>=1.11.4,!=1.18.0", # 1.18.0 breaks pickling/deepcopy of interpolators, fixed in 1.18.1 (scipy#25489)

packages/pybamm/src/pybamm/solvers/idaklu_solver.py

@@ -677,7 +677,8 @@ def to_idaklu(fn):
 
         self._setup = {
             "number_of_states": len(y0),
-            "inputs": len(stacked_inputs),
+            # Not len: an empty stack is 1 x 0, and each solved row must be this wide
+            "inputs": stacked_inputs.size,
             "jac_bandwidth_upper": jac_bw_upper,
             "jac_bandwidth_lower": jac_bw_lower,
             "jac_times_cjmass_colptrs": jac_times_cjmass_colptrs,

packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py

@@ -380,6 +380,33 @@ def test_model_events(self):
             atol=1e-5,
         )
 
+    def test_an_input_free_model_takes_zero_width_inputs_rows(self):
+        # An empty input stack is 1 x 0, so each solver must copy no values from a row
+        model = pybamm.BaseModel()
+        model.convert_to_format = "casadi"
+        u = pybamm.Variable("u")
+        model.rhs = {u: -u}
+        model.initial_conditions = {u: 1}
+        pybamm.Discretisation().process_model(model)
+        solver = pybamm.IDAKLUSolver()
+        t_eval = np.array([0.0, 1.0])
+        solution = solver.solve(model, t_eval)
+        np.testing.assert_allclose(
+            solution["u"](t_eval), np.exp(-t_eval), rtol=1e-3, atol=1e-6
+        )
+
+        group = solver._setup["solver"]
+        with pytest.raises(
+            ValueError, match=r"wrong number of cols\. Expected 0 but got 1"
+        ):
+            group.solve(
+                t_eval,
+                np.array([]),
+                np.ones((1, 1)),
+                np.zeros((1, 1)),
+                np.ones((1, 1)),
+            )
+
     def test_input_params(self):
         # test a mix of scalar and vector input params
         model = pybamm.BaseModel()

packages/pybammsolvers/CMakeLists.txt

@@ -124,6 +124,9 @@ pybind11_add_module(idaklu
         src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp
         src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp
         src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionTypes.hpp
+        # IDAKLU expressions - Rust backend (header-only)
+        src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp
+        src/pybammsolvers/idaklu_source/Expressions/Rust/pybamm_rust_ffi.h
         # IDAKLU expressions - concrete implementations
         ${IDAKLU_EXPR_CASADI_SOURCE_FILES}
 )

packages/pybammsolvers/pyproject.toml

@@ -18,7 +18,7 @@ requires-python = ">=3.10,<3.16"
 license-files = ["LICENSE"]
 readme = "README.md"
 # version.py stays the source of truth; a pre-commit hook mirrors it here.
-version = "0.11.1"
+version = "0.12.0"
 dependencies = [
     "casadi==3.8.1",
     "numpy>=2.0.0",

packages/pybammsolvers/src/pybammsolvers/idaklu.cpp

@@ -15,6 +15,7 @@
 #include "idaklu_source/IdakluJax.hpp"
 #include "idaklu_source/common.hpp"
 #include "idaklu_source/Expressions/Casadi/CasadiFunctions.hpp"
+#include "idaklu_source/Expressions/Rust/RustFunctions.hpp"
 #include "idaklu_source/sundials_error_handler.hpp"
 #include "idaklu_source/reduce.hpp"
 #include "idaklu_source/StandaloneNewtonSolver.hpp"
@@ -119,8 +120,110 @@ IDAKLUSolverGroup *create_casadi_solver_group(
     );
   }
 
+  // No owner: CasADi's expression sets hold their own functions.
   return new IDAKLUSolverGroup(
-    std::move(solvers), number_of_states, number_of_parameters);
+    std::move(solvers), number_of_states, number_of_parameters, inputs_length, py::none());
+}
+
+// An int leads: without pybind11 headers, CodeQL parses a leading `py::object` as a variable initialiser
+IDAKLUSolverGroup *create_rust_solver_group(
+    int number_of_states,
+    py::object rust_evaluators,
+    int number_of_inputs,
+    int number_of_sens_params,
+    int number_of_events,
+    int number_of_algebraic_states,
+    const std::vector<int> &output_lens,
+    const np_array_int &jac_colptrs,
+    const np_array_int &jac_rowvals,
+    int jac_nnz,
+    int jac_bandwidth_lower,
+    int jac_bandwidth_upper,
+    const np_array_int &alg_jac_rowvals,
+    const np_array_int &alg_jac_colvals,
+    int alg_jac_nnz,
+    const np_array &rhs_alg_id,
+    const np_array &atol_np,
+    double rel_tol,
+    py::dict py_opts
+) {
+    auto setup_opts = SetupOptions(py_opts);
+    auto solver_opts = SolverOptions(py_opts);
+
+    // Resolve before any parallel region: a throwing static initialiser is retried per thread.
+    rust_ffi();
+
+    const int num_solvers = setup_opts.num_solvers;
+    const int pool_size = static_cast<int>(py::len(rust_evaluators));
+    if (pool_size != num_solvers) {
+        throw std::invalid_argument(
+            "rust_evaluators holds " + std::to_string(pool_size) +
+            " evaluator(s) but num_solvers is " + std::to_string(num_solvers) +
+            "; each solver needs its own evaluator.");
+    }
+    // The sub-block IC hands Rust a residual buffer sized from the algebraic ids.
+    if (number_of_algebraic_states != 0) {
+        auto ids = rhs_alg_id.unchecked<1>();
+        int n_alg_ids = 0;
+        for (py::ssize_t i = 0; i < ids.shape(0); ++i) {
+            n_alg_ids += is_algebraic(ids(i)) ? 1 : 0;
+        }
+        if (n_alg_ids != number_of_algebraic_states) {
+            throw std::invalid_argument(
+                "n_alg is " + std::to_string(number_of_algebraic_states) +
+                " but rhs_alg_id marks " + std::to_string(n_alg_ids) + " states algebraic.");
+        }
+    }
+
+    auto colptrs_ptr = jac_colptrs.unchecked<1>();
+    auto rowvals_ptr = jac_rowvals.unchecked<1>();
+    std::vector<int64_t> colptrs(colptrs_ptr.data(0),
+                                  colptrs_ptr.data(0) + colptrs_ptr.size());
+    std::vector<int64_t> rowvals(rowvals_ptr.data(0),
+                                  rowvals_ptr.data(0) + rowvals_ptr.size());
+    auto alg_rowvals_ptr = alg_jac_rowvals.unchecked<1>();
+    auto alg_colvals_ptr = alg_jac_colvals.unchecked<1>();
+    std::vector<int64_t> alg_rowvals(
+        alg_rowvals_ptr.data(0),
+        alg_rowvals_ptr.data(0) + alg_rowvals_ptr.size()
+    );
+    std::vector<int64_t> alg_colvals(
+        alg_colvals_ptr.data(0),
+        alg_colvals_ptr.data(0) + alg_colvals_ptr.size()
+    );
+
+    std::vector<std::unique_ptr<IDAKLUSolver>> solvers;
+    solvers.reserve(num_solvers);
+    for (int i = 0; i < num_solvers; i++) {
+        // One evaluator per solver: they share the compiled tape but not their scratch.
+        void* model_ptr = pybamm_rust_detail::model_from_address(
+            rust_evaluators.attr("as_ptr")(i).cast<uintptr_t>());
+
+        auto functions = std::make_unique<RustFunctions>(
+            model_ptr, number_of_states, number_of_inputs, number_of_sens_params,
+            number_of_algebraic_states, number_of_events, jac_nnz, colptrs, rowvals,
+            alg_jac_nnz, alg_rowvals, alg_colvals, output_lens, setup_opts
+        );
+
+        solvers.emplace_back(
+            std::unique_ptr<IDAKLUSolver>(
+                create_idaklu_solver(
+                    std::move(functions),
+                    number_of_sens_params,
+                    jac_colptrs, jac_rowvals, jac_nnz,
+                    jac_bandwidth_lower, jac_bandwidth_upper,
+                    number_of_events,
+                    rhs_alg_id, atol_np, rel_tol,
+                    number_of_inputs,
+                    solver_opts, setup_opts
+                )
+            )
+        );
+    }
+
+    return new IDAKLUSolverGroup(
+        std::move(solvers), number_of_states, number_of_sens_params, number_of_inputs,
+        std::move(rust_evaluators));
 }
 
 PYBIND11_MAKE_OPAQUE(std::vector<np_array>);
@@ -176,6 +279,29 @@ PYBIND11_MODULE(idaklu, m)
     py::arg("alg_jac"),
     py::return_value_policy::take_ownership);
 
+  m.def("create_rust_solver_group", &create_rust_solver_group,
+      "Create a Rust-backed IDAKLU solver group",
+      py::arg("number_of_states"),
+      py::arg("rust_evaluators"),
+      py::arg("number_of_inputs"),
+      py::arg("number_of_sens_params"),
+      py::arg("number_of_events"),
+      py::arg("n_alg"),
+      py::arg("output_lens"),
+      py::arg("jac_colptrs"),
+      py::arg("jac_rowvals"),
+      py::arg("jac_nnz"),
+      py::arg("jac_bandwidth_lower"),
+      py::arg("jac_bandwidth_upper"),
+      py::arg("alg_jac_rowvals"),
+      py::arg("alg_jac_colvals"),
+      py::arg("alg_jac_nnz"),
+      py::arg("rhs_alg_id"),
+      py::arg("atol"),
+      py::arg("rtol"),
+      py::arg("options"),
+      py::return_value_policy::take_ownership);
+
   m.def("observe", &observe,
     "Observe variables",
     py::arg("ts"),
@@ -266,6 +392,16 @@ PYBIND11_MODULE(idaklu, m)
          py::arg("max_iter"), py::arg("max_backtracks"),
          py::arg("eps_newt"), py::arg("use_sparse"),
          py::arg("klu_pivot_tolerance") = 0.1)
+    .def(py::init<std::uintptr_t, int, int,
+                  std::vector<expr_int>, std::vector<expr_int>,
+                  std::vector<sunrealtype>, sunrealtype, sunrealtype,
+                  int, int, sunrealtype, bool, sunrealtype>(),
+         py::arg("rust_model"), py::arg("n_rhs"), py::arg("n_alg"),
+         py::arg("jac_rows"), py::arg("jac_cols"),
+         py::arg("atol"), py::arg("rtol"), py::arg("step_tol"),
+         py::arg("max_iter"), py::arg("max_backtracks"),
+         py::arg("eps_newt"), py::arg("use_sparse"),
+         py::arg("klu_pivot_tolerance") = 0.1)
     .def("solve", &StandaloneNewtonSolver::solve,
          py::arg("t"), py::arg("y0"), py::arg("inputs"),
          py::arg("logger") = py::none(),

packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSet.hpp

@@ -10,6 +10,11 @@ class ExpressionSet
 {
 public:
 
+  // Compile-time markers, overridden by RustFunctions.
+  static constexpr bool kNativeOutputSensitivities = false;
+  static constexpr bool kNativeBatchedOutputs = false;
+  static constexpr bool kNativeBatchedJacActions = false;
+
   /**
    * @brief Constructor
    */

packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Rust/RustFunctions.hpp

@@ -15,6 +15,11 @@
 
 namespace pybamm_rust_detail {
 
+/* Rust models reach C++ as the integer address Python holds, so the cast is unavoidable. */
+inline void* model_from_address(std::uintptr_t address) {
+    return reinterpret_cast<void*>(address);  // NOLINT(performance-no-int-to-ptr)
+}
+
 /* Throw unless `given`, a count C++ sizes buffers from, equals the Rust model's `rust`. */
 inline void check_dim(const char* name, int64_t given, int64_t rust) {
     if (given != rust) {

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp

@@ -100,6 +100,12 @@ std::vector<Solution> IDAKLUSolverGroup::solve(
       "inputs has wrong number of rows. Expected " + std::to_string(number_of_groups) +
       " but got " + std::to_string(inputs.shape()[0]));
 
+  // Each solver copies `number_of_inputs` values from its row, so a narrower row is over-read
+  if (inputs.shape()[1] != number_of_inputs)
+    throw std::domain_error(
+      "inputs has wrong number of cols. Expected " + std::to_string(number_of_inputs) +
+      " but got " + std::to_string(inputs.shape()[1]));
+
   // pbar is optional: an empty array leaves IDAS at its unit default.
   const bool has_pbar = pbar.size() > 0;
   if (has_pbar) {

packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.hpp

@@ -13,11 +13,18 @@ class IDAKLUSolverGroup
 
   /**
    * @brief Default constructor
+   *
+   * `number_of_inputs` is the width each solver copies from an inputs row. `rust_owner` keeps
+   * the Python object owning the expression sets' memory alive with the group (none for
+   * CasADi); pybind11 destroys the group with the GIL held.
    */
-  IDAKLUSolverGroup(std::vector<std::unique_ptr<IDAKLUSolver>> solvers, int number_of_states, int number_of_parameters):
+  IDAKLUSolverGroup(std::vector<std::unique_ptr<IDAKLUSolver>> solvers, int number_of_states, int number_of_parameters,
+                    int number_of_inputs, py::object rust_owner):
+    m_rust_owner(std::move(rust_owner)),
     m_solvers(std::move(solvers)),
     number_of_states(number_of_states),
-    number_of_parameters(number_of_parameters)
+    number_of_parameters(number_of_parameters),
+    number_of_inputs(number_of_inputs)
     {}
 
   // no copy constructor (unique_ptr cannot be copied)
@@ -54,9 +61,12 @@ class IDAKLUSolverGroup
      */
     void set_loggers(const py::object &logger);
 
+    // Declared first so it is destroyed last, after the solvers that borrow its memory.
+    py::object m_rust_owner;
     std::vector<std::unique_ptr<IDAKLUSolver>> m_solvers;
     int number_of_states;
     int number_of_parameters;
+    int number_of_inputs;
 };
 
 #endif // PYBAMM_IDAKLU_SOLVER_GROUP_HPP