Repository: pybamm-team/PyBaMM

Issue:
Shape tests for all expression tree objects
When creating a model, at the moment some errors are only raised when we try to discretise, which makes it a bit difficult to debug. 
The aim of this issue is to find a way to do shape tests at the node creation stage for nodes that don't have an obvious shape, i.e. can't be evaluated (e.g. `Variable`, `Concatenation`, etc).

The best I can come with so far is an `evaluate_for_shape` function that, by default, calls `evaluate`, but if this fails returns a numpy array of the appropriate shape. For example, always give variables in the negative electrode shape '17', variables in the separator shape '23', etc.
Whenever a new object is created, the function `evaluate_for_shape` is called to make sure that it's possible to create that object without getting shape errors further down the line.

Discussion:

Patch:
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 27a576d59..0b57d7a4e 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -17,7 +17,7 @@ if sys.version_info[0] < 3:
     del (x)  # Before Python3, list comprehension iterators leaked
 
 #
-# Expose pints version
+# Expose pybamm version
 #
 
 
@@ -44,11 +44,12 @@ ABSOLUTE_PATH = os.path.join(os.path.split(script_path)[0], "..")
 from .util import Timer
 from .util import load_function
 from .logger import logger, set_logging_level
+from .settings import settings
 
 #
 # Classes for the Expression Tree
 #
