Revert lru_cache to functional form; Add lru-dict for integrators

Author: jsbrittain

CHANGELOG.md

@@ -8,13 +8,15 @@
 - Added "Open-circuit voltage [V]", which is the open-circuit voltage as calculated from the bulk particle concentrations. The old variable "Measured open circuit voltage [V]", which referred to the open-circuit potential as calculated from the surface particle concentrations, has been renamed to "Surface open-circuit voltage [V]". ([#2740](https://github.com/pybamm-team/PyBaMM/pull/2740))
 - Added an example for `plot_voltage_components`, explaining what the different voltage components are. ([#2740](https://github.com/pybamm-team/PyBaMM/pull/2740))
 - Added method to calculate maximum theoretical energy. ([#2777](https://github.com/pybamm-team/PyBaMM/pull/2777)) and add to summary variables ([#2781](https://github.com/pybamm-team/PyBaMM/pull/2781))
+- Added option to limit the number of integrators stored in CasadiSolver, which is particularly relevant when running simulations back-to-back [#2823](https://github.com/pybamm-team/PyBaMM/pull/2823)
 
 ## Bug fixes
 
 - Fixed a bug where variable bounds could not contain `InputParameters` ([#2795](https://github.com/pybamm-team/PyBaMM/pull/2795))
 - Improved `model.latexify()` to have a cleaner and more readable output ([#2764](https://github.com/pybamm-team/PyBaMM/pull/2764))
 - Fixed electrolyte conservation in the case of concentration-dependent transference number ([#2758](https://github.com/pybamm-team/PyBaMM/pull/2758))
 - Fixed `plot_voltage_components` so that the sum of overpotentials is now equal to the voltage ([#2740](https://github.com/pybamm-team/PyBaMM/pull/2740))
+- Fixed excessive RAM consumption when running multiple simulations ([#2823](https://github.com/pybamm-team/PyBaMM/pull/2823))
 
 ## Optimizations
 

pybamm/models/full_battery_models/lithium_ion/electrode_soh.py

@@ -205,14 +205,18 @@ def __init__(
 
         self.lims_ocp = (x0_min, x100_max, y100_min, y0_max)
         self.OCV_function = None
+        self._get_electrode_soh_sims_full = lru_cache()(
+            self.__get_electrode_soh_sims_full
+        )
+        self._get_electrode_soh_sims_split = lru_cache()(
+            self.__get_electrode_soh_sims_split
+        )
 
-    @lru_cache
-    def _get_electrode_soh_sims_full(self):
+    def __get_electrode_soh_sims_full(self):
         full_model = _ElectrodeSOH(param=self.param, known_value=self.known_value)
         return pybamm.Simulation(full_model, parameter_values=self.parameter_values)
 
-    @lru_cache
-    def _get_electrode_soh_sims_split(self):
+    def __get_electrode_soh_sims_split(self):
         x100_model = _ElectrodeSOH(
             param=self.param, solve_for=["x_100"], known_value=self.known_value
         )

pybamm/simulation.py

@@ -128,6 +128,23 @@ def __init__(
 
             warnings.filterwarnings("ignore")
 
+        self.get_esoh_solver = lru_cache()(self._get_esoh_solver)
+
+    def __getstate__(self):
+        """
+        Return dictionary of picklable items
+        """
+        result = self.__dict__.copy()
+        result["get_esoh_solver"] = None  # Exclude LRU cache
+        return result
+
+    def __setstate__(self, state):
+        """
+        Unpickle, restoring unpicklable relationships
+        """
+        self.__dict__ = state
+        self.get_esoh_solver = lru_cache()(self._get_esoh_solver)
+
     def set_up_and_parameterise_experiment(self):
         """
         Set up a simulation to run with an experiment. This creates a dictionary of
@@ -900,8 +917,7 @@ def step(
 
         return self.solution
 
-    @lru_cache
-    def get_esoh_solver(self, calc_esoh):
+    def _get_esoh_solver(self, calc_esoh):
         if (
             calc_esoh is False
             or isinstance(self.model, pybamm.lead_acid.BaseModel)

pybamm/solvers/casadi_solver.py

@@ -6,6 +6,7 @@
 import numpy as np
 import warnings
 from scipy.interpolate import interp1d
+from .lrudict import LRUDict
 
 
 class CasadiSolver(pybamm.BaseSolver):
@@ -67,6 +68,10 @@ class CasadiSolver(pybamm.BaseSolver):
         Whether to perturb algebraic initial conditions to avoid a singularity. This
         can sometimes slow down the solver, but is kept True as default for "safe" mode
         as it seems to be more robust (False by default for other modes).
+    integrators_maxcount : int, optional
+        The maximum number of integrators that the solver will retain before
+        ejecting past integrators using an LRU methodology. A value of 0 or
+        None (default) leaves the number of integrators unbound.
     """
 
     def __init__(
@@ -83,6 +88,7 @@ def __init__(
         extra_options_call=None,
         return_solution_if_failed_early=False,
         perturb_algebraic_initial_conditions=None,
+        integrators_maxcount=None,
     ):
         super().__init__(
             "problem dependent",
@@ -123,8 +129,9 @@ def __init__(
         self.name = "CasADi solver with '{}' mode".format(mode)
 
         # Initialize
-        self.integrators = {}
-        self.integrator_specs = {}
+        self.integrators_maxcount = integrators_maxcount
+        self.integrators = LRUDict(maxsize=self.integrators_maxcount)
+        self.integrator_specs = LRUDict(maxsize=self.integrators_maxcount)
         self.y_sols = {}
 
         pybamm.citations.register("Andersson2019")

pybamm/solvers/lrudict.py

@@ -0,0 +1,32 @@
+from collections import OrderedDict
+
+
+class LRUDict(OrderedDict):
+    """LRU extension of a dictionary"""
+
+    def __init__(self, maxsize=None):
+        """maxsize limits the item count based on an LRU strategy
+
+        The dictionary remains unbound when maxsize = 0 | None
+        """
+        super().__init__()
+        self.maxsize = maxsize
+
+    def __setitem__(self, key, value):
+        super().__setitem__(key, value)
+        while self.maxsize and self.__len__() > self.maxsize:
+            self.popitem(last=False)
+
+    def __getitem__(self, key):
+        try:
+            self.move_to_end(key, last=True)
+        except KeyError:
+            pass
+        return super().__getitem__(key)
+
+    def get(self, key):
+        try:
+            self.move_to_end(key, last=True)
+        except KeyError:
+            pass
+        return super().get(key)

tests/unit/test_solvers/test_lrudict.py

@@ -0,0 +1,60 @@
+#
+# Tests for the LRUDict class
+#
+import unittest
+from pybamm.solvers.lrudict import LRUDict
+from collections import OrderedDict
+
+
+class TestLRUDict(unittest.TestCase):
+    def test_lrudict_defaultbehaviour(self):
+        """Default behaviour [no LRU] mimics Dict"""
+        d = LRUDict()
+        dd = dict()
+        for count in range(1, 100):
+            d[count] = f"v{count}"
+            dd[count] = f"v{count}"
+            if count % 5 == 0:
+                # LRU will reorder list, but not remove entries
+                d.get(count - 2)
+                dd.get(count - 2)
+            assert set(d.keys()) == set(dd.keys())
+            assert set(d.values()) == set(dd.values())
+
+    def test_lrudict_noitems(self):
+        """Edge case: no items in LRU, raises KeyError on assignment"""
+        d = LRUDict(maxsize=-1)
+        with self.assertRaises(KeyError):
+            d["a"] = 1
+
+    def test_lrudict_singleitem(self):
+        """Only the last added element should be present"""
+        d = LRUDict(maxsize=1)
+        item_list = range(1, 100)
+        assert len(d) == 0
+        for item in item_list:
+            d[item] = item
+            assert len(d) == 1
+            assert d[item]
+        d.popitem()
+        assert len(d) == 0
+
+    def test_lrudict_multiitem(self):
+        """Check that the correctly ordered items are always present"""
+        for maxsize in range(1, 101, 5):
+            d = LRUDict(maxsize=maxsize)
+            expected = OrderedDict()
+            for key in range(1, 100):
+                value = f"v{key}"
+                d[key] = value
+                expected[key] = value
+                if key % 5 == 0 and maxsize > 3:
+                    # Push item to front of the list
+                    key_to_read = key - 3
+                    d.get(key_to_read)
+                    expected.move_to_end(key_to_read)
+                expected = OrderedDict(
+                    (k, expected[k]) for k in list(expected.keys())[-maxsize:]
+                )
+                assert list(d.keys()) == list(expected.keys())
+                assert list(d.values()) == list(expected.values())