feat(rust): pybamm.DiffsolSolver, integrating rust models with diffsol

CHANGELOG.md

@@ -24,6 +24,7 @@
 - `IDAKLUSolver` solves `convert_to_format = "rust"` models through a Rust `CompiledModel`, one evaluator per thread, so `num_threads` parallelises input sets. `output_variables` support sensitivities and time integrals, and `NonlinearSolver` compiles its own model. ([#5962](https://github.com/pybamm-team/PyBaMM/pull/5962))
 - `IDAKLUSolver` observes rust-model solutions natively: variables and their sensitivities are evaluated on tapes compiled into the solve's Rust graph, with no CasADi function, and `Solution.update([...])` compiles a set of variables into one tape per segment. ([#5963](https://github.com/pybamm-team/PyBaMM/pull/5963))
 - `pybamm-core` gains an in-process DAE solver built on diffsol's BDF, with forward sensitivities under error control, batched one lane per parameter, and a rayon batch layer for input-set sweeps. ([#5964](https://github.com/pybamm-team/PyBaMM/pull/5964))
+- Added `pybamm.DiffsolSolver`, which integrates rust models entirely in Rust, releasing the GIL. Like `IDAKLUSolver`, it finds consistent initial conditions, solves input sets on `num_threads` threads, and computes forward sensitivities natively under error control. ([#5965](https://github.com/pybamm-team/PyBaMM/pull/5965))
 
 ## Bug fixes
 

docs/source/api/solvers/diffsol_solver.rst

@@ -0,0 +1,5 @@
+Diffsol Solver
+==============
+
+.. autoclass:: pybamm.DiffsolSolver
+  :members:

docs/source/api/solvers/index.rst

@@ -4,6 +4,7 @@ Solvers
 .. toctree::
 
   base_solver
+  diffsol_solver
   dummy_solver
   scipy_solver
   jax_solver

docs/source/developer/rust_idaklu_build.md

@@ -16,7 +16,7 @@ variables to configure.
 **How the pieces connect:**
 
 - `pybamm.rust` provides the Python API that lowers a discretised model to a
-  `CompiledModel` for the IDAKLU solver.
+  `CompiledModel` and drives the diffsol solver.
 - pybammsolvers does **not** link against the Rust core (the package
   dependency only flows the other way). The IDAKLU C++ consumer resolves the
   `extern "C"` entry points at runtime — `dlsym(RTLD_DEFAULT)` on POSIX, a
@@ -105,6 +105,9 @@ cd packages/pybamm-rust && cargo test --workspace
 # Rust-IDAKLU integration tests (from the repo root)
 uv run --group dev pytest packages/pybamm/tests/integration/test_rust_idaklu_parity.py
 uv run --group dev pytest packages/pybamm/tests/integration/test_rust_idaklu_spm.py
+
+# Diffsol solver unit tests
+uv run --group dev pytest packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py
 ```
 
 ## Verifying FFI symbols

packages/pybamm-rust/pybamm-core/src/error.rs

@@ -189,6 +189,14 @@ pub enum CoreError {
     )]
     Y0SensLength { got: usize, expected: usize },
 
+    #[error(
+        "a batch needs one entry per input set, but got {y0} initial state(s) and {inputs} input vector(s)"
+    )]
+    BatchWidths { y0: usize, inputs: usize },
+
+    #[error("a batch of {expected} input set(s) needs {expected} dy0/dp seed(s), but got {got}")]
+    BatchSensWidth { got: usize, expected: usize },
+
     #[error("sens_atol_factor must be finite and > 0, got {got}")]
     SensAtolFactor { got: f64 },
 

packages/pybamm-rust/pybamm-core/src/solver/batch.rs

@@ -8,6 +8,23 @@ use rayon::prelude::*;
 use super::solve::{InputSet, PreparedSolver, SolveOutcome, SolveRequest};
 use crate::error::CoreError;
 
+/// Check that every per-set argument agrees on the batch width.
+pub const fn check_batch_widths(
+    y0: usize,
+    inputs: usize,
+    y0_sens: Option<usize>,
+) -> Result<(), CoreError> {
+    if y0 != inputs {
+        return Err(CoreError::BatchWidths { y0, inputs });
+    }
+    if let Some(got) = y0_sens
+        && got != y0
+    {
+        return Err(CoreError::BatchSensWidth { got, expected: y0 });
+    }
+    Ok(())
+}
+
 impl PreparedSolver {
     /// Solve one trajectory per input set, concurrently, each factorising serially.
     ///
@@ -175,4 +192,19 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn mismatched_batch_widths_are_rejected_before_any_solve() {
+        let err = check_batch_widths(2, 1, None).expect_err("mismatched widths accepted");
+        assert!(matches!(err, CoreError::BatchWidths { y0: 2, inputs: 1 }));
+
+        let err = check_batch_widths(2, 2, Some(1)).expect_err("mismatched seed count accepted");
+        assert!(matches!(
+            err,
+            CoreError::BatchSensWidth {
+                got: 1,
+                expected: 2
+            }
+        ));
+    }
 }

packages/pybamm-rust/pybamm-core/src/solver/solve.rs

@@ -81,6 +81,31 @@ pub enum Termination {
     },
 }
 
+impl Termination {
+    /// The flag the solver boundary carries: `0` completed, `1` event root.
+    pub const fn as_flag(self) -> i32 {
+        match self {
+            Self::Completed => 0,
+            Self::Event { .. } => 1,
+        }
+    }
+
+    pub const fn event_time(self) -> Option<f64> {
+        match self {
+            Self::Completed => None,
+            Self::Event { t, .. } => Some(t),
+        }
+    }
+
+    /// The registration index of the event whose root ended the solve.
+    pub const fn event_index(self) -> Option<usize> {
+        match self {
+            Self::Completed => None,
+            Self::Event { index, .. } => Some(index),
+        }
+    }
+}
+
 /// What one solve returns, whichever payloads were asked for.
 ///
 /// On an event root, `t` ends at the root time rather than the last requested point.

packages/pybamm-rust/pybamm-python/Cargo.toml

@@ -9,14 +9,16 @@ name = "_core"
 crate-type = ["cdylib"]
 
 [features]
-default = ["serialize"]
+default = ["serialize", "diffsol"]
 serialize = ["pybamm-core/serialize", "dep:bincode"]
+diffsol = ["pybamm-core/diffsol", "dep:rayon"]
 
 [dependencies]
 pybamm-core = { path = "../pybamm-core" }
 pyo3 = { version = "0.28" }
 numpy = "0.28"
 bincode = { version = "1", optional = true }
+rayon = { version = "1", optional = true }
 
 [lints]
 workspace = true

packages/pybamm-rust/pybamm-python/src/errors.rs

@@ -4,11 +4,16 @@ use pybamm_core::CoreError;
 use pyo3::PyErr;
 use pyo3::exceptions::PyValueError;
 
-/// Map a core error to a Python exception: caller-argument and invalid-data
-/// problems (a malformed matrix or interpolant crossing the boundary) become
-/// `ValueError`.
-// By value, because every caller hands it over through `map_err`.
-#[allow(clippy::needless_pass_by_value)]
+/// Map a core error to `RuntimeError` for integration failures, else `ValueError`.
 pub fn core_err_to_py(err: CoreError) -> PyErr {
-    PyValueError::new_err(err.to_string())
+    match err {
+        #[cfg(feature = "diffsol")]
+        CoreError::Diffsol(source) => {
+            pyo3::exceptions::PyRuntimeError::new_err(format!("diffsol error: {source}"))
+        },
+        retry @ CoreError::SensRetryFailed { .. } => {
+            pyo3::exceptions::PyRuntimeError::new_err(retry.to_string())
+        },
+        other => PyValueError::new_err(other.to_string()),
+    }
 }

packages/pybamm-rust/pybamm-python/src/lib.rs

@@ -12,7 +12,11 @@ mod expr;
 mod function_metadata;
 mod jacobian;
 mod model;
+#[cfg(feature = "diffsol")]
+mod rayon_pool;
 mod scratch_pool;
+#[cfg(feature = "diffsol")]
+mod solver;
 
 #[pymodule]
 fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -23,5 +27,14 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<model::CompiledModel>()?;
     m.add_class::<evaluator_pool::EvaluatorPool>()?;
 
+    #[cfg(feature = "diffsol")]
+    {
+        m.add_class::<solver::PreparedSolver>()?;
+        m.add_class::<solver::SolverStatistics>()?;
+        m.add_class::<solver::SolveOutcome>()?;
+        m.add_function(wrap_pyfunction!(solver::default_solver_options, m)?)?;
+        m.add_function(wrap_pyfunction!(rayon_pool::_pool_ids, m)?)?;
+    }
+
     Ok(())
 }

packages/pybamm-rust/pybamm-python/src/model.rs

@@ -80,9 +80,9 @@ pub struct CompiledModel {
     pub(crate) compiled: Arc<core_model::CompiledModel>,
     /// Workspaces for the direct `eval_*`/`assemble_*` helpers.
     scratch_pool: ScratchPool<core_model::Workspace>,
-    /// Retained derivation source for the bundle views.
-    graph: Py<ExprGraph>,
-    rhs_root: NodeId,
+    /// Retained derivation source for the bundle views and `PreparedSolver`.
+    pub(crate) graph: Py<ExprGraph>,
+    pub(crate) rhs_root: NodeId,
     output_roots: Vec<NodeId>,
     event_roots: Vec<NodeId>,
     /// The algebraic block's root and its strictly-ascending global state indices.

packages/pybamm-rust/pybamm-python/src/rayon_pool.rs

@@ -0,0 +1,175 @@
+//! Process-wide rayon pools, keyed by thread count and shared by all solvers.
+
+use std::sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError};
+
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use rayon::{ThreadPool, ThreadPoolBuilder};
+
+/// Pools kept at once; building another drops the least recently used from the cache.
+const MAX_CACHED_POOLS: usize = 4;
+
+/// The pools built by process `pid`, least recently used first.
+struct PoolCache {
+    pid: u32,
+    pools: Vec<(usize, Arc<ThreadPool>)>,
+}
+
+impl PoolCache {
+    const fn new(pid: u32) -> Self {
+        Self {
+            pid,
+            pools: Vec::new(),
+        }
+    }
+
+    /// The cached pool of `threads` workers, now the most recently used.
+    fn get(&mut self, threads: usize) -> Option<Arc<ThreadPool>> {
+        let index = self.pools.iter().position(|(n, _)| *n == threads)?;
+        let entry = self.pools.remove(index);
+        let pool = Arc::clone(&entry.1);
+        self.pools.push(entry);
+        Some(pool)
+    }
+
+    /// Cache `pool` unless a racing build got there first, returning the pool to use and
+    /// any pool evicted to make room.
+    fn insert(
+        &mut self,
+        threads: usize,
+        pool: Arc<ThreadPool>,
+    ) -> (Arc<ThreadPool>, Option<Arc<ThreadPool>>) {
+        if let Some(cached) = self.get(threads) {
+            return (cached, None);
+        }
+        let evicted = (self.pools.len() == MAX_CACHED_POOLS).then(|| self.pools.remove(0).1);
+        self.pools.push((threads, Arc::clone(&pool)));
+        (pool, evicted)
+    }
+}
+
+static POOLS: LazyLock<Mutex<PoolCache>> =
+    LazyLock::new(|| Mutex::new(PoolCache::new(std::process::id())));
+
+/// Lock this process's pool cache, ignoring poisoning (a panic cannot leave it
+/// observably corrupt).
+fn lock_cache() -> MutexGuard<'static, PoolCache> {
+    let mut cache = POOLS.lock().unwrap_or_else(PoisonError::into_inner);
+    let pid = std::process::id();
+    if cache.pid != pid {
+        // A forked child inherits the pools but not their workers, so a job would wait
+        // forever. Leak them: dropping one can block on a lock a lost worker held.
+        for (_, pool) in cache.pools.drain(..) {
+            std::mem::forget(pool);
+        }
+        cache.pid = pid;
+    }
+    cache
+}
+
+/// A fresh pool of `threads` named workers.
+fn build_pool(threads: usize) -> PyResult<ThreadPool> {
+    ThreadPoolBuilder::new()
+        .num_threads(threads)
+        .thread_name(|i| format!("pybamm-solve-{i}"))
+        .build()
+        .map_err(|e| {
+            PyRuntimeError::new_err(format!(
+                "could not build a rayon pool of {threads} thread(s): {e}"
+            ))
+        })
+}
+
+/// The shared pool of `threads` workers, built on first request.
+pub fn pool_for(threads: usize) -> PyResult<Arc<ThreadPool>> {
+    pool_for_with(threads, build_pool)
+}
+
+/// [`pool_for`], building a missing pool with `build`.
+fn pool_for_with(
+    threads: usize,
+    build: impl FnOnce(usize) -> PyResult<ThreadPool>,
+) -> PyResult<Arc<ThreadPool>> {
+    let cached = lock_cache().get(threads);
+    if let Some(pool) = cached {
+        return Ok(pool);
+    }
+    // Built unlocked: a fork while another thread holds the lock deadlocks the child.
+    let built = Arc::new(build(threads)?);
+    let (pool, evicted) = lock_cache().insert(threads, built);
+    // A solve still running on an evicted pool keeps it alive until it returns.
+    drop(evicted);
+    Ok(pool)
+}
+
+/// The cached pools as `{thread_count: pool identity}`, for tests.
+#[pyfunction]
+pub fn _pool_ids(py: Python<'_>) -> PyResult<Py<PyDict>> {
+    let dict = PyDict::new(py);
+    // Integer keys run no Python, so `set_item` cannot re-enter `pool_for` under the lock.
+    for (threads, pool) in &lock_cache().pools {
+        dict.set_item(threads, Arc::as_ptr(pool).addr())?;
+    }
+    Ok(dict.unbind())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A one-worker pool, standing in for one of any width.
+    fn small_pool() -> Arc<ThreadPool> {
+        Arc::new(build_pool(1).expect("pool build failed"))
+    }
+
+    fn cached_widths(cache: &PoolCache) -> Vec<usize> {
+        cache.pools.iter().map(|(threads, _)| *threads).collect()
+    }
+
+    #[test]
+    fn the_cache_keeps_only_the_most_recently_used_pools() {
+        let mut cache = PoolCache::new(std::process::id());
+        let first = small_pool();
+        assert!(cache.insert(1, Arc::clone(&first)).1.is_none());
+        for threads in 2..=MAX_CACHED_POOLS {
+            assert!(cache.insert(threads, small_pool()).1.is_none());
+        }
+
+        // Reading the oldest makes width 2 the least recently used.
+        let read = cache.get(1).expect("width 1 cached");
+        assert!(Arc::ptr_eq(&read, &first));
+        let newest = MAX_CACHED_POOLS + 1;
+        let (_, evicted) = cache.insert(newest, small_pool());
+        assert!(evicted.is_some(), "a full cache must evict to grow");
+        let widths = cached_widths(&cache);
+        assert_eq!(widths.len(), MAX_CACHED_POOLS);
+        assert!(!widths.contains(&2), "{widths:?}");
+        assert_eq!(widths[MAX_CACHED_POOLS - 2..], [1, newest]);
+    }
+
+    #[test]
+    fn a_racing_build_hands_back_the_pool_already_cached() {
+        let mut cache = PoolCache::new(std::process::id());
+        let first = small_pool();
+        cache.insert(2, Arc::clone(&first));
+        let (pool, evicted) = cache.insert(2, small_pool());
+        assert!(Arc::ptr_eq(&pool, &first));
+        assert!(evicted.is_none());
+        assert_eq!(cached_widths(&cache), [2]);
+    }
+
+    #[test]
+    fn a_pool_is_built_without_holding_the_cache_lock() {
+        // The only test on the process-wide cache, so no other test holds its lock.
+        let pool = pool_for_with(3, |threads| {
+            assert!(
+                POOLS.try_lock().is_ok(),
+                "the cache lock is held across the build"
+            );
+            build_pool(threads)
+        })
+        .expect("pool");
+        assert!(Arc::ptr_eq(&pool, &pool_for(3).expect("cached pool")));
+    }
+}