-from .expression_tree.symbol import Symbol
+from .expression_tree.symbol import Symbol, evaluate_for_shape_using_domain
 from .expression_tree.binary_operators import (
     is_scalar_zero,
     is_matrix_zero,
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index efb60253c..c0ade17cd 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -388,7 +388,6 @@ class Discretisation(object):
         elif isinstance(symbol, pybamm.UnaryOperator):
             child = symbol.child
             disc_child = self.process_symbol(child)
-            pybamm.SpatialMethod.test_shape(disc_child)
             if child.domain != []:
                 child_spatial_method = self._spatial_methods[child.domain[0]]
             if isinstance(symbol, pybamm.Gradient):
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 2e75336b3..853131326 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -141,6 +141,12 @@ class BinaryOperator(pybamm.Symbol):
             right = self.right.evaluate(t, y)
             return self._binary_evaluate(left, right)
 
+    def evaluate_for_shape(self, t=None, y=None):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape()`. """
+        left = self.children[0].evaluate_for_shape()
+        right = self.children[1].evaluate_for_shape()
+        return self._binary_evaluate(left, right)
+
 
 class Power(BinaryOperator):
     """A node in the expression tree representing a `**` power operator
@@ -478,7 +484,10 @@ class Division(BinaryOperator):
         if issparse(left):
             return left.multiply(1 / right)
         else:
-            return left / right
+            if isinstance(right, numbers.Number) and right == 0:
+                return left * np.inf
+            else:
+                return left / right
 
     def _binary_simplify(self, left, right):
         """ See :meth:`pybamm.BinaryOperator.simplify()`. """
diff --git a/pybamm/expression_tree/broadcasts.py b/pybamm/expression_tree/broadcasts.py
index 8094ed237..ce02af6eb 100644
--- a/pybamm/expression_tree/broadcasts.py
+++ b/pybamm/expression_tree/broadcasts.py
@@ -2,6 +2,7 @@
 # Unary operator classes and methods
 #
 import numbers
+import numpy as np
 import pybamm
 
 
@@ -53,3 +54,15 @@ class Broadcast(pybamm.SpatialOperator):
         """ See :meth:`pybamm.UnaryOperator.simplify()`. """
 
         return Broadcast(child, self.domain)
+
+    def evaluate_for_shape(self):
+        """
+        Returns a vector of NaNs to represent the shape of a Broadcast.
+        See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()`
+        """
+        child_eval = self.children[0].evaluate_for_shape()
+        vec = pybamm.evaluate_for_shape_using_domain(self.domain)
+        if self.children[0].domain == ["current collector"]:
+            return np.outer(child_eval, vec).reshape(-1, 1)
+        else:
+            return child_eval * vec
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index 1bef73f75..7d5dc62cf 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -20,13 +20,14 @@ class Concatenation(pybamm.Symbol):
 
     """
 
-    def __init__(self, *children, name=None, check_domain=True):
+    def __init__(self, *children, name=None, check_domain=True, concat_fun=None):
         if name is None:
             name = "concatenation"
         if check_domain:
             domain = self.get_children_domains(children)
         else:
             domain = []
+        self.concatenation_function = concat_fun
         super().__init__(name, children, domain=domain)
 
     def get_children_domains(self, children):
@@ -41,8 +42,11 @@ class Concatenation(pybamm.Symbol):
         return domain
 
     def _concatenation_evaluate(self, children_eval):
-        """ Concatenate the evaluated children. """
-        raise NotImplementedError
+        """ See :meth:`Concatenation._concatenation_evaluate()`. """
+        if len(children_eval) == 0:
+            return np.array([])
+        else:
+            return self.concatenation_function(children_eval)
 
     def evaluate(self, t=None, y=None, known_evals=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
@@ -76,6 +80,17 @@ class Concatenation(pybamm.Symbol):
         new_symbol.domain = []
         return new_symbol
 
+    def evaluate_for_shape(self):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape` """
+        if len(self.children) == 0:
+            return np.array([])
+        else:
+            # Default: use np.concatenate
+            concatenation_function = self.concatenation_function or np.concatenate
+            return concatenation_function(
+                [child.evaluate_for_shape() for child in self.children]
+            )
+
 
 class NumpyConcatenation(Concatenation):
     """A node in the expression tree representing a concatenation of equations, when we
@@ -101,14 +116,12 @@ class NumpyConcatenation(Concatenation):
         for i, child in enumerate(children):
             if child.evaluates_to_number():
                 children[i] = child * pybamm.Vector(np.array([1]))
-        super().__init__(*children, name="numpy concatenation", check_domain=False)
-
-    def _concatenation_evaluate(self, children_eval):
-        """ See :meth:`Concatenation._concatenation_evaluate()`. """
-        if len(children_eval) == 0:
-            return np.array([])
-        else:
-            return np.concatenate(children_eval)
+        super().__init__(
+            *children,
+            name="numpy concatenation",
+            check_domain=False,
+            concat_fun=np.concatenate
+        )
 
     def jac(self, variable):
         """ See :meth:`pybamm.Symbol.jac()`. """
@@ -267,8 +280,6 @@ class SparseStack(Concatenation):
 
     def __init__(self, *children):
         children = list(children)
-        super().__init__(*children, name="sparse stack", check_domain=False)
-
-    def _concatenation_evaluate(self, children_eval):
-        """ See :meth:`Concatenation.evaluate()`. """
-        return vstack(children_eval)
+        super().__init__(
+            *children, name="sparse stack", check_domain=False, concat_fun=vstack
+        )
diff --git a/pybamm/expression_tree/independent_variable.py b/pybamm/expression_tree/independent_variable.py
index 348370e8a..cbef38655 100644
--- a/pybamm/expression_tree/independent_variable.py
+++ b/pybamm/expression_tree/independent_variable.py
@@ -27,6 +27,10 @@ class IndependentVariable(pybamm.Symbol):
     def __init__(self, name, domain=[]):
         super().__init__(name, domain=domain)
 
+    def evaluate_for_shape(self):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()` """
+        return pybamm.evaluate_for_shape_using_domain(self.domain)
+
 
 class Time(IndependentVariable):
     """A node in the expression tree representing time
@@ -47,6 +51,13 @@ class Time(IndependentVariable):
             raise ValueError("t must be provided")
         return t
 
+    def evaluate_for_shape(self):
+        """
+        Return the scalar '0' to represent the shape of the independent variable `Time`.
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return 0
+
 
 class SpatialVariable(IndependentVariable):
     """A node in the expression tree representing a spatial variable
diff --git a/pybamm/expression_tree/parameter.py b/pybamm/expression_tree/parameter.py
index 9e373f842..a235abf17 100644
--- a/pybamm/expression_tree/parameter.py
+++ b/pybamm/expression_tree/parameter.py
@@ -1,6 +1,7 @@
 #
 # Parameter classes
 #
+import numpy as np
 import pybamm
 
 
@@ -26,6 +27,13 @@ class Parameter(pybamm.Symbol):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
         return Parameter(self.name, self.domain)
 
+    def evaluate_for_shape(self, t=None, y=None):
+        """
+        Returns the scalar 'NaN' to represent the shape of a parameter.
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return np.nan
+
 
 class FunctionParameter(pybamm.UnaryOperator):
     """A node in the expression tree representing a function parameter
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 16a54905b..364e7d3df 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -11,6 +11,18 @@ import autograd.numpy as np
 from anytree.exporter import DotExporter
 
 
+def evaluate_for_shape_using_domain(domain):
+    """
+    Return a vector of the appropriate shape, based on the domain.
+    Domain 'sizes' can clash, but are unlikely to, and won't cause failures if they do.
+    """
+    if domain == []:
+        size = 1
+    else:
+        size = sum(hash(dom) % 100 for dom in domain)
+    return np.nan * np.ones((size, 1))
+
+
 class Symbol(anytree.NodeMixin):
     """Base node class for the expression tree
 
@@ -42,6 +54,16 @@ class Symbol(anytree.NodeMixin):
         # Set domain (and hence id)
         self.domain = domain
 
+        # Test shape on everything but nodes that contain the base Symbol class or
+        # the base BinaryOperator class
+        if pybamm.settings.debug_mode is True:
+            if not any(
+                issubclass(pybamm.Symbol, type(x))
+                or issubclass(pybamm.BinaryOperator, type(x))
+                for x in self.pre_order()
+            ):
+                self.test_shape()
+
     @property
     def children(self):
         """
@@ -366,7 +388,7 @@ class Symbol(anytree.NodeMixin):
         )
 
     def evaluate(self, t=None, y=None, known_evals=None):
-        """Evaluate expression tree (wrapper for dict of known values).
+        """Evaluate expression tree (wrapper to allow using dict of known values).
         If the dict 'known_evals' is provided, the dict is searched for self.id; if
         self.id is in the keys, return that value; otherwise, evaluate using
         :meth:`_base_evaluate()` and add that value to known_evals
@@ -394,6 +416,14 @@ class Symbol(anytree.NodeMixin):
         else:
             return self._base_evaluate(t, y)
 
+    def evaluate_for_shape(self):
+        """Evaluate expression tree to find its shape. For symbols that cannot be
+        evaluated directly (e.g. `Variable` or `Parameter`), a vector of the appropriate
+        shape is returned instead, using the symbol's domain.
+        See :meth:`pybamm.Symbol.evaluate()`
+        """
+        return self.evaluate()
+
     def is_constant(self):
         """returns true if evaluating the expression is not dependent on `t` or `y`
 
@@ -496,13 +526,9 @@ class Symbol(anytree.NodeMixin):
     @property
     def shape(self):
         """
-        Shape of an object, found by evaluating it with appropriate t and y
-
-        Raises
-        ------
-        NotImplementedError
-            If trying to find the shape of an object that cannot be evaluated
+        Shape of an object, found by evaluating it with appropriate t and y.
         """
+        # Default behaviour is to try to evaluate the object directly
         state_vectors_in_node = [
             x for x in self.pre_order() if isinstance(x, pybamm.StateVector)
         ]
@@ -512,8 +538,36 @@ class Symbol(anytree.NodeMixin):
             min_y_size = max(x.y_slice.stop for x in state_vectors_in_node)
             # Pick a y that won't cause RuntimeWarnings
             y = np.linspace(0.1, 0.9, min_y_size)
-        evaluated_self = self.evaluate(t=0, y=y)
+        evaluated_self = self.evaluate(0, y)
         if isinstance(evaluated_self, numbers.Number):
             return ()
         else:
             return evaluated_self.shape
+
+    @property
+    def shape_for_testing(self):
+        """
+        Shape of an object for cases where it cannot be evaluated directly. If a symbol
+        cannot be evaluated directly (e.g. it is a `Variable` or `Parameter`), it is
+        instead given an arbitrary domain-dependent shape from the dictionary
+        `pybamm.DOMAIN_SIZES_FOR_TESTING` (note that this only works for some domains)
+        """
+        evaluated_self = self.evaluate_for_shape()
+        if isinstance(evaluated_self, numbers.Number):
+            return ()
+        else:
+            return evaluated_self.shape
+
+    def test_shape(self):
+        """
+        Check that the discretised self has a pybamm `shape`, i.e. can be evaluated
+
+        Raises
+        ------
+        pybamm.ShapeError
+            If the shape of the object cannot be found
+        """
+        try:
+            self.shape_for_testing
+        except ValueError as e:
+            raise pybamm.ShapeError("Cannot find shape (original error: {})".format(e))
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index ede89b1a1..daf1b36ed 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -66,6 +66,13 @@ class UnaryOperator(pybamm.Symbol):
             child = self.child.evaluate(t, y)
             return self._unary_evaluate(child)
 
+    def evaluate_for_shape(self, t=None, y=None):
+        """
+        Default behaviour: unary operator has same shape as child
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return self.children[0].evaluate_for_shape()
+
 
 class Negate(UnaryOperator):
     """A node in the expression tree representing a `-` negation operator
@@ -265,6 +272,9 @@ class Index(UnaryOperator):
 
         return self.__class__(child, self.index)
 
+    def evaluate_for_shape(self):
+        return self.children[0].evaluate_for_shape()[self.slice]
+
 
 class SpatialOperator(UnaryOperator):
     """A node in the expression tree representing a unary spatial operator
@@ -401,6 +411,10 @@ class Integral(SpatialOperator):
 
         return self.__class__(child, self.integration_variable)
 
+    def evaluate_for_shape(self):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()` """
+        return pybamm.evaluate_for_shape_using_domain(self.domain)
+
 
 class IndefiniteIntegral(Integral):
     """A node in the expression tree representing an indefinite integral operator
@@ -432,6 +446,9 @@ class IndefiniteIntegral(Integral):
         # the integrated variable has the same domain as the child
         self.domain = child.domain
 
+    def evaluate_for_shape(self):
+        return self.children[0].evaluate_for_shape()
+
 
 class BoundaryOperator(SpatialOperator):
     """A node in the expression tree which gets the boundary value of a variable.
@@ -470,6 +487,10 @@ class BoundaryOperator(SpatialOperator):
         """ See :meth:`UnaryOperator._unary_new_copy()`. """
         return self.__class__(child, self.side)
 
+    def evaluate_for_shape(self):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()` """
+        return pybamm.evaluate_for_shape_using_domain(self.domain)
+
 
 class BoundaryValue(BoundaryOperator):
     """A node in the expression tree which gets the boundary value of a variable.
diff --git a/pybamm/expression_tree/variable.py b/pybamm/expression_tree/variable.py
index 5db545b01..76ecd2ccf 100644
--- a/pybamm/expression_tree/variable.py
+++ b/pybamm/expression_tree/variable.py
@@ -28,3 +28,7 @@ class Variable(pybamm.Symbol):
     def new_copy(self):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
         return Variable(self.name, self.domain)
+
+    def evaluate_for_shape(self):
+        """ See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()` """
+        return pybamm.evaluate_for_shape_using_domain(self.domain)
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index 24840dc91..be46f4c06 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -142,3 +142,11 @@ class StateVector(pybamm.Symbol):
     def new_copy(self):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
         return StateVector(self.y_slice, self.name, domain=[])
+
+    def evaluate_for_shape(self):
+        """
+        Returns a vector of NaNs to represent the shape of a StateVector.
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        start = self.y_slice.start or 0
+        return np.nan * np.ones((self.y_slice.stop - start, 1))
diff --git a/pybamm/models/submodels/electrolyte_current.py b/pybamm/models/submodels/electrolyte_current.py
index 51c59b114..aac1de999 100644
--- a/pybamm/models/submodels/electrolyte_current.py
+++ b/pybamm/models/submodels/electrolyte_current.py
@@ -202,7 +202,6 @@ class ElectrolyteCurrentBaseModel(pybamm.SubModel):
                 + (1 - l_p) / kappa_s
             )
         )
-
         phi_e = pybamm.Concatenation(phi_e_n, phi_e_s, phi_e_p)
 
         "Ohmic losses and overpotentials"
diff --git a/pybamm/settings.py b/pybamm/settings.py
new file mode 100644
index 000000000..3ca3f6528
--- /dev/null
+++ b/pybamm/settings.py
@@ -0,0 +1,19 @@
+#
+# Settings class for PyBaMM
+#
+
+
+class Settings(object):
+    _debug_mode = True
+
+    @property
+    def debug_mode(self):
+        return self._debug_mode
+
+    @debug_mode.setter
+    def debug_mode(self, value):
+        assert isinstance(value, bool)
+        self._debug_mode = value
+
+
+settings = Settings()
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index e82fb5231..1b3513d9f 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -81,7 +81,6 @@ class FiniteVolume(pybamm.SpatialMethod):
         gradient_matrix = self.gradient_matrix(domain)
 
         out = gradient_matrix @ discretised_symbol
-        self.test_shape(out)
         return out
 
     def gradient_matrix(self, domain):
@@ -149,7 +148,6 @@ class FiniteVolume(pybamm.SpatialMethod):
         else:
             out = divergence_matrix @ discretised_symbol
 
-        self.test_shape(out)
         return out
 
     def divergence_matrix(self, domain):
@@ -204,7 +202,6 @@ class FiniteVolume(pybamm.SpatialMethod):
             out = integration_vector @ discretised_symbol
         out.domain = []
 
-        self.test_shape(out)
         return out
 
     def definite_integral_vector(self, domain):
@@ -253,7 +250,6 @@ class FiniteVolume(pybamm.SpatialMethod):
 
         out.domain = domain
 
-        self.test_shape(out)
         return out
 
     def indefinite_integral_matrix_edges(self, domain):
@@ -519,7 +515,6 @@ class FiniteVolume(pybamm.SpatialMethod):
         else:
             boundary_value.domain = []
 
-        self.test_shape(boundary_value)
         return boundary_value
 
     def process_binary_operators(self, bin_op, left, right, disc_left, disc_right):
@@ -560,7 +555,6 @@ class FiniteVolume(pybamm.SpatialMethod):
             disc_left = self.node_to_edge(disc_left)
         # Return new binary operator with appropriate class
         out = bin_op.__class__(disc_left, disc_right)
-        self.test_shape(out)
         return out
 
     def concatenation(self, disc_children):
@@ -678,5 +672,4 @@ class FiniteVolume(pybamm.SpatialMethod):
         else:
             out = arithmetic_mean(discretised_symbol)
 
-        self.test_shape(out)
         return out
diff --git a/pybamm/spatial_methods/spatial_method.py b/pybamm/spatial_methods/spatial_method.py
index 1349108c5..47d7e8a84 100644
--- a/pybamm/spatial_methods/spatial_method.py
+++ b/pybamm/spatial_methods/spatial_method.py
@@ -80,7 +80,6 @@ class SpatialMethod:
             )
         else:
             out = symbol * pybamm.Vector(np.ones(vector_size_2D), domain=domain)
-        self.test_shape(out)
         return out
 
     def gradient(self, symbol, discretised_symbol, boundary_conditions):
@@ -204,6 +203,7 @@ class SpatialMethod:
             # as above, but now we want a single point with value 1 at (0, n-1)
             right_vector = coo_matrix(([1], ([0], [n - 1])), shape=(1, n))
             bv_vector = pybamm.Matrix(right_vector)
+
         out = bv_vector @ discretised_child
         # boundary value removes domain
         out.domain = []
@@ -285,30 +285,3 @@ class SpatialMethod:
             Concatenation of the discretised children
         """
         return pybamm.DomainConcatenation(disc_children, self.mesh)
-
-    @staticmethod
-    def test_shape(symbol):
-        """
-        Check that the discretised symbol has a pybamm `shape`, i.e. can be evaluated
-
-        Parameters
-        ----------
-        symbol : :class:`pybamm.Symbol`
-            The symbol to be tested
-
-        Raises
-        ------
-        pybamm.ShapeError
-            If the shape of the object cannot be found
-        """
-        try:
-            symbol.shape
-        except ValueError as e:
-            raise pybamm.ShapeError(
-                """
-                Cannot evaluate the discretised symbol to find its shape
-                (original error: {})
-                """.format(
-                    e
-                )
-            )


Test patch:
diff --git a/run-tests.py b/run-tests.py
index 1549f2fbb..d7e3c2f8f 100755
--- a/run-tests.py
+++ b/run-tests.py
@@ -7,6 +7,7 @@
 #
 import re
 import os
+import pybamm
 import sys
 import argparse
 import unittest
@@ -30,6 +31,8 @@ def run_code_tests(executable=None, folder: str = "unit"):
         tests = "tests/"
     else:
         tests = "tests/" + folder
+        if folder == "unit":
+            pybamm.settings.debug_mode = True
     if executable is None:
         suite = unittest.defaultTestLoader.discover(tests, pattern="test*.py")
         unittest.TextTestRunner(verbosity=2).run(suite)
@@ -203,9 +206,9 @@ def test_notebook(path, executable="python"):
                 j = str(1 + i)
                 print(j + " " * (5 - len(j)) + line)
             print("-- stdout " + "-" * (79 - 10))
-            print(str(stdout, 'utf-8'))
+            print(str(stdout, "utf-8"))
             print("-- stderr " + "-" * (79 - 10))
-            print(str(stderr, 'utf-8'))
+            print(str(stderr, "utf-8"))
             print("-" * 79)
             return False
     except KeyboardInterrupt:
@@ -244,9 +247,9 @@ def test_script(path, executable="python"):
             # Show failing code, output and errors before returning
             print("ERROR")
             print("-- stdout " + "-" * (79 - 10))
-            print(str(stdout, 'utf-8'))
+            print(str(stdout, "utf-8"))
             print("-- stderr " + "-" * (79 - 10))
-            print(str(stderr, 'utf-8'))
+            print(str(stderr, "utf-8"))
             print("-" * 79)
             return False
     except KeyboardInterrupt:
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index bbc6a4b41..feff31a4c 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -116,7 +116,7 @@ class TestDiscretise(unittest.TestCase):
         # variable
         var = pybamm.Variable("var")
         var_vec = pybamm.Variable("var vec", domain=["negative electrode"])
-        disc._y_slices = {var.id: slice(53), var_vec.id: slice(53, 102)}
+        disc._y_slices = {var.id: slice(53), var_vec.id: slice(53, 93)}
         var_disc = disc.process_symbol(var)
         self.assertIsInstance(var_disc, pybamm.StateVector)
         self.assertEqual(var_disc._y_slice, disc._y_slices[var.id])
@@ -189,7 +189,7 @@ class TestDiscretise(unittest.TestCase):
         # create discretisation
         disc = get_discretisation_for_testing()
 
-        disc._y_slices = {var1.id: slice(53), var2.id: slice(53, 59)}
+        disc._y_slices = {var1.id: slice(53), var2.id: slice(53, 106)}
         exp_disc = disc.process_symbol(expression)
         self.assertIsInstance(exp_disc, pybamm.Division)
         # left side
diff --git a/tests/unit/test_expression_tree/test_binary_operators.py b/tests/unit/test_expression_tree/test_binary_operators.py
index 1bdf8d33c..8369e9a80 100644
--- a/tests/unit/test_expression_tree/test_binary_operators.py
+++ b/tests/unit/test_expression_tree/test_binary_operators.py
@@ -233,10 +233,10 @@ class TestBinaryOperators(unittest.TestCase):
         np.testing.assert_array_equal(
             (pybammS2 * pybammD2).evaluate().toarray(), S2.toarray() * D2
         )
-        with self.assertRaisesRegex(ValueError, "inconsistent shapes"):
-            (pybammS1 * pybammS2).evaluate()
-        with self.assertRaisesRegex(ValueError, "inconsistent shapes"):
-            (pybammS2 * pybammS1).evaluate()
+        with self.assertRaisesRegex(pybamm.ShapeError, "inconsistent shapes"):
+            (pybammS1 * pybammS2).test_shape()
+        with self.assertRaisesRegex(pybamm.ShapeError, "inconsistent shapes"):
+            (pybammS2 * pybammS1).test_shape()
 
         # Matrix multiplication is normal matrix multiplication
         np.testing.assert_array_equal(
@@ -249,10 +249,10 @@ class TestBinaryOperators(unittest.TestCase):
         np.testing.assert_array_equal((pybammD2 @ pybammS1).evaluate(), D2 * S1)
         np.testing.assert_array_equal((pybammS2 @ pybammD1).evaluate(), S2 * D1)
         np.testing.assert_array_equal((pybammD1 @ pybammS2).evaluate(), D1 * S2)
-        with self.assertRaisesRegex(ValueError, "dimension mismatch"):
-            (pybammS1 @ pybammS1).evaluate()
-        with self.assertRaisesRegex(ValueError, "dimension mismatch"):
-            (pybammS2 @ pybammS2).evaluate()
+        with self.assertRaisesRegex(pybamm.ShapeError, "dimension mismatch"):
+            (pybammS1 @ pybammS1).test_shape()
+        with self.assertRaisesRegex(pybamm.ShapeError, "dimension mismatch"):
+            (pybammS2 @ pybammS2).test_shape()
 
     def test_sparse_divide(self):
         row = np.array([0, 3, 1, 0])
diff --git a/tests/unit/test_expression_tree/test_concatenations.py b/tests/unit/test_expression_tree/test_concatenations.py
index aefbe885b..e11310878 100644
--- a/tests/unit/test_expression_tree/test_concatenations.py
+++ b/tests/unit/test_expression_tree/test_concatenations.py
@@ -18,11 +18,11 @@ class TestConcatenations(unittest.TestCase):
         self.assertEqual(conc.children[0].name, "a")
         self.assertEqual(conc.children[1].name, "b")
         self.assertEqual(conc.children[2].name, "c")
-        d = pybamm.Scalar(2)
-        e = pybamm.Scalar(1)
-        f = pybamm.Scalar(3)
+        d = pybamm.Vector(np.array([2]))
+        e = pybamm.Vector(np.array([1]))
+        f = pybamm.Vector(np.array([3]))
         conc2 = pybamm.Concatenation(d, e, f)
-        with self.assertRaises(NotImplementedError):
+        with self.assertRaises(TypeError):
             conc2.evaluate()
 
     def test_concatenation_domains(self):
@@ -160,9 +160,9 @@ class TestConcatenations(unittest.TestCase):
         )
 
     def test_concatenation_orphans(self):
-        a = pybamm.Variable("a")
-        b = pybamm.Variable("b")
-        c = pybamm.Variable("c")
+        a = pybamm.Variable("a", domain=["negative electrode"])
+        b = pybamm.Variable("b", domain=["separator"])
+        c = pybamm.Variable("c", domain=["positive electrode"])
         conc = pybamm.Concatenation(a, b, c)
         a_new, b_new, c_new = conc.orphans
 
diff --git a/tests/unit/test_expression_tree/test_copy.py b/tests/unit/test_expression_tree/test_copy.py
index 83fda91d4..4c993b93f 100644
--- a/tests/unit/test_expression_tree/test_copy.py
+++ b/tests/unit/test_expression_tree/test_copy.py
@@ -32,7 +32,7 @@ class TestSimplify(unittest.TestCase):
             pybamm.BoundaryValue(a, "right"),
             pybamm.BoundaryFlux(a, "right"),
             pybamm.Broadcast(a, "domain"),
-            pybamm.Concatenation(a, b, v_n),
+            pybamm.Concatenation(v_n, v_s),
             pybamm.NumpyConcatenation(a, b, v_s),
             pybamm.DomainConcatenation([v_n, v_s], mesh),
             pybamm.Parameter("param"),
diff --git a/tests/unit/test_expression_tree/test_evaluate.py b/tests/unit/test_expression_tree/test_evaluate.py
index 9a22242d7..aabef48ef 100644
--- a/tests/unit/test_expression_tree/test_evaluate.py
+++ b/tests/unit/test_expression_tree/test_evaluate.py
@@ -32,13 +32,14 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(variable_symbols.keys())[2], expr.id)
 
         # test values of variable_symbols
