Repository: tinosulzer/PyBaMM

Issue:
smarter domain concatenation and broadcast
At the moment, the DomainConcatentation class will convert any Scalar children to vectors. This should be generalised, so that if the `evaluate` of its children return a scalar, this is appropriately mapped to a vector at run-time (i.e. for time dependent functions)

Discussion:

Patch:
diff --git a/docs/source/expression_tree/broadcasts.rst b/docs/source/expression_tree/broadcasts.rst
new file mode 100644
index 00000000..98db1fe2
--- /dev/null
+++ b/docs/source/expression_tree/broadcasts.rst
@@ -0,0 +1,8 @@
+Broadcasting Operators
+======================
+
+.. autoclass:: pybamm.Broadcast
+  :members:
+
+.. autoclass:: pybamm.NumpyBroadcast
+  :members:
diff --git a/docs/source/expression_tree/concatenations.rst b/docs/source/expression_tree/concatenations.rst
index b0bbd364..3456d9e6 100644
--- a/docs/source/expression_tree/concatenations.rst
+++ b/docs/source/expression_tree/concatenations.rst
@@ -4,11 +4,8 @@ Concatenations
 .. autoclass:: pybamm.Concatenation
   :members:
 
-.. autoclass:: pybamm.NumpyModelConcatenation
+.. autoclass:: pybamm.NumpyConcatenation
   :members:
 
 .. autoclass:: pybamm.DomainConcatenation
   :members:
-
-.. autoclass:: pybamm.PiecewiseConstant
-  :members:
diff --git a/docs/source/expression_tree/index.rst b/docs/source/expression_tree/index.rst
index 8b63c475..15f8422c 100644
--- a/docs/source/expression_tree/index.rst
+++ b/docs/source/expression_tree/index.rst
@@ -7,9 +7,10 @@ Expression Tree
   parameter
   variable
   independent_variable
-  binary_operator
-  unary_operator
   scalar
   matrix
   vector
+  binary_operator
+  unary_operator
   concatenations
+  broadcasts
diff --git a/docs/source/expression_tree/unary_operator.rst b/docs/source/expression_tree/unary_operator.rst
index 13675c49..c547d1af 100644
--- a/docs/source/expression_tree/unary_operator.rst
+++ b/docs/source/expression_tree/unary_operator.rst
@@ -19,12 +19,6 @@ Unary Operators
 .. autoclass:: pybamm.Divergence
   :members:
 
-.. autoclass:: pybamm.Broadcast
-  :members:
-
-.. autoclass:: pybamm.NumpyBroadcast
-  :members:
-
 .. autoclass:: pybamm.SurfaceValue
   :members:
 
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index eae5d948..3a879f1b 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -57,13 +57,13 @@ from .expression_tree.binary_operators import (
     Power,
     Subtraction,
     Multiplication,
+    MatrixMultiplication,
     Division,
 )
 from .expression_tree.concatenations import (
     Concatenation,
-    NumpyModelConcatenation,
+    NumpyConcatenation,
     DomainConcatenation,
-    PiecewiseConstant,
 )
 from .expression_tree.array import Array
 from .expression_tree.matrix import Matrix
@@ -76,13 +76,12 @@ from .expression_tree.unary_operators import (
     SpatialOperator,
     Gradient,
     Divergence,
-    Broadcast,
-    NumpyBroadcast,
     SurfaceValue,
     grad,
     div,
     surf,
 )
+from .expression_tree.broadcasts import Broadcast, NumpyBroadcast
 from .expression_tree.function_parameter import FunctionParameter
 from .expression_tree.scalar import Scalar
 from .expression_tree.variable import Variable
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 14d7e507..de5a745e 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -164,12 +164,12 @@ class Discretisation(object):
         """
 
         for eqn_key, eqn in var_eqn_dict.items():
-
+            # Broadcast if the equation evaluates to a number(e.g. Scalar)
             if eqn.evaluates_to_number():
                 if isinstance(eqn_key, str):
-                    eqn = pybamm.NumpyBroadcast(eqn, [], {})
+                    eqn = pybamm.Broadcast(eqn, [])
                 elif eqn_key.domain == []:
-                    eqn = pybamm.NumpyBroadcast(eqn, eqn_key.domain, {})
+                    eqn = pybamm.Broadcast(eqn, eqn_key.domain)
                 else:
                     eqn = self._spatial_methods[eqn_key.domain[0]].broadcast(
                         eqn, eqn_key.domain
@@ -215,7 +215,7 @@ class Discretisation(object):
             # Broadcast new_child to the domain specified by symbol.domain
             # Different discretisations may broadcast differently
             if symbol.domain == []:
-                symbol = pybamm.NumpyBroadcast(symbol, symbol.domain, {})
+                symbol = pybamm.NumpyBroadcast(new_child, symbol.domain, {})
             else:
                 symbol = self._spatial_methods[symbol.domain[0]].broadcast(
                     new_child, symbol.domain
@@ -241,7 +241,7 @@ class Discretisation(object):
             return symbol.__class__(new_child)
 
         elif isinstance(symbol, pybamm.Variable):
-            return pybamm.StateVector(self._y_slices[symbol.id])
+            return pybamm.StateVector(self._y_slices[symbol.id], domain=symbol.domain)
 
         elif isinstance(symbol, pybamm.SpatialVariable):
             return self._spatial_methods[symbol.domain[0]].spatial_variable(symbol)
@@ -251,7 +251,8 @@ class Discretisation(object):
             new_symbol = pybamm.DomainConcatenation(new_children, self.mesh)
 
             if new_symbol.is_constant():
-                return pybamm.Vector(new_symbol.evaluate())
+                value = new_symbol.evaluate()
+                return pybamm.Vector(value)
             return new_symbol
 
         else:
@@ -309,7 +310,7 @@ class Discretisation(object):
         return bin_op.__class__(new_left, new_right)
 
     def concatenate(self, *symbols):
-        return pybamm.NumpyModelConcatenation(*symbols)
+        return pybamm.NumpyConcatenation(*symbols)
 
     def _concatenate_init(self, var_eqn_dict):
         """
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index e1437969..87173591 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -118,10 +118,23 @@ class Multiplication(BinaryOperator):
 
     def evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
-        if isinstance(self.children[0], pybamm.Matrix):
-            return self.children[0].evaluate(t, y) @ self.children[1].evaluate(t, y)
-        else:
-            return self.children[0].evaluate(t, y) * self.children[1].evaluate(t, y)
+        return self.children[0].evaluate(t, y) * self.children[1].evaluate(t, y)
+
+
+class MatrixMultiplication(BinaryOperator):
+    """A node in the expression tree representing a matrix multiplication operator
+
+    **Extends:** :class:`BinaryOperator`
+    """
+
+    def __init__(self, left, right):
+        """ See :meth:`pybamm.BinaryOperator.__init__()`. """
+
+        super().__init__("*", left, right)
+
+    def evaluate(self, t=None, y=None):
+        """ See :meth:`pybamm.Symbol.evaluate()`. """
+        return self.children[0].evaluate(t, y) @ self.children[1].evaluate(t, y)
 
 
 class Division(BinaryOperator):
