Repository: pybamm-team/PyBaMM

Issue:
Matrix-simplify bug
The following code gives a `ValueError` (shape mismatch):
```python3
m2 = pybamm.Matrix(np.ones((299, 300)))
v2 = pybamm.StateVector(slice(0, 300))
v3 = pybamm.Vector(np.ones(299))
expr = (m2 @ v2) * v3
expr.evaluate(y=np.ones(300)) # fine
expr.simplify().evaluate(y=np.ones(300)) # error (expr.simplify() == m2 @ (v2 * v3))
```
Note that 
```python3
expr = v3 * (m2 @ v2)
expr.evaluate(y=np.ones(300)) # fine
expr.simplify().evaluate(y=np.ones(300)) # fine (expr.simplify() == (v3 * m2) @ v2)
```
is fine. @martinjrobins could you have a look at this?

Discussion:

Patch:
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index c7dfda40a..b2ce2eb8e 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -236,6 +236,15 @@ def simplify_multiplication_division(myclass, left, right):
                 or isinstance(other_child, (pybamm.Scalar, pybamm.Vector))
             ):
                 left, right = child.orphans
+                if child == left_child and this_class == pybamm.Multiplication:
+                    # change (m @ v1) * v2 -> v2 * m @ v so can simplify correctly
+                    # (#341)
+                    numerator.append(other_child)
+                    numerator_types.append(previous_class)
+                    flatten(
+                        this_class, child.__class__, left, right, in_numerator, False
+                    )
+                    break
                 if child == left_child:
                     flatten(
                         previous_class, child.__class__, left, right, in_numerator, True


Test patch:
diff --git a/tests/unit/test_expression_tree/test_matrix.py b/tests/unit/test_expression_tree/test_matrix.py
index dbaf50f8e..41c1de5b3 100644
--- a/tests/unit/test_expression_tree/test_matrix.py
+++ b/tests/unit/test_expression_tree/test_matrix.py
@@ -179,6 +179,19 @@ class TestMatrix(unittest.TestCase):
         expr4simp = expr4.simplify()
         self.assertEqual(expr4.id, expr4simp.id)
 
+        m2 = pybamm.Matrix(np.ones((299, 300)))
+        v2 = pybamm.StateVector(slice(0, 300))
+        v3 = pybamm.Vector(np.ones(299))
+        exprs = [(m2 @ v2) * v3, (m2 @ v2) / v3]
+        for expr in exprs:
+            exprsimp = expr.simplify()
+            self.assertIsInstance(exprsimp, pybamm.MatrixMultiplication)
+            self.assertIsInstance(exprsimp.children[0], pybamm.Matrix)
+            self.assertIsInstance(exprsimp.children[1], pybamm.StateVector)
+            np.testing.assert_array_equal(
+                expr.evaluate(y=np.ones(300)), exprsimp.evaluate(y=np.ones(300))
+            )
+
     def test_matrix_modification(self):
         exp = self.mat @ self.mat + self.mat
         self.A[0, 0] = -1