-        self.assertEqual(list(variable_symbols.values())[0], 'y[0:1]')
-        self.assertEqual(list(variable_symbols.values())[1], 'y[1:2]')
+        self.assertEqual(list(variable_symbols.values())[0], "y[0:1]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[1:2]")
 
         var_a = pybamm.id_to_python_variable(a.id)
         var_b = pybamm.id_to_python_variable(b.id)
-        self.assertEqual(list(variable_symbols.values())[
-                         2], '{} + {}'.format(var_a, var_b))
+        self.assertEqual(
+            list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b)
+        )
 
         # test identical subtree
         constant_symbols = OrderedDict()
@@ -54,14 +55,16 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(variable_symbols.keys())[3], expr.id)
 
         # test values of variable_symbols
-        self.assertEqual(list(variable_symbols.values())[0], 'y[0:1]')
-        self.assertEqual(list(variable_symbols.values())[1], 'y[1:2]')
-        self.assertEqual(list(variable_symbols.values())[
-                         2], '{} + {}'.format(var_a, var_b))
+        self.assertEqual(list(variable_symbols.values())[0], "y[0:1]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[1:2]")
+        self.assertEqual(
+            list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b)
+        )
 
         var_child = pybamm.id_to_python_variable(expr.children[0].id)
-        self.assertEqual(list(variable_symbols.values())[
-                         3], '{} + {}'.format(var_child, var_b))
+        self.assertEqual(
+            list(variable_symbols.values())[3], "{} + {}".format(var_child, var_b)
+        )
 
         # test unary op
         constant_symbols = OrderedDict()