diff --git a/pybamm/expression_tree/broadcasts.py b/pybamm/expression_tree/broadcasts.py
new file mode 100644
index 00000000..ec98e261
--- /dev/null
+++ b/pybamm/expression_tree/broadcasts.py
@@ -0,0 +1,119 @@
+#
+# Unary operator classes and methods
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+import pybamm
+import numbers
+import numpy as np
+
+
+class Broadcast(pybamm.SpatialOperator):
+    """A node in the expression tree representing a broadcasting operator.
+    Broadcasts a child (which *must* have empty domain) to a specified domain. After
+    discretisation, this will evaluate to an array of the right shape for the specified
+    domain.
+
+    Parameters
+    ----------
+    child : :class:`Symbol`
+        child node
+    domain : iterable of string
+        the domain to broadcast the child to
+    name : string
+        name of the node
+
+    **Extends:** :class:`SpatialOperator`
+    """
+
+    def __init__(self, child, domain, name=None):
+        # Convert child to Scalar if it is a number
+        if isinstance(child, numbers.Number):
+            child = pybamm.Scalar(child)
+        # Check domain
+        if child.domain not in [[], domain]:
+            raise pybamm.DomainError(
+                """
+                Domain of a broadcasted child must be []
+                or same as 'domain' but is '{}'
+                """.format(
+                    child.domain
+                )
+            )
+        if name is None:
+            name = "broadcast"
+        super().__init__(name, child)
+        # overwrite child domain ([]) with specified broadcasting domain
+        self.domain = domain
+
+
+class NumpyBroadcast(Broadcast):
+    """A node in the expression tree implementing a broadcasting operator using numpy.
+    Broadcasts a child (which *must* have empty domain) to a specified domain. To do
+    this, creates a np array of ones of the same shape as the submesh domain, and then
+    multiplies the child by that array upon evaluation
+
+    Parameters
+    ----------
+    child : :class:`Symbol`
+        child node
+    domain : iterable of string
+        the domain to broadcast the child to
+    mesh : :class:`pybamm.Mesh`
+        the mesh on which to broadcast
+
+    **Extends:** :class:`SpatialOperator`
+    """
+
+    def __init__(self, child, domain, mesh):
+        # Only accept a 'constant' input if it evaluates to a number (i.e. no vectors
+        # and matrices)
+        if child.is_constant() and not child.evaluates_to_number():
+            raise TypeError("cannot Broadcast a constant Vector or Matrix")
+
+        super().__init__(child, domain, name="numpy broadcast")
+        # determine broadcasting vector size (size 1 if the domain is empty)
+        if domain == []:
+            self.broadcasting_vector_size = 1
+        else:
+            self.broadcasting_vector_size = sum(
+                [mesh[dom].npts_for_broadcast for dom in domain]
+            )
+        # create broadcasting vector (vector of ones with shape determined by the
+        # domain)
+        self.broadcasting_vector = np.ones(self.broadcasting_vector_size)
+
+    def evaluate(self, t=None, y=None):
+        """ See :meth:`pybamm.Symbol.evaluate()`. """
+        child = self.children[0]
+        child_eval = child.evaluate(t, y)
+
+        # Different broadcasting based on the shape of child_eval
+        try:
+            child_eval_size = child_eval.size
+        except AttributeError:
+            child_eval_size = 0
+
+        if child_eval_size <= 1:
+            return child_eval * self.broadcasting_vector
+        if child_eval_size > 1:
+            # Possible shapes for a child with a shape:
+            # (n,) -> (e.g. time-like object) broadcast to (n, broadcasting_size)
+            # (1,n) -> (e.g. state-vector-like object) broadcast to
+            #          (n, broadcasting_size)
+            # (n,1) -> error
+            # (n,m) -> error
+            # (n,m,k,...) -> error
+            if child_eval.ndim == 1:
+                # shape (n,)
+                return np.repeat(
+                    child_eval[np.newaxis, :], self.broadcasting_vector_size, axis=0
+                )
+            elif child_eval.ndim == 2:
+                if child_eval.shape[0] == 1:
+                    # shape (1, m) since size > 1
+                    return np.repeat(child_eval, self.broadcasting_vector_size, axis=0)
+            # All other cases
+            raise ValueError(
+                "cannot broadcast child with shape '{}'".format(child_eval.shape)
+            )
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index c17c7c4d..9bb84ac7 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -6,7 +6,6 @@ from __future__ import print_function, unicode_literals
 import pybamm
 
 import numpy as np
-import warnings
 
 
 class Concatenation(pybamm.Symbol):
@@ -48,12 +47,13 @@ class Concatenation(pybamm.Symbol):
         return domain
 
 
-class NumpyModelConcatenation(pybamm.Symbol):
-    """A node in the expression tree representing a concatenation of equations.
+class NumpyConcatenation(pybamm.Symbol):
+    """A node in the expression tree representing a concatenation of equations, when we
+    *don't* care about domains. The class :class:`pybamm.DomainConcatenation`, which
+    *is* careful about domains and uses broadcasting where appropriate, should be used
+    whenever possible instead.
+
     Upon evaluation, equations are concatenated using numpy concatenation.
-    Unlike :class:`pybamm.Concatenation`, this doesn't check domains, as its only use
-    is to concatenate model equations (e.g. rhs equations or initial conditions, in
-    :class:`pybamm.Discretisation`), which might have common domains
 
     **Extends**: :class:`pybamm.Symbol`
 
@@ -80,7 +80,8 @@ class NumpyModelConcatenation(pybamm.Symbol):
 
 
 class DomainConcatenation(Concatenation):
-    """A node in the expression tree representing a concatenation of symbols.
+    """A node in the expression tree representing a concatenation of symbols, being
+    careful about domains.
 
     It is assumed that each child has a domain, and the final concatenated vector will
     respect the sizes and ordering of domains established in pybamm.KNOWN_DOMAINS
@@ -105,10 +106,6 @@ class DomainConcatenation(Concatenation):
 
         children = list(children)
 
-        for i, child in enumerate(children):
-            if child.is_constant() or child.evaluates_to_number():
-                children[i] = self.process_node_for_concatenate(child, mesh)
-
         # Allow the base class to sort the domains into the correct order
         super().__init__(*children, name="domain concatenation")
 
@@ -133,37 +130,8 @@ class DomainConcatenation(Concatenation):
             start = end
         return slices
 
-    def process_node_for_concatenate(self, node, mesh):
-        """
-        the node is assumed to be constant in time. this function replaces it with a
-        single Vector node with the correct length vector (according to its domain)
-
-        Parameters
-        ----------
-        node: derived from :class:`Symbol`
-            the sub-expression to process (node.is_constant() is true)
-
-        """
-        try:
-            node_size = node.size
-        except AttributeError:
-            node_size = 0
-
-        if node_size > 1:
-            if node_size == sum([mesh[dom].npts for dom in node.domain]):
-                return node
-            else:
-                raise ValueError(
-                    "Error: expression evaluated to a vector of incorrect length"
-                )
-        else:
-            warnings.warn("Implementation not robust")
-            npts = {dom: submesh.npts for dom, submesh in mesh.items()}
-            return pybamm.NumpyBroadcast(node, node.domain, npts)
-
     def evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
-
         # preallocate vector
         vector = np.empty(self._size)
 
@@ -174,31 +142,3 @@ class DomainConcatenation(Concatenation):
                 vector[self._slices[dom]] = child_vector[slices[dom]]
 
         return vector
