Repository: tinosulzer/PyBaMM

Issue:
overloaded binary operators on `Symbol` take numbers
see #36. It might be useful for the overloaded binary operators on`Symbol` (e.g. `__mul__`) be able to take `Numbers` (e.g. `int'). They would then implicitly convert this to a `Scalar` node.

Note: also need to add `__pow__` to this list.

Discussion:

Patch:
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 890b01e7..78cfee5b 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -5,6 +5,8 @@ from __future__ import absolute_import, division
 from __future__ import print_function, unicode_literals
 import pybamm
 
+import numbers
+
 
 class BinaryOperator(pybamm.Symbol):
     """A node in the expression tree representing a binary operator (e.g. `+`, `*`)
@@ -18,14 +20,27 @@ class BinaryOperator(pybamm.Symbol):
 
     name : str
         name of the node
-    left : :class:`Symbol`
-        lhs child node
-    right : :class:`Symbol`
-        rhs child node
+    left : :class:`Symbol` or :class:`Number`
+        lhs child node (converted to :class:`Scalar` if Number)
+    right : :class:`Symbol` or :class:`Number`
+        rhs child node (converted to :class:`Scalar` if Number)
 
     """
 
     def __init__(self, name, left, right):
+        assert isinstance(left, (pybamm.Symbol, numbers.Number)) and isinstance(
+            right, (pybamm.Symbol, numbers.Number)
+        ), TypeError(
+            """left and right must both be Symbols or Numbers
+                but they are {} and {}""".format(
+                type(left), type(right)
+            )
+        )
+        if isinstance(left, numbers.Number):
+            left = pybamm.Scalar(left)
+        if isinstance(right, numbers.Number):
+            right = pybamm.Scalar(right)
+
         super().__init__(name, children=[left, right])
 
     def __str__(self):
@@ -45,7 +60,7 @@ class Power(BinaryOperator):
 
     def evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
-        return self.children[0].evaluate(t, y)**self.children[1].evaluate(t, y)
+        return self.children[0].evaluate(t, y) ** self.children[1].evaluate(t, y)
 
 
 class Addition(BinaryOperator):
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 79491358..0c9a93b5 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -4,7 +4,9 @@
 from __future__ import absolute_import, division
 from __future__ import print_function, unicode_literals
 import pybamm
+
 import anytree
+import numbers
 import copy
 
 
@@ -81,37 +83,88 @@ class Symbol(anytree.NodeMixin):
 
     def __repr__(self):
         """returns the string `__class__(id, name, parent expression)`"""
-        return "{!s}({}, {!s}, {!s})".format(self.__class__.__name__,
-                                             hex(self.id), self._name, self.parent)
+        return "{!s}({}, {!s}, {!s})".format(
+            self.__class__.__name__, hex(self.id), self._name, self.parent
+        )
 
     def __add__(self, other):
         """return an :class:`Addition` object"""
-        if isinstance(other, Symbol):
+        if isinstance(other, (Symbol, numbers.Number)):
             return pybamm.Addition(self, other)
         else:
             raise NotImplementedError
 
+    def __radd__(self, other):
+        """return an :class:`Addition` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Addition(other, self)
+        else:
+            raise NotImplementedError
+
     def __sub__(self, other):
-        """return an :class:`Subtraction` object"""
-        if isinstance(other, Symbol):
+        """return a :class:`Subtraction` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
             return pybamm.Subtraction(self, other)
         else:
             raise NotImplementedError
 
+    def __rsub__(self, other):
+        """return a :class:`Subtraction` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Subtraction(other, self)
+        else:
+            raise NotImplementedError
+
     def __mul__(self, other):
-        """return an :class:`Multiplication` object"""
-        if isinstance(other, Symbol):
+        """return a :class:`Multiplication` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
             return pybamm.Multiplication(self, other)
         else:
             raise NotImplementedError
 
+    def __rmul__(self, other):
+        """return a :class:`Multiplication` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Multiplication(other, self)
+        else:
+            raise NotImplementedError
+
     def __truediv__(self, other):
-        """return an :class:`Division` object"""
-        if isinstance(other, Symbol):
+        """return a :class:`Division` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
             return pybamm.Division(self, other)
         else:
             raise NotImplementedError
 