@@ -77,12 +80,13 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(variable_symbols.keys())[3], expr.id)
 
         # test values of variable_symbols
-        self.assertEqual(list(variable_symbols.values())[0], 'y[0:1]')
-        self.assertEqual(list(variable_symbols.values())[1], 'y[1:2]')
-        self.assertEqual(list(variable_symbols.values())[2], '-{}'.format(var_b))
+        self.assertEqual(list(variable_symbols.values())[0], "y[0:1]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[1:2]")
+        self.assertEqual(list(variable_symbols.values())[2], "-{}".format(var_b))
         var_child = pybamm.id_to_python_variable(expr.children[1].id)
-        self.assertEqual(list(variable_symbols.values())[3],
-                         '{} + {}'.format(var_a, var_child))
+        self.assertEqual(
+            list(variable_symbols.values())[3], "{} + {}".format(var_a, var_child)
+        )
 
         # test function
         constant_symbols = OrderedDict()
@@ -93,10 +97,11 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(constant_symbols.values())[0], test_function)
         self.assertEqual(list(variable_symbols.keys())[0], a.id)
         self.assertEqual(list(variable_symbols.keys())[1], expr.id)
-        self.assertEqual(list(variable_symbols.values())[0], 'y[0:1]')
+        self.assertEqual(list(variable_symbols.values())[0], "y[0:1]")
         var_funct = pybamm.id_to_python_variable(expr.id, True)