-
-
-class PiecewiseConstant(Concatenation):
-    """Piecewise constant concatenation of three symbols.
-    This is useful when we don't want to assign a domain to the inputs
-
-    Parameters
-    ----------
-    neg_value: :class:`numbers.Number` or :class:`pybamm.Symbol`
-        The value in the negative electrode
-    sep_value: :class:`numbers.Number` or :class:`pybamm.Symbol`
-        The value in the separator
-    pos_value: :class:`numbers.Number` or :class:`pybamm.Symbol`
-        The value in the positive electrode
-
-    """
-
-    def __init__(self, neg_value, sep_value, pos_value):
-        neg_value_with_domain = neg_value * pybamm.Scalar(
-            1, domain=["negative electrode"]
-        )
-        sep_value_with_domain = sep_value * pybamm.Scalar(1, domain=["separator"])
-        pos_value_with_domain = pos_value * pybamm.Scalar(
-            1, domain=["positive electrode"]
-        )
-        super().__init__(
-            neg_value_with_domain, sep_value_with_domain, pos_value_with_domain
-        )
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 8a7eeff8..250bd3ab 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -246,6 +246,20 @@ class Symbol(anytree.NodeMixin):
         else:
             raise NotImplementedError
 
+    def __matmul__(self, other):
+        """return a :class:`MatrixMultiplication` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.MatrixMultiplication(self, other)
+        else:
+            raise NotImplementedError
+
+    def __rmatmul__(self, other):
+        """return a :class:`MatrixMultiplication` object"""
+        if isinstance(other, (Symbol, numbers.Number)):
+            return pybamm.MatrixMultiplication(other, self)
+        else:
+            raise NotImplementedError
+
     def __truediv__(self, other):
         """return a :class:`Division` object"""
         if isinstance(other, (Symbol, numbers.Number)):
@@ -324,7 +338,8 @@ class Symbol(anytree.NodeMixin):
 
     def evaluates_to_number(self):
         """Returns True if evaluating the expression returns a number.
-        Returns False otherwise, including if NotImplementedError is raised.
+        Returns False otherwise, including if NotImplementedError or TyperError
+        is raised.
         !Not to be confused with isinstance(self, pybamm.Scalar)!
 
         See Also
@@ -339,10 +354,13 @@ class Symbol(anytree.NodeMixin):
             # return false if NotImplementedError is raised
             # (there is a e.g. Parameter, Variable, ... in the tree)
             return False
-        except TypeError:
-            # return false if TypeError is raised
+        except TypeError as error:
+            # return false if specific TypeError is raised
             # (there is a e.g. StateVector in the tree)
-            return False
+            if error.args[0] == "StateVector cannot evaluate input 'y=None'":
+                return False
+            else:
+                raise error
 
     def has_spatial_derivatives(self):
         """Returns True if equation has spatial derivatives (grad or div)."""
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index cf42c215..88a939ce 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -128,91 +128,6 @@ class Divergence(SpatialOperator):
         super().__init__("div", child)
 
 
-class Broadcast(SpatialOperator):
-    """A node in the expression tree representing a broadcasting operator.
-    Broadcasts a child (which *must* have empty domain) to a specified domain. After
-    discretisation, this will evaluate to an array of the right shape for the specified
-    domain.
-
-    Parameters
-    ----------
-    child : :class:`Symbol`
-        child node
-    domain : iterable of string
-        the domain to broadcast the child to
-    name : string
-        name of the node
-
-    **Extends:** :class:`SpatialOperator`
-    """
-
-    def __init__(self, child, domain, name=None):
-        if child.domain not in [[], domain]:
-            raise pybamm.DomainError(
-                """
-                Domain of a broadcasted child must be []
-                or same as 'domain' but is '{}'
-                """.format(
-                    child.domain
-                )
-            )
-        if name is None:
-            name = "broadcast"
-        super().__init__(name, child)
-        # overwrite child domain ([]) with specified broadcasting domain
-        self.domain = domain
-
-
-class NumpyBroadcast(Broadcast):
-    """A node in the expression tree implementing a broadcasting operator using numpy.
-    Broadcasts a child (which *must* have empty domain) to a specified domain. To do
-    this, creates a np array of ones of the same shape as the submesh domain, and then
-    multiplies the child by that array upon evaluation
-
-    Parameters
-    ----------
-    child : :class:`Symbol`
-        child node
-    domain : iterable of string
-        the domain to broadcast the child to
-    npts : dict
-        contains the number of points on each domain
-
-    **Extends:** :class:`SpatialOperator`
-    """
-
-    def __init__(self, child, domain, npts):
-        super().__init__(child, domain, name="numpy broadcast")
-        # determine broadcasting vector size (size 1 if the domain is empty)
-        if domain == []:
-            self.broadcasting_vector_size = 1
-        else:
-            self.broadcasting_vector_size = sum([npts[dom] for dom in domain])
-        # create broadcasting vector (vector of ones with shape determined by the
-        # domain)
-        self.broadcasting_vector = np.ones(self.broadcasting_vector_size)
-
-    def evaluate(self, t=None, y=None):
-        """ See :meth:`pybamm.Symbol.evaluate()`. """
-        child = self.children[0]
-        child_eval = child.evaluate(t, y)
-        # if child is a vector, add a dimension for broadcasting
-        if isinstance(child, pybamm.Vector):
-            return child_eval[:, np.newaxis] * self.broadcasting_vector
-        # if child is a state vector, check that it has the right shape and then
-        # broadcast
-        elif isinstance(child, pybamm.StateVector):
-            assert child_eval.shape[0] == 1, ValueError(
-                """child_eval should have shape (1,n), not {}""".format(
-                    child_eval.shape
-                )
-            )
-            return np.repeat(child_eval, self.broadcasting_vector_size, axis=0)
-        # otherwise just do normal multiplication
-        else:
-            return child_eval * self.broadcasting_vector
-
-
 class SurfaceValue(SpatialOperator):
     """A node in the expression tree which gets the surface value of a variable.
 
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index 0fc0e06f..10f4eff7 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -27,7 +27,7 @@ class Vector(pybamm.Array):
         # make sure that entries are a vector
         if entries.ndim != 1:
             raise ValueError(
-                """Entries must have 1 dimension,  not {}""".format(entries.ndim)
+                """Entries must have 1 dimension, not {}""".format(entries.ndim)
             )
         if name is None:
             name = "Vector of length {!s}".format(entries.shape[0])
@@ -62,6 +62,13 @@ class StateVector(pybamm.Symbol):
         """Slice of an external y to read"""
         return self._y_slice
 
-    def evaluate(self, t, y):
+    def evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol.evaluate()`. """
-        return y[self._y_slice]
+        if y is None:
+            raise TypeError("StateVector cannot evaluate input 'y=None'")
+        if y.shape[0] < self.y_slice.stop:
+            raise ValueError(
+                "y is too short, so value with slice is smaller than expected"
+            )
+        else:
+            return y[self._y_slice]
diff --git a/pybamm/models/submodels/interface.py b/pybamm/models/submodels/interface.py
index 77dd3226..75fac64d 100644
--- a/pybamm/models/submodels/interface.py
+++ b/pybamm/models/submodels/interface.py
@@ -11,9 +11,15 @@ def homogeneous_reaction(current):
     Homogeneous reaction at the electrode-electrolyte interface
     """
     # hack to make the concatenation work. Concatenation needs some work
-    current_neg = current / pybamm.standard_parameters.ln
-    current_pos = -current / pybamm.standard_parameters.lp
-    return pybamm.PiecewiseConstant(current_neg, 0, current_pos)
+    current_neg = pybamm.Broadcast(
+        current / pybamm.standard_parameters.ln, ["negative electrode"]
+    )
+    current_pos = pybamm.Broadcast(
+        -current / pybamm.standard_parameters.lp, ["positive electrode"]
+    )
+    return pybamm.Concatenation(
+        current_neg, pybamm.Broadcast(0, ["separator"]), current_pos
+    )
 
 
 def butler_volmer_lead_acid(c, phi, domain=None):
@@ -82,7 +88,7 @@ def butler_volmer_lead_acid(c, phi, domain=None):
         # Negative electrode
         current_neg = butler_volmer_lead_acid(cn, phin, domain=["negative electrode"])
         # Separator
-        current_sep = pybamm.Scalar(0, domain=["separator"])
+        current_sep = pybamm.Broadcast(0, ["separator"])
         # Positive electrode
         current_pos = butler_volmer_lead_acid(cp, phip, domain=["positive electrode"])
         # Concatenate
diff --git a/pybamm/parameters/standard_parameters_lead_acid.py b/pybamm/parameters/standard_parameters_lead_acid.py
index c80a0d3d..f91b1a81 100644
--- a/pybamm/parameters/standard_parameters_lead_acid.py
+++ b/pybamm/parameters/standard_parameters_lead_acid.py
@@ -210,6 +210,18 @@ epsp_init = epsp_max - epsDelta_p * (1 - q_init)  # Initial pororsity (pos) [-]
 
 
 # Concatenated symbols
-s = pybamm.PiecewiseConstant(sn, 0, sp)
-beta_surf = pybamm.PiecewiseConstant(beta_surf_n, 0, beta_surf_p)
-eps_init = pybamm.PiecewiseConstant(epsn_init, epss_init, epsp_init)
+s = pybamm.Concatenation(
+    pybamm.Broadcast(sn, ["negative electrode"]),
+    pybamm.Broadcast(0, ["separator"]),
+    pybamm.Broadcast(sp, ["positive electrode"]),
+)
+beta_surf = pybamm.Concatenation(
+    pybamm.Broadcast(beta_surf_n, ["negative electrode"]),
+    pybamm.Broadcast(0, ["separator"]),
+    pybamm.Broadcast(beta_surf_p, ["positive electrode"]),
+)
+eps_init = pybamm.Concatenation(
+    pybamm.Broadcast(epsn_init, ["negative electrode"]),
+    pybamm.Broadcast(epss_init, ["separator"]),
+    pybamm.Broadcast(epsp_init, ["positive electrode"]),
+)
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index 18968ad5..c0553017 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -23,6 +23,9 @@ class FiniteVolume(pybamm.SpatialMethod):
     """
 
     def __init__(self, mesh):
