diff --git a/CHANGELOG.md b/CHANGELOG.md index 912966ce2..c4c19793b 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. - [#965](https://github.com/pybop-team/PyBOP/pull/965) - Adds synthetic data and example scripts for OCV parameterisation. - [#963](https://github.com/pybop-team/PyBOP/pull/963) - Adds an example for generating synthetic data from a specification and exporting it to a PyProBE-compatible parquet file. - [#928](https://github.com/pybop-team/PyBOP/pull/928) - Adds the option to choose `matplotlib` as the plotting library for plotting functions. Additionally, figures and axes can be created manually and passed as keyword arguments to plotting functions. An example notebook `plotting.ipynb` was added to the `getting_started` directory to demonstrate usage of the new features. diff --git a/pybop/pybamm/__init__.py b/pybop/pybamm/__init__.py index a1cfcce9a..f4dcb51ee 100644 --- a/pybop/pybamm/__init__.py +++ b/pybop/pybamm/__init__.py @@ -2,5 +2,5 @@ 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 from .synthetic_utils import archive_data, convert_to_half_cell_parameters, simulate_procedure diff --git a/pybop/pybamm/utils.py b/pybop/pybamm/utils.py index 6a144007d..41bd5e4e8 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 9e7540668..9d10fbeb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "numpy>=1.26", "scipy>=1.12", "pints>=0.6.0", + "pebble", "polars>=1.18.0", ] 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()