-        self.assertEqual(list(variable_symbols.values())[1], '{}({})'.format(var_funct,
-                                                                             var_a))
+        self.assertEqual(
+            list(variable_symbols.values())[1], "{}({})".format(var_funct, var_a)
+        )
 
         # test matrix
         constant_symbols = OrderedDict()
@@ -105,8 +110,9 @@ class TestEvaluate(unittest.TestCase):
         pybamm.find_symbols(A, constant_symbols, variable_symbols)
         self.assertEqual(len(variable_symbols), 0)
         self.assertEqual(list(constant_symbols.keys())[0], A.id)
-        np.testing.assert_allclose(list(constant_symbols.values())[0],
-                                   np.array([[1, 2], [3, 4]]))
+        np.testing.assert_allclose(
+            list(constant_symbols.values())[0], np.array([[1, 2], [3, 4]])
+        )
 
         # test sparse matrix
         constant_symbols = OrderedDict()
@@ -115,8 +121,9 @@ class TestEvaluate(unittest.TestCase):
         pybamm.find_symbols(A, constant_symbols, variable_symbols)
         self.assertEqual(len(variable_symbols), 0)
         self.assertEqual(list(constant_symbols.keys())[0], A.id)
-        np.testing.assert_allclose(list(constant_symbols.values())[
-                                   0].toarray(), A.entries.toarray())
+        np.testing.assert_allclose(
+            list(constant_symbols.values())[0].toarray(), A.entries.toarray()
+        )
 
         # test numpy concatentate
         constant_symbols = OrderedDict()