+        # add npts_for_broadcast to mesh domains for this particular discretisation
+        for dom in mesh.keys():
+            mesh[dom].npts_for_broadcast = mesh[dom].npts
         super().__init__(mesh)
 
     def spatial_variable(self, symbol):
@@ -56,8 +59,7 @@ class FiniteVolume(pybamm.SpatialMethod):
         """
 
         # for finite volume we send variables to cells and so use number_of_cells
-        number_of_cells = {dom: submesh.npts for dom, submesh in self.mesh.items()}
-        broadcasted_symbol = pybamm.NumpyBroadcast(symbol, domain, number_of_cells)
+        broadcasted_symbol = pybamm.NumpyBroadcast(symbol, domain, self.mesh)
 
         # if the broadcasted symbol evaluates to a constant value, replace the
         # symbol-Vector multiplication with a single array
@@ -92,7 +94,7 @@ class FiniteVolume(pybamm.SpatialMethod):
 
         # note in 1D spherical grad and normal grad are the same
         gradient_matrix = self.gradient_matrix(domain)
-        return gradient_matrix * discretised_symbol
+        return gradient_matrix @ discretised_symbol
 
     def gradient_matrix(self, domain):
         """
@@ -137,9 +139,7 @@ class FiniteVolume(pybamm.SpatialMethod):
             # and also a "positive particle" left and right.
             lbc = boundary_conditions[symbol.id]["left"]
             rbc = boundary_conditions[symbol.id]["right"]
-            discretised_symbol = pybamm.NumpyModelConcatenation(
-                lbc, discretised_symbol, rbc
-            )
+            discretised_symbol = pybamm.NumpyConcatenation(lbc, discretised_symbol, rbc)
 
         domain = symbol.domain
         # check for spherical domains
@@ -153,12 +153,12 @@ class FiniteVolume(pybamm.SpatialMethod):
             r_edges = pybamm.Vector(submesh.edges)
 
             out = (1 / (r ** 2)) * (
-                divergence_matrix * ((r_edges ** 2) * discretised_symbol)
+                divergence_matrix @ ((r_edges ** 2) * discretised_symbol)
             )
 
         else:
             divergence_matrix = self.divergence_matrix(domain)
-            out = divergence_matrix * discretised_symbol
+            out = divergence_matrix @ discretised_symbol
         return out
 
     def divergence_matrix(self, domain):
@@ -235,7 +235,7 @@ class FiniteVolume(pybamm.SpatialMethod):
         last_node = pybamm.StateVector(slice(y_slice_stop - 1, y_slice_stop))
         right_ghost_cell = 2 * rbc - last_node
         # concatenate
-        return pybamm.NumpyModelConcatenation(
+        return pybamm.NumpyConcatenation(
             left_ghost_cell, discretised_symbol, right_ghost_cell
         )
 
@@ -263,7 +263,7 @@ class FiniteVolume(pybamm.SpatialMethod):
         y_slice_stop = discretised_symbol.y_slice.stop
         last_node = pybamm.StateVector(slice(y_slice_stop - 1, y_slice_stop))
         penultimate_node = pybamm.StateVector(slice(y_slice_stop - 2, y_slice_stop - 1))
-        surface_value = (last_node + (last_node - penultimate_node) / 2)
+        surface_value = last_node + (last_node - penultimate_node) / 2
         return surface_value
 
     #######################################################


Test patch:
diff --git a/tests/shared.py b/tests/shared.py
index 4b1f44e8..6333f29a 100644
--- a/tests/shared.py
+++ b/tests/shared.py
@@ -10,17 +10,18 @@ class SpatialMethodForTesting(pybamm.SpatialMethod):
     """Identity operators, no boundary conditions."""
 
     def __init__(self, mesh):
-        self._mesh = mesh
+        for dom in mesh.keys():
+            mesh[dom].npts_for_broadcast = mesh[dom].npts
         super().__init__(mesh)
 
     def spatial_variable(self, symbol):
-        symbol_mesh = self._mesh.combine_submeshes(*symbol.domain)
+        # for finite volume we use the cell centres
+        symbol_mesh = self.mesh.combine_submeshes(*symbol.domain)
         return pybamm.Vector(symbol_mesh.nodes)
 
     def broadcast(self, symbol, domain):
         # for finite volume we send variables to cells and so use number_of_cells
-        number_of_cells = {dom: submesh.npts for dom, submesh in self._mesh.items()}
-        broadcasted_symbol = pybamm.NumpyBroadcast(symbol, domain, number_of_cells)
+        broadcasted_symbol = pybamm.NumpyBroadcast(symbol, domain, self.mesh)
 
         # if the broadcasted symbol evaluates to a constant value, replace the
         # symbol-Vector multiplication with a single array
@@ -36,14 +37,14 @@ class SpatialMethodForTesting(pybamm.SpatialMethod):
         for domain in symbol.domain:
             n += self.mesh[domain].npts
         gradient_matrix = pybamm.Matrix(np.eye(n))
-        return gradient_matrix * discretised_symbol
+        return gradient_matrix @ discretised_symbol
 
     def divergence(self, symbol, discretised_symbol, boundary_conditions):
         n = 0
         for domain in symbol.domain:
             n += self.mesh[domain].npts
         divergence_matrix = pybamm.Matrix(np.eye(n))
-        return divergence_matrix * discretised_symbol
+        return divergence_matrix @ discretised_symbol
 
     def compute_diffusivity(self, symbol):
         return symbol
diff --git a/tests/test_discretisations/test_discretisation.py b/tests/test_discretisations/test_discretisation.py
index 938b6174..7327f60b 100644
--- a/tests/test_discretisations/test_discretisation.py
+++ b/tests/test_discretisations/test_discretisation.py
@@ -26,7 +26,7 @@ class TestDiscretise(unittest.TestCase):
         disc._y_slices = {c.id: slice(0, 1), a.id: slice(2, 3), b.id: slice(3, 4)}
         result = disc._concatenate_init(initial_conditions)
 
-        self.assertIsInstance(result, pybamm.NumpyModelConcatenation)
+        self.assertIsInstance(result, pybamm.NumpyConcatenation)
         self.assertEqual(result.children[0].evaluate(), 1)
         self.assertEqual(result.children[1].evaluate(), 2)
         self.assertEqual(result.children[2].evaluate(), 3)
@@ -182,7 +182,7 @@ class TestDiscretise(unittest.TestCase):
         for eqn in [pybamm.grad(var), pybamm.div(var)]:
             eqn_disc = disc.process_symbol(eqn)
 
-            self.assertIsInstance(eqn_disc, pybamm.Multiplication)
+            self.assertIsInstance(eqn_disc, pybamm.MatrixMultiplication)
             self.assertIsInstance(eqn_disc.children[0], pybamm.Matrix)
             self.assertIsInstance(eqn_disc.children[1], pybamm.StateVector)
 
@@ -200,7 +200,7 @@ class TestDiscretise(unittest.TestCase):
 
             self.assertIsInstance(eqn_disc, pybamm.Multiplication)
             self.assertIsInstance(eqn_disc.children[0], pybamm.StateVector)
-            self.assertIsInstance(eqn_disc.children[1], pybamm.Multiplication)
+            self.assertIsInstance(eqn_disc.children[1], pybamm.MatrixMultiplication)
             self.assertIsInstance(eqn_disc.children[1].children[0], pybamm.Matrix)
             self.assertIsInstance(eqn_disc.children[1].children[1], pybamm.StateVector)
 
@@ -441,7 +441,6 @@ class TestDiscretise(unittest.TestCase):
         whole_cell = ["negative electrode", "separator", "positive electrode"]
 
         a = pybamm.Scalar(7)
-        vec = pybamm.Vector(np.linspace(0, 1))
         var = pybamm.Variable("var")
 
         # create discretisation
@@ -458,28 +457,17 @@ class TestDiscretise(unittest.TestCase):
         )
         self.assertEqual(broad.domain, whole_cell)
 