+    def __rtruediv__(self, other):
+        """return a :class:`Division` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Division(other, self)
+        else:
+            raise NotImplementedError
+
+    def __pow__(self, other):
+        """return a :class:`Power` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Power(self, other)
+        else:
+            raise NotImplementedError
+
+    def __rpow__(self, other):
+        """return a :class:`Power` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.Power(other, self)
+        else:
+            raise NotImplementedError
+
+    def __neg__(self):
+        """return a :class:`Negate` object"""
+        return pybamm.Negate(self)
+
+    def __abs__(self):
+        """return an :class:`AbsoluteValue` object"""
+        return pybamm.AbsoluteValue(self)
+
     def evaluate(self, t=None, y=None):
         """evaluate expression tree
 
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 926cdc90..a0a729f4 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -90,7 +90,6 @@ class SpatialOperator(UnaryOperator):
 
     def __init__(self, name, child):
         super().__init__(name, child)
-        # self.domain = child.domain
 
 
 class Gradient(SpatialOperator):


Test patch:
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index a8044d8f..e2f873d5 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -71,30 +71,38 @@ class TestDiscretise(unittest.TestCase):
         # variable
         var = pybamm.Variable("var")
         y_slices = {var.id: slice(53)}
-        var_disc = disc.process_symbol(var, None, y_slices, None)
+        var_disc = disc.process_symbol(var, None, y_slices)
         self.assertTrue(isinstance(var_disc, pybamm.StateVector))
         self.assertEqual(var_disc._y_slice, y_slices[var.id])
         # scalar
         scal = pybamm.Scalar(5)
-        scal_disc = disc.process_symbol(scal, None, None, None)
+        scal_disc = disc.process_symbol(scal, None)
         self.assertTrue(isinstance(scal_disc, pybamm.Scalar))
         self.assertEqual(scal_disc.value, scal.value)
 
         # parameter
         par = pybamm.Parameter("par")
-        par_disc = disc.process_symbol(par, None, None, None)
+        par_disc = disc.process_symbol(par, None)
         self.assertTrue(isinstance(par_disc, pybamm.Parameter))
         self.assertEqual(par_disc.name, par.name)
 
         # binary operator
         bin = var + scal
-        bin_disc = disc.process_symbol(bin, None, y_slices, None)
+        bin_disc = disc.process_symbol(bin, None, y_slices)
         self.assertTrue(isinstance(bin_disc, pybamm.Addition))
         self.assertTrue(isinstance(bin_disc.children[0], pybamm.StateVector))
         self.assertTrue(isinstance(bin_disc.children[1], pybamm.Scalar))
 
         # non-spatial unary operator
-        # TODO: none of these implemented yet
+        un1 = -var
+        un1_disc = disc.process_symbol(un1, None, y_slices)
+        self.assertTrue(isinstance(un1_disc, pybamm.Negate))
+        self.assertTrue(isinstance(un1_disc.children[0], pybamm.StateVector))
+
+        un2 = abs(scal)
+        un2_disc = disc.process_symbol(un2, None)
+        self.assertTrue(isinstance(un2_disc, pybamm.AbsoluteValue))
+        self.assertTrue(isinstance(un2_disc.children[0], pybamm.Scalar))
 
     def test_process_complex_expression(self):
         var1 = pybamm.Variable("var1")
@@ -151,7 +159,7 @@ class TestDiscretise(unittest.TestCase):
             self.assertTrue(isinstance(eqn_disc.children[1], pybamm.StateVector))
 
             y = mesh["whole cell"].nodes ** 2
