From a835d03d5bf5c8b8af5cbbfe4511ab8484cbd410 Mon Sep 17 00:00:00 2001 From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com> Date: Fri, 5 Jun 2026 09:54:39 +0100 Subject: [PATCH 1/3] Add SafeSolver that uses pebble --- pybop/pybamm/__init__.py | 2 +- pybop/pybamm/utils.py | 83 ++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + tests/unit/test_simulator.py | 18 ++++++++ 4 files changed, 103 insertions(+), 1 deletion(-) diff --git a/pybop/pybamm/__init__.py b/pybop/pybamm/__init__.py index 3f541b31c..b13bed7de 100644 --- a/pybop/pybamm/__init__.py +++ b/pybop/pybamm/__init__.py @@ -2,4 +2,4 @@ from .eis_simulator import EISSimulator from .parameter_utils import set_formation_concentrations, cell_mass, cell_volume from .design_variables import add_variable_to_model -from .utils import RecommendedSolver, SymbolReplacer +from .utils import RecommendedSolver, SafeSolver, SymbolReplacer diff --git a/pybop/pybamm/utils.py b/pybop/pybamm/utils.py index 96e7b23e5..c9996a80f 100644 --- a/pybop/pybamm/utils.py +++ b/pybop/pybamm/utils.py @@ -2,6 +2,7 @@ import platform import pybamm +from pebble import ProcessPool class RecommendedSolver(pybamm.IDAKLUSolver): @@ -22,6 +23,88 @@ def __init__(self, output_variables: list[str] | None = None): ) +class SafeSolver(pybamm.CasadiSolver): + """ + A version of PyBaMM's CasadiSolver with a timeout option. + + Additional parameters + --------------------- + timeout : float, optional + If timeout is a positive number, simulations are terminated after timeout + seconds if not completed successfully within this time. Default is None. + """ + + def __init__(self, timeout: float | None = None, **kwargs): + super().__init__(**kwargs) + self.timeout = timeout + + def _integrate( + self, + model: pybamm.BaseModel, + t_eval, + inputs_list: list[dict] | None = None, + t_interp=None, + nproc=1, + ): + """ + Solve a DAE model defined by residuals with initial conditions y0. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model whose solution to calculate. + t_eval : numeric type + The times at which to compute the solution + inputs_list : list of dict, optional + Any input parameters to pass to the model when solving + """ + + inputs_list = inputs_list or [{}] + + ninputs = len(inputs_list) + if ninputs == 1 and self.timeout is None: + new_solution = self._integrate_single( + model, + t_eval, + inputs_list[0], + model.y0_list[0], + ) + new_solutions = [new_solution] + else: + with ProcessPool( + context=mp.get_context(self._mp_context), + max_workers=nproc or mp.cpu_count(), + ) as p: + model_list = [model] * ninputs + t_eval_list = [t_eval] * ninputs + y0_list = model.y0_list + + futures = p.map( + self._integrate_single, + model_list, + t_eval_list, + inputs_list, + y0_list, + timeout=self.timeout, + ) + iterator = futures.result() + + new_solutions = [] + while True: + try: + new_solutions.append(next(iterator)) + except StopIteration: + break + except TimeoutError as e: + raise pybamm.SolverError( + f"Timeout after {e.args[1]:.1f} seconds." + ) from e + except Exception as e: + raise pybamm.SolverError(str(e)) from None + + return new_solutions + + class SymbolReplacer: """ Helper class to replace all instances of one or more symbols in an expression tree diff --git a/pyproject.toml b/pyproject.toml index acd729723..1ffc2be83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "numpy>=1.26", "scipy>=1.12", "pints>=0.6.0", + "pebble", ] [project.optional-dependencies] diff --git a/tests/unit/test_simulator.py b/tests/unit/test_simulator.py index 748617d17..888af2302 100644 --- a/tests/unit/test_simulator.py +++ b/tests/unit/test_simulator.py @@ -78,3 +78,21 @@ def test_set_output_variables(self): ValueError, match="Not a variable is not a variable in the model." ): simulator.set_output_variables(["Not a variable"]) + + def test_safe_solver(self): + model = pybamm.lithium_ion.DFN() + + experiment = pybamm.Experiment(["Discharge at C/10 for 10 seconds"]) + simulator = pybop.pybamm.Simulator( + model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=10) + ) + solution = simulator.solve() + assert isinstance(solution, pybamm.Solution) + + experiment = pybamm.Experiment(["Discharge at C/10 for 10 hours"]) + simulator = pybop.pybamm.Simulator( + model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=0.1) + ) + simulator.debug_mode = True + with pytest.raises(pybamm.SolverError, match="Timeout after 0.1 seconds."): + simulator.solve() From 5d7eb23e92ea0798295d41df932c60e8c0565be5 Mon Sep 17 00:00:00 2001 From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:30:19 +0100 Subject: [PATCH 2/3] Add test on short simulation --- pybop/pybamm/utils.py | 3 ++- tests/unit/test_simulator.py | 10 +++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/pybop/pybamm/utils.py b/pybop/pybamm/utils.py index c9996a80f..b4ed8f184 100644 --- a/pybop/pybamm/utils.py +++ b/pybop/pybamm/utils.py @@ -1,5 +1,6 @@ import multiprocessing as mp import platform +from concurrent.futures import TimeoutError as CFTimeoutError import pybamm from pebble import ProcessPool @@ -95,7 +96,7 @@ def _integrate( new_solutions.append(next(iterator)) except StopIteration: break - except TimeoutError as e: + except (TimeoutError, CFTimeoutError) as e: raise pybamm.SolverError( f"Timeout after {e.args[1]:.1f} seconds." ) from e diff --git a/tests/unit/test_simulator.py b/tests/unit/test_simulator.py index 888af2302..b2dbbf751 100644 --- a/tests/unit/test_simulator.py +++ b/tests/unit/test_simulator.py @@ -82,17 +82,25 @@ def test_set_output_variables(self): def test_safe_solver(self): model = pybamm.lithium_ion.DFN() + # Test SafeSolver on short simulations experiment = pybamm.Experiment(["Discharge at C/10 for 10 seconds"]) + simulator = pybop.pybamm.Simulator( + model, protocol=experiment, solver=pybop.pybamm.SafeSolver() + ) + solution = simulator.solve() + assert isinstance(solution, pybamm.Solution) + simulator = pybop.pybamm.Simulator( model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=10) ) solution = simulator.solve() assert isinstance(solution, pybamm.Solution) + # Test SafeSolver timeout on long simulation experiment = pybamm.Experiment(["Discharge at C/10 for 10 hours"]) simulator = pybop.pybamm.Simulator( model, protocol=experiment, solver=pybop.pybamm.SafeSolver(timeout=0.1) ) simulator.debug_mode = True - with pytest.raises(pybamm.SolverError, match="Timeout after 0.1 seconds."): + with pytest.raises(pybamm.SolverError, match="Timeout"): simulator.solve() From 2d0acbd58fe1d9e66c3eca32a9bc2ac81e3c18bb Mon Sep 17 00:00:00 2001 From: NicolaCourtier <45851982+NicolaCourtier@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:31:05 +0100 Subject: [PATCH 3/3] Update CHANGELOG.md --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7400fe662..07ce776ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Features +- [#948](https://github.com/pybop-team/PyBOP/pull/948) - Adds `SafeSolver`, which uses pebble multiprocessing with a timeout option. - [#918](https://github.com/pybop-team/PyBOP/pull/918) - Adds a plot for predictions sampled from a posterior distribution (`pybop.plot.predictive`). - [#940](https://github.com/pybop-team/PyBOP/pull/940) - Adds support for Python 3.14 (EP-BOLFI optimiser and PyProBE still restricted to Python 3.12 or below).