-        # vector
-        broad = disc._spatial_methods[whole_cell[0]].broadcast(vec, ["separator"])
-        self.assertIsInstance(broad, pybamm.Array)
-        np.testing.assert_array_equal(
-            broad.evaluate(),
-            np.linspace(0, 1)[:, np.newaxis] * np.ones_like(mesh["separator"].nodes),
-        )
-        self.assertEqual(broad.domain, ["separator"])
+        broad_disc = disc.process_symbol(broad)
+        # type of broad will be array as broad is constant
+        self.assertIsInstance(broad_disc, pybamm.Array)
 
-        # process Broadcast symbol
+        # process Broadcast variable
         disc._y_slices = {var.id: slice(53)}
         broad1 = pybamm.Broadcast(var, ["negative electrode"])
         broad1_disc = disc.process_symbol(broad1)
         self.assertIsInstance(broad1_disc, pybamm.NumpyBroadcast)
         self.assertIsInstance(broad1_disc.children[0], pybamm.StateVector)
 
-        scal = pybamm.Scalar(3)
-        broad2 = pybamm.Broadcast(scal, ["negative electrode"])
-        broad2_disc = disc.process_symbol(broad2)
-        # type of broad2 will be array as broad2 is constant
-        self.assertIsInstance(broad2_disc, pybamm.Array)
-
     def test_concatenation(self):
         a = pybamm.Symbol("a")
         b = pybamm.Symbol("b")
@@ -489,12 +477,12 @@ class TestDiscretise(unittest.TestCase):
         disc = get_discretisation_for_testing()
 
         conc = disc.concatenate(a, b, c)
-        self.assertIsInstance(conc, pybamm.NumpyModelConcatenation)
+        self.assertIsInstance(conc, pybamm.NumpyConcatenation)
 
     def test_concatenation_of_scalars(self):
         whole_cell = ["negative electrode", "separator", "positive electrode"]
-        a = pybamm.Scalar(5, domain=["negative electrode"])
-        b = pybamm.Scalar(4, domain=["positive electrode"])
+        a = pybamm.Broadcast(5, ["negative electrode"])
+        b = pybamm.Broadcast(4, ["positive electrode"])
 
         # create discretisation
         disc = get_discretisation_for_testing()
diff --git a/tests/test_expression_tree/test_broadcasts.py b/tests/test_expression_tree/test_broadcasts.py
new file mode 100644
index 00000000..ed4828e3
--- /dev/null
+++ b/tests/test_expression_tree/test_broadcasts.py
@@ -0,0 +1,96 @@
+#
+# Tests for the Broadcast classes
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+import pybamm
+from tests import get_discretisation_for_testing
+import unittest
+import numpy as np
+
+
+class TestBroadcasts(unittest.TestCase):
+    def test_broadcast(self):
+        a = pybamm.Symbol("a")
+        broad_a = pybamm.Broadcast(a, ["negative electrode"])
+        self.assertEqual(broad_a.name, "broadcast")
+        self.assertEqual(broad_a.children[0].name, a.name)
+        self.assertEqual(broad_a.domain, ["negative electrode"])
+
+        b = pybamm.Symbol("b", domain=["negative electrode"])
+        with self.assertRaises(pybamm.DomainError):
+            pybamm.Broadcast(b, ["separator"])
+
+    def test_broadcast_number(self):
+        broad_a = pybamm.Broadcast(1, ["negative electrode"])
+        self.assertEqual(broad_a.name, "broadcast")
+        self.assertIsInstance(broad_a.children[0], pybamm.Symbol)
+        self.assertEqual(broad_a.children[0].name, str(1))
+        self.assertEqual(broad_a.domain, ["negative electrode"])
+
+        b = pybamm.Symbol("b", domain=["negative electrode"])
+        with self.assertRaises(pybamm.DomainError):
+            pybamm.Broadcast(b, ["separator"])
+
+    def test_numpy_broadcast(self):
+        # create discretisation
+        disc = get_discretisation_for_testing()
+        mesh = disc.mesh
+        for dom in mesh.keys():
+            mesh[dom].npts_for_broadcast = mesh[dom].npts
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        combined_submeshes = mesh.combine_submeshes(*whole_cell)
+
+        # scalar
+        a = pybamm.Scalar(7)
+        broad = pybamm.NumpyBroadcast(a, whole_cell, mesh)
+        np.testing.assert_array_equal(
+            broad.evaluate(), 7 * np.ones_like(combined_submeshes.nodes)
+        )
+        self.assertEqual(broad.domain, whole_cell)
+
+        # time
+        t = 3 * pybamm.t + 4
+        broad = pybamm.NumpyBroadcast(t, whole_cell, mesh)
+        self.assertEqual(broad.domain, whole_cell)
+        np.testing.assert_array_equal(
+            broad.evaluate(t=3), 13 * np.ones_like(combined_submeshes.nodes)
+        )
+        np.testing.assert_array_equal(
+            broad.evaluate(t=np.linspace(0, 1)),
+            (
+                (3 * np.linspace(0, 1) + 4)[:, np.newaxis]
+                * np.ones_like(combined_submeshes.nodes)
+            ).T,
+        )
+
+        # state vector
+        state_vec = pybamm.StateVector(slice(1, 2))
+        broad = pybamm.NumpyBroadcast(state_vec, whole_cell, mesh)
+        y = np.vstack([np.linspace(0, 1), np.linspace(0, 2)])
+        np.testing.assert_array_equal(
+            broad.evaluate(y=y), (y[1:2].T * np.ones_like(combined_submeshes.nodes)).T
+        )
+
+        # state vector - bad input
+        state_vec = pybamm.StateVector(slice(1, 5))
+        broad = pybamm.NumpyBroadcast(state_vec, whole_cell, mesh)
+        y = np.vstack([np.linspace(0, 1), np.linspace(0, 2)]).T
+        with self.assertRaisesRegex(ValueError, "cannot broadcast child with shape"):
+            broad.evaluate(y=y)
+
+        # vector - not accepted
+        vec = pybamm.Vector(np.ones(5))
+        with self.assertRaisesRegex(
+            TypeError, "cannot Broadcast a constant Vector or Matrix"
+        ):
+            broad = pybamm.NumpyBroadcast(vec, whole_cell, mesh)
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()
diff --git a/tests/test_expression_tree/test_concatenations.py b/tests/test_expression_tree/test_concatenations.py
index c351b1dd..e25515f9 100644
--- a/tests/test_expression_tree/test_concatenations.py
+++ b/tests/test_expression_tree/test_concatenations.py
@@ -43,36 +43,36 @@ class TestConcatenations(unittest.TestCase):
             ["negative electrode", "separator", "positive electrode", "test"],
         )
 
