#709 first go at simplifying on creation

Author: valentinsulzer

pybamm/expression_tree/operations/simplify.py

@@ -8,11 +8,17 @@
 from scipy.sparse import issparse
 
 
-def simplify_if_constant(symbol):
+def simplify_if_constant(symbol, keep_domains=False):
     """
     Utility function to simplify an expression tree if it evalutes to a constant
     scalar, vector or matrix
     """
+    if keep_domains is True:
+        domain = symbol.domain
+        auxiliary_domains = symbol.auxiliary_domains
+    else:
+        domain = None
+        auxiliary_domains = None
     if symbol.is_constant():
         result = symbol.evaluate_ignoring_errors()
         if result is not None:
@@ -22,9 +28,13 @@ def simplify_if_constant(symbol):
                 return pybamm.Scalar(result)
             elif isinstance(result, np.ndarray) or issparse(result):
                 if result.ndim == 1 or result.shape[1] == 1:
-                    return pybamm.Vector(result)
+                    return pybamm.Vector(
+                        result, domain=domain, auxiliary_domains=auxiliary_domains
+                    )
                 else:
-                    return pybamm.Matrix(result)
+                    return pybamm.Matrix(
+                        result, domain=domain, auxiliary_domains=auxiliary_domains
+                    )
 
     return symbol
 

pybamm/expression_tree/symbol.py

@@ -360,79 +360,109 @@ def __repr__(self):
 
     def __add__(self, other):
         """return an :class:`Addition` object"""
-        return pybamm.Addition(self, other)
+        return pybamm.simplify_if_constant(
+            pybamm.Addition(self, other), keep_domains=True
+        )
 
     def __radd__(self, other):
         """return an :class:`Addition` object"""
-        return pybamm.Addition(other, self)
+        return pybamm.simplify_if_constant(
+            pybamm.Addition(other, self), keep_domains=True
+        )
 
     def __sub__(self, other):
         """return a :class:`Subtraction` object"""
-        return pybamm.Subtraction(self, other)
+        return pybamm.simplify_if_constant(
+            pybamm.Subtraction(self, other), keep_domains=True
+        )
 
     def __rsub__(self, other):
         """return a :class:`Subtraction` object"""
-        return pybamm.Subtraction(other, self)
+        return pybamm.simplify_if_constant(
+            pybamm.Subtraction(other, self), keep_domains=True
+        )
 
     def __mul__(self, other):
         """return a :class:`Multiplication` object"""
-        return pybamm.Multiplication(self, other)
+        return pybamm.simplify_if_constant(
+            pybamm.Multiplication(self, other), keep_domains=True
+        )
 
     def __rmul__(self, other):
         """return a :class:`Multiplication` object"""
-        return pybamm.Multiplication(other, self)
+        return pybamm.simplify_if_constant(
+            pybamm.Multiplication(other, self), keep_domains=True
+        )
 
     def __matmul__(self, other):
         """return a :class:`MatrixMultiplication` object"""
-        return pybamm.MatrixMultiplication(self, other)
+        return pybamm.simplify_if_constant(
+            pybamm.MatrixMultiplication(self, other), keep_domains=True
+        )
 
     def __rmatmul__(self, other):
         """return a :class:`MatrixMultiplication` object"""
-        return pybamm.MatrixMultiplication(other, self)
+        return pybamm.simplify_if_constant(
+            pybamm.MatrixMultiplication(other, self), keep_domains=True
+        )
 
     def __truediv__(self, other):
         """return a :class:`Division` object"""
-        return pybamm.Division(self, other)
+        return pybamm.simplify_if_constant(
+            pybamm.Division(self, other), keep_domains=True
+        )
 
     def __rtruediv__(self, other):
         """return a :class:`Division` object"""
-        return pybamm.Division(other, self)
+        return pybamm.simplify_if_constant(
+            pybamm.Division(other, self), keep_domains=True
+        )
 
     def __pow__(self, other):
         """return a :class:`Power` object"""