-            var_disc = disc.process_symbol(var, None, y_slices, None)
+            var_disc = disc.process_symbol(var, None, y_slices)
             # grad and var are identity operators here (for testing purposes)
             np.testing.assert_array_equal(
                 eqn_disc.evaluate(None, y), var_disc.evaluate(None, y)
diff --git a/tests/test_expression_tree/test_binary_operators.py b/tests/test_expression_tree/test_binary_operators.py
index 716a1786..fb0c00d8 100644
--- a/tests/test_expression_tree/test_binary_operators.py
+++ b/tests/test_expression_tree/test_binary_operators.py
@@ -58,6 +58,12 @@ class TestBinaryOperators(unittest.TestCase):
         bin5 = pybamm.BinaryOperator("test", a, d)
         self.assertNotEqual(bin1.id, bin5.id)
 
+    def test_number_overloading(self):
+        a = pybamm.Scalar(4)
+        prod = a * 3
+        self.assertTrue(isinstance(prod.children[1], pybamm.Scalar))
+        self.assertEqual(prod.evaluate(), 12)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_expression_tree/test_symbol.py b/tests/test_expression_tree/test_symbol.py
index ac9df67d..b787d7f7 100644
--- a/tests/test_expression_tree/test_symbol.py
+++ b/tests/test_expression_tree/test_symbol.py
@@ -18,18 +18,57 @@ class TestSymbol(unittest.TestCase):
         a = pybamm.Symbol("a")
         b = pybamm.Symbol("b")
 
-        self.assertTrue(isinstance(a + b, pybamm.Symbol))
-        self.assertTrue(isinstance(a - b, pybamm.Symbol))
-        self.assertTrue(isinstance(a * b, pybamm.Symbol))
-        self.assertTrue(isinstance(a / b, pybamm.Symbol))
+        # unary
+        self.assertTrue(isinstance(-a, pybamm.Negate))
+        self.assertTrue(isinstance(abs(a), pybamm.AbsoluteValue))
+
+        # binary - two symbols
+        self.assertTrue(isinstance(a + b, pybamm.Addition))
+        self.assertTrue(isinstance(a - b, pybamm.Subtraction))
+        self.assertTrue(isinstance(a * b, pybamm.Multiplication))
+        self.assertTrue(isinstance(a / b, pybamm.Division))
+        self.assertTrue(isinstance(a ** b, pybamm.Power))
+
+        # binary - symbol and number
+        self.assertTrue(isinstance(a + 2, pybamm.Addition))
+        self.assertTrue(isinstance(a - 2, pybamm.Subtraction))
+        self.assertTrue(isinstance(a * 2, pybamm.Multiplication))
+        self.assertTrue(isinstance(a / 2, pybamm.Division))
+        self.assertTrue(isinstance(a ** 2, pybamm.Power))
+
+        # binary - number and symbol
+        self.assertTrue(isinstance(3 + b, pybamm.Addition))
+        self.assertEqual((3 + b).children[1].id, b.id)
+        self.assertTrue(isinstance(3 - b, pybamm.Subtraction))
+        self.assertEqual((3 - b).children[1].id, b.id)
+        self.assertTrue(isinstance(3 * b, pybamm.Multiplication))
+        self.assertEqual((3 * b).children[1].id, b.id)
+        self.assertTrue(isinstance(3 / b, pybamm.Division))
+        self.assertEqual((3 / b).children[1].id, b.id)
+        self.assertTrue(isinstance(3 ** b, pybamm.Power))
+        self.assertEqual((3 ** b).children[1].id, b.id)
+
+        # error raising
+        with self.assertRaises(NotImplementedError):
+            a + "two"
+        with self.assertRaises(NotImplementedError):
+            a - "two"
+        with self.assertRaises(NotImplementedError):
+            a * "two"
+        with self.assertRaises(NotImplementedError):
+            a / "two"
+        with self.assertRaises(NotImplementedError):
+            a ** "two"
+        with self.assertRaises(NotImplementedError):
+            "two" + a
         with self.assertRaises(NotImplementedError):
-            a + 2
+            "two" - a
         with self.assertRaises(NotImplementedError):
-            a - 2
+            "two" * a
         with self.assertRaises(NotImplementedError):
-            a * 2
+            "two" / a
         with self.assertRaises(NotImplementedError):
-            a / 2
+            "two" ** a
 
     def test_multiple_symbols(self):
         a = pybamm.Symbol("a")
diff --git a/tests/test_expression_tree/test_unary_operators.py b/tests/test_expression_tree/test_unary_operators.py
index 2f92ae6c..fcf0952b 100644
--- a/tests/test_expression_tree/test_unary_operators.py
+++ b/tests/test_expression_tree/test_unary_operators.py
@@ -39,8 +39,9 @@ class TestUnaryOperators(unittest.TestCase):
         grad = pybamm.Gradient(a)
         self.assertEqual(grad.children[0].name, a.name)
 
-    def test_gradient_printing(self):
+    def test_printing(self):
         a = pybamm.Symbol("a")
+        self.assertEqual(str(-a), "-a")
         grad = pybamm.Gradient(a)
         self.assertEqual(grad.name, "grad")
         self.assertEqual(str(grad), "grad(a)")