Repository: pybamm-team/PyBaMM

Issue:
Cache children
Predefine children, such as `left` and `right` in `BinaryOperator`, to get around `self.children` being slow

Discussion:

Patch:
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 80388c682..bab0dea5a 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -526,10 +526,12 @@ class BinaryOperator(pybamm.Symbol):
             right = pybamm.Scalar(right)
         domain = self.get_children_domains(left.domain, right.domain)
         super().__init__(name, children=[left, right], domain=domain)
+        self.left = self.children[0]
+        self.right = self.children[1]
 
     def __str__(self):
         """ See :meth:`pybamm.Symbol.__str__()`. """
-        return "{!s} {} {!s}".format(self.children[0], self.name, self.children[1])
+        return "{!s} {} {!s}".format(self.left, self.name, self.right)
 
     def get_children_domains(self, ldomain, rdomain):
         if ldomain == rdomain:
@@ -550,8 +552,8 @@ class BinaryOperator(pybamm.Symbol):
 
     def simplify(self):
         """ See :meth:`pybamm.Symbol.simplify()`. """
-        left = self.children[0].simplify()
-        right = self.children[1].simplify()
+        left = self.left.simplify()
+        right = self.right.simplify()
 
         # _binary_simplify defined in derived classes for specific rules
         new_node = self._binary_simplify(left, right)
@@ -569,14 +571,14 @@ class BinaryOperator(pybamm.Symbol):
             try:
                 return known_evals[id], known_evals
             except KeyError:
-                left, known_evals = self.children[0].evaluate(t, y, known_evals)
-                right, known_evals = self.children[1].evaluate(t, y, known_evals)
+                left, known_evals = self.left.evaluate(t, y, known_evals)
+                right, known_evals = self.right.evaluate(t, y, known_evals)
                 value = self._binary_evaluate(left, right)
                 known_evals[id] = value
                 return value, known_evals
         else:
-            left = self.children[0].evaluate(t, y)
-            right = self.children[1].evaluate(t, y)
+            left = self.left.evaluate(t, y)
+            right = self.right.evaluate(t, y)
             return self._binary_evaluate(left, right)
 
 
@@ -656,14 +658,14 @@ class Addition(BinaryOperator):
         if variable.id == self.id:
             return pybamm.Scalar(1)
         else:
-            return self.children[0].diff(variable) + self.children[1].diff(variable)
+            return self.left.diff(variable) + self.right.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
         if variable.id == self.id:
             return pybamm.Scalar(1)
         else:
-            return self.children[0].jac(variable) + self.children[1].jac(variable)
+            return self.left.jac(variable) + self.right.jac(variable)
 
     def _binary_evaluate(self, left, right):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
@@ -697,11 +699,11 @@ class Subtraction(BinaryOperator):
         if variable.id == self.id:
             return pybamm.Scalar(1)
         else:
-            return self.children[0].diff(variable) - self.children[1].diff(variable)
+            return self.left.diff(variable) - self.right.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
-        return self.children[0].jac(variable) - self.children[1].jac(variable)
+        return self.left.jac(variable) - self.right.jac(variable)
 
     def _binary_evaluate(self, left, right):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index c9d5ad7e1..d0af9a374 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -53,16 +53,17 @@ class Concatenation(pybamm.Symbol):
 
     def evaluate(self, t=None, y=None, known_evals=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
+        children = self.cached_children
         if known_evals is not None:
             if self.id not in known_evals:
-                children_eval = [None] * len(self.children)
-                for idx, child in enumerate(self.children):
+                children_eval = [None] * len(children)
+                for idx, child in enumerate(children):
                     children_eval[idx], known_evals = child.evaluate(t, y, known_evals)
                 known_evals[self.id] = self._concatenation_evaluate(children_eval)
             return known_evals[self.id], known_evals
         else:
-            children_eval = [None] * len(self.children)
-            for idx, child in enumerate(self.children):
+            children_eval = [None] * len(children)
+            for idx, child in enumerate(children):
                 children_eval[idx] = child.evaluate(t, y)
             return self._concatenation_evaluate(children_eval)
 
@@ -102,14 +103,15 @@ class NumpyConcatenation(Concatenation):
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
-        if len(self.children) == 0:
+        children = self.cached_children
+        if len(children) == 0:
             return pybamm.Scalar(0)
         else:
-            return SparseStack(*[child.jac(variable) for child in self.children])
+            return SparseStack(*[child.jac(variable) for child in children])
 
     def simplify(self):
         """ See :meth:`pybamm.Symbol.simplify()`. """
-        children = [child.simplify() for child in self.children]
+        children = [child.simplify() for child in self.cached_children]
 
         new_node = self.__class__(*children)
 
@@ -171,7 +173,7 @@ class DomainConcatenation(Concatenation):
 
             # create disc of domain => slice for each child
             self._children_slices = [
-                self.create_slices(child) for child in self.children
+                self.create_slices(child) for child in self.cached_children
             ]
         else:
             self._mesh = copy.copy(copy_this._mesh)
@@ -217,14 +219,15 @@ class DomainConcatenation(Concatenation):
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
-        if len(self.children) == 0:
+        children = self.cached_children
+        if len(children) == 0:
             return pybamm.Scalar(0)
         else:
-            return SparseStack(*[child.jac(variable) for child in self.children])
+            return SparseStack(*[child.jac(variable) for child in children])
 
     def simplify(self):
         """ See :meth:`pybamm.Symbol.simplify()`. """
-        children = self.children
+        children = self.cached_children
         # Simplify Concatenation of StateVectors to a single StateVector
         if all([isinstance(x, pybamm.StateVector) for x in children]) and all(
             [
@@ -270,16 +273,17 @@ class SparseStack(Concatenation):
 
     def evaluate(self, t=None, y=None, known_evals=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
+        children = self.cached_children
         if known_evals is not None:
             if self.id not in known_evals:
-                children_eval = [None] * len(self.children)
-                for idx, child in enumerate(self.children):
+                children_eval = [None] * len(children)
+                for idx, child in enumerate(children):
                     children_eval[idx], known_evals = child.evaluate(t, y, known_evals)
                 known_evals[self.id] = self._concatenation_evaluate(children_eval)
             return known_evals[self.id], known_evals
         else:
-            children_eval = [None] * len(self.children)
-            for idx, child in enumerate(self.children):
+            children_eval = [None] * len(children)
+            for idx, child in enumerate(children):
                 children_eval[idx] = child.evaluate(t, y)
             return self._concatenation_evaluate(children_eval)
 
@@ -292,7 +296,7 @@ class SparseStack(Concatenation):
 
     def simplify(self):
         """ See :meth:`pybamm.Symbol.simplify()`. """
-        children = [child.simplify() for child in self.children]
+        children = [child.simplify() for child in self.cached_children]
 
         new_node = self.__class__(*children)
 
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 7611b986f..978a7d802 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -57,6 +57,9 @@ class Symbol(anytree.NodeMixin):
             # this also adds copy.copy(child) to self.children
             copy.copy(child).parent = self
 
+        # cache children
+        self.cached_children = super(Symbol, self).children
+
         # Set domain (and hence id)
         self.domain = domain
 
@@ -64,6 +67,10 @@ class Symbol(anytree.NodeMixin):
         self._has_left_ghost_cell = False
         self._has_right_ghost_cell = False
 
+    @property
+    def children(self):
+        return self.cached_children
+
     @property
     def name(self):
         """name of the node"""
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 6711fb1cf..52c74fca5 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -31,14 +31,15 @@ class UnaryOperator(pybamm.Symbol):
 
     def __init__(self, name, child):
         super().__init__(name, children=[child], domain=child.domain)
+        self.child = self.children[0]
 
     def __str__(self):
         """ See :meth:`pybamm.Symbol.__str__()`. """
-        return "{}({!s})".format(self.name, self.children[0])
+        return "{}({!s})".format(self.name, self.child)
 
     def simplify(self):
         """ See :meth:`pybamm.Symbol.simplify()`. """
-        child = self.children[0].simplify()
+        child = self.child.simplify()
 
         # _binary_simplify defined in derived classes for specific rules
         new_node = self._unary_simplify(child)
@@ -58,11 +59,11 @@ class UnaryOperator(pybamm.Symbol):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
         if known_evals is not None:
             if self.id not in known_evals:
-                child, known_evals = self.children[0].evaluate(t, y, known_evals)
+                child, known_evals = self.child.evaluate(t, y, known_evals)
                 known_evals[self.id] = self._unary_evaluate(child)
             return known_evals[self.id], known_evals
         else:
-            child = self.children[0].evaluate(t, y)
+            child = self.child.evaluate(t, y)
             return self._unary_evaluate(child)
 
 
@@ -78,18 +79,18 @@ class Negate(UnaryOperator):
 
     def __str__(self):
         """ See :meth:`pybamm.Symbol.__str__()`. """
-        return "{}{!s}".format(self.name, self.children[0])
+        return "{}{!s}".format(self.name, self.child)
 
     def diff(self, variable):
         """ See :meth:`pybamm.Symbol.diff()`. """
         if variable.id == self.id:
             return pybamm.Scalar(1)
         else:
-            return -self.children[0].diff(variable)
+            return -self.child.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
-        return -self.children[0].jac(variable)
+        return -self.child.jac(variable)
 
     def _unary_evaluate(self, child):
         """ See :meth:`UnaryOperator._unary_evaluate()`. """
@@ -192,7 +193,7 @@ class Function(UnaryOperator):
             # If self.func() takes no parameters then we can always simplify it
             return pybamm.Scalar(self.func())
         else:
-            child = self.children[0].simplify()
+            child = self.child.simplify()
 
             new_node = pybamm.Function(self.func, child)
 


Test patch:
diff --git a/tests/unit/test_expression_tree/test_symbol.py b/tests/unit/test_expression_tree/test_symbol.py
index 40e157886..a929e6698 100644
--- a/tests/unit/test_expression_tree/test_symbol.py
+++ b/tests/unit/test_expression_tree/test_symbol.py
@@ -17,6 +17,28 @@ class TestSymbol(unittest.TestCase):
         self.assertEqual(sym.name, "a symbol")
         self.assertEqual(str(sym), "a symbol")
 
+    def test_cached_children(self):
+        symc1 = pybamm.Symbol("child1")
+        symc2 = pybamm.Symbol("child2")
+        symc3 = pybamm.Symbol("child3")
+        symp = pybamm.Symbol("parent", children=[symc1, symc2])
+
+        # test tuples of children for equality based on their name
+        def check_are_equal(children1, children2):
+            self.assertEqual(len(children1), len(children2))
+            for i in range(len(children1)):
+                self.assertEqual(children1[i].name, children2[i].name)
+
+        check_are_equal(symp.children, super(pybamm.Symbol, symp).children)
+        check_are_equal(symp.children, (symc1, symc2))
+
+        # update children, since we cache the children they will be unchanged
+        symc3.parent = symp
+        check_are_equal(symp.children, (symc1, symc2))
+
+        # check that the *actual* children are updated
+        check_are_equal(super(pybamm.Symbol, symp).children, (symc1, symc2, symc3))
+
     def test_symbol_simplify(self):
         a = pybamm.Scalar(0)
         b = pybamm.Scalar(1)