-    def test_numpy_model_concatenation_vectors(self):
+    def test_numpy_concatenation_vectors(self):
         # with entries
         y = np.linspace(0, 1, 15)
         a = pybamm.Vector(y[:5])
         b = pybamm.Vector(y[5:9])
         c = pybamm.Vector(y[9:])
-        conc = pybamm.NumpyModelConcatenation(a, b, c)
+        conc = pybamm.NumpyConcatenation(a, b, c)
         np.testing.assert_array_equal(conc.evaluate(None, y), y)
         # with y_slice
         a = pybamm.StateVector(slice(0, 10))
         b = pybamm.StateVector(slice(10, 15))
         c = pybamm.StateVector(slice(15, 23))
-        conc = pybamm.NumpyModelConcatenation(a, b, c)
+        conc = pybamm.NumpyConcatenation(a, b, c)
         y = np.linspace(0, 1, 23)
         np.testing.assert_array_equal(conc.evaluate(None, y), y)
 
-    def test_numpy_model_concatenation_vector_scalar(self):
+    def test_numpy_concatenation_vector_scalar(self):
         # with entries
         y = np.linspace(0, 1, 10)
         a = pybamm.Vector(y)
         b = pybamm.Scalar(16)
         c = pybamm.Scalar(3)
-        conc = pybamm.NumpyModelConcatenation(a, b, c)
+        conc = pybamm.NumpyConcatenation(a, b, c)
         np.testing.assert_array_equal(
             conc.evaluate(None, y), np.concatenate([y, np.array([16]), np.array([3])])
         )
 
         # with y_slice
         a = pybamm.StateVector(slice(0, 10))
-        conc = pybamm.NumpyModelConcatenation(a, b, c)
+        conc = pybamm.NumpyConcatenation(a, b, c)
         np.testing.assert_array_equal(
             conc.evaluate(None, y), np.concatenate([y, np.array([16]), np.array([3])])
         )
@@ -84,7 +84,7 @@ class TestConcatenations(unittest.TestCase):
 
         a_dom = ["negative electrode"]
         b_dom = ["positive electrode"]
-        a = pybamm.Scalar(2, domain=a_dom)
+        a = pybamm.NumpyBroadcast(pybamm.Scalar(2, a_dom), a_dom, mesh)
         b = pybamm.Vector(np.ones_like(mesh[b_dom[0]].nodes), domain=b_dom)
 
         # concatenate them the "wrong" way round to check they get reordered correctly
@@ -96,15 +96,10 @@ class TestConcatenations(unittest.TestCase):
             ),
         )
 
-        # vector child of wrong size will throw
-        b = pybamm.Vector(np.full(mesh[b_dom[0]].npts - 5, 1), domain=b_dom)
-        with self.assertRaises(ValueError):
-            conc = pybamm.DomainConcatenation([b, a], mesh)
-
         # check the reordering in case a child vector has to be split up
         a_dom = ["separator"]
         b_dom = ["negative electrode", "positive electrode"]
