Repository: pybamm-team/PyBaMM

Issue:
simplify A - A = 0 and A + A = 2A

Discussion:

Patch:
diff --git a/pybamm/expression_tree/simplify.py b/pybamm/expression_tree/simplify.py
index 602fbeba4..e72a03708 100644
--- a/pybamm/expression_tree/simplify.py
+++ b/pybamm/expression_tree/simplify.py
@@ -139,6 +139,36 @@ def simplify_addition_subtraction(myclass, left, right):
                     ret -= child
         return ret
 
+    # simplify identical terms
+    i = 0
+    while i < len(numerator) - 1:
+        if isinstance(numerator[i], pybamm.Multiplication) and \
+                isinstance(numerator[i].children[0], pybamm.Scalar):
+            term_i = numerator[i].orphans[1]
+            term_i_count = numerator[i].children[0].evaluate()
+        else:
+            term_i = numerator[i]
+            term_i_count = 1
+
+        # loop through rest of numerator counting up and deleting identical terms
+        for j, (term_j, typ_j) in \
+                enumerate(zip(numerator[i + 1:], numerator_types[i + 1:])):
+            if term_i.id == term_j.id:
+                if typ_j == pybamm.Addition:
+                    term_i_count += 1
+                elif typ_j == pybamm.Subtraction:
+                    term_i_count -= 1
+                del numerator[j + i + 1]
+                del numerator_types[j + i + 1]
+
+        # replace this term by count * term if count > 1
+        if term_i_count != 1:
+            # simplify the result just in case
+            # (e.g. count == 0, or can fold constant into the term)
+            numerator[i] = (term_i_count * term_i).simplify()
+
+        i += 1
+
     # can reorder the numerator
     (constant, nonconstant, constant_types, nonconstant_types) = partition_by_constant(
         numerator, numerator_types


Test patch:
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index dde29c087..ca263d9a6 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -148,8 +148,9 @@ class TestSimplify(unittest.TestCase):
         self.assertIsInstance(expr, pybamm.Addition)
         self.assertIsInstance(expr.children[0], pybamm.Scalar)
         self.assertEqual(expr.children[0].evaluate(), 4.0)
-        self.assertIsInstance(expr.children[1], pybamm.Addition)
-        self.assertIsInstance(expr.children[1].children[0], pybamm.Parameter)
+        self.assertIsInstance(expr.children[1], pybamm.Multiplication)
+        self.assertIsInstance(expr.children[1].children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[1].children[0].evaluate(), 2)
         self.assertIsInstance(expr.children[1].children[1], pybamm.Parameter)
 
         expr = ((-1 + c) - (c + 1) + (c - 1)).simplify()
@@ -210,6 +211,31 @@ class TestSimplify(unittest.TestCase):
         with self.assertRaises(NotImplementedError):
             sym.simplify()
 
+        # A + A = 2A (#323)
+        a = pybamm.Variable('A')
+        expr = (a + a).simplify()
+        self.assertIsInstance(expr, pybamm.Multiplication)
+        self.assertIsInstance(expr.children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[0].evaluate(), 2)
+        self.assertIsInstance(expr.children[1], pybamm.Variable)
+
+        expr = (a + a + a + a).simplify()
+        self.assertIsInstance(expr, pybamm.Multiplication)
+        self.assertIsInstance(expr.children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[0].evaluate(), 4)
+        self.assertIsInstance(expr.children[1], pybamm.Variable)
+
+        expr = (a - a + a - a + a + a).simplify()
+        self.assertIsInstance(expr, pybamm.Multiplication)
+        self.assertIsInstance(expr.children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[0].evaluate(), 2)
+        self.assertIsInstance(expr.children[1], pybamm.Variable)
+
+        # A - A = 0 (#323)
+        expr = (a - a).simplify()
+        self.assertIsInstance(expr, pybamm.Scalar)
+        self.assertEqual(expr.evaluate(), 0)
+
     def test_function_simplify(self):
         a = pybamm.Parameter("a")
         funca = pybamm.Function(test_const_function, a).simplify()
@@ -377,6 +403,25 @@ class TestSimplify(unittest.TestCase):
                 expr.evaluate(y=np.ones(300)), exprsimp.evaluate(y=np.ones(300))
             )
 
+        # A + A = 2A (#323)
+        a = pybamm.StateVector(slice(0, 300))
+        expr = (a + a).simplify()
+        self.assertIsInstance(expr, pybamm.Multiplication)
+        self.assertIsInstance(expr.children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[0].evaluate(), 2)
+        self.assertIsInstance(expr.children[1], pybamm.StateVector)
+
+        expr = (a + a + a + a).simplify()
+        self.assertIsInstance(expr, pybamm.Multiplication)
+        self.assertIsInstance(expr.children[0], pybamm.Scalar)
+        self.assertEqual(expr.children[0].evaluate(), 4)
+        self.assertIsInstance(expr.children[1], pybamm.StateVector)
+
+        # A - A = 0 (#323)
+        expr = (a - a).simplify()
+        self.assertIsInstance(expr, pybamm.Scalar)
+        self.assertEqual(expr.evaluate(), 0)
+
     def test_domain_concatenation_simplify(self):
         # create discretisation
         disc = get_discretisation_for_testing()