From 519f514d38d56990fb2f8676966613cc07bf76ca Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Tue, 22 Sep 2026 22:00:56 +0100 Subject: [PATCH] feat: pybamm.DiffsolSolver PreparedSolver bindings, pybamm.DiffsolSolver and its docs. --- CHANGELOG.md | 1 + docs/source/api/solvers/diffsol_solver.rst | 5 + docs/source/api/solvers/index.rst | 1 + docs/source/developer/rust_idaklu_build.md | 5 +- packages/pybamm-rust/Cargo.lock | 1 + packages/pybamm-rust/pybamm-core/src/error.rs | 8 + .../pybamm-core/src/solver/batch.rs | 32 + .../pybamm-core/src/solver/solve.rs | 25 + packages/pybamm-rust/pybamm-python/Cargo.toml | 4 +- .../pybamm-rust/pybamm-python/src/errors.rs | 17 +- packages/pybamm-rust/pybamm-python/src/lib.rs | 13 + .../pybamm-rust/pybamm-python/src/model.rs | 6 +- .../pybamm-python/src/rayon_pool.rs | 175 +++ .../pybamm-rust/pybamm-python/src/solver.rs | 516 ++++++++ packages/pybamm/hatch_build.py | 2 +- packages/pybamm/src/pybamm/__init__.py | 1 + packages/pybamm/src/pybamm/rust/__init__.py | 8 + packages/pybamm/src/pybamm/rust/_core.pyi | 155 +++ .../pybamm/src/pybamm/solvers/__init__.py | 1 + .../pybamm/src/pybamm/solvers/base_solver.py | 57 +- .../src/pybamm/solvers/diffsol_solver.py | 597 +++++++++ .../src/pybamm/solvers/idaklu_solver.py | 84 +- .../pybamm/src/pybamm/solvers/observation.py | 22 + .../pybamm/src/pybamm/solvers/solution.py | 48 +- .../test_convert_to_format_rust.py | 8 + .../test_diffsol_discontinuous_current.py | 216 +++ .../test_diffsol_initial_conditions.py | 131 ++ .../tests/integration/test_diffsol_spm.py | 293 +++++ .../integration/test_rust_observation.py | 119 +- packages/pybamm/tests/shared.py | 2 +- packages/pybamm/tests/strategies/__init__.py | 13 +- .../pybamm/tests/strategies/input_sweeps.py | 17 + .../tests/strategies/serialise_values.py | 6 + .../tests/unit/test_rust_function_api.py | 1 + .../test_diffsol_batched_inputs.py | 53 + .../test_diffsol_output_variables.py | 190 +++ .../test_diffsol_sensitivities.py | 999 ++++++++++++++ .../unit/test_solvers/test_diffsol_solver.py | 1166 +++++++++++++++++ .../unit/test_solvers/test_idaklu_solver.py | 22 +- .../test_solvers/test_native_observation.py | 23 + .../test_solvers/test_observation_layout.py | 19 + .../test_solvers/test_rust_gil_release.py | 84 ++ .../unit/test_solvers/test_rust_lowering.py | 19 +- .../tests/unit/test_solvers/test_solution.py | 20 + 44 files changed, 5050 insertions(+), 135 deletions(-) create mode 100644 docs/source/api/solvers/diffsol_solver.rst create mode 100644 packages/pybamm-rust/pybamm-python/src/rayon_pool.rs create mode 100644 packages/pybamm-rust/pybamm-python/src/solver.rs create mode 100644 packages/pybamm/src/pybamm/solvers/diffsol_solver.py create mode 100644 packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py create mode 100644 packages/pybamm/tests/integration/test_diffsol_initial_conditions.py create mode 100644 packages/pybamm/tests/integration/test_diffsol_spm.py create mode 100644 packages/pybamm/tests/strategies/input_sweeps.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ea4f696ec..ccfe051af8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/source/api/solvers/diffsol_solver.rst b/docs/source/api/solvers/diffsol_solver.rst new file mode 100644 index 0000000000..517845572e --- /dev/null +++ b/docs/source/api/solvers/diffsol_solver.rst @@ -0,0 +1,5 @@ +Diffsol Solver +============== + +.. autoclass:: pybamm.DiffsolSolver + :members: diff --git a/docs/source/api/solvers/index.rst b/docs/source/api/solvers/index.rst index 2bb15503bc..bd108708d3 100644 --- a/docs/source/api/solvers/index.rst +++ b/docs/source/api/solvers/index.rst @@ -4,6 +4,7 @@ Solvers .. toctree:: base_solver + diffsol_solver dummy_solver scipy_solver jax_solver diff --git a/docs/source/developer/rust_idaklu_build.md b/docs/source/developer/rust_idaklu_build.md index 8e81fa0eea..dbaf312f0e 100644 --- a/docs/source/developer/rust_idaklu_build.md +++ b/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 diff --git a/packages/pybamm-rust/Cargo.lock b/packages/pybamm-rust/Cargo.lock index 177327fbcf..2b4e96f92d 100644 --- a/packages/pybamm-rust/Cargo.lock +++ b/packages/pybamm-rust/Cargo.lock @@ -1332,6 +1332,7 @@ dependencies = [ "numpy", "pybamm-core", "pyo3", + "rayon", ] [[package]] diff --git a/packages/pybamm-rust/pybamm-core/src/error.rs b/packages/pybamm-rust/pybamm-core/src/error.rs index c95ebf8b61..fbc69d488c 100644 --- a/packages/pybamm-rust/pybamm-core/src/error.rs +++ b/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 }, diff --git a/packages/pybamm-rust/pybamm-core/src/solver/batch.rs b/packages/pybamm-rust/pybamm-core/src/solver/batch.rs index cfb3312a8e..87157c09b9 100644 --- a/packages/pybamm-rust/pybamm-core/src/solver/batch.rs +++ b/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, +) -> 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 + } + )); + } } diff --git a/packages/pybamm-rust/pybamm-core/src/solver/solve.rs b/packages/pybamm-rust/pybamm-core/src/solver/solve.rs index f62a6da16c..49a3663340 100644 --- a/packages/pybamm-rust/pybamm-core/src/solver/solve.rs +++ b/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 { + 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 { + 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. diff --git a/packages/pybamm-rust/pybamm-python/Cargo.toml b/packages/pybamm-rust/pybamm-python/Cargo.toml index 327f9fc972..2da2848701 100644 --- a/packages/pybamm-rust/pybamm-python/Cargo.toml +++ b/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 diff --git a/packages/pybamm-rust/pybamm-python/src/errors.rs b/packages/pybamm-rust/pybamm-python/src/errors.rs index 18ab87f001..f8ad00d525 100644 --- a/packages/pybamm-rust/pybamm-python/src/errors.rs +++ b/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()), + } } diff --git a/packages/pybamm-rust/pybamm-python/src/lib.rs b/packages/pybamm-rust/pybamm-python/src/lib.rs index b50360ec33..e1e7f48cfb 100644 --- a/packages/pybamm-rust/pybamm-python/src/lib.rs +++ b/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::()?; m.add_class::()?; + #[cfg(feature = "diffsol")] + { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(solver::default_solver_options, m)?)?; + m.add_function(wrap_pyfunction!(rayon_pool::_pool_ids, m)?)?; + } + Ok(()) } diff --git a/packages/pybamm-rust/pybamm-python/src/model.rs b/packages/pybamm-rust/pybamm-python/src/model.rs index d82ac0a307..2d89fccf81 100644 --- a/packages/pybamm-rust/pybamm-python/src/model.rs +++ b/packages/pybamm-rust/pybamm-python/src/model.rs @@ -80,9 +80,9 @@ pub struct CompiledModel { pub(crate) compiled: Arc, /// Workspaces for the direct `eval_*`/`assemble_*` helpers. scratch_pool: ScratchPool, - /// Retained derivation source for the bundle views. - graph: Py, - rhs_root: NodeId, + /// Retained derivation source for the bundle views and `PreparedSolver`. + pub(crate) graph: Py, + pub(crate) rhs_root: NodeId, output_roots: Vec, event_roots: Vec, /// The algebraic block's root and its strictly-ascending global state indices. diff --git a/packages/pybamm-rust/pybamm-python/src/rayon_pool.rs b/packages/pybamm-rust/pybamm-python/src/rayon_pool.rs new file mode 100644 index 0000000000..8149938523 --- /dev/null +++ b/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)>, +} + +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> { + 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, + ) -> (Arc, Option>) { + 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> = + 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 { + 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> { + 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, +) -> PyResult> { + 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> { + 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 { + Arc::new(build_pool(1).expect("pool build failed")) + } + + fn cached_widths(cache: &PoolCache) -> Vec { + 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"))); + } +} diff --git a/packages/pybamm-rust/pybamm-python/src/solver.rs b/packages/pybamm-rust/pybamm-python/src/solver.rs new file mode 100644 index 0000000000..88ca9f96f6 --- /dev/null +++ b/packages/pybamm-rust/pybamm-python/src/solver.rs @@ -0,0 +1,516 @@ +// PyO3 bindings require specific argument types that clippy flags incorrectly +#![allow(clippy::needless_pass_by_value)] + +use std::sync::Arc; + +use numpy::ndarray::{Array2, ShapeBuilder}; +use numpy::{ + IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods, +}; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use pybamm_core::CoreError; +use pybamm_core::solver::SolverOptions; +use pybamm_core::solver::batch::check_batch_widths; +use pybamm_core::solver::solve as core_solve; + +use crate::errors::core_err_to_py; +use crate::model::CompiledModel; +use crate::rayon_pool::pool_for; + +/// Declare `PySolverOptions` and both conversions from core's option list, so the +/// defaults dict carries exactly the keys an `options` dict must supply. +macro_rules! solver_options { + ($($(#[$doc:meta])* $field:ident: $ty:ty,)*) => { + /// Integrator tuning from `PyBaMM`'s `options` dict; every key is required. + #[derive(Debug, FromPyObject, IntoPyObject)] + #[pyo3(from_item_all)] + pub struct PySolverOptions { + $($field: $ty,)* + } + + impl From for SolverOptions { + fn from(options: PySolverOptions) -> Self { + Self { $($field: options.$field,)* } + } + } + + impl From for PySolverOptions { + fn from(options: SolverOptions) -> Self { + Self { $($field: options.$field,)* } + } + } + }; +} + +pybamm_core::for_each_solver_option!(solver_options); + +/// The integrator defaults, as the `options` dict `PyBaMM` overlays onto. +#[pyfunction] +pub fn default_solver_options(py: Python<'_>) -> PyResult> { + PySolverOptions::from(SolverOptions::default()).into_pyobject(py) +} + +/// BDF solver statistics exposed to Python. +#[derive(Debug, Clone)] +// Output-only type: opt out of the automatic `FromPyObject` derive. +#[pyclass(get_all, skip_from_py_object, module = "pybamm.rust")] +pub struct SolverStatistics { + number_of_steps: usize, + number_of_linear_solver_setups: usize, + number_of_nonlinear_solver_iterations: usize, + number_of_nonlinear_solver_fails: usize, + number_of_error_test_failures: usize, + number_of_linear_solver_setups_from_checkpoint: usize, + number_of_linear_solver_setups_from_first_convergence_fail: usize, + number_of_linear_solver_setups_from_second_convergence_fail: usize, + number_of_linear_solver_setups_from_error_test_fail: usize, + number_of_linear_solver_setups_from_step_success: usize, + problem_setup_time_secs: f64, + ic_time_secs: f64, + solver_setup_time_secs: f64, + integration_time_secs: f64, + sens_error_control_relaxed: bool, +} + +#[pymethods] +impl SolverStatistics { + fn __repr__(&self) -> String { + format!( + "SolverStatistics(steps={}, linear_setups={} \ + [checkpoint={}, 1st_conv_fail={}, 2nd_conv_fail={}, err_fail={}, heuristic={}], \ + nl_iters={}, nl_fails={}, err_fails={}, \ + problem_setup={:.3}ms, ic_time={:.3}ms, solver_setup={:.3}ms, \ + integration={:.3}ms)", + self.number_of_steps, + self.number_of_linear_solver_setups, + self.number_of_linear_solver_setups_from_checkpoint, + self.number_of_linear_solver_setups_from_first_convergence_fail, + self.number_of_linear_solver_setups_from_second_convergence_fail, + self.number_of_linear_solver_setups_from_error_test_fail, + self.number_of_linear_solver_setups_from_step_success, + self.number_of_nonlinear_solver_iterations, + self.number_of_nonlinear_solver_fails, + self.number_of_error_test_failures, + self.problem_setup_time_secs * 1000.0, + self.ic_time_secs * 1000.0, + self.solver_setup_time_secs * 1000.0, + self.integration_time_secs * 1000.0, + ) + } +} + +impl From for SolverStatistics { + fn from(s: core_solve::SolverStatistics) -> Self { + let counters = s.solver; + Self { + number_of_steps: counters.number_of_steps, + number_of_linear_solver_setups: counters.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations: counters.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails: counters.number_of_nonlinear_solver_fails, + number_of_error_test_failures: counters.number_of_error_test_failures, + number_of_linear_solver_setups_from_checkpoint: counters + .number_of_linear_solver_setups_from_checkpoint, + number_of_linear_solver_setups_from_first_convergence_fail: counters + .number_of_linear_solver_setups_from_first_convergence_fail, + number_of_linear_solver_setups_from_second_convergence_fail: counters + .number_of_linear_solver_setups_from_second_convergence_fail, + number_of_linear_solver_setups_from_error_test_fail: counters + .number_of_linear_solver_setups_from_error_test_fail, + number_of_linear_solver_setups_from_step_success: counters + .number_of_linear_solver_setups_from_step_success, + problem_setup_time_secs: s.problem_setup_time_secs, + ic_time_secs: s.ic_time_secs, + solver_setup_time_secs: s.solver_setup_time_secs, + integration_time_secs: s.integration_time_secs, + sens_error_control_relaxed: s.sens_error_control_relaxed, + } + } +} + +/// Result of a diffsol solve; payloads the solve was not asked for read as `None`. +#[derive(Debug)] +#[pyclass(module = "pybamm.rust")] +pub struct SolveOutcome { + /// Solver flag: 0 = success, 1 = root found + #[pyo3(get)] + flag: i32, + /// Time at which an event was triggered, if any + #[pyo3(get)] + t_event: Option, + /// Index of the termination event that ended the solve, in registration order + #[pyo3(get)] + event_index: Option, + #[pyo3(get)] + statistics: SolverStatistics, + t_vec: Vec, + /// Flat column-major trajectory of states, or of outputs when requested. + y_flat: Vec, + /// Flat row-derivative trajectory matching `y_flat`, when stored. + yp_flat: Option>, + sens_blocks: Option>>, + n_rows: usize, + n_times: usize, + y_event_vec: Option>, + y_start_vec: Vec, + y_start_sens: Option>>, + y_event_sens: Option>>, + y_cache: Option>>, + yp_cache: Option>>, +} + +#[pymethods] +impl SolveOutcome { + /// Time points as a numpy array. + #[getter] + fn t<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + PyArray1::from_slice(py, &self.t_vec) + } + + /// Trajectory matrix `(n_rows, n_times)` of states, or of outputs when requested. + #[getter] + fn y<'py>(&mut self, py: Python<'py>) -> PyResult>> { + take_cached_matrix( + py, + &mut self.y_cache, + &mut self.y_flat, + self.n_rows, + self.n_times, + ) + } + + /// Row-derivative matrix `(n_rows, n_times)`, or `None` without `store_yp` or for outputs. + #[getter] + fn yp<'py>(&mut self, py: Python<'py>) -> PyResult>>> { + take_cached_matrix_opt( + py, + &mut self.yp_cache, + &mut self.yp_flat, + self.n_rows, + self.n_times, + ) + } + + /// Per-parameter sensitivities in the flat layout of `y`, or `None` if not requested. + #[getter] + #[pyo3(name = "yS")] + fn param_sensitivities<'py>(&self, py: Python<'py>) -> Option>>> { + self.sens_blocks + .as_ref() + .map(|blocks| sens_blocks(blocks, py)) + } + + /// Full state where the trajectory ends, never an outputs row. + #[getter] + fn y_event<'py>(&self, py: Python<'py>) -> Option>> { + self.y_event_vec + .as_ref() + .map(|v| PyArray1::from_slice(py, v)) + } + + /// Full state after the initial-condition solve, never an outputs row. + #[getter] + fn y_start<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + PyArray1::from_slice(py, &self.y_start_vec) + } + + /// Per-parameter state sensitivities at `y_start`, or `None` if not requested. + #[getter] + #[pyo3(name = "yS_start")] + fn start_sensitivities<'py>(&self, py: Python<'py>) -> Option>>> { + self.y_start_sens + .as_ref() + .map(|blocks| sens_blocks(blocks, py)) + } + + /// Per-parameter state sensitivities at `y_event`, or `None` if either is absent. + #[getter] + #[pyo3(name = "yS_event")] + fn event_sensitivities<'py>(&self, py: Python<'py>) -> Option>>> { + self.y_event_sens + .as_ref() + .map(|blocks| sens_blocks(blocks, py)) + } +} + +impl From for SolveOutcome { + fn from(outcome: core_solve::SolveOutcome) -> Self { + Self { + flag: outcome.termination.as_flag(), + t_event: outcome.termination.event_time(), + event_index: outcome.termination.event_index(), + statistics: SolverStatistics::from(outcome.statistics), + n_times: outcome.t.len(), + t_vec: outcome.t, + y_flat: outcome.y, + yp_flat: outcome.yp, + sens_blocks: outcome.sensitivities, + n_rows: outcome.n_rows, + y_event_vec: outcome.y_event, + y_start_vec: outcome.y_start, + y_start_sens: outcome.y_start_sens, + y_event_sens: outcome.y_event_sens, + y_cache: None, + yp_cache: None, + } + } +} + +/// Move a column-major flat buffer into a cached `(nrows, ncols)` F-order array. +fn take_cached_matrix<'py>( + py: Python<'py>, + cache: &mut Option>>, + flat: &mut Vec, + nrows: usize, + ncols: usize, +) -> PyResult>> { + if let Some(cached) = cache { + return Ok(cached.bind(py).clone()); + } + let data = std::mem::take(flat); + let arr = Array2::from_shape_vec((nrows, ncols).f(), data) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + let py_arr = arr.into_pyarray(py); + *cache = Some(py_arr.clone().unbind()); + Ok(py_arr) +} + +/// [`take_cached_matrix`] for an optional flat buffer: `None` stays `None`. +fn take_cached_matrix_opt<'py>( + py: Python<'py>, + cache: &mut Option>>, + flat: &mut Option>, + nrows: usize, + ncols: usize, +) -> PyResult>>> { + flat.as_mut().map_or_else( + || Ok(None), + |flat| take_cached_matrix(py, cache, flat, nrows, ncols).map(Some), + ) +} + +/// The `dy0/dp` seed as a flat slice; a missing seed is the empty (all-zero) one. +/// +/// A seed without `sensitivities` is rejected rather than silently ignored. +fn y0_sens_slice<'a>( + y0_sens: Option<&'a PyReadonlyArray1<'_, f64>>, + sensitivities: bool, +) -> PyResult<&'a [f64]> { + match y0_sens { + Some(_) if !sensitivities => Err(seed_without_sensitivities()), + Some(arr) => Ok(arr.as_slice()?), + None => Ok(&[]), + } +} + +fn seed_without_sensitivities() -> PyErr { + PyValueError::new_err("y0_sens was given but sensitivities=False") +} + +/// Per-parameter flat sensitivity blocks as a list of 1-D numpy arrays. +fn sens_blocks<'py>(blocks: &[Vec], py: Python<'py>) -> Vec>> { + blocks + .iter() + .map(|blk| PyArray1::from_slice(py, blk)) + .collect() +} + +/// Absolute tolerance: one value shared by every state, or one per state. +/// +/// `PerState` comes first because a length-1 array would also extract as `Uniform`. +#[derive(Debug, FromPyObject)] +pub enum PyAtol<'py> { + PerState(PyReadonlyArray1<'py, f64>), + Uniform(f64), +} + +impl PyAtol<'_> { + /// Widen to the per-state vector the core takes. + fn into_vec(self, n_states: usize) -> Vec { + match self { + Self::PerState(atol) => atol.as_array().to_vec(), + Self::Uniform(atol) => vec![atol; n_states], + } + } +} + +/// Prepare-once/execute-many solver for repeated integrations of one model. +#[derive(Debug)] +#[pyclass(frozen, module = "pybamm.rust")] +pub struct PreparedSolver { + inner: core_solve::PreparedSolver, +} + +#[pymethods] +impl PreparedSolver { + /// Prepare a solver for repeated use with the given model and tolerances. + /// + /// `store_yp` also stores the state time derivatives (`result.yp`) for Hermite interpolation. + #[new] + // Keep `sens_atol_factor` equal to `core_solve::DEFAULT_SENS_ATOL_FACTOR`; a + // path default renders as `...` in `__text_signature__`, which stubtest rejects. + #[pyo3(signature = (model, rtol=1e-6, atol=PyAtol::Uniform(1e-6), sens_atol_factor=1.0, options=None, store_yp=false))] + fn new( + py: Python<'_>, + model: &CompiledModel, + rtol: f64, + atol: PyAtol<'_>, + sens_atol_factor: f64, + options: Option, + store_yp: bool, + ) -> PyResult { + let compiled = Arc::clone(&model.compiled); + let atol_vec = atol.into_vec(compiled.n_states()); + let options = options.map_or_else(SolverOptions::default, SolverOptions::from); + + let mut prepared = core_solve::PreparedSolver::new(compiled, rtol, &atol_vec) + .and_then(|p| p.with_sens_atol_factor(sens_atol_factor)) + .and_then(|p| p.with_options(options)) + .map_err(core_err_to_py)? + .with_store_yp(store_yp); + // Only stored slopes read dF/dt, so other solves never compile it. + if store_yp { + let graph = model + .graph + .try_borrow(py) + .map_err(|_| PyValueError::new_err("PreparedSolver: graph borrow conflict"))?; + prepared = prepared.with_time_derivative(graph.arena(), model.rhs_root); + } + + Ok(Self { inner: prepared }) + } + + /// Solve the model over the given time span. + /// + /// Every `t_stop` restart time must also appear in `t_eval`; `y0_sens` is the + /// flattened, column-major `dy0/dp` seed. `num_threads` workers factorise the LU. + #[pyo3(signature = (t_eval, t_stop, y0, inputs, *, outputs=false, sensitivities=false, y0_sens=None, num_threads=1))] + #[allow(clippy::too_many_arguments)] + fn solve( + &self, + py: Python<'_>, + t_eval: PyReadonlyArray1<'_, f64>, + t_stop: PyReadonlyArray1<'_, f64>, + y0: PyReadonlyArray1<'_, f64>, + inputs: PyReadonlyArray1<'_, f64>, + outputs: bool, + sensitivities: bool, + y0_sens: Option>, + num_threads: usize, + ) -> PyResult { + let request = + solve_request(&t_eval, &t_stop, outputs, sensitivities)?.with_lu_threads(num_threads); + let set = core_solve::InputSet::new(y0.as_slice()?, inputs.as_slice()?) + .with_sens_seed(y0_sens_slice(y0_sens.as_ref(), sensitivities)?); + + let inner = &self.inner; + let outcome = py + .detach(|| { + if num_threads <= 1 { + return Ok(inner.solve(request, set)); + } + // `solve_batch`'s cached pool, so the LU never exceeds `num_threads`. + let pool = pool_for(num_threads)?; + Ok::<_, PyErr>(pool.install(|| inner.solve(request, set))) + })? + .map_err(core_err_to_py)?; + + Ok(SolveOutcome::from(outcome)) + } + + /// Solve every input set in `y0`/`inputs`, `num_threads` at a time. + /// + /// Arrays are C-contiguous with one row per set; returns one result per row, or the + /// unraised exception for a failed set. + #[pyo3(signature = (t_eval, t_stop, y0, inputs, num_threads, *, outputs=false, sensitivities=false, y0_sens=None))] + #[allow(clippy::too_many_arguments)] + fn solve_batch( + &self, + py: Python<'_>, + t_eval: PyReadonlyArray1<'_, f64>, + t_stop: PyReadonlyArray1<'_, f64>, + y0: PyReadonlyArray2<'_, f64>, + inputs: PyReadonlyArray2<'_, f64>, + num_threads: usize, + outputs: bool, + sensitivities: bool, + y0_sens: Option>, + ) -> PyResult>> { + let request = solve_request(&t_eval, &t_stop, outputs, sensitivities)?; + let y0_rows = batch_rows(&y0)?; + let inputs_rows = batch_rows(&inputs)?; + let seed_rows = batch_sens_rows(y0_sens.as_ref(), sensitivities, y0_rows.len())?; + let seeds_given = y0_sens.is_some().then_some(seed_rows.len()); + check_batch_widths(y0_rows.len(), inputs_rows.len(), seeds_given) + .map_err(core_err_to_py)?; + let sets: Vec> = y0_rows + .iter() + .zip(&inputs_rows) + .zip(&seed_rows) + .map(|((y0, inputs), seed)| core_solve::InputSet::new(y0, inputs).with_sens_seed(seed)) + .collect(); + + let inner = &self.inner; + let outcomes = py.detach(|| { + let pool = pool_for(num_threads)?; + Ok::<_, PyErr>(pool.install(|| inner.solve_batch(request, &sets))) + })?; + batch_entries(py, outcomes) + } +} + +/// Build the shared half of a solve from the arguments Python sent. +fn solve_request<'a>( + t_eval: &'a PyReadonlyArray1<'_, f64>, + t_stop: &'a PyReadonlyArray1<'_, f64>, + outputs: bool, + sensitivities: bool, +) -> PyResult> { + Ok(core_solve::SolveRequest { + t_eval: t_eval.as_slice()?, + t_stop: t_stop.as_slice()?, + outputs, + sensitivities, + lu_threads: 1, + }) +} + +/// One borrowed row per input set of a C-contiguous `(n_sets, width)` array. +fn batch_rows<'a>(array: &'a PyReadonlyArray2<'_, f64>) -> PyResult> { + let shape = array.shape(); + let (n_sets, width) = (shape[0], shape[1]); + // `as_slice` enforces C-contiguity, which makes the row split sound. + let flat = array.as_slice()?; + Ok((0..n_sets) + .map(|i| &flat[i * width..(i + 1) * width]) + .collect()) +} + +/// The per-set `dy0/dp` seeds, or `n_sets` empty (all-zero) seeds when omitted. +fn batch_sens_rows<'a>( + y0_sens: Option<&'a PyReadonlyArray2<'_, f64>>, + sensitivities: bool, + n_sets: usize, +) -> PyResult> { + match y0_sens { + Some(_) if !sensitivities => Err(seed_without_sensitivities()), + Some(array) => batch_rows(array), + None => Ok(vec![&[][..]; n_sets]), + } +} + +/// Convert per-set core outcomes into the Python list `solve_batch` returns. +fn batch_entries( + py: Python<'_>, + outcomes: Vec>, +) -> PyResult>> { + outcomes + .into_iter() + .map(|outcome| match outcome { + Ok(value) => Ok(Py::new(py, SolveOutcome::from(value))?.into_any()), + Err(error) => Ok(core_err_to_py(error).into_value(py).into_any()), + }) + .collect() +} diff --git a/packages/pybamm/hatch_build.py b/packages/pybamm/hatch_build.py index 444d9f7e5b..cccfe9814b 100644 --- a/packages/pybamm/hatch_build.py +++ b/packages/pybamm/hatch_build.py @@ -12,7 +12,7 @@ from hatchling.builders.hooks.plugin.interface import BuildHookInterface # One wheel per platform for CPython 3.10 and later, at a small per-FFI-call cost. -FEATURES = "pyo3/abi3-py310" +FEATURES = "pyo3/abi3-py310,diffsol" CRATE_MANIFEST = Path("pybamm-python") / "Cargo.toml" DESTINATION = Path("src") / "pybamm" / "rust" diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index e46e9485ee..5893015291 100644 --- a/packages/pybamm/src/pybamm/__init__.py +++ b/packages/pybamm/src/pybamm/__init__.py @@ -229,6 +229,7 @@ from .solvers.idaklu_jax import IDAKLUJax from .solvers.idaklu_solver import IDAKLUSolver +from .solvers.diffsol_solver import DiffsolSolver from .solvers.nonlinear_solver import NonlinearSolver # Experiments diff --git a/packages/pybamm/src/pybamm/rust/__init__.py b/packages/pybamm/src/pybamm/rust/__init__.py index 876f138b19..ba8fb6278b 100644 --- a/packages/pybamm/src/pybamm/rust/__init__.py +++ b/packages/pybamm/src/pybamm/rust/__init__.py @@ -28,6 +28,10 @@ def _global_symbol_visibility(): EvaluatorPool, Expr, ExprGraph, + PreparedSolver, + SolveOutcome, + SolverStatistics, + default_solver_options, ) __all__ = [ @@ -37,4 +41,8 @@ def _global_symbol_visibility(): "EvaluatorPool", "Expr", "ExprGraph", + "PreparedSolver", + "SolveOutcome", + "SolverStatistics", + "default_solver_options", ] diff --git a/packages/pybamm/src/pybamm/rust/_core.pyi b/packages/pybamm/src/pybamm/rust/_core.pyi index 254fab3430..84f0709db8 100644 --- a/packages/pybamm/src/pybamm/rust/_core.pyi +++ b/packages/pybamm/src/pybamm/rust/_core.pyi @@ -14,6 +14,11 @@ __all__ = [ "EvaluatorPool", "Expr", "ExprGraph", + "PreparedSolver", + "SolveOutcome", + "SolverStatistics", + "_pool_ids", + "default_solver_options", ] _FloatArray: TypeAlias = npt.NDArray[np.float64] @@ -25,6 +30,25 @@ _TimeGrid: TypeAlias = ( Sequence[float] | npt.NDArray[np.floating] | npt.NDArray[np.integer] ) +class _SolverOptions(TypedDict): + """Integrator tuning dict; every key is required (defaults live Python-side).""" + + max_nonlinear_solver_iterations: int + max_error_test_failures: int + max_nonlinear_solver_failures: int + nonlinear_solver_tolerance: float + min_timestep: float + max_timestep_growth: float | None + min_timestep_growth: float | None + max_timestep_shrink: float | None + min_timestep_shrink: float | None + update_jacobian_after_steps: int + update_rhs_jacobian_after_steps: int + threshold_to_update_jacobian: float + threshold_to_update_rhs_jacobian: float + pi_control_proportional: float + pi_control_integral: float + class _JacobianStats(TypedDict): """Assembly stats returned by :meth:`CompiledModel.jacobian_stats`.""" @@ -506,3 +530,134 @@ class EvaluatorPool: """Address of evaluator ``index``, handed out once; the pool must outlive it.""" def __len__(self) -> int: ... + +@final +class SolverStatistics: + """BDF solver statistics for one solve.""" + + @property + def number_of_steps(self) -> int: ... + @property + def number_of_linear_solver_setups(self) -> int: ... + @property + def number_of_nonlinear_solver_iterations(self) -> int: ... + @property + def number_of_nonlinear_solver_fails(self) -> int: ... + @property + def number_of_error_test_failures(self) -> int: ... + @property + def number_of_linear_solver_setups_from_checkpoint(self) -> int: ... + @property + def number_of_linear_solver_setups_from_first_convergence_fail(self) -> int: ... + @property + def number_of_linear_solver_setups_from_second_convergence_fail(self) -> int: ... + @property + def number_of_linear_solver_setups_from_error_test_fail(self) -> int: ... + @property + def number_of_linear_solver_setups_from_step_success(self) -> int: ... + @property + def problem_setup_time_secs(self) -> float: ... + @property + def ic_time_secs(self) -> float: ... + @property + def solver_setup_time_secs(self) -> float: ... + @property + def integration_time_secs(self) -> float: ... + @property + def sens_error_control_relaxed(self) -> bool: ... + +@final +class SolveOutcome: + """What one diffsol solve returns, whichever payloads it was asked for.""" + + @property + def flag(self) -> int: + """0 = success, 1 = root found.""" + + @property + def t_event(self) -> float | None: ... + @property + def event_index(self) -> int | None: + """Index of the termination event that ended the solve, in registration order.""" + + @property + def statistics(self) -> SolverStatistics: ... + @property + def t(self) -> _FloatArray: ... + @property + def y(self) -> _FloatArray: + """Trajectory ``(n_rows, n_times)`` of states, or of outputs when requested.""" + + @property + def yp(self) -> _FloatArray | None: + """Row derivatives ``(n_rows, n_times)``, or ``None`` without ``store_yp`` or for outputs.""" + + @property + def yS(self) -> list[_FloatArray] | None: + """Per-parameter sensitivities in the flat layout of ``y``, or ``None``.""" + + @property + def y_event(self) -> _FloatArray | None: + """The full state where the trajectory ends, never an outputs row.""" + + @property + def y_start(self) -> _FloatArray: + """The full state after the initial-condition solve, never an outputs row.""" + + @property + def yS_start(self) -> list[_FloatArray] | None: + """Per-parameter state sensitivities at ``y_start``, or ``None``.""" + + @property + def yS_event(self) -> list[_FloatArray] | None: + """Per-parameter state sensitivities at ``y_event``, or ``None``.""" + +@final +class PreparedSolver: + """Prepare-once/execute-many diffsol solver for repeated integrations.""" + + def __new__( + cls, + model: CompiledModel, + rtol: float = 1e-6, + atol: float | _FloatArray = ..., + sens_atol_factor: float = 1.0, + options: _SolverOptions | None = None, + store_yp: bool = False, + ) -> PreparedSolver: ... + def solve( + self, + t_eval: _FloatArray, + t_stop: _FloatArray, + y0: _FloatArray, + inputs: _FloatArray, + *, + outputs: bool = False, + sensitivities: bool = False, + y0_sens: _FloatArray | None = None, + num_threads: int = 1, + ) -> SolveOutcome: + """Integrate over ``t_eval``, landing exactly on each ``t_stop``. + + ``num_threads`` workers of the shared pool factorise the sparse LU. + """ + + def solve_batch( + self, + t_eval: _FloatArray, + t_stop: _FloatArray, + y0: _FloatArray, + inputs: _FloatArray, + num_threads: int, + *, + outputs: bool = False, + sensitivities: bool = False, + y0_sens: _FloatArray | None = None, + ) -> list[SolveOutcome | Exception]: + """Solve one row per set; a failed set yields its exception, unraised.""" + +def default_solver_options() -> _SolverOptions: + """Diffsol's own integrator defaults, the dict PyBaMM overlays onto.""" + +def _pool_ids() -> dict[int, int]: + """Cached rayon pools as ``{thread_count: pool identity}`` (test introspection).""" diff --git a/packages/pybamm/src/pybamm/solvers/__init__.py b/packages/pybamm/src/pybamm/solvers/__init__.py index 6b9c0821a6..1f59e632e5 100644 --- a/packages/pybamm/src/pybamm/solvers/__init__.py +++ b/packages/pybamm/src/pybamm/solvers/__init__.py @@ -4,6 +4,7 @@ "casadi_algebraic_solver", "casadi_solver", "composite_solver", + "diffsol_solver", "dummy_solver", "idaklu_jax", "idaklu_solver", diff --git a/packages/pybamm/src/pybamm/solvers/base_solver.py b/packages/pybamm/src/pybamm/solvers/base_solver.py index 58ad664a2e..ce57cee85a 100644 --- a/packages/pybamm/src/pybamm/solvers/base_solver.py +++ b/packages/pybamm/src/pybamm/solvers/base_solver.py @@ -29,8 +29,9 @@ class BaseSolver: The method to use for integration, specific to each solver rtol : float, optional The relative tolerance for the solver (default is 1e-6). - atol : float, optional - The absolute tolerance for the solver (default is 1e-6). + atol : float or array-like, optional + The absolute tolerance for the solver, one value or one per state + (default is 1e-6). Only IDAKLU and Diffsol solvers honour per-state values. root_method : str or pybamm algebraic solver class, optional The method to use to find initial conditions (for DAE solvers). Default is "nonlinear_solver", which uses a custom Newton solver for consistent @@ -258,6 +259,50 @@ def copy(self): new_solver._model_set_up = {} return new_solver + def _check_atol_type( + self, atol: float | list | tuple | np.ndarray, model: pybamm.BaseModel + ) -> np.ndarray: + """Widen an absolute tolerance to a float64 array with one value per state.""" + number_of_states = model.len_rhs_and_alg + if isinstance(atol, bool) or not isinstance( + atol, numbers.Real | list | tuple | np.ndarray + ): + raise pybamm.SolverError( + "Absolute tolerances must be a float, or a list, tuple or array " + f"of floats with one value per state, not {type(atol).__name__}" + ) + if isinstance(atol, numbers.Real): + return np.full(number_of_states, float(atol)) + + try: + values = np.asarray(atol) + except ValueError as error: + raise pybamm.SolverError( + "Absolute tolerances must be a flat list, tuple or array" + ) from error + # NumPy upcasts a bool beside numbers to 1.0 or 0.0, losing it + if isinstance(atol, list | tuple) and any( + isinstance(value, bool | np.bool_) + for value in np.asarray(atol, dtype=object).ravel() + ): + raise pybamm.SolverError( + "Absolute tolerances must be real numbers, not booleans" + ) + # Bool and complex arrays would otherwise cast to float silently + if values.dtype.kind not in "iuf": + raise pybamm.SolverError( + "Absolute tolerances must be real numbers, not an array of " + f"dtype {values.dtype}" + ) + if values.ndim == 0: + return np.full(number_of_states, float(values)) + if values.shape not in ((number_of_states,), (number_of_states, 1)): + raise pybamm.SolverError( + f"Absolute tolerances have shape {values.shape} but " + f"({number_of_states},) was expected (one value per state)" + ) + return np.ascontiguousarray(values.ravel(), dtype=np.float64) + def _lowers_via_compiled_model(self, model) -> bool: """Whether one shared Rust lowering serves every evaluator of this solve.""" return self._integrates_via_compiled_model and model.convert_to_format == "rust" @@ -1450,7 +1495,15 @@ def evaluate(function): return start_indices, end_indices, t_eval + @property + def _internal_initialisation(self) -> bool: + """Whether the integrator finds consistent initial conditions itself.""" + return False + def _check_event_violation_on_initialisation(self, *args, **kwargs): + # An integrator that initialises itself is handed only a guess. + if self._internal_initialisation: + return self._check_event_violation(*args, **kwargs) def _check_event_violation_post_solve(self, *args, **kwargs): diff --git a/packages/pybamm/src/pybamm/solvers/diffsol_solver.py b/packages/pybamm/src/pybamm/solvers/diffsol_solver.py new file mode 100644 index 0000000000..caf161867a --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/diffsol_solver.py @@ -0,0 +1,597 @@ +"""Solver class wrapping the Rust diffsol BDF integrator.""" + +import numbers +import sys + +import numpy as np +import numpy.typing as npt + +import pybamm +from pybamm.solvers.base_solver import ( + flatten_inputs, + validate_rust_sensitivity_widths, +) +from pybamm.solvers.idaklu_solver import _state_sensitivities +from pybamm.solvers.observation import ObserverCache, OutputAssembly +from pybamm.solvers.rust_lowering import RustModelLowering + +# Native handles and set-up products that the next set_up() rebuilds +_REBUILDABLE_STATE_KEYS = ( + "_prepared_solver", + "_rust_model", + "_output_assembly", + "_model_set_up", + "_time_integral_vars", +) + + +def _as_flat_float64(y0) -> npt.NDArray[np.float64]: + """Flatten an initial state to a float64 vector (handles ``casadi.DM``).""" + if hasattr(y0, "full"): + y0 = y0.full() + return np.asarray(y0, dtype=np.float64).flatten() + + +def _flatten_y0_sens(y0S, n_states: int, n_params: int) -> npt.NDArray[np.float64]: + """Flatten ``dy0/dp`` to the column-major block the Rust solver expects. + + Takes one ``(n_states, 1)`` or ``(n_states,)`` column per parameter, or one + ``(n_states, n_params)`` matrix; an empty input is the all-zero seed. + """ + blocks = [y0S] if hasattr(y0S, "full") or isinstance(y0S, np.ndarray) else list(y0S) + if not blocks: + # A shape complaint here would mask the solver's clearer error + return np.zeros(0, dtype=np.float64) + matrix = np.column_stack( + [ + np.asarray( + block.full() if hasattr(block, "full") else block, dtype=np.float64 + ) + for block in blocks + ] + ) + if matrix.shape != (n_states, n_params): + raise pybamm.SolverError( + f"Initial-condition sensitivities have shape {matrix.shape} but " + f"({n_states}, {n_params}) was expected (states x sensitivity " + "parameters)." + ) + return matrix.ravel(order="F") + + +class DiffsolSolver(pybamm.BaseSolver): + """Solve a discretised model using the Rust diffsol BDF solver. + + States are stored only at the output times, not at internal steps. Without + ``t_interp``, these are ``t_eval`` merged with a uniform 100-point grid over its + span; an explicit ``t_interp`` (merged with ``t_eval``) sets them exactly. + + Parameters + ---------- + rtol : float, optional + The relative tolerance for the solver (default is 1e-6). + atol : float or :class:`numpy.ndarray`, optional + The absolute tolerance, one value or one per state (default is 1e-6). + root_method : str or pybamm algebraic solver class, optional + The initial-condition method; if None (default), ``"nonlinear_solver"`` + for DAEs when ``calc_ic=False``, otherwise skipped. + root_tol : float, optional + The tolerance for the initial-condition solver. Default is 1e-6. + extrap_tol : float, optional + The tolerance to assert whether extrapolation occurs or not. + on_extrapolation : str, optional + What to do on extrapolation: "warn" (default), "error" or "ignore". + on_failure : str, optional + Only "error" (default): diffsol returns no partial solution, so an + integration failure always raises :class:`pybamm.SolverError`. + output_variables : list[str], optional + Variables to return; if empty (default), the full state is returned. + calc_ic : bool, optional + If True (default), diffsol finds consistent initial conditions itself; + if False, PyBaMM first solves for them with ``root_method``. + hermite_interpolation : bool, optional + If True (default), store state derivatives for cubic Hermite off-grid reads. + Ignored with ``output_variables`` or more than 4096 output points. + sens_atol_factor : float, optional + Multiplier on ``atol`` for the differential states' sensitivity tolerance + (default 1.0, matching :class:`IDAKLUSolver`); lower it for tighter gradients. + options : dict, optional + diffsol ``OdeSolverOptions`` overrides plus ``num_threads``: input sets solved + concurrently, or the threads a single set's sparse LU uses (worthwhile from + ~50k states). Defaults are in :attr:`DEFAULT_OPTIONS`; unknown keys and + values of the wrong type raise :class:`pybamm.SolverError`. + """ + + _integrates_via_compiled_model = True + # By design, a variable that cannot be lowered to Rust hard-fails on read + _observes_via_compiled_model = True + + #: diffsol's defaults, bar the per-solve ``max_nonlinear_solver_failures``: 50 + #: caps solve length, not divergence. Literal so ``import pybamm`` needs no extension. + _INTEGRATOR_DEFAULTS = { + "max_nonlinear_solver_iterations": 10, + "max_error_test_failures": 40, + "max_nonlinear_solver_failures": 100000, + "nonlinear_solver_tolerance": 0.2, + "min_timestep": 1e-13, + "max_timestep_growth": None, + "min_timestep_growth": None, + "max_timestep_shrink": None, + "min_timestep_shrink": None, + "update_jacobian_after_steps": 20, + "update_rhs_jacobian_after_steps": 50, + "threshold_to_update_jacobian": 0.3, + "threshold_to_update_rhs_jacobian": 0.2, + "pi_control_proportional": 0.0, + "pi_control_integral": 0.5, + } + + #: The integrator knobs plus ``num_threads``, which never reaches diffsol. + DEFAULT_OPTIONS = _INTEGRATOR_DEFAULTS | {"num_threads": 1} + + def __init__( + self, + rtol=1e-6, + atol=1e-6, + root_method=None, + root_tol=1e-6, + extrap_tol=None, + on_extrapolation=None, + on_failure=None, + output_variables=None, + calc_ic=True, + hermite_interpolation=True, + sens_atol_factor=1.0, + options=None, + ): + super().__init__( + "problem dependent", + rtol, + atol, + root_method, + root_tol, + extrap_tol, + on_extrapolation, + on_failure, + output_variables, + ) + self.name = "diffsol solver (bdf)" + self._calc_ic = calc_ic + self._hermite_interpolation = bool(hermite_interpolation) + self._supports_interp = True + self._supports_t_eval_discontinuities = True + self._options = self._combine_options(options) + + try: + factor = float(sens_atol_factor) + except (TypeError, ValueError) as exc: + raise pybamm.SolverError( + f"sens_atol_factor must be a finite number > 0, got {sens_atol_factor!r}" + ) from exc + if not np.isfinite(factor) or factor <= 0: + raise pybamm.SolverError( + f"sens_atol_factor must be a finite number > 0, got {sens_atol_factor!r}" + ) + self._sens_atol_factor = factor + + if root_method is None and not calc_ic: + self._use_default_root_method = True + else: + self._use_default_root_method = False + + @property + def on_failure(self) -> str: + return self._on_failure + + @on_failure.setter + def on_failure(self, value: str) -> None: + if value in ("warn", "ignore"): + raise pybamm.SolverError( + "DiffsolSolver returns no partial solution, so on_failure must be " + f"'error', not {value!r}" + ) + pybamm.BaseSolver.on_failure.__set__(self, value) + + @classmethod + def _combine_options(cls, user_options: dict | None) -> dict: + """Overlay ``user_options`` on :attr:`DEFAULT_OPTIONS`, validating each value.""" + options = cls._overlay_options( + cls.DEFAULT_OPTIONS, user_options, solver_name="diffsol" + ) + for name, default in cls._INTEGRATOR_DEFAULTS.items(): + options[name] = cls._checked_option(name, options[name], default) + options["num_threads"] = cls._checked_count( + "num_threads", options["num_threads"], minimum=1 + ) + return options + + @classmethod + def _checked_option(cls, name: str, value, default): + """Validate an integrator option as the type of its default.""" + if isinstance(default, int): + return cls._checked_count(name, value, minimum=0) + if value is None and default is None: + return None + if isinstance(value, bool) or not isinstance(value, numbers.Real): + kind = "a real number or None" if default is None else "a real number" + raise pybamm.SolverError(f"{name} must be {kind}, got {value!r}") + return float(value) + + @staticmethod + def _checked_count(name: str, value, minimum: int) -> int: + """Validate ``value`` as an integer from ``minimum`` to ``sys.maxsize``.""" + if ( + isinstance(value, bool) + or not isinstance(value, numbers.Integral) + or not minimum <= value <= sys.maxsize + ): + raise pybamm.SolverError( + f"{name} must be an integer from {minimum} to {sys.maxsize}, " + f"got {value!r}" + ) + return int(value) + + def _integrator_options(self) -> dict: + """The subset of ``self._options`` diffsol's ``OdeSolverOptions`` takes.""" + return {key: self._options[key] for key in self._INTEGRATOR_DEFAULTS} + + @property + def _internal_initialisation(self) -> bool: + """Return True if using native diffsol IC calculation.""" + return self._calc_ic + + def __getstate__(self): + state = self.__dict__.copy() + for key in _REBUILDABLE_STATE_KEYS: + state.pop(key, None) + return state + + def __setstate__(self, state): + self.__dict__.update(state) + # Read as "not yet set up", so the next solve() rebuilds the prepared solver + self._model_set_up = {} + + def set_up(self, model, inputs=None, t_eval=None, ics_only=False): + """Set up the solver, building the Rust compiled model.""" + if model.convert_to_format != "rust": + pybamm.logger.info( + f"Converting {model.name} to Rust for solving with DiffsolSolver" + ) + model.convert_to_format = "rust" + + if self._use_default_root_method and model.len_alg > 0: + self.root_method = "nonlinear_solver" + + base_set_up_return = super().set_up(model, inputs, t_eval, ics_only) + + if ics_only: + return base_set_up_return + + if isinstance(inputs, list): + inputs_dict = inputs[0] + else: + inputs_dict = inputs or {} + + self._build_rust_model(model, inputs_dict) + + return base_set_up_return + + def _build_rust_model(self, model, inputs_dict): + """Lower the model to a Rust ``CompiledModel`` and prepare the solver.""" + lowering = RustModelLowering(model, inputs_dict) + lowering.state_residual() + + if model.calculate_sensitivities: + validate_rust_sensitivity_widths(model, model.calculate_sensitivities) + lowering.sensitivity_indices(model.calculate_sensitivities) + + # Time-integral outputs lower to their integrand; the sum runs in the assembly + _, output_lens = lowering.outputs( + self.output_variables, time_integral_vars=self._time_integral_vars + ) + self._output_assembly = OutputAssembly.dense( + self.output_variables, + output_lens, + time_integrals=self._time_integral_vars, + ) + + lowering.termination_events() + + self._rust_model = lowering.compile() + lowering.bind_generic_evaluators(self._rust_model, root_method=self.root_method) + if self._observes_via_compiled_model: + ObserverCache.of(model).bind_rust_model(self._rust_model) + + from pybamm.rust import PreparedSolver + + try: + self._prepared_solver = PreparedSolver( + self._rust_model, + float(self.rtol), + self._check_atol_type(self.atol, model), + self._sens_atol_factor, + self._integrator_options(), + # Outputs-only solves store no state trajectory to Hermite between. + self._hermite_interpolation and not self.output_variables, + ) + except ValueError as error: + # The binding rejects an invalid tolerance or option with ValueError + raise pybamm.SolverError(str(error)) from error + + def _integrate( + self, + model, + t_eval, + inputs_list=None, + t_interp=None, + nproc=1, + ): + """Integrate the model using the diffsol BDF solver. + + ``t_eval`` and ``t_interp`` merge into one output grid, with ``t_eval`` also + the stop times. ``nproc`` is ignored: concurrency is the ``num_threads`` + option, as ``PreparedSolver`` cannot be pickled to worker processes. + """ + if not hasattr(self, "_prepared_solver"): + raise pybamm.SolverError("DiffsolSolver requires set_up() before solve()") + + inputs_list = inputs_list or [{}] + t_solve = np.asarray(self._output_times(t_eval, t_interp), dtype=np.float64) + t_stop = np.asarray(t_eval, dtype=np.float64) + y0S_list = getattr(model, "y0S_list", None) or [None] * len(model.y0_list) + + n_sets = len(inputs_list) + if self._options["num_threads"] > 1 and n_sets > 1: + results = self._solve_batch(model, t_solve, t_stop, inputs_list, y0S_list) + else: + results = [ + self._solve_one(model, t_solve, t_stop, inputs_dict, y0, y0S) + for inputs_dict, y0, y0S in zip( + inputs_list, model.y0_list, y0S_list, strict=True + ) + ] + + failures = [ + f"input set {index}: {result}" + for index, result in enumerate(results) + if isinstance(result, BaseException) + ] + if failures: + raise pybamm.SolverError("; ".join(failures)) + + return [ + self._build_solution(model, inputs_dict, result) + for inputs_dict, result in zip(inputs_list, results, strict=True) + ] + + def _payload_flags(self, model) -> dict[str, bool]: + """The ``outputs``/``sensitivities`` flags for ``PreparedSolver.solve``.""" + return { + "outputs": bool(self.output_variables), + "sensitivities": bool(model.calculate_sensitivities), + } + + def process_t_interp(self, t_interp, t_eval=None) -> npt.NDArray[np.float64]: + """Process ``t_interp``, defaulting to 100 points uniform over ``t_eval``. + + The default applies when ``t_interp`` is None, or empty with a two-point + ``t_eval``. + + Parameters + ---------- + t_interp : None, list or ndarray + The requested interpolation times. + t_eval : None, list or ndarray, optional + The requested stop times. + + Returns + ------- + ndarray + The interpolation times. + """ + requested = t_interp + t_interp = super().process_t_interp(t_interp, t_eval=t_eval) + if t_interp.size > 0 or t_eval is None or len(t_eval) < 2: + return t_interp + if requested is None or len(t_eval) == 2: + # Distant output times would make off-grid reads one long chord or cubic + return np.linspace(t_eval[0], t_eval[-1], 100) + return t_interp + + @staticmethod + def _output_times(t_eval, t_interp) -> npt.NDArray[np.float64]: + """The grid the solver reports its dense output on over one segment.""" + if t_interp is None or len(t_interp) == 0: + return np.asarray(t_eval, dtype=np.float64) + t_interp = np.asarray(t_interp, dtype=np.float64) + # A segment between discontinuities reports only its own share of t_interp + inside = (t_interp >= t_eval[0]) & (t_interp <= t_eval[-1]) + return np.union1d(t_eval, t_interp[inside]) + + def _solve_one(self, model, t_eval, t_stop, inputs_dict, y0, y0S=None): + """Integrate one input set, returning the raw Rust result or its failure. + + ``t_stop`` entries not also in ``t_eval`` are integrated through. + """ + y0_np = _as_flat_float64(y0) + y0_sens = None + if model.calculate_sensitivities and y0S is not None: + y0_sens = _flatten_y0_sens( + y0S, y0_np.size, len(model.calculate_sensitivities) + ) + + try: + return self._prepared_solver.solve( + t_eval, + t_stop, + y0_np, + flatten_inputs(inputs_dict), + y0_sens=y0_sens, + num_threads=self._options["num_threads"], + **self._payload_flags(model), + ) + except RuntimeError as error: + # Integration failures cross the FFI boundary as RuntimeError + return error + except ValueError as error: + # An invalid request, such as sensitivities the model was compiled without + raise pybamm.SolverError(str(error)) from error + + def _solve_batch(self, model, t_eval, t_stop, inputs_list, y0S_list): + """Integrate every input set concurrently; failed sets return their error.""" + y0 = np.vstack([_as_flat_float64(y0) for y0 in model.y0_list]) + inputs = np.vstack([flatten_inputs(d) for d in inputs_list]) + y0_sens = None + if model.calculate_sensitivities: + y0_sens = self._stack_y0_sens(model, y0S_list, y0.shape[1]) + + try: + return self._prepared_solver.solve_batch( + t_eval, + t_stop, + y0, + inputs, + # The configured width, so the process-wide pool cache is reused + self._options["num_threads"], + y0_sens=y0_sens, + **self._payload_flags(model), + ) + except ValueError as error: + raise pybamm.SolverError(str(error)) from error + + @staticmethod + def _stack_y0_sens(model, y0S_list, n_states): + """The batch's ``dy0/dp`` seeds, one row per set, or ``None`` if all zero.""" + if not model.calculate_sensitivities: + return None + n_params = len(model.calculate_sensitivities) + + def seed(y0S): + if y0S is None: + return None + flat = _flatten_y0_sens(y0S, n_states, n_params) + return flat if flat.size else None + + rows = [seed(y0S) for y0S in y0S_list] + if all(row is None for row in rows): + return None + zeros = np.zeros(n_states * n_params, dtype=np.float64) + return np.vstack([zeros if row is None else row for row in rows]) + + def _build_solution(self, model, inputs_dict, result): + """Turn one raw Rust result into a :class:`pybamm.Solution`.""" + t = result.t + if result.flag == 1 and t.size == 1: + # Stopped at t0 on an event violated after diffsol's own initialisation + self._check_event_violation(t, model, result.y_event, inputs_dict) + termination = "event" if result.flag == 1 else "final time" + t_event = None + if result.t_event is not None: + t_event = np.array([result.t_event]) + + if self.output_variables: + # The terminal full state, the only one outputs-only callers restart from + y_event = np.asarray(result.y_event).reshape(-1, 1) + sensitivity_names = model.calculate_sensitivities or [] + # The states have no sensitivities, as all_ys has no rows; the names stay + state_sensitivities = {} + if sensitivity_names: + state_sensitivities = { + name: np.zeros((0, 1)) for name in sensitivity_names + } + state_sensitivities["all"] = np.zeros((0, len(sensitivity_names))) + + sol = pybamm.Solution( + t, + # Zero-row, not None: experiment-step stitching slices all_ys + np.zeros((0, t.size)), + model, + inputs_dict, + t_event, + y_event, + termination, + all_sensitivities=state_sensitivities, + variables_returned=True, + ) + # first_state rebuilds from it, as all_ys holds no states + sol._y0 = result.y_start + sensitivities = None + if sensitivity_names: + # So first_state, last_state and a continued step have them + sol._y0_sensitivities = _state_sensitivities( + np.vstack(result.yS_start), sensitivity_names + ) + sol._y_event_sensitivities = _state_sensitivities( + np.vstack(result.yS_event), sensitivity_names + ) + sensitivities = self._output_assembly.stack_parameter_blocks( + result.yS, t.size, sensitivity_names + ) + # `result.y` holds output rows, output-major; the assembly wants time-major + self._output_assembly.attach( + sol, + np.asarray(result.y).T, + sensitivities=sensitivities, + sensitivity_names=sensitivity_names, + ) + else: + y = result.y # shape (n_states, n_times) from Rust + + y_event = None + if t_event is not None and result.y_event is not None: + y_event = np.asarray(result.y_event).reshape(-1, 1) + + yS_out = {} + if model.calculate_sensitivities: + sensitivity_names = model.calculate_sensitivities + # Bind yS once — the getter rebuilds the list on every access. + yS_list = result.yS + yS_out = { + name: np.asarray(yS_list[i]).reshape(-1, 1) + for i, name in enumerate(sensitivity_names) + } + yS_out["all"] = np.hstack([yS_out[name] for name in sensitivity_names]) + + sol = pybamm.Solution( + t, + y, + model, + inputs_dict, + t_event, + y_event, + termination, + all_sensitivities=yS_out, + all_yps=result.yp, + ) + + # Names the event, so BaseSolver need not re-evaluate them all at y_event + sol.closest_event_idx = result.event_index + + # Bind statistics once — the getter clones the struct on every access. + statistics = result.statistics + # Measured in Rust, as batched sets' wall clocks overlap + sol.integration_time = statistics.integration_time_secs + sol.solver_statistics = pybamm.SolverStatistics( + number_of_steps=statistics.number_of_steps, + number_of_linear_solver_setups=statistics.number_of_linear_solver_setups, + number_of_nonlinear_solver_iterations=statistics.number_of_nonlinear_solver_iterations, + number_of_nonlinear_solver_fails=statistics.number_of_nonlinear_solver_fails, + number_of_error_test_failures=statistics.number_of_error_test_failures, + number_of_linear_solver_setups_from_checkpoint=statistics.number_of_linear_solver_setups_from_checkpoint, + number_of_linear_solver_setups_from_first_convergence_fail=statistics.number_of_linear_solver_setups_from_first_convergence_fail, + number_of_linear_solver_setups_from_second_convergence_fail=statistics.number_of_linear_solver_setups_from_second_convergence_fail, + number_of_linear_solver_setups_from_error_test_fail=statistics.number_of_linear_solver_setups_from_error_test_fail, + number_of_linear_solver_setups_from_step_success=statistics.number_of_linear_solver_setups_from_step_success, + problem_setup_time_secs=statistics.problem_setup_time_secs, + ic_time_secs=statistics.ic_time_secs, + solver_setup_time_secs=statistics.solver_setup_time_secs, + sens_error_control_relaxed=statistics.sens_error_control_relaxed, + ) + if sol.solver_statistics.sens_error_control_relaxed: + pybamm.logger.warning( + "The diffsol sensitivity solve failed under error control and was " + "retried with sensitivities excluded from it, so their accuracy is " + "not controlled. The likely cause is a `sens_atol_factor` too small " + "for any step to meet; raise it if the gradients matter." + ) + return sol diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 3b1cec717b..094c394a96 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -1,5 +1,4 @@ # mypy: ignore-errors -import dataclasses import logging import numbers import warnings @@ -23,6 +22,15 @@ _UNSET = object() +# The SolverStatistics fields IDAS reports; the others keep their defaults +_IDAS_STATISTICS = ( + "number_of_steps", + "number_of_linear_solver_setups", + "number_of_nonlinear_solver_iterations", + "number_of_nonlinear_solver_fails", + "number_of_error_test_failures", +) + def _state_sensitivities(yS: np.ndarray, sensitivity_names: list[str]) -> dict: """Sensitivities of one state vector, keyed as in ``Solution``. @@ -477,70 +485,6 @@ def _check_options(self, options: dict): _check_klu_pivot_tolerance(options["klu_pivot_tolerance"]) - def _check_atol_type( - self, atol: float | list | tuple | np.ndarray, model: pybamm.BaseModel - ) -> np.ndarray: - """Widen an absolute tolerance to one value per state. - - Parameters - ---------- - atol : float, list, tuple or :class:`numpy.ndarray` - One tolerance for every state, or one per state with shape ``(n,)`` - or ``(n, 1)``, where ``n`` is the number of states. - model : :class:`pybamm.BaseModel` - The discretised model. - - Returns - ------- - :class:`numpy.ndarray` - Float64 tolerances with shape ``(n,)``. - - Raises - ------ - :class:`pybamm.SolverError` - If ``atol`` is not a real number, or a list, tuple or array of real - numbers with one value per state. - """ - number_of_states = model.len_rhs_and_alg - if isinstance(atol, bool) or not isinstance( - atol, numbers.Real | list | tuple | np.ndarray - ): - raise pybamm.SolverError( - "Absolute tolerances must be a float, or a list, tuple or array " - f"of floats with one value per state, not {type(atol).__name__}" - ) - if isinstance(atol, numbers.Real): - return np.full(number_of_states, float(atol)) - - try: - values = np.asarray(atol) - except ValueError as error: - raise pybamm.SolverError( - "Absolute tolerances must be a flat list, tuple or array" - ) from error - # NumPy upcasts a bool beside numbers to 1.0 or 0.0, losing it - if isinstance(atol, list | tuple) and any( - isinstance(value, bool | np.bool_) - for value in np.asarray(atol, dtype=object).ravel() - ): - raise pybamm.SolverError( - "Absolute tolerances must be real numbers, not booleans" - ) - # Bool and complex arrays would otherwise cast to float silently - if values.dtype.kind not in "iuf": - raise pybamm.SolverError( - "Absolute tolerances must be real numbers, not an array of " - f"dtype {values.dtype}" - ) - if values.ndim == 0: - return np.full(number_of_states, float(values)) - if values.shape not in ((number_of_states,), (number_of_states, 1)): - raise pybamm.SolverError( - f"Absolute tolerances have shape {values.shape} but " - f"({number_of_states},) was expected (one value per state)" - ) - return np.ascontiguousarray(values.ravel(), dtype=np.float64) - def _set_up_rust(self, model, inputs_dict, y0): """Set up solver with Rust evaluation backend.""" options = self._validate_rust_compatibility(model) @@ -1052,11 +996,6 @@ def _integrate( def _internal_initialisation(self) -> bool: return bool(self._options["calc_ic"]) - def _check_event_violation_on_initialisation(self, *args, **kwargs): - if self._internal_initialisation: - return - return self._check_event_violation(*args, **kwargs) - def _check_event_violation_post_solve(self, *args, **kwargs): if not self._internal_initialisation: return @@ -1172,10 +1111,7 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev newsol.integration_time = integration_time newsol.solver_statistics = pybamm.SolverStatistics( - **{ - field.name: getattr(sol.stats, field.name) - for field in dataclasses.fields(pybamm.SolverStatistics) - } + **{name: getattr(sol.stats, name) for name in _IDAS_STATISTICS} ) if not save_outputs_only: return newsol diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index 44b1805c02..2bf766788e 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -497,6 +497,28 @@ def attach( ) solution._variables[name] = variable + def stack_parameter_blocks( + self, + blocks: Sequence[npt.ArrayLike], + n_timesteps: int, + sensitivity_names: Sequence[str], + ) -> npt.NDArray[np.float64]: + """``(n_t, n_rows, n_p)`` sensitivities from one flat block per parameter. + + Each block holds ``n_t * n_rows`` values, time-outer/output-inner. + """ + if len(blocks) != len(sensitivity_names): + raise pybamm.SolverError( + f"Sensitivity block count mismatch: expected " + f"{len(sensitivity_names)} parameter blocks (from " + f"model.calculate_sensitivities) but the solver returned " + f"{len(blocks)}." + ) + return np.stack( + [np.asarray(block).reshape(n_timesteps, self.n_rows) for block in blocks], + axis=-1, + ) + def _checked_rows(self, data: npt.ArrayLike) -> np.ndarray: """``data`` as a ``(n_t, n_rows)`` array, or a complaint about its width.""" array = np.asarray(data) diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index cfdddddc29..2048b1aa0d 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -6,7 +6,7 @@ import json import numbers import pickle -from dataclasses import astuple, dataclass +from dataclasses import astuple, dataclass, fields from itertools import chain import casadi @@ -50,6 +50,27 @@ class SolverStatistics: Convergence failures of the nonlinear solver. number_of_error_test_failures : int Steps rejected by the local error test. + number_of_linear_solver_setups_from_checkpoint : int + Linear solver setups triggered by a checkpoint or reinitialisation. + number_of_linear_solver_setups_from_first_convergence_fail : int + Linear solver setups triggered by a first nonlinear convergence failure. + number_of_linear_solver_setups_from_second_convergence_fail : int + Linear solver setups triggered by a second nonlinear convergence failure. + number_of_linear_solver_setups_from_error_test_fail : int + Linear solver setups triggered by a local error test failure. + number_of_linear_solver_setups_from_step_success : int + Linear solver setups triggered after a successful step. + problem_setup_time_secs : float + Seconds spent building the problem, before any initialisation. + ic_time_secs : float + Seconds spent computing consistent initial conditions. + solver_setup_time_secs : float + Seconds spent creating the solver instance. + sens_error_control_relaxed : bool + Whether sensitivities were excluded from error control after a failure. + + Only :class:`pybamm.DiffsolSolver` reports the fields after + ``number_of_error_test_failures``. """ number_of_steps: int = 0 @@ -57,15 +78,38 @@ class SolverStatistics: number_of_nonlinear_solver_iterations: int = 0 number_of_nonlinear_solver_fails: int = 0 number_of_error_test_failures: int = 0 + number_of_linear_solver_setups_from_checkpoint: int = 0 + number_of_linear_solver_setups_from_first_convergence_fail: int = 0 + number_of_linear_solver_setups_from_second_convergence_fail: int = 0 + number_of_linear_solver_setups_from_error_test_fail: int = 0 + number_of_linear_solver_setups_from_step_success: int = 0 + problem_setup_time_secs: float = 0.0 + ic_time_secs: float = 0.0 + solver_setup_time_secs: float = 0.0 + sens_error_control_relaxed: bool = False def __add__(self, other: SolverStatistics) -> SolverStatistics: if not isinstance(other, SolverStatistics): return NotImplemented return SolverStatistics( - *(a + b for a, b in zip(astuple(self), astuple(other), strict=True)) + *( + (a or b) if isinstance(a, bool) else a + b + for a, b in zip(astuple(self), astuple(other), strict=True) + ) ) +def _set_solver_statistics_state(statistics: SolverStatistics, state: list) -> None: + # Pickles from before the diffsol fields carry only the leading values + for index, field in enumerate(fields(statistics)): + value = state[index] if index < len(state) else field.default + object.__setattr__(statistics, field.name, value) + + +# Python 3.10's slots=True would swap a frozen class's own __setstate__ for its default +SolverStatistics.__setstate__ = _set_solver_statistics_state + + # Shared by every first and last state, which do no solver work _NO_SOLVER_STATISTICS = SolverStatistics() diff --git a/packages/pybamm/tests/integration/test_convert_to_format_rust.py b/packages/pybamm/tests/integration/test_convert_to_format_rust.py index 953b20bd67..5bca7d5467 100644 --- a/packages/pybamm/tests/integration/test_convert_to_format_rust.py +++ b/packages/pybamm/tests/integration/test_convert_to_format_rust.py @@ -158,6 +158,14 @@ def test_default_root_method_keeps_the_rust_format(self): assert model.convert_to_format == "rust" assert isinstance(solver.root_method, pybamm.NonlinearSolver) + def test_diffsol_normalises_model_to_rust(self): + from tests.unit.test_solvers.test_process_rust import _toy_dae + + model = _toy_dae("casadi") + solver = pybamm.DiffsolSolver() + solver.solve(model, np.linspace(0, 1, 10), inputs={"a": 0.5}) + assert model.convert_to_format == "rust" + def test_dae_consistent_ic_parity_rust_newton_vs_casadi(self): from tests.unit.test_solvers.test_process_rust import _toy_dae diff --git a/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py b/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py new file mode 100644 index 0000000000..16379eb8ad --- /dev/null +++ b/packages/pybamm/tests/integration/test_diffsol_discontinuous_current.py @@ -0,0 +1,216 @@ +"""diffsol under a current profile with many discontinuities.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import pybamm + +PULSE_AMPLITUDE_A = 5.0 +PULSE_ON_S = 60.0 +PULSE_REST_S = 120.0 +PULSE_RAMP_S = 2.0 +DURATION_S = 1800.0 +SOLVER_TOL = 1e-6 + + +def _pulse_train_breakpoints() -> tuple[np.ndarray, np.ndarray]: + """Corners of a ramped pulse/rest train; finite ramps keep it Lipschitz.""" + times, values = [0.0], [0.0] + start = 0.0 + period = PULSE_ON_S + PULSE_REST_S + while start < DURATION_S: + times.extend( + [ + start + PULSE_RAMP_S, + start + PULSE_ON_S, + start + PULSE_ON_S + PULSE_RAMP_S, + start + period, + ] + ) + values.extend([PULSE_AMPLITUDE_A, PULSE_AMPLITUDE_A, 0.0, 0.0]) + start += period + times_arr = np.asarray(times, dtype=np.float64) + keep = times_arr <= DURATION_S + return times_arr[keep], np.asarray(values, dtype=np.float64)[keep] + + +def _pulse_train_parameter_values(inputs: dict[str, str] | None = None): + parameter_values = pybamm.ParameterValues("Chen2020") + times, values = _pulse_train_breakpoints() + parameter_values["Current function [A]"] = pybamm.Interpolant( + times, values, pybamm.t, interpolator="linear" + ) + for name, input_name in (inputs or {}).items(): + parameter_values[name] = pybamm.InputParameter(input_name) + return parameter_values + + +INFERENCE_INPUTS = { + "Negative particle diffusivity [m2.s-1]": "D_n", + "Positive particle diffusivity [m2.s-1]": "D_p", + "Negative electrode active material volume fraction": "eps_n", + "Positive electrode active material volume fraction": "eps_p", +} +INPUT_VALUES = {"D_n": 3.3e-14, "D_p": 4.0e-15, "eps_n": 0.75, "eps_p": 0.665} + + +def _solve( + solver, + calculate_sensitivities=False, + inputs=None, + options=None, + t_eval=None, + model_factory=pybamm.lithium_ion.DFN, +): + model = model_factory() + model.convert_to_format = "casadi" if solver == "casadi_idaklu" else "rust" + parameter_values = _pulse_train_parameter_values( + INFERENCE_INPUTS if inputs else None + ) + if solver == "diffsol": + instance = pybamm.DiffsolSolver( + rtol=SOLVER_TOL, atol=SOLVER_TOL, options=options + ) + else: + instance = pybamm.IDAKLUSolver(rtol=SOLVER_TOL, atol=SOLVER_TOL) + breakpoints, _ = _pulse_train_breakpoints() + simulation = pybamm.Simulation( + model, parameter_values=parameter_values, solver=instance + ) + return simulation.solve( + breakpoints if t_eval is None else t_eval, + t_interp=np.linspace(0.0, DURATION_S, 100), + initial_soc=0.5, + inputs=inputs, + calculate_sensitivities=calculate_sensitivities, + ) + + +class TestDiffsolPulseTrain: + def test_the_workload_still_costs_nonlinear_solver_failures(self): + # The failure-budget test below can only bite on a workload that spends some. + solution = _solve("diffsol") + assert solution.solver_statistics.number_of_nonlinear_solver_fails > 5 + + def test_a_long_pulse_train_completes(self): + solution = _solve("diffsol") + assert solution.t[-1] == pytest.approx(DURATION_S) + assert np.all(np.isfinite(np.asarray(solution["Voltage [V]"](solution.t)))) + + def test_the_failure_budget_option_reaches_diffsol(self): + with pytest.raises(pybamm.SolverError, match=r"nonlinear solver failures"): + _solve("diffsol", options={"max_nonlinear_solver_failures": 5}) + + def test_values_match_idaklu(self): + native = _solve("diffsol") + reference = _solve("casadi_idaklu") + t = np.linspace(0.0, DURATION_S, 100) + np.testing.assert_allclose( + np.asarray(native["Voltage [V]"](t)), + np.asarray(reference["Voltage [V]"](t)), + rtol=1e-4, + atol=1e-5, + ) + + def test_t_eval_breakpoints_are_stop_times_not_just_output_times(self): + # Equal cost either way would mean t_eval never reached the step control. + given = _solve("diffsol") + hidden = _solve("diffsol", t_eval=np.array([0.0, DURATION_S])) + assert ( + given.solver_statistics.number_of_nonlinear_solver_fails + < hidden.solver_statistics.number_of_nonlinear_solver_fails + ) + assert ( + given.solver_statistics.number_of_error_test_failures + < hidden.solver_statistics.number_of_error_test_failures + ) + + def test_gradients_hold_error_control_across_the_corners(self): + # Corners must be met with consistent sensitivity derivatives, not just states. + # Whether one solve this tight relaxes is down to rounding, so most must hold. + relaxed = [] + for k in range(6): + inputs = {**INPUT_VALUES, "eps_n": INPUT_VALUES["eps_n"] * (1 + k * 1e-9)} + solution = _solve( + "diffsol", calculate_sensitivities=sorted(inputs), inputs=inputs + ) + relaxed.append(solution.solver_statistics.sens_error_control_relaxed) + assert sum(relaxed) <= 2, relaxed + + def test_a_pure_ode_model_meets_the_corners_too(self): + # With no algebraic state diffsol's `set_consistent` returns early. + native = _solve("diffsol", model_factory=pybamm.lithium_ion.SPM) + reference = _solve("casadi_idaklu", model_factory=pybamm.lithium_ion.SPM) + t = np.linspace(0.0, DURATION_S, 100) + np.testing.assert_allclose( + np.asarray(native["Voltage [V]"](t)), + np.asarray(reference["Voltage [V]"](t)), + rtol=1e-4, + atol=1e-5, + ) + + def test_a_long_pulse_train_completes_with_sensitivities(self): + inputs = INPUT_VALUES + names = sorted(inputs) + solution = _solve("diffsol", calculate_sensitivities=names, inputs=inputs) + + assert solution.t[-1] == pytest.approx(DURATION_S) + for name in names: + gradient = np.asarray(solution["Voltage [V]"].sensitivities[name]).ravel() + assert np.all(np.isfinite(gradient)) + assert np.abs(gradient).max() > 0.0 + + +STEP_TIME_S = 1800.0 + + +def _solve_stepped_current(solver, t_eval): + """An SPM whose current steps down at ``STEP_TIME_S``.""" + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Current function [A]"] = lambda t: 5 * (t < STEP_TIME_S) + 1 + simulation = pybamm.Simulation( + pybamm.lithium_ion.SPM(), parameter_values=parameter_values, solver=solver + ) + return simulation.solve(t_eval) + + +class TestDiffsolSteppedCurrent: + """A Heaviside step adds a stop time to ``t_eval``, which leaves t0 coarse.""" + + @pytest.mark.parametrize( + "t_eval", + [[0, 2 * STEP_TIME_S], np.array([0, STEP_TIME_S, 2 * STEP_TIME_S])], + ids=["bare span", "three points"], + ) + @pytest.mark.parametrize( + ("hermite", "read_atol"), + [(True, 3e-4), (False, 3e-3)], + ids=["hermite", "linear"], + ) + def test_voltage_matches_idaklu(self, t_eval, hermite, read_atol): + native = _solve_stepped_current( + pybamm.DiffsolSolver( + rtol=SOLVER_TOL, atol=SOLVER_TOL, hermite_interpolation=hermite + ), + t_eval, + ) + reference = _solve_stepped_current( + pybamm.IDAKLUSolver(rtol=SOLVER_TOL, atol=SOLVER_TOL), + [0, 2 * STEP_TIME_S], + ) + np.testing.assert_allclose( + native["Voltage [V]"](native.t), + reference["Voltage [V]"](native.t), + rtol=0, + atol=1e-5, + ) + # Off the 100-point grid a read is only as good as its interpolant. + t = np.linspace(0.0, 2 * STEP_TIME_S, 50) + np.testing.assert_allclose( + native["Voltage [V]"](t), + reference["Voltage [V]"](t), + rtol=0, + atol=read_atol, + ) diff --git a/packages/pybamm/tests/integration/test_diffsol_initial_conditions.py b/packages/pybamm/tests/integration/test_diffsol_initial_conditions.py new file mode 100644 index 0000000000..bac19a6e46 --- /dev/null +++ b/packages/pybamm/tests/integration/test_diffsol_initial_conditions.py @@ -0,0 +1,131 @@ +"""``DiffsolSolver`` native initial conditions against PyBaMM's root solve.""" + +import re + +import numpy as np +import pytest + +import pybamm + +UNSTRUCTURED_VAR_PTS = { + "r_n": 5, + "r_p": 5, + "x_n": 10, + "x_s": 10, + "x_p": 10, + "y": 5, + "z": 5, +} + + +def _chen2020(c_rate): + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Current function [A]"] = ( + c_rate * parameter_values["Nominal cell capacity [A.h]"] + ) + return parameter_values + + +def _solve(calc_ic, parameter_values=None, model=None, var_pts=None, **solve_kwargs): + simulation = pybamm.Simulation( + model or pybamm.lithium_ion.DFN(), + parameter_values=parameter_values, + solver=pybamm.DiffsolSolver(calc_ic=calc_ic), + var_pts=var_pts, + ) + return simulation.solve(**solve_kwargs) + + +class TestDiffsolInitialConditions: + @pytest.mark.parametrize( + "parameter_values", + [None, _chen2020(1.0)], + ids=["Marquis2019", "Chen2020"], + ) + def test_native_initialisation_matches_the_root_solve(self, parameter_values): + t_eval = np.linspace(0, 3000, 51) + native = _solve(True, parameter_values, t_eval=[0, 3000], t_interp=t_eval) + python = _solve(False, parameter_values, t_eval=[0, 3000], t_interp=t_eval) + np.testing.assert_allclose( + native["Voltage [V]"](t_eval), python["Voltage [V]"](t_eval), atol=1e-5 + ) + + def test_off_grid_reads_in_the_first_interval_match_the_root_solve(self): + # Hermite reads in the first interval depend on the slope stored at t0. + grid = np.linspace(0, 600, 11) + t_check = np.array([2.0, 10.0, 17.8, 30.0, 50.0]) + native = _solve(True, t_eval=[0, 600], t_interp=grid) + python = _solve(False, t_eval=[0, 600], t_interp=grid) + np.testing.assert_allclose( + native["Voltage [V]"](t_check), python["Voltage [V]"](t_check), atol=1e-5 + ) + + def test_charging_from_empty_starts_at_the_consistent_state(self): + # The guessed potentials sit exactly on the 2.5 V cut-off; the consistent ones don't. + native = _solve(True, _chen2020(-1.0), t_eval=[0, 3000], initial_soc=0) + python = _solve(False, _chen2020(-1.0), t_eval=[0, 3000], initial_soc=0) + assert native.termination == "event: Maximum voltage [V]" + assert native.termination == python.termination + np.testing.assert_allclose(native.t[-1], python.t[-1], atol=1.0) + + @pytest.mark.parametrize("calc_ic", [True, False]) + def test_discharging_from_empty_is_infeasible_at_initial_conditions(self, calc_ic): + with pytest.raises( + pybamm.SolverError, + match=re.escape( + "Events ['Minimum voltage [V]'] are non-positive at initial conditions" + ), + ): + _solve(calc_ic, _chen2020(1.0), t_eval=[0, 600], initial_soc=0) + + @pytest.mark.parametrize("calc_ic", [True, False]) + def test_an_experiment_skips_a_step_infeasible_at_its_start(self, calc_ic): + # From a full cell the charge step's 4.1 V cut-off is already exceeded. + experiment = pybamm.Experiment( + [("Charge at 1C until 4.1 V", "Rest for 1 minute")] + ) + simulation = pybamm.Simulation( + pybamm.lithium_ion.DFN(), + parameter_values=pybamm.ParameterValues("Chen2020"), + experiment=experiment, + solver=pybamm.DiffsolSolver(calc_ic=calc_ic), + ) + solution = simulation.solve(initial_soc=1) + charge, rest = solution.cycles[0].steps + assert isinstance(charge, pybamm.EmptySolution) + assert charge.termination == "Event exceeded in initial conditions" + assert rest.termination == "final time" + + def test_an_outputs_only_first_state_is_the_consistent_state(self): + # The local ECM resistance reads the potentials, so it is ~0 at the guess. + name = "Local ECM resistance [Ohm]" + outputs = pybamm.Simulation( + pybamm.lithium_ion.DFN(), + solver=pybamm.DiffsolSolver(output_variables=["Voltage [V]"]), + ).solve([0, 60]) + full = _solve(True, t_eval=[0, 60]) + np.testing.assert_allclose( + outputs.first_state[name](0), full.first_state[name](0), rtol=1e-5 + ) + + def test_a_3d_model_initialises_natively(self): + model_options = {"dimensionality": 2} + var_pts = { + **pybamm.lithium_ion.BasicDFNUnstructured(model_options).default_var_pts, + **UNSTRUCTURED_VAR_PTS, + } + grid = np.linspace(0, 600, 11) + native, python = ( + _solve( + calc_ic, + model=pybamm.lithium_ion.BasicDFNUnstructured(model_options), + var_pts=var_pts, + t_eval=[0, 600], + t_interp=grid, + ) + for calc_ic in (True, False) + ) + assert native.termination == "final time" + np.testing.assert_allclose( + native["Voltage [V]"](grid), python["Voltage [V]"](grid), atol=1e-5 + ) diff --git a/packages/pybamm/tests/integration/test_diffsol_spm.py b/packages/pybamm/tests/integration/test_diffsol_spm.py new file mode 100644 index 0000000000..495b511985 --- /dev/null +++ b/packages/pybamm/tests/integration/test_diffsol_spm.py @@ -0,0 +1,293 @@ +"""Integration tests: SPM solved via the Rust diffsol BDF solver.""" + +import numpy as np +import pytest + +import pybamm + + +class TestDiffsolSPM: + def test_spm_basic_solve(self): + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 3600, 100) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + assert len(solution.t) > 0 + + voltage = solution["Voltage [V]"] + assert voltage is not None + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_spm_voltage_trajectory(self): + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 1800, 100) + sol = sim.solve(t_eval) + + voltage = sol["Voltage [V]"].entries + assert voltage[0] > 3.5, "Initial voltage should be > 3.5V" + assert voltage[-1] > 3.0, "Final voltage at 1800s should be > 3V" + assert voltage[-1] < voltage[0], "Voltage should decrease during discharge" + + def test_diffsol_matches_casadi(self): + model_ds = pybamm.lithium_ion.SPM() + model_cs = pybamm.lithium_ion.SPM() + model_cs.convert_to_format = "casadi" + + t_eval = np.linspace(0, 3600, 100) + + sol_ds = pybamm.Simulation( + model_ds, solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + ).solve(t_eval) + sol_cs = pybamm.Simulation( + model_cs, solver=pybamm.CasadiSolver(rtol=1e-8, atol=1e-8) + ).solve(t_eval) + + v_ds = sol_ds["Voltage [V]"].entries + v_cs = sol_cs["Voltage [V]"].entries + + np.testing.assert_allclose( + v_ds, + v_cs, + rtol=1e-6, + atol=1e-7, + err_msg="DiffsolSolver voltage differs from CasadiSolver", + ) + + def test_diffsol_explicit_nonlinear_solver_no_crash(self): + solver = pybamm.DiffsolSolver(root_method="nonlinear_solver") + assert solver.root_method is not None + assert solver.root_tol == 1e-6 + + def test_diffsol_dfn_default_solves(self): + model = pybamm.lithium_ion.DFN() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + voltage = solution["Voltage [V]"] + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_diffsol_explicit_casadi_dfn(self): + model = pybamm.lithium_ion.DFN() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6, root_method="casadi") + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + def test_diffsol_calc_ic_native_spm(self): + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6, calc_ic=True) + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 100, 10) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + + def test_diffsol_spm_unchanged_after_dae_fix(self): + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver() + sim = pybamm.Simulation(model, solver=solver) + + t_eval = np.linspace(0, 3600, 100) + solution = sim.solve(t_eval) + + assert solution is not None + assert solution.termination == "final time" + voltage = solution["Voltage [V]"] + assert voltage.entries.min() > 2.5 + assert voltage.entries.max() < 4.5 + + def test_diffsol_output_variables(self): + model = pybamm.lithium_ion.SPM() + solver = pybamm.DiffsolSolver(output_variables=["Voltage [V]"]) + sim = pybamm.Simulation(model, solver=solver) + + sol = sim.solve([0, 3600]) + + voltage = sol["Voltage [V]"].data + assert len(voltage) > 0 + assert voltage[0] > 3.0 + assert voltage[0] < 5.0 + + def test_diffsol_output_variables_with_event(self): + model = pybamm.lithium_ion.SPM() + model.events.append( + pybamm.Event( + "Voltage cutoff", + model.variables["Voltage [V]"] - 3.2, + pybamm.EventType.TERMINATION, + ) + ) + + solver = pybamm.DiffsolSolver(output_variables=["Voltage [V]"]) + sim = pybamm.Simulation(model, solver=solver) + + sol = sim.solve([0, 7200]) + + assert sol.termination.startswith("event") + assert sol.t[-1] < 7200 + voltage_end = sol["Voltage [V]"].data[-1] + assert abs(voltage_end - 3.2) < 0.01 + + +class TestDiffsolCubicOCP: + @staticmethod + def _cubic_ocp_parameter_values(): + pv = pybamm.ParameterValues("Chen2020") + sto = np.linspace(0.0, 1.0, 200) + un_cb = pv["Negative electrode OCP [V]"] + up_cb = pv["Positive electrode OCP [V]"] + un_data = np.array([float(un_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + up_data = np.array([float(up_cb(pybamm.Scalar(s)).evaluate()) for s in sto]) + + def make_cubic_ocp(x_data, y_data, name): + def ocp(stoich): + return pybamm.Interpolant( + x_data, y_data, stoich, name=name, interpolator="cubic" + ) + + return ocp + + pv["Negative electrode OCP [V]"] = make_cubic_ocp(sto, un_data, "Un") + pv["Positive electrode OCP [V]"] = make_cubic_ocp(sto, up_data, "Up") + return pv + + def test_cubic_ocp_diffsol_vs_casadi(self): + model = pybamm.lithium_ion.SPM() + model.events = [] + pv = self._cubic_ocp_parameter_values() + t_eval = np.linspace(0, 1800, 50) + + # Pass t_interp so both solvers output exactly on t_eval. + ref_model = model.new_copy() + ref_model.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + ref_model, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval, t_interp=t_eval) + + sol_diffsol = pybamm.Simulation( + model.new_copy(), + parameter_values=pv, + solver=pybamm.DiffsolSolver(), + ).solve(t_eval, t_interp=t_eval) + + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-4, + atol=1e-6, + err_msg="Diffsol cubic-OCP voltage differs from CasADi", + ) + + +class TestDiffsolECMNDInterpolants: + def test_ecm_diffsol_vs_casadi(self): + model = pybamm.equivalent_circuit.Thevenin() + model.events = [] + pv = pybamm.ParameterValues("ECM_Example") + t_eval = np.linspace(0, 600, 50) + + # Pass t_interp so both solvers output exactly on t_eval. + ref_model = model.new_copy() + ref_model.convert_to_format = "casadi" + sol_casadi = pybamm.Simulation( + ref_model, + parameter_values=pv, + solver=pybamm.IDAKLUSolver(), + ).solve(t_eval, t_interp=t_eval) + + sol_diffsol = pybamm.Simulation( + model.new_copy(), + parameter_values=pv, + solver=pybamm.DiffsolSolver(), + ).solve(t_eval, t_interp=t_eval) + + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"].data, + sol_casadi["Voltage [V]"].data, + rtol=1e-4, + atol=1e-6, + err_msg="Diffsol ECM ND-interpolant voltage differs from CasADi", + ) + + +class TestDiffsolExperiment: + def test_two_step_experiment_matches_idaklu(self): + experiment = pybamm.Experiment( + [ + "Discharge at 1C for 5 minutes", + "Rest for 5 minutes", + ] + ) + + def run(solver): + sim = pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment, solver=solver + ) + return sim.solve() + + sol_diffsol = run(pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8)) + sol_idaklu = run(pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8)) + + # A flat instruction list makes one cycle per instruction. + assert len(sol_diffsol.cycles) == 2 + assert sol_diffsol.t[-1] == pytest.approx(600.0) + # Off diffsol's grid this would measure interpolation error, not solver error. + t_common = sol_diffsol.t + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"](t_common), + sol_idaklu["Voltage [V]"](t_common), + rtol=1e-5, + atol=1e-6, + ) + + +class TestDiffsolCurrentSteps: + def test_a_pulsed_dfn_matches_idaklu(self): + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Current function [A]"] = lambda t: ( + 5.0 * ((t < 60) + (t >= 120) * (t < 180)) + ) + # Odd seconds stay off the corners, where the solvers' knots may differ. + t_interp = np.arange(1.0, 240.0, 2.0) + + def run(solver): + sim = pybamm.Simulation( + pybamm.lithium_ion.DFN(), + parameter_values=parameter_values, + solver=solver, + ) + return sim.solve([0, 240], t_interp=t_interp) + + sol_diffsol = run(pybamm.DiffsolSolver()) + sol_idaklu = run(pybamm.IDAKLUSolver()) + + assert sol_diffsol.termination == "final time" + np.testing.assert_allclose( + sol_diffsol["Voltage [V]"](t_interp), + sol_idaklu["Voltage [V]"](t_interp), + atol=1e-4, + ) diff --git a/packages/pybamm/tests/integration/test_rust_observation.py b/packages/pybamm/tests/integration/test_rust_observation.py index f317e91940..5686603771 100644 --- a/packages/pybamm/tests/integration/test_rust_observation.py +++ b/packages/pybamm/tests/integration/test_rust_observation.py @@ -26,6 +26,7 @@ def _native_idaklu(native=True): SOLVERS = [ ("idaklu", pybamm.IDAKLUSolver, {}), + ("diffsol", pybamm.DiffsolSolver, {}), ] NATIVE_IDS = [s[0] for s in SOLVERS] @@ -193,28 +194,69 @@ def test_experiment_parity(self): ) -def test_multistep_time_integral_parity(): - # A time integral over a multi-segment solution integrates once, not per segment. - def _build(native): +class TestDiffsolNativeObservation: + def test_diffsol_native_by_default(self): + t_eval = np.linspace(0, 3600, 51) + model = pybamm.lithium_ion.SPM() - model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( - model.variables["Current [A]"], pybamm.Scalar(0.0) + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve(t_eval=t_eval, t_interp=t_eval) + assert _native_caches(sol) is not None + + # Cross-solver comparison at default tolerances, hence rtol=1e-4. + cmodel = pybamm.lithium_ion.SPM() + cmodel.convert_to_format = "casadi" + csim = pybamm.Simulation(cmodel, solver=pybamm.IDAKLUSolver()) + csol = csim.solve(t_eval=t_eval, t_interp=t_eval) + np.testing.assert_allclose( + sol["Terminal voltage [V]"].entries, + csol["Terminal voltage [V]"].entries, + rtol=1e-4, + atol=1e-6, ) - solver = pybamm.IDAKLUSolver() - model.convert_to_format = "rust" - solver._observes_via_compiled_model = native - sim = pybamm.Simulation(model, solver=solver) - for _ in range(3): - sim.step(360) - return sim.solution - native_sol = _build(native=True) - casadi_sol = _build(native=False) - assert _native_caches(native_sol) is not None - nv = native_sol["Integrated current [A.s]"].entries - cv = casadi_sol["Integrated current [A.s]"].entries - assert nv.shape == (1,), f"expected single scalar, got shape {nv.shape}" - np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) + def test_diffsol_does_not_build_casadi_observation(self, monkeypatch): + import pybamm.solvers.solution as solution_mod + + model = pybamm.lithium_ion.SPM() + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + sol = sim.solve([0, 600]) + + called = {"casadi": False} + orig = solution_mod.Solution._convert_to_casadi + + def spy(self, *a, **k): + called["casadi"] = True + return orig(self, *a, **k) + + monkeypatch.setattr(solution_mod.Solution, "_convert_to_casadi", spy) + _ = sol["Terminal voltage [V]"] + assert called["casadi"] is False + + +class TestMultiStepTimeIntegralParity: + def test_multistep_time_integral_parity(self): + # A time integral over a multi-segment solution integrates once, not per segment. + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( + model.variables["Current [A]"], pybamm.Scalar(0.0) + ) + solver = pybamm.IDAKLUSolver() + model.convert_to_format = "rust" + solver._observes_via_compiled_model = native + sim = pybamm.Simulation(model, solver=solver) + for _ in range(3): + sim.step(360) + return sim.solution + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert _native_caches(native_sol) is not None + nv = native_sol["Integrated current [A.s]"].entries + cv = casadi_sol["Integrated current [A.s]"].entries + assert nv.shape == (1,), f"expected single scalar, got shape {nv.shape}" + np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) class TestDiscreteTimeSumParity: @@ -295,27 +337,28 @@ def build(native): ) -@pytest.mark.parametrize("solver_cls", [pybamm.IDAKLUSolver]) -def test_dense_trajectory_parity(solver_cls): - t_interp = np.linspace(0, 3600, 500) +class TestDenseTrajectoryParity: + @pytest.mark.parametrize("solver_cls", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver]) + def test_dense_trajectory_parity(self, solver_cls): + t_interp = np.linspace(0, 3600, 500) - def build(native): - model = pybamm.lithium_ion.SPMe() - model.convert_to_format = "rust" - solver = solver_cls() - solver._observes_via_compiled_model = native - return pybamm.Simulation(model, solver=solver).solve( - t_eval=[0, 3600], t_interp=t_interp - ) + def build(native): + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + solver = solver_cls() + solver._observes_via_compiled_model = native + return pybamm.Simulation(model, solver=solver).solve( + t_eval=[0, 3600], t_interp=t_interp + ) - native_sol = build(native=True) - casadi_sol = build(native=False) - assert _native_caches(native_sol) is not None - assert native_sol["Terminal voltage [V]"].entries.shape[-1] == len(t_interp) - for var in ("Terminal voltage [V]", "Electrolyte concentration [mol.m-3]"): - np.testing.assert_allclose( - native_sol[var].entries, casadi_sol[var].entries, rtol=1e-8, atol=1e-10 - ) + native_sol = build(native=True) + casadi_sol = build(native=False) + assert _native_caches(native_sol) is not None + assert native_sol["Terminal voltage [V]"].entries.shape[-1] == len(t_interp) + for var in ("Terminal voltage [V]", "Electrolyte concentration [mol.m-3]"): + np.testing.assert_allclose( + native_sol[var].entries, casadi_sol[var].entries, rtol=1e-8, atol=1e-10 + ) class TestFieldLayoutParity: diff --git a/packages/pybamm/tests/shared.py b/packages/pybamm/tests/shared.py index 19185bfcfb..3e50d6679a 100644 --- a/packages/pybamm/tests/shared.py +++ b/packages/pybamm/tests/shared.py @@ -434,7 +434,7 @@ def get_base_model_with_battery_geometry(**kwargs): return model -def get_broken_input_model(convert_to_format="casadi"): +def get_broken_input_model(convert_to_format="rust"): """Return the discretised model ``dv/dt = -sqrt(k) * v``. Its residual is NaN for any input ``k < 0``, whatever the state. diff --git a/packages/pybamm/tests/strategies/__init__.py b/packages/pybamm/tests/strategies/__init__.py index b0ddbb5fda..d0df33f464 100644 --- a/packages/pybamm/tests/strategies/__init__.py +++ b/packages/pybamm/tests/strategies/__init__.py @@ -1,4 +1,4 @@ -"""Hypothesis strategies and shared settings for the serialisation property tests. +"""Hypothesis strategies and shared settings for the property tests. ``serialisation_settings`` is applied as a decorator to the individual serialisation property/smoke tests rather than registered as a suite-wide @@ -11,7 +11,7 @@ import os -from hypothesis import settings +from hypothesis import HealthCheck, settings #: Shared Hypothesis settings for the serialisation property tests # Increase `max_examples` for a larger amount of randomised draws @@ -19,3 +19,12 @@ max_examples=500, deadline=None, ) + +#: Shared Hypothesis settings for properties whose example runs a solve. +solve_settings = settings( + max_examples=20, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) + +__all__ = ["serialisation_settings", "solve_settings"] diff --git a/packages/pybamm/tests/strategies/input_sweeps.py b/packages/pybamm/tests/strategies/input_sweeps.py new file mode 100644 index 0000000000..e1c023b640 --- /dev/null +++ b/packages/pybamm/tests/strategies/input_sweeps.py @@ -0,0 +1,17 @@ +"""Hypothesis strategies for sweeps over solver input sets.""" + +from __future__ import annotations + +import hypothesis.strategies as st + + +def decay_rate_sweeps( + min_sets: int = 2, max_sets: int = 8 +) -> st.SearchStrategy[list[float]]: + """Sweeps of distinct positive decay rates spanning two decades of solve cost.""" + return st.lists( + st.floats(min_value=0.05, max_value=5.0, allow_nan=False, allow_infinity=False), + min_size=min_sets, + max_size=max_sets, + unique=True, + ) diff --git a/packages/pybamm/tests/strategies/serialise_values.py b/packages/pybamm/tests/strategies/serialise_values.py index fece2bb239..80b1db557e 100644 --- a/packages/pybamm/tests/strategies/serialise_values.py +++ b/packages/pybamm/tests/strategies/serialise_values.py @@ -31,6 +31,7 @@ pybamm.CasadiAlgebraicSolver, pybamm.NonlinearSolver, pybamm.CompositeSolver, + pybamm.DiffsolSolver, ) @@ -87,6 +88,11 @@ def _strategy(draw): elif cls is pybamm.NonlinearSolver: kwargs["rtol"] = draw(_POSITIVE_TOL) kwargs["atol"] = draw(_POSITIVE_TOL) + elif cls is pybamm.DiffsolSolver: + kwargs["rtol"] = draw(_POSITIVE_TOL) + kwargs["atol"] = draw(_POSITIVE_TOL) + kwargs["root_tol"] = draw(_POSITIVE_TOL) + kwargs["calc_ic"] = draw(st.booleans()) elif cls is pybamm.CompositeSolver: # Draw two simple solver classes (possibly the same) for sub_solvers. sub_cls_a, sub_cls_b = draw( diff --git a/packages/pybamm/tests/unit/test_rust_function_api.py b/packages/pybamm/tests/unit/test_rust_function_api.py index cecd2c5142..6c081cc5f0 100644 --- a/packages/pybamm/tests/unit/test_rust_function_api.py +++ b/packages/pybamm/tests/unit/test_rust_function_api.py @@ -862,6 +862,7 @@ def test_public_names_resolve_through_pybamm_rust(self): "CompiledFunction", "CompiledModel", "ExprGraph", + "PreparedSolver", ): assert hasattr(pybamm.rust, name), name diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py new file mode 100644 index 0000000000..970c215e00 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_batched_inputs.py @@ -0,0 +1,53 @@ +"""Diffsol with several input sets: one solution per set.""" + +import numpy as np + +import pybamm + +_SOLVER_TOL = 1e-8 +_CURRENTS = [0.5, 1.5] + + +class TestDiffsolBatchedInputs: + def _build(self): + model = pybamm.lithium_ion.SPM() + params = model.default_parameter_values + params["Current function [A]"] = "[input]" + sim = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ) + sim.build() + return sim + + def test_returns_one_solution_per_input_set(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + assert isinstance(solutions, list) + assert len(solutions) == len(inputs) + + def test_each_solution_carries_its_own_input_set(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + for solution, current in zip(solutions, _CURRENTS, strict=True): + assert solution.all_inputs[0]["Current function [A]"] == np.array([current]) + + def test_distinct_input_sets_give_distinct_trajectories(self): + sim = self._build() + t_eval = np.linspace(0, 300, 11) + inputs = [{"Current function [A]": current} for current in _CURRENTS] + + solutions = sim.solver.solve(sim.built_model, t_eval, inputs=inputs) + + first = solutions[0]["Voltage [V]"](t_eval) + second = solutions[1]["Voltage [V]"](t_eval) + assert np.all(second < first) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py new file mode 100644 index 0000000000..e6edf820a0 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_output_variables.py @@ -0,0 +1,190 @@ +"""Diffsol ``output_variables``: layout parity against the full-state solve.""" + +import numpy as np +import pytest + +import pybamm + +# A vector variable between scalars: an ordinal-indexed reader shifts later rows. +_VECTOR_VAR = "Negative particle surface concentration [mol.m-3]" +_OUTPUT_VARIABLES = ["Voltage [V]", _VECTOR_VAR, "Current [A]"] + +_SOLVER_TOL = 1e-8 + + +class TestDiffsolOutputVariables: + @pytest.fixture(scope="class") + def solutions(self): + """Solve SPM twice: once output-only, once for the full state.""" + t_eval = np.linspace(0, 600, 11) + + model_outputs = pybamm.lithium_ion.SPM() + sol_outputs = pybamm.Simulation( + model_outputs, + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, + atol=_SOLVER_TOL, + output_variables=_OUTPUT_VARIABLES, + ), + ).solve(t_eval, t_interp=t_eval) + + model_full = pybamm.lithium_ion.SPM() + sol_full = pybamm.Simulation( + model_full, + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ).solve(t_eval, t_interp=t_eval) + + return sol_outputs, sol_full, t_eval + + def test_vector_output_keeps_every_component(self, solutions): + sol_outputs, sol_full, _ = solutions + + expected = np.asarray(sol_full[_VECTOR_VAR].entries) + actual = np.asarray(sol_outputs[_VECTOR_VAR].entries) + + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-8) + + def test_scalar_after_a_vector_is_not_shifted(self, solutions): + sol_outputs, sol_full, _ = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs["Current [A]"].entries), + np.asarray(sol_full["Current [A]"].entries), + rtol=1e-6, + atol=1e-8, + ) + + def test_scalar_before_a_vector_is_unaffected(self, solutions): + sol_outputs, sol_full, _ = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs["Voltage [V]"].entries), + np.asarray(sol_full["Voltage [V]"].entries), + rtol=1e-6, + atol=1e-8, + ) + + def test_vector_output_primal_and_sensitivity_layouts_agree(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + + model = pybamm.lithium_ion.SPM() + model.events = [] + sol = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, atol=_SOLVER_TOL, output_variables=[_VECTOR_VAR] + ), + ).solve( + np.linspace(0, 100, 6), inputs={"I": 0.5}, calculate_sensitivities=["I"] + ) + + variable = sol[_VECTOR_VAR] + assert variable.entries.size == variable.sensitivities["I"].shape[0] + + def test_output_rows_stay_aligned_across_the_batch_window(self): + """200 output points cross Rust's 128-lane staging window.""" + t_eval = np.linspace(0, 600, 200) + + sol_outputs = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + solver=pybamm.DiffsolSolver( + rtol=_SOLVER_TOL, + atol=_SOLVER_TOL, + output_variables=_OUTPUT_VARIABLES, + ), + ).solve(t_eval, t_interp=t_eval) + sol_full = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + solver=pybamm.DiffsolSolver(rtol=_SOLVER_TOL, atol=_SOLVER_TOL), + ).solve(t_eval, t_interp=t_eval) + + for name in _OUTPUT_VARIABLES: + np.testing.assert_allclose( + np.asarray(sol_outputs[name](t_eval)), + np.asarray(sol_full[name](t_eval)), + rtol=1e-6, + atol=1e-8, + err_msg=name, + ) + + @pytest.mark.parametrize("name", _OUTPUT_VARIABLES) + def test_output_variables_are_callable(self, solutions, name): + """Compared to the full-state solve: spatial calls add the boundary points.""" + sol_outputs, sol_full, t_eval = solutions + + np.testing.assert_allclose( + np.asarray(sol_outputs[name](t_eval)), + np.asarray(sol_full[name](t_eval)), + rtol=1e-6, + atol=1e-8, + ) + + +class TestDiffsolOutputVariablesAcrossExperimentSteps: + @staticmethod + def _solve(solver_cls, convert_to_format, **solver_kwargs): + experiment = pybamm.Experiment( + ["Discharge at 1C for 200 seconds", "Rest for 100 seconds"] + ) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format + return pybamm.Simulation( + model, + parameter_values=pybamm.ParameterValues("Chen2020"), + experiment=experiment, + solver=solver_cls(rtol=_SOLVER_TOL, atol=_SOLVER_TOL, **solver_kwargs), + ).solve() + + def test_experiment_with_output_variables_stitches(self): + sol = self._solve( + pybamm.DiffsolSolver, "rust", output_variables=["Voltage [V]"] + ) + assert len(sol.all_ts) == 2 + assert sol.variables_returned + assert all(y.shape[0] == 0 for y in sol.all_ys) + assert np.all(np.isfinite(np.asarray(sol["Voltage [V]"](sol.t)))) + + def test_experiment_outputs_match_the_full_state_solve(self): + sol_out = self._solve( + pybamm.DiffsolSolver, "rust", output_variables=["Voltage [V]"] + ) + sol_full = self._solve(pybamm.DiffsolSolver, "rust") + t = sol_full.t + np.testing.assert_allclose( + np.asarray(sol_out["Voltage [V]"](t)), + np.asarray(sol_full["Voltage [V]"](t)), + rtol=1e-6, + atol=1e-8, + ) + + def test_restarting_with_sensitivities_seeds_from_the_state_sensitivities(self): + def two_steps(output_variables): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values["Negative particle diffusivity [m2.s-1]"] = ( + pybamm.InputParameter("D_n") + ) + solver = pybamm.DiffsolSolver(output_variables=output_variables) + simulation = pybamm.Simulation( + model, parameter_values=parameter_values, solver=solver + ) + simulation.build() + kwargs = {"inputs": {"D_n": 3.3e-14}, "calculate_sensitivities": ["D_n"]} + first = solver.step( + pybamm.EmptySolution(), simulation.built_model, dt=100.0, **kwargs + ) + second = solver.step(first, simulation.built_model, dt=100.0, **kwargs) + return second["Voltage [V]"].sensitivities["D_n"] + + from_outputs = two_steps(["Voltage [V]"]) + from_states = two_steps(None) + np.testing.assert_allclose( + from_outputs, + from_states, + rtol=1e-4, + atol=1e-4 * np.max(np.abs(from_states)), + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py new file mode 100644 index 0000000000..57087d683e --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_sensitivities.py @@ -0,0 +1,999 @@ +"""Parity tests: diffsol native state sensitivities vs CasADi-IDAKLU oracle.""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import _native_caches + +_PARITY_RTOL = 1e-5 +_PARITY_ATOL = 1e-8 + +# Loosest tolerance keeping BDF/IDA trajectory differences within parity bounds. +_SCALAR_SOLVER_TOL = 1e-9 + +# Tighter for D_n, whose small |p * dV/dp| leaves little oracle margin. +_TWO_PARAM_SOLVER_TOL = 1e-11 + +# Spatial outputs: the concentration Jacobian amplifies cross-integrator error. +_SPATIAL_SOLVER_TOL = 1e-12 + + +class TestDiffsolSensitivities: + def _solve_both( + self, output_name, calc, extra_inputs=None, solver_tol=_SCALAR_SOLVER_TOL + ): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + if extra_inputs: + for pybamm_name, (input_name, value) in extra_inputs.items(): + params[pybamm_name] = pybamm.InputParameter(input_name) + inputs[input_name] = value + t_eval = np.linspace(0, 100, 15) + + m_native = pybamm.lithium_ion.SPM() + m_native.events = [] + m_native.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m_native, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=solver_tol, atol=solver_tol), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=calc, t_interp=t_eval) + + m_casadi = pybamm.lithium_ion.SPM() + m_casadi.events = [] + m_casadi.convert_to_format = "casadi" + # t_interp=t_eval puts both solutions on the same time grid. + sol_c = pybamm.Simulation( + m_casadi, + parameter_values=params, + solver=pybamm.IDAKLUSolver(rtol=solver_tol, atol=solver_tol), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=calc, t_interp=t_eval) + return sol_n[output_name], sol_c[output_name] + + def test_sensitivities_match_casadi_0d(self): + var_n, var_c = self._solve_both("Terminal voltage [V]", ["I"]) + sens_n, sens_c = var_n.sensitivities, var_c.sensitivities + assert set(sens_n) == set(sens_c) + assert sens_c["I"].size > 1 and np.any(sens_c["I"] != 0) + for key in sens_c: + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_match_casadi_spatial(self): + name = "Negative particle concentration [mol.m-3]" + var_n, var_c = self._solve_both(name, ["I"], solver_tol=_SPATIAL_SOLVER_TOL) + assert var_c.entries.shape[0] > 1 + sens_n, sens_c = var_n.sensitivities, var_c.sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_all_block_two_params(self): + var_n, var_c = self._solve_both( + "Terminal voltage [V]", + ["D_n", "I"], + extra_inputs={"Negative particle diffusivity [m2.s-1]": ("D_n", 3.3e-14)}, + solver_tol=_TWO_PARAM_SOLVER_TOL, + ) + a_n, a_c = var_n.sensitivities["all"], var_c.sensitivities["all"] + assert a_n.shape == a_c.shape and a_n.shape[1] == 2 + assert np.any(var_n.sensitivities["D_n"] != 0) + np.testing.assert_allclose( + a_n[:, 0], var_n.sensitivities["D_n"], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + for k in ("D_n", "I"): + np.testing.assert_allclose( + var_n.sensitivities[k], + var_c.sensitivities[k], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + np.testing.assert_allclose(a_n, a_c, rtol=_PARITY_RTOL, atol=_PARITY_ATOL) + + def test_state_sensitivity_matches_finite_difference(self): + # FD of diffsol's own values avoids cross-integrator noise and tight tols. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + i0, h = 0.5, 1e-4 + + def voltage(current): + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8), + ).solve(t_eval, inputs={"I": current}) + return sol[name].entries.ravel() + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8), + ).solve(t_eval, inputs={"I": i0}, calculate_sensitivities=["I"]) + analytic = sol[name].sensitivities["I"].ravel() + fd = (voltage(i0 + h) - voltage(i0 - h)) / (2 * h) + assert np.any(analytic != 0) + np.testing.assert_allclose(analytic, fd, rtol=2e-3, atol=1e-6) + + def test_output_variable_sensitivities_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL, output_variables=[name] + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval + ) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_subset_of_inputs_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL, output_variables=[name] + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval + ) + + # Oracle via the IDAKLU state path: its output path mislabels a strict subset. + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + assert set(sol_n[name].sensitivities) == {"I", "all"} + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_state_sensitivities_subset_of_inputs_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + assert set(sol_n[name].sensitivities) == {"I", "all"} + np.testing.assert_allclose( + sol_n[name].sensitivities["I"], + sol_c[name].sensitivities["I"], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + @staticmethod + def _build_time_integral_model(fmt): + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + integral = pybamm.ExplicitTimeIntegral(c, 0) ** 2 + model = pybamm.BaseModel(name="time_integral_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["integral"] = integral + model.variables["c"] = c + model.convert_to_format = fmt + pybamm.Discretisation().process_model(model) + return model + + def test_time_integral_sensitivities_match_casadi(self): + model_check = self._build_time_integral_model("casadi") + var = model_check.get_processed_variable_or_event("integral") + time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var(var, 2) + assert time_integral is not None + assert time_integral.post_sum_node is not None + + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + + # The squaring amplifies BDF-vs-IDA differences, so use the tight tolerance. + m_native = self._build_time_integral_model("rust") + sol_n = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + m_native, + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + assert _native_caches(sol_n) is not None + + m_casadi = self._build_time_integral_model("casadi") + sol_c = pybamm.IDAKLUSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + m_casadi, + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + + sens_n, sens_c = ( + sol_n["integral"].sensitivities, + sol_c["integral"].sensitivities, + ) + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_output_variables_time_integral_matches_the_full_state_solve(self): + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + solve_kwargs = { + "t_eval": [times[0], times[-1]], + "t_interp": times, + "inputs": inputs, + "calculate_sensitivities": ["a", "b"], + } + + sol_out = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, + atol=_SPATIAL_SOLVER_TOL, + output_variables=["integral"], + ).solve(self._build_time_integral_model("rust"), **solve_kwargs) + sol_full = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve(self._build_time_integral_model("rust"), **solve_kwargs) + + np.testing.assert_allclose( + np.asarray(sol_out["integral"].entries), + np.asarray(sol_full["integral"].entries), + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + sens_out, sens_full = ( + sol_out["integral"].sensitivities, + sol_full["integral"].sensitivities, + ) + for name in inputs: + np.testing.assert_allclose( + sens_out[name], + sens_full[name], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + err_msg=f"sensitivity mismatch for param '{name}'", + ) + + @staticmethod + def _build_two_integrals_model(fmt): + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + model = pybamm.BaseModel(name="two_integrals_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["sq"] = pybamm.ExplicitTimeIntegral(c, 0) ** 2 + model.variables["cube"] = pybamm.ExplicitTimeIntegral(c, 0) ** 3 + model.variables["c"] = c + model.convert_to_format = fmt + pybamm.Discretisation().process_model(model) + return model + + def test_distinct_time_integrals_same_n_inner_dont_alias(self): + times = np.linspace(0, 1, 15) + inputs = {"a": 0.7, "b": 1.0} + + sol_n = pybamm.DiffsolSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + self._build_two_integrals_model("rust"), + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + sol_c = pybamm.IDAKLUSolver( + rtol=_SPATIAL_SOLVER_TOL, atol=_SPATIAL_SOLVER_TOL + ).solve( + self._build_two_integrals_model("casadi"), + t_eval=[times[0], times[-1]], + t_interp=times, + inputs=inputs, + calculate_sensitivities=["a", "b"], + ) + + for var in ("sq", "cube"): + sens_n, sens_c = sol_n[var].sensitivities, sol_c[var].sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + np.testing.assert_allclose( + sens_n[key], sens_c[key], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + def test_sensitivities_with_event_termination_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 5.0} # high current trips the lower voltage cutoff early + t_eval = np.linspace(0, 3600, 50) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_n.termination.startswith("event") + + m_c = pybamm.lithium_ion.SPM() + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, atol=_SCALAR_SOLVER_TOL + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_c.termination.startswith("event") + + # Allow at most one grid point of difference in the event window. + len_n = sol_n[name].sensitivities["I"].shape[0] + len_c = sol_c[name].sensitivities["I"].shape[0] + assert abs(len_n - len_c) <= 1, ( + f"solvers disagree on event window length: diffsol={len_n}, idaklu={len_c}" + ) + n = min(len_n, len_c) + assert n > 1 + np.testing.assert_allclose( + sol_n[name].sensitivities["I"][:n], + sol_c[name].sensitivities["I"][:n], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_with_event_termination_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 5.0} + t_eval = np.linspace(0, 3600, 50) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.convert_to_format = "rust" + sol_n = pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + assert sol_n.termination.startswith("event") + + m_c = pybamm.lithium_ion.SPM() + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_SCALAR_SOLVER_TOL, + atol=_SCALAR_SOLVER_TOL, + output_variables=[name], + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"], t_interp=t_eval) + + len_n = sol_n[name].sensitivities["I"].shape[0] + len_c = sol_c[name].sensitivities["I"].shape[0] + assert abs(len_n - len_c) <= 1, ( + f"solvers disagree on event window length: diffsol={len_n}, idaklu={len_c}" + ) + n = min(len_n, len_c) + assert n > 1 + np.testing.assert_allclose( + sol_n[name].sensitivities["I"][:n], + sol_c[name].sensitivities["I"][:n], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + ) + + def test_output_variable_sensitivities_two_params_match_casadi(self): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + params["Negative particle diffusivity [m2.s-1]"] = pybamm.InputParameter("D_n") + # Insertion order differs from sorted order to expose column mislabelling. + inputs = {"I": 0.5, "D_n": 3.3e-14} + t_eval = np.linspace(0, 100, 15) + name = "Terminal voltage [V]" + + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + solver = pybamm.DiffsolSolver( + rtol=_TWO_PARAM_SOLVER_TOL, + atol=_TWO_PARAM_SOLVER_TOL, + output_variables=[name], + ) + sol_n = pybamm.Simulation(m, parameter_values=params, solver=solver).solve( + t_eval, inputs=inputs, calculate_sensitivities=["D_n", "I"], t_interp=t_eval + ) + + m_c = pybamm.lithium_ion.SPM() + m_c.events = [] + m_c.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_c, + parameter_values=params, + solver=pybamm.IDAKLUSolver( + rtol=_TWO_PARAM_SOLVER_TOL, + atol=_TWO_PARAM_SOLVER_TOL, + output_variables=[name], + ), + ).solve( + t_eval, inputs=inputs, calculate_sensitivities=["D_n", "I"], t_interp=t_eval + ) + + for param_name in ("I", "D_n"): + np.testing.assert_allclose( + sol_n[name].sensitivities[param_name], + sol_c[name].sensitivities[param_name], + rtol=_PARITY_RTOL, + atol=_PARITY_ATOL, + err_msg=f"sensitivity mislabelled for param '{param_name}'", + ) + + def test_discrete_time_sum_guard(self): + data_times = np.linspace(0, 1, 10) + ref = pybamm.BaseModel(name="ref") + c_ref = pybamm.Variable("c") + ref.rhs = {c_ref: -2 * c_ref} + ref.initial_conditions = {c_ref: 1} + ref.variables["c"] = c_ref + pybamm.Discretisation().process_model(ref) + data_values = ( + pybamm.IDAKLUSolver() + .solve(ref, t_eval=[data_times[0], data_times[-1]], t_interp=data_times)[ + "c" + ] + .entries + ) + data = pybamm.DiscreteTimeData(data_times, data_values, "test_data") + + c = pybamm.Variable("c") + c2 = pybamm.Variable("c2") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + model = pybamm.BaseModel(name="dts_model") + model.rhs = {c: b * -a * c, c2: -2 * c2} + model.initial_conditions = {c: 1, c2: 1} + model.variables["data_comparison"] = pybamm.DiscreteTimeSum((c - data) ** 2) + model.variables["c"] = c + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + # Solve on a mismatched grid (t_interp != data_times) so the guard fires. + sol = pybamm.DiffsolSolver().solve( + model, + t_eval=[0, 1], + t_interp=np.linspace(0, 1, 5), + inputs={"a": 0.7, "b": 1.0}, + calculate_sensitivities=["a", "b"], + ) + assert _native_caches(sol) is not None + with pytest.raises(pybamm.SolverError, match="discrete times"): + _ = sol["data_comparison"].sensitivities + + def test_calculate_sensitivities_rejects_vector_width_input_parameter(self): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + b = pybamm.InputParameter("b", expected_size=2) + model.rhs = {u: -(pybamm.Index(b, 0) + pybamm.Index(b, 1)) * u} + model.initial_conditions = {u: 1} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + with pytest.raises( + pybamm.SolverError, match=r"vector-width input parameters.*'b'" + ): + pybamm.DiffsolSolver().solve( + model, + t_eval=np.linspace(0, 1, 10), + inputs={"b": np.array([0.2, 0.3])}, + calculate_sensitivities=["b"], + ) + + def test_sensitivity_of_an_input_behind_a_vector_one(self): + """``k`` is at packed offset 2 but registration index 1, behind vector ``b``.""" + model = pybamm.BaseModel() + u = pybamm.Variable("u") + b = pybamm.InputParameter("b", expected_size=2) + k = pybamm.InputParameter("k") + total = pybamm.Index(b, 0) + pybamm.Index(b, 1) + model.rhs = {u: -total * k * u} + model.initial_conditions = {u: pybamm.Scalar(1)} + model.variables = {"u": u} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + t_eval = np.linspace(0, 2, 21) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + model, + t_eval=t_eval, + inputs={"b": np.array([0.2, 0.3]), "k": 2.0}, + calculate_sensitivities=["k"], + t_interp=t_eval, + ) + + # u = exp(-(b0 + b1) * k * t), so du/dk = -(b0 + b1) * t * u. + decay = 0.5 * 2.0 + np.testing.assert_allclose( + sol["u"](t_eval), np.exp(-decay * t_eval), rtol=1e-6, atol=1e-8 + ) + analytic = -0.5 * t_eval * np.exp(-decay * t_eval) + np.testing.assert_allclose( + sol["u"].sensitivities["k"].ravel(), analytic, rtol=1e-5, atol=1e-8 + ) + + def _chen_spm_params(self, extra): + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + for pybamm_name, input_name in extra.items(): + params[pybamm_name] = pybamm.InputParameter(input_name) + return params + + def _solve( + self, + fmt, + tol, + params, + inputs, + t_eval, + calc, + t_interp=None, + sens_atol_factor=None, + ): + model = pybamm.lithium_ion.SPM() + model.events = [] + model.convert_to_format = fmt + extra = ( + {} if sens_atol_factor is None else {"sens_atol_factor": sens_atol_factor} + ) + solver = ( + pybamm.DiffsolSolver(rtol=tol, atol=tol, **extra) + if fmt == "rust" + else pybamm.IDAKLUSolver(rtol=tol, atol=tol) + ) + kwargs = {"inputs": inputs, "calculate_sensitivities": calc} + # Both backends must land on the same grid for the arrays to be comparable. + kwargs["t_interp"] = t_eval if t_interp is None else t_interp + return pybamm.Simulation(model, parameter_values=params, solver=solver).solve( + t_eval, **kwargs + ) + + def test_weakly_influential_parameter_sensitivity_is_accurate(self): + # A state-sized atol floor swamps D_n's small |p * dV/dp|. + params = self._chen_spm_params( + {"Negative particle diffusivity [m2.s-1]": "D_n"} + ) + inputs = {"D_n": 3.3e-14, "I": 0.5} + t_eval = np.linspace(0, 100, 15) + calc = ["D_n", "I"] + + reference = self._solve("casadi", 1e-12, params, inputs, t_eval, calc) + native = self._solve( + "rust", 1e-9, params, inputs, t_eval, calc, sens_atol_factor=1e-3 + ) + + ref = np.asarray(reference["Terminal voltage [V]"].sensitivities["D_n"]).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities["D_n"]).ravel() + assert np.any(ref != 0) + np.testing.assert_allclose(got, ref, rtol=1e-6, atol=0.0) + + def test_early_transient_sensitivity_is_accurate(self): + # Error peaks near t=0.1s, which a uniform grid samples weakly; use a log grid. + params = self._chen_spm_params( + {"Positive electrode active material volume fraction": "eps_p"} + ) + inputs = {"I": 5.0, "eps_p": 0.665} + grid = np.unique( + np.concatenate([np.array([0.0]), np.logspace(-3, np.log10(30.0), 25)]) + ) + span = np.array([0.0, 30.0]) + calc = ["I", "eps_p"] + + reference = self._solve( + "casadi", 1e-11, params, inputs, span, calc, t_interp=grid + ) + native = self._solve( + "rust", + 1e-6, + params, + inputs, + span, + calc, + t_interp=grid, + sens_atol_factor=1e-3, + ) + + for key in calc: + ref = np.asarray( + reference["Terminal voltage [V]"].sensitivities[key] + ).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities[key]).ravel() + peak = np.max(np.abs(ref)) + assert peak > 0 + assert np.max(np.abs(got - ref)) / peak < 1e-6, ( + f"early-time {key} sensitivity error too large" + ) + + def test_lowering_sens_atol_factor_tightens_the_gradients(self): + params = self._chen_spm_params( + {"Negative particle diffusivity [m2.s-1]": "D_n"} + ) + inputs = {"D_n": 3.3e-14, "I": 0.5} + t_eval = np.linspace(0, 100, 15) + calc = ["D_n", "I"] + + reference = self._solve("casadi", 1e-12, params, inputs, t_eval, calc) + ref = np.asarray(reference["Terminal voltage [V]"].sensitivities["D_n"]).ravel() + assert np.any(ref != 0) + + factors = [1.0, 1e-1, 1e-2, 1e-3] + errors = [] + for factor in factors: + native = self._solve( + "rust", 1e-9, params, inputs, t_eval, calc, sens_atol_factor=factor + ) + got = np.asarray( + native["Terminal voltage [V]"].sensitivities["D_n"] + ).ravel() + # Peak-normalised: the column crosses zero, so element-wise ratios fail. + errors.append(float(np.max(np.abs(got - ref)) / np.max(np.abs(ref)))) + + # A tighter floor changes the step sequence, so a rung may be slightly worse. + for looser, tighter, error_looser, error_tighter in zip( + factors, factors[1:], errors, errors[1:], strict=False + ): + assert error_tighter <= 4.0 * error_looser, ( + f"sens_atol_factor {tighter} was much less accurate than " + f"{looser}: {error_tighter:.2e} against {error_looser:.2e}" + ) + # A flat ladder would pass the pairwise check above on its own. + assert errors[0] / errors[-1] > 10.0 + + def test_parameter_entering_initial_conditions_is_seeded(self): + name = "Maximum concentration in negative electrode [mol.m-3]" + params = self._chen_spm_params({name: "c_n_max"}) + inputs = {"I": 0.5, "c_n_max": 33133.0} + t_eval = np.linspace(0, 100, 15) + calc = ["I", "c_n_max"] + + reference = self._solve("casadi", 1e-9, params, inputs, t_eval, calc) + native = self._solve("rust", 1e-9, params, inputs, t_eval, calc) + + ref = np.asarray( + reference["Terminal voltage [V]"].sensitivities["c_n_max"] + ).ravel() + got = np.asarray( + native["Terminal voltage [V]"].sensitivities["c_n_max"] + ).ravel() + assert abs(ref[0]) > 0 + np.testing.assert_allclose(got, ref, rtol=1e-5, atol=0.0) + + def test_dy0_dp_seed_does_not_smear_across_columns(self): + # I never reaches y0, so its seed column is zero while c_n_max's is not. + name = "Maximum concentration in negative electrode [mol.m-3]" + params = self._chen_spm_params({name: "c_n_max"}) + inputs = {"I": 0.5, "c_n_max": 33133.0} + t_eval = np.linspace(0, 100, 15) + calc = ["I", "c_n_max"] + + reference = self._solve("casadi", 1e-9, params, inputs, t_eval, calc) + native = self._solve("rust", 1e-9, params, inputs, t_eval, calc) + + ref = np.asarray(reference["Terminal voltage [V]"].sensitivities["I"]).ravel() + got = np.asarray(native["Terminal voltage [V]"].sensitivities["I"]).ravel() + assert np.any(ref != 0) + np.testing.assert_allclose(got, ref, rtol=_PARITY_RTOL, atol=_PARITY_ATOL) + + @pytest.mark.parametrize("parameter_set", ["default", "Chen2020"]) + @pytest.mark.parametrize( + ("sens_atol_factor", "relaxed"), + [(1.0, False), (1e-30, True)], + ids=["controlled", "relaxed"], + ) + def test_dfn_sensitivities_solve(self, parameter_set, sens_atol_factor, relaxed): + model = pybamm.lithium_ion.DFN() + params = ( + model.default_parameter_values.copy() + if parameter_set == "default" + else pybamm.ParameterValues("Chen2020") + ) + sens = { + "Current function [A]": "I", + "Positive electrode active material volume fraction": "eps_p", + } + inputs = {} + for pybamm_name, input_name in sens.items(): + inputs[input_name] = float(params[pybamm_name]) + params[pybamm_name] = pybamm.InputParameter(input_name) + model.convert_to_format = "rust" + + grid = np.linspace(0.0, 600.0, 50) + solution = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=1e-6, atol=1e-6, sens_atol_factor=sens_atol_factor + ), + ).solve( + [float(grid[0]), float(grid[-1])], + t_interp=grid, + inputs=inputs, + calculate_sensitivities=sorted(sens.values()), + ) + + # No step meets a floor of 1e-30, so only the relaxed retry can finish. + assert solution.solver_statistics.sens_error_control_relaxed is relaxed + for key in sorted(sens.values()): + block = np.asarray(solution["Voltage [V]"].sensitivities[key]).ravel() + assert block.size > 1 + assert np.all(np.isfinite(block)) + assert np.any(block != 0) + + def test_dfn_sensitivities_hold_error_control_at_a_tighter_floor(self): + model = pybamm.lithium_ion.DFN() + params = pybamm.ParameterValues("Chen2020") + sens = { + "Current function [A]": "I", + "Positive electrode active material volume fraction": "eps_p", + } + inputs = {} + for pybamm_name, input_name in sens.items(): + inputs[input_name] = float(params[pybamm_name]) + params[pybamm_name] = pybamm.InputParameter(input_name) + model.convert_to_format = "rust" + + grid = np.linspace(0.0, 600.0, 50) + solution = pybamm.Simulation( + model, + parameter_values=params, + solver=pybamm.DiffsolSolver(rtol=1e-6, atol=1e-6, sens_atol_factor=1e-2), + ).solve( + [float(grid[0]), float(grid[-1])], + t_interp=grid, + inputs=inputs, + calculate_sensitivities=sorted(sens.values()), + ) + + assert solution.solver_statistics.sens_error_control_relaxed is False + for key in sorted(sens.values()): + block = np.asarray(solution["Voltage [V]"].sensitivities[key]).ravel() + assert np.all(np.isfinite(block)) + assert np.any(block != 0) + + def test_state_sensitivities_subset_selects_the_extreme_scale_input(self): + params = self._chen_spm_params( + {"Negative particle diffusivity [m2.s-1]": "D_n"} + ) + inputs = {"D_n": 3.3e-14, "I": 0.5} + t_eval = np.linspace(0, 100, 15) + + native = self._solve( + "rust", _TWO_PARAM_SOLVER_TOL, params, inputs, t_eval, ["D_n"] + ) + casadi = self._solve( + "casadi", _TWO_PARAM_SOLVER_TOL, params, inputs, t_eval, ["D_n"] + ) + + sens_n = native["Terminal voltage [V]"].sensitivities + sens_c = casadi["Terminal voltage [V]"].sensitivities + assert "D_n" in sens_n + assert "I" not in sens_n + assert np.any(np.asarray(sens_n["D_n"]) != 0) + np.testing.assert_allclose( + sens_n["D_n"], sens_c["D_n"], rtol=_PARITY_RTOL, atol=_PARITY_ATOL + ) + + +class TestDiffsolSensAtolFactor: + def test_default_sens_atol_factor(self): + # Factor 1.0 matches IDAS's IDASensEEtolerances (atol / pbar_i). + solver = pybamm.DiffsolSolver() + assert solver._sens_atol_factor == pytest.approx(1.0) + + def test_sens_atol_factor_is_configurable(self): + solver = pybamm.DiffsolSolver(sens_atol_factor=1e-2) + assert solver._sens_atol_factor == pytest.approx(1e-2) + + @pytest.mark.parametrize( + "bad", [0, -1.0, float("nan"), float("inf"), "not-a-number", None] + ) + def test_sens_atol_factor_rejects_invalid(self, bad): + with pytest.raises(pybamm.SolverError, match=r"sens_atol_factor"): + pybamm.DiffsolSolver(sens_atol_factor=bad) + + def test_sens_atol_factor_reaches_rust(self): + # A tighter floor forces more BDF steps under error control. + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + t_eval = np.linspace(0, 100, 15) + + def solve(sens_atol_factor): + m = pybamm.lithium_ion.SPM() + m.events = [] + m.convert_to_format = "rust" + return pybamm.Simulation( + m, + parameter_values=params, + solver=pybamm.DiffsolSolver( + rtol=1e-9, atol=1e-9, sens_atol_factor=sens_atol_factor + ), + ).solve(t_eval, inputs=inputs, calculate_sensitivities=["I"]) + + sol_default = solve(1e-3) + sol_relaxed = solve(1.0) + assert ( + sol_default.solver_statistics.number_of_steps + > sol_relaxed.solver_statistics.number_of_steps + ) + + +class TestFlattenY0Sens: + """``(n, 1)`` and bare ``(n,)`` ``dy0/dp`` columns become one column-major block.""" + + @staticmethod + def _flatten(*args, **kwargs): + return pybamm.solvers.diffsol_solver._flatten_y0_sens(*args, **kwargs) + + def test_a_list_of_column_vectors_is_column_major(self): + blocks = [np.array([[1.0], [2.0], [3.0]]), np.array([[4.0], [5.0], [6.0]])] + np.testing.assert_array_equal( + self._flatten(blocks, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_tuple_of_bare_columns_is_the_same_block(self): + blocks = (np.array([1.0, 2.0, 3.0]), np.array([4.0, 5.0, 6.0])) + np.testing.assert_array_equal( + self._flatten(blocks, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_single_matrix_is_the_same_block(self): + matrix = np.array([[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]]) + np.testing.assert_array_equal( + self._flatten(matrix, 3, 2), [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + + def test_a_casadi_column_is_densified(self): + casadi = pybamm.import_optional_dependency("casadi") + blocks = [casadi.DM([1.0, 2.0]), casadi.DM([3.0, 4.0])] + np.testing.assert_array_equal(self._flatten(blocks, 2, 2), [1.0, 2.0, 3.0, 4.0]) + + def test_an_empty_seed_is_the_all_zero_case(self): + assert self._flatten([], 3, 2).size == 0 + + def test_a_wrong_shape_is_a_solver_error(self): + # Silently reshaping would attach one parameter's seed to another. + with pytest.raises(pybamm.SolverError, match=r"\(3, 2\) was expected"): + self._flatten([np.array([1.0, 2.0, 3.0])], 3, 2) + + def test_a_wrong_state_count_is_a_solver_error(self): + with pytest.raises(pybamm.SolverError, match=r"\(3, 1\) was expected"): + self._flatten([np.array([1.0, 2.0])], 3, 1) + + +class TestDiffsolSensitivitiesAcrossExperimentSteps: + @staticmethod + def _solve(convert_to_format): + parameter_values = pybamm.ParameterValues("Chen2020") + name = "Negative particle diffusivity [m2.s-1]" + parameter_values[name] = pybamm.InputParameter("D_n") + model = pybamm.lithium_ion.SPM() + model.convert_to_format = convert_to_format + solver = ( + pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + if convert_to_format == "rust" + else pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8) + ) + return pybamm.Simulation( + model, + parameter_values=parameter_values, + experiment=pybamm.Experiment( + ["Discharge at 1C for 200 seconds", "Rest for 100 seconds"], + period="20 seconds", + ), + solver=solver, + ).solve(inputs={"D_n": 3.3e-14}, calculate_sensitivities=["D_n"]) + + def test_stepped_sensitivities_match_casadi(self): + native = self._solve("rust") + reference = self._solve("casadi") + assert len(native.all_ts) == 2 + + got = np.asarray(native["Voltage [V]"].sensitivities["D_n"]).ravel() + ref = np.asarray(reference["Voltage [V]"].sensitivities["D_n"]).ravel() + assert np.any(ref != 0) + assert got.shape == ref.shape + # Scaled by |p|: dV/dD_n is ~1e10, so a bare rtol reads as noise. + scale = np.abs(ref).max() + np.testing.assert_allclose(got / scale, ref / scale, rtol=0.0, atol=1e-4) diff --git a/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py b/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py new file mode 100644 index 0000000000..2a11e08b52 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_diffsol_solver.py @@ -0,0 +1,1166 @@ +"""General ``DiffsolSolver`` behaviour on small models.""" + +import multiprocessing +import pickle +import re + +import numpy as np +import pytest +from hypothesis import given + +import pybamm +from pybamm.solvers.observation import _native_caches +from pybamm.solvers.variable_observer import NativeObserver +from tests.shared import get_broken_input_model as _broken_input_model +from tests.strategies import solve_settings +from tests.strategies.input_sweeps import decay_rate_sweeps + + +def _decay_model(with_event=False, with_input=False): + """``dv/dt = -v`` (or ``-k*v``), so every value is analytically known.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + rate = -pybamm.InputParameter("k") * v if with_input else -v + model.rhs = {v: rate} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v, "2v": 2 * v} + if with_event: + model.events = [pybamm.Event("v = 0.5", v - 0.5)] + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +class TestDiffsolOutputsOnlyStates: + """Outputs-only solves must still expose observable first and last states.""" + + def _solve_outputs(self, model, t_eval): + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + return solver.solve(model, t_eval) + + def test_last_state_carries_the_terminal_state(self): + sol = self._solve_outputs(_decay_model(), np.linspace(0, 1, 5)) + assert sol.termination == "final time" + assert _native_caches(sol) is not None + last = sol.last_state + np.testing.assert_allclose(last["v"].data, np.exp(-1.0), rtol=1e-8) + np.testing.assert_allclose(last["2v"].data, 2 * np.exp(-1.0), rtol=1e-8) + + def test_first_state_rebuilds_from_initial_conditions(self): + sol = self._solve_outputs(_decay_model(), np.linspace(0, 1, 5)) + np.testing.assert_allclose(sol.first_state["v"].data, 1.0, rtol=1e-10) + + def test_event_terminated_last_state_is_the_root_state(self): + sol = self._solve_outputs(_decay_model(with_event=True), np.linspace(0, 2, 9)) + assert sol.termination.startswith("event") + np.testing.assert_allclose(sol.t[-1], np.log(2.0), rtol=1e-6) + np.testing.assert_allclose(sol.last_state["v"].data, 0.5, rtol=1e-6) + + def test_last_state_matches_the_full_state_solve(self): + t_eval = np.linspace(0, 1, 5) + sol_out = self._solve_outputs(_decay_model(), t_eval) + sol_full = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval + ) + # Integrator state vs dense interpolation: each errs at the solve tolerance. + for solution in (sol_out, sol_full): + np.testing.assert_allclose( + solution.last_state["v"].data, np.exp(-1.0), rtol=1e-9 + ) + np.testing.assert_allclose( + sol_out.last_state["v"].data, + sol_full.last_state["v"].data, + rtol=1e-8, + ) + + def _solve_outputs_with_sensitivities(self, model, t_eval): + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + return solver.solve( + model, t_eval, inputs={"k": 0.7}, calculate_sensitivities=True + ) + + def test_first_and_last_state_carry_state_sensitivities(self): + # v = exp(-k t), so dv/dk = -t exp(-k t) + sol = self._solve_outputs_with_sensitivities( + _decay_model(with_input=True), np.linspace(0, 1, 5) + ) + np.testing.assert_allclose( + sol.first_state["v"].sensitivities["k"], 0.0, atol=1e-12 + ) + np.testing.assert_allclose( + sol.last_state["v"].sensitivities["k"], -np.exp(-0.7), rtol=1e-6 + ) + + def test_event_terminated_last_state_carries_the_root_sensitivity(self): + sol = self._solve_outputs_with_sensitivities( + _decay_model(with_event=True, with_input=True), np.linspace(0, 2, 9) + ) + assert sol.termination.startswith("event") + t_root = np.log(2.0) / 0.7 + np.testing.assert_allclose( + sol.last_state["v"].sensitivities["k"], -t_root * 0.5, rtol=1e-5 + ) + + def test_a_sensitivity_solve_continues_from_an_outputs_only_step(self): + model = _decay_model(with_input=True) + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + kwargs = {"inputs": {"k": 0.7}, "calculate_sensitivities": True} + sol = solver.step(None, model, 1.0, t_interp=np.linspace(0, 1, 5), **kwargs) + sol = solver.step(sol, model, 1.0, t_interp=np.linspace(0, 1, 5), **kwargs) + # 2v = 2 exp(-k t), so d(2v)/dk = -4 exp(-2 k) at t = 2 + np.testing.assert_allclose( + sol["2v"].sensitivities["k"][-1], -4 * np.exp(-1.4), rtol=1e-6 + ) + + +def _rate_switch_model(t_switch): + """``dv/dt = -v``, doubling its rate at ``t_switch``.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: -(1 + (pybamm.t > t_switch)) * v} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +def _rate_switch_exact(t, t_switch): + return np.where(t <= t_switch, np.exp(-t), np.exp(-t_switch - 2 * (t - t_switch))) + + +class TestDiffsolTInterp: + def test_t_interp_merges_with_a_multi_point_t_eval(self): + t_eval = np.array([0.0, 0.5, 1.0]) + t_interp = np.linspace(0.0, 1.0, 21) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval, t_interp=t_interp + ) + np.testing.assert_array_equal(sol.t, np.union1d(t_eval, t_interp)) + np.testing.assert_allclose(sol["v"](t_interp), np.exp(-t_interp), rtol=1e-8) + + def test_a_multi_point_t_eval_is_merged_with_the_default_grid(self): + # Off-grid reads between distant stop times would be one long cubic. + t_eval = np.array([0.0, 0.5, 1.0]) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve(_decay_model(), t_eval) + np.testing.assert_array_equal( + sol.t, np.union1d(t_eval, np.linspace(0.0, 1.0, 100)) + ) + t_check = np.array([0.13, 0.31, 0.77]) + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=1e-8 + ) + + def test_empty_t_interp_keeps_t_eval(self): + t_eval = np.linspace(0, 1, 5) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _decay_model(), t_eval, t_interp=np.array([]) + ) + np.testing.assert_array_equal(sol.t, t_eval) + + def test_bare_span_t_eval_densifies_the_output_grid(self): + # A two-point solution would make every later sol[...](t) read a single chord. + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve(_decay_model(), [0, 1]) + np.testing.assert_array_equal(sol.t, np.linspace(0.0, 1.0, 100)) + np.testing.assert_allclose(sol["v"](sol.t), np.exp(-sol.t), rtol=1e-8) + + def test_bare_span_off_grid_reads_are_not_a_chord(self): + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve(_decay_model(), [0, 1]) + t_check = np.array([0.31, 0.5, 0.77]) + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=2e-5 + ) + + def test_bare_span_stays_dense_across_a_time_discontinuity(self): + # Its stop time must not make the span read as an explicit three-point grid. + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _rate_switch_model(0.5), [0, 1] + ) + assert np.isin(np.linspace(0.0, 1.0, 100), sol.t).all() + t_check = np.array([0.13, 0.49, 0.51, 0.77]) + np.testing.assert_allclose( + sol["v"](t_check), _rate_switch_exact(t_check, 0.5), rtol=1e-8 + ) + + def test_t_interp_is_shared_out_between_discontinuity_segments(self): + # An input-dependent switch time is solved as two segments. + t_interp = np.linspace(0.0, 1.0, 11) + sol = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10).solve( + _rate_switch_model(pybamm.InputParameter("t_switch")), + np.array([0.0, 1.0]), + inputs={"t_switch": 0.5}, + t_interp=t_interp, + ) + assert len(sol.sub_solutions) == 2 + assert np.all(np.diff(sol.t) > 0) + np.testing.assert_allclose( + sol["v"](t_interp), _rate_switch_exact(t_interp, 0.5), rtol=1e-5 + ) + + def test_outputs_only_rows_follow_the_merged_grid(self): + t_eval = np.array([0.0, 1.0]) + t_interp = np.linspace(0.0, 1.0, 33) + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + sol = solver.solve(_decay_model(), t_eval, t_interp=t_interp) + np.testing.assert_allclose( + sol["2v"](t_interp), 2.0 * np.exp(-t_interp), rtol=1e-8 + ) + + +class TestDiffsolHermiteInterpolation: + def _solve(self, model, t_eval, hermite=True, **solve_kwargs): + solver = pybamm.DiffsolSolver( + rtol=1e-10, atol=1e-10, hermite_interpolation=hermite + ) + return solver.solve(model, t_eval, **solve_kwargs) + + def test_off_grid_reads_beat_the_chord(self): + t_eval = np.linspace(0, 1, 6) + t_check = t_eval[:-1] + 0.1 # knot-interval midpoints + hermite = self._solve(_decay_model(), t_eval, t_interp=t_eval) + linear = self._solve(_decay_model(), t_eval, hermite=False, t_interp=t_eval) + hermite_error = np.max(np.abs(hermite["v"](t_check) - np.exp(-t_check))) + linear_error = np.max(np.abs(linear["v"](t_check) - np.exp(-t_check))) + assert hermite_error < 1e-5 + assert linear_error > 1e-3 + + def test_yp_matches_the_analytic_derivative(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5)) + assert sol.hermite_interpolation + # The t0 knot stores the consistent derivative; later knots are ~1e-6 off. + np.testing.assert_allclose(sol.yp[0], -np.exp(-sol.t), rtol=1e-5) + np.testing.assert_allclose(sol.yp[0][0], -1.0, rtol=1e-12) + + def test_disabling_hermite_reads_through_the_native_observer(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5), hermite=False) + assert not sol.hermite_interpolation + assert sol.all_yps is None + assert isinstance(sol["v"]._observer, NativeObserver) + np.testing.assert_allclose(sol["v"].entries, np.exp(-sol.t), rtol=1e-8) + + def test_a_diffsol_solve_leaves_how_an_idaklu_solution_is_read(self): + model = _decay_model() + idaklu = pybamm.IDAKLUSolver(options={"hermite_interpolation": False}) + earlier = idaklu.solve(model, np.linspace(0, 1, 5)) + before = earlier["v"] + + self._solve(model, np.linspace(0, 1, 5), hermite=False) + earlier.update("v") + + assert type(earlier["v"]) is type(before) + np.testing.assert_array_equal(earlier["v"].entries, before.entries) + + @pytest.mark.parametrize( + ("points", "stored"), [(4096, True), (4097, False)], ids=["limit", "past"] + ) + def test_a_dense_output_grid_gives_up_the_derivatives(self, points, stored): + # Past the limit the chord is already at the integration error floor. + t_eval = np.linspace(0, 1, points) + sol = self._solve(_decay_model(), t_eval, t_interp=t_eval) + assert sol.hermite_interpolation is stored + assert (sol.all_yps is not None) is stored + np.testing.assert_allclose(sol["v"].data, np.exp(-sol.t), rtol=1e-7) + + def test_a_dense_grid_reads_off_grid_to_the_solver_tolerance(self): + t_check = np.linspace(0.05, 0.95, 41) + sol = self._solve(_decay_model(), np.linspace(0, 1, 8193)) + assert not sol.hermite_interpolation + np.testing.assert_allclose(sol["v"](t_check), np.exp(-t_check), atol=1e-8) + + def test_hermite_observation_is_the_native_processed_variable(self): + sol = self._solve(_decay_model(), np.linspace(0, 1, 5)) + assert not isinstance(sol["v"], pybamm.ProcessedVariableComputed) + np.testing.assert_allclose(sol["v"].data, np.exp(-sol.t), rtol=1e-8) + + def test_outputs_only_solves_store_no_derivatives(self): + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10, output_variables=["2v"]) + sol = solver.solve(_decay_model(), np.linspace(0, 1, 5)) + assert sol.all_yps is None + assert _native_caches(sol) is not None + assert isinstance(sol.first_state["v"]._observer, NativeObserver) + + def test_the_event_root_column_carries_the_wound_back_slope(self): + # Coarse grid, so the root at ln 2 lands mid-interval. + sol = self._solve(_decay_model(with_event=True), np.linspace(0, 2, 5)) + assert sol.termination.startswith("event") + np.testing.assert_allclose(sol.yp[0, -1], -0.5, rtol=1e-6) + t_check = np.array([0.6, 0.69]) + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=1e-5 + ) + + def test_sensitivities_survive_the_native_hermite_path(self): + sol = self._solve( + _decay_model(with_input=True), + np.linspace(0, 1, 9), + inputs={"k": 0.5}, + calculate_sensitivities=True, + ) + t = sol.t + np.testing.assert_allclose( + sol["v"].sensitivities["k"], + -t * np.exp(-0.5 * t), + rtol=1e-5, + atol=1e-8, + ) + + def test_stepped_solve_keeps_hermite_across_segments(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10) + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=0.5, npts=6) + assert sol.hermite_interpolation + t_check = np.array([0.25, 0.75]) + np.testing.assert_allclose( + sol["v"](t_check), np.exp(-t_check), rtol=0, atol=1e-6 + ) + + def test_a_time_discontinuity_corner_is_not_smoothed(self): + # Each side of the corner's ULP bracket must use its own branch's slope. + solver = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8) + sol = solver.solve(_ramp_then_hold_model(), np.linspace(0, 10, 11)) + t_check = np.array([4.75, 5.25]) + np.testing.assert_allclose( + sol["v"](t_check), np.array([4.75, 5.0]), rtol=0, atol=1e-6 + ) + + +class TestDiffsolStep: + def test_stepped_solve_stays_observable_across_segments(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(rtol=1e-10, atol=1e-10) + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=0.5, npts=6) + + assert len(sol.sub_solutions) == 2 + assert _native_caches(sol) is not None + # One on-grid point per segment, so grid interpolation adds no error. + t_check = np.array([0.2, 0.7]) + np.testing.assert_allclose(sol["v"](t_check), np.exp(-t_check), rtol=1e-6) + np.testing.assert_allclose(sol.last_state["v"].data, np.exp(-1.0), rtol=1e-6) + + +class TestDiffsolTerminationEvent: + @pytest.mark.parametrize("output_variables", [None, ["2v"]]) + def test_the_solution_records_which_event_ended_it(self, output_variables): + # The second event is crossed first, so the first one is not a lucky default. + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: -v} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v, "2v": 2 * v} + model.events = [ + pybamm.Event("v = 0.1", v - 0.1), + pybamm.Event("v = 0.5", v - 0.5), + ] + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + solver = pybamm.DiffsolSolver(output_variables=output_variables) + solution = solver.solve(model, np.linspace(0, 3, 31)) + assert solution.closest_event_idx == 1 + assert solution.termination == "event: v = 0.5" + + +class TestDiffsolFailureHandling: + @pytest.mark.parametrize("on_failure", ["warn", "ignore"]) + def test_an_on_failure_needing_a_partial_solution_is_rejected(self, on_failure): + with pytest.raises(pybamm.SolverError, match=r"on_failure must be 'error'"): + pybamm.DiffsolSolver(on_failure=on_failure) + + def test_integrating_before_set_up_raises_solver_error(self): + with pytest.raises(pybamm.SolverError, match=r"requires set_up\(\)"): + pybamm.DiffsolSolver()._integrate(_decay_model(), np.array([0.0, 1.0])) + + def test_integration_failure_raises_solver_error(self): + # dv/dt = v^2 with v(0) = 1 blows up at t = 1. + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: v**2} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + with pytest.raises(pybamm.SolverError, match=r"diffsol error"): + pybamm.DiffsolSolver().solve(model, np.linspace(0, 2, 10)) + + def test_a_failed_sensitivity_retry_is_a_solver_error(self): + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: pybamm.InputParameter("k") * v**2} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + with pytest.raises( + pybamm.SolverError, match=r"^input set 0: sensitivity solve failed" + ): + pybamm.DiffsolSolver().solve( + model, + np.linspace(0, 2, 10), + inputs={"k": 1.0}, + calculate_sensitivities=["k"], + ) + + +def _algebraic_model(z_guess, event=None): + """``dv/dt = -v`` with ``0 = 2v - z - z^3``, consistent at ``z = 1``.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + z = pybamm.Variable("z") + model.rhs = {v: -v} + model.algebraic = {z: 2 * v - z - z**3} + model.initial_conditions = {v: 1.0, z: z_guess} + model.variables = {"v": v, "z": z} + if event is not None: + name, expression = event + model.events = [pybamm.Event(name, expression(z))] + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +def _exact_z(t): + """The real root of ``z^3 + z = 2 exp(-t)``, by Cardano.""" + q = np.exp(-np.asarray(t)) + root = np.sqrt(q**2 + 1 / 27) + return np.cbrt(q + root) + np.cbrt(q - root) + + +class TestDiffsolInitialConditions: + """Events and stored derivatives use the consistent state, not the guess.""" + + def test_native_initialisation_is_the_default(self): + model = _algebraic_model(5.0) + native = pybamm.DiffsolSolver() + python = pybamm.DiffsolSolver(calc_ic=False) + native.solve(model, [0, 1]) + python.solve(_algebraic_model(5.0), [0, 1]) + + assert native._internal_initialisation + assert native.root_method is None + assert not python._internal_initialisation + assert isinstance(python.root_method, pybamm.NonlinearSolver) + + @pytest.mark.parametrize("calc_ic", [True, False]) + def test_both_initialisations_reach_the_consistent_state(self, calc_ic): + t_eval = np.linspace(0, 1, 6) + solver = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8, calc_ic=calc_ic) + sol = solver.solve(_algebraic_model(5.0), t_eval) + np.testing.assert_allclose(sol["z"](t_eval), _exact_z(t_eval), rtol=1e-6) + + @pytest.mark.parametrize("root_method", ["lm", "hybr"]) + def test_a_generic_root_method_reaches_the_consistent_state(self, root_method): + t_eval = np.linspace(0, 1, 6) + solver = pybamm.DiffsolSolver( + rtol=1e-8, atol=1e-8, calc_ic=False, root_method=root_method + ) + sol = solver.solve(_algebraic_model(5.0), t_eval) + np.testing.assert_allclose(sol["z"](t_eval), _exact_z(t_eval), rtol=1e-6) + + @pytest.mark.parametrize("calc_ic", [True, False]) + def test_off_grid_reads_in_the_first_interval_follow_the_solution(self, calc_ic): + t_eval = np.linspace(0, 1, 6) + t_check = np.linspace(0.02, 0.18, 9) + solver = pybamm.DiffsolSolver(calc_ic=calc_ic) + sol = solver.solve(_algebraic_model(5.0), t_eval) + np.testing.assert_allclose(sol["z"](t_check), _exact_z(t_check), atol=2e-5) + # dz/dt = 2 dv/dt / (1 + 3 z^2) = -0.5 at t = 0 + np.testing.assert_allclose(sol.yp[1, 0], -0.5, rtol=1e-6) + + @pytest.mark.parametrize("output_variables", [None, ["z"]]) + def test_a_guess_on_an_event_boundary_does_not_trip_it(self, output_variables): + # The guess z = 0.5 sits exactly on the event; the consistent z = 1 does not. + model = _algebraic_model(0.5, event=("z above 0.5", lambda z: z - 0.5)) + solver = pybamm.DiffsolSolver(output_variables=output_variables) + sol = solver.solve(model, np.linspace(0, 0.1, 3)) + assert sol.termination == "final time" + + @pytest.mark.parametrize("calc_ic", [True, False]) + @pytest.mark.parametrize("output_variables", [None, ["z"]]) + def test_an_event_violated_at_the_consistent_state_raises( + self, calc_ic, output_variables + ): + # The guess z = 0.5 clears the event; the consistent z = 1 violates it. + model = _algebraic_model(0.5, event=("z below 0.9", lambda z: 0.9 - z)) + solver = pybamm.DiffsolSolver( + calc_ic=calc_ic, output_variables=output_variables + ) + with pytest.raises( + pybamm.SolverError, + match=re.escape("Events ['z below 0.9'] are non-positive at initial"), + ): + solver.solve(model, np.linspace(0, 1, 6)) + + @pytest.mark.parametrize("output_variables", [None, ["z"]]) + def test_a_violation_only_diffsol_sees_raises_without_calc_ic( + self, monkeypatch, output_variables + ): + # Stands in for a root solve that clears the event by less than diffsol's + # own refinement then moves the state, so only the solve can name it. + model = _algebraic_model(0.5, event=("z below 0.9", lambda z: 0.9 - z)) + solver = pybamm.DiffsolSolver(calc_ic=False, output_variables=output_variables) + monkeypatch.setattr( + solver, "_check_event_violation_on_initialisation", lambda *_: None + ) + with pytest.raises( + pybamm.SolverError, + match=re.escape("Events ['z below 0.9'] are non-positive at initial"), + ): + solver.solve(model, np.linspace(0, 1, 6)) + + def test_outputs_only_first_state_is_the_consistent_state(self): + solver = pybamm.DiffsolSolver(output_variables=["z"]) + sol = solver.solve(_algebraic_model(5.0), np.linspace(0, 1, 6)) + np.testing.assert_allclose(sol.first_state["z"](0), 1.0, rtol=1e-6) + np.testing.assert_allclose(sol.first_state["v"](0), 1.0, rtol=1e-12) + + @pytest.mark.parametrize("sensitivities", [False, ["k"]]) + def test_time_dependent_constraint_has_the_analytic_initial_slope( + self, sensitivities + ): + model = pybamm.BaseModel() + v = pybamm.Variable("v") + z = pybamm.Variable("z") + model.rhs = {v: -pybamm.InputParameter("k") * v} + # dz/dt = 1000 at t0 comes only from the constraint's time derivative + model.algebraic = {z: z - 1000 * pybamm.t * (1 + 1000 * pybamm.t)} + model.initial_conditions = {v: 1.0, z: 0.0} + model.variables = {"v": v, "z": z} + + solution = pybamm.DiffsolSolver(rtol=1e-9, atol=1e-9).solve( + model, + np.array([0.0, 0.001, 0.01]), + inputs={"k": 1.0}, + calculate_sensitivities=sensitivities, + ) + + np.testing.assert_allclose(solution.yp[1, 0], 1000.0, rtol=1e-12) + np.testing.assert_allclose(solution["z"](0.0005), 0.75, rtol=1e-9, atol=1e-9) + + +class TestDiffsolPickling: + def test_a_solved_solver_pickles_and_solves_again(self): + t_eval = np.linspace(0, 1, 10) + solver = pybamm.DiffsolSolver() + first = solver.solve(_decay_model(), t_eval) + + restored = pickle.loads(pickle.dumps(solver)) + second = restored.solve(_decay_model(), t_eval) + + np.testing.assert_array_equal(second.y, first.y) + + +class TestDiffsolSolverOptions: + def test_defaults_are_used_when_no_options_are_given(self): + solver = pybamm.DiffsolSolver() + assert solver._options == pybamm.DiffsolSolver.DEFAULT_OPTIONS + + def test_an_override_leaves_the_other_defaults_alone(self): + solver = pybamm.DiffsolSolver(options={"max_error_test_failures": 7}) + assert solver._options["max_error_test_failures"] == 7 + untouched = set(pybamm.DiffsolSolver.DEFAULT_OPTIONS) - { + "max_error_test_failures" + } + for key in untouched: + assert solver._options[key] == pybamm.DiffsolSolver.DEFAULT_OPTIONS[key] + + def test_defaults_match_the_rust_defaults(self): + # Catches silent drift in diffsol's OdeSolverOptions defaults on a bump. + from pybamm.rust import default_solver_options + + assert pybamm.DiffsolSolver._INTEGRATOR_DEFAULTS == default_solver_options() + + def test_an_unknown_option_is_rejected(self): + # A typo would otherwise be dropped silently, leaving the default. + with pytest.raises(pybamm.SolverError, match=r"Unknown diffsol solver option"): + pybamm.DiffsolSolver(options={"max_error_test_failure": 7}) + + def test_the_failure_budget_default_is_not_diffsols_own(self): + # diffsol counts failures cumulatively, so its own 50 caps solve length. + assert ( + pybamm.DiffsolSolver.DEFAULT_OPTIONS["max_nonlinear_solver_failures"] + > 10000 + ) + + def test_options_reach_the_integrator(self): + model = _decay_model() + starved = pybamm.DiffsolSolver( + options={"max_nonlinear_solver_failures": 0, "min_timestep": 0.5} + ) + with pytest.raises(pybamm.SolverError, match=r"diffsol error"): + starved.solve(model, np.linspace(0, 1, 5)) + + assert pybamm.DiffsolSolver().solve(model, np.linspace(0, 1, 5)) is not None + + def test_an_out_of_range_option_is_rejected_by_the_solver(self): + model = _decay_model() + solver = pybamm.DiffsolSolver(options={"nonlinear_solver_tolerance": -1.0}) + with pytest.raises(pybamm.SolverError, match=r"nonlinear_solver_tolerance"): + solver.solve(model, np.linspace(0, 1, 5)) + + @pytest.mark.parametrize( + ("name", "value"), + [ + ("max_nonlinear_solver_iterations", -1), + ("max_error_test_failures", 2.5), + ("max_nonlinear_solver_failures", True), + ("update_jacobian_after_steps", 10**30), + ("min_timestep", None), + ("min_timestep", "1e-3"), + ("max_timestep_growth", "2"), + ("pi_control_integral", False), + ], + ) + def test_an_option_of_the_wrong_type_is_rejected_up_front(self, name, value): + with pytest.raises(pybamm.SolverError, match=rf"^{name} must be"): + pybamm.DiffsolSolver(options={name: value}) + + def test_numpy_scalars_and_an_unset_bound_reach_the_integrator(self): + solver = pybamm.DiffsolSolver( + options={ + "max_error_test_failures": np.int64(7), + "min_timestep": np.float32(1e-10), + "max_timestep_growth": None, + } + ) + options = solver._integrator_options() + assert type(options["max_error_test_failures"]) is int + assert type(options["min_timestep"]) is float + assert options["max_timestep_growth"] is None + solution = solver.solve(_decay_model(), np.linspace(0, 1, 5)) + np.testing.assert_allclose(solution["v"].data, np.exp(-solution.t), rtol=1e-5) + + @pytest.mark.parametrize( + ("rtol", "atol"), [(-1e-6, 1e-6), (1e-6, np.nan)], ids=["rtol", "atol"] + ) + def test_an_invalid_tolerance_is_rejected_by_the_solver(self, rtol, atol): + model = _decay_model() + solver = pybamm.DiffsolSolver(rtol=rtol, atol=atol) + with pytest.raises(pybamm.SolverError, match=r"not a valid tolerance"): + solver.solve(model, np.linspace(0, 1, 5)) + + +def _solve_both_ways(inputs_list, num_threads, t_eval, with_event=True): + """The same sweep solved at ``num_threads`` and serially.""" + parallel = pybamm.DiffsolSolver( + rtol=1e-8, atol=1e-10, options={"num_threads": num_threads} + ).solve( + _decay_model(with_event=with_event, with_input=True), t_eval, inputs=inputs_list + ) + serial = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-10).solve( + _decay_model(with_event=with_event, with_input=True), t_eval, inputs=inputs_list + ) + return parallel, serial + + +def _assert_sweeps_identical(parallel, serial): + """Bit-identical, not close: the batch fans out over the same entry point.""" + for i, (got, want) in enumerate(zip(parallel, serial, strict=True)): + np.testing.assert_array_equal(got.t, want.t, err_msg=f"set {i} times") + np.testing.assert_array_equal(got.y, want.y, err_msg=f"set {i} states") + assert got.termination == want.termination, f"set {i} termination" + assert ( + got.solver_statistics.number_of_steps + == want.solver_statistics.number_of_steps + ), f"set {i} step count" + + +class TestDiffsolNumThreads: + """``num_threads`` solves that many input sets at once, or one set's LU with them.""" + + @staticmethod + def _sweep(n): + return [{"k": 0.3 * (i + 1)} for i in range(n)] + + @pytest.mark.parametrize("num_threads", [1, 2, 8]) + @pytest.mark.parametrize("with_event", [False, True]) + def test_a_parallel_sweep_is_identical_to_the_serial_one( + self, num_threads, with_event + ): + parallel, serial = _solve_both_ways( + self._sweep(6), num_threads, np.linspace(0, 3, 40), with_event=with_event + ) + _assert_sweeps_identical(parallel, serial) + + @solve_settings + @given(rates=decay_rate_sweeps()) + def test_a_batch_is_indistinguishable_from_the_serial_loop(self, rates): + parallel, serial = _solve_both_ways( + [{"k": rate} for rate in rates], 4, np.linspace(0, 3, 40) + ) + _assert_sweeps_identical(parallel, serial) + + def test_results_follow_input_order_not_completion_order(self): + # Descending solve cost, so returning completion order would reverse it. + t_eval = np.linspace(0, 3, 40) + inputs_list = [{"k": k} for k in (0.1, 0.5, 2.0, 9.0)] + solutions = pybamm.DiffsolSolver( + rtol=1e-10, atol=1e-12, options={"num_threads": 4} + ).solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + + for inputs, solution in zip(inputs_list, solutions, strict=True): + assert solution.all_inputs[0] == inputs + np.testing.assert_allclose( + solution["v"](t_eval), + np.exp(-inputs["k"] * t_eval), + rtol=1e-6, + atol=1e-12, + ) + + @pytest.mark.parametrize("num_threads", [1, 4]) + def test_every_failing_set_is_named_whether_batched_or_serial(self, num_threads): + model = _broken_input_model() + inputs_list = [{"k": k} for k in (1.0, -1.0, 2.0, -2.0, 3.0)] + solver = pybamm.DiffsolSolver(options={"num_threads": num_threads}) + with pytest.raises(pybamm.SolverError) as excinfo: + solver.solve(model, np.linspace(0, 1, 10), inputs=inputs_list) + message = str(excinfo.value) + assert re.findall(r"input set (\d+): diffsol error", message) == ["1", "3"] + + def test_a_single_input_set_failure_is_named(self): + with pytest.raises(pybamm.SolverError, match=r"^input set 0: diffsol error"): + pybamm.DiffsolSolver().solve( + _broken_input_model(), np.linspace(0, 1, 10), inputs=[{"k": -1.0}] + ) + + def test_one_pool_is_shared_by_every_solver_of_the_same_width(self): + from pybamm.rust._core import _pool_ids + + t_eval = np.linspace(0, 1, 10) + inputs_list = self._sweep(8) + first = pybamm.DiffsolSolver(options={"num_threads": 8}) + first.solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + pool_id = _pool_ids()[8] + + second = pybamm.DiffsolSolver(options={"num_threads": 8}) + second.solve(_decay_model(with_input=True), t_eval, inputs=inputs_list) + assert _pool_ids()[8] == pool_id + + def test_a_single_set_factorises_on_the_shared_pool_unchanged(self): + from pybamm.rust._core import _pool_ids + + # A width no other test uses, so the pool is this solve's. + parallel, serial = _solve_both_ways({"k": 0.6}, 5, np.linspace(0, 3, 40)) + _assert_sweeps_identical([parallel], [serial]) + assert 5 in _pool_ids() + + @pytest.mark.skipif( + "fork" not in multiprocessing.get_all_start_methods(), reason="needs fork" + ) + @pytest.mark.filterwarnings("ignore:This process:DeprecationWarning") + def test_a_forked_child_builds_its_own_pool(self): + t_eval = np.linspace(0, 1, 10) + inputs_list = self._sweep(4) + + def solve(): + pybamm.DiffsolSolver(options={"num_threads": 3}).solve( + _decay_model(with_input=True), t_eval, inputs=inputs_list + ) + + # Solve first, so the child inherits the parent's pool and not its workers + solve() + child = multiprocessing.get_context("fork").Process(target=solve) + child.start() + child.join(timeout=60) + if child.is_alive(): + child.kill() + child.join() + pytest.fail("the forked child hung on the pool it inherited") + assert child.exitcode == 0 + + def test_the_default_builds_no_pool_at_all(self): + from pybamm.rust._core import _pool_ids + + pybamm.DiffsolSolver().solve( + _decay_model(with_input=True), np.linspace(0, 1, 10), inputs=self._sweep(4) + ) + assert 1 not in _pool_ids() + + def test_pools_are_keyed_on_the_configured_width_not_the_sweep(self): + from pybamm.rust._core import _pool_ids + + # The cache is process-wide, so only the keys this solver adds are ours. + before = set(_pool_ids()) + model = _decay_model(with_input=True) + solver = pybamm.DiffsolSolver(options={"num_threads": 7}) + for n_sets in (2, 3): + solver.solve(model, np.linspace(0, 1, 10), inputs=self._sweep(n_sets)) + added = set(_pool_ids()) - before + assert 7 in _pool_ids() + assert added.isdisjoint({2, 3}) + + @pytest.mark.parametrize("num_threads", [0, -1, 2.5, "4", True, None, 10**30]) + def test_a_num_threads_that_is_not_a_count_is_rejected(self, num_threads): + with pytest.raises(pybamm.SolverError, match=r"num_threads must be an integer"): + pybamm.DiffsolSolver(options={"num_threads": num_threads}) + + def test_num_threads_does_not_reach_the_integrator_options(self): + solver = pybamm.DiffsolSolver(options={"num_threads": 4}) + assert "num_threads" not in solver._integrator_options() + assert set(solver._integrator_options()) == set( + pybamm.DiffsolSolver._INTEGRATOR_DEFAULTS + ) + + +class TestDiffsolIntegrationTime: + def test_each_set_gets_its_own_integration_time(self): + # Batched wall clocks overlap, so a Python timer would report the batch time. + inputs_list = [{"k": 0.3 * (i + 1)} for i in range(4)] + solutions = pybamm.DiffsolSolver(options={"num_threads": 4}).solve( + _decay_model(with_input=True), np.linspace(0, 2, 30), inputs=inputs_list + ) + times = [solution.integration_time for solution in solutions] + assert all(time > 0 for time in times) + assert len(set(times)) > 1 + + def test_the_reported_time_covers_setup_and_integration(self): + solution = pybamm.DiffsolSolver().solve(_decay_model(), np.linspace(0, 2, 30)) + statistics = solution.solver_statistics + assert solution.integration_time >= ( + statistics.ic_time_secs + statistics.solver_setup_time_secs + ) + + +class TestDiffsolSensitivityConfigErrors: + def test_unknown_sensitivity_parameter_is_a_single_clear_error(self): + model = _decay_model(with_input=True) + with pytest.raises( + pybamm.SolverError, + match=r"Sensitivities were requested for \['not_an_input'\]", + ) as excinfo: + pybamm.DiffsolSolver().solve( + model, + np.linspace(0, 1, 5), + inputs={"k": 1.0}, + calculate_sensitivities=["not_an_input"], + ) + # A config error must not go through the relaxed retry and be double-reported. + assert "retry" not in str(excinfo.value) + + +def _two_scale_model(): + """``du/dt = -u`` and ``dw/dt = -w`` with ``u`` at O(1) and ``w`` at O(1e4).""" + model = pybamm.BaseModel() + u = pybamm.Variable("u") + w = pybamm.Variable("w") + model.rhs = {u: -u, w: -w} + model.initial_conditions = {u: 1.0, w: 1e4} + model.variables = {"u": u, "w": w} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +@pytest.mark.parametrize("solver_class", [pybamm.DiffsolSolver, pybamm.IDAKLUSolver]) +class TestPerStateAtol: + def _steps(self, solver_class, atol): + solution = solver_class(rtol=1e-6, atol=atol).solve( + _two_scale_model(), np.linspace(0, 1, 5) + ) + return solution.solver_statistics.number_of_steps + + def test_a_uniform_array_matches_the_scalar(self, solver_class): + t_eval = np.linspace(0, 1, 5) + scalar = solver_class(rtol=1e-8, atol=1e-8).solve(_two_scale_model(), t_eval) + array = solver_class(rtol=1e-8, atol=np.full(2, 1e-8)).solve( + _two_scale_model(), t_eval + ) + np.testing.assert_array_equal(array["u"](t_eval), scalar["u"](t_eval)) + np.testing.assert_array_equal(array["w"](t_eval), scalar["w"](t_eval)) + + def test_every_entry_reaches_its_own_state(self, solver_class): + loose = self._steps(solver_class, np.full(2, 1e-1)) + assert self._steps(solver_class, 1e-1) == loose + assert self._steps(solver_class, np.array([1e-12, 1e-1])) > loose + assert self._steps(solver_class, np.array([1e-1, 1e-12])) > loose + + def test_a_per_state_atol_survives_a_dae_solve(self, solver_class): + # The initial-condition root solver takes a scalar tolerance, not the array. + model = pybamm.BaseModel() + u = pybamm.Variable("u") + z = pybamm.Variable("z") + model.rhs = {u: -u} + model.algebraic = {z: z - 2 * u} + model.initial_conditions = {u: 1.0, z: 2.0} + model.variables = {"u": u, "z": z} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + + t_eval = np.linspace(0, 1, 5) + solution = solver_class(rtol=1e-8, atol=np.array([1e-9, 1e-7])).solve( + model, t_eval + ) + np.testing.assert_allclose( + solution["z"](t_eval), 2 * np.exp(-t_eval), rtol=1e-6 + ) + + def test_a_config_round_trip_keeps_the_per_state_atol(self, solver_class): + # to_config puts the array through JSON, which has no arrays to put back. + atol = np.array([1e-1, 1e-12]) + restored = pybamm.BaseSolver.from_config( + solver_class(rtol=1e-6, atol=atol).to_config() + ) + assert self._steps(solver_class, restored.atol) == self._steps( + solver_class, atol + ) + + def test_a_wrong_width_atol_is_rejected(self, solver_class): + solver = solver_class(atol=np.full(3, 1e-6)) + with pytest.raises(pybamm.SolverError, match=r"shape \(3,\) but \(2,\)"): + solver.solve(_two_scale_model(), np.linspace(0, 1, 5)) + + @pytest.mark.parametrize( + ("bad_atol", "message"), + [ + ("tight", r"must be a float, or a list, tuple or array"), + (["tight", "loose"], r"must be real numbers"), + (np.array(["tight", "loose"]), r"must be real numbers"), + ([True, 1e-6], r"not booleans"), + ], + ) + def test_a_non_numeric_atol_is_rejected(self, solver_class, bad_atol, message): + solver = solver_class(atol=bad_atol) + with pytest.raises(pybamm.SolverError, match=message): + solver.solve(_two_scale_model(), np.linspace(0, 1, 5)) + + +def _ramp_then_hold_model(): + """``dv/dt = 1`` until t = 5 and 0 after, so ``v(t) = min(t, 5)``.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: pybamm.t < 5} + model.initial_conditions = {v: pybamm.Scalar(0)} + model.variables = {"v": v} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +def _algebraic_jump_model(rate=0.1, terminate=False): + """``dv/dt = -rate v`` with ``z = v`` until t = 0.5 and ``z = -v`` after.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + z = pybamm.Variable("z") + model.rhs = {v: -rate * v} + model.algebraic = {z: z - v * (1 - 2 * (pybamm.t > 0.5))} + model.initial_conditions = {v: 1.0, z: 1.0} + model.variables = {"v": v, "z": z} + if terminate: + model.events = [pybamm.Event("positive z", z)] + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +def _switch_at_the_end_model(switch): + """``dv/dt = -v/10`` with ``z = v * switch(t)``, for a switch that flips at the end.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + z = pybamm.Variable("z") + model.rhs = {v: -0.1 * v} + model.algebraic = {z: z - v * switch(pybamm.t)} + model.initial_conditions = {v: 1.0, z: 1.0} + model.variables = {"v": v, "z": z} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +def _jump_just_after_start_model(): + """``dv/dt = -k v`` with ``z`` stepping from 0 to 1 at t = 1e-8.""" + model = pybamm.BaseModel() + v = pybamm.Variable("v") + z = pybamm.Variable("z") + model.rhs = {v: -pybamm.InputParameter("k") * v} + model.algebraic = {z: z - (pybamm.t > 1e-8)} + model.initial_conditions = {v: 1.0, z: 0.0} + model.variables = {"v": v, "z": z} + model.convert_to_format = "rust" + pybamm.Discretisation().process_model(model) + return model + + +class TestDiffsolTimeDiscontinuities: + """A constant time discontinuity reaches diffsol as a stop time, as for IDAKLU.""" + + @pytest.mark.parametrize( + "solver_class", [pybamm.DiffsolSolver, pybamm.IDAKLUSolver] + ) + def test_a_heaviside_in_time_is_solved_in_one_pass(self, solver_class): + solver = solver_class(rtol=1e-8, atol=1e-8) + solution = solver.solve(_ramp_then_hold_model(), np.linspace(0, 10, 21)) + + assert len(solution.sub_solutions) == 1 + np.testing.assert_allclose( + solution["v"](np.array([2.5, 5.0, 7.5])), + np.array([2.5, 5.0, 5.0]), + rtol=1e-6, + atol=1e-6, + ) + + @pytest.mark.parametrize( + "solver_class", [pybamm.DiffsolSolver, pybamm.IDAKLUSolver] + ) + def test_an_algebraic_jump_restarts_on_the_new_branch(self, solver_class): + solver = solver_class(rtol=1e-8, atol=1e-8) + solution = solver.solve(_algebraic_jump_model(), np.linspace(0, 2, 9)) + + t_check = np.array([0.4, 0.5, 0.6, 1.0, 2.0]) + v = np.exp(-0.1 * t_check) + np.testing.assert_allclose(solution["v"](t_check), v, rtol=1e-6) + np.testing.assert_allclose( + solution["z"](t_check), np.where(t_check > 0.5, -v, v), rtol=1e-6 + ) + + @pytest.mark.parametrize("output_variables", [None, ["z"]]) + @pytest.mark.parametrize("sensitivities", [False, ["k"]]) + def test_an_algebraic_jump_terminates_when_it_exceeds_an_event( + self, output_variables, sensitivities + ): + model = _algebraic_jump_model(rate=pybamm.InputParameter("k"), terminate=True) + solution = pybamm.DiffsolSolver(output_variables=output_variables).solve( + model, + np.linspace(0, 2, 9), + inputs={"k": 0.1}, + calculate_sensitivities=sensitivities, + ) + + assert solution.termination == "event: positive z" + assert solution.t[-1] > 0.5 + np.testing.assert_allclose(solution.t[-1], 0.5, rtol=0, atol=1e-15) + np.testing.assert_allclose( + solution["z"](solution.t[-1]), -np.exp(-0.05), rtol=1e-5 + ) + np.testing.assert_allclose( + solution.last_state["v"](solution.t[-1]), np.exp(-0.05), rtol=1e-5 + ) + + @pytest.mark.parametrize( + ("switch", "t_end"), + [ + (lambda t: 1 - 2 * (t >= 0.5), 0.5), + # A sawtooth drive whose span is a whole number of periods. + (lambda t: 1 + t % 0.25, 1.0), + ], + ids=["heaviside", "sawtooth"], + ) + def test_a_switch_at_the_final_time_restarts_on_its_far_side(self, switch, t_end): + # IDAKLU reports the pre-switch value at t_end instead. + solution = pybamm.DiffsolSolver(rtol=1e-8, atol=1e-8).solve( + _switch_at_the_end_model(switch), [0, t_end] + ) + + t_check = np.array([0.1, t_end]) + v = np.exp(-0.1 * t_check) + np.testing.assert_allclose(solution["v"](t_check), v, rtol=1e-6) + np.testing.assert_allclose( + solution["z"](t_check), v * [switch(t) for t in t_check], rtol=1e-6 + ) + + @pytest.mark.parametrize("sensitivities", [False, ["k"]]) + def test_a_jump_just_after_the_start_stays_out_of_its_slope(self, sensitivities): + solution = pybamm.DiffsolSolver().solve( + _jump_just_after_start_model(), + np.linspace(0, 1, 5), + inputs={"k": 1.0}, + calculate_sensitivities=sensitivities, + ) + + np.testing.assert_allclose(solution.yp[1, 0], 0.0, atol=1e-12) + np.testing.assert_allclose( + solution["z"](np.array([0.0, 2.5e-9, 5e-9, 0.5, 1.0])), + [0.0, 0.0, 0.0, 1.0, 1.0], + atol=1e-12, + ) + + +def _native_prepared_solver(with_sens=False): + """``dy/dt = -k*y`` natively, with two outputs against its single state.""" + from pybamm.rust import CompiledModel, ExprGraph, PreparedSolver + + g = ExprGraph() + y = g.state_vector(0, 1) + rate = g.mul(g.mul(g.scalar(-1.0), g.input_parameter("k")), y) + model = CompiledModel.from_expr( + g, + rate, + np.ones(1), + np.arange(2, dtype=np.int64), + np.zeros(1, dtype=np.int64), + n_inputs=1, + sens_param_indices=[0] if with_sens else [], + output_exprs=[g.mul(g.scalar(2.0), y), g.mul(g.scalar(3.0), y)], + event_exprs=[], + ) + return PreparedSolver(model, 1e-10, 1e-10) + + +class TestNativeSolveRequest: + _T_EVAL = np.linspace(0.0, 1.0, 11) + _NO_STOPS = np.array([], dtype=np.float64) + + def _solve(self, solver, **flags): + return solver.solve( + self._T_EVAL, self._NO_STOPS, np.array([1.0]), np.array([1.0]), **flags + ) + + def test_the_row_space_follows_the_outputs_flag(self): + solver = _native_prepared_solver() + states = self._solve(solver) + outputs = self._solve(solver, outputs=True) + + # One state, two outputs (2y, 3y): the flag swaps the rows, not the columns. + assert states.y.shape == (1, self._T_EVAL.size) + assert outputs.y.shape == (2, self._T_EVAL.size) + np.testing.assert_allclose(outputs.y[0], 2.0 * states.y[0], rtol=1e-8) + np.testing.assert_allclose(outputs.y[1], 3.0 * states.y[0], rtol=1e-8) + + def test_payloads_not_asked_for_read_as_none(self): + solver = _native_prepared_solver(with_sens=True) + + assert self._solve(solver).yS is None + assert self._solve(solver, sensitivities=True).yS is not None + # yp is the state trajectory's slopes, and this solver does not store them. + assert self._solve(solver).yp is None + + def test_a_seed_without_sensitivities_is_rejected(self): + solver = _native_prepared_solver(with_sens=True) + + # Ignoring the seed would return zero sensitivities that look computed. + with pytest.raises(ValueError, match=r"y0_sens was given but"): + self._solve(solver, y0_sens=np.array([0.5])) + + def test_a_batch_answers_one_shared_request(self): + solver = _native_prepared_solver(with_sens=True) + rates = [0.5, 1.0, 2.0] + + results = solver.solve_batch( + self._T_EVAL, + self._NO_STOPS, + np.ones((len(rates), 1)), + np.array(rates).reshape(-1, 1), + 1, + outputs=True, + sensitivities=True, + ) + + assert len(results) == len(rates) + for i, (result, rate) in enumerate(zip(results, rates, strict=True)): + serial = solver.solve( + self._T_EVAL, + self._NO_STOPS, + np.array([1.0]), + np.array([rate]), + outputs=True, + sensitivities=True, + ) + assert result.yS is not None, f"set {i} dropped its blocks" + np.testing.assert_array_equal(result.y, serial.y, err_msg=f"set {i}") + np.testing.assert_array_equal( + result.yS[0], serial.yS[0], err_msg=f"set {i} sensitivities" + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py index 307b14eb41..f957ec64b0 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -1,4 +1,3 @@ -import dataclasses import io import itertools import json @@ -2157,6 +2156,21 @@ def test_rust_events_survive_compiled_model_pickle(self): for a, b in zip(before, after, strict=True): np.testing.assert_array_equal(a, b) + def test_rust_diffsol_event_termination(self): + params = pybamm.ParameterValues("Chen2020") + params["Lower voltage cut-off [V]"] = 3.5 + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + sim = pybamm.Simulation( + model, parameter_values=params, solver=pybamm.DiffsolSolver() + ) + sol = sim.solve(np.linspace(0, 3600, 100)) + assert sol.termination.startswith("event:"), ( + f"expected diffsol root-finding to terminate on an event, " + f"got {sol.termination!r}" + ) + assert float(sol.t[-1]) < 3600.0 + def test_simulation_period(self): model = pybamm.lithium_ion.DFN() parameter_values = pybamm.ParameterValues("Chen2020") @@ -3717,7 +3731,11 @@ def printed_count(label): "Number of error test failures" ), ) - assert all(count > 0 for count in dataclasses.astuple(stats)) + # Only the fields IDAS reports; the diffsol-only ones stay at zero + assert all( + getattr(stats, name) > 0 + for name in pybamm.solvers.idaklu_solver._IDAS_STATISTICS + ) assert stats.number_of_nonlinear_solver_iterations >= stats.number_of_steps def test_statistics_accumulate_across_a_breakpoint(self, decay_model): diff --git a/packages/pybamm/tests/unit/test_solvers/test_native_observation.py b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py index be9cce10e7..0b8bbe7042 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_native_observation.py +++ b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py @@ -350,6 +350,29 @@ def test_pickle_round_trip_after_observe(self): ) +class TestDiffsolObservationCache: + def test_compiled_fn_reused_across_solves(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + model.events = [] + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=pybamm.DiffsolSolver()) + + sol1 = sim.solve(np.linspace(0, 600, 50)) + _ = sol1["Terminal voltage [V]"] + cache = ObserverCache.of(sol1.all_models[0]) + assert cache.rust_model is sim._solver._rust_model + fn1 = cache._native_leaves["Terminal voltage [V]"][0] + nodes_after_first = sim._solver._rust_model.graph.n_nodes + + sol2 = sim.solve(np.linspace(0, 600, 50)) + _ = sol2["Terminal voltage [V]"] + assert ObserverCache.of(sol2.all_models[0]) is cache + assert cache._native_leaves["Terminal voltage [V]"][0] is fn1 + assert sim._solver._rust_model.graph.n_nodes == nodes_after_first + + class TestBatchedNativeObservation: """`update([...])` shares one tape per segment across the whole set.""" diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py index 1f9d9068cf..6b73fd0660 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py @@ -148,6 +148,25 @@ def test_a_dense_layout_needs_one_length_per_name(self): with pytest.raises(pybamm.SolverError, match=r"Output layout mismatch"): OutputAssembly.dense(["Voltage [V]", "Current [A]"], [1]) + def test_parameter_blocks_stack_into_the_attachable_layout(self, solution): + assembly, _ = self._assembly(solution, ["Voltage [V]"]) + n_t = len(solution.t) + blocks = [np.arange(n_t, dtype=float), np.arange(n_t, dtype=float) * -1.0] + + stacked = assembly.stack_parameter_blocks(blocks, n_t, ["a", "b"]) + + assert stacked.shape == (n_t, 1, 2) + np.testing.assert_allclose(stacked[:, 0, 0], blocks[0]) + np.testing.assert_allclose(stacked[:, 0, 1], blocks[1]) + + def test_one_block_per_parameter_is_required(self, solution): + assembly, _ = self._assembly(solution, ["Voltage [V]"]) + n_t = len(solution.t) + with pytest.raises( + pybamm.SolverError, match=r"Sensitivity block count mismatch" + ): + assembly.stack_parameter_blocks([np.zeros(n_t)], n_t, ["a", "b"]) + def test_a_solve_without_sensitivities_leaves_an_empty_mapping(self, solution): # Empty, not None: an outputs-only solve keeps no state to compute them from assembly, data = self._assembly(solution, ["Voltage [V]"]) diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py b/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py new file mode 100644 index 0000000000..f414e49398 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_gil_release.py @@ -0,0 +1,84 @@ +"""The native diffsol solver must release the GIL during integration.""" + +import threading +import time + +import numpy as np + +from pybamm.rust import CompiledModel, ExprGraph, PreparedSolver + + +def _decay_prepared_solver(n_states): + """Native prepared solver for a decoupled linear decay system dy_i/dt=-y_i.""" + g = ExprGraph() + terms = [g.mul(g.scalar(-1.0), g.state_vector(i, i + 1)) for i in range(n_states)] + rhs = g.concat(terms) + model = CompiledModel.from_expr( + g, + rhs, + np.ones(n_states), # identity mass matrix (CSR) + np.arange(n_states + 1, dtype=np.int64), + np.arange(n_states, dtype=np.int64), + n_inputs=0, + sens_param_indices=[], + output_exprs=[], + event_exprs=[], + ) + return PreparedSolver(model, 1e-8, 1e-8) + + +class TestRustGilRelease: + def test_native_solve_releases_gil(self): + n_states = 100 + ps = _decay_prepared_solver(n_states) + y0 = np.ones(n_states) + inputs = np.array([], dtype=np.float64) + no_stops = np.array([], dtype=np.float64) + # A fine output grid makes one native solve take hundreds of milliseconds. + t_eval = np.linspace(0.0, 5.0, 400_000) + + ps.solve( + t_eval, no_stops, y0, inputs + ) # warm up (first solve builds internal state) + + stop = threading.Event() + ticks = [] + + def worker(): + # Sleeping keeps this thread off the CPU so it never competes with the solve. + while not stop.is_set(): + ticks.append(time.perf_counter()) + time.sleep(0.001) + + t = threading.Thread(target=worker) + t.start() + calls = [] + try: + # Baseline tick cadence: sleep granularity varies by platform. + time.sleep(0.1) + quiet_cadence = np.percentile(np.diff(ticks[:]), 95) + + # A coarse sleep timer needs several solves to span enough ticks. + while sum(end - start for start, end in calls) < 8 * quiet_cadence: + start = time.perf_counter() + ps.solve(t_eval, no_stops, y0, inputs) + calls.append((start, time.perf_counter())) + finally: + stop.set() + t.join() + + for start, end in calls: + solve_wall = end - start + # A held GIL blocks the worker outright, so it records no tick inside the call. + during = [tick for tick in ticks if start < tick < end] + assert during, ( + f"worker ran zero times during a {solve_wall * 1e3:.1f} ms native " + "solve: the GIL was not released" + ) + + # Bound the longest quiet stretch, not throughput, which tracks CPU contention. + longest_gap = np.diff([start, *during, end]).max() + assert longest_gap < 0.5 * solve_wall, ( + f"worker stalled for {longest_gap * 1e3:.1f} ms inside a " + f"{solve_wall * 1e3:.1f} ms native solve: the GIL was held for part of it" + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py b/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py index 1a9d6c899c..f7c3c71617 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py +++ b/packages/pybamm/tests/unit/test_solvers/test_rust_lowering.py @@ -54,6 +54,17 @@ def test_tensor_field_output_is_rejected(self): with pytest.raises(pybamm.SolverError, match=r"'field' is a tensor field"): lowering.outputs(["field"]) + def test_idaklu_and_diffsol_agree_on_output_lengths(self): + t_eval = np.linspace(0, 100, 5) + + idaklu = pybamm.IDAKLUSolver(output_variables=_OUTPUTS) + idaklu.solve(_built_model(), t_eval) + + diffsol = pybamm.DiffsolSolver(output_variables=_OUTPUTS) + diffsol.solve(_built_model(), t_eval) + + assert idaklu._setup["output_assembly"]._rows == diffsol._output_assembly._rows + class TestInputRegistration: def test_graph_input_order_follows_the_inputs_dict(self): @@ -193,7 +204,9 @@ def test_bound_events_follow_model_events_order(self): rtol=1e-12, ) - @pytest.mark.parametrize("solver_class", [pybamm.IDAKLUSolver]) + @pytest.mark.parametrize( + "solver_class", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver] + ) def test_every_native_solver_binds_both_slots(self, solver_class): """An unbound events slot would silently skip the event-violation check.""" model = _built_model() @@ -245,7 +258,9 @@ def counting(*args, **kwargs): solver.set_up(model, inputs={}, t_eval=np.linspace(0, 100, 3)) return len(built) - @pytest.mark.parametrize("solver_class", [pybamm.IDAKLUSolver]) + @pytest.mark.parametrize( + "solver_class", [pybamm.IDAKLUSolver, pybamm.DiffsolSolver] + ) def test_native_set_up_builds_one_graph_beside_the_initial_conditions( self, monkeypatch, solver_class ): diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 1234b11b64..1632612f9e 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -2438,6 +2438,26 @@ def test_statistics_survive_pickling(self): solution = self._solution(0, statistics) assert pickle.loads(pickle.dumps(solution)).solver_statistics == statistics + def test_statistics_pickled_before_the_diffsol_fields_load(self): + # Such a pickle holds a value for each of the five earlier fields only + earlier_fields = ( + "number_of_steps", + "number_of_linear_solver_setups", + "number_of_nonlinear_solver_iterations", + "number_of_nonlinear_solver_fails", + "number_of_error_test_failures", + ) + statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) + state = [getattr(statistics, name) for name in earlier_fields] + + class EarlierPickle: + def __reduce__(self): + return object.__new__, (pybamm.SolverStatistics,), state + + loaded = pickle.loads(pickle.dumps(EarlierPickle())) + assert loaded == statistics + assert loaded + statistics == pybamm.SolverStatistics(2, 4, 6, 8, 10) + def test_solution_pickled_before_slots_loads(self, caplog): solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) first_state = solution.first_state