Repository: pybamm-team/PyBaMM

Issue:
discretising 2D bug
when objects that are concatenated are discretised in 2D, they come out the right shape but wrong order
to do:
- [x] state vector
- [x] domain concatenation
- [x] discretisation
- [x] jacobians
- [x] evaluators

Discussion:

Patch:
diff --git a/docs/source/expression_tree/index.rst b/docs/source/expression_tree/index.rst
index bac81d756..8c2099ac2 100644
--- a/docs/source/expression_tree/index.rst
+++ b/docs/source/expression_tree/index.rst
@@ -10,6 +10,7 @@ Expression Tree
   scalar
   matrix
   vector
+  state_vector
   binary_operator
   unary_operator
   concatenations
diff --git a/docs/source/expression_tree/state_vector.rst b/docs/source/expression_tree/state_vector.rst
new file mode 100644
index 000000000..4a0792b3b
--- /dev/null
+++ b/docs/source/expression_tree/state_vector.rst
@@ -0,0 +1,5 @@
+State Vector
+============
+
+.. autoclass:: pybamm.StateVector
+  :members:
diff --git a/docs/source/expression_tree/vector.rst b/docs/source/expression_tree/vector.rst
index de412c8e8..27528dd92 100644
--- a/docs/source/expression_tree/vector.rst
+++ b/docs/source/expression_tree/vector.rst
@@ -3,6 +3,3 @@ Vector
 
 .. autoclass:: pybamm.Vector
   :members:
-
-.. autoclass:: pybamm.StateVector
-  :members:
diff --git a/examples/notebooks/expression_tree/expression-tree.ipynb b/examples/notebooks/expression_tree/expression-tree.ipynb
index 961e0f7a2..04fd8f737 100644
--- a/examples/notebooks/expression_tree/expression-tree.ipynb
+++ b/examples/notebooks/expression_tree/expression-tree.ipynb
@@ -13,7 +13,7 @@
     "2. To perform the time-stepping, it must be able to be evaluated, given the current state vector $\\mathbf{y}$ and the current time $t$\n",
     "3. For solvers that require it, its gradient with respect to a given variable must be able to be evaluated (once again given $\\mathbf{y}$ and $t$)\n",
     "\n",
