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). 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..b4ed8f184 100644 --- a/pybop/pybamm/utils.py +++ b/pybop/pybamm/utils.py @@ -1,7 +1,9 @@ import multiprocessing as mp import platform +from concurrent.futures import TimeoutError as CFTimeoutError import pybamm +from pebble import ProcessPool class RecommendedSolver(pybamm.IDAKLUSolver): @@ -22,6 +24,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, CFTimeoutError) 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..b2dbbf751 100644 --- a/tests/unit/test_simulator.py +++ b/tests/unit/test_simulator.py @@ -78,3 +78,29 @@ 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() + + # 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"): + simulator.solve()