@@ -127,9 +134,10 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(variable_symbols.keys())[0], a.id)
         self.assertEqual(list(variable_symbols.keys())[1], b.id)
         self.assertEqual(list(variable_symbols.keys())[2], expr.id)
-        self.assertEqual(list(variable_symbols.values())[2],
-                         "np.concatenate(({},{}))".format(var_a, var_b)
-                         )
+        self.assertEqual(
+            list(variable_symbols.values())[2],
+            "np.concatenate(({},{}))".format(var_a, var_b),
+        )
 
         # test domain concatentate
         constant_symbols = OrderedDict()
@@ -140,22 +148,10 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(list(variable_symbols.keys())[0], a.id)
         self.assertEqual(list(variable_symbols.keys())[1], b.id)
         self.assertEqual(list(variable_symbols.keys())[2], expr.id)
-        self.assertEqual(list(variable_symbols.values())[2],
-                         "np.concatenate(({},{}))".format(var_a, var_b)
-                         )
-
-        # test sparse stack
-        constant_symbols = OrderedDict()
-        variable_symbols = OrderedDict()
-        expr = pybamm.SparseStack(a, b)
-        pybamm.find_symbols(expr, constant_symbols, variable_symbols)
-        self.assertEqual(len(constant_symbols), 0)
-        self.assertEqual(list(variable_symbols.keys())[0], a.id)
-        self.assertEqual(list(variable_symbols.keys())[1], b.id)
-        self.assertEqual(list(variable_symbols.keys())[2], expr.id)
-        self.assertEqual(list(variable_symbols.values())[2],
-                         "scipy.sparse.vstack(({},{}))".format(var_a, var_b)
-                         )
+        self.assertEqual(
+            list(variable_symbols.values())[2],
+            "np.concatenate(({},{}))".format(var_a, var_b),
+        )
 
         # test that Concatentation throws
         expr = pybamm.Concatenation(a, b)
@@ -163,7 +159,7 @@ class TestEvaluate(unittest.TestCase):
             pybamm.find_symbols(expr, constant_symbols, variable_symbols)
 
         # test that these nodes throw
-        for expr in (pybamm.Variable('a'), pybamm.Parameter('a')):
+        for expr in (pybamm.Variable("a"), pybamm.Parameter("a")):
             with self.assertRaises(NotImplementedError):
                 pybamm.find_symbols(expr, constant_symbols, variable_symbols)
 
@@ -194,10 +190,10 @@ class TestEvaluate(unittest.TestCase):
         var_a = pybamm.id_to_python_variable(a.id)
         var_b = pybamm.id_to_python_variable(b.id)
         self.assertEqual(len(constant_symbols), 0)
-        self.assertEqual(list(variable_symbols.values())[2],
-                         "np.concatenate(({}[0:{}],{}[0:{}]))".format(
-                             var_a, a_pts, var_b, b_pts)
-                         )
+        self.assertEqual(
+            list(variable_symbols.values())[2],
+            "np.concatenate(({}[0:{}],{}[0:{}]))".format(var_a, a_pts, var_b, b_pts),
+        )
 
         evaluator = pybamm.EvaluatorPython(expr)
         result = evaluator.evaluate(y=y)
@@ -217,9 +213,7 @@ class TestEvaluate(unittest.TestCase):
         b1_pts = mesh[b_dom[1]][0].npts
 
         a = pybamm.StateVector(slice(0, a0_pts), domain=a_dom)
-        b = pybamm.StateVector(slice(a0_pts, a0_pts + b0_pts + b1_pts),
-                               domain=b_dom,
-                               )
+        b = pybamm.StateVector(slice(a0_pts, a0_pts + b0_pts + b1_pts), domain=b_dom)
 
         y = np.empty((a0_pts + b0_pts + b1_pts, 1))
         for i in range(len(y)):
