Repository: pybamm-team/PyBaMM

Issue:
only do discretisation of gradient and divergence once
for an expression with more than one `grad(c)` or `div(c)`, the BaseDiscretisation will discretise each of them individually. We should identify these sub-expressions and replace them all with the same Matrix

note that identifying sub-expressions is also in #55, so this will overlap a bit

Discussion:

Patch:
diff --git a/docs/source/expression_tree/simplify.rst b/docs/source/expression_tree/simplify.rst
index 8e2ab29c8..cfbf11e25 100644
--- a/docs/source/expression_tree/simplify.rst
+++ b/docs/source/expression_tree/simplify.rst
@@ -1,7 +1,8 @@
 Simplify
 ========
 
-.. autofunction:: pybamm.simplify
+.. autoclass:: pybamm.Simplification
+  :members:
 
 .. autofunction:: pybamm.simplify_if_constant
 
diff --git a/examples/notebooks/spatial_methods/finite-volumes.ipynb b/examples/notebooks/spatial_methods/finite-volumes.ipynb
index 51c71d0b1..2c6b8c1da 100644
--- a/examples/notebooks/spatial_methods/finite-volumes.ipynb
+++ b/examples/notebooks/spatial_methods/finite-volumes.ipynb
@@ -630,7 +630,7 @@
     }
    ],
    "source": [
-    "disc._bcs = {u.id: {\"left\": (pybamm.Scalar(1), \"Dirichlet\"), \"right\": (pybamm.Scalar(2), \"Dirichlet\")}}\n",
+    "disc.bcs = {u.id: {\"left\": (pybamm.Scalar(1), \"Dirichlet\"), \"right\": (pybamm.Scalar(2), \"Dirichlet\")}}\n",
     "grad_u_disc = disc.process_symbol(grad_u)\n",
     "print(\"The gradient object is:\")\n",
     "(grad_u_disc.render())\n",
@@ -681,7 +681,7 @@
     }
    ],
    "source": [
-    "disc._bcs = {u.id: {\"left\": (pybamm.Scalar(3), \"Neumann\"), \"right\": (pybamm.Scalar(4), \"Neumann\")}}\n",
+    "disc.bcs = {u.id: {\"left\": (pybamm.Scalar(3), \"Neumann\"), \"right\": (pybamm.Scalar(4), \"Neumann\")}}\n",
     "grad_u_disc = disc.process_symbol(grad_u)\n",
     "print(\"The gradient object is:\")\n",
     "(grad_u_disc.render())\n",
@@ -729,7 +729,7 @@
     }
    ],
    "source": [
-    "disc._bcs = {u.id: {\"left\": (pybamm.Scalar(5), \"Dirichlet\"), \"right\": (pybamm.Scalar(6), \"Neumann\")}}\n",
+    "disc.bcs = {u.id: {\"left\": (pybamm.Scalar(5), \"Dirichlet\"), \"right\": (pybamm.Scalar(6), \"Neumann\")}}\n",
     "grad_u_disc = disc.process_symbol(grad_u)\n",
     "print(\"The gradient object is:\")\n",
     "(grad_u_disc.render())\n",
@@ -766,7 +766,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "disc._bcs = {u.id: {\"left\": (pybamm.Scalar(-1), \"Neumann\"), \"right\": (pybamm.Scalar(1), \"Neumann\")}}"
+    "disc.bcs = {u.id: {\"left\": (pybamm.Scalar(-1), \"Neumann\"), \"right\": (pybamm.Scalar(1), \"Neumann\")}}"
    ]
   },
   {
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 4e51a4ed0..58138b8d0 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -108,7 +108,7 @@ from .expression_tree.exceptions import (
     ModelWarning,
 )
 from .expression_tree.simplify import (
-    simplify,
+    Simplification,
     simplify_if_constant,
     simplify_addition_subtraction,
     simplify_multiplication_division,
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index ac8857417..587a739d5 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -34,13 +34,24 @@ class Discretisation(object):
             self._spatial_methods = {
                 dom: method(mesh) for dom, method in spatial_methods.items()
             }
-        self._bcs = {}
+        self.bcs = {}
         self._y_slices = {}
+        self._discretised_symbols = {}
 
     @property
     def mesh(self):
         return self._mesh
 
+    @property
+    def bcs(self):
+        return self._bcs
+
+    @bcs.setter
+    def bcs(self, value):
+        self._bcs = value
+        # reset discretised_symbols
+        self._discretised_symbols = {}
+
     def process_model(self, model, inplace=True):
         """Discretise a model.
         Currently inplace, could be changed to return a new model.
@@ -72,7 +83,7 @@ class Discretisation(object):
         self.set_variable_slices(variables)
 
         # set boundary conditions (only need key ids for boundary_conditions)
-        self._bcs = self.process_boundary_conditions(model)
+        self.bcs = self.process_boundary_conditions(model)
 
         # set up inplace vs not inplace
         if inplace:
@@ -149,6 +160,9 @@ class Discretisation(object):
             """y_slices should be dict, not {}""".format(type(self._y_slices))
         )
 
+        # reset discretised_symbols
+        self._discretised_symbols = {}
+
     def process_initial_conditions(self, model):
         """Discretise model initial_conditions.
 
@@ -287,7 +301,7 @@ class Discretisation(object):
             else:
                 mass_list.append(
                     self._spatial_methods[var.domain[0]]
-                    .mass_matrix(var, self._bcs)
+                    .mass_matrix(var, self.bcs)
                     .entries
                 )
 
@@ -341,6 +355,7 @@ class Discretisation(object):
 
     def process_symbol(self, symbol):
         """Discretise operators in model equations.
+        If a symbol has already been discretised, the stored value is returned.
 
         Parameters
         ----------
@@ -354,6 +369,15 @@ class Discretisation(object):
 
         """
         pybamm.logger.debug("Discretise {!s}".format(symbol))
+        try:
+            return self._discretised_symbols[symbol.id]
+        except KeyError:
+            discretised_symbol = self._process_symbol(symbol)
+            self._discretised_symbols[symbol.id] = discretised_symbol
+            return discretised_symbol
+
+    def _process_symbol(self, symbol):
+        """ See :meth:`Discretisation.process_symbol()`. """
 
         if symbol.domain != []:
             spatial_method = self._spatial_methods[symbol.domain[0]]
@@ -375,10 +399,10 @@ class Discretisation(object):
             if child.domain != []:
                 child_spatial_method = self._spatial_methods[child.domain[0]]
             if isinstance(symbol, pybamm.Gradient):
-                return child_spatial_method.gradient(child, disc_child, self._bcs)
+                return child_spatial_method.gradient(child, disc_child, self.bcs)
 
             elif isinstance(symbol, pybamm.Divergence):
-                return child_spatial_method.divergence(child, disc_child, self._bcs)
+                return child_spatial_method.divergence(child, disc_child, self.bcs)
 
             elif isinstance(symbol, pybamm.IndefiniteIntegral):
                 return child_spatial_method.indefinite_integral(
diff --git a/pybamm/expression_tree/parameter.py b/pybamm/expression_tree/parameter.py
index 814918105..9e373f842 100644
--- a/pybamm/expression_tree/parameter.py
+++ b/pybamm/expression_tree/parameter.py
@@ -49,9 +49,17 @@ class FunctionParameter(pybamm.UnaryOperator):
     """
 
     def __init__(self, name, child, diff_variable=None):
-        super().__init__(name, child)
         # assign diff variable
         self.diff_variable = diff_variable
+        super().__init__(name, child)
+
+    def set_id(self):
+        """See :meth:`pybamm.Symbol.set_id` """
+        self._id = hash(
+            (self.__class__, self.name, self.diff_variable)
+            + tuple([child.id for child in self.children])
+            + tuple(self.domain)
+        )
 
     def diff(self, variable):
         """ See :meth:`pybamm.Symbol.diff()`. """
diff --git a/pybamm/expression_tree/simplify.py b/pybamm/expression_tree/simplify.py
index 9c0df7f71..b5105cb60 100644
--- a/pybamm/expression_tree/simplify.py
+++ b/pybamm/expression_tree/simplify.py
@@ -77,7 +77,7 @@ def simplify_addition_subtraction(myclass, left, right):
         1 - (2 + 3)       -> [1, 2, 3]    and [None, Subtraction, Subtraction]
         (1 + 2) - (2 + 3) -> [1, 2, 2, 3] and [None, Addition, Subtraction, Subtraction]
         """
-        for child in [left_child, right_child]:
+        for side, child in [("left", left_child), ("right", right_child)]:
             if isinstance(child, (pybamm.Addition, pybamm.Subtraction)):
                 left, right = child.orphans
                 flatten(child.__class__, left, right, in_subtraction)
@@ -91,7 +91,7 @@ def simplify_addition_subtraction(myclass, left, right):
                 else:
                     numerator_types.append(pybamm.Addition)
 
-            if child == left_child:
+            if side == "left":
                 if in_subtraction is None:
                     in_subtraction = this_class == pybamm.Subtraction
                 elif this_class == pybamm.Subtraction:
@@ -270,9 +270,9 @@ def simplify_multiplication_division(myclass, left, right):
         (1 @ 2) / 3 ->  [1, 2]       [3]       [None, MatrixMultiplication]
         1 / (c / 2) ->  [1, 2]       [c]       [None, Multiplication]
         """
-        for child in [left_child, right_child]:
+        for side, child in [("left", left_child), ("right", right_child)]:
 
-            if child == left_child:
+            if side == "left":
                 other_child = right_child
             else:
                 other_child = left_child
@@ -286,7 +286,7 @@ def simplify_multiplication_division(myclass, left, right):
             ):
                 left, right = child.orphans
                 if (
-                    child == left_child
+                    side == "left"
                     and this_class == pybamm.Multiplication
                     and isinstance(other_child, pybamm.Vector)
                 ):
@@ -298,7 +298,7 @@ def simplify_multiplication_division(myclass, left, right):
                         this_class, child.__class__, left, right, in_numerator, True
                     )
                     break
-                if child == left_child:
+                if side == "left":
                     flatten(
                         previous_class, child.__class__, left, right, in_numerator, True
                     )
@@ -312,7 +312,7 @@ def simplify_multiplication_division(myclass, left, right):
                 and not in_matrix_multiplication
             ):
                 left, right = child.orphans
-                if child == left_child:
+                if side == "left":
                     flatten(
                         previous_class,
                         child.__class__,
@@ -329,18 +329,18 @@ def simplify_multiplication_division(myclass, left, right):
             else:
                 if in_numerator:
                     numerator.append(child)
-                    if child == left_child:
+                    if side == "left":
                         numerator_types.append(previous_class)
                     else:
                         numerator_types.append(this_class)
                 else:
                     denominator.append(child)
-                    if child == left_child:
+                    if side == "left":
                         denominator_types.append(previous_class)
                     else:
                         denominator_types.append(this_class)
 
-            if child == left_child and this_class == pybamm.Division:
+            if side == "left" and this_class == pybamm.Division:
                 in_numerator = not in_numerator
 
     flatten(None, myclass, left, right, True, myclass == pybamm.MatrixMultiplication)
@@ -532,52 +532,66 @@ def simplify_multiplication_division(myclass, left, right):
     return result
 
 
-def simplify(symbol):
-    """
-    This function recurses down the tree, applying any simplifications defined in
-    classes derived from pybamm.Symbol. E.g. any expression multiplied by a
-    pybamm.Scalar(0) will be simplified to a pybamm.Scalar(0)
+class Simplification(object):
+    def __init__(self, simplified_symbols=None):
+        self._simplified_symbols = simplified_symbols or {}
 
-    Parameters
-    ----------
-    symbol : :class:`pybamm.Symbol`
+    def simplify(self, symbol):
+        """
+        This function recurses down the tree, applying any simplifications defined in
+        classes derived from pybamm.Symbol. E.g. any expression multiplied by a
+        pybamm.Scalar(0) will be simplified to a pybamm.Scalar(0).
+        If a symbol has already been simplified, the stored value is returned.
+
+        Parameters
+        ----------
+        symbol : :class:`pybamm.Symbol`
         The symbol to simplify
 
-    Returns
-    -------
-    :class:`pybamm.Symbol`
+        Returns
+        -------
+        :class:`pybamm.Symbol`
         Simplified symbol
-    """
-    pybamm.logger.debug("Simplify {!s}".format(symbol))
-
-    if isinstance(symbol, pybamm.BinaryOperator):
-        left, right = symbol.children
-        # process children
-        new_left = simplify(left)
-        new_right = simplify(right)
-        # make new symbol, ensure domain remains the same
-        # _binary_simplify defined in derived classes for specific rules
-        new_symbol = symbol._binary_simplify(new_left, new_right)
-        new_symbol.domain = symbol.domain
-        return simplify_if_constant(new_symbol)
-
-    elif isinstance(symbol, pybamm.UnaryOperator):
-        new_child = simplify(symbol.child)
-        new_symbol = symbol._unary_simplify(new_child)
-        new_symbol.domain = symbol.domain
-        return simplify_if_constant(new_symbol)
-
-    elif isinstance(symbol, pybamm.Concatenation):
-        new_children = [simplify(child) for child in symbol.cached_children]
-        new_symbol = symbol._concatenation_simplify(new_children)
-
-        return simplify_if_constant(new_symbol)
-
-    else:
-        # Backup option: return new copy of the object
+        """
+        pybamm.logger.debug("Simplify {!s}".format(symbol))
         try:
-            return symbol.new_copy()
-        except NotImplementedError:
-            raise NotImplementedError(
-                "Cannot simplify symbol of type '{}'".format(type(symbol))
-            )
+            return self._simplified_symbols[symbol.id]
+        except KeyError:
+            simplified_symbol = self._simplify(symbol)
+            self._simplified_symbols[symbol.id] = simplified_symbol
+            return simplified_symbol
+
+    def _simplify(self, symbol):
+        """ See :meth:`Simplification.simplify()`. """
+
+        if isinstance(symbol, pybamm.BinaryOperator):
+            left, right = symbol.children
+            # process children
+            new_left = self.simplify(left)
+            new_right = self.simplify(right)
+            # make new symbol, ensure domain remains the same
+            # _binary_simplify defined in derived classes for specific rules
+            new_symbol = symbol._binary_simplify(new_left, new_right)
+            new_symbol.domain = symbol.domain
+            return simplify_if_constant(new_symbol)
+
+        elif isinstance(symbol, pybamm.UnaryOperator):
+            new_child = self.simplify(symbol.child)
+            new_symbol = symbol._unary_simplify(new_child)
+            new_symbol.domain = symbol.domain
+            return simplify_if_constant(new_symbol)
+
+        elif isinstance(symbol, pybamm.Concatenation):
+            new_children = [self.simplify(child) for child in symbol.children]
+            new_symbol = symbol._concatenation_simplify(new_children)
+
+            return simplify_if_constant(new_symbol)
+
+        else:
+            # Backup option: return new copy of the object
+            try:
+                return symbol.new_copy()
+            except NotImplementedError:
+                raise NotImplementedError(
+                    "Cannot simplify symbol of type '{}'".format(type(symbol))
+                )
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 6b63a493c..1fbfbdc54 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -467,9 +467,9 @@ class Symbol(anytree.NodeMixin):
         """Returns True if equation has a Divergence term."""
         return any(isinstance(symbol, pybamm.Divergence) for symbol in self.pre_order())
 
-    def simplify(self):
-        """ Simplify the expression tree. See :meth:`pybamm.simplify()`. """
-        return pybamm.simplify(self)
+    def simplify(self, simplified_symbols=None):
+        """ Simplify the expression tree. See :class:`pybamm.Simplification`. """
+        return pybamm.Simplification(simplified_symbols).simplify(self)
 
     def new_copy(self):
         """
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index f1b6c7243..d486fa6a1 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -26,8 +26,8 @@ class UnaryOperator(pybamm.Symbol):
     """
 
     def __init__(self, name, child):
+        self.child = child
         super().__init__(name, children=[child], domain=child.domain)
-        self.child = self.children[0]
 
     def __str__(self):
         """ See :meth:`pybamm.Symbol.__str__()`. """
@@ -139,8 +139,8 @@ class Function(UnaryOperator):
 
     def __init__(self, func, child):
         """ See :meth:`pybamm.UnaryOperator.__init__()`. """
-        super().__init__("function ({})".format(func.__name__), child)
         self.func = func
+        super().__init__("function ({})".format(func.__name__), child)
         # hack to work out whether function takes any params
         # (signature doesn't work for numpy)
         if isinstance(func, np.ufunc):
@@ -207,8 +207,14 @@ class Index(UnaryOperator):
     def __init__(self, child, index, name=None):
         if name is None:
             name = child.name + "[" + str(index) + "]"
-        super().__init__(name, child)
         self.index = index
+        super().__init__(name, child)
+
+    def set_id(self):
+        """ See :meth:`pybamm.Symbol.set_id()` """
+        self._id = hash(
+            (self.__class__, self.name, self.index, self.child.id) + tuple(self.domain)
+        )
 
     def _unary_evaluate(self, child):
         """ See :meth:`UnaryOperator._unary_evaluate()`. """
@@ -324,8 +330,8 @@ class Integral(SpatialOperator):
         name = "integral d{}".format(integration_variable.name)
         if isinstance(integration_variable, pybamm.SpatialVariable):
             name += " {}".format(integration_variable.domain)
-        super().__init__(name, child)
         self._integration_variable = integration_variable
+        super().__init__(name, child)
         # integrating removes the domain
         self.domain = []
 
@@ -333,6 +339,13 @@ class Integral(SpatialOperator):
     def integration_variable(self):
         return self._integration_variable
 
+    def set_id(self):
+        """ See :meth:`pybamm.Symbol.set_id()` """
+        self._id = hash(
+            (self.__class__, self.name, self.integration_variable.id, self.child.id)
+            + tuple(self.domain)
+        )
+
     def _unary_simplify(self, simplified_child):
         """ See :meth:`UnaryOperator._unary_simplify()`. """
 
@@ -391,12 +404,18 @@ class BoundaryOperator(SpatialOperator):
     """
 
     def __init__(self, name, child, side):
-        super().__init__(name, child)
         self.side = side
         # Domain of Boundary must be ([]) so that expressions can be formed
         # of boundary values of variables in different domains
+        super().__init__(name, child)
         self.domain = []
 
+    def set_id(self):
+        """ See :meth:`pybamm.Symbol.set_id()` """
+        self._id = hash(
+            (self.__class__, self.name, self.side, self.child.id) + tuple(self.domain)
+        )
+
     def _unary_simplify(self, simplified_child):
         """ See :meth:`UnaryOperator._unary_simplify()`. """
         return self.__class__(simplified_child, self.side)
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index ec2aa08d9..24840dc91 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -73,14 +73,21 @@ class StateVector(pybamm.Symbol):
                 name = "y[:{:d}]".format(y_slice.stop)
             else:
                 name = "y[{:d}:{:d}]".format(y_slice.start, y_slice.stop)
-        super().__init__(name=name, domain=domain)
         self._y_slice = y_slice
+        super().__init__(name=name, domain=domain)
 
     @property
     def y_slice(self):
         """Slice of an external y to read"""
         return self._y_slice
 
+    def set_id(self):
+        """ See :meth:`pybamm.Symbol.set_id()` """
+        self._id = hash(
+            (self.__class__, self.name, self.y_slice.start, self.y_slice.stop)
+            + tuple(self.domain)
+        )
+
     def _base_evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol._base_evaluate()`. """
         if y is None:
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index 153855482..05a551580 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -38,6 +38,8 @@ class ParameterValues(dict):
         # doing parameter studies
         self.update(optional_parameters)
 
+        self._processed_symbols = {}
+
     def read_parameters_csv(self, filename):
         """Reads parameters from csv file into dict.
 
@@ -57,6 +59,12 @@ class ParameterValues(dict):
         df.dropna(how="all", inplace=True)
         return {k: v for (k, v) in zip(df["Name [units]"], df["Value"])}
 
+    def update(self, values):
+        for k, v in values.items():
+            self[k] = v
+        # reset processed symbols
+        self._processed_symbols = {}
+
     def process_model(self, model, processing="process"):
         """Assign parameter values to a model.
         Currently inplace, could be changed to return a new model.
@@ -152,6 +160,7 @@ class ParameterValues(dict):
 
     def process_symbol(self, symbol):
         """Walk through the symbol and replace any Parameter with a Value.
+        If a symbol has already been processed, the stored value is returned.
 
         Parameters
         ----------
@@ -165,6 +174,15 @@ class ParameterValues(dict):
 
         """
         pybamm.logger.debug("Set parameters for {!s}".format(symbol))
+        try:
+            return self._processed_symbols[symbol.id]
+        except KeyError:
+            processed_symbol = self._process_symbol(symbol)
+            self._processed_symbols[symbol.id] = processed_symbol
+            return processed_symbol
+
+    def _process_symbol(self, symbol):
+        """ See :meth:`ParameterValues.process_symbol()`. """
 
         if isinstance(symbol, pybamm.Parameter):
             value = self[symbol.name]
diff --git a/pybamm/solvers/dae_solver.py b/pybamm/solvers/dae_solver.py
index 1c03a648b..1d17ff7e1 100644
--- a/pybamm/solvers/dae_solver.py
+++ b/pybamm/solvers/dae_solver.py
@@ -59,9 +59,11 @@ class DaeSolver(pybamm.BaseSolver):
         concatenated_algebraic = model.concatenated_algebraic
         events = model.events
         if model.use_simplify:
-            concatenated_rhs = concatenated_rhs.simplify()
-            concatenated_algebraic = concatenated_algebraic.simplify()
-            events = [event.simplify() for event in events]
+            # set up simplification object, for re-use of dict
+            simp = pybamm.Simplification()
+            concatenated_rhs = simp.simplify(concatenated_rhs)
+            concatenated_algebraic = simp.simplify(concatenated_algebraic)
+            events = [simp.simplify(event) for event in events]
 
         def residuals(t, y, ydot):
             pybamm.logger.debug(
@@ -101,8 +103,8 @@ class DaeSolver(pybamm.BaseSolver):
             jac_rhs = concatenated_rhs.jac(y)
             jac_algebraic = concatenated_algebraic.jac(y)
             if model.use_simplify:
-                jac_rhs = jac_rhs.simplify()
-                jac_algebraic = jac_algebraic.simplify()
+                jac_rhs = simp.simplify(jac_rhs)
+                jac_algebraic = simp.simplify(jac_algebraic)
 
             jac = pybamm.SparseStack(jac_rhs, jac_algebraic)
 
diff --git a/pybamm/solvers/ode_solver.py b/pybamm/solvers/ode_solver.py
index 38e93010c..ee11d3971 100644
--- a/pybamm/solvers/ode_solver.py
+++ b/pybamm/solvers/ode_solver.py
@@ -34,9 +34,11 @@ class OdeSolver(pybamm.BaseSolver):
         concatenated_rhs = model.concatenated_rhs
         events = model.events
         if model.use_simplify:
+            # set up simplification object, for re-use of dict
+            simp = pybamm.Simplification()
             # create simplified rhs and event expressions
-            concatenated_rhs = concatenated_rhs.simplify()
-            events = [event.simplify() for event in events]
+            concatenated_rhs = simp.simplify(concatenated_rhs)
+            events = [simp.simplify(event) for event in events]
 
         def dydt(t, y):
             pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
@@ -60,7 +62,7 @@ class OdeSolver(pybamm.BaseSolver):
             y = pybamm.StateVector(slice(0, np.size(y0)))
             jac_rhs = concatenated_rhs.jac(y)
             if model.use_simplify:
-                jac_rhs = jac_rhs.simplify()
+                jac_rhs = simp.simplify(jac_rhs)
 
             def jacobian(t, y):
                 return jac_rhs.evaluate(t, y, known_evals={})[0]


Test patch:
diff --git a/tests/integration/test_spatial_methods/test_finite_volume.py b/tests/integration/test_spatial_methods/test_finite_volume.py
index 99f52d42c..7392f18f1 100644
--- a/tests/integration/test_spatial_methods/test_finite_volume.py
+++ b/tests/integration/test_spatial_methods/test_finite_volume.py
@@ -29,7 +29,7 @@ class TestFiniteVolumeConvergence(unittest.TestCase):
             # create mesh and discretisation
             mesh = get_mesh_for_testing(n)
             disc = pybamm.Discretisation(mesh, spatial_methods)
-            disc._bcs = boundary_conditions
+            disc.bcs = boundary_conditions
             disc.set_variable_slices([var])
 
             # Define exact solutions
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index 5d35f43f3..77d094613 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -274,7 +274,7 @@ class TestDiscretise(unittest.TestCase):
         combined_submesh = mesh.combine_submeshes(*whole_cell)
 
         y = combined_submesh[0].nodes[:, np.newaxis] ** 2
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         disc.set_variable_slices(list(rhs.keys()))
         # rhs - grad and div are identity operators here
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index 5e1986cc9..ed95137bc 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -290,7 +290,7 @@ class TestSimplify(unittest.TestCase):
                 expr.children[0].entries, np.array([[3, 0], [0, 3]])
             )
 
-        expr = (v @ v / 2).simplify()
+        expr = ((v @ v) / 2).simplify()
         self.assertIsInstance(expr, pybamm.Multiplication)
         self.assertIsInstance(expr.children[0], pybamm.Scalar)
         self.assertEqual(expr.children[0].evaluate(), 0.5)
diff --git a/tests/unit/test_expression_tree/test_unary_operators.py b/tests/unit/test_expression_tree/test_unary_operators.py
index ab415fd22..27c923645 100644
--- a/tests/unit/test_expression_tree/test_unary_operators.py
+++ b/tests/unit/test_expression_tree/test_unary_operators.py
@@ -108,6 +108,11 @@ class TestUnaryOperators(unittest.TestCase):
         with self.assertRaises(ValueError):
             pybamm.Integral(a, y)
 
+    def test_index(self):
+        vec = pybamm.Vector(np.array([1, 2, 3, 4, 5]))
+        ind = pybamm.Index(vec, 3)
+        self.assertEqual(ind.evaluate(), 4)
+
     def test_diff(self):
         a = pybamm.StateVector(slice(0, 1))
         y = np.array([5])
diff --git a/tests/unit/test_spatial_methods/test_finite_volume.py b/tests/unit/test_spatial_methods/test_finite_volume.py
index 18dfd4529..fcfd6f15f 100644
--- a/tests/unit/test_spatial_methods/test_finite_volume.py
+++ b/tests/unit/test_spatial_methods/test_finite_volume.py
@@ -125,7 +125,7 @@ class TestFiniteVolume(unittest.TestCase):
         ]:
             # Check that the equation can be evaluated in each case
             # Dirichlet
-            disc._bcs = {
+            disc.bcs = {
                 var.id: {
                     "left": (pybamm.Scalar(0), "Dirichlet"),
                     "right": (pybamm.Scalar(1), "Dirichlet"),
@@ -134,7 +134,7 @@ class TestFiniteVolume(unittest.TestCase):
             eqn_disc = disc.process_symbol(eqn)
             eqn_disc.evaluate(None, y_test)
             # Neumann
-            disc._bcs = {
+            disc.bcs = {
                 var.id: {
                     "left": (pybamm.Scalar(0), "Neumann"),
                     "right": (pybamm.Scalar(1), "Neumann"),
@@ -143,7 +143,7 @@ class TestFiniteVolume(unittest.TestCase):
             eqn_disc = disc.process_symbol(eqn)
             eqn_disc.evaluate(None, y_test)
             # One of each
-            disc._bcs = {
+            disc.bcs = {
                 var.id: {
                     "left": (pybamm.Scalar(0), "Dirichlet"),
                     "right": (pybamm.Scalar(1), "Neumann"),
@@ -151,7 +151,7 @@ class TestFiniteVolume(unittest.TestCase):
             }
             eqn_disc = disc.process_symbol(eqn)
             eqn_disc.evaluate(None, y_test)
-            disc._bcs = {
+            disc.bcs = {
                 var.id: {
                     "left": (pybamm.Scalar(0), "Neumann"),
                     "right": (pybamm.Scalar(1), "Dirichlet"),
@@ -350,7 +350,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         disc.set_variable_slices([var])
         grad_eqn_disc = disc.process_symbol(grad_eqn)
@@ -372,7 +372,7 @@ class TestFiniteVolume(unittest.TestCase):
             }
         }
 
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         grad_eqn_disc = disc.process_symbol(grad_eqn)
         np.testing.assert_array_almost_equal(
@@ -408,7 +408,7 @@ class TestFiniteVolume(unittest.TestCase):
             }
         }
 
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         disc.set_variable_slices([var])
         grad_eqn_disc = disc.process_symbol(grad_eqn)
@@ -425,7 +425,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         y_linear = combined_submesh[0].nodes
         grad_eqn_disc = disc.process_symbol(grad_eqn)
@@ -445,7 +445,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(6), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         div_eqn_disc = disc.process_symbol(div_eqn)
         np.testing.assert_array_almost_equal(
@@ -479,7 +479,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         disc.set_variable_slices([var])
         grad_eqn_disc = disc.process_symbol(grad_eqn)
@@ -505,7 +505,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(6), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         div_eqn_disc = disc.process_symbol(div_eqn)
         div_eval = div_eqn_disc.evaluate(None, const)
@@ -544,7 +544,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Neumann"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
         div_eqn_disc = disc.process_symbol(div_eqn)
 
         # Linear y should have laplacian zero
@@ -586,7 +586,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(0), "Neumann"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
         grad_eqn_disc = disc.process_symbol(grad_eqn)
         div_eqn_disc = disc.process_symbol(div_eqn)
 
@@ -607,7 +607,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
         grad_eqn_disc = disc.process_symbol(grad_eqn)
         div_eqn_disc = disc.process_symbol(div_eqn)
 
@@ -659,7 +659,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(0), "Neumann"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
         div_eqn_disc = disc.process_symbol(div_eqn)
 
         linear_y = combined_submesh[0].nodes
@@ -709,7 +709,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(0), "Neumann"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
         div_eqn_disc = disc.process_symbol(div_eqn)
 
         const = 6 * np.ones(sec_pts * prim_pts)
@@ -735,7 +735,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(1), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         disc.set_variable_slices([var])
 
@@ -758,7 +758,7 @@ class TestFiniteVolume(unittest.TestCase):
                 "right": (pybamm.Scalar(combined_submesh[0].edges[-1]), "Dirichlet"),
             }
         }
-        disc._bcs = boundary_conditions
+        disc.bcs = boundary_conditions
 
         grad_eqn_disc = disc.process_symbol(grad_eqn)
         np.testing.assert_array_almost_equal(
@@ -1074,7 +1074,7 @@ class TestFiniteVolume(unittest.TestCase):
 
         # grad
         eqn = pybamm.grad(var)
-        disc._bcs = {
+        disc.bcs = {
             var.id: {
                 "left": (pybamm.Scalar(1), "Dirichlet"),
                 "right": (pybamm.Scalar(2), "Dirichlet"),
@@ -1101,7 +1101,7 @@ class TestFiniteVolume(unittest.TestCase):
         # div(grad)
         flux = pybamm.grad(var)
         eqn = pybamm.div(flux)
-        disc._bcs = {
+        disc.bcs = {
             var.id: {
                 "left": (pybamm.Scalar(1), "Neumann"),
                 "right": (pybamm.Scalar(2), "Neumann"),
@@ -1114,7 +1114,7 @@ class TestFiniteVolume(unittest.TestCase):
         # div(grad) with averaging
         flux = var * pybamm.grad(var)
         eqn = pybamm.div(flux)
-        disc._bcs = {
+        disc.bcs = {
             var.id: {
                 "left": (pybamm.Scalar(1), "Neumann"),
                 "right": (pybamm.Scalar(2), "Neumann"),