-        return pybamm.Power(self, other)
+        return pybamm.simplify_if_constant(pybamm.Power(self, other), keep_domains=True)
 
     def __rpow__(self, other):
         """return a :class:`Power` object"""
-        return pybamm.Power(other, self)
+        return pybamm.simplify_if_constant(pybamm.Power(other, self), keep_domains=True)
 
     def __lt__(self, other):
         """return a :class:`Heaviside` object"""
-        return pybamm.Heaviside(self, other, equal=False)
+        return pybamm.simplify_if_constant(
+            pybamm.Heaviside(self, other, equal=False), keep_domains=True
+        )
 
     def __le__(self, other):
         """return a :class:`Heaviside` object"""
-        return pybamm.Heaviside(self, other, equal=True)
+        return pybamm.simplify_if_constant(
+            pybamm.Heaviside(self, other, equal=True), keep_domains=True
+        )
 
     def __gt__(self, other):
         """return a :class:`Heaviside` object"""
-        return pybamm.Heaviside(other, self, equal=False)
+        return pybamm.simplify_if_constant(
+            pybamm.Heaviside(other, self, equal=False), keep_domains=True
+        )
 
     def __ge__(self, other):
         """return a :class:`Heaviside` object"""
-        return pybamm.Heaviside(other, self, equal=True)
+        return pybamm.simplify_if_constant(
+            pybamm.Heaviside(other, self, equal=True), keep_domains=True
+        )
 
     def __neg__(self):
         """return a :class:`Negate` object"""
-        return pybamm.Negate(self)
+        return pybamm.simplify_if_constant(pybamm.Negate(self), keep_domains=True)
 
     def __abs__(self):
         """return an :class:`AbsoluteValue` object"""
-        return pybamm.AbsoluteValue(self)
+        return pybamm.simplify_if_constant(
+            pybamm.AbsoluteValue(self), keep_domains=True
+        )
 
     def __getitem__(self, key):
         """return a :class:`Index` object"""
