diff --git a/CHANGELOG.md b/CHANGELOG.md index 7533ecf0d3..74fe907b1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,18 @@ - 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, 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 + +- 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 - 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/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) 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/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/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/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/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 2862367a36..735306a91a 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,69 @@ 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 numpy array or float" + "Absolute tolerances must be a float, or a list, tuple or array " + f"of floats with one value per state, not {type(atol).__name__}" ) + if isinstance(atol, numbers.Real): + return np.full(number_of_states, float(atol)) + + try: + values = np.asarray(atol) + except ValueError as error: + raise pybamm.SolverError( + "Absolute tolerances must be a flat list, tuple or array" + ) from error + # NumPy upcasts a bool beside numbers to 1.0 or 0.0, losing it + if isinstance(atol, list | tuple) and any( + isinstance(value, bool | np.bool_) + for value in np.asarray(atol, dtype=object).ravel() + ): + raise pybamm.SolverError( + "Absolute tolerances must be real numbers, not booleans" + ) + # Bool and complex arrays would otherwise cast to float silently + if values.dtype.kind not in "iuf": + raise pybamm.SolverError( + "Absolute tolerances must be real numbers, not an array of " + f"dtype {values.dtype}" + ) + if values.ndim == 0: + return np.full(number_of_states, float(values)) + if values.shape not in ((number_of_states,), (number_of_states, 1)): + raise pybamm.SolverError( + f"Absolute tolerances have shape {values.shape} but " + f"({number_of_states},) was expected (one value per state)" + ) + return np.ascontiguousarray(values.ravel(), dtype=np.float64) def set_up(self, model, inputs=None, t_eval=None, ics_only=False): if model.convert_to_format != "casadi": @@ -737,6 +793,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 @@ -802,17 +871,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 @@ -861,6 +921,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 +1274,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..6440e76cd2 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,61 @@ def default(self, obj): return json.JSONEncoder.default(self, obj) # pragma: no cover +@dataclass(frozen=True, slots=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)) + ) + + +# 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``. + + 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 +297,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 +679,7 @@ def first_state(self): new_sol.solve_time = 0 new_sol.integration_time = 0 + new_sol.solver_statistics = _NO_SOLVER_STATISTICS new_sol.set_up_time = 0 return new_sol @@ -663,6 +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 = _NO_SOLVER_STATISTICS new_sol.set_up_time = 0 return new_sol @@ -1110,6 +1171,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 @@ -1177,6 +1239,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 @@ -1215,12 +1278,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: @@ -1291,6 +1349,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)) + # 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): @@ -1329,6 +1389,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/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/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_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() 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..c5dc20dafb 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]]: @@ -89,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 @@ -170,6 +195,86 @@ 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_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): + 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(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", + "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. + 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 +807,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 +842,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 +855,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 +866,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 +883,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 +892,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 +916,50 @@ 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. 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 + ) + assert started < n_inputs + def test_setup_options(self): model = pybamm.BaseModel() u = pybamm.Variable("u") @@ -2139,3 +2288,271 @@ 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"), + ([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"), + ([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_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}) + 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..e7e5f79cdd 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 # nosec B403 - used in tests with trusted input import subprocess # nosec B404 - used in tests with trusted input import sys from unittest import mock @@ -1839,3 +1840,110 @@ 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_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() + 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) + 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 # nosec B301 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/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..a802a1e3fa 100644 --- a/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp +++ b/packages/pybammsolvers/src/pybammsolvers/idaklu_source/IDAKLUSolverGroup.cpp @@ -1,5 +1,8 @@ #include "IDAKLUSolverGroup.hpp" +#include "sundials_error_handler.hpp" #include +#include +#include #include #include @@ -97,9 +100,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 +122,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 +131,39 @@ 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()); + // cppcheck-suppress unreadVariable + const int team_size = std::max(1, std::min(m_solvers.size(), number_of_groups)); + // 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) + { + // 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,30 @@ std::vector IDAKLUSolverGroup::solve( std::rethrow_exception(python_exception); } - if (exception_message.has_value()) { - py::set_error(PyExc_ValueError, exception_message->c_str()); + 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; + } + py::set_error(PyExc_ValueError, failures.c_str()); throw py::error_already_set(); } - // 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); - } - - 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..9998f883e4 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, + 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) + 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..e1b2808f98 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, + const 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; @@ -68,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). @@ -81,7 +88,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 +105,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..4b7b1bcf22 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_; // cppcheck-suppress unusedStructMember + std::unique_ptr jac_; // cppcheck-suppress unusedStructMember int n_vars_; bool use_sparse_; 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/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. 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" },