Merge pull request #2465 from agriyakhetarpal/memoization

Use `@cache` and `@cached_property` for memoization

Author: valentinsulzer

CHANGELOG.md

@@ -8,6 +8,7 @@
 
 ## Optimizations
 
+- Implemented memoization via `cache` and `cached_property` from functools ([#2465](https://github.com/pybamm-team/PyBaMM/pull/2465))
 - `ParameterValues` now avoids trying to process children if a function parameter is an object that doesn't depend on its children ([#2477](https://github.com/pybamm-team/PyBaMM/pull/2477))
 - Added more rules for simplifying expressions, especially around Concatenations. Also, meshes constructed from multiple domains are now cached ([#2443](https://github.com/pybamm-team/PyBaMM/pull/2443))
 - Added more rules for simplifying expressions. Constants in binary operators are now moved to the left by default (e.g. `x*2` returns `2*x`) ([#2424](https://github.com/pybamm-team/PyBaMM/pull/2424))

pybamm/expression_tree/symbol.py

@@ -8,6 +8,7 @@
 import sympy
 from anytree.exporter import DotExporter
 from scipy.sparse import csr_matrix, issparse
+from functools import lru_cache, cached_property
 
 import pybamm
 from pybamm.expression_tree.printing.print_name import prettify_print_name
@@ -827,6 +828,7 @@ def evaluates_to_number(self):
     def evaluates_to_constant_number(self):
         return self.evaluates_to_number() and self.is_constant()
 
+    @lru_cache
     def evaluates_on_edges(self, dimension):
         """
         Returns True if a symbol evaluates on an edge, i.e. symbol contains a gradient
@@ -845,12 +847,9 @@ def evaluates_on_edges(self, dimension):
             Whether the symbol evaluates on edges (in the finite volume discretisation
             sense)
         """
-        try:
-            return self._saved_evaluates_on_edges[dimension]
-        except KeyError:
-            eval_on_edges = self._evaluates_on_edges(dimension)
-            self._saved_evaluates_on_edges[dimension] = eval_on_edges
-            return eval_on_edges
+        eval_on_edges = self._evaluates_on_edges(dimension)
+        self._saved_evaluates_on_edges[dimension] = eval_on_edges
+        return eval_on_edges
 
     def _evaluates_on_edges(self, dimension):
         # Default behaviour: return False
@@ -894,48 +893,39 @@ def new_copy(self):
         obj._print_name = self.print_name
         return obj
 
-    @property
+    @cached_property
     def size(self):
         """
         Size of an object, found by evaluating it with appropriate t and y
         """
-        try:
-            return self._saved_size
-        except AttributeError:
-            self._saved_size = np.prod(self.shape)
-            return self._saved_size
+        return np.prod(self.shape)
 
-    @property
+    @cached_property
     def shape(self):
         """
         Shape of an object, found by evaluating it with appropriate t and y.
         """
+        # Default behaviour is to try to evaluate the object directly
+        # Try with some large y, to avoid having to unpack (slow)
         try:
-            return self._saved_shape
-        except AttributeError:
-            # Default behaviour is to try to evaluate the object directly
-            # Try with some large y, to avoid having to unpack (slow)
-            try:
-                y = np.nan * np.ones((1000, 1))
-                evaluated_self = self.evaluate(0, y, y, inputs="shape test")
-            # If that fails, fall back to calculating how big y should really be
-            except ValueError:
-                unpacker = pybamm.SymbolUnpacker(pybamm.StateVector)
-                state_vectors_in_node = unpacker.unpack_symbol(self)
-                min_y_size = max(
-                    max(len(x._evaluation_array) for x in state_vectors_in_node), 1
-                )
-                # Pick a y that won't cause RuntimeWarnings
-                y = np.nan * np.ones((min_y_size, 1))
-                evaluated_self = self.evaluate(0, y, y, inputs="shape test")
-
-            # Return shape of evaluated object
-            if isinstance(evaluated_self, numbers.Number):
-                self._saved_shape = ()
-            else:
-                self._saved_shape = evaluated_self.shape
+            y = np.nan * np.ones((1000, 1))
+            evaluated_self = self.evaluate(0, y, y, inputs="shape test")
+        # If that fails, fall back to calculating how big y should really be
+        except ValueError:
+            unpacker = pybamm.SymbolUnpacker(pybamm.StateVector)
+            state_vectors_in_node = unpacker.unpack_symbol(self)
+            min_y_size = max(
+                max(len(x._evaluation_array) for x in state_vectors_in_node), 1
+            )
+            # Pick a y that won't cause RuntimeWarnings
+            y = np.nan * np.ones((min_y_size, 1))
+            evaluated_self = self.evaluate(0, y, y, inputs="shape test")
 
-        return self._saved_shape
+        # Return shape of evaluated object
+        if isinstance(evaluated_self, numbers.Number):
+            return ()
+        else:
+            return evaluated_self.shape
 
     @property
     def size_for_testing(self):

pybamm/models/full_battery_models/base_battery_model.py

@@ -4,6 +4,7 @@
 
 import pybamm
 import numbers
+from functools import cached_property
 
 
 class BatteryModelOptions(pybamm.FuzzyDict):
@@ -600,19 +601,14 @@ def phases(self):
                 self._phases[domain] = phases
             return self._phases
 
-    @property
+    @cached_property
     def whole_cell_domains(self):
-        try:
-            return self._whole_cell_domains
-        except AttributeError:
-            if self["working electrode"] == "positive":
-                wcd = ["separator", "positive electrode"]
-            elif self["working electrode"] == "negative":
-                wcd = ["negative electrode", "separator"]
-            elif self["working electrode"] == "both":
-                wcd = ["negative electrode", "separator", "positive electrode"]
-            self._whole_cell_domains = wcd
-            return wcd
+        if self["working electrode"] == "positive":
+            return ["separator", "positive electrode"]
+        elif self["working electrode"] == "negative":
+            return ["negative electrode", "separator"]
+        elif self["working electrode"] == "both":
+            return ["negative electrode", "separator", "positive electrode"]
 
     @property
     def electrode_types(self):

pybamm/models/full_battery_models/lithium_ion/electrode_soh.py

@@ -3,6 +3,7 @@
 #
 import pybamm
 import numpy as np
+from functools import lru_cache
 
 
 class ElectrodeSOH(pybamm.BaseModel):
@@ -246,28 +247,18 @@ def __init__(self, parameter_values, param=None):
             - self.param.n.prim.U_dimensional(x, T)
         )
 
+    @lru_cache
     def _get_electrode_soh_sims_full(self):
-        try:
-            return self._full_sim
-        except AttributeError:
-            full_model = ElectrodeSOH(param=self.param)
-            self._full_sim = pybamm.Simulation(
-                full_model, parameter_values=self.parameter_values
-            )
-            return self._full_sim
+        full_model = ElectrodeSOH(param=self.param)
+        return pybamm.Simulation(full_model, parameter_values=self.parameter_values)
 
+    @lru_cache
     def _get_electrode_soh_sims_split(self):
-        try:
-            return self._split_sims
-        except AttributeError:
-            x100_model = ElectrodeSOHx100(param=self.param)
-            x100_sim = pybamm.Simulation(
-                x100_model, parameter_values=self.parameter_values
-            )
-            x0_model = ElectrodeSOHx0(param=self.param)
-            x0_sim = pybamm.Simulation(x0_model, parameter_values=self.parameter_values)
-            self._split_sims = [x100_sim, x0_sim]
-            return self._split_sims
+        x100_model = ElectrodeSOHx100(param=self.param)
+        x100_sim = pybamm.Simulation(x100_model, parameter_values=self.parameter_values)
+        x0_model = ElectrodeSOHx0(param=self.param)
+        x0_sim = pybamm.Simulation(x0_model, parameter_values=self.parameter_values)
+        return [x100_sim, x0_sim]
 
     def solve(self, inputs):
         ics = self._set_up_solve(inputs)

pybamm/simulation.py

@@ -7,6 +7,7 @@
 import copy
 import warnings
 import sys
+from functools import lru_cache
 
 
 def is_notebook():
@@ -889,6 +890,7 @@ def step(
 
         return self.solution
 
+    @lru_cache
     def get_esoh_solver(self, calc_esoh):
         if (
             calc_esoh is False
@@ -897,13 +899,9 @@ def get_esoh_solver(self, calc_esoh):
         ):
             return None
 
-        try:
-            return self._esoh_solver
-        except AttributeError:
-            self._esoh_solver = pybamm.lithium_ion.ElectrodeSOHSolver(
-                self.parameter_values, self.model.param
-            )
-            return self._esoh_solver
+        return pybamm.lithium_ion.ElectrodeSOHSolver(
+            self.parameter_values, self.model.param
+        )
 
     def plot(self, output_variables=None, **kwargs):
         """

pybamm/solvers/solution.py

@@ -9,6 +9,7 @@
 import pybamm
 import pandas as pd
 from scipy.io import savemat
+from functools import cached_property
 
 
 class NumpyEncoder(json.JSONEncoder):
@@ -344,15 +345,9 @@ def all_models(self):
         """Model(s) used for solution"""
         return self._all_models
 
-    @property
+    @cached_property
     def all_inputs_casadi(self):
-        try:
-            return self._all_inputs_casadi
-        except AttributeError:
-            self._all_inputs_casadi = [
-                casadi.vertcat(*inp.values()) for inp in self.all_inputs
-            ]
-            return self._all_inputs_casadi
+        return [casadi.vertcat(*inp.values()) for inp in self.all_inputs]
 
     @property
     def t_event(self):
@@ -374,63 +369,55 @@ def termination(self, value):
         """Updates the reason for termination"""
         self._termination = value
 
-    @property
+    @cached_property
     def first_state(self):
         """
         A Solution object that only contains the first state. This is faster to evaluate
         than the full solution when only the first state is needed (e.g. to initialize
         a model with the solution)
         """
-        try:
-            return self._first_state
-        except AttributeError:
-            new_sol = Solution(
-                self.all_ts[0][:1],
-                self.all_ys[0][:, :1],
-                self.all_models[:1],
-                self.all_inputs[:1],
-                None,
-                None,
-                "success",
-            )
-            new_sol._all_inputs_casadi = self.all_inputs_casadi[:1]
-            new_sol._sub_solutions = self.sub_solutions[:1]
+        new_sol = Solution(
+            self.all_ts[0][:1],
+            self.all_ys[0][:, :1],
+            self.all_models[:1],
+            self.all_inputs[:1],
+            None,
+            None,
+            "success",
+        )
+        new_sol._all_inputs_casadi = self.all_inputs_casadi[:1]
+        new_sol._sub_solutions = self.sub_solutions[:1]
 
-            new_sol.solve_time = 0
-            new_sol.integration_time = 0
-            new_sol.set_up_time = 0
+        new_sol.solve_time = 0
+        new_sol.integration_time = 0
+        new_sol.set_up_time = 0
 
-            self._first_state = new_sol
-            return self._first_state
+        return new_sol
 
-    @property
+    @cached_property
     def last_state(self):
         """
         A Solution object that only contains the final state. This is faster to evaluate
         than the full solution when only the final state is needed (e.g. to initialize
         a model with the solution)
         """
-        try:
-            return self._last_state
-        except AttributeError:
-            new_sol = Solution(
-                self.all_ts[-1][-1:],
-                self.all_ys[-1][:, -1:],
-                self.all_models[-1:],
-                self.all_inputs[-1:],
-                self.t_event,
-                self.y_event,
-                self.termination,
-            )
-            new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:]
-            new_sol._sub_solutions = self.sub_solutions[-1:]
+        new_sol = Solution(
+            self.all_ts[-1][-1:],
+            self.all_ys[-1][:, -1:],
+            self.all_models[-1:],
+            self.all_inputs[-1:],
+            self.t_event,
+            self.y_event,
+            self.termination,
+        )
+        new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:]
+        new_sol._sub_solutions = self.sub_solutions[-1:]
 
-            new_sol.solve_time = 0
-            new_sol.integration_time = 0
-            new_sol.set_up_time = 0
+        new_sol.solve_time = 0
+        new_sol.integration_time = 0
+        new_sol.set_up_time = 0
 
-            self._last_state = new_sol
-            return self._last_state
+        return new_sol
 
     @property
     def total_time(self):