@@ -237,9 +231,10 @@ class TestEvaluate(unittest.TestCase):
         b1_str = "{}[{}:{}]".format(var_b, b0_pts, b0_pts + b1_pts)
 
         self.assertEqual(len(constant_symbols), 0)
-        self.assertEqual(list(variable_symbols.values())[2],
-                         "np.concatenate(({},{},{}))".format(b0_str, a0_str, b1_str)
-                         )
+        self.assertEqual(
+            list(variable_symbols.values())[2],
+            "np.concatenate(({},{},{}))".format(b0_str, a0_str, b1_str),
+        )
 
         evaluator = pybamm.EvaluatorPython(expr)
         result = evaluator.evaluate(y=y)
@@ -252,10 +247,11 @@ class TestEvaluate(unittest.TestCase):
         # test a * b
         expr = a + b
         constant_str, variable_str = pybamm.to_python(expr)
-        expected_str = \
-            "self\.var_[0-9m]+ = y\[0:1\].*\\n" \
-            "self\.var_[0-9m]+ = y\[1:2\].*\\n" \
+        expected_str = (
+            "self\.var_[0-9m]+ = y\[0:1\].*\\n"
+            "self\.var_[0-9m]+ = y\[1:2\].*\\n"
             "self\.var_[0-9m]+ = self\.var_[0-9m]+ \+ self\.var_[0-9m]+"
+        )
 
         self.assertRegex(variable_str, expected_str)
 
@@ -287,7 +283,7 @@ class TestEvaluate(unittest.TestCase):
         self.assertEqual(result, 6)
 
         # test a larger expression
-        expr = a * b + b + a**2 / b + 2 * a + b / 2 + 4
+        expr = a * b + b + a ** 2 / b + 2 * a + b / 2 + 4
         evaluator = pybamm.EvaluatorPython(expr)
         for y in y_tests:
             result = evaluator.evaluate(t=None, y=y)
diff --git a/tests/unit/test_expression_tree/test_jac.py b/tests/unit/test_expression_tree/test_jac.py
index a2b13c690..325dc35b6 100644
--- a/tests/unit/test_expression_tree/test_jac.py
+++ b/tests/unit/test_expression_tree/test_jac.py
@@ -43,7 +43,7 @@ class TestJacobian(unittest.TestCase):
         dfunc_dy = func.jac(y).simplify().evaluate(y=y0)
         np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
 
-        func = u @ A
+        func = u @ pybamm.StateVector(slice(0, 1))
         with self.assertRaises(NotImplementedError):
             func.jac(y)
 
diff --git a/tests/unit/test_expression_tree/test_parameter.py b/tests/unit/test_expression_tree/test_parameter.py
index bcae87603..e2bbb394a 100644
--- a/tests/unit/test_expression_tree/test_parameter.py
+++ b/tests/unit/test_expression_tree/test_parameter.py
@@ -1,8 +1,8 @@
 #
 # Tests for the Parameter class
 #
+import numbers
 import pybamm
-
 import unittest
 
 
@@ -14,6 +14,10 @@ class TestParameter(unittest.TestCase):
         c = pybamm.Parameter("c", domain=["test"])
         self.assertEqual(c.domain[0], "test")
 
+    def test_evaluate_for_shape(self):
+        a = pybamm.Parameter("a")
+        self.assertIsInstance(a.evaluate_for_shape(), numbers.Number)
+
 
 class TestFunctionParameter(unittest.TestCase):
     def test_function_parameter_init(self):
@@ -29,6 +33,11 @@ class TestFunctionParameter(unittest.TestCase):
         func = pybamm.FunctionParameter("a", var).diff(var)
         self.assertEqual(func.diff_variable, var)
 
+    def test_evaluate_for_shape(self):
+        a = pybamm.Parameter("a")
+        func = pybamm.FunctionParameter("func", 2 * a)
+        self.assertIsInstance(func.evaluate_for_shape(), numbers.Number)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index 3a1bb2a4e..7a5476e6f 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -256,18 +256,15 @@ class TestSimplify(unittest.TestCase):
         self.assertEqual(funca.evaluate(), 1)
 
     def test_matrix_simplifications(self):
-        a = pybamm.Scalar(0)
-        b = pybamm.Scalar(1)
-        c = pybamm.Parameter("c")
+        a = pybamm.Matrix(np.zeros((2, 2)))
+        b = pybamm.Matrix(np.ones((2, 2)))
 
         # matrix multiplication
         A = pybamm.Matrix(np.array([[1, 0], [0, 1]]))
-        self.assertIsInstance((a @ A).simplify(), pybamm.Scalar)
-        self.assertEqual((a @ A).simplify().evaluate(), 0)
-        self.assertIsInstance((A @ a).simplify(), pybamm.Scalar)
-        self.assertEqual((A @ a).simplify().evaluate(), 0)
-
-        self.assertIsInstance((A @ c).simplify(), pybamm.MatrixMultiplication)
+        self.assertIsInstance((a @ A).simplify(), pybamm.Matrix)
+        np.testing.assert_array_equal((a @ A).simplify().evaluate(), 0)
+        self.assertIsInstance((A @ a).simplify(), pybamm.Matrix)
+        np.testing.assert_array_equal((A @ a).simplify().evaluate(), 0)
 
         # matrix * matrix
         m1 = pybamm.Matrix(np.array([[2, 0], [0, 2]]))
@@ -290,11 +287,11 @@ 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)
-        self.assertIsInstance(expr.children[1], pybamm.MatrixMultiplication)
+        self.assertIsInstance(expr.children[1], pybamm.Multiplication)
 
         # mat-mul on numerator and denominator
         expr = (m2 @ (m1 @ v) / (m2 @ (m1 @ v))).simplify()
@@ -321,6 +318,7 @@ class TestSimplify(unittest.TestCase):
         )
 
         # scalar * matrix
+        b = pybamm.Scalar(1)
         for expr in [
             ((b * m1) @ v).simplify(),
             (b * (m1 @ v)).simplify(),
@@ -366,7 +364,7 @@ class TestSimplify(unittest.TestCase):
         )
         self.assertEqual(expr2.id, expr2simp.id)
 
-        expr3 = m1 @ ((m2 @ v1) * (m3 @ v2))
+        expr3 = m1 @ ((m2 @ v2) * (m2 @ v2))
         expr3simp = expr3.simplify()
         self.assertEqual(expr3.id, expr3simp.id)
 