-    "As an initial example, the code below shows how to construct an expression tree of the equation $2y(1 - y) + t$. We use the `pybamm.StateVector` to represent $\\mathbf{y}$, which in this case will be a vector of size 1. The time variable $t$ is already provided by PyBaMM and is of class `pybamm.IndependentVariable`."
+    "As an initial example, the code below shows how to construct an expression tree of the equation $2y(1 - y) + t$. We use the `pybamm.StateVector` to represent $\\mathbf{y}$, which in this case will be a vector of size 1. The time variable $t$ is already provided by PyBaMM and is of class `pybamm.Time`."
    ]
   },
   {
@@ -189,7 +189,7 @@
    "source": [
     "from tests import get_discretisation_for_testing\n",
     "disc = get_discretisation_for_testing()\n",
-    "disc._y_slices = {c.id: slice(0, 40)}\n",
+    "disc.y_slices = {c.id: [slice(0, 40)]}\n",
     "dcdt = disc.process_symbol(dcdt)\n",
     "dcdt.visualise('expression_tree5.png')"
    ]
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index cb5133c68..d00027663 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -113,7 +113,8 @@ from .expression_tree.independent_variable import (
     SpatialVariable,
 )
 from .expression_tree.independent_variable import t
-from .expression_tree.vector import Vector, StateVector
+from .expression_tree.vector import Vector
+from .expression_tree.state_vector import StateVector
 
 from .expression_tree.exceptions import (
     DomainError,
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index 122b7647b..f6cad3892 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -2,8 +2,8 @@
 # Interface for discretisation
 #
 import pybamm
-
 import numpy as np
+from collections import defaultdict, OrderedDict
 from scipy.sparse import block_diag, csr_matrix
 
 
@@ -35,13 +35,28 @@ class Discretisation(object):
                 dom: method(mesh) for dom, method in spatial_methods.items()
             }
         self.bcs = {}
-        self._y_slices = {}
+        self.y_slices = {}
         self._discretised_symbols = {}
 
     @property
     def mesh(self):
         return self._mesh
 
+    @property
+    def y_slices(self):
+        return self._y_slices
+
+    @y_slices.setter
+    def y_slices(self, value):
+        if not isinstance(value, dict):
+            raise TypeError("""y_slices should be dict, not {}""".format(type(value)))
+
+        self._y_slices = value
+
+    @property
+    def spatial_methods(self):
+        return self._spatial_methods
+
     @property
     def bcs(self):
         return self._bcs
@@ -145,57 +160,47 @@ class Discretisation(object):
 
         return model_disc
 
-    def _flatten_variables(self, variables):
-        """
-        Extract variable objects from concatenations
-
-        Parameters
-        ----------
-        variables : iterable of :class:`pybamm.Symbol`
-            Iterable of symbols from which to extract variables
-
-        Returns
-        -------
-        flattened_variables : list of :class:`pybamm.Variable`
-            List of variable objects
-        """
-        # Unpack symbols in variables that are concatenations of variables
-        flattened_variables = []
-        for symbol in variables:
-            if isinstance(symbol, pybamm.Concatenation):
-                flattened_variables.extend([child for child in symbol.children])
-            else:
-                flattened_variables.append(symbol)
-        return flattened_variables
-
     def set_variable_slices(self, variables):
         """Sets the slicing for variables.
 
         variables : iterable of :class:`pybamm.Variables`
-            The variables for which to set slices
+        The variables for which to set slices
         """
-        flattened_variables = self._flatten_variables(variables)
         # Set up y_slices
-        y_slices = {variable.id: None for variable in flattened_variables}
+        y_slices = defaultdict(list)
         start = 0
         end = 0
-        # Iterate through flattened variables, adding appropriate slices to y_slices
-        for variable in flattened_variables:
+        # Iterate through unpacked variables, adding appropriate slices to y_slices
+        for variable in variables:
             # If domain is empty then variable has size 1
             if variable.domain == []:
                 end += 1
+                y_slices[variable.id].append(slice(start, end))
+                start = end
             # Otherwise, add up the size of all the domains in variable.domain
+            elif isinstance(variable, pybamm.Concatenation):
+                children = variable.children
+                meshes = OrderedDict()
+                for child in children:
+                    meshes[child] = [
+                        self.spatial_methods[dom].mesh[dom] for dom in child.domain
+                    ]
+                sec_points = len(list(meshes.values())[0][0])
+                for i in range(sec_points):
+                    for child, mesh in meshes.items():
+                        for domain_mesh in mesh:
+                            submesh = domain_mesh[i]
+                            end += submesh.npts_for_broadcast
+                        y_slices[child.id].append(slice(start, end))
+                        start = end
             else:
                 for dom in variable.domain:
-                    for submesh in self._spatial_methods[dom].mesh[dom]:
+                    for submesh in self.spatial_methods[dom].mesh[dom]:
                         end += submesh.npts_for_broadcast
-            y_slices[variable.id] = slice(start, end)
-            start = end
-        self._y_slices = y_slices
+                y_slices[variable.id].append(slice(start, end))
+                start = end
 
-        assert isinstance(self._y_slices, dict), ValueError(
-            """y_slices should be dict, not {}""".format(type(self._y_slices))
-        )
+        self.y_slices = y_slices
 
         # reset discretised_symbols
         self._discretised_symbols = {}
@@ -217,12 +222,13 @@ class Discretisation(object):
             left_domain = left_symbol.domain[0]
             right_domain = right_symbol.domain[0]
 
-            left_mesh = self._spatial_methods[left_domain].mesh[left_domain]
-            right_mesh = self._spatial_methods[right_domain].mesh[right_domain]
+            left_mesh = self.spatial_methods[left_domain].mesh[left_domain]
+            right_mesh = self.spatial_methods[right_domain].mesh[right_domain]
 
             left_symbol_disc = self.process_symbol(left_symbol)
             right_symbol_disc = self.process_symbol(right_symbol)
-            return self._spatial_methods[left_domain].internal_neumann_condition(
+
+            return self.spatial_methods[left_domain].internal_neumann_condition(
                 left_symbol_disc, right_symbol_disc, left_mesh, right_mesh
             )
 
@@ -394,9 +400,18 @@ class Discretisation(object):
 
         # get a list of model rhs variables that are sorted according to
         # where they are in the state vector
-        model_variables = self._flatten_variables(model.rhs.keys())
-        model_slices = [self._y_slices[var.id] for var in model_variables]
-
+        model_variables = model.rhs.keys()
+        model_slices = []
+        for v in model_variables:
+            if isinstance(v, pybamm.Concatenation):
+                model_slices.append(
+                    slice(
+                        self.y_slices[v.children[0].id][0].start,
+                        self.y_slices[v.children[-1].id][0].stop,
+                    )
+                )
+            else:
+                model_slices.append(self.y_slices[v.id][0])
         sorted_model_variables = [
             v for _, v in sorted(zip(model_slices, model_variables))
         ]
@@ -408,7 +423,7 @@ class Discretisation(object):
                 mass_list.append(1.0)
             else:
                 mass_list.append(
-                    self._spatial_methods[var.domain[0]]
+                    self.spatial_methods[var.domain[0]]
                     .mass_matrix(var, self.bcs)
                     .entries
                 )
@@ -484,7 +499,7 @@ class Discretisation(object):
         """ See :meth:`Discretisation.process_symbol()`. """
 
         if symbol.domain != []:
-            spatial_method = self._spatial_methods[symbol.domain[0]]
+            spatial_method = self.spatial_methods[symbol.domain[0]]
 
         if isinstance(symbol, pybamm.BinaryOperator):
             # Pre-process children
@@ -502,7 +517,7 @@ class Discretisation(object):
             child = symbol.child
             disc_child = self.process_symbol(child)
             if child.domain != []:
-                child_spatial_method = self._spatial_methods[child.domain[0]]
+                child_spatial_method = self.spatial_methods[child.domain[0]]
             if isinstance(symbol, pybamm.Gradient):
                 return child_spatial_method.gradient(child, disc_child, self.bcs)
 
@@ -549,7 +564,7 @@ class Discretisation(object):
 
         elif isinstance(symbol, pybamm.Variable):
             return pybamm.StateVector(
-                self._y_slices[symbol.id],
+                *self._y_slices[symbol.id],
                 domain=symbol.domain,
                 auxiliary_domains=symbol.auxiliary_domains,
             )
@@ -577,12 +592,12 @@ class Discretisation(object):
 
     def _concatenate_in_order(self, var_eqn_dict, check_complete=False):
         """
-        Concatenate a dictionary of {variable: equation} using self._y_slices
+        Concatenate a dictionary of {variable: equation} using self.y_slices
 
         The keys/variables in `var_eqn_dict` must be the same as the ids in
-        `self._y_slices`.
+        `self.y_slices`.
         The resultant concatenation is ordered according to the ordering of the slice
-        values in `self._y_slices`
+        values in `self.y_slices`
 
         Parameters
         ----------
@@ -605,18 +620,18 @@ class Discretisation(object):
                 # get sorted correctly
                 slices.append(
                     slice(
-                        self._y_slices[symbol.children[0].id].start,
-                        self._y_slices[symbol.children[-1].id].stop,
+                        self.y_slices[symbol.children[0].id][0].start,
+                        self.y_slices[symbol.children[-1].id][0].stop,
                     )
                 )
             else:
                 unpacked_variables.append(symbol)
-                slices.append(self._y_slices[symbol.id])
+                slices.append(self.y_slices[symbol.id][0])
 
         if check_complete:
-            # Check keys from the given var_eqn_dict against self._y_slices
+            # Check keys from the given var_eqn_dict against self.y_slices
             ids = {v.id for v in unpacked_variables}
-            if ids != self._y_slices.keys():
+            if ids != self.y_slices.keys():
                 given_variable_names = [v.name for v in var_eqn_dict.keys()]
                 raise pybamm.ModelError(
                     "Initial conditions are insufficient. Only "
diff --git a/pybamm/expression_tree/array.py b/pybamm/expression_tree/array.py
index cd560a70e..55852d327 100644
--- a/pybamm/expression_tree/array.py
+++ b/pybamm/expression_tree/array.py
@@ -47,11 +47,6 @@ class Array(pybamm.Symbol):
         """ returns the number of entries along each dimension"""
         return self._entries.shape
 
-    @property
-    def size(self):
-        """ returns the total number of entries in the tensor"""
-        return self._entries.size
-
     @property
     def entries_string(self):
         return self._entries_string
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 412184054..3b74f869c 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -357,10 +357,10 @@ class Multiplication(BinaryOperator):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
 
         if issparse(left):
-            return left.multiply(right)
+            return csr_matrix(left.multiply(right))
         elif issparse(right):
             # Hadamard product is commutative, so we can switch right and left
-            return right.multiply(left)
+            return csr_matrix(right.multiply(left))
         else:
             return left * right
 
@@ -482,7 +482,7 @@ class Division(BinaryOperator):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
 
         if issparse(left):
-            return left.multiply(1 / right)
+            return csr_matrix(left.multiply(1 / right))
         else:
             if isinstance(right, numbers.Number) and right == 0:
                 return left * np.inf
@@ -690,7 +690,7 @@ class Kron(BinaryOperator):
 
     def _binary_evaluate(self, left, right):
         """ See :meth:`pybamm.BinaryOperator._binary_evaluate()`. """
-        return kron(left, right)
+        return csr_matrix(kron(left, right))
 
     def _binary_simplify(self, left, right):
         """ See :meth:`pybamm.BinaryOperator.simplify()`. """
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index a91efa305..2c5c59aa9 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -1,11 +1,11 @@
 #
 # Concatenation classes
 #
-import pybamm
-
+import copy
 import numpy as np
+import pybamm
 from scipy.sparse import vstack
-import copy
+from collections import defaultdict
 
 
 class Concatenation(pybamm.Symbol):
@@ -223,10 +223,11 @@ class DomainConcatenation(Concatenation):
                 )
 
             # create dict of domain => slice of final vector
+            self.secondary_dimensions_npts = len(self.mesh[self.domain[0]])
             self._slices = self.create_slices(self)
 
             # store size of final vector
-            self._size = self._slices[self.domain[-1]].stop
+            self._size = self._slices[self.domain[-1]][-1].stop
 
             # create disc of domain => slice for each child
             self._children_slices = [
@@ -237,21 +238,27 @@ class DomainConcatenation(Concatenation):
             self._slices = copy.copy(copy_this._slices)
             self._size = copy.copy(copy_this._size)
             self._children_slices = copy.copy(copy_this._children_slices)
+            self.secondary_dimensions_npts = copy_this.secondary_dimensions_npts
 
     @property
     def mesh(self):
         return self._mesh
 
     def create_slices(self, node):
-        slices = {}
+        slices = defaultdict(list)
         start = 0
         end = 0
-        for dom in node.domain:
-            prim_pts = self.mesh[dom][0].npts
-            second_pts = len(self.mesh[dom])
-            end += prim_pts * second_pts
-            slices[dom] = slice(start, end)
-            start = end
+        second_pts = len(self.mesh[node.domain[0]])
+        if second_pts != self.secondary_dimensions_npts:
+            raise ValueError(
+                """Concatenation and children must have the same number of
+                points in secondary dimensions"""
+            )
+        for i in range(second_pts):
+            for dom in node.domain:
+                end += self.mesh[dom][i].npts
+                slices[dom].append(slice(start, end))
+                start = end
         return slices
 
     def _concatenation_evaluate(self, children_eval):
@@ -262,17 +269,27 @@ class DomainConcatenation(Concatenation):
         # loop through domains of children writing subvectors to final vector
         for child_vector, slices in zip(children_eval, self._children_slices):
             for child_dom, child_slice in slices.items():
-                vector[self._slices[child_dom]] = child_vector[child_slice]
+                for i, _slice in enumerate(child_slice):
+                    vector[self._slices[child_dom][i]] = child_vector[_slice]
 
         return vector
 
     def _jac(self, variable):
         """ See :meth:`pybamm.Symbol._jac()`. """
-        children = self.cached_children
-        if len(children) == 0:
-            return pybamm.Scalar(0)
-        else:
-            return SparseStack(*[child.jac(variable) for child in children])
+        # note that this assumes that the children are in the right order and only have
+        # one domain each
+        jacs = []
+        child_jacs = [child.jac(variable) for child in self.cached_children]
+        for i in range(self.secondary_dimensions_npts):
+            for child_jac, slices in zip(child_jacs, self._children_slices):
+                if len(slices) > 1:
+                    raise NotImplementedError(
+                        """jacobian only implemented for when each child has
+                        a single domain"""
+                    )
+                child_slice = next(iter(slices.values()))
+                jacs.append(child_jac[child_slice[i]])
+        return SparseStack(*jacs)
 
     def _concatenation_new_copy(self, children):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
@@ -282,15 +299,21 @@ class DomainConcatenation(Concatenation):
     def _concatenation_simplify(self, children):
         """ See :meth:`pybamm.Symbol.simplify()`. """
         # Simplify Concatenation of StateVectors to a single StateVector
-        if all([isinstance(x, pybamm.StateVector) for x in children]) and all(
-            [
-                children[idx].y_slice.stop == children[idx + 1].y_slice.start
-                for idx in range(len(children) - 1)
-            ]
-        ):
-            return pybamm.StateVector(
-                slice(children[0].y_slice.start, children[-1].y_slice.stop)
-            )
+        # The sum of the evalation arrays of the StateVectors must be exactly 1
+        if all([isinstance(child, pybamm.StateVector) for child in children]):
+            longest_eval_array = len(children[-1]._evaluation_array)
+            eval_arrays = {}
+            for child in children:
+                eval_arrays[child] = np.concatenate(
+                    [
+                        child.evaluation_array,
+                        np.zeros(longest_eval_array - len(child.evaluation_array)),
+                    ]
+                )
+            if all(sum(array for array in eval_arrays.values()) == 1):
+                return pybamm.StateVector(
+                    slice(children[0].y_slices[0].start, children[-1].y_slices[-1].stop)
+                )
 
         new_symbol = self.__class__(children, self.mesh, self)
 
diff --git a/pybamm/expression_tree/evaluate.py b/pybamm/expression_tree/evaluate.py
index 745cc810a..4539e8b7e 100644
--- a/pybamm/expression_tree/evaluate.py
+++ b/pybamm/expression_tree/evaluate.py
@@ -72,13 +72,16 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
         # the right line, avoiding these checks
         if isinstance(symbol, pybamm.Multiplication):
             symbol_str = (
-                "{0}.multiply({1}) if scipy.sparse.issparse({0}) else "
-                "{1}.multiply({0}) if scipy.sparse.issparse({1}) else "
+                "scipy.sparse.csr_matrix({0}.multiply({1})) "
+                "if scipy.sparse.issparse({0}) else "
+                "scipy.sparse.csr_matrix({1}.multiply({0})) "
+                "if scipy.sparse.issparse({1}) else "
                 "{0} * {1}".format(children_vars[0], children_vars[1])
             )
         elif isinstance(symbol, pybamm.Division):
             symbol_str = (
-                "{0}.multiply(1/{1}) if scipy.sparse.issparse({0}) else "
+                "scipy.sparse.csr_matrix({0}.multiply(1/{1})) "
+                "if scipy.sparse.issparse({0}) else "
                 "{0} / {1}".format(children_vars[0], children_vars[1])
             )
         elif isinstance(symbol, pybamm.Outer):
@@ -86,7 +89,7 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
                 children_vars[0], children_vars[1]
             )
         elif isinstance(symbol, pybamm.Kron):
-            symbol_str = "scipy.sparse.kron({}, {})".format(
+            symbol_str = "scipy.sparse.csr_matrix(scipy.sparse.kron({}, {}))".format(
                 children_vars[0], children_vars[1]
             )
         else:
@@ -134,18 +137,22 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
         # which we must follow
         elif isinstance(symbol, pybamm.DomainConcatenation):
             slice_starts = []
-            child_vectors = []
-            for child_var, slices in zip(children_vars, symbol._children_slices):
-                for child_dom, child_slice in slices.items():
-                    slice_starts.append(symbol._slices[child_dom].start)
-                    child_vectors.append(
-                        "{}[{}:{}]".format(
-                            child_var, child_slice.start, child_slice.stop
+            all_child_vectors = []
+            for i in range(symbol.secondary_dimensions_npts):
+                child_vectors = []
+                for child_var, slices in zip(children_vars, symbol._children_slices):
+                    for child_dom, child_slice in slices.items():
+                        slice_starts.append(symbol._slices[child_dom][i].start)
+                        child_vectors.append(
+                            "{}[{}:{}]".format(
+                                child_var, child_slice[i].start, child_slice[i].stop
+                            )
                         )
-                    )
-            child_vectors = [v for _, v in sorted(zip(slice_starts, child_vectors))]
-            if len(children_vars) > 1:
-                symbol_str = "np.concatenate(({}))".format(",".join(child_vectors))
+                all_child_vectors.extend(
+                    [v for _, v in sorted(zip(slice_starts, child_vectors))]
+                )
+            if len(children_vars) > 1 or symbol.secondary_dimensions_npts > 1:
+                symbol_str = "np.concatenate(({}))".format(",".join(all_child_vectors))
             else:
                 symbol_str = "{}".format(",".join(children_vars))
         else:
@@ -153,7 +160,9 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
 
     # Note: we assume that y is being passed as a column vector
     elif isinstance(symbol, pybamm.StateVector):
-        symbol_str = symbol.name
+        symbol_str = "y[:{}][{}]".format(
+            len(symbol.evaluation_array), symbol.evaluation_array
+        )
 
     elif isinstance(symbol, pybamm.Time):
         symbol_str = "t"
diff --git a/pybamm/expression_tree/functions.py b/pybamm/expression_tree/functions.py
index 7625ad73e..f7a508a46 100644
--- a/pybamm/expression_tree/functions.py
+++ b/pybamm/expression_tree/functions.py
@@ -95,10 +95,7 @@ class Function(pybamm.Symbol):
         if all(child.evaluates_to_number() for child in self.children):
             # if children all evaluate to numbers the return zeros
             # of right size
-            variable_y_indices = np.arange(
-                variable.y_slice.start, variable.y_slice.stop
-            )
-            jac = csr_matrix((1, np.size(variable_y_indices)))
+            jac = csr_matrix((1, variable.evaluation_array.count(True)))
             jacobian = pybamm.Matrix(jac)
         else:
 
diff --git a/pybamm/expression_tree/simplify.py b/pybamm/expression_tree/simplify.py
index d816ab4c7..97ac0ab59 100644
--- a/pybamm/expression_tree/simplify.py
+++ b/pybamm/expression_tree/simplify.py
@@ -17,12 +17,12 @@ def simplify_if_constant(symbol):
         result = symbol.evaluate_ignoring_errors()
         if result is not None:
             if isinstance(result, numbers.Number):
-                return pybamm.Scalar(result, domain=symbol.domain)
+                return pybamm.Scalar(result)
             elif isinstance(result, np.ndarray) or issparse(result):
                 if result.ndim == 1 or result.shape[1] == 1:
-                    return pybamm.Vector(result, domain=symbol.domain)
+                    return pybamm.Vector(result)
                 else:
-                    return pybamm.Matrix(result, domain=symbol.domain)
+                    return pybamm.Matrix(result)
 
     return symbol
 
@@ -578,6 +578,8 @@ class Simplification(object):
 
     def _simplify(self, symbol):
         """ See :meth:`Simplification.simplify()`. """
+        symbol.domain = []
+        symbol.auxiliary_domains = {}
 
         if isinstance(symbol, pybamm.BinaryOperator):
             left, right = symbol.children
@@ -605,13 +607,10 @@ class Simplification(object):
             # Backup option: return new copy of the object
             try:
                 new_symbol = symbol.new_copy()
-                new_symbol.domain = []
                 return new_symbol
             except NotImplementedError:
                 raise NotImplementedError(
                     "Cannot simplify symbol of type '{}'".format(type(symbol))
                 )
 
-        new_symbol.domain = []
-        new_symbol.auxiliary_domains = {}
         return simplify_if_constant(new_symbol)
diff --git a/pybamm/expression_tree/state_vector.py b/pybamm/expression_tree/state_vector.py
new file mode 100644
index 000000000..1023764a6
--- /dev/null
+++ b/pybamm/expression_tree/state_vector.py
@@ -0,0 +1,168 @@
+#
+# State Vector class
+#
+import pybamm
+
+import numpy as np
+from scipy.sparse import csr_matrix, vstack
+
+
+class StateVector(pybamm.Symbol):
+    """
+    node in the expression tree that holds a slice to read from an external vector type
+
+    Parameters
+    ----------
+
+    y_slice: slice
+        the slice of an external y to read
+    name: str, optional
+        the name of the node
+    domain : iterable of str, optional
+        list of domains the parameter is valid over, defaults to empty list
+    auxiliary_domains : dict of str, optional
+        dictionary of auxiliary domains
+
+    *Extends:* :class:`Array`
+    """
+
+    def __init__(self, *y_slices, name=None, domain=None, auxiliary_domains=None):
+        for y_slice in y_slices:
+            if not isinstance(y_slice, slice):
+                raise TypeError("all y_slices must be slice objects")
+        if name is None:
+            if y_slices[0].start is None:
+                name = "y[:{:d}]".format(y_slice.stop)
+            else:
+                name = "y[{:d}:{:d}".format(y_slices[0].start, y_slices[0].stop)
+            if len(y_slices) > 1:
+                name += ",{:d}:{:d}".format(y_slices[1].start, y_slices[1].stop)
+                if len(y_slices) > 2:
+                    name += ",...,{:d}:{:d}]".format(
+                        y_slices[-1].start, y_slices[-1].stop
+                    )
+                else:
+                    name += "]"
+            else:
+                name += "]"
+        if domain is None:
+            domain = []
+        if auxiliary_domains is None:
+            auxiliary_domains = {}
+        self._y_slices = y_slices
+        self._first_point = y_slices[0].start
+        self._last_point = y_slices[-1].stop
+        self.set_evaluation_array(y_slices)
+        super().__init__(name=name, domain=domain, auxiliary_domains=auxiliary_domains)
+
+    @property
+    def y_slices(self):
+        return self._y_slices
+
+    @property
+    def first_point(self):
+        return self._first_point
+
+    @property
+    def last_point(self):
+        return self._last_point
+
+    @property
+    def evaluation_array(self):
+        """Array to use for evaluating"""
+        return self._evaluation_array
+
+    @property
+    def size(self):
+        return self.evaluation_array.count(True)
+
+    def set_evaluation_array(self, y_slices):
+        "Set evaluation array using slices"
+        array = np.zeros(y_slices[-1].stop)
+        for y_slice in y_slices:
+            array[y_slice] = True
+        self._evaluation_array = [bool(x) for x in array]
+
+    def set_id(self):
+        """ See :meth:`pybamm.Symbol.set_id()` """
+        self._id = hash(
+            (self.__class__, self.name, tuple(self.evaluation_array))
+            + tuple(self.domain)
+        )
+
+    def _base_evaluate(self, t=None, y=None):
+        """ See :meth:`pybamm.Symbol._base_evaluate()`. """
+        if y is None:
+            raise TypeError("StateVector cannot evaluate input 'y=None'")
+        if y.shape[0] < len(self.evaluation_array):
+            raise ValueError(
+                "y is too short, so value with slice is smaller than expected"
+            )
+        else:
+            out = (y[: len(self._evaluation_array)])[self._evaluation_array]
+            if out.ndim == 1:
+                out = out[:, np.newaxis]
+            return out
+
+    def jac(self, variable):
+        """
+        Differentiate a slice of a StateVector of size m with respect to another
+        slice of a StateVector of size n. This returns a (sparse) matrix of size
+        m x n with ones where the y slices match, and zeros elsewhere.
+
+        Parameters
+        ----------
+        variable : :class:`pybamm.Symbol`
+            The variable with respect to which to differentiate
+
+        """
+        if len(variable.y_slices) > 1:
+            raise NotImplementedError(
+                "Jacobian only implemented for a single-slice StateVector"
+            )
+        variable_y_indices = np.arange(variable.first_point, variable.last_point)
+
+        jac = csr_matrix((0, np.size(variable_y_indices)))
+        for y_slice in self.y_slices:
+            # Get indices of state vectors
+            slice_indices = np.arange(y_slice.start, y_slice.stop)
+
+            # Return zeros of correct size if no entries match
+            if np.size(np.intersect1d(slice_indices, variable_y_indices)) == 0:
+                jac = csr_matrix((np.size(slice_indices), np.size(variable_y_indices)))
+            else:
+                # Populate entries corresponding to matching y slices, and shift so
+                # that the matrix is the correct size
+                row = np.intersect1d(slice_indices, variable_y_indices) - y_slice.start
+                col = (
+                    np.intersect1d(slice_indices, variable_y_indices)
+                    - variable.first_point
+                )
+                data = np.ones_like(row)
+                jac = vstack(
+                    [
+                        jac,
+                        csr_matrix(
+                            (data, (row, col)),
+                            shape=(np.size(slice_indices), np.size(variable_y_indices)),
+                        ),
+                    ]
+                )
+        return pybamm.Matrix(jac)
+
+    def new_copy(self):
+        """ See :meth:`pybamm.Symbol.new_copy()`. """
+        return StateVector(
+            *self.y_slices,
+            name=self.name,
+            domain=self.domain,
+            auxiliary_domains=self.auxiliary_domains
+        )
+
+    def evaluate_for_shape(self):
+        """
+        Returns a vector of NaNs to represent the shape of a StateVector.
+        The size of a StateVector is the number of True elements in its evaluation_array
+        See :meth:`pybamm.Symbol.evaluate_for_shape()`
+        """
+        return np.nan * np.ones((self.size, 1))
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 0fe5908e3..c9914911e 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -615,9 +615,14 @@ class Symbol(anytree.NodeMixin):
             state_vectors_in_node = [
                 x for x in self.pre_order() if isinstance(x, pybamm.StateVector)
             ]
-            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)
+            if state_vectors_in_node == []:
+                y = None
+            else:
+                min_y_size = max(
+                    len(x._evaluation_array) 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(0, y)
 
         # Return shape of evaluated object
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index da99d7b1f..5192cacce 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -191,10 +191,7 @@ class Index(UnaryOperator):
         # tree for StateVectors and return a matrix of zeros of the correct size
         # if none are found.
         if all([not (isinstance(n, pybamm.StateVector)) for n in self.pre_order()]):
-            variable_y_indices = np.arange(
-                variable.y_slice.start, variable.y_slice.stop
-            )
-            jac = csr_matrix((1, np.size(variable_y_indices)))
+            jac = csr_matrix((1, variable.evaluation_array.count(True)))
             return pybamm.Matrix(jac)
         else:
             child_jac = self.child.jac(variable)
@@ -223,7 +220,7 @@ class Index(UnaryOperator):
         return self.__class__(child, self.index)
 
     def evaluate_for_shape(self):
-        return self.children[0].evaluate_for_shape()[self.slice]
+        return self._unary_evaluate(self.children[0].evaluate_for_shape())
 
     def evaluates_on_edges(self):
         """ See :meth:`pybamm.Symbol.evaluates_on_edges()`. """
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index 787b56bf5..8fcebbb5d 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -1,5 +1,5 @@
 #
-# Vector classes
+# Vector class
 #
 import pybamm
 
@@ -43,114 +43,6 @@ class Vector(pybamm.Array):
 
     def _jac(self, variable):
         """ See :meth:`pybamm.Symbol._jac()`. """
-        # Get indices of state vector
-        variable_y_indices = np.arange(variable.y_slice.start, variable.y_slice.stop)
         # Return zeros of correct size
-        jac = csr_matrix((np.size(self), np.size(variable_y_indices)))
+        jac = csr_matrix((self.size, variable.evaluation_array.count(True)))
         return pybamm.Matrix(jac)
-
-
-class StateVector(pybamm.Symbol):
-    """
-    node in the expression tree that holds a slice to read from an external vector type
-
-    Parameters
-    ----------
-
-    y_slice: slice
-        the slice of an external y to read
-    name: str, optional
-        the name of the node
-    domain : iterable of str, optional
-        list of domains the parameter is valid over, defaults to empty list
-
-    *Extends:* :class:`Array`
-    """
-
-    def __init__(self, y_slice, name=None, domain=None, auxiliary_domains=None):
-        if name is None:
-            if y_slice.start is None:
-                name = "y[:{:d}]".format(y_slice.stop)
-            else:
-                name = "y[{:d}:{:d}]".format(y_slice.start, y_slice.stop)
-        if domain is None:
-            domain = []
-        if auxiliary_domains is None:
-            auxiliary_domains = {}
-        self._y_slice = y_slice
-        super().__init__(name=name, domain=domain, auxiliary_domains=auxiliary_domains)
-
-    @property
-    def y_slice(self):
-        """Slice of an external y to read"""
-        return self._y_slice
-
-    def set_id(self):
-        """ See :meth:`pybamm.Symbol.set_id()` """
-        self._id = hash(
-            (self.__class__, self.name, self.y_slice.start, self.y_slice.stop)
-            + tuple(self.domain)
-        )
-
-    def _base_evaluate(self, t=None, y=None):
-        """ See :meth:`pybamm.Symbol._base_evaluate()`. """
-        if y is None:
-            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:
-            out = y[self._y_slice]
-            if out.ndim == 1:
-                out = out[:, np.newaxis]
-            return out
-
-    def jac(self, variable):
-        """
-        Differentiate a slice of a StateVector of size m with respect to another
-        slice of a StateVector of size n. This returns a (sparse) matrix of size
-        m x n with ones where the y slices match, and zeros elsewhere.
-
-        Parameters
-        ----------
-        variable : :class:`pybamm.Symbol`
-            The variable with respect to which to differentiate
-
-        """
-
-        # Get inices of state vectors
-        self_y_indices = np.arange(self.y_slice.start, self.y_slice.stop)
-        variable_y_indices = np.arange(variable.y_slice.start, variable.y_slice.stop)
-
-        # Return zeros of correct size if no entries match
-        if np.size(np.intersect1d(self_y_indices, variable_y_indices)) == 0:
-            jac = csr_matrix((np.size(self_y_indices), np.size(variable_y_indices)))
-        else:
-            # Populate entries corresponding to matching y slices, and shift so
-            # that the matrix is the correct size
-            row = (
-                np.intersect1d(self_y_indices, variable_y_indices) - self.y_slice.start
-            )
-            col = (
-                np.intersect1d(self_y_indices, variable_y_indices)
-                - variable.y_slice.start
-            )
-            data = np.ones_like(row)
-            jac = csr_matrix(
-                (data, (row, col)),
-                shape=(np.size(self_y_indices), np.size(variable_y_indices)),
-            )
-        return pybamm.Matrix(jac)
-
-    def new_copy(self):
-        """ See :meth:`pybamm.Symbol.new_copy()`. """
-        return StateVector(self.y_slice, self.name, domain=self.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/spatial_methods/spatial_method.py b/pybamm/spatial_methods/spatial_method.py
index 54b55ce03..868661266 100644
--- a/pybamm/spatial_methods/spatial_method.py
+++ b/pybamm/spatial_methods/spatial_method.py
@@ -249,11 +249,12 @@ class SpatialMethod:
             # coo_matrix takes inputs (data, (row, col)) and puts data[i] at the point
             # (row[i], col[i]) for each index of data. Here we just want a single point
             # with value 1 at (0,0).
-            left_vector = coo_matrix(([1], ([0], [0])), shape=(1, n))
+            # Convert to a csr_matrix to allow indexing and other functionality
+            left_vector = csr_matrix(coo_matrix(([1], ([0], [0])), shape=(1, n)))
             bv_vector = pybamm.Matrix(left_vector)
         elif symbol.side == "right":
             # 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))
+            right_vector = csr_matrix(coo_matrix(([1], ([0], [n - 1])), shape=(1, n)))
             bv_vector = pybamm.Matrix(right_vector)
 
         out = bv_vector @ discretised_child


Test patch:
diff --git a/tests/shared.py b/tests/shared.py
index 35bf20c9e..98c5a6a48 100644
--- a/tests/shared.py
+++ b/tests/shared.py
@@ -95,9 +95,11 @@ def get_p2d_mesh_for_testing(xpts=None, rpts=10):
     return get_mesh_for_testing(xpts=xpts, rpts=rpts, geometry=geometry)
 
 
-def get_1p1d_mesh_for_testing(xpts=None, zpts=15):
+def get_1p1d_mesh_for_testing(xpts=None, zpts=15, cc_submesh=pybamm.Uniform1DSubMesh):
     geometry = pybamm.Geometry("1+1D macro")
-    return get_mesh_for_testing(xpts=xpts, zpts=zpts, geometry=geometry)
+    return get_mesh_for_testing(
+        xpts=xpts, zpts=zpts, geometry=geometry, cc_submesh=cc_submesh
+    )
 
 
 def get_2p1d_mesh_for_testing(
diff --git a/tests/unit/test_discretisations/test_discretisation.py b/tests/unit/test_discretisations/test_discretisation.py
index 6631f0e47..5f3bb7de0 100644
--- a/tests/unit/test_discretisations/test_discretisation.py
+++ b/tests/unit/test_discretisations/test_discretisation.py
@@ -30,7 +30,7 @@ class TestDiscretise(unittest.TestCase):
         # create discretisation
         disc = get_discretisation_for_testing()
 
-        disc._y_slices = {c.id: slice(0, 1), a.id: slice(2, 3), b.id: slice(3, 4)}
+        disc.y_slices = {c.id: [slice(0, 1)], a.id: [slice(2, 3)], b.id: [slice(3, 4)]}
         result = disc._concatenate_in_order(initial_conditions)
 
         self.assertIsInstance(result, pybamm.NumpyConcatenation)
@@ -77,13 +77,13 @@ class TestDiscretise(unittest.TestCase):
         variables = [c]
         disc.set_variable_slices(variables)
 
-        self.assertEqual(disc._y_slices, {c.id: slice(0, 100)})
+        self.assertEqual(disc.y_slices, {c.id: [slice(0, 100)]})
 
         combined_submesh = mesh.combine_submeshes(*whole_cell)
 
         c_true = combined_submesh[0].nodes ** 2
         y = c_true
-        np.testing.assert_array_equal(y[disc._y_slices[c.id]], c_true)
+        np.testing.assert_array_equal(y[disc.y_slices[c.id][0]], c_true)
 
         # Several variables
         d = pybamm.Variable("d", domain=whole_cell)
@@ -92,15 +92,15 @@ class TestDiscretise(unittest.TestCase):
         disc.set_variable_slices(variables)
 
         self.assertEqual(
-            disc._y_slices,
-            {c.id: slice(0, 100), d.id: slice(100, 200), jn.id: slice(200, 240)},
+            disc.y_slices,
+            {c.id: [slice(0, 100)], d.id: [slice(100, 200)], jn.id: [slice(200, 240)]},
         )
         d_true = 4 * combined_submesh[0].nodes
         jn_true = mesh["negative electrode"][0].nodes ** 3
         y = np.concatenate([c_true, d_true, jn_true])
-        np.testing.assert_array_equal(y[disc._y_slices[c.id]], c_true)
-        np.testing.assert_array_equal(y[disc._y_slices[d.id]], d_true)
-        np.testing.assert_array_equal(y[disc._y_slices[jn.id]], jn_true)
+        np.testing.assert_array_equal(y[disc.y_slices[c.id][0]], c_true)
+        np.testing.assert_array_equal(y[disc.y_slices[d.id][0]], d_true)
+        np.testing.assert_array_equal(y[disc.y_slices[jn.id][0]], jn_true)
 
         # Variables with a concatenation
         js = pybamm.Variable("js", domain=["separator"])
@@ -109,21 +109,24 @@ class TestDiscretise(unittest.TestCase):
         variables = [c, d, j]
         disc.set_variable_slices(variables)
         self.assertEqual(
-            disc._y_slices,
+            disc.y_slices,
             {
-                c.id: slice(0, 100),
-                d.id: slice(100, 200),
-                jn.id: slice(200, 240),
-                js.id: slice(240, 265),
-                jp.id: slice(265, 300),
+                c.id: [slice(0, 100)],
+                d.id: [slice(100, 200)],
+                jn.id: [slice(200, 240)],
+                js.id: [slice(240, 265)],
+                jp.id: [slice(265, 300)],
             },
         )
         d_true = 4 * combined_submesh[0].nodes
         jn_true = mesh["negative electrode"][0].nodes ** 3
         y = np.concatenate([c_true, d_true, jn_true])
-        np.testing.assert_array_equal(y[disc._y_slices[c.id]], c_true)
-        np.testing.assert_array_equal(y[disc._y_slices[d.id]], d_true)
-        np.testing.assert_array_equal(y[disc._y_slices[jn.id]], jn_true)
+        np.testing.assert_array_equal(y[disc.y_slices[c.id][0]], c_true)
+        np.testing.assert_array_equal(y[disc.y_slices[d.id][0]], d_true)
+        np.testing.assert_array_equal(y[disc.y_slices[jn.id][0]], jn_true)
+
+        with self.assertRaisesRegex(TypeError, "y_slices should be"):
+            disc.y_slices = 1
 
     def test_process_symbol_base(self):
         # create discretisation
@@ -139,10 +142,10 @@ 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, 93)}
+        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])
+        self.assertEqual(var_disc.y_slices[0], disc.y_slices[var.id][0])
         # scalar
         scal = pybamm.Scalar(5)
         scal_disc = disc.process_symbol(scal)
@@ -236,7 +239,7 @@ class TestDiscretise(unittest.TestCase):
         # create discretisation
         disc = get_discretisation_for_testing()
 
-        disc._y_slices = {var1.id: slice(53), var2.id: slice(53, 106)}
+        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
@@ -250,8 +253,8 @@ class TestDiscretise(unittest.TestCase):
             isinstance(exp_disc.children[0].children[1].children[1], pybamm.StateVector)
         )
         self.assertEqual(
-            exp_disc.children[0].children[1].children[1].y_slice,
-            disc._y_slices[var2.id],
+            exp_disc.children[0].children[1].children[1].y_slices[0],
+            disc.y_slices[var2.id][0],
         )
         # right side
         self.assertIsInstance(exp_disc.children[1], pybamm.Addition)
@@ -262,8 +265,8 @@ class TestDiscretise(unittest.TestCase):
             isinstance(exp_disc.children[1].children[0].children[0], pybamm.StateVector)
         )
         self.assertEqual(
-            exp_disc.children[1].children[0].children[0].y_slice,
-            disc._y_slices[var1.id],
+            exp_disc.children[1].children[0].children[0].y_slices[0],
+            disc.y_slices[var1.id][0],
         )
         self.assertTrue(
             isinstance(exp_disc.children[1].children[0].children[1], pybamm.Scalar)
@@ -362,10 +365,10 @@ class TestDiscretise(unittest.TestCase):
         # rhs
         processed_rhs = disc.process_dict(rhs)
         np.testing.assert_array_equal(
-            y[disc._y_slices[c.id]], processed_rhs[c].evaluate(None, y)
+            y[disc.y_slices[c.id][0]], processed_rhs[c].evaluate(None, y)
         )
         np.testing.assert_array_equal(
-            y[disc._y_slices[T.id]], processed_rhs[T].evaluate(None, y)
+            y[disc.y_slices[T.id][0]], processed_rhs[T].evaluate(None, y)
         )
         # initial conditions
         y0 = disc.process_dict(initial_conditions)
@@ -454,7 +457,7 @@ class TestDiscretise(unittest.TestCase):
         disc.process_model(model)
         y0 = model.concatenated_initial_conditions
         y0_expect = np.empty((0, 1))
-        for var_id, _ in sorted(disc._y_slices.items(), key=lambda kv: kv[1]):
+        for var_id, _ in sorted(disc.y_slices.items(), key=lambda kv: kv[1]):
             if var_id == c.id:
                 vect = 2 * np.ones_like(combined_submesh[0].nodes[:, np.newaxis])
             elif var_id == T.id:
@@ -684,7 +687,7 @@ class TestDiscretise(unittest.TestCase):
         self.assertIsInstance(broad_disc.children[1], pybamm.Vector)
 
         # process Broadcast variable
-        disc._y_slices = {var.id: slice(1)}
+        disc.y_slices = {var.id: [slice(1)]}
         broad1 = pybamm.Broadcast(var, ["negative electrode"])
         broad1_disc = disc.process_symbol(broad1)
         self.assertIsInstance(broad1_disc, pybamm.Multiplication)
@@ -761,6 +764,33 @@ class TestDiscretise(unittest.TestCase):
         )[:, np.newaxis]
         np.testing.assert_allclose(eqn_disc.evaluate(), expected_vector)
 
+    def test_concatenation_2D(self):
+        disc = get_1p1d_discretisation_for_testing(zpts=3)
+
+        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)
+        disc.set_variable_slices([conc])
+        self.assertEqual(
+            disc.y_slices[a.id], [slice(0, 40), slice(100, 140), slice(200, 240)]
+        )
+        self.assertEqual(
+            disc.y_slices[b.id], [slice(40, 65), slice(140, 165), slice(240, 265)]
+        )
+        self.assertEqual(
+            disc.y_slices[c.id], [slice(65, 100), slice(165, 200), slice(265, 300)]
+        )
+        expr = disc.process_symbol(conc)
+        self.assertIsInstance(expr, pybamm.DomainConcatenation)
+
+        # Evaulate
+        y = np.linspace(0, 1, 300)
+        self.assertEqual(expr.children[0].evaluate(0, y).shape, (120, 1))
+        self.assertEqual(expr.children[1].evaluate(0, y).shape, (75, 1))
+        self.assertEqual(expr.children[2].evaluate(0, y).shape, (105, 1))
+        np.testing.assert_equal(expr.evaluate(0, y), y[:, np.newaxis])
+
     def test_exceptions(self):
         c_n = pybamm.Variable("c", domain=["negative electrode"])
         N_n = pybamm.grad(c_n)
diff --git a/tests/unit/test_expression_tree/test_concatenations.py b/tests/unit/test_expression_tree/test_concatenations.py
index ad527d2a0..ed70947ae 100644
--- a/tests/unit/test_expression_tree/test_concatenations.py
+++ b/tests/unit/test_expression_tree/test_concatenations.py
@@ -82,9 +82,8 @@ class TestConcatenations(unittest.TestCase):
         )
 
     def test_numpy_domain_concatenation(self):
-        # create discretisation
-        disc = get_discretisation_for_testing()
-        mesh = disc.mesh
+        # create mesh
+        mesh = get_mesh_for_testing()
 
         a_dom = ["negative electrode"]
         b_dom = ["positive electrode"]
diff --git a/tests/unit/test_expression_tree/test_evaluate.py b/tests/unit/test_expression_tree/test_evaluate.py
index 87b617668..2ba951cec 100644
--- a/tests/unit/test_expression_tree/test_evaluate.py
+++ b/tests/unit/test_expression_tree/test_evaluate.py
@@ -3,7 +3,7 @@
 #
 import pybamm
 
-from tests import get_discretisation_for_testing
+from tests import get_discretisation_for_testing, get_1p1d_discretisation_for_testing
 import unittest
 import numpy as np
 import scipy.sparse
@@ -32,8 +32,8 @@ 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[:1][[True]]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[:2][[False, True]]")
 
         var_a = pybamm.id_to_python_variable(a.id)
         var_b = pybamm.id_to_python_variable(b.id)
@@ -55,8 +55,8 @@ 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())[0], "y[:1][[True]]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[:2][[False, True]]")
         self.assertEqual(
             list(variable_symbols.values())[2], "{} + {}".format(var_a, var_b)
         )
@@ -80,8 +80,8 @@ 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())[0], "y[:1][[True]]")
+        self.assertEqual(list(variable_symbols.values())[1], "y[:2][[False, True]]")
         self.assertEqual(list(variable_symbols.values())[2], "-{}".format(var_b))
         var_child = pybamm.id_to_python_variable(expr.children[1].id)
         self.assertEqual(
@@ -97,7 +97,7 @@ 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[:1][[True]]")
         var_funct = pybamm.id_to_python_variable(expr.id, True)
         self.assertEqual(
             list(variable_symbols.values())[1], "{}({})".format(var_funct, var_a)
@@ -240,6 +240,44 @@ class TestEvaluate(unittest.TestCase):
         result = evaluator.evaluate(y=y)
         np.testing.assert_allclose(result, expr.evaluate(y=y))
 
+    def test_domain_concatenation_2D(self):
+        disc = get_1p1d_discretisation_for_testing()
+
+        a_dom = ["negative electrode"]
+        b_dom = ["positive electrode"]
+        a = pybamm.Variable("a", domain=a_dom)
+        b = pybamm.Variable("b", domain=b_dom)
+        conc = pybamm.Concatenation(a, b)
+        disc.set_variable_slices([conc])
+        expr = disc.process_symbol(conc)
+        self.assertIsInstance(expr, pybamm.DomainConcatenation)
+        a_disc = expr.children[0]
+        b_disc = expr.children[1]
+
+        y = np.empty((expr._size, 1))
+        for i in range(len(y)):
+            y[i] = i
+
+        constant_symbols = OrderedDict()
+        variable_symbols = OrderedDict()
+        pybamm.find_symbols(expr, constant_symbols, variable_symbols)
+
+        self.assertEqual(list(variable_symbols.keys())[0], a_disc.id)
+        self.assertEqual(list(variable_symbols.keys())[1], b_disc.id)
+        self.assertEqual(list(variable_symbols.keys())[2], expr.id)
+
+        self.assertEqual(len(constant_symbols), 0)
+
+        evaluator = pybamm.EvaluatorPython(expr)
+        result = evaluator.evaluate(y=y)
+        np.testing.assert_allclose(result, expr.evaluate(y=y))
+
+        # check that concatenating a single domain is consistent
+        expr = disc.process_symbol(pybamm.Concatenation(a))
+        evaluator = pybamm.EvaluatorPython(expr)
+        result = evaluator.evaluate(y=y)
+        np.testing.assert_allclose(result, expr.evaluate(y=y))
+
     def test_to_python(self):
         a = pybamm.StateVector(slice(0, 1))
         b = pybamm.StateVector(slice(1, 2))
@@ -248,8 +286,8 @@ class TestEvaluate(unittest.TestCase):
         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"
+            "self\.var_[0-9m]+ = y\[:1\]\[\[True\]\].*\\n"
+            "self\.var_[0-9m]+ = y\[:2\]\[\[False, True\]\].*\\n"
             "self\.var_[0-9m]+ = self\.var_[0-9m]+ \+ self\.var_[0-9m]+"
         )
 
diff --git a/tests/unit/test_expression_tree/test_jac.py b/tests/unit/test_expression_tree/test_jac.py
index 4237c5c8d..bea917bf3 100644
--- a/tests/unit/test_expression_tree/test_jac.py
+++ b/tests/unit/test_expression_tree/test_jac.py
@@ -7,6 +7,7 @@ import numpy as np
 import autograd.numpy as auto_np
 import unittest
 from scipy.sparse import eye
+from tests import get_mesh_for_testing
 
 
 def test_multi_var_function(arg1, arg2):
@@ -271,6 +272,57 @@ class TestJacobian(unittest.TestCase):
         jac = pybamm.inner(a * vec, b * vec).jac(vec).evaluate(y=np.ones(2)).toarray()
         np.testing.assert_array_equal(jac, 4 * np.eye(2))
 
+    def test_jac_of_domain_concatenation(self):
+        # create mesh
+        mesh = get_mesh_for_testing()
+        y = pybamm.StateVector(slice(0, 100))
+
+        # Jacobian of a DomainConcatenation of constants is a zero matrix of the
+        # appropriate size
+        a_dom = ["negative electrode"]
+        b_dom = ["separator"]
+        c_dom = ["positive electrode"]
+        a_npts = mesh[a_dom[0]][0].npts
+        b_npts = mesh[b_dom[0]][0].npts
+        c_npts = mesh[c_dom[0]][0].npts
+        a = 2 * pybamm.Vector(np.ones(a_npts), domain=a_dom)
+        b = pybamm.Vector(np.ones(b_npts), domain=b_dom)
+        c = 3 * pybamm.Vector(np.ones(c_npts), domain=c_dom)
+
+        conc = pybamm.DomainConcatenation([a, b, c], mesh)
+        jac = conc.jac(y).evaluate().toarray()
+        np.testing.assert_array_equal(jac, np.zeros((100, 100)))
+
+        # Jacobian of a DomainConcatenation of StateVectors
+        a = 2 * pybamm.StateVector(slice(0, a_npts), domain=a_dom)
+        b = pybamm.StateVector(slice(a_npts, a_npts + b_npts), domain=b_dom)
+        c = 3 * pybamm.StateVector(
+            slice(a_npts + b_npts, a_npts + b_npts + c_npts), domain=c_dom
+        )
+        conc = pybamm.DomainConcatenation([a, b, c], mesh)
+
+        y0 = np.ones(100)
+        jac = conc.jac(y).evaluate(y=y0).toarray()
+        np.testing.assert_array_equal(
+            jac,
+            np.diag(
+                np.concatenate(
+                    [2 * np.ones(a_npts), np.ones(b_npts), 3 * np.ones(c_npts)]
+                )
+            ),
+        )
+
+        # multi=domain case not implemented
+        a = 2 * pybamm.StateVector(slice(0, a_npts), domain=a_dom)
+        b = pybamm.StateVector(
+            slice(a_npts, a_npts + b_npts + c_npts), domain=b_dom + c_dom
+        )
+        conc = pybamm.DomainConcatenation([a, b], mesh)
+        with self.assertRaisesRegex(
+            NotImplementedError, "jacobian only implemented for when each child has"
+        ):
+            conc.jac(y)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_expression_tree/test_jac_2D.py b/tests/unit/test_expression_tree/test_jac_2D.py
new file mode 100644
index 000000000..5dac7929d
--- /dev/null
+++ b/tests/unit/test_expression_tree/test_jac_2D.py
@@ -0,0 +1,291 @@
+#
+# Tests for the jacobian methods for two-dimensional objects
+#
+import pybamm
+
+import numpy as np
+import unittest
+from scipy.sparse import eye
+from tests import get_1p1d_discretisation_for_testing
+
+
+def test_multi_var_function(arg1, arg2):
+    return arg1 + arg2
+
+
+class TestJacobian(unittest.TestCase):
+    def test_linear(self):
+        y = pybamm.StateVector(slice(0, 8))
+        u = pybamm.StateVector(slice(0, 2), slice(4, 6))
+        v = pybamm.StateVector(slice(2, 4), slice(6, 8))
+
+        y0 = np.ones(8)
+
+        func = u
+        jacobian = np.array(
+            [
+                [1, 0, 0, 0, 0, 0, 0, 0],
+                [0, 1, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 1, 0, 0, 0],
+                [0, 0, 0, 0, 0, 1, 0, 0],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = -v
+        jacobian = np.array(
+            [
+                [0, 0, -1, 0, 0, 0, 0, 0],
+                [0, 0, 0, -1, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, -1, 0],
+                [0, 0, 0, 0, 0, 0, 0, -1],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = 3 * u + 4 * v
+        jacobian = np.array(
+            [
+                [3, 0, 4, 0, 0, 0, 0, 0],
+                [0, 3, 0, 4, 0, 0, 0, 0],
+                [0, 0, 0, 0, 3, 0, 4, 0],
+                [0, 0, 0, 0, 0, 3, 0, 4],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = 7 * u - v * 9
+        jacobian = np.array(
+            [
+                [7, 0, -9, 0, 0, 0, 0, 0],
+                [0, 7, 0, -9, 0, 0, 0, 0],
+                [0, 0, 0, 0, 7, 0, -9, 0],
+                [0, 0, 0, 0, 0, 7, 0, -9],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        A = pybamm.Matrix(2 * eye(4))
+        func = A @ u
+        jacobian = np.array(
+            [
+                [2, 0, 0, 0, 0, 0, 0, 0],
+                [0, 2, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 2, 0, 0, 0],
+                [0, 0, 0, 0, 0, 2, 0, 0],
+            ]
+        )
+        dfunc_dy = func.jac(y).simplify().evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        # when differentiating by independent part of the state vector
+        with self.assertRaises(NotImplementedError):
+            u.jac(v)
+
+        # test Jacobian of Outer (must set domain to be 'current collector')
+        u.domain = ["current collector"]
+        func = pybamm.Outer(u, pybamm.Scalar(4))
+        jacobian = np.array(
+            [
+                [4, 0, 0, 0, 0, 0, 0, 0],
+                [0, 4, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 4, 0, 0, 0],
+                [0, 0, 0, 0, 0, 4, 0, 0],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = pybamm.Outer(u, pybamm.Vector(np.array([1, 2, 3])))
+        jacobian = np.array(
+            [
+                [1, 0, 0, 0, 0, 0, 0, 0],
+                [2, 0, 0, 0, 0, 0, 0, 0],
+                [3, 0, 0, 0, 0, 0, 0, 0],
+                [0, 1, 0, 0, 0, 0, 0, 0],
+                [0, 2, 0, 0, 0, 0, 0, 0],
+                [0, 3, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 1, 0, 0, 0],
+                [0, 0, 0, 0, 2, 0, 0, 0],
+                [0, 0, 0, 0, 3, 0, 0, 0],
+                [0, 0, 0, 0, 0, 1, 0, 0],
+                [0, 0, 0, 0, 0, 2, 0, 0],
+                [0, 0, 0, 0, 0, 3, 0, 0],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+    def test_nonlinear(self):
+        y = pybamm.StateVector(slice(0, 8))
+        u = pybamm.StateVector(slice(0, 2), slice(4, 6))
+        v = pybamm.StateVector(slice(2, 4), slice(6, 8))
+
+        y0 = np.arange(1, 9)
+
+        func = v ** 2
+        jacobian = np.array(
+            [
+                [0, 0, 6, 0, 0, 0, 0, 0],
+                [0, 0, 0, 8, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 14, 0],
+                [0, 0, 0, 0, 0, 0, 0, 16],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = 2 ** v
+        jacobian = np.array(
+            [
+                [0, 0, 2 ** 3 * np.log(2), 0, 0, 0, 0, 0],
+                [0, 0, 0, 2 ** 4 * np.log(2), 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 2 ** 7 * np.log(2), 0],
+                [0, 0, 0, 0, 0, 0, 0, 2 ** 8 * np.log(2)],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = v ** v
+        jacobian = [
+            [0, 0, 27 * (1 + np.log(3)), 0, 0, 0, 0, 0],
+            [0, 0, 0, 256 * (1 + np.log(4)), 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 823543 * (1 + np.log(7)), 0],
+            [0, 0, 0, 0, 0, 0, 0, 16777216 * (1 + np.log(8))],
+        ]
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_almost_equal(jacobian, dfunc_dy.toarray())
+
+        func = u * v
+        jacobian = np.array(
+            [
+                [3, 0, 1, 0, 0, 0, 0, 0],
+                [0, 4, 0, 2, 0, 0, 0, 0],
+                [0, 0, 0, 0, 7, 0, 5, 0],
+                [0, 0, 0, 0, 0, 8, 0, 6],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = u * (u + v)
+        jacobian = np.array(
+            [
+                [5, 0, 1, 0, 0, 0, 0, 0],
+                [0, 8, 0, 2, 0, 0, 0, 0],
+                [0, 0, 0, 0, 17, 0, 5, 0],
+                [0, 0, 0, 0, 0, 20, 0, 6],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+    def test_functions(self):
+        y = pybamm.StateVector(slice(0, 8))
+        u = pybamm.StateVector(slice(0, 2), slice(4, 6))
+        v = pybamm.StateVector(slice(2, 4), slice(6, 8))
+
+        y0 = np.arange(1, 9)
+        const = pybamm.Scalar(1)
+
+        func = pybamm.sin(u)
+        jacobian = np.array(
+            [
+                [np.cos(1), 0, 0, 0, 0, 0, 0, 0],
+                [0, np.cos(2), 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, np.cos(5), 0, 0, 0],
+                [0, 0, 0, 0, 0, np.cos(6), 0, 0],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = pybamm.cos(v)
+        jacobian = np.array(
+            [
+                [0, 0, -np.sin(3), 0, 0, 0, 0, 0],
+                [0, 0, 0, -np.sin(4), 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, -np.sin(7), 0],
+                [0, 0, 0, 0, 0, 0, 0, -np.sin(8)],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        func = pybamm.sin(3 * u * v)
+        jacobian = np.array(
+            [
+                [9 * np.cos(9), 0, 3 * np.cos(9), 0, 0, 0, 0, 0],
+                [0, 12 * np.cos(24), 0, 6 * np.cos(24), 0, 0, 0, 0],
+                [0, 0, 0, 0, 21 * np.cos(105), 0, 15 * np.cos(105), 0],
+                [0, 0, 0, 0, 0, 24 * np.cos(144), 0, 18 * np.cos(144)],
+            ]
+        )
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        # when child evaluates to number
+        func = pybamm.sin(const)
+        jacobian = np.array([[0, 0, 0, 0, 0, 0, 0, 0]])
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+        # several children
+        func = pybamm.Function(test_multi_var_function, 2 * y, 3 * y)
+        jacobian = np.diag(5 * np.ones(8))
+        dfunc_dy = func.jac(y).evaluate(y=y0)
+        np.testing.assert_array_equal(jacobian, dfunc_dy.toarray())
+
+    def test_jac_of_domain_concatenation(self):
+        # create mesh
+        disc = get_1p1d_discretisation_for_testing()
+        mesh = disc.mesh
+        y = pybamm.StateVector(slice(0, 1500))
+
+        # Jacobian of a DomainConcatenation of constants is a zero matrix of the
+        # appropriate size
+        a_dom = ["negative electrode"]
+        b_dom = ["separator"]
+        c_dom = ["positive electrode"]
+        a_npts = mesh[a_dom[0]][0].npts
+        b_npts = mesh[b_dom[0]][0].npts
+        c_npts = mesh[c_dom[0]][0].npts
+        cc_npts = mesh["current collector"][0].npts
+        curr_coll_vector = pybamm.Vector(np.ones(cc_npts), domain="current collector")
+        a = 2 * pybamm.Outer(
+            curr_coll_vector, pybamm.Vector(np.ones(a_npts), domain=a_dom)
+        )
+        b = pybamm.Outer(curr_coll_vector, pybamm.Vector(np.ones(b_npts), domain=b_dom))
+        c = 3 * pybamm.Outer(
+            curr_coll_vector, pybamm.Vector(np.ones(c_npts), domain=c_dom)
+        )
+
+        conc = pybamm.DomainConcatenation([a, b, c], mesh)
+        jac = conc.jac(y).evaluate().toarray()
+        np.testing.assert_array_equal(jac, np.zeros((1500, 1500)))
+
+        # Jacobian of a DomainConcatenation of StateVectors
+        a = pybamm.Variable("a", domain=a_dom)
+        b = pybamm.Variable("b", domain=b_dom)
+        c = pybamm.Variable("c", domain=c_dom)
+        conc = pybamm.Concatenation(a, b, c)
+        disc.set_variable_slices([conc])
+        conc_disc = disc.process_symbol(conc)
+        y0 = np.ones(1500)
+        jac = conc_disc.jac(y).evaluate(y=y0).toarray()
+        np.testing.assert_array_equal(jac, np.eye(1500))
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    pybamm.settings.debug_mode = True
+    unittest.main()
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index 04069db39..372b44b42 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -227,10 +227,6 @@ class TestSimplify(unittest.TestCase):
         self.assertIsInstance((b / b).simplify(), pybamm.Scalar)
         self.assertEqual((b / b).simplify().evaluate(), 1)
 
-        # Spatial variable
-        x = pybamm.SpatialVariable("x", ["negative electrode"])
-        self.assertIsInstance(x.simplify(), pybamm.SpatialVariable)
-
         # not implemented for Symbol
         sym = pybamm.Symbol("sym")
         with self.assertRaises(NotImplementedError):
@@ -545,8 +541,9 @@ class TestSimplify(unittest.TestCase):
 
         y = mesh.combine_submeshes(*conc.domain)[0].nodes ** 2
         self.assertIsInstance(conc_simp, pybamm.StateVector)
-        self.assertEqual(conc_simp.y_slice.start, 0)
-        self.assertEqual(conc_simp.y_slice.stop, len(y))
+        self.assertEqual(len(conc_simp.y_slices), 1)
+        self.assertEqual(conc_simp.y_slices[0].start, 0)
+        self.assertEqual(conc_simp.y_slices[0].stop, len(y))
         np.testing.assert_array_equal(conc_disc.evaluate(y=y), conc_simp.evaluate(y=y))
 
     def test_simplify_outer(self):
diff --git a/tests/unit/test_expression_tree/test_state_vector.py b/tests/unit/test_expression_tree/test_state_vector.py
new file mode 100644
index 000000000..853fcbdaa
--- /dev/null
+++ b/tests/unit/test_expression_tree/test_state_vector.py
@@ -0,0 +1,60 @@
+#
+# Tests for the Vector class
+#
+import pybamm
+import numpy as np
+
+import unittest
+
+
+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)[:, np.newaxis]
+        )
+
+        # 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_list(self):
+        sv = pybamm.StateVector(slice(0, 11), slice(20, 31))
+        y = np.linspace(0, 3, 31)
+        np.testing.assert_array_almost_equal(
+            sv.evaluate(y=y),
+            np.concatenate([np.linspace(0, 1, 11), np.linspace(2, 3, 11)])[
+                :, np.newaxis
+            ],
+        )
+        sv = pybamm.StateVector(slice(0, 11), slice(11, 20), slice(20, 31))
+        y = np.linspace(0, 3, 31)
+        np.testing.assert_array_almost_equal(sv.evaluate(y=y), y[:, np.newaxis])
+
+    def test_name(self):
+        sv = pybamm.StateVector(slice(0, 10))
+        self.assertEqual(sv.name, "y[0:10]")
+        sv = pybamm.StateVector(slice(0, 10), slice(20, 30))
+        self.assertEqual(sv.name, "y[0:10,20:30]")
+        sv = pybamm.StateVector(
+            slice(0, 10), slice(20, 30), slice(40, 50), slice(60, 70)
+        )
+        self.assertEqual(sv.name, "y[0:10,20:30,...,60:70]")
+
+    def test_failure(self):
+        with self.assertRaisesRegex(TypeError, "all y_slices must be slice objects"):
+            pybamm.StateVector(slice(0, 10), 1)
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    pybamm.settings.debug_mode = True
+    unittest.main()
diff --git a/tests/unit/test_expression_tree/test_vector.py b/tests/unit/test_expression_tree/test_vector.py
index 1e19de495..86e54a052 100644
--- a/tests/unit/test_expression_tree/test_vector.py
+++ b/tests/unit/test_expression_tree/test_vector.py
@@ -45,34 +45,6 @@ class TestVector(unittest.TestCase):
             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)[:, np.newaxis]
-        )
-
-        # 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")
     import sys
diff --git a/tests/unit/test_spatial_methods/test_finite_volume.py b/tests/unit/test_spatial_methods/test_finite_volume.py
index af3e5862e..528f52ccc 100644
--- a/tests/unit/test_spatial_methods/test_finite_volume.py
+++ b/tests/unit/test_spatial_methods/test_finite_volume.py
@@ -253,7 +253,7 @@ class TestFiniteVolume(unittest.TestCase):
         whole_cell = ["negative electrode", "separator", "positive electrode"]
         var = pybamm.Variable("var", domain=whole_cell)
         disc.set_variable_slices([var])
-        discretised_symbol = pybamm.StateVector(disc._y_slices[var.id])
+        discretised_symbol = pybamm.StateVector(*disc.y_slices[var.id])
         bcs = {
             "left": (pybamm.Scalar(0), "Dirichlet"),
             "right": (pybamm.Scalar(3), "Dirichlet"),
@@ -348,10 +348,10 @@ class TestFiniteVolume(unittest.TestCase):
         c_s_p = pybamm.Variable("c_s_p", domain=["positive particle"])
 
         disc.set_variable_slices([c_s_n])
-        disc_c_s_n = pybamm.StateVector(disc._y_slices[c_s_n.id])
+        disc_c_s_n = pybamm.StateVector(*disc.y_slices[c_s_n.id])
 
         disc.set_variable_slices([c_s_p])
-        disc_c_s_p = pybamm.StateVector(disc._y_slices[c_s_p.id])
+        disc_c_s_p = pybamm.StateVector(*disc.y_slices[c_s_p.id])
         bcs = {
             "left": (pybamm.Scalar(0), "Dirichlet"),
             "right": (pybamm.Scalar(3), "Dirichlet"),
@@ -463,6 +463,39 @@ class TestFiniteVolume(unittest.TestCase):
             np.zeros_like(combined_submesh[0].nodes[:, np.newaxis]),
         )
 
+    def test_grad_1plus1d(self):
+        mesh = get_1p1d_mesh_for_testing()
+        spatial_methods = {"macroscale": pybamm.FiniteVolume}
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+
+        a = pybamm.Variable("a", domain=["negative electrode"])
+        b = pybamm.Variable("b", domain=["separator"])
+        c = pybamm.Variable("c", domain=["positive electrode"])
+        var = pybamm.Concatenation(a, b, c)
+        boundary_conditions = {
+            var.id: {
+                "left": (pybamm.Vector(np.linspace(0, 1, 15)), "Neumann"),
+                "right": (pybamm.Vector(np.linspace(0, 1, 15)), "Neumann"),
+            }
+        }
+
+        disc.bcs = boundary_conditions
+        disc.set_variable_slices([var])
+        grad_eqn_disc = disc.process_symbol(pybamm.grad(var))
+
+        # Evaulate
+        combined_submesh = mesh.combine_submeshes(*var.domain)
+        linear_y = np.outer(np.linspace(0, 1, 15), combined_submesh[0].nodes).reshape(
+            -1, 1
+        )
+
+        expected = np.outer(
+            np.linspace(0, 1, 15), np.ones_like(combined_submesh[0].edges)
+        ).reshape(-1, 1)
+        np.testing.assert_array_almost_equal(
+            grad_eqn_disc.evaluate(None, linear_y), expected
+        )
+
     def test_spherical_grad_div_shapes_Dirichlet_bcs(self):
         """
         Test grad and div with Dirichlet boundary conditions (applied by grad on var)