Repository: pybamm-team/PyBaMM

Issue:
Faster `diff`
Taking `diff` can be a bit slow. Might be able to speed up by first checking that the symbol with respect to which we want to differentiate is not in the expression tree of the symbol being differentiated (in which case return 0)

Discussion:

Patch:
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 853131326..c8c9f119c 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -110,10 +110,6 @@ class BinaryOperator(pybamm.Symbol):
                 )
             )
 
-    def _binary_evaluate(self, left, right):
-        """ Perform binary operation on nodes 'left' and 'right'. """
-        raise NotImplementedError
-
     def new_copy(self):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
         # process children
@@ -147,6 +143,10 @@ class BinaryOperator(pybamm.Symbol):
         right = self.children[1].evaluate_for_shape()
         return self._binary_evaluate(left, right)
 
+    def _binary_evaluate(self, left, right):
+        """ Perform binary operation on nodes 'left' and 'right'. """
+        raise NotImplementedError
+
 
 class Power(BinaryOperator):
     """A node in the expression tree representing a `**` power operator
@@ -158,17 +158,14 @@ class Power(BinaryOperator):
         """ See :meth:`pybamm.BinaryOperator.__init__()`. """
         super().__init__("**", left, right)
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            # apply chain rule and power rule
-            base, exponent = self.orphans
-            return base ** (exponent - 1) * (
-                exponent * base.diff(variable)
-                + base * pybamm.Function(np.log, base) * exponent.diff(variable)
-            )
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        # apply chain rule and power rule
+        base, exponent = self.orphans
+        return base ** (exponent - 1) * (
+            exponent * base.diff(variable)
+            + base * pybamm.Function(np.log, base) * exponent.diff(variable)
+        )
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -219,12 +216,9 @@ class Addition(BinaryOperator):
         """ See :meth:`pybamm.BinaryOperator.__init__()`. """
         super().__init__("+", left, right)
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            return self.left.diff(variable) + self.right.diff(variable)
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        return self.left.diff(variable) + self.right.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -273,12 +267,9 @@ class Subtraction(BinaryOperator):
 
         super().__init__("-", left, right)
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            return self.left.diff(variable) - self.right.diff(variable)
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        return self.left.diff(variable) - self.right.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -330,14 +321,11 @@ class Multiplication(BinaryOperator):
 
         super().__init__("*", left, right)
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            # apply product rule
-            left, right = self.orphans
-            return left.diff(variable) * right + left * right.diff(variable)
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        # apply product rule
+        left, right = self.orphans
+        return left.diff(variable) * right + left * right.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -449,16 +437,11 @@ class Division(BinaryOperator):
         """ See :meth:`pybamm.BinaryOperator.__init__()`. """
         super().__init__("/", left, right)
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            # apply quotient rule
-            top, bottom = self.orphans
-            return (
-                top.diff(variable) * bottom - top * bottom.diff(variable)
-            ) / bottom ** 2
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        # apply quotient rule
+        top, bottom = self.orphans
+        return (top.diff(variable) * bottom - top * bottom.diff(variable)) / bottom ** 2
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 364e7d3df..eff3a2769 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -331,9 +331,10 @@ class Symbol(anytree.NodeMixin):
 
     def diff(self, variable):
         """
-        Differentiate a symbol with respect to a variable. Default behaviour is to
-        return `1` if differentiating with respect to yourself and zero otherwise.
-        Binary and Unary Operators override this.
+        Differentiate a symbol with respect to a variable. For any symbol that can be
+        differentiated, return `1` if differentiating with respect to yourself,
+        `self._diff(variable)` if `variable` is in the expression tree of the symbol,
+        and zero otherwise.
 
         Parameters
         ----------
@@ -343,9 +344,15 @@ class Symbol(anytree.NodeMixin):
         """
         if variable.id == self.id:
             return pybamm.Scalar(1)
+        elif any(variable.id == x.id for x in self.pre_order()):
+            return self._diff(variable)
         else:
             return pybamm.Scalar(0)
 
+    def _diff(self, variable):
+        "Default behaviour for differentiation, overriden by Binary and Unary Operators"
+        raise NotImplementedError
+
     def jac(self, variable):
         """
         Differentiate a symbol with respect to a (slice of) a State Vector.
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index daf1b36ed..6563d9462 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -88,12 +88,9 @@ class Negate(UnaryOperator):
         """ See :meth:`pybamm.Symbol.__str__()`. """
         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.child.diff(variable)
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        return -self.child.diff(variable)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -159,19 +156,10 @@ class Function(UnaryOperator):
         else:
             self.takes_no_params = len(signature(func).parameters) == 0
 
-    def diff(self, variable):
-        """ See :meth:`pybamm.Symbol.diff()`. """
-        if variable.id == self.id:
-            return pybamm.Scalar(1)
-        else:
-            child = self.orphans[0]
-            if variable.id in [symbol.id for symbol in child.pre_order()]:
-                # if variable appears in the function,use autograd to differentiate
-                # function, and apply chain rule
-                return child.diff(variable) * Function(autograd.grad(self.func), child)
-            else:
-                # otherwise the derivative of the function is zero
-                return pybamm.Scalar(0)
+    def _diff(self, variable):
+        """ See :meth:`pybamm.Symbol._diff()`. """
+        child = self.orphans[0]
+        return child.diff(variable) * Function(autograd.grad(self.func), child)
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """


Test patch:
diff --git a/tests/unit/test_expression_tree/test_symbolic_diff.py b/tests/unit/test_expression_tree/test_symbolic_diff.py
index 97af4e539..75d06d614 100644
--- a/tests/unit/test_expression_tree/test_symbolic_diff.py
+++ b/tests/unit/test_expression_tree/test_symbolic_diff.py
@@ -57,6 +57,19 @@ class TestSymbolicDifferentiation(unittest.TestCase):
         func = pybamm.Function(np.sin, a * b)
         self.assertEqual(func.diff(a).evaluate(y=y), 3 * np.cos(15))
 
+    def test_diff_zero(self):
+        a = pybamm.StateVector(slice(0, 1))
+        b = pybamm.StateVector(slice(1, 2))
+        func = (a * 2 + 5 * (-a)) / (a * a)
+        self.assertEqual(func.diff(b).id, pybamm.Scalar(0).id)
+        self.assertNotEqual(func.diff(a).id, pybamm.Scalar(0).id)
+
+    def test_exceptions(self):
+        a = pybamm.Symbol("a")
+        b = pybamm.Symbol("b")
+        with self.assertRaises(NotImplementedError):
+            a._diff(b)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")