diff --git a/tests/unit/test_expression_tree/test_symbol.py b/tests/unit/test_expression_tree/test_symbol.py
index b141b563d..dbea2d478 100644
--- a/tests/unit/test_expression_tree/test_symbol.py
+++ b/tests/unit/test_expression_tree/test_symbol.py
@@ -289,14 +289,54 @@ class TestSymbol(unittest.TestCase):
 
     def test_shape(self):
         scal = pybamm.Scalar(1)
-        self.assertEqual(scal.shape, ())
+        self.assertEqual(scal.shape_for_testing, ())
+        self.assertEqual(scal.size, 1)
 
+        state = pybamm.StateVector(slice(10))
+        self.assertEqual(state.shape_for_testing, (10, 1))
+        self.assertEqual(state.size, 10)
         state = pybamm.StateVector(slice(10, 25))
-        self.assertEqual(state.shape, (15, 1))
+        self.assertEqual(state.shape_for_testing, (15, 1))
+
+    def test_shape_for_testing(self):
+        scal = pybamm.Scalar(1)
+        self.assertEqual(scal.shape_for_testing, scal.shape)
+
+        state = pybamm.StateVector(slice(10, 25))
+        self.assertEqual(state.shape_for_testing, state.shape)
+
+        param = pybamm.Parameter("a")
+        self.assertEqual(param.shape_for_testing, ())
+
+        func = pybamm.FunctionParameter("func", state)
+        self.assertEqual(func.shape_for_testing, state.shape_for_testing)
+
+        concat = pybamm.Concatenation()
+        self.assertEqual(concat.shape_for_testing, (0,))
+        concat = pybamm.Concatenation(state, state)
+        self.assertEqual(concat.shape_for_testing, (30, 1))
+
+        var = pybamm.Variable("var", domain="negative electrode")
+        broadcast = pybamm.Broadcast(0, domain="negative electrode")
+        self.assertEqual(var.shape_for_testing, broadcast.shape_for_testing)
+        self.assertEqual(
+            (var + broadcast).shape_for_testing, broadcast.shape_for_testing
+        )
+        broadcast2 = pybamm.Broadcast(0, domain="separator")
+        self.assertNotEqual(var.shape_for_testing, broadcast2.shape_for_testing)
 
         sym = pybamm.Symbol("sym")
         with self.assertRaises(NotImplementedError):
-            sym.shape
+            sym.shape_for_testing
+
+    def test_test_shape(self):
+        # right shape, passes
+        y1 = pybamm.StateVector(slice(0, 10))
+        y1.test_shape()
+        # bad shape, fails
+        y2 = pybamm.StateVector(slice(0, 5))
+        with self.assertRaises(pybamm.ShapeError):
+            (y1 + y2).test_shape()
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/test_models/test_base_model.py b/tests/unit/test_models/test_base_model.py
index 9a34df1cd..92361c5fe 100644
--- a/tests/unit/test_models/test_base_model.py
+++ b/tests/unit/test_models/test_base_model.py
@@ -241,7 +241,7 @@ class TestBaseModel(unittest.TestCase):
         model = pybamm.BaseModel()
         model.algebraic = {
             c: 1,
-            d: pybamm.StateVector(slice(0, 15)) - pybamm.StateVector(slice(15, 25)),
+            d: pybamm.StateVector(slice(0, 15)) - pybamm.StateVector(slice(15, 30)),
         }
         with self.assertRaisesRegex(
             pybamm.ModelError,
diff --git a/tests/unit/test_parameters/test_parameter_values.py b/tests/unit/test_parameters/test_parameter_values.py
index 1fc47baf3..5121588ed 100644
--- a/tests/unit/test_parameters/test_parameter_values.py
+++ b/tests/unit/test_parameters/test_parameter_values.py
@@ -107,12 +107,14 @@ class TestParameterValues(unittest.TestCase):
         self.assertEqual(processed_broad.children[0].evaluate(), np.array([1]))
 
         # process concatenation
-        conc = pybamm.Concatenation(a, b)
+        conc = pybamm.Concatenation(
+            pybamm.Vector(np.ones(10)), pybamm.Vector(2 * np.ones(15))
+        )
         processed_conc = parameter_values.process_symbol(conc)
-        self.assertIsInstance(processed_conc.children[0], pybamm.Scalar)
-        self.assertIsInstance(processed_conc.children[1], pybamm.Scalar)
-        self.assertEqual(processed_conc.children[0].value, 1)
-        self.assertEqual(processed_conc.children[1].value, 2)
+        self.assertIsInstance(processed_conc.children[0], pybamm.Vector)
+        self.assertIsInstance(processed_conc.children[1], pybamm.Vector)
+        np.testing.assert_array_equal(processed_conc.children[0].entries, 1)
+        np.testing.assert_array_equal(processed_conc.children[1].entries, 2)
 
         # process domain concatenation
         c_e_n = pybamm.Variable("c_e_n", ["negative electrode"])
diff --git a/tests/unit/test_spatial_methods/test_base_spatial_method.py b/tests/unit/test_spatial_methods/test_base_spatial_method.py
index 61491e119..c69664285 100644
--- a/tests/unit/test_spatial_methods/test_base_spatial_method.py
+++ b/tests/unit/test_spatial_methods/test_base_spatial_method.py
@@ -61,17 +61,6 @@ class TestSpatialMethod(unittest.TestCase):
         with self.assertRaisesRegex(NotImplementedError, "Cannot process 2D symbol"):
             spatial_method.boundary_value_or_flux(symbol, child)
 
-    def test_test_shape(self):
-        mesh = get_mesh_for_testing()
-        spatial_method = pybamm.SpatialMethod(mesh)
-        # right shape, passes
-        y1 = pybamm.StateVector(slice(0, 10))
-        spatial_method.test_shape(y1)
-        # bad shape, fails
-        y2 = pybamm.StateVector(slice(0, 5))
-        with self.assertRaises(pybamm.ShapeError):
-            spatial_method.test_shape(y1 + y2)
-
 
 if __name__ == "__main__":
     print("Add -v for more debug output")