-        return pybamm.Index(self, key)
+        return pybamm.simplify_if_constant(pybamm.Index(self, key), keep_domains=True)
 
     def diff(self, variable):
         """

pybamm/spatial_methods/finite_volume.py

@@ -1025,7 +1025,9 @@ def process_binary_operators(self, bin_op, left, right, disc_left, disc_right):
                 method = "arithmetic"
             disc_left = self.node_to_edge(disc_left, method=method)
         # Return new binary operator with appropriate class
-        out = bin_op.__class__(disc_left, disc_right)
+        out = pybamm.simplify_if_constant(
+            bin_op.__class__(disc_left, disc_right), keep_domains=True
+        )
         return out
 
     def concatenation(self, disc_children):

tests/unit/test_discretisations/test_discretisation.py

@@ -251,10 +251,10 @@ def test_process_symbol_base(self):
         self.assertIsInstance(un1_disc, pybamm.Negate)
         self.assertIsInstance(un1_disc.children[0], pybamm.StateVector)
 
-        un2 = abs(scal)
+        un2 = abs(var)
         un2_disc = disc.process_symbol(un2)
         self.assertIsInstance(un2_disc, pybamm.AbsoluteValue)
-        self.assertIsInstance(un2_disc.children[0], pybamm.Scalar)
+        self.assertIsInstance(un2_disc.children[0], pybamm.StateVector)
 
         # function of one variable
         def myfun(x):
@@ -749,7 +749,7 @@ def test_process_empty_model(self):
     def test_broadcast(self):
         whole_cell = ["negative electrode", "separator", "positive electrode"]
 
-        a = pybamm.Scalar(7)
+        a = pybamm.InputParameter("a")
         var = pybamm.Variable("var")
 
         # create discretisation
@@ -761,13 +761,14 @@ def test_broadcast(self):
         # scalar
         broad = disc.process_symbol(pybamm.FullBroadcast(a, whole_cell, {}))
         np.testing.assert_array_equal(
-            broad.evaluate(), 7 * np.ones_like(combined_submesh[0].nodes[:, np.newaxis])
+            broad.evaluate(u={"a": 7}),
+            7 * np.ones_like(combined_submesh[0].nodes[:, np.newaxis]),
         )
         self.assertEqual(broad.domain, whole_cell)
 
         broad_disc = disc.process_symbol(broad)
         self.assertIsInstance(broad_disc, pybamm.Multiplication)
-        self.assertIsInstance(broad_disc.children[0], pybamm.Scalar)
+        self.assertIsInstance(broad_disc.children[0], pybamm.InputParameter)
         self.assertIsInstance(broad_disc.children[1], pybamm.Vector)
 
         # process Broadcast variable
@@ -804,16 +805,14 @@ def test_secondary_broadcast_2D(self):
         # secondary broadcast in 2D --> Matrix multiplication
         disc = get_discretisation_for_testing()
         mesh = disc.mesh
-        var = pybamm.Vector(
-            mesh["negative particle"][0].nodes, domain=["negative particle"]
-        )
+        var = pybamm.Variable("var", domain=["negative particle"])
         broad = pybamm.SecondaryBroadcast(var, "negative electrode")
 
         disc.set_variable_slices([var])
         broad_disc = disc.process_symbol(broad)
         self.assertIsInstance(broad_disc, pybamm.MatrixMultiplication)
         self.assertIsInstance(broad_disc.children[0], pybamm.Matrix)
-        self.assertIsInstance(broad_disc.children[1], pybamm.Vector)
+        self.assertIsInstance(broad_disc.children[1], pybamm.StateVector)
         self.assertEqual(
             broad_disc.shape,
             (mesh["negative particle"][0].npts * mesh["negative electrode"][0].npts, 1),
@@ -905,7 +904,9 @@ def test_exceptions(self):
 
         # check doesn't raise if broadcast
         model.variables = {
-            c_n.name: pybamm.PrimaryBroadcast(pybamm.Scalar(2), ["negative electrode"])
+            c_n.name: pybamm.PrimaryBroadcast(
+                pybamm.InputParameter("a"), ["negative electrode"]
+            )
         }
         disc.process_model(model)
 

tests/unit/test_expression_tree/test_binary_operators.py

@@ -45,6 +45,10 @@ def test_addition(self):
         self.assertEqual(summ.children[0].name, a.name)
         self.assertEqual(summ.children[1].name, b.name)
 
+        # test simplifying
+        summ2 = pybamm.Scalar(1) + pybamm.Scalar(3)
+        self.assertEqual(summ2.id, pybamm.Scalar(4).id)
+
     def test_power(self):
         a = pybamm.Symbol("a")
         b = pybamm.Symbol("b")
@@ -61,22 +65,28 @@ def test_power(self):
     def test_known_eval(self):
         # Scalars
         a = pybamm.Scalar(4)
-        b = pybamm.Scalar(2)
+        b = pybamm.StateVector(slice(0, 1))
         expr = (a + b) - (a + b) * (a + b)
-        value = expr.evaluate()
-        self.assertEqual(expr.evaluate(known_evals={})[0], value)
-        self.assertIn((a + b).id, expr.evaluate(known_evals={})[1])
-        self.assertEqual(expr.evaluate(known_evals={})[1][(a + b).id], 6)
+        value = expr.evaluate(y=np.array([2]))
+        self.assertEqual(expr.evaluate(y=np.array([2]), known_evals={})[0], value)
+        self.assertIn((a + b).id, expr.evaluate(y=np.array([2]), known_evals={})[1])
+        self.assertEqual(
+            expr.evaluate(y=np.array([2]), known_evals={})[1][(a + b).id], 6
+        )
 
         # Matrices
         a = pybamm.Matrix(np.random.rand(5, 5))
-        b = pybamm.Matrix(np.random.rand(5, 5))
+        b = pybamm.StateVector(slice(0, 5))
         expr2 = (a @ b) - (a @ b) * (a @ b) + (a @ b)
-        value = expr2.evaluate()
-        np.testing.assert_array_equal(expr2.evaluate(known_evals={})[0], value)
-        self.assertIn((a @ b).id, expr2.evaluate(known_evals={})[1])
+        y_test = np.linspace(0, 1, 5)
+        value = expr2.evaluate(y=y_test)
+        np.testing.assert_array_equal(
+            expr2.evaluate(y=y_test, known_evals={})[0], value
+        )
+        self.assertIn((a @ b).id, expr2.evaluate(y=y_test, known_evals={})[1])
         np.testing.assert_array_equal(
-            expr2.evaluate(known_evals={})[1][(a @ b).id], (a @ b).evaluate()
+            expr2.evaluate(y=y_test, known_evals={})[1][(a @ b).id],
+            (a @ b).evaluate(y=y_test),
         )
 
     def test_diff(self):
@@ -158,7 +168,7 @@ def test_id(self):
     def test_number_overloading(self):
         a = pybamm.Scalar(4)
         prod = a * 3
-        self.assertIsInstance(prod.children[1], pybamm.Scalar)
+        self.assertIsInstance(prod, pybamm.Scalar)
         self.assertEqual(prod.evaluate(), 12)
 
     def test_sparse_multiply(self):

tests/unit/test_expression_tree/test_symbol.py

@@ -194,7 +194,7 @@ def test_symbol_evaluates_to_number(self):
         a = pybamm.Parameter("a")
         self.assertFalse(a.evaluates_to_number())
 
-        a = pybamm.Scalar(3) * pybamm.Scalar(2)
+        a = pybamm.Scalar(3) * pybamm.Time()
         self.assertTrue(a.evaluates_to_number())
         # highlight difference between this function and isinstance(a, Scalar)
         self.assertNotIsInstance(a, pybamm.Scalar)
@@ -339,10 +339,10 @@ def test_has_spatial_derivatives(self):
 
     def test_orphans(self):
         a = pybamm.Scalar(1)
-        b = pybamm.Scalar(2)
-        sum = a + b
+        b = pybamm.Parameter("b")
+        summ = a + b
 
-        a_orp, b_orp = sum.orphans
+        a_orp, b_orp = summ.orphans
         self.assertIsNone(a_orp.parent)
         self.assertIsNone(b_orp.parent)
         self.assertEqual(a.id, a_orp.id)

tests/unit/test_expression_tree/test_unary_operators.py

@@ -124,22 +124,23 @@ def test_integral(self):
             pybamm.Integral(a, y)
 
     def test_index(self):
-        vec = pybamm.Vector(np.array([1, 2, 3, 4, 5]))
+        vec = pybamm.StateVector(slice(0, 5))
+        y_test = np.array([1, 2, 3, 4, 5])
         # with integer
         ind = vec[3]
         self.assertIsInstance(ind, pybamm.Index)
         self.assertEqual(ind.slice, slice(3, 4))
-        self.assertEqual(ind.evaluate(), 4)
+        self.assertEqual(ind.evaluate(y=y_test), 4)
         # with slice
         ind = vec[1:3]
         self.assertIsInstance(ind, pybamm.Index)
         self.assertEqual(ind.slice, slice(1, 3))
-        np.testing.assert_array_equal(ind.evaluate(), np.array([[2], [3]]))
+        np.testing.assert_array_equal(ind.evaluate(y=y_test), np.array([[2], [3]]))
         # with only stop slice
         ind = vec[:3]
         self.assertIsInstance(ind, pybamm.Index)
         self.assertEqual(ind.slice, slice(3))
-        np.testing.assert_array_equal(ind.evaluate(), np.array([[1], [2], [3]]))
+        np.testing.assert_array_equal(ind.evaluate(y=y_test), np.array([[1], [2], [3]]))
 
         # errors
         with self.assertRaisesRegex(TypeError, "index must be integer or slice"):