From cc9a58a141ec9a9088af9d51d109fbd507a8a1c4 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 23 Sep 2026 19:36:04 +0100 Subject: [PATCH 01/13] feat(idaklu): solver statistics, self-scheduled input-set sweeps and polymorphic Newton Solutions carry IDAKLU's counters as SolverStatistics, summed as solutions combine; input sets are self-scheduled with every failing set named; the standalone Newton owns its Expressions; per-state atol takes lists and tuples. Requires pybammsolvers 0.10.0. --- CHANGELOG.md | 10 + docs/source/api/solvers/solution.rst | 2 + packages/pybamm/src/pybamm/__init__.py | 1 + .../src/pybamm/solvers/idaklu_solver.py | 87 +++- .../pybamm/src/pybamm/solvers/solution.py | 60 +++ packages/pybamm/tests/__init__.py | 1 + packages/pybamm/tests/shared.py | 14 + .../unit/test_solvers/test_idaklu_solver.py | 393 +++++++++++++++++- .../tests/unit/test_solvers/test_solution.py | 71 ++++ .../src/pybammsolvers/idaklu.cpp | 14 +- .../Expressions/Base/Expression.hpp | 5 + .../idaklu_source/IDAKLUSolverGroup.cpp | 89 ++-- .../idaklu_source/IDAKLUSolverOpenMP.inl | 15 +- .../pybammsolvers/idaklu_source/Solution.hpp | 7 +- .../idaklu_source/SolutionData.hpp | 10 +- .../idaklu_source/StandaloneNewtonSolver.cpp | 54 ++- .../idaklu_source/StandaloneNewtonSolver.hpp | 22 +- .../pybammsolvers/tests/test_integration.py | 32 ++ 18 files changed, 775 insertions(+), 112 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7533ecf0d3..e436f0e80d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ - Setting `Array.entries_string` or `Interpolant.entries_string`, and calling `StateVector.set_evaluation_array`, are in-place symbol updates and now emit `SymbolMutationDeprecationWarning`. Construct a new symbol instead. ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825)) +## Features + +- `Solution.solver_statistics` reports the integrator's step, linear-solver-setup, nonlinear-iteration, nonlinear-failure and error-test-failure counts as a `pybamm.SolverStatistics` dataclass. `IDAKLUSolver` fills it for every solution, counting across `t_eval` breakpoints and separately for each input set, and the counts are summed when solutions are combined, so a stepped or experiment solution reports its total. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) + +## Bug fixes + +- A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) + ## Optimizations - Symbol mutation is no longer guarded by intercepting every attribute write, and `children` and `domain` return the stored immutable sequences instead of building a view on each access. Writes to a symbol's slots after construction are instead rejected by a static check of PyBaMM's code, tests, examples and documentation. Under the test suite's mutation guard, constructing an expression is twice as fast and parameterising it is 25% faster, which recovers most of the `test_parameterise` benchmark regression from [#5779](https://github.com/pybamm-team/PyBaMM/pull/5779). ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825)) diff --git a/docs/source/api/solvers/solution.rst b/docs/source/api/solvers/solution.rst index 022d276577..51ff856c60 100644 --- a/docs/source/api/solvers/solution.rst +++ b/docs/source/api/solvers/solution.rst @@ -4,4 +4,6 @@ Solutions .. autoclass:: pybamm.Solution :members: +.. autoclass:: pybamm.SolverStatistics + .. footbibliography:: diff --git a/packages/pybamm/src/pybamm/__init__.py b/packages/pybamm/src/pybamm/__init__.py index 95286250d7..e46e9485ee 100644 --- a/packages/pybamm/src/pybamm/__init__.py +++ b/packages/pybamm/src/pybamm/__init__.py @@ -204,6 +204,7 @@ # Solver classes from .solvers.solution import ( + SolverStatistics, SolutionBase, Solution, EISSolution, diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 2862367a36..68934ae8df 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -1,4 +1,5 @@ # mypy: ignore-errors +import dataclasses import logging import math import numbers @@ -88,8 +89,9 @@ class IDAKLUSolver(pybamm.BaseSolver): ---------- rtol : float, optional The relative tolerance for the solver (default is 1e-4). - 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, either one value for every state + or a list, tuple or array with one value per state (default is 1e-6). root_method : str or pybamm algebraic solver class, optional The method to use to find initial conditions (for DAE solvers). Default is None, which uses a custom Newton solver for consistent @@ -169,7 +171,7 @@ class IDAKLUSolver(pybamm.BaseSolver): "max_order_bdf": 5, # Maximum number of steps to be taken by the solver in its attempt to # reach the next output time. - # Note: this value differs from the IDA default of 500 + # This value differs from the IDA default of 500 "max_num_steps": 100000, # Initial step size. The solver default is used if this is left at 0.0 "dt_init": 0.0, @@ -182,17 +184,17 @@ class IDAKLUSolver(pybamm.BaseSolver): # Maximum number of error test failures in attempting one step "max_error_test_failures": 10, # Maximum number of nonlinear solver iterations at one step - # Note: this value differs from the IDA default of 4 + # This value differs from the IDA default of 4 "max_nonlinear_iterations": 40, # Maximum number of nonlinear solver convergence failures at one step - # Note: this value differs from the IDA default of 10 + # This value differs from the IDA default of 10 "max_convergence_failures": 100, # Safety factor in the nonlinear convergence test "nonlinear_convergence_coefficient": 0.33, # Suppress algebraic variables from error test "suppress_algebraic_error": False, # Store Hermite interpolation data for the solution. - # Note: this option is always disabled if output_variables are given + # This option is always disabled if output_variables are given # or if t_interp values are specified "hermite_interpolation": True, # Setting hermite_reduction_factor > 1.0 compresses the solution size @@ -208,15 +210,15 @@ class IDAKLUSolver(pybamm.BaseSolver): # initial condition calculation "nonlinear_convergence_coefficient_ic": 0.0033, # Maximum number of steps allowed when `init_all_y_ic = False` - # Note: this value differs from the IDA default of 5 + # This value differs from the IDA default of 5 "max_num_steps_ic": 50, # Maximum number of the approximate Jacobian or preconditioner evaluations # allowed when the Newton iteration appears to be slowly converging - # Note: this value differs from the IDA default of 4 + # This value differs from the IDA default of 4 "max_num_jacobians_ic": 40, # Maximum number of Newton iterations allowed in any one attempt to solve # the initial conditions calculation problem - # Note: this value differs from the IDA default of 10 + # This value differs from the IDA default of 10 "max_num_iterations_ic": 100, # Maximum number of linesearch backtracks allowed in any Newton iteration, # when solving the initial conditions calculation problem @@ -371,15 +373,61 @@ def _check_options(self, options: dict): if not isinstance(options["compile"], bool): raise pybamm.SolverError("compile must be a bool") - def _check_atol_type(self, atol, model): - if isinstance(atol, float): - return np.full(model.len_rhs_and_alg, atol) - elif isinstance(atol, np.ndarray): - return atol - else: + 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 numpy array or float" + "Absolute tolerances must be a flat list, tuple or array" + ) from error + # 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(self, model, inputs=None, t_eval=None, ics_only=False): if model.convert_to_format != "casadi": @@ -861,6 +909,12 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev newsol.closest_event_idx = int(np.nanargmin(np.abs(event_values))) newsol.integration_time = integration_time + newsol.solver_statistics = pybamm.SolverStatistics( + **{ + field.name: getattr(sol.stats, field.name) + for field in dataclasses.fields(pybamm.SolverStatistics) + } + ) if not save_outputs_only: return newsol @@ -1208,6 +1262,7 @@ def reduce_solution( new_sol.solve_time = solution.solve_time new_sol.integration_time = solution.integration_time + new_sol.solver_statistics = solution.solver_statistics new_sol.set_up_time = solution.set_up_time return new_sol diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index b67ce33c7e..ebb37ec5d0 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -6,6 +6,7 @@ import json import numbers import pickle +from dataclasses import astuple, dataclass from functools import cached_property from itertools import chain @@ -32,6 +33,57 @@ def default(self, obj): return json.JSONEncoder.default(self, obj) # pragma: no cover +@dataclass(frozen=True) +class SolverStatistics: + """Counts of the work an integrator did to produce a solution. + + Attributes + ---------- + number_of_steps : int + Internal time steps taken. + number_of_linear_solver_setups : int + Calls to the linear solver setup, such as Jacobian factorisations. + number_of_nonlinear_solver_iterations : int + Iterations of the nonlinear (Newton) solver. + number_of_nonlinear_solver_fails : int + Convergence failures of the nonlinear solver. + number_of_error_test_failures : int + Steps rejected by the local error test. + """ + + number_of_steps: int = 0 + number_of_linear_solver_setups: int = 0 + number_of_nonlinear_solver_iterations: int = 0 + number_of_nonlinear_solver_fails: int = 0 + number_of_error_test_failures: int = 0 + + 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)) + ) + + +def _sum_solver_statistics(solutions) -> SolverStatistics | None: + """Total the solver statistics of ``solutions``. + + Parameters + ---------- + solutions : list of :class:`pybamm.Solution` + The solutions whose statistics are summed. + + Returns + ------- + :class:`pybamm.SolverStatistics` or None + The field-wise sum, or None if any solution has no statistics. + """ + statistics = [solution.solver_statistics for solution in solutions] + if any(stats is None for stats in statistics): + return None + return sum(statistics[1:], statistics[0]) + + class SolutionBase: """Base class for all PyBaMM solution types (time-series and EIS). @@ -241,6 +293,9 @@ class Solution(SolutionBase): """ + # A class default, so solutions pickled before this attribute existed load + solver_statistics: SolverStatistics | None = None + def __init__( self, all_ts, @@ -620,6 +675,7 @@ def first_state(self): new_sol.solve_time = 0 new_sol.integration_time = 0 + new_sol.solver_statistics = SolverStatistics() new_sol.set_up_time = 0 return new_sol @@ -663,6 +719,7 @@ def last_state(self): new_sol._sub_solutions = self.sub_solutions[-1:] new_sol.solve_time = 0 new_sol.integration_time = 0 + new_sol.solver_statistics = SolverStatistics() new_sol.set_up_time = 0 return new_sol @@ -1177,6 +1234,7 @@ def __add__(self, other): and getattr(other, attr, None) is not None ): setattr(new_sol, attr, getattr(self, attr) + getattr(other, attr)) + new_sol.solver_statistics = _sum_solver_statistics([self, other]) # Set sub_solutions new_sol._sub_solutions = self.sub_solutions + other.sub_solutions @@ -1291,6 +1349,7 @@ def from_sub_solutions(cls, sub_solutions): vals = [getattr(s, attr, None) for s in segments] if all(v is not None for v in vals): setattr(new_sol, attr, sum(vals)) + new_sol.solver_statistics = _sum_solver_statistics(segments) # output_variables path: reproduce __add__'s pairwise left-fold. if any(s.variables_returned for s in segments): @@ -1329,6 +1388,7 @@ def copy(self): new_sol.solve_time = self.solve_time new_sol.integration_time = self.integration_time + new_sol.solver_statistics = self.solver_statistics new_sol.set_up_time = self.set_up_time # copy over variables which were derived at the solver stage diff --git a/packages/pybamm/tests/__init__.py b/packages/pybamm/tests/__init__.py index ca1aeddfef..0ad83917cb 100644 --- a/packages/pybamm/tests/__init__.py +++ b/packages/pybamm/tests/__init__.py @@ -47,6 +47,7 @@ get_unit_2p1D_mesh_for_testing, get_cylindrical_discretisation_for_testing, get_base_model_with_battery_geometry, + get_broken_input_model, get_required_distribution_deps, get_optional_distribution_deps, get_present_optional_import_deps, diff --git a/packages/pybamm/tests/shared.py b/packages/pybamm/tests/shared.py index e80526d5aa..e599288699 100644 --- a/packages/pybamm/tests/shared.py +++ b/packages/pybamm/tests/shared.py @@ -434,6 +434,20 @@ def get_base_model_with_battery_geometry(**kwargs): return model +def get_broken_input_model(): + """Return the discretised model ``dv/dt = -sqrt(k) * v``. + + Its residual is NaN for any input ``k < 0``, whatever the state. + """ + model = pybamm.BaseModel() + v = pybamm.Variable("v") + model.rhs = {v: -pybamm.sqrt(pybamm.InputParameter("k")) * v} + model.initial_conditions = {v: 1.0} + model.variables = {"v": v} + pybamm.Discretisation().process_model(model) + return model + + def get_required_distribution_deps(package_name): pattern = re.compile(r"(?!.*extra\b)^([^<>=;\[]+)\b.*$") if json_deps := importlib_metadata.metadata(package_name).json.get("requires_dist"): 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 24dc9edb70..41e0d62ac4 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -1,7 +1,13 @@ +import dataclasses import io import itertools +import json import logging +import os +import re +import subprocess # nosec B404 - runs this interpreter on a fixed script import sys +import textwrap import warnings from contextlib import redirect_stdout @@ -13,7 +19,11 @@ from scipy.sparse import csc_matrix import pybamm -from tests import get_discretisation_for_testing, no_internet_connection +from tests import ( + get_broken_input_model, + get_discretisation_for_testing, + no_internet_connection, +) def _hermite_wrms(sol_base, sol_reduced, atol, rtol) -> list[tuple[int, float]]: @@ -170,6 +180,55 @@ def test_multiple_inputs(self): rtol=1e-4, ) + def test_every_failing_input_set_is_named(self): + model = get_broken_input_model() + solver = pybamm.IDAKLUSolver(options={"num_threads": 4}) + inputs_list = [{"k": k} for k in (1.0, -1.0, 2.0, -2.0, 3.0)] + with pytest.raises(pybamm.SolverError) as error: + solver.solve(model, np.linspace(0, 1, 10), inputs=inputs_list) + assert re.findall(r"input set (\d+): ", str(error.value)) == ["1", "3"] + + def test_a_single_input_set_failure_is_named(self): + model = get_broken_input_model() + with pytest.raises(pybamm.SolverError, match=r"^input set 0: "): + pybamm.IDAKLUSolver().solve(model, [0, 1], inputs={"k": -1.0}) + + def test_every_input_set_is_solved_by_a_smaller_team(self): + # The OpenMP runtime reads OMP_THREAD_LIMIT once, so this needs a fresh + # process; the team then has fewer threads than the solver has solvers. + script = textwrap.dedent( + """ + import json + + import numpy as np + + import pybamm + + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -pybamm.InputParameter("a") * u} + model.initial_conditions = {u: 1} + model.variables = {"u": u} + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver( + rtol=1e-8, atol=1e-10, options={"num_threads": 4} + ) + inputs = [{"a": 1.0 + i} for i in range(6)] + solutions = solver.solve(model, [0, 1], inputs=inputs) + print(json.dumps([float(np.squeeze(s["u"](1.0))) for s in solutions])) + """ + ) + result = subprocess.run( # nosec B603 - fixed arguments, no shell + [sys.executable, "-c", script], + env={**os.environ, "OMP_THREAD_LIMIT": "2"}, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + values = json.loads(result.stdout.strip().splitlines()[-1]) + np.testing.assert_allclose(values, np.exp(-(1.0 + np.arange(6))), rtol=1e-6) + def test_model_events(self): # Create model model = pybamm.BaseModel() @@ -702,19 +761,19 @@ def test_banded(self): np.testing.assert_allclose(soln.y, soln_banded.y, rtol=1e-6, atol=1e-5) @pytest.mark.parametrize( - ("num_threads", "n_inputs", "all_on_calling_thread"), + ("num_threads", "n_inputs"), [ # One solver, so every solve runs on the GIL-holding thread - (1, 2, True), - # Fewer input sets than solvers, so all of them land in the - # serial remainder loop, which also runs on that thread - (4, 2, True), - # Two solves per solver, so a worker thread must buffer its share - (2, 4, False), + (1, 2), + # More solvers than input sets, so the team is capped at one thread + # per set and the worker buffers its set + (4, 2), + # More input sets than threads, so each thread solves several + (2, 4), ], ) def test_diagnostics_emitted_once_per_input_set( - self, decay_model, num_threads, n_inputs, all_on_calling_thread, caplog, capsys + self, decay_model, num_threads, n_inputs, caplog, capsys ): t_eval = np.linspace(0, 1, 3) inputs = [{"a": 1.0 + i} for i in range(n_inputs)] @@ -737,9 +796,9 @@ def test_diagnostics_emitted_once_per_input_set( assert len(starts) == len(stats) == n_inputs # Values are printed through py::print, so the tab prefix is preserved assert "\tNumber of steps =" in "\n".join(lines) - if all_on_calling_thread: - # Buffering would emit every trace before the first statistics block - assert stats[0] < starts[1] + # The calling thread streams its first set live, so that set's + # statistics come before every later trace + assert stats[0] < starts[1] def test_debug_log_emitted_when_solve_fails(self, caplog): model = pybamm.BaseModel() @@ -750,8 +809,8 @@ def test_debug_log_emitted_when_solve_fails(self, caplog): model.variables = {"u": u} pybamm.Discretisation().process_model(model) - # Two solves per solver, so one solver buffers on a worker thread and - # only the flush at the end of the parallel region can emit its trace + # More input sets than threads, so the worker thread buffers the sets it + # takes and only the flush after the parallel region can emit them solver = pybamm.IDAKLUSolver(options={"num_threads": 2}) inputs = [{"a": 1.0 + i} for i in range(4)] # u' = a u^2 blows up at t = 1/a, so integrating to t = 5 fails @@ -761,8 +820,8 @@ def test_debug_log_emitted_when_solve_fails(self, caplog): ): solver.solve(model, np.array([0.0, 5.0]), inputs=inputs) - # A partial solution is still returned, so every solve runs and the two - # traces beyond the calling thread's own prove its buffer was drained + # A partial solution is still returned, so every set is solved and the + # worker's traces prove its buffer was drained starts = [m for m in caplog.messages if m.startswith("Integrating from t =")] assert len(starts) == len(inputs) assert any(m.startswith("Step ") for m in caplog.messages) @@ -778,7 +837,7 @@ def test_debug_log_flushed_when_solve_raises(self, caplog): model.variables = {"u": u} pybamm.Discretisation().process_model(model) - # Two solves per solver, so one solver buffers on a worker thread + # More input sets than threads, so the worker thread buffers its sets solver = pybamm.IDAKLUSolver(options={"num_threads": 2}) inputs = [{"a": 1.0 + i} for i in range(4)] with ( @@ -787,10 +846,10 @@ def test_debug_log_flushed_when_solve_raises(self, caplog): ): solver.solve(model, np.array([0.0, 5.0]), inputs=inputs) - # Each solver throws on its first solve, and the calling thread streams - # its own, so the second trace can only come from the pre-rethrow flush + # A failing set does not stop the others, and the worker's traces can + # only come from the flush before the rethrow starts = [m for m in caplog.messages if m.startswith("Integrating from t =")] - assert len(starts) == 2 + assert len(starts) == 4 def test_solve_interrupted_from_debug_logger(self, caplog, monkeypatch): sim = pybamm.Simulation(pybamm.lithium_ion.SPM()) @@ -811,6 +870,49 @@ def logger_interrupts_on_first_step(msg, *args, **kwargs): assert not any(m.startswith("Integration complete") for m in caplog.messages) + @staticmethod + def _count_sets_started_after_interrupt( + model, num_threads, n_inputs, caplog, monkeypatch + ): + _debug_logger = pybamm.logger.debug + interrupted = [] + + # Interrupts only once, so the flush after the parallel region still + # emits the traces other threads buffered + def logger_interrupts_on_first_step(msg, *args, **kwargs): + _debug_logger(msg, *args, **kwargs) + if not interrupted and isinstance(msg, str) and msg.startswith("Step "): + interrupted.append(msg) + raise KeyboardInterrupt + + monkeypatch.setattr(pybamm.logger, "debug", logger_interrupts_on_first_step) + solver = pybamm.IDAKLUSolver(options={"num_threads": num_threads}) + inputs = [{"a": 1.0 + i} for i in range(n_inputs)] + with ( + caplog.at_level(logging.DEBUG, logger=pybamm.logger.name), + pytest.raises(KeyboardInterrupt), + ): + solver.solve(model, np.array([0.0, 5.0]), inputs=inputs) + return sum(m.startswith("Integrating from t =") for m in caplog.messages) + + def test_solve_interrupted_stops_later_input_sets( + self, decay_model, caplog, monkeypatch + ): + started = self._count_sets_started_after_interrupt( + decay_model, 1, 5, caplog, monkeypatch + ) + assert started == 1 + + def test_solve_interrupted_stops_every_thread( + self, decay_model, caplog, monkeypatch + ): + # The other thread may finish the set it holds, but takes no more + n_inputs = 16 + started = self._count_sets_started_after_interrupt( + decay_model, 2, n_inputs, caplog, monkeypatch + ) + assert started < n_inputs + def test_setup_options(self): model = pybamm.BaseModel() u = pybamm.Variable("u") @@ -2139,3 +2241,254 @@ def test_a_tight_dfn_diffusivity_gradient_matches_finite_differences(self): rtol=1e-4, atol=1e-5 * np.abs(finite_difference).max(), ) + + +class TestIDAKLUAtol: + def _two_state_model(self, atol=None): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + w = pybamm.Variable("w") + model.rhs = {u: -u, w: -2 * w} + model.initial_conditions = {u: 1.0, w: 1.0} + model.variables = {"u": u, "w": w} + pybamm.Discretisation().process_model(model) + if atol is not None: + model.atol = atol + return model + + def _steps(self, model, atol): + solution = pybamm.IDAKLUSolver(rtol=1e-6, atol=atol).solve( + model, np.linspace(0, 1, 5) + ) + return solution.solver_statistics.number_of_steps + + def test_a_per_state_model_atol_wins_over_the_solvers(self): + tight = self._steps(self._two_state_model(), 1e-12) + loose = self._steps(self._two_state_model(), 1e-1) + assert self._steps(self._two_state_model(np.full(2, 1e-1)), 1e-12) == loose + assert loose < tight + + def test_a_wrong_width_model_atol_is_rejected(self): + model = self._two_state_model(np.full(3, 1e-6)) + with pytest.raises(pybamm.SolverError, match=r"shape \(3,\) but \(2,\)"): + self._steps(model, 1e-6) + + @pytest.mark.parametrize( + "atol", + [ + 1e-3, + 1, + np.float32(1e-3), + np.array(1e-3), + [1e-3, 1e-3], + (1e-3, 1e-3), + np.full(2, 1e-3), + np.full((2, 1), 1e-3), + ], + ) + def test_accepted_atol_widens_to_one_value_per_state(self, atol): + model = self._two_state_model() + widened = pybamm.IDAKLUSolver()._check_atol_type(atol, model) + assert widened.dtype == np.float64 + assert widened.shape == (2,) + np.testing.assert_allclose(widened, np.full(2, float(np.ravel(atol)[0]))) + + @pytest.mark.parametrize( + ("atol", "match"), + [ + (True, r"must be a float, or a list, tuple or array"), + ({"u": 1e-3}, r"must be a float, or a list, tuple or array"), + ("1e-3", r"must be a float, or a list, tuple or array"), + ([True, False], r"must be real numbers"), + (["a", "b"], r"must be real numbers"), + (np.array([1e-3 + 1j, 1e-3]), r"must be real numbers"), + ([[1e-3], [1e-3, 1e-3]], r"must be a flat list"), + ([1e-3, 1e-3, 1e-3], r"shape \(3,\) but \(2,\)"), + (np.full((1, 2), 1e-3), r"shape \(1, 2\) but \(2,\)"), + (np.full((2, 2), 1e-3), r"shape \(2, 2\) but \(2,\)"), + ], + ) + def test_invalid_atol_is_rejected(self, atol, match): + model = self._two_state_model() + with pytest.raises(pybamm.SolverError, match=match): + pybamm.IDAKLUSolver()._check_atol_type(atol, model) + + @pytest.mark.parametrize("sequence", [list, tuple]) + def test_solve_with_a_sequence_atol(self, sequence): + model = self._two_state_model() + t_eval = np.linspace(0, 1, 5) + expected = pybamm.IDAKLUSolver(atol=np.array([1e-3, 1e-9])).solve(model, t_eval) + solution = pybamm.IDAKLUSolver(atol=sequence([1e-3, 1e-9])).solve(model, t_eval) + assert solution.solver_statistics == expected.solver_statistics + np.testing.assert_allclose(solution.y, expected.y, rtol=1e-12, atol=0) + + def test_per_state_atol_survives_a_config_round_trip(self): + model = self._two_state_model() + t_eval = np.linspace(0, 1, 5) + solver = pybamm.IDAKLUSolver(atol=np.array([1e-3, 1e-9])) + restored = pybamm.BaseSolver.from_config(solver.to_config()) + assert isinstance(restored.atol, list) + expected = solver.solve(model, t_eval) + solution = restored.solve(model, t_eval) + assert solution.solver_statistics == expected.solver_statistics + + +def _van_der_pol_model(): + """Discretised stiff van der Pol oscillator.""" + model = pybamm.BaseModel() + x = pybamm.Variable("x") + y = pybamm.Variable("y") + model.rhs = {x: y, y: 1000 * (1 - x**2) * y - x} + model.initial_conditions = {x: 2, y: 0} + model.variables = {"x": x} + pybamm.Discretisation().process_model(model) + return model + + +class TestIDAKLUSolverStatistics: + def test_statistics_match_the_printed_statistics(self): + # A cap of two Newton iterations makes every counter nonzero + solver = pybamm.IDAKLUSolver( + options={"print_stats": True, "max_nonlinear_iterations": 2} + ) + printed = io.StringIO() + with redirect_stdout(printed): + solution = solver.solve(_van_der_pol_model(), [0, 3000]) + + def printed_count(label): + return int(re.search(rf"\t{label} = (\d+)", printed.getvalue())[1]) + + stats = solution.solver_statistics + assert stats == pybamm.SolverStatistics( + number_of_steps=printed_count("Number of steps"), + number_of_linear_solver_setups=printed_count( + "Number of linear solver setup calls" + ), + number_of_nonlinear_solver_iterations=printed_count( + "Number of nonlinear iterations performed" + ), + number_of_nonlinear_solver_fails=printed_count( + "Number of nonlinear convergence failures" + ), + number_of_error_test_failures=printed_count( + "Number of error test failures" + ), + ) + assert all(count > 0 for count in dataclasses.astuple(stats)) + assert stats.number_of_nonlinear_solver_iterations >= stats.number_of_steps + + def test_statistics_accumulate_across_a_breakpoint(self, decay_model): + solver = pybamm.IDAKLUSolver() + inputs = {"a": 50.0} + first = solver.solve(decay_model, [0, 1], inputs=inputs).solver_statistics + # The breakpoint at t = 1 reinitialises the integrator, which resets + # its own counters; the tail after it takes far fewer steps than [0, 1] + both = solver.solve( + decay_model, np.array([0.0, 1.0, 2.0]), inputs=inputs + ).solver_statistics + assert both.number_of_steps > first.number_of_steps + assert ( + both.number_of_nonlinear_solver_iterations + > first.number_of_nonlinear_solver_iterations + ) + + def test_each_input_set_has_its_own_statistics(self, decay_model): + rates = (0.1, 1.0, 10.0, 100.0) + solutions = pybamm.IDAKLUSolver(options={"num_threads": 2}).solve( + decay_model, [0, 1], inputs=[{"a": a} for a in rates] + ) + alone = [ + pybamm.IDAKLUSolver().solve(decay_model, [0, 1], inputs={"a": a}) + for a in rates + ] + assert [s.solver_statistics for s in solutions] == [ + s.solver_statistics for s in alone + ] + assert len({s.solver_statistics.number_of_steps for s in solutions}) > 1 + + def test_reused_solver_resets_statistics(self, decay_model): + solver = pybamm.IDAKLUSolver() + first = solver.solve(decay_model, [0, 1], inputs={"a": 10.0}) + second = solver.solve(decay_model, [0, 1], inputs={"a": 10.0}) + assert second.solver_statistics == first.solver_statistics + + def test_step_with_save_sums_statistics(self, decay_model): + solver = pybamm.IDAKLUSolver() + first = solver.step(None, decay_model, 1, inputs={"a": 10.0}) + second = solver.step(first, decay_model, 1, inputs={"a": 10.0}, save=False) + combined = solver.step(first, decay_model, 1, inputs={"a": 10.0}) + assert combined.solver_statistics == ( + first.solver_statistics + second.solver_statistics + ) + + def test_experiment_solution_sums_its_steps(self): + experiment = pybamm.Experiment( + [ + ( + "Discharge at 1C for 10 minutes", + "Rest for 5 minutes", + "Charge at 1C until 4.1 V", + ) + ] + * 2 + ) + solution = pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment + ).solve() + steps = [step for cycle in solution.cycles for step in cycle.steps] + assert all(step.solver_statistics is not None for step in steps) + total = sum( + (step.solver_statistics for step in steps), pybamm.SolverStatistics() + ) + assert solution.solver_statistics == total + assert solution.cycles[0].solver_statistics.number_of_steps > 0 + + def test_reduced_solution_keeps_statistics(self, decay_model): + solver = pybamm.IDAKLUSolver(options={"hermite_reduction_factor": 1.0}) + solution = solver.solve(decay_model, [0, 1], inputs={"a": 10.0}) + reduced = solver.reduce_solution(solution) + assert reduced.solver_statistics == solution.solver_statistics + + @pytest.mark.parametrize( + ("preconditioner", "uses_preconditioner"), [("none", False), ("BBDP", True)] + ) + def test_iterative_linear_solver_statistics( + self, preconditioner, uses_preconditioner + ): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + model.rhs = {u: -pybamm.InputParameter("a") * u} + model.algebraic = {v: v - u} + model.initial_conditions = {u: 1, v: 1} + model.variables = {"u": u} + pybamm.Discretisation().process_model(model) + # Output variables with sensitivities fill the model data that a BBD + # counter read without a BBD preconditioner would alias + solver = pybamm.IDAKLUSolver( + output_variables=["u"], + options={ + "linear_solver": "SUNLinSol_SPBCGS", + "preconditioner": preconditioner, + "print_stats": True, + }, + ) + printed = io.StringIO() + with redirect_stdout(printed): + solution = solver.solve( + model, + np.linspace(0, 1, 3), + inputs={"a": 0.1}, + calculate_sensitivities=True, + ) + + evaluations = re.search( + r"residual function in preconditioner = (-?\d+)", printed.getvalue() + ) + if uses_preconditioner: + assert int(evaluations[1]) > 0 + else: + assert int(evaluations[1]) == 0 + stats = solution.solver_statistics + assert stats.number_of_steps > 0 + assert stats.number_of_nonlinear_solver_iterations >= stats.number_of_steps diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index b81cb19afe..f4a6b3858f 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -4,6 +4,7 @@ import io import json import logging +import pickle import subprocess # nosec B404 - used in tests with trusted input import sys from unittest import mock @@ -1839,3 +1840,73 @@ def test_from_sub_solutions_skipped_duplicate_preserves_hermite(self): assert out.hermite_interpolation == folded.hermite_interpolation assert out.hermite_interpolation is True np.testing.assert_array_equal(out.yp, folded.yp) + + +class TestSolutionSolverStatistics: + @staticmethod + def _solution(start, statistics): + t = np.linspace(start, start + 1, 5) + solution = pybamm.Solution(t, np.tile(t, (2, 1)), pybamm.BaseModel(), {}) + solution.solver_statistics = statistics + return solution + + def test_statistics_add_field_wise(self): + total = pybamm.SolverStatistics(1, 2, 3, 4, 5) + pybamm.SolverStatistics( + 10, 20, 30, 40, 50 + ) + assert total == pybamm.SolverStatistics(11, 22, 33, 44, 55) + assert pybamm.SolverStatistics() == pybamm.SolverStatistics(0, 0, 0, 0, 0) + with pytest.raises(TypeError): + pybamm.SolverStatistics() + 1 + + def test_a_new_solution_has_no_statistics(self): + t = np.linspace(0, 1) + solution = pybamm.Solution(t, np.tile(t, (2, 1)), pybamm.BaseModel(), {}) + assert solution.solver_statistics is None + + def test_add_sums_statistics(self): + first = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) + second = self._solution(1, pybamm.SolverStatistics(10, 20, 30, 40, 50)) + assert (first + second).solver_statistics == pybamm.SolverStatistics( + 11, 22, 33, 44, 55 + ) + + def test_add_is_none_when_a_side_has_no_statistics(self): + statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) + with_none_after = self._solution(0, statistics) + self._solution(1, None) + assert with_none_after.solver_statistics is None + with_none_before = self._solution(0, None) + self._solution(1, statistics) + assert with_none_before.solver_statistics is None + + def test_copy_and_adding_nothing_keep_statistics(self): + statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) + solution = self._solution(0, statistics) + assert solution.copy().solver_statistics == statistics + assert (solution + None).solver_statistics == statistics + assert (None + solution).solver_statistics == statistics + assert (pybamm.EmptySolution() + solution).solver_statistics == statistics + + def test_from_sub_solutions_sums_statistics(self): + solutions = [ + self._solution(i, pybamm.SolverStatistics(i, i, i, i, i)) for i in range(4) + ] + combined = pybamm.Solution.from_sub_solutions(solutions) + assert combined.solver_statistics == pybamm.SolverStatistics(6, 6, 6, 6, 6) + solutions[2].solver_statistics = None + assert pybamm.Solution.from_sub_solutions(solutions).solver_statistics is None + + def test_first_and_last_state_report_no_work(self): + solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) + assert solution.first_state.solver_statistics == pybamm.SolverStatistics() + assert solution.last_state.solver_statistics == pybamm.SolverStatistics() + + def test_statistics_survive_pickling(self): + statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) + solution = self._solution(0, statistics) + assert pickle.loads(pickle.dumps(solution)).solver_statistics == statistics + + def test_solution_pickled_without_statistics_loads(self): + solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) + # A solution pickled before the attribute existed has no such state + del solution.__dict__["solver_statistics"] + assert pickle.loads(pickle.dumps(solution)).solver_statistics is None diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp index d2c764da4e..26c9c5e5e0 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu.cpp @@ -274,6 +274,17 @@ PYBIND11_MODULE(idaklu, m) py::class_(m, "Function"); + py::class_(m, "SolverStats") + .def_readonly("number_of_steps", &IDAKLUStats::nsteps) + .def_readonly("number_of_residual_evaluations", &IDAKLUStats::nrevals) + .def_readonly("number_of_linear_solver_setups", &IDAKLUStats::nlinsetups) + .def_readonly("number_of_error_test_failures", &IDAKLUStats::netfails) + .def_readonly("number_of_nonlinear_solver_iterations", &IDAKLUStats::nniters) + .def_readonly("number_of_nonlinear_solver_fails", &IDAKLUStats::nncfails) + .def_readonly("number_of_jacobian_evaluations", &IDAKLUStats::njevals) + .def_readonly("number_of_linear_iterations", &IDAKLUStats::nliters) + .def_readonly("number_of_linear_convergence_failures", &IDAKLUStats::nlcfails); + py::class_(m, "solution") .def_readwrite("t", &Solution::t) .def_readwrite("y", &Solution::y) @@ -281,6 +292,7 @@ PYBIND11_MODULE(idaklu, m) .def_readwrite("yS", &Solution::yS) .def_readwrite("ypS", &Solution::ypS) .def_readwrite("y_term", &Solution::y_term) - .def_readwrite("flag", &Solution::flag); + .def_readwrite("flag", &Solution::flag) + .def_readonly("stats", &Solution::stats); } diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp index 6b87e78ec6..a552cb75d7 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Expressions/Base/Expression.hpp @@ -12,6 +12,11 @@ class Expression { */ Expression() = default; + /** + * @brief Virtual destructor (required for safe polymorphic deletion via base pointer) + */ + virtual ~Expression() = default; + /** * @brief Evaluation operator (for use after setting input and output data references) */ diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index 8f5ef017a4..32bacac2eb 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -1,5 +1,7 @@ #include "IDAKLUSolverGroup.hpp" #include +#include +#include #include #include @@ -97,9 +99,6 @@ std::vector IDAKLUSolverGroup::solve( "inputs has wrong number of rows. Expected " + std::to_string(number_of_groups) + " but got " + std::to_string(inputs.shape()[0])); - const std::size_t solves_per_thread = number_of_groups / m_solvers.size(); - const std::size_t remainder_solves = number_of_groups % m_solvers.size(); - // pbar is optional: an empty array leaves IDAS at its unit default. const bool has_pbar = pbar.size() > 0; if (has_pbar) { @@ -122,7 +121,8 @@ std::vector IDAKLUSolverGroup::solve( std::vector results(number_of_groups); - std::optional exception_message; + // One slot per input set, so the rethrow below can name every set that failed + std::vector> errors(number_of_groups); // Python exceptions carry their own type, which the string path below loses std::exception_ptr python_exception; @@ -130,29 +130,40 @@ std::vector IDAKLUSolverGroup::solve( // may log directly or must buffer until flush_logs(). set_loggers(logger); - omp_set_num_threads(m_solvers.size()); - #pragma omp parallel for - for (int i = 0; i < m_solvers.size(); i++) { - try { - for (int j = 0; j < solves_per_thread; j++) { - const std::size_t index = i * solves_per_thread + j; - const sunrealtype *y = y0 + index * y0_np.shape(1); - const sunrealtype *yp = yp0 + index * yp0_np.shape(1); - const sunrealtype *input = inputs_data + index * inputs.shape(1); - const sunrealtype *scales = pbar_data ? pbar_data + index * number_of_parameters : nullptr; - results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps); - } - } catch (py::error_already_set &) { - #pragma omp critical - { - python_exception = std::current_exception(); - } - } catch (std::exception &e) { - // If an exception is thrown, we need to catch it and rethrow it outside the parallel region - #pragma omp critical - { - exception_message = std::string(e.what()); + // num_threads() scopes the team to this region, unlike the process-wide + // omp_set_num_threads; the runtime may still start fewer threads than asked. + const int team_size = std::max(1, std::min(m_solvers.size(), number_of_groups)); + // A plain atomic, not schedule(dynamic): the macOS wheels' libomp lacks + // __kmpc_dispatch_deinit, and MSVC needs -openmp:llvm for omp atomic capture. + std::atomic next_group{1}; + std::atomic interrupted{false}; + #pragma omp parallel num_threads(team_size) + { + // Thread 0 is the calling thread, which holds the GIL, so it always takes a + // set and streams that set's diagnostics instead of buffering them. + const int thread = omp_get_thread_num(); + int i = thread == 0 ? 0 : next_group.fetch_add(1, std::memory_order_relaxed); + while (i < number_of_groups && !interrupted.load(std::memory_order_relaxed)) { + const sunrealtype *y = y0 + i * y0_np.shape(1); + const sunrealtype *yp = yp0 + i * yp0_np.shape(1); + const sunrealtype *input = inputs_data + i * inputs.shape(1); + const sunrealtype *scales = pbar_data ? pbar_data + i * number_of_parameters : nullptr; + try { + results[i] = m_solvers[thread]->solve( + t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps); + } catch (py::error_already_set &) { + // A Python exception such as KeyboardInterrupt stops the whole sweep + interrupted.store(true, std::memory_order_relaxed); + #pragma omp critical + { + if (!python_exception) { + python_exception = std::current_exception(); + } + } + } catch (std::exception &e) { + errors[i] = e.what(); } + i = next_group.fetch_add(1, std::memory_order_relaxed); } } @@ -163,23 +174,21 @@ std::vector IDAKLUSolverGroup::solve( std::rethrow_exception(python_exception); } - if (exception_message.has_value()) { - py::set_error(PyExc_ValueError, exception_message->c_str()); - throw py::error_already_set(); + std::string failures; + for (int i = 0; i < number_of_groups; i++) { + if (!errors[i].has_value()) { + continue; + } + if (!failures.empty()) { + failures += "; "; + } + failures += "input set " + std::to_string(i) + ": " + *errors[i]; } - - // Runs on this thread, so these solves log directly rather than buffering - for (int i = 0; i < remainder_solves; i++) { - const std::size_t index = number_of_groups - remainder_solves + i; - const sunrealtype *y = y0 + index * y0_np.shape(1); - const sunrealtype *yp = yp0 + index * yp0_np.shape(1); - const sunrealtype *input = inputs_data + index * inputs.shape(1); - const sunrealtype *scales = pbar_data ? pbar_data + index * number_of_parameters : nullptr; - results[index] = m_solvers[i]->solve(t_eval, t_interp, y, yp, input, scales, save_adaptive_steps, save_interp_steps); + if (!failures.empty()) { + py::set_error(PyExc_ValueError, failures.c_str()); + throw py::error_already_set(); } - flush_logs(); - // create solutions (needs to be serial as we're using the Python GIL) std::vector solutions(number_of_groups); for (int i = 0; i < number_of_groups; i++) { diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl index bc03a24cb7..f1a40dcadd 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverOpenMP.inl @@ -643,10 +643,8 @@ void IDAKLUSolverOpenMP::HandleBreakpoint( i_eval++; t_eval_next = t_eval[i_eval]; CheckErrors(IDASetStopTime(ida_mem, t_eval_next), "IDASetStopTime"); - if (solver_opts.print_stats) { - // Save stats before reinitializing (reinit resets IDA counters) - SaveStats(); - } + // Save stats before reinitializing (reinit resets IDA counters) + SaveStats(); // Reinitialize the solver to deal with the discontinuity at t = t_val ReinitializeIntegrator(t_val); @@ -661,8 +659,8 @@ template SolutionData IDAKLUSolverOpenMP::BuildSolutionData(int retval) { DEBUG("IDAKLUSolver::BuildSolutionData"); + SaveStats(); if (solver_opts.print_stats) { - SaveStats(); CaptureStats(); } @@ -707,7 +705,8 @@ SolutionData IDAKLUSolverOpenMP::BuildSolutionData(int retval) { arg_sens0, arg_sens1, arg_sens2, - save_hermite + save_hermite, + accumulated_stats ); } @@ -1140,6 +1139,10 @@ PendingStats IDAKLUSolverOpenMP::GetStats() { if (setup_opts.using_iterative_solver) { CheckErrors(IDAGetNumLinIters(ida_mem, &stats.nliters), "IDAGetNumLinIters"); CheckErrors(IDAGetNumLinConvFails(ida_mem, &stats.nlcfails), "IDAGetNumLinConvFails"); + } + // Without IDABBDPrecInit, IDA's preconditioner data is the user data, which + // this call would misread as the BBD preconditioner's. + if (setup_opts.preconditioner != "none") { CheckErrors(IDABBDPrecGetNumGfnEvals(ida_mem, &stats.ngevalsBBDP), "IDABBDPrecGetNumGfnEvals"); } diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp index ab7cbf15f5..e6b3e7b238 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp @@ -2,6 +2,7 @@ #define PYBAMM_IDAKLU_SOLUTION_HPP #include "common.hpp" +#include "IDAKLUStats.hpp" /** * @brief Solution class @@ -18,9 +19,10 @@ class Solution * @brief Constructor */ Solution(int retval, np_array t_np, np_array y_np, np_array yp_np, - np_array yS_np, np_array ypS_np, np_array y_term_np) + np_array yS_np, np_array ypS_np, np_array y_term_np, + IDAKLUStats stats) : flag(retval), t(t_np), y(y_np), yp(yp_np), yS(yS_np), ypS(ypS_np), - y_term(y_term_np) + y_term(y_term_np), stats(stats) { } @@ -36,6 +38,7 @@ class Solution np_array yS; np_array ypS; np_array y_term; + IDAKLUStats stats; }; #endif // PYBAMM_IDAKLU_SOLUTION_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp index 2cb7a8b565..620b8d786e 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp @@ -48,7 +48,8 @@ class SolutionData ptrdiff_t arg_sens0, ptrdiff_t arg_sens1, ptrdiff_t arg_sens2, - bool save_hermite) + bool save_hermite, + IDAKLUStats stats) : flag(flag), t_vec(std::move(t)), y_vec(std::move(y)), @@ -59,7 +60,8 @@ class SolutionData arg_sens0(arg_sens0), arg_sens1(arg_sens1), arg_sens2(arg_sens2), - save_hermite(save_hermite) + save_hermite(save_hermite), + stats(stats) {} ~SolutionData() = default; @@ -81,7 +83,8 @@ class SolutionData vector_to_numpy_3d(std::move(yS_vec), arg_sens0, arg_sens1, arg_sens2), vector_to_numpy_3d(std::move(ypS_vec), save_hermite ? arg_sens0 : 0, arg_sens1, arg_sens2), - vector_to_numpy(std::move(yterm_vec)) + vector_to_numpy(std::move(yterm_vec)), + stats ); } @@ -97,6 +100,7 @@ class SolutionData ptrdiff_t arg_sens1 = 0; ptrdiff_t arg_sens2 = 0; bool save_hermite = false; + IDAKLUStats stats; }; #endif // PYBAMM_IDAKLU_SOLUTION_DATA_HPP diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp index 8f6a6f649f..e2cad93314 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.cpp @@ -5,12 +5,13 @@ // ────────────────────── StandaloneAlgebraicSystem ────────────────────── StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( - casadi::Function res_fn, - casadi::Function jac_fn, + std::unique_ptr res_fn, + std::unique_ptr jac_fn, + int n_vars, bool use_sparse) - : res_cf_(res_fn), - jac_cf_(jac_fn), - n_vars_(static_cast(res_fn.nnz_out(0))), + : res_(std::move(res_fn)), + jac_(std::move(jac_fn)), + n_vars_(n_vars), use_sparse_(use_sparse), sunctx_(nullptr), J_(nullptr), LS_(nullptr), res_nvec_(nullptr), delta_nvec_(nullptr) @@ -20,7 +21,7 @@ StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( res_nvec_ = N_VNew_Serial(n_vars_, sunctx_); delta_nvec_ = N_VNew_Serial(n_vars_, sunctx_); - int jac_nnz = static_cast(jac_cf_.nnz_out()); + int jac_nnz = static_cast(jac_->nnz_out()); jac_buf_.resize(jac_nnz > 0 ? jac_nnz : n_vars_ * n_vars_); if (use_sparse_ && jac_nnz > 0) { @@ -31,6 +32,17 @@ StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( } } +StandaloneAlgebraicSystem::StandaloneAlgebraicSystem( + casadi::Function res_fn, + casadi::Function jac_fn, + bool use_sparse) + : StandaloneAlgebraicSystem( + std::make_unique(res_fn), + std::make_unique(jac_fn), + static_cast(res_fn.nnz_out(0)), + use_sparse) +{} + StandaloneAlgebraicSystem::~StandaloneAlgebraicSystem() { if (res_nvec_) N_VDestroy(res_nvec_); if (delta_nvec_) N_VDestroy(delta_nvec_); @@ -42,22 +54,22 @@ StandaloneAlgebraicSystem::~StandaloneAlgebraicSystem() { void StandaloneAlgebraicSystem::eval_residual( sunrealtype t, const sunrealtype* y, sunrealtype* res) { - res_cf_.m_arg[0] = &t; - res_cf_.m_arg[1] = y; - res_cf_.m_arg[2] = inputs_.data(); - res_cf_.m_res[0] = res; - res_cf_(); + res_->m_arg[0] = &t; + res_->m_arg[1] = y; + res_->m_arg[2] = inputs_.data(); + res_->m_res[0] = res; + (*res_)(); } int StandaloneAlgebraicSystem::solve_linear( sunrealtype t, const sunrealtype* y, sunrealtype* res, sunrealtype* delta) { - jac_cf_.m_arg[0] = &t; - jac_cf_.m_arg[1] = y; - jac_cf_.m_arg[2] = inputs_.data(); - jac_cf_.m_res[0] = jac_buf_.data(); - jac_cf_(); + jac_->m_arg[0] = &t; + jac_->m_arg[1] = y; + jac_->m_arg[2] = inputs_.data(); + jac_->m_res[0] = jac_buf_.data(); + (*jac_)(); if (use_sparse_) { sunrealtype* mat_data = SUNSparseMatrix_Data(J_); @@ -88,8 +100,8 @@ int StandaloneAlgebraicSystem::solve_linear( } void StandaloneAlgebraicSystem::BuildSparseResources(int jac_nnz) { - const auto& rows = jac_cf_.get_row(); - const auto& cols = jac_cf_.get_col(); + const auto& rows = jac_->get_row(); + const auto& cols = jac_->get_col(); int nnz_total = jac_nnz; // Build CSC from COO (same O(nnz) algorithm as AlgebraicICBuilder) @@ -132,9 +144,9 @@ void StandaloneAlgebraicSystem::BuildSparseResources(int jac_nnz) { } void StandaloneAlgebraicSystem::BuildDenseResources() { - const auto& rows = jac_cf_.get_row(); - const auto& cols = jac_cf_.get_col(); - int nnz_total = static_cast(jac_cf_.nnz_out()); + const auto& rows = jac_->get_row(); + const auto& cols = jac_->get_col(); + int nnz_total = static_cast(jac_->nnz_out()); // Build CSC structure for dense scatter std::vector col_count(n_vars_ + 1, 0); diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp index b4e1c6e5be..8f05d05a14 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp @@ -8,7 +8,7 @@ #include /** - * @brief NonlinearSystem backed by standalone CasadiFunctions + SUNLinSol. + * @brief NonlinearSystem backed by standalone Expression functions + SUNLinSol. * * Self-contained: no dependency on IDAKLUSolverOpenMP or IDA memory. * Residual signature: F(t, y_alg, inputs) -> res (n_vars outputs) @@ -16,6 +16,22 @@ */ class StandaloneAlgebraicSystem : public NonlinearSystem { public: + /** + * @brief Construct from residual and Jacobian Expressions. + * @param res_fn Residual function F(t, y, inputs) -> res + * @param jac_fn Jacobian function J(t, y, inputs) -> data (COO) + * @param n_vars Number of algebraic variables (residual output length) + * @param use_sparse Whether to use a sparse (KLU) linear solver + */ + StandaloneAlgebraicSystem( + std::unique_ptr res_fn, + std::unique_ptr jac_fn, + int n_vars, + bool use_sparse); + + /** + * @brief Casadi convenience constructor; delegates to the Expression ctor. + */ StandaloneAlgebraicSystem( casadi::Function res_fn, casadi::Function jac_fn, @@ -39,8 +55,8 @@ class StandaloneAlgebraicSystem : public NonlinearSystem { void BuildSparseResources(int jac_nnz); void BuildDenseResources(); - CasadiFunction res_cf_; - CasadiFunction jac_cf_; + std::unique_ptr res_; + std::unique_ptr jac_; int n_vars_; bool use_sparse_; diff --git a/packages/pybammsolvers/tests/test_integration.py b/packages/pybammsolvers/tests/test_integration.py index 50b5246720..425457aa81 100644 --- a/packages/pybammsolvers/tests/test_integration.py +++ b/packages/pybammsolvers/tests/test_integration.py @@ -115,6 +115,38 @@ def test_solution_has_termination_flag(self, exponential_decay_solver): # IDA_SUCCESS=0, IDA_TSTOP_RETURN=1, IDA_ROOT_RETURN=2 are all success codes assert sol.flag in [0, 1, 2], f"Solver failed with flag {sol.flag}" + def test_solution_has_solver_statistics(self, exponential_decay_solver): + """ + Verify Solution object carries the integrator's counters for its solve. + """ + solver_data = exponential_decay_solver + solver = solver_data["solver"] + t_eval = solver_data["model"]["t_eval"] + + solution = solver.solve( + t_eval, t_eval, solver_data["y0"], solver_data["yp0"], solver_data["inputs"] + ) + stats = solution[0].stats + + counters = [ + "number_of_steps", + "number_of_residual_evaluations", + "number_of_linear_solver_setups", + "number_of_error_test_failures", + "number_of_nonlinear_solver_iterations", + "number_of_nonlinear_solver_fails", + "number_of_jacobian_evaluations", + "number_of_linear_iterations", + "number_of_linear_convergence_failures", + ] + for counter in counters: + value = getattr(stats, counter) + assert isinstance(value, int) + assert value >= 0 + assert stats.number_of_steps > 0 + assert stats.number_of_residual_evaluations > 0 + assert stats.number_of_nonlinear_solver_iterations >= stats.number_of_steps + def test_solution_accuracy_exponential_decay(self, exponential_decay_solver): """ Verify Solution matches exact solution for exponential decay. From 9ac563a5c9002a5eed979b14dbfec0f1d73b3194 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Thu, 24 Sep 2026 14:15:44 +0100 Subject: [PATCH 02/13] test: give the interrupted sweep more sets than a slow thread 0 can race With 16 sets the GIL-free thread could finish all 15 of its sets before thread 0 reached its first step on a loaded macOS runner. --- .../pybamm/tests/unit/test_solvers/test_idaklu_solver.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 41e0d62ac4..5777f0bdb4 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -906,8 +906,9 @@ def test_solve_interrupted_stops_later_input_sets( def test_solve_interrupted_stops_every_thread( self, decay_model, caplog, monkeypatch ): - # The other thread may finish the set it holds, but takes no more - n_inputs = 16 + # The other thread may finish the set it holds, but takes no more. Enough + # sets that it cannot finish them all before thread 0's first step + n_inputs = 4000 started = self._count_sets_started_after_interrupt( decay_model, 2, n_inputs, caplog, monkeypatch ) From 3d540d5352aff6299743b3b97db29e3911444f93 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Thu, 24 Sep 2026 14:37:06 +0100 Subject: [PATCH 03/13] style: clear Codacy's findings on the new statistics and scheduling code Pass IDAKLUStats by const reference; suppress cppcheck where it cannot see a use (the omp num_threads clause, members used only in the .cpp); mark the tests' pickle round trips nosec, as main does. --- packages/pybamm/tests/unit/test_solvers/test_solution.py | 6 +++--- .../src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp | 1 + .../src/pybammsolvers/idaklu_source/Solution.hpp | 2 +- .../src/pybammsolvers/idaklu_source/SolutionData.hpp | 2 +- .../pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp | 4 ++-- 5 files changed, 8 insertions(+), 7 deletions(-) diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index f4a6b3858f..9de0720412 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -4,7 +4,7 @@ import io import json import logging -import pickle +import pickle # nosec B403 - used in tests with trusted input import subprocess # nosec B404 - used in tests with trusted input import sys from unittest import mock @@ -1903,10 +1903,10 @@ def test_first_and_last_state_report_no_work(self): def test_statistics_survive_pickling(self): statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) solution = self._solution(0, statistics) - assert pickle.loads(pickle.dumps(solution)).solver_statistics == statistics + assert pickle.loads(pickle.dumps(solution)).solver_statistics == statistics # nosec B301 def test_solution_pickled_without_statistics_loads(self): solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) # A solution pickled before the attribute existed has no such state del solution.__dict__["solver_statistics"] - assert pickle.loads(pickle.dumps(solution)).solver_statistics is None + assert pickle.loads(pickle.dumps(solution)).solver_statistics is None # nosec B301 diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index 32bacac2eb..a5d763312f 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -132,6 +132,7 @@ std::vector IDAKLUSolverGroup::solve( // num_threads() scopes the team to this region, unlike the process-wide // omp_set_num_threads; the runtime may still start fewer threads than asked. + // cppcheck-suppress unreadVariable const int team_size = std::max(1, std::min(m_solvers.size(), number_of_groups)); // A plain atomic, not schedule(dynamic): the macOS wheels' libomp lacks // __kmpc_dispatch_deinit, and MSVC needs -openmp:llvm for omp atomic capture. diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp index e6b3e7b238..9998f883e4 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/Solution.hpp @@ -20,7 +20,7 @@ class Solution */ Solution(int retval, np_array t_np, np_array y_np, np_array yp_np, np_array yS_np, np_array ypS_np, np_array y_term_np, - IDAKLUStats stats) + const IDAKLUStats &stats) : flag(retval), t(t_np), y(y_np), yp(yp_np), yS(yS_np), ypS(ypS_np), y_term(y_term_np), stats(stats) { diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp index 620b8d786e..f0e03071fc 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp @@ -49,7 +49,7 @@ class SolutionData ptrdiff_t arg_sens1, ptrdiff_t arg_sens2, bool save_hermite, - IDAKLUStats stats) + const IDAKLUStats &stats) : flag(flag), t_vec(std::move(t)), y_vec(std::move(y)), diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp index 8f05d05a14..4b7b1bcf22 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/StandaloneNewtonSolver.hpp @@ -55,8 +55,8 @@ class StandaloneAlgebraicSystem : public NonlinearSystem { void BuildSparseResources(int jac_nnz); void BuildDenseResources(); - std::unique_ptr res_; - std::unique_ptr jac_; + std::unique_ptr res_; // cppcheck-suppress unusedStructMember + std::unique_ptr jac_; // cppcheck-suppress unusedStructMember int n_vars_; bool use_sparse_; From ef211b1021df565bf6ed71896c7fd2d1b7a1c2ed Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 11:43:06 +0100 Subject: [PATCH 04/13] style: say why the input-set sweep emulates schedule(dynamic) Drop the num_threads() comment, which described the change rather than the code, and name what the atomic counter stands in for. --- .../src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index a5d763312f..15f626e451 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -130,12 +130,10 @@ std::vector IDAKLUSolverGroup::solve( // may log directly or must buffer until flush_logs(). set_loggers(logger); - // num_threads() scopes the team to this region, unlike the process-wide - // omp_set_num_threads; the runtime may still start fewer threads than asked. // cppcheck-suppress unreadVariable const int team_size = std::max(1, std::min(m_solvers.size(), number_of_groups)); - // A plain atomic, not schedule(dynamic): the macOS wheels' libomp lacks - // __kmpc_dispatch_deinit, and MSVC needs -openmp:llvm for omp atomic capture. + // Emulates schedule(dynamic) with an atomic counter: the macOS wheels' libomp + // lacks __kmpc_dispatch_deinit, and MSVC needs -openmp:llvm for omp atomic capture. std::atomic next_group{1}; std::atomic interrupted{false}; #pragma omp parallel num_threads(team_size) From 3b47fa5982506a4fc0e224a0764e17ddaf361a64 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 11:43:07 +0100 Subject: [PATCH 05/13] fix: count a single-sample duplicate's solver statistics when adding solutions The __add__ short-circuit, and from_sub_solutions skipping the same segment, dropped that operand's statistics. Its samples add nothing, but the solver may still have worked for it. --- .../pybamm/src/pybamm/solvers/solution.py | 11 +++----- .../tests/unit/test_solvers/test_solution.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index ebb37ec5d0..ee16ab30d2 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -1167,6 +1167,7 @@ def __add__(self, other): new_sol._termination = other.termination new_sol._t_event = other._t_event new_sol._y_event = other._y_event + new_sol.solver_statistics = _sum_solver_statistics([self, other]) return new_sol # Append other onto self (the already-merged left side); only other's @@ -1273,12 +1274,7 @@ def from_sub_solutions(cls, sub_solutions): if len(sols) == 1: return sols[0].copy() - # Decide once which segments contribute. __add__ short-circuits a - # single-sample segment whose only sample duplicates the running - # boundary to a copy, so it contributes nothing to the merge: skip it - # here and derive every quantity below from the kept segments only. - # `repeated` records whether a kept segment's leading sample duplicates - # the previous boundary (dropped once on concatenation). + # Skip a single-sample duplicate of the running boundary, as __add__ does kept = [] prev_last_t = None for s in sols: @@ -1349,7 +1345,8 @@ def from_sub_solutions(cls, sub_solutions): vals = [getattr(s, attr, None) for s in segments] if all(v is not None for v in vals): setattr(new_sol, attr, sum(vals)) - new_sol.solver_statistics = _sum_solver_statistics(segments) + # sols, not segments: __add__ counts a skipped duplicate's solver work + new_sol.solver_statistics = _sum_solver_statistics(sols) # output_variables path: reproduce __add__'s pairwise left-fold. if any(s.variables_returned for s in segments): diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 9de0720412..b2fa8d342f 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -1895,6 +1895,34 @@ def test_from_sub_solutions_sums_statistics(self): solutions[2].solver_statistics = None assert pybamm.Solution.from_sub_solutions(solutions).solver_statistics is None + def test_add_sums_statistics_of_a_single_sample_duplicate(self): + first = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) + duplicate = pybamm.Solution( + np.array([1.0]), np.ones((2, 1)), pybamm.BaseModel(), {} + ) + duplicate.solver_statistics = pybamm.SolverStatistics(10, 20, 30, 40, 50) + combined = first + duplicate + assert combined.t[-1] == pytest.approx(1.0) + assert combined.solver_statistics == pybamm.SolverStatistics(11, 22, 33, 44, 55) + + def test_from_sub_solutions_sums_statistics_of_a_skipped_duplicate(self): + import functools + import operator + + first = self._solution(0, pybamm.SolverStatistics(1, 1, 1, 1, 1)) + duplicate = pybamm.Solution( + np.array([1.0]), np.ones((2, 1)), pybamm.BaseModel(), {} + ) + duplicate.solver_statistics = pybamm.SolverStatistics(10, 10, 10, 10, 10) + second = self._solution(1, pybamm.SolverStatistics(100, 100, 100, 100, 100)) + solutions = [first, duplicate, second] + combined = pybamm.Solution.from_sub_solutions(solutions) + folded = functools.reduce(operator.add, solutions) + assert combined.solver_statistics == pybamm.SolverStatistics( + 111, 111, 111, 111, 111 + ) + assert folded.solver_statistics == combined.solver_statistics + def test_first_and_last_state_report_no_work(self): solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) assert solution.first_state.solver_statistics == pybamm.SolverStatistics() From 6c9d1ec2e446498f7fc630c1ad2eb64dabaf4351 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 13:39:52 +0100 Subject: [PATCH 06/13] fix(idaklu): name every input set that fails part-way A set whose integration fails part-way returns a partial solution with a negative flag instead of throwing, so the per-set post-processing raised on the first such set and hid the rest. Collect them all after the sweep, and when another set throws, name the part-way failures alongside it, so the SolverError and each on_failure="warn" warning carry "input set N: ". --- CHANGELOG.md | 2 +- .../src/pybamm/solvers/idaklu_solver.py | 24 ++++++---- .../unit/test_solvers/test_idaklu_solver.py | 46 +++++++++++++++++++ .../idaklu_source/IDAKLUSolverGroup.cpp | 30 ++++++++---- .../idaklu_source/SolutionData.hpp | 5 ++ 5 files changed, 86 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e436f0e80d..2b72173f4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ## Features - `Solution.solver_statistics` reports the integrator's step, linear-solver-setup, nonlinear-iteration, nonlinear-failure and error-test-failure counts as a `pybamm.SolverStatistics` dataclass. `IDAKLUSolver` fills it for every solution, counting across `t_eval` breakpoints and separately for each input set, and the counts are summed when solutions are combined, so a stepped or experiment solution reports its total. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) -- A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, whether it failed from the start or part-way through, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. Each `on_failure="warn"` warning names its set the same way. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) ## Bug fixes diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 68934ae8df..db3c7a6057 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -785,6 +785,19 @@ def _integrate( raise pybamm.SolverError(str(e)) from None integration_time = timer.time() + # A set that fails part-way returns a partial solution with a negative + # flag instead of throwing, so its failure is reported here + failures = [ + f"input set {index}: {idaklu.sundials_error_message(soln.flag)}" + for index, soln in enumerate(solns) + if soln.flag < 0 + ] + if failures and self._on_failure == "error": + raise pybamm.SolverError("; ".join(failures)) + if self._on_failure == "warn": + for failure in failures: + warnings.warn(failure + ", returning a partial solution.", stacklevel=2) + return [ self._post_process_solution( soln, model, integration_time, inputs_dict, t_eval @@ -850,17 +863,8 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev termination = "event" elif sol.flag >= 0: termination = "final time" - elif sol.flag < 0: + else: termination = "failure" - msg = idaklu.sundials_error_message(sol.flag) - match self._on_failure: - case "warn": - warnings.warn( - msg + ", returning a partial solution.", - stacklevel=2, - ) - case "error": - raise pybamm.SolverError(msg) if sol.yp.size > 0: yp = sol.yp.reshape((number_of_timesteps, number_of_states)).T 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 5777f0bdb4..d1c50e3651 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -99,6 +99,21 @@ def decay_model(): return model +def _part_way_failure_model(): + """Discretised ``du/dt = sqrt(k - t)``, with ``k`` as the only input parameter. + + Its residual is NaN once ``t > k``, so a solve past ``t = k > 0`` fails + part-way and one with ``k < 0`` fails from the start. + """ + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: pybamm.sqrt(pybamm.InputParameter("k") - pybamm.t)} + model.initial_conditions = {u: 0} + model.variables = {"u": u} + pybamm.Discretisation().process_model(model) + return model + + class TestIDAKLUSolver: def test_ida_roberts_klu(self): # this test implements a python version of the ida Roberts @@ -193,6 +208,37 @@ def test_a_single_input_set_failure_is_named(self): with pytest.raises(pybamm.SolverError, match=r"^input set 0: "): pybamm.IDAKLUSolver().solve(model, [0, 1], inputs={"k": -1.0}) + def test_every_input_set_failing_part_way_is_named(self): + solver = pybamm.IDAKLUSolver(options={"num_threads": 4}) + inputs_list = [{"k": k} for k in (2.0, 0.5, 2.0, 0.3)] + with pytest.raises(pybamm.SolverError) as error: + solver.solve(_part_way_failure_model(), [0, 1], inputs=inputs_list) + assert re.findall(r"input set (\d+): IDA_", str(error.value)) == ["1", "3"] + + def test_each_partial_solution_warning_names_its_input_set(self): + solver = pybamm.IDAKLUSolver(options={"num_threads": 4}, on_failure="warn") + inputs_list = [{"k": k} for k in (2.0, 0.5, 2.0, 0.3)] + with pytest.warns(UserWarning, match="returning a partial solution") as record: + solutions = solver.solve( + _part_way_failure_model(), [0, 1], inputs=inputs_list + ) + named = [re.match(r"input set (\d+): ", str(w.message)) for w in record] + assert [match.group(1) for match in named] == ["1", "3"] + assert [solution.termination for solution in solutions] == [ + "final time", + "failure", + "final time", + "failure", + ] + + def test_a_throwing_input_set_also_names_the_part_way_failures(self): + solver = pybamm.IDAKLUSolver(options={"num_threads": 4}) + # Set 2 fails from the start, sets 1 and 4 part-way + inputs_list = [{"k": k} for k in (2.0, 0.5, -1.0, 2.0, 0.3)] + with pytest.raises(pybamm.SolverError) as error: + solver.solve(_part_way_failure_model(), [0, 1], inputs=inputs_list) + assert re.findall(r"input set (\d+): ", str(error.value)) == ["1", "2", "4"] + def test_every_input_set_is_solved_by_a_smaller_team(self): # The OpenMP runtime reads OMP_THREAD_LIMIT once, so this needs a fresh # process; the team then has fewer threads than the solver has solvers. diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp index 15f626e451..a802a1e3fa 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -1,4 +1,5 @@ #include "IDAKLUSolverGroup.hpp" +#include "sundials_error_handler.hpp" #include #include #include @@ -173,17 +174,26 @@ std::vector IDAKLUSolverGroup::solve( std::rethrow_exception(python_exception); } - std::string failures; - for (int i = 0; i < number_of_groups; i++) { - if (!errors[i].has_value()) { - continue; - } - if (!failures.empty()) { - failures += "; "; + const bool any_thrown = std::any_of( + errors.begin(), errors.end(), [](const auto &error) { return error.has_value(); }); + if (any_thrown) { + // No solution is returned, so also name the sets whose integration failed + // part-way, which would otherwise surface as partial solutions + std::string failures; + for (int i = 0; i < number_of_groups; i++) { + std::string error; + if (errors[i].has_value()) { + error = *errors[i]; + } else if (results[i].get_flag() < 0) { + error = sundials_error_message(results[i].get_flag()); + } else { + continue; + } + if (!failures.empty()) { + failures += "; "; + } + failures += "input set " + std::to_string(i) + ": " + error; } - failures += "input set " + std::to_string(i) + ": " + *errors[i]; - } - if (!failures.empty()) { py::set_error(PyExc_ValueError, failures.c_str()); throw py::error_already_set(); } diff --git a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp index f0e03071fc..e1b2808f98 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/SolutionData.hpp @@ -70,6 +70,11 @@ class SolutionData SolutionData(SolutionData&&) noexcept = default; SolutionData& operator=(SolutionData&&) noexcept = default; + /** + * @brief The IDA return flag, negative if the integration failed. + */ + int get_flag() const { return flag; } + /** * @brief Convert raw vectors to numpy arrays and create Solution. * MUST be called with GIL held (i.e., in serial section). From 407f7cbf6d18c3fd2161307a09a7ee5e641b2518 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 13:40:30 +0100 Subject: [PATCH 07/13] fix: count the solver statistics of the cycles save_at_cycles drops An experiment solution is folded from the cycles it keeps, so the solver work of a cycle that save_at_cycles leaves out went uncounted. Total those cycles' statistics as they are dropped, keeping none of their solutions alive, and add them to the folded solution's. --- CHANGELOG.md | 2 +- .../src/pybamm/simulation/simulation.py | 20 +++++++++++ .../pybamm/tests/memory/test_memory_leaks.py | 34 +++++++++++++++++++ .../unit/test_solvers/test_idaklu_solver.py | 14 ++++++++ 4 files changed, 69 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b72173f4b..169521f0ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ ## Features -- `Solution.solver_statistics` reports the integrator's step, linear-solver-setup, nonlinear-iteration, nonlinear-failure and error-test-failure counts as a `pybamm.SolverStatistics` dataclass. `IDAKLUSolver` fills it for every solution, counting across `t_eval` breakpoints and separately for each input set, and the counts are summed when solutions are combined, so a stepped or experiment solution reports its total. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- `Solution.solver_statistics` reports the integrator's step, linear-solver-setup, nonlinear-iteration, nonlinear-failure and error-test-failure counts as a `pybamm.SolverStatistics` dataclass. `IDAKLUSolver` fills it for every solution, counting across `t_eval` breakpoints and separately for each input set, and the counts are summed when solutions are combined, so a stepped or experiment solution reports its total, including the cycles that `save_at_cycles` does not keep. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, whether it failed from the start or part-way through, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. Each `on_failure="warn"` warning names its set the same way. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) ## Bug fixes diff --git a/packages/pybamm/src/pybamm/simulation/simulation.py b/packages/pybamm/src/pybamm/simulation/simulation.py index de112cb39b..c92689e138 100644 --- a/packages/pybamm/src/pybamm/simulation/simulation.py +++ b/packages/pybamm/src/pybamm/simulation/simulation.py @@ -987,6 +987,9 @@ def solve( # Track running termination (real cycles + skipped steps), applied to the # folded solution below to match the old left-fold's final termination. last_termination = None + # Solver work of the cycles left out of the fold, totalled so their + # solutions can be freed; None once one of them has no statistics + discarded_cycle_statistics = pybamm.SolverStatistics() for cycle_num, cycle_length in enumerate( cycle_lengths, @@ -1172,6 +1175,16 @@ def solve( cross_cycle_segments.append(cycle_solution) if not isinstance(cycle_solution, pybamm.EmptySolution): last_termination = cycle_solution.termination + elif ( + isinstance(cycle_solution, pybamm.Solution) + and discarded_cycle_statistics is not None + ): + statistics = cycle_solution.solver_statistics + discarded_cycle_statistics = ( + None + if statistics is None + else discarded_cycle_statistics + statistics + ) if steps: if all( @@ -1242,6 +1255,13 @@ def solve( self._solution = pybamm.Solution.from_sub_solutions(cross_cycle_segments) if last_termination is not None: self._solution.termination = last_termination + if isinstance(self._solution, pybamm.Solution): + statistics = self._solution.solver_statistics + self._solution.solver_statistics = ( + None + if statistics is None or discarded_cycle_statistics is None + else statistics + discarded_cycle_statistics + ) if self._solution is not None and len(all_cycle_solutions) > 0: self._solution.cycles = all_cycle_solutions diff --git a/packages/pybamm/tests/memory/test_memory_leaks.py b/packages/pybamm/tests/memory/test_memory_leaks.py index 4dc0c7c12f..063bc669a6 100644 --- a/packages/pybamm/tests/memory/test_memory_leaks.py +++ b/packages/pybamm/tests/memory/test_memory_leaks.py @@ -1,5 +1,6 @@ import gc import tracemalloc +import weakref import pybamm @@ -288,6 +289,39 @@ def test_gitt_memory(self): peak_mb = peak / 1024 / 1024 assert peak_mb < 6, f"Peak memory {peak_mb:.1f} MB for GITT is excessive." + def test_unsaved_cycles_are_freed_before_the_experiment_ends(self): + # Each cycle is folded from its steps by Solution.__add__, so a cycle + # that save_at_cycles leaves out must not keep its sum alive + def live_sums_at_experiment_end(number_of_cycles): + sums = [] + original = pybamm.Solution.__add__ + + def recording(self, other): + result = original(self, other) + sums.append(weakref.ref(result)) + return result + + class CountLiveSums(pybamm.callbacks.Callback): + def on_experiment_end(self, logs): + gc.collect() + self.live = sum(ref() is not None for ref in sums) + + callback = CountLiveSums() + experiment = pybamm.Experiment( + [("Discharge at 1C for 5 minutes", "Rest for 2 minutes")] + * number_of_cycles + ) + pybamm.Solution.__add__ = recording + try: + pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment + ).solve(save_at_cycles=[1, number_of_cycles], callbacks=callback) + finally: + pybamm.Solution.__add__ = original + return callback.live + + assert live_sums_at_experiment_end(6) == live_sums_at_experiment_end(3) + def test_processed_variable_computed_initialise_is_lazy(self): # initialise_* (np.concatenate / flatten / xr.DataArray build) must # defer until the data is actually read, otherwise per-step 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 d1c50e3651..3c85a0734a 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -2490,6 +2490,20 @@ def test_experiment_solution_sums_its_steps(self): assert solution.solver_statistics == total assert solution.cycles[0].solver_statistics.number_of_steps > 0 + def test_experiment_solution_counts_the_cycles_it_does_not_save(self): + experiment = pybamm.Experiment( + [("Discharge at 1C for 10 minutes", "Rest for 5 minutes")] * 4 + ) + every_cycle = pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment + ).solve() + some_cycles = pybamm.Simulation( + pybamm.lithium_ion.SPM(), experiment=experiment + ).solve(save_at_cycles=3) + # Cycle 2 is solved but not saved + assert some_cycles.cycles[1] is None + assert some_cycles.solver_statistics == every_cycle.solver_statistics + def test_reduced_solution_keeps_statistics(self, decay_model): solver = pybamm.IDAKLUSolver(options={"hermite_reduction_factor": 1.0}) solution = solver.solve(decay_model, [0, 1], inputs={"a": 10.0}) From 05f63ed4044936358d41e925574c702761391498 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 13:40:42 +0100 Subject: [PATCH 08/13] fix: reject a boolean among numeric per-state atol values NumPy upcasts a bool beside numbers, so [True, 1e-6] became a float array that passed the dtype check as a tolerance of 1.0. --- CHANGELOG.md | 2 +- packages/pybamm/src/pybamm/solvers/idaklu_solver.py | 8 ++++++++ .../pybamm/tests/unit/test_solvers/test_idaklu_solver.py | 3 +++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 169521f0ad..7d6327f8a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,7 @@ ## Bug fixes -- A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans, even one among numbers such as `[True, 1e-6]`, and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) ## Optimizations diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index db3c7a6057..735306a91a 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -414,6 +414,14 @@ def _check_atol_type( 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( 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 3c85a0734a..763e5fe617 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -2347,6 +2347,9 @@ def test_accepted_atol_widens_to_one_value_per_state(self, atol): ({"u": 1e-3}, r"must be a float, or a list, tuple or array"), ("1e-3", r"must be a float, or a list, tuple or array"), ([True, False], r"must be real numbers"), + ([True, 1e-3], r"must be real numbers, not booleans"), + ((1e-3, np.False_), r"must be real numbers, not booleans"), + ([[1e-3], [True]], r"must be real numbers, not booleans"), (["a", "b"], r"must be real numbers"), (np.array([1e-3 + 1j, 1e-3]), r"must be real numbers"), ([[1e-3], [1e-3, 1e-3]], r"must be a flat list"), From 6844c7f0cf7c3c56925efc51c1c9569c29515a48 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 13:40:43 +0100 Subject: [PATCH 09/13] chore: bump pybammsolvers to 0.11.0 The solver statistics, self-scheduled sweep and per-set failure messages need a pybammsolvers newer than the released 0.10.0, so pin pybamm to 0.11. --- packages/pybamm/pyproject.toml | 2 +- packages/pybammsolvers/pyproject.toml | 2 +- packages/pybammsolvers/src/pybammsolvers/version.py | 2 +- packages/pybammsolvers/vcpkg.json | 2 +- uv.lock | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index bc68dbf5e1..c00fca568e 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -27,7 +27,7 @@ classifiers = [ "Topic :: Scientific/Engineering", ] dependencies = [ - "pybammsolvers>=0.10.0,<0.11.0", + "pybammsolvers>=0.11.0,<0.12.0", "black", "numpy>=2.0.0", "scipy>=1.11.4,!=1.18.0", # 1.18.0 breaks pickling/deepcopy of interpolators, fixed in 1.18.1 (scipy#25489) diff --git a/packages/pybammsolvers/pyproject.toml b/packages/pybammsolvers/pyproject.toml index 5770392621..4faae285b2 100644 --- a/packages/pybammsolvers/pyproject.toml +++ b/packages/pybammsolvers/pyproject.toml @@ -18,7 +18,7 @@ requires-python = ">=3.10,<3.15" license-files = ["LICENSE"] readme = "README.md" # version.py stays the source of truth; a pre-commit hook mirrors it here. -version = "0.10.0" +version = "0.11.0" dependencies = [ "casadi==3.8.1", "numpy>=2.0.0", diff --git a/packages/pybammsolvers/src/pybammsolvers/version.py b/packages/pybammsolvers/src/pybammsolvers/version.py index 0883771731..3716aedb4c 100644 --- a/packages/pybammsolvers/src/pybammsolvers/version.py +++ b/packages/pybammsolvers/src/pybammsolvers/version.py @@ -1,3 +1,3 @@ # Single source of truth for the pybammsolvers version. Bump this at release; # a pre-commit hook mirrors it into pyproject.toml (see .pre-commit-config.yaml). -__version__ = "0.10.0" +__version__ = "0.11.0" diff --git a/packages/pybammsolvers/vcpkg.json b/packages/pybammsolvers/vcpkg.json index 9c096e37dc..58f55c53ea 100644 --- a/packages/pybammsolvers/vcpkg.json +++ b/packages/pybammsolvers/vcpkg.json @@ -1,6 +1,6 @@ { "name": "pybammsolvers", - "version-string": "0.10.0", + "version-string": "0.11.0", "dependencies": [ "casadi", { diff --git a/uv.lock b/uv.lock index eb5fbd75e1..8fbd1ff6d1 100644 --- a/uv.lock +++ b/uv.lock @@ -3364,7 +3364,7 @@ provides-extras = ["zoo-all"] [[package]] name = "pybammsolvers" -version = "0.10.0" +version = "0.11.0" source = { editable = "packages/pybammsolvers" } dependencies = [ { name = "casadi" }, From 8056a6074529438a2337d65e9d405ee6e8132689 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 16:33:08 +0100 Subject: [PATCH 10/13] fix: keep the mass matrix sparse when every block is dense A 0D domain's mass matrix block is dense, so a model whose differential variables are all scalars or 0D handed block_diag only dense blocks, which SciPy 1.18 deprecates in favour of returning a sparse array. Convert each block to csr_matrix first. --- CHANGELOG.md | 1 + .../pybamm/discretisations/discretisation.py | 5 +++- .../test_discretisation.py | 23 ++++++++++++++++++- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d6327f8a5..956fa0573b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ - A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans, even one among numbers such as `[True, 1e-6]`, and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- Discretising a model whose differential variables are all scalars or on zero-dimensional domains (such as a lumped current collector) no longer raises SciPy 1.18's `block_diag` `DeprecationWarning`, and its mass matrix stays a sparse matrix rather than becoming a sparse array once SciPy changes `block_diag`'s return type. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) ## Optimizations diff --git a/packages/pybamm/src/pybamm/discretisations/discretisation.py b/packages/pybamm/src/pybamm/discretisations/discretisation.py index aa8ebb502e..d74874250a 100644 --- a/packages/pybamm/src/pybamm/discretisations/discretisation.py +++ b/packages/pybamm/src/pybamm/discretisations/discretisation.py @@ -769,7 +769,10 @@ def create_mass_matrix(self, model): N_alg = len(model.algebraic) if N_rhs > 0 or N_alg > 0: - mass_matrix = pybamm.Matrix(block_diag(mass_list, format="csr")) + # If every block were dense, block_diag would return a sparse array, + # not a matrix, and a 0D domain's mass matrix is dense + blocks = [csr_matrix(mass) for mass in mass_list] + mass_matrix = pybamm.Matrix(block_diag(blocks, format="csr")) else: mass_matrix = None diff --git a/packages/pybamm/tests/unit/test_discretisations/test_discretisation.py b/packages/pybamm/tests/unit/test_discretisations/test_discretisation.py index a010fb6d22..6246b27c8f 100644 --- a/packages/pybamm/tests/unit/test_discretisations/test_discretisation.py +++ b/packages/pybamm/tests/unit/test_discretisations/test_discretisation.py @@ -3,10 +3,11 @@ # import random +import warnings import numpy as np import pytest -from scipy.sparse import block_diag, issparse +from scipy.sparse import block_diag, csr_matrix, issparse import pybamm from tests import ( @@ -1156,6 +1157,26 @@ def test_mass_matrix_with_finite_element(self): assert issparse(model.mass_matrix.entries) assert not model.is_standard_form_dae + @pytest.mark.parametrize("domain", [[], "current collector"]) + def test_mass_matrix_of_dense_blocks_is_a_sparse_matrix(self, domain): + u = pybamm.Variable("u", domain=domain) + v = pybamm.Variable("v", domain=domain) + model = pybamm.BaseModel() + model.rhs = {u: -u, v: -v} + model.initial_conditions = {u: 1, v: 2} + model.variables = {"u": u, "v": v} + disc = get_discretisation_for_testing( + cc_method=pybamm.ZeroDimensionalSpatialMethod + ) + + with warnings.catch_warnings(): + # scipy 1.18 deprecates block_diag of only dense blocks + warnings.simplefilter("error", DeprecationWarning) + disc.process_model(model) + + assert isinstance(model.mass_matrix.entries, csr_matrix) + np.testing.assert_array_equal(model.mass_matrix.entries.toarray(), np.eye(2)) + def test_process_input_variable(self): disc = get_discretisation_for_testing() From ef11defe44a60a655b7f6f169485469a1cd5580a Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 16:33:38 +0100 Subject: [PATCH 11/13] ci: install the CI matrix from uv.lock The prebuilt-wheel path resolved pybamm's dependencies afresh with uv pip install --no-sources, so CI tested the latest releases rather than the locked set local nox sessions use. Sync the lock as usual, skipping pybammsolvers, then install this interpreter's wheel with --no-deps. --- CHANGELOG.md | 1 + noxfile.py | 100 ++++++++++++++++++++++----------------------------- 2 files changed, 44 insertions(+), 57 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 956fa0573b..74fe907b1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ - A per-state `atol` given to `IDAKLUSolver` (or as `model.atol`) can now be a list, tuple or array, so it survives a `to_config`/`from_config` round trip, which returns a JSON list. One without exactly one value per state is rejected with a `SolverError` instead of reaching the integrator, as are booleans, even one among numbers such as `[True, 1e-6]`, and non-real values. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - Discretising a model whose differential variables are all scalars or on zero-dimensional domains (such as a lumped current collector) no longer raises SciPy 1.18's `block_diag` `DeprecationWarning`, and its mass matrix stays a sparse matrix rather than becoming a sparse array once SciPy changes `block_diag`'s return type. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) +- The CI test matrix now installs every dependency from `uv.lock`, as local nox sessions do, instead of resolving the latest releases, so only the prebuilt `pybammsolvers` wheel is installed outside the lock. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) ## Optimizations diff --git a/noxfile.py b/noxfile.py index e194448936..ef2e10892c 100644 --- a/noxfile.py +++ b/noxfile.py @@ -46,70 +46,58 @@ def is_macos_intel(): return sys.platform == "darwin" and platform.machine() in ("x86_64", "i386") -def install_locked(session, *, extras=None, groups=None, zoo=False, zoo_extras=None): - """Install pybamm and its dependencies into the session environment. - - Two modes, selected by the ``PYBAMM_SOLVER_WHEELS`` environment variable: - - * **Unset:** ``uv sync --frozen``, so ``pybammsolvers`` builds from source. - * **Set to a wheel directory (the CI matrix):** install this interpreter's - wheel by path, then pybamm with ``--no-sources``. The in-repo solver - version collides with the PyPI release, so nothing weaker picks the right - artifact, and Windows has no from-source build. - - ``zoo=True`` also installs the zoo with ``zoo_extras``, which only the - prebuilt-wheel path needs — a workspace sync installs every member anyway. +def install_locked(session, *, extras=None, groups=None, zoo_extras=None): + """Install every workspace member and its ``uv.lock`` dependencies. + + ``uv sync --frozen`` installs the locked environment, building + ``pybammsolvers`` from source. When ``PYBAMM_SOLVER_WHEELS`` names a wheel + or a directory of wheels (the CI matrix), the sync skips ``pybammsolvers`` + and this interpreter's prebuilt wheel is installed in its place, since + Windows has no from-source build. """ env = {"UV_PROJECT_ENVIRONMENT": session.virtualenv.location} - wheels = os.getenv("PYBAMM_SOLVER_WHEELS") - if wheels: - wheels_path = Path(wheels) - if wheels_path.is_dir(): - # The per-OS artifact holds wheels for every Python version; pick the - # one whose interpreter tag matches this cell. The trailing dash in - # "*-cpXY-*" avoids matching free-threaded "cpXYt-" builds. - tag = f"cp{sys.version_info.major}{sys.version_info.minor}" - matches = sorted(wheels_path.glob(f"*-{tag}-*.whl")) - if not matches: - session.error( - f"PYBAMM_SOLVER_WHEELS={wheels} contains no pybammsolvers " - f"wheel for {tag}" - ) - wheel = matches[0].as_posix() - else: - wheel = wheels_path.as_posix() - python_bin = os.path.join( - session.bin, "python.exe" if sys.platform == "win32" else "python" - ) - extras_str = f"[{','.join(extras)}]" if extras else "" - zoo_extras_str = f"[{','.join(zoo_extras)}]" if zoo_extras else "" - cmd = [ - "uv", - "pip", - "install", - "--python", - python_bin, - "--no-sources", - wheel, - "-e", - f"./packages/pybamm{extras_str}", - ] - if zoo: - cmd.extend(["-e", f"./packages/pybamm-model-zoo{zoo_extras_str}"]) - for group in groups or []: - # Groups (dev, docs) are defined in the pybamm package, not the root. - cmd.extend(["--group", f"packages/pybamm/pyproject.toml:{group}"]) - session.run(*cmd, env=env, external=True) - return cmd = ["uv", "sync", "--frozen"] for extra in [*(extras or []), *(zoo_extras or [])]: cmd.extend(["--extra", extra]) for group in groups or []: cmd.extend(["--group", group]) + if wheels: + cmd.extend(["--no-install-package", "pybammsolvers"]) session.run(*cmd, env=env, external=True) + if not wheels: + return + wheels_path = Path(wheels) + if wheels_path.is_dir(): + # The trailing dash keeps free-threaded "cpXYt-" wheels out of the match + tag = f"cp{sys.version_info.major}{sys.version_info.minor}" + matches = sorted(wheels_path.glob(f"*-{tag}-*.whl")) + if not matches: + session.error( + f"PYBAMM_SOLVER_WHEELS={wheels} contains no pybammsolvers " + f"wheel for {tag}" + ) + wheel = matches[0].as_posix() + else: + wheel = wheels_path.as_posix() + python_bin = os.path.join( + session.bin, "python.exe" if sys.platform == "win32" else "python" + ) + # The sync already installed the solver's locked dependencies + session.run( + "uv", + "pip", + "install", + "--python", + python_bin, + "--no-deps", + wheel, + env=env, + external=True, + ) + @nox.session(name="coverage", default=False) def run_coverage(session): @@ -333,9 +321,7 @@ def build_docs(session): def install_zoo(session): """Install pybamm plus the zoo and every model's declared dependencies.""" set_environment_variables(PYBAMM_ENV, session=session) - install_locked( - session, extras=["all"], groups=["dev"], zoo=True, zoo_extras=["zoo-all"] - ) + install_locked(session, extras=["all"], groups=["dev"], zoo_extras=["zoo-all"]) def zoo_pytest(session, marker, *args, allow_empty=False): @@ -400,7 +386,7 @@ def run_zoo_new(session): """Create a new model zoo entry from the template.""" # Only the zoo itself and pybamm's version metadata are needed, so this skips # the extras and the dev group that install_zoo pulls in. - install_locked(session, zoo=True) + install_locked(session) session.run("python", f"{ZOO_TESTS}/scripts/new_model.py", *session.posargs) From 06e7136244c7d202edd866fb3eb99bddd1a2c050 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 16:34:03 +0100 Subject: [PATCH 12/13] test: match only the partial-solution warnings per input set pytest.warns records every warning raised in its block, not just the matching ones, so an unrelated warning made re.match return None. Build the model outside the block and ignore records that are not a partial-solution warning. --- .../tests/unit/test_solvers/test_idaklu_solver.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 763e5fe617..c5dc20dafb 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -216,14 +216,14 @@ def test_every_input_set_failing_part_way_is_named(self): assert re.findall(r"input set (\d+): IDA_", str(error.value)) == ["1", "3"] def test_each_partial_solution_warning_names_its_input_set(self): + model = _part_way_failure_model() solver = pybamm.IDAKLUSolver(options={"num_threads": 4}, on_failure="warn") inputs_list = [{"k": k} for k in (2.0, 0.5, 2.0, 0.3)] with pytest.warns(UserWarning, match="returning a partial solution") as record: - solutions = solver.solve( - _part_way_failure_model(), [0, 1], inputs=inputs_list - ) - named = [re.match(r"input set (\d+): ", str(w.message)) for w in record] - assert [match.group(1) for match in named] == ["1", "3"] + solutions = solver.solve(model, [0, 1], inputs=inputs_list) + pattern = r"input set (\d+): .*, returning a partial solution\.$" + named = [re.match(pattern, str(w.message)) for w in record] + assert [match.group(1) for match in named if match] == ["1", "3"] assert [solution.termination for solution in solutions] == [ "final time", "failure", From a5d4fd3e11691868ef12011c7d64da7d1adc70b9 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 30 Sep 2026 18:44:28 +0100 Subject: [PATCH 13/13] perf: reduce solver statistics storage for retained states --- packages/pybamm/src/pybamm/solvers/solution.py | 10 +++++++--- .../pybamm/tests/unit/test_solvers/test_solution.py | 9 +++++++++ 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index ee16ab30d2..6440e76cd2 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -33,7 +33,7 @@ def default(self, obj): return json.JSONEncoder.default(self, obj) # pragma: no cover -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class SolverStatistics: """Counts of the work an integrator did to produce a solution. @@ -65,6 +65,10 @@ def __add__(self, other: SolverStatistics) -> SolverStatistics: ) +# Shared by every first and last state, which do no solver work +_NO_SOLVER_STATISTICS = SolverStatistics() + + def _sum_solver_statistics(solutions) -> SolverStatistics | None: """Total the solver statistics of ``solutions``. @@ -675,7 +679,7 @@ def first_state(self): new_sol.solve_time = 0 new_sol.integration_time = 0 - new_sol.solver_statistics = SolverStatistics() + new_sol.solver_statistics = _NO_SOLVER_STATISTICS new_sol.set_up_time = 0 return new_sol @@ -719,7 +723,7 @@ def last_state(self): new_sol._sub_solutions = self.sub_solutions[-1:] new_sol.solve_time = 0 new_sol.integration_time = 0 - new_sol.solver_statistics = SolverStatistics() + new_sol.solver_statistics = _NO_SOLVER_STATISTICS new_sol.set_up_time = 0 return new_sol diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index b2fa8d342f..e7e5f79cdd 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -1928,6 +1928,15 @@ def test_first_and_last_state_report_no_work(self): assert solution.first_state.solver_statistics == pybamm.SolverStatistics() assert solution.last_state.solver_statistics == pybamm.SolverStatistics() + def test_statistics_stay_small_across_retained_states(self): + # An experiment keeps a first and last state for every step + solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5)) + assert ( + solution.first_state.solver_statistics + is solution.last_state.solver_statistics + ) + assert not hasattr(solution.solver_statistics, "__dict__") + def test_statistics_survive_pickling(self): statistics = pybamm.SolverStatistics(1, 2, 3, 4, 5) solution = self._solution(0, statistics)