-        a = pybamm.Scalar(2, domain=a_dom)
+        a = pybamm.NumpyBroadcast(pybamm.Scalar(2, a_dom), a_dom, mesh)
         b = pybamm.Vector(
             np.concatenate(
                 [np.full(mesh[b_dom[0]].npts, 1), np.full(mesh[b_dom[1]].npts, 3)]
@@ -142,22 +137,107 @@ class TestConcatenations(unittest.TestCase):
         self.assertEqual(c.id, c_new.id)
         self.assertEqual(conc.id, pybamm.Concatenation(a_new, b_new, c_new).id)
 
-    def test_piecewise_constant(self):
-        a = pybamm.Scalar(1)
-        b = pybamm.Scalar(2)
-        c = pybamm.Scalar(3)
-        conc = pybamm.PiecewiseConstant(a, b, c)
+    def test_broadcast_and_concatenate(self):
+        # create discretisation
+        disc = get_discretisation_for_testing()
+        mesh = disc.mesh
+
+        # Piecewise constant scalars
+        a = pybamm.Broadcast(1, domain=["negative electrode"])
+        b = pybamm.Broadcast(2, domain=["separator"])
+        c = pybamm.Broadcast(3, domain=["positive electrode"])
+        conc = pybamm.Concatenation(a, b, c)
+
+        self.assertEqual(
+            conc.domain, ["negative electrode", "separator", "positive electrode"]
+        )
+        self.assertEqual(conc.children[0].domain, ["negative electrode"])
+        self.assertEqual(conc.children[1].domain, ["separator"])
+        self.assertEqual(conc.children[2].domain, ["positive electrode"])
+        processed_conc = disc.process_symbol(conc)
+        np.testing.assert_array_equal(
+            processed_conc.evaluate(),
+            np.concatenate(
+                [
+                    np.ones(mesh["negative electrode"].npts),
+                    2 * np.ones(mesh["separator"].npts),
+                    3 * np.ones(mesh["positive electrode"].npts),
+                ]
+            ),
+        )
+
+        # Piecewise constant functions of time
+        a_t = pybamm.Broadcast(pybamm.t, domain=["negative electrode"])
+        b_t = pybamm.Broadcast(2 * pybamm.t, domain=["separator"])
+        c_t = pybamm.Broadcast(3 * pybamm.t, domain=["positive electrode"])
+        conc = pybamm.Concatenation(a_t, b_t, c_t)
+
+        self.assertEqual(
+            conc.domain, ["negative electrode", "separator", "positive electrode"]
+        )
+        self.assertEqual(conc.children[0].domain, ["negative electrode"])
+        self.assertEqual(conc.children[1].domain, ["separator"])
+        self.assertEqual(conc.children[2].domain, ["positive electrode"])
+
+        processed_conc = disc.process_symbol(conc)
+        np.testing.assert_array_equal(
+            processed_conc.evaluate(t=2),
+            np.concatenate(
+                [
+                    2 * np.ones(mesh["negative electrode"].npts),
+                    4 * np.ones(mesh["separator"].npts),
+                    6 * np.ones(mesh["positive electrode"].npts),
+                ]
+            ),
+        )
+
+        # Piecewise constant state vectors
+        a_sv = pybamm.Broadcast(pybamm.StateVector(slice(0, 1)), ["negative electrode"])
+        b_sv = pybamm.Broadcast(pybamm.StateVector(slice(1, 2)), ["separator"])
+        c_sv = pybamm.Broadcast(pybamm.StateVector(slice(2, 3)), ["positive electrode"])
+        conc = pybamm.Concatenation(a_sv, b_sv, c_sv)
 
-        self.assertIsInstance(conc, pybamm.Concatenation)
         self.assertEqual(
             conc.domain, ["negative electrode", "separator", "positive electrode"]
         )
         self.assertEqual(conc.children[0].domain, ["negative electrode"])
         self.assertEqual(conc.children[1].domain, ["separator"])
         self.assertEqual(conc.children[2].domain, ["positive electrode"])
-        self.assertEqual(conc.children[0].evaluate(), 1)
-        self.assertEqual(conc.children[1].evaluate(), 2)
-        self.assertEqual(conc.children[2].evaluate(), 3)
+
+        processed_conc = disc.process_symbol(conc)
+        y = np.array([1, 2, 3])
+        np.testing.assert_array_equal(
+            processed_conc.evaluate(y=y),
+            np.concatenate(
+                [
+                    np.ones(mesh["negative electrode"].npts),
+                    2 * np.ones(mesh["separator"].npts),
+                    3 * np.ones(mesh["positive electrode"].npts),
+                ]
+            ),
+        )
+
+        # Mixed
+        conc = pybamm.Concatenation(a, b_t, c_sv)
+
+        self.assertEqual(
+            conc.domain, ["negative electrode", "separator", "positive electrode"]
+        )
+        self.assertEqual(conc.children[0].domain, ["negative electrode"])
+        self.assertEqual(conc.children[1].domain, ["separator"])
+        self.assertEqual(conc.children[2].domain, ["positive electrode"])
+
+        processed_conc = disc.process_symbol(conc)
+        np.testing.assert_array_equal(
+            processed_conc.evaluate(t=2, y=y),
+            np.concatenate(
+                [
+                    np.ones(mesh["negative electrode"].npts),
+                    4 * np.ones(mesh["separator"].npts),
+                    3 * np.ones(mesh["positive electrode"].npts),
+                ]
+            ),
+        )
 
 
 if __name__ == "__main__":
diff --git a/tests/test_expression_tree/test_matrix.py b/tests/test_expression_tree/test_matrix.py
index 65a73f38..a399665b 100644
--- a/tests/test_expression_tree/test_matrix.py
+++ b/tests/test_expression_tree/test_matrix.py
@@ -28,11 +28,11 @@ class TestMatrix(unittest.TestCase):
         np.testing.assert_array_equal((self.mat + self.mat).evaluate(), 2 * self.A)
         np.testing.assert_array_equal((self.mat - self.mat).evaluate(), 0 * self.A)
         np.testing.assert_array_equal(
-            (self.mat * self.vect).evaluate(), np.array([5, 2, 3])
+            (self.mat @ self.vect).evaluate(), np.array([5, 2, 3])
         )
 
     def test_matrix_modification(self):
-        exp = self.mat * self.mat + self.mat
+        exp = self.mat @ self.mat + self.mat
         self.A[0, 0] = -1
         self.assertTrue(exp.children[1]._entries[0, 0], -1)
         self.assertTrue(exp.children[0].children[0]._entries[0, 0], -1)
diff --git a/tests/test_expression_tree/test_symbol.py b/tests/test_expression_tree/test_symbol.py
index e35c0749..1ebe056c 100644
--- a/tests/test_expression_tree/test_symbol.py
+++ b/tests/test_expression_tree/test_symbol.py
@@ -41,6 +41,7 @@ class TestSymbol(unittest.TestCase):
         self.assertIsInstance(a + b, pybamm.Addition)
         self.assertIsInstance(a - b, pybamm.Subtraction)
         self.assertIsInstance(a * b, pybamm.Multiplication)
+        self.assertIsInstance(a @ b, pybamm.MatrixMultiplication)
         self.assertIsInstance(a / b, pybamm.Division)
         self.assertIsInstance(a ** b, pybamm.Power)
 
@@ -48,6 +49,7 @@ class TestSymbol(unittest.TestCase):
         self.assertIsInstance(a + 2, pybamm.Addition)
         self.assertIsInstance(a - 2, pybamm.Subtraction)
         self.assertIsInstance(a * 2, pybamm.Multiplication)
+        self.assertIsInstance(a @ 2, pybamm.MatrixMultiplication)
         self.assertIsInstance(a / 2, pybamm.Division)
         self.assertIsInstance(a ** 2, pybamm.Power)
 
@@ -58,6 +60,8 @@ class TestSymbol(unittest.TestCase):
         self.assertEqual((3 - b).children[1].id, b.id)
         self.assertIsInstance(3 * b, pybamm.Multiplication)
         self.assertEqual((3 * b).children[1].id, b.id)
+        self.assertIsInstance(3 @ b, pybamm.MatrixMultiplication)
+        self.assertEqual((3 @ b).children[1].id, b.id)
         self.assertIsInstance(3 / b, pybamm.Division)
         self.assertEqual((3 / b).children[1].id, b.id)
         self.assertIsInstance(3 ** b, pybamm.Power)
@@ -70,6 +74,8 @@ class TestSymbol(unittest.TestCase):
             a - "two"
         with self.assertRaises(NotImplementedError):
             a * "two"
+        with self.assertRaises(NotImplementedError):
+            a @ "two"
         with self.assertRaises(NotImplementedError):
             a / "two"
         with self.assertRaises(NotImplementedError):
@@ -80,6 +86,8 @@ class TestSymbol(unittest.TestCase):
             "two" - a
         with self.assertRaises(NotImplementedError):
             "two" * a
+        with self.assertRaises(NotImplementedError):
+            "two" @ a
         with self.assertRaises(NotImplementedError):
             "two" / a
         with self.assertRaises(NotImplementedError):
@@ -158,6 +166,9 @@ class TestSymbol(unittest.TestCase):
         a = pybamm.Matrix(np.ones((4, 6)))
         self.assertFalse(a.evaluates_to_number())
 
+        a = pybamm.StateVector(slice(0, 10))
+        self.assertFalse(a.evaluates_to_number())
+
         # Time variable returns true
         a = 3 * pybamm.t + 2
         self.assertTrue(a.evaluates_to_number())
diff --git a/tests/test_expression_tree/test_unary_operators.py b/tests/test_expression_tree/test_unary_operators.py
index 452833d6..afe3bf9e 100644
--- a/tests/test_expression_tree/test_unary_operators.py
+++ b/tests/test_expression_tree/test_unary_operators.py
@@ -4,7 +4,7 @@
 from __future__ import absolute_import, division
 from __future__ import print_function, unicode_literals
 import pybamm
-from tests import get_discretisation_for_testing
+
 import unittest
 import numpy as np
 
@@ -86,58 +86,6 @@ class TestUnaryOperators(unittest.TestCase):
         un5 = pybamm.UnaryOperator("test", d)
         self.assertNotEqual(un1.id, un5.id)
 
-    def test_broadcast(self):
-        a = pybamm.Symbol("a")
-        broad_a = pybamm.Broadcast(a, ["negative electrode"])
-        self.assertEqual(broad_a.name, "broadcast")
-        self.assertEqual(broad_a.children[0].name, a.name)
-        self.assertEqual(broad_a.domain, ["negative electrode"])
-
-        b = pybamm.Symbol("b", domain=["negative electrode"])
-        with self.assertRaises(pybamm.DomainError):
-            pybamm.Broadcast(b, ["separator"])
-
-    def test_numpy_broadcast(self):
-        # create discretisation
-        disc = get_discretisation_for_testing()
-        mesh = disc.mesh
-
-        whole_cell = ["negative electrode", "separator", "positive electrode"]
-        combined_submeshes = mesh.combine_submeshes(*whole_cell)
-        # scalar
-        a = pybamm.Scalar(7)
-        npts = {dom: submesh.npts for dom, submesh in mesh.items()}
-        broad = pybamm.NumpyBroadcast(a, whole_cell, npts)
-        np.testing.assert_array_equal(
-            broad.evaluate(), 7 * np.ones_like(combined_submeshes.nodes)
-        )
-        self.assertEqual(broad.domain, whole_cell)
-
-        # vector
-        vec = pybamm.Vector(np.linspace(0, 1))
-        broad = pybamm.NumpyBroadcast(vec, whole_cell, npts)
-        np.testing.assert_array_equal(
-            broad.evaluate(),
-            np.linspace(0, 1)[:, np.newaxis] * np.ones_like(combined_submeshes.nodes),
-        )
-
-        self.assertEqual(broad.domain, whole_cell)
-
-        # state vector
-        state_vec = pybamm.StateVector(slice(1, 2))
-        broad = pybamm.NumpyBroadcast(state_vec, whole_cell, npts)
-        y = np.vstack([np.linspace(0, 1), np.linspace(0, 2)])
-        np.testing.assert_array_equal(
-            broad.evaluate(y=y), (y[1:2].T * np.ones_like(combined_submeshes.nodes)).T
-        )
-
-        # state vector - bad input
-        state_vec = pybamm.StateVector(slice(1, 5))
-        broad = pybamm.NumpyBroadcast(state_vec, whole_cell, npts)
-        y = np.vstack([np.linspace(0, 1), np.linspace(0, 2)]).T
-        with self.assertRaises(AssertionError):
-            broad.evaluate(y=y)
-
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_expression_tree/test_vector.py b/tests/test_expression_tree/test_vector.py
index bef5c96a..58f0104e 100644
--- a/tests/test_expression_tree/test_vector.py
+++ b/tests/test_expression_tree/test_vector.py
@@ -36,6 +36,36 @@ class TestVector(unittest.TestCase):
         self.assertTrue(exp.children[0].children[0]._entries[0], -1)
         self.assertTrue(exp.children[0].children[1]._entries[0], -1)
 
+    def test_wrong_size_entries(self):
+        with self.assertRaisesRegex(ValueError, "Entries must have 1 dimension, not 2"):
+            pybamm.Vector(np.ones((4, 5)))
+
+
+class TestStateVector(unittest.TestCase):
+    def test_evaluate(self):
+        sv = pybamm.StateVector(slice(0, 10))
+        y = np.linspace(0, 2, 19)
+        np.testing.assert_array_equal(sv.evaluate(y=y), np.linspace(0, 1, 10))
+
+        # Try evaluating with a y that is too short
+        y2 = np.ones(5)
+        with self.assertRaisesRegex(
+            ValueError, "y is too short, so value with slice is smaller than expected"
+        ):
+            sv.evaluate(y=y2)
+
+    def test_evaluate_2D(self):
+        sv = pybamm.StateVector(slice(0, 10))
+        y = np.ones((20, 40))
+        np.testing.assert_array_equal(sv.evaluate(y=y), np.ones((10, 40)))
+
+        # Try evaluating with a y that is too short
+        y2 = np.ones((5, 40))
+        with self.assertRaisesRegex(
+            ValueError, "y is too short, so value with slice is smaller than expected"
+        ):
+            sv.evaluate(y=y2)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/test_models/test_interface.py b/tests/test_models/test_interface.py
index b23ced54..a2a84a9c 100644
--- a/tests/test_models/test_interface.py
+++ b/tests/test_models/test_interface.py
@@ -42,18 +42,12 @@ class TestHomogeneousReaction(unittest.TestCase):
         self.assertEqual(processed_rxn.children[1].domain, ["separator"])
         self.assertEqual(processed_rxn.children[2].domain, ["positive electrode"])
 
-        # test values
-        ln = param.process_symbol(pybamm.standard_parameters.ln)
-        lp = param.process_symbol(pybamm.standard_parameters.lp)
-        self.assertEqual(processed_rxn.children[0].evaluate() * ln.evaluate(), 1)
-        self.assertEqual(processed_rxn.children[1].evaluate(), 0)
-        self.assertEqual(processed_rxn.children[2].evaluate() * lp.evaluate(), -1)
-
     def test_discretisation(self):
         param = pybamm.ParameterValues(
             "input/parameters/lithium-ion/parameters/LCO.csv"
         )
         disc = get_discretisation_for_testing()
+        mesh = disc.mesh
 
         rxn = pybamm.interface.homogeneous_reaction(1)
 
@@ -67,6 +61,21 @@ class TestHomogeneousReaction(unittest.TestCase):
         self.assertIsInstance(processed_rxn, pybamm.Vector)
         self.assertEqual(processed_rxn.shape, combined_submeshes.nodes.shape)
 
+        # test values
+        ln = param.process_symbol(pybamm.standard_parameters.ln)
+        lp = param.process_symbol(pybamm.standard_parameters.lp)
+        npts_n = mesh["negative electrode"].npts
+        npts_s = mesh["separator"].npts
+        np.testing.assert_array_equal(
+            processed_rxn.evaluate()[:npts_n] * ln.evaluate(), 1
+        )
+        np.testing.assert_array_equal(
+            processed_rxn.evaluate()[npts_n : npts_n + npts_s], 0
+        )
+        np.testing.assert_array_equal(
+            processed_rxn.evaluate()[npts_n + npts_s :] * lp.evaluate(), -1
+        )
+
 
 class TestButlerVolmerLeadAcid(unittest.TestCase):
     def setUp(self):
@@ -183,7 +192,6 @@ class TestButlerVolmerLeadAcid(unittest.TestCase):
             mesh["positive electrode"].nodes.shape,
         )
 
-    @unittest.skip("needs smarter domain concatenation")
     def test_discretisation_whole(self):
         bv_whole = pybamm.interface.butler_volmer_lead_acid(self.c, self.phi)
 
diff --git a/tests/test_spatial_methods/test_spatial_methods.py b/tests/test_spatial_methods/test_spatial_methods.py
index e69de29b..51f59f0a 100644
--- a/tests/test_spatial_methods/test_spatial_methods.py
+++ b/tests/test_spatial_methods/test_spatial_methods.py
@@ -0,0 +1,35 @@
+#
+# Test for the operator class
+#
+import pybamm
+from tests import get_mesh_for_testing
+
+import unittest
+
+
+class TestSpatialMethod(unittest.TestCase):
+    def test_basics(self):
+        mesh = get_mesh_for_testing()
+        spatial_method = pybamm.SpatialMethod(mesh)
+        self.assertEqual(spatial_method.mesh, mesh)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.spatial_variable(None)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.broadcast(None, None)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.gradient(None, None, None)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.divergence(None, None, None)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.surface_value(None)
+        with self.assertRaises(NotImplementedError):
+            spatial_method.compute_diffusivity()
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()