Repository: pybamm-team/PyBaMM

Issue:
ProcessedVariable for 1+1D models
Allow ProcessedVariable to take an object with domain electrode + current collector and return an object that can be evaluated at any t, x, z

Discussion:

Patch:
diff --git a/.flake8 b/.flake8
index 9acc35b75..5915461f2 100644
--- a/.flake8
+++ b/.flake8
@@ -10,7 +10,8 @@ exclude=
     lib,
     lib64,
     share,
-    pyvenv.cfg
+    pyvenv.cfg,
+    third-party,
 ignore=
     # False positive for white space before ':' on list slice
     # black should format these correctly
diff --git a/.gitignore b/.gitignore
index 2c47eb938..59aaf8425 100644
--- a/.gitignore
+++ b/.gitignore
@@ -55,6 +55,7 @@ pyproject.toml
 # virtual enviroment
 venv/
 venv3.5/
+PyBaMM-env/
 bin/
 etc/
 lib/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 342d6bd17..41b14edb4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,19 +2,21 @@
 
 ## Features
 
-- Add method to evaluate parameters more easily (#669 )
-- Add `Interpolant` class to interpolate experimental data (e.g. OCP curves) (#661 )
-- Allow parameters to be set by material or by specifying a particular paper (#647 )
-- Set relative and absolute tolerances independently in solvers (#645 )
-- Add some non-uniform meshes in 1D and 2D (#617 )
+- Add method to evaluate parameters more easily ([#669](https://github.com/pybamm-team/PyBaMM/pull/669))
+- Add `Interpolant` class to interpolate experimental data (e.g. OCP curves) ([#661](https://github.com/pybamm-team/PyBaMM/pull/661))
+- Allow parameters to be set by material or by specifying a particular paper ([#647](https://github.com/pybamm-team/PyBaMM/pull/647))
+- Set relative and absolute tolerances independently in solvers ([#645](https://github.com/pybamm-team/PyBaMM/pull/645))
+- Add some non-uniform meshes in 1D and 2D ([#617](https://github.com/pybamm-team/PyBaMM/pull/617))
 
 ## Optimizations
 
-- Avoid re-checking size when making a copy of an `Index` object (#656 )
-- Avoid recalculating `_evaluation_array` when making a copy of a `StateVector` object (#653 )
+- Avoid re-checking size when making a copy of an `Index` object ([#656](https://github.com/pybamm-team/PyBaMM/pull/656))
+- Avoid recalculating `_evaluation_array` when making a copy of a `StateVector` object ([#653](https://github.com/pybamm-team/PyBaMM/pull/653))
 
 ## Bug fixes
 
+- Improve the way `ProcessedVariable` objects are created in higher dimensions ([#581](https://github.com/pybamm-team/PyBaMM/pull/581))
+
 # [v0.1.0](https://github.com/pybamm-team/PyBaMM/tree/v0.1.0) - 2019-10-08
 
 This is the first official version of PyBaMM.
diff --git a/examples/scripts/compare_lead_acid_3D.py b/examples/scripts/compare_lead_acid_3D.py
index 5fa825328..b7b8cfb66 100644
--- a/examples/scripts/compare_lead_acid_3D.py
+++ b/examples/scripts/compare_lead_acid_3D.py
@@ -83,6 +83,7 @@ t_eval = np.linspace(0, 3, 1000)
 for i, model in enumerate(models):
     solution = model.default_solver.solve(model, t_eval)
     solutions[i] = solution
+    pybamm.post_process_variables(model.variables, solution.t, solution.y, mesh=mesh)
 
 # plot
 output_variables = [
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index a00e8f6ae..d8dc62f93 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -525,8 +525,6 @@ class Discretisation(object):
 
             new_var_eqn_dict[eqn_key] = self.process_symbol(eqn)
 
-            new_var_eqn_dict[eqn_key].test_shape()
-
         return new_var_eqn_dict
 
     def process_symbol(self, symbol):
@@ -549,6 +547,7 @@ class Discretisation(object):
         except KeyError:
             discretised_symbol = self._process_symbol(symbol)
             self._discretised_symbols[symbol.id] = discretised_symbol
+            discretised_symbol.test_shape()
             return discretised_symbol
 
     def _process_symbol(self, symbol):
diff --git a/pybamm/expression_tree/array.py b/pybamm/expression_tree/array.py
index 55852d327..befa06a14 100644
--- a/pybamm/expression_tree/array.py
+++ b/pybamm/expression_tree/array.py
@@ -19,11 +19,22 @@ class Array(pybamm.Symbol):
         the name of the node
     domain : iterable of str, optional
         list of domains the parameter is valid over, defaults to empty list
+    auxiliary_domainds : dict, optional
+        dictionary of auxiliary domains, defaults to empty dict
+    entries_string : str
+        String representing the entries (slow to recalculate when copying)
 
     *Extends:* :class:`Symbol`
     """
 
-    def __init__(self, entries, name=None, domain=[], entries_string=None):
+    def __init__(
+        self,
+        entries,
+        name=None,
+        domain=None,
+        auxiliary_domains=None,
+        entries_string=None,
+    ):
         if entries.ndim == 1:
             entries = entries[:, np.newaxis]
         if name is None:
@@ -31,7 +42,7 @@ class Array(pybamm.Symbol):
         self._entries = entries
         # Use known entries string to avoid re-hashing, where possible
         self.entries_string = entries_string
-        super().__init__(name, domain=domain)
+        super().__init__(name, domain=domain, auxiliary_domains=auxiliary_domains)
 
     @property
     def entries(self):
@@ -73,7 +84,13 @@ class Array(pybamm.Symbol):
 
     def new_copy(self):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
-        return self.__class__(self.entries, self.name, self.domain, self.entries_string)
+        return self.__class__(
+            self.entries,
+            self.name,
+            self.domain,
+            self.auxiliary_domains,
+            self.entries_string,
+        )
 
     def _base_evaluate(self, t=None, y=None):
         """ See :meth:`pybamm.Symbol._base_evaluate()`. """
diff --git a/pybamm/expression_tree/binary_operators.py b/pybamm/expression_tree/binary_operators.py
index 7031134e9..e177b8ac5 100644
--- a/pybamm/expression_tree/binary_operators.py
+++ b/pybamm/expression_tree/binary_operators.py
@@ -82,11 +82,10 @@ class BinaryOperator(pybamm.Symbol):
         # right child.
         if isinstance(self, (pybamm.Outer, pybamm.Kron)):
             domain = right.domain
+            auxiliary_domains = {"secondary": left.domain}
         else:
             domain = self.get_children_domains(left.domain, right.domain)
-        auxiliary_domains = self.get_children_auxiliary_domains(
-            left.auxiliary_domains, right.auxiliary_domains
-        )
+            auxiliary_domains = self.get_children_auxiliary_domains([left, right])
         super().__init__(
             name,
             children=[left, right],
@@ -118,15 +117,6 @@ class BinaryOperator(pybamm.Symbol):
                 )
             )
 
-    def get_children_auxiliary_domains(self, l_aux_domains, r_aux_domains):
-        "Combine auxiliary domains from children, at all levels"
-        aux_domains = {}
-        for level in set(l_aux_domains.keys()).union(r_aux_domains.keys()):
-            ldomain = l_aux_domains.get(level, [])
-            rdomain = r_aux_domains.get(level, [])
-            aux_domains[level] = self.get_children_domains(ldomain, rdomain)
-        return aux_domains
-
     def new_copy(self):
         """ See :meth:`pybamm.Symbol.new_copy()`. """
 
diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py
index 2c5c59aa9..3f832a0a4 100644
--- a/pybamm/expression_tree/concatenations.py
+++ b/pybamm/expression_tree/concatenations.py
@@ -45,27 +45,6 @@ class Concatenation(pybamm.Symbol):
                 raise pybamm.DomainError("""domain of children must be disjoint""")
         return domain
 
-    def get_children_auxiliary_domains(self, children):
-        "Combine auxiliary domains from children, at all levels"
-        aux_domains = {}
-        for child in children:
-            for level in child.auxiliary_domains.keys():
-                if (
-                    not hasattr(aux_domains, level)
-                    or aux_domains[level] == []
-                    or child.auxiliary_domains[level] == aux_domains[level]
-                ):
-                    aux_domains[level] = child.auxiliary_domains[level]
-                else:
-                    raise pybamm.DomainError(
-                        """children must have same or empty auxiliary domains,
-                        not {!s} and {!s}""".format(
-                            aux_domains[level], child.auxiliary_domains[level]
-                        )
-                    )
-
-        return aux_domains
-
     def _concatenation_evaluate(self, children_eval):
         """ See :meth:`Concatenation._concatenation_evaluate()`. """
         if len(children_eval) == 0:
diff --git a/pybamm/expression_tree/functions.py b/pybamm/expression_tree/functions.py
index 20f6ae46e..f90319659 100644
--- a/pybamm/expression_tree/functions.py
+++ b/pybamm/expression_tree/functions.py
@@ -35,6 +35,7 @@ class Function(pybamm.Symbol):
                 name = "function ({})".format(function.__class__)
         children_list = list(children)
         domain = self.get_children_domains(children_list)
+        auxiliary_domains = self.get_children_auxiliary_domains(children)
 
         self.function = function
         self.derivative = derivative
@@ -46,7 +47,12 @@ class Function(pybamm.Symbol):
         else:
             self.takes_no_params = len(signature(function).parameters) == 0
 
-        super().__init__(name, children=children_list, domain=domain)
+        super().__init__(
+            name,
+            children=children_list,
+            domain=domain,
+            auxiliary_domains=auxiliary_domains,
+        )
 
     def get_children_domains(self, children_list):
         """Obtains the unique domain of the children. If the
diff --git a/pybamm/expression_tree/matrix.py b/pybamm/expression_tree/matrix.py
index 8c70b75dd..cd613266d 100644
--- a/pybamm/expression_tree/matrix.py
+++ b/pybamm/expression_tree/matrix.py
@@ -11,19 +11,18 @@ class Matrix(pybamm.Array):
 
     **Extends:** :class:`Array`
 
-    Parameters
-    ----------
-
-    entries : numpy.array
-        the array associated with the node
-    name : str, optional
-        the name of the node
-
     """
 
-    def __init__(self, entries, name=None, domain=[], entries_string=None):
+    def __init__(
+        self,
+        entries,
+        name=None,
+        domain=None,
+        auxiliary_domains=None,
+        entries_string=None,
+    ):
         if name is None:
             name = "Matrix {!s}".format(entries.shape)
             if issparse(entries):
                 name = "Sparse " + name
-        super().__init__(entries, name, domain, entries_string)
+        super().__init__(entries, name, domain, auxiliary_domains, entries_string)
diff --git a/pybamm/expression_tree/parameter.py b/pybamm/expression_tree/parameter.py
index 6656f13e8..f99bb0991 100644
--- a/pybamm/expression_tree/parameter.py
+++ b/pybamm/expression_tree/parameter.py
@@ -61,7 +61,10 @@ class FunctionParameter(pybamm.Symbol):
         self.diff_variable = diff_variable
         children_list = list(children)
         domain = self.get_children_domains(children_list)
-        super().__init__(name, children=children, domain=domain)
+        auxiliary_domains = self.get_children_auxiliary_domains(children)
+        super().__init__(
+            name, children=children, domain=domain, auxiliary_domains=auxiliary_domains
+        )
 
     def set_id(self):
         """See :meth:`pybamm.Symbol.set_id` """
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 8bbbb8db9..c5575e9e5 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -74,6 +74,14 @@ class Symbol(anytree.NodeMixin):
     domain : iterable of str, or str
         list of domains over which the node is valid (empty list indicates the symbol
         is valid over all domains)
+    auxiliary_domains : dict of str
+        dictionary of auxiliary domains over which the node is valid (empty dictionary
+        indicates no auxiliary domains). Keys can be "secondary" or "tertiary". The
+        symbol is broadcast over its auxiliary domains.
+        For example, a symbol might have domain "negative particle", secondary domain
+        "separator" and tertiary domain "current collector" (`domain="negative
+        particle", auxiliary_domains={"secondary": "separator", "tertiary": "current
+        collector"}`).
 
     """
 
@@ -160,6 +168,27 @@ class Symbol(anytree.NodeMixin):
             # Update id since domain has changed
             self.set_id()
 
+    def get_children_auxiliary_domains(self, children):
+        "Combine auxiliary domains from children, at all levels"
+        aux_domains = {}
+        for child in children:
+            for level in child.auxiliary_domains.keys():
+                if (
+                    level not in aux_domains
+                    or aux_domains[level] == []
+                    or child.auxiliary_domains[level] == aux_domains[level]
+                ):
+                    aux_domains[level] = child.auxiliary_domains[level]
+                else:
+                    raise pybamm.DomainError(
+                        """children must have same or empty auxiliary domains,
+                        not {!s} and {!s}""".format(
+                            aux_domains[level], child.auxiliary_domains[level]
+                        )
+                    )
+
+        return aux_domains
+
     @property
     def id(self):
         return self._id
diff --git a/pybamm/expression_tree/vector.py b/pybamm/expression_tree/vector.py
index 8fcebbb5d..26fff832d 100644
--- a/pybamm/expression_tree/vector.py
+++ b/pybamm/expression_tree/vector.py
@@ -12,19 +12,16 @@ class Vector(pybamm.Array):
 
     **Extends:** :class:`Array`
 
-    Parameters
-    ----------
-
-    entries : numpy.array
-        the array associated with the node
-    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
-
     """
 
-    def __init__(self, entries, name=None, domain=[], entries_string=None):
+    def __init__(
+        self,
+        entries,
+        name=None,
+        domain=None,
+        auxiliary_domains=None,
+        entries_string=None,
+    ):
         # make sure that entries are a vector (can be a column vector)
         if entries.ndim == 1:
             entries = entries[:, np.newaxis]
@@ -39,7 +36,7 @@ class Vector(pybamm.Array):
         if name is None:
             name = "Column vector of length {!s}".format(entries.shape[0])
 
-        super().__init__(entries, name, domain, entries_string)
+        super().__init__(entries, name, domain, auxiliary_domains, entries_string)
 
     def _jac(self, variable):
         """ See :meth:`pybamm.Symbol._jac()`. """
diff --git a/pybamm/processed_variable.py b/pybamm/processed_variable.py
index 2e6507459..3a41872e8 100644
--- a/pybamm/processed_variable.py
+++ b/pybamm/processed_variable.py
@@ -210,37 +210,90 @@ class ProcessedVariable(object):
 
     def initialise_3D(self):
         """
-        Initialise a 3D object that depends on x and r.
-        Needs to be generalised to deal with other domains
+        Initialise a 3D object that depends on x and r, or x and z.
+        Needs to be generalised to deal with other domains.
+
+        Notes
+        -----
+        There is different behaviour between a variable on an electrode domain
+        broadcast to a particle (such as temperature) and a variable on a particle
+        domain broadcast to an electrode (such as particle concentration). We deal with
+        this by reshaping the former with the Fortran order ("F") and the latter with
+        the C order ("C"). These are transposes of each other, so this approach simply
+        avoids having to transpose later.
         """
-        if self.domain in [["negative particle"], ["negative electrode"]]:
+        # Dealt with weird particle/electrode case
+        if self.domain in [
+            ["negative electrode"],
+            ["positive electrode"],
+        ] and self.auxiliary_domains["secondary"] in [
+            ["negative particle"],
+            ["positive particle"],
+        ]:
+            # Switch domain and auxiliary domains and set order to Fortran order ("F")
+            dom = self.domain
+            self.domain = self.auxiliary_domains["secondary"]
+            self.auxiliary_domains["secondary"] = dom
+            order = "F"
+        else:
+            # Set order to C order ("C")
+            order = "C"
+
+        # Process x-r or x-z
+        if self.domain == ["negative particle"] and self.auxiliary_domains[
+            "secondary"
+        ] == ["negative electrode"]:
             x_sol = self.mesh["negative electrode"][0].nodes
             r_nodes = self.mesh["negative particle"][0].nodes
             r_edges = self.mesh["negative particle"][0].edges
-        elif self.domain in [["positive particle"], ["positive electrode"]]:
+            set_up_r = True
+        elif self.domain == ["positive particle"] and self.auxiliary_domains[
+            "secondary"
+        ] == ["positive electrode"]:
             x_sol = self.mesh["positive electrode"][0].nodes
             r_nodes = self.mesh["positive particle"][0].nodes
             r_edges = self.mesh["positive particle"][0].edges
+            set_up_r = True
+        elif self.domain[0] in [
+            "negative electrode",
+            "separator",
+            "positive electrode",
+        ] and self.auxiliary_domains["secondary"] == ["current collector"]:
+            x_nodes = self.mesh.combine_submeshes(*self.domain)[0].nodes
+            x_edges = self.mesh.combine_submeshes(*self.domain)[0].edges
+            z_sol = self.mesh["current collector"][0].nodes
+            r_sol = None
+            self.first_dimension = "x"
+            self.second_dimension = "z"
+
+            if self.base_eval.size // len(z_sol) == len(x_nodes):
+                x_sol = x_nodes
+            elif self.base_eval.size // len(z_sol) == len(x_edges):
+                x_sol = x_edges
+            first_dim_nodes = x_sol
+            second_dim_nodes = z_sol
+            set_up_r = False
         else:
             raise pybamm.DomainError(
-                """ Can only create 3D objects on electrodes and particles. Current
-                collector not yet implemented"""
+                """ Cannot process 3D object with domain '{}'
+                and auxiliary_domains '{}'""".format(
+                    self.domain, self.auxiliary_domains
+                )
             )
-        len_x = len(x_sol)
-        len_r = self.base_eval.shape[0] // len_x
-        if self.domain in [["negative particle"], ["positive particle"]]:
+        if set_up_r:
+            z_sol = None
             self.first_dimension = "x"
             self.second_dimension = "r"
-            first_dim_size = len_x
-            second_dim_size = len_r
-            transpose = False
-        else:
-            self.first_dimension = "r"
-            self.second_dimension = "x"
-            first_dim_size = len_r
-            second_dim_size = len_x
-            transpose = True
-        entries = np.empty((len_x, len_r, len(self.t_sol)))
+            if self.base_eval.size // len(x_sol) == len(r_nodes):
+                r_sol = r_nodes
+            elif self.base_eval.size // len(x_sol) == len(r_edges):
+                r_sol = r_edges
+            first_dim_nodes = x_sol
+            second_dim_nodes = r_sol
+
+        first_dim_size = len(first_dim_nodes)
+        second_dim_size = len(second_dim_nodes)
+        entries = np.empty((first_dim_size, second_dim_size, len(self.t_sol)))
 
         # Evaluate the base_variable index-by-index
         for idx in range(len(self.t_sol)):
@@ -250,36 +303,29 @@ class ProcessedVariable(object):
                 eval_and_known_evals = self.base_variable.evaluate(
                     t, u, self.known_evals[t]
                 )
-                temporary = np.reshape(
-                    eval_and_known_evals[0], [first_dim_size, second_dim_size]
+                entries[:, :, idx] = np.reshape(
+                    eval_and_known_evals[0],
+                    [first_dim_size, second_dim_size],
+                    order=order,
                 )
                 self.known_evals[t] = eval_and_known_evals[1]
             else:
-                temporary = np.reshape(
-                    self.base_variable.evaluate(t, u), [first_dim_size, second_dim_size]
+                entries[:, :, idx] = np.reshape(
+                    self.base_variable.evaluate(t, u),
+                    [first_dim_size, second_dim_size],
+                    order=order,
                 )
-            if transpose is True:
-                entries[:, :, idx] = np.transpose(temporary)
-            else:
-                entries[:, :, idx] = temporary
-
-        # Assess whether on nodes or edges
-        if entries.shape[1] == len(r_nodes):
-            r_sol = r_nodes
-        elif entries.shape[1] == len(r_edges):
-            r_sol = r_edges
-        else:
-            raise ValueError("3D variable shape does not match domain shape")
 
         # assign attributes for reference
         self.entries = entries
         self.dimensions = 3
         self.x_sol = x_sol
         self.r_sol = r_sol
+        self.z_sol = z_sol
 
         # set up interpolation
         self._interpolation_function = interp.RegularGridInterpolator(
-            (x_sol, r_sol, self.t_sol),
+            (first_dim_nodes, second_dim_nodes, self.t_sol),
             entries,
             method=self.interp_kind,
             fill_value=np.nan,
@@ -346,7 +392,9 @@ class ProcessedVariable(object):
         )
 
     def __call__(self, t=None, x=None, r=None, y=None, z=None):
-        "Evaluate the variable at arbitrary t (and x and/or r), using interpolation"
+        """
+        Evaluate the variable at arbitrary t (and x, r, y and/or z), using interpolation
+        """
         if self.dimensions == 1:
             return self._interpolation_function(t)
         elif self.dimensions == 2:
@@ -359,32 +407,13 @@ class ProcessedVariable(object):
 
     def call_2D(self, t, x, r, z):
         "Evaluate a 2D variable"
-        if self.spatial_var_name == "r":
-            if r is not None:
-                return self._interpolation_function(t, r)
-            else:
-                raise ValueError("r cannot be None for microscale variable")
-        elif self.spatial_var_name == "x":
-            if x is not None:
-                return self._interpolation_function(t, x)
-            else:
-                raise ValueError("x cannot be None for macroscale variable")
-        else:
-            if z is not None:
-                return self._interpolation_function(t, z)
-            else:
-                raise ValueError("z cannot be None for macroscale variable")
+        spatial_var = eval_dimension_name(self.spatial_var_name, x, r, None, z)
+        return self._interpolation_function(t, spatial_var)
 
     def call_3D(self, t, x, r, y, z):
         "Evaluate a 3D variable"
-        if (self.first_dimension == "x" and self.second_dimension == "r") or (
-            self.first_dimension == "r" and self.second_dimension == "x"
-        ):
-            first_dim = x
-            second_dim = r
-        elif self.first_dimension == "y" and self.second_dimension == "z":
-            first_dim = y
-            second_dim = z
+        first_dim = eval_dimension_name(self.first_dimension, x, r, y, z)
+        second_dim = eval_dimension_name(self.second_dimension, x, r, y, z)
         if isinstance(first_dim, np.ndarray):
             if isinstance(second_dim, np.ndarray) and isinstance(t, np.ndarray):
                 first_dim = first_dim[:, np.newaxis, np.newaxis]
@@ -396,3 +425,19 @@ class ProcessedVariable(object):
                 second_dim = second_dim[:, np.newaxis]
 
         return self._interpolation_function((first_dim, second_dim, t))
+
+
+def eval_dimension_name(name, x, r, y, z):
+    if name == "x":
+        out = x
+    elif name == "r":
+        out = r
+    elif name == "y":
+        out = y
+    elif name == "z":
+        out = z
+
+    if out is None:
+        raise ValueError("inputs {} cannot be None".format(name))
+    else:
+        return out
diff --git a/pybamm/spatial_methods/spatial_method.py b/pybamm/spatial_methods/spatial_method.py
index 7919a8204..1a8fbd7a2 100644
--- a/pybamm/spatial_methods/spatial_method.py
+++ b/pybamm/spatial_methods/spatial_method.py
@@ -84,11 +84,11 @@ class SpatialMethod:
             out = pybamm.Outer(
                 symbol, pybamm.Vector(np.ones(primary_pts_for_broadcast), domain=domain)
             )
-            out.auxiliary_domains = auxiliary_domains
 
         elif broadcast_type == "full":
             out = symbol * pybamm.Vector(np.ones(full_pts_for_broadcast), domain=domain)
 
+        out.auxiliary_domains = auxiliary_domains
         return out
 
     def gradient(self, symbol, discretised_symbol, boundary_conditions):


Test patch:
diff --git a/tests/unit/test_expression_tree/test_concatenations.py b/tests/unit/test_expression_tree/test_concatenations.py
index ed70947ae..a010df470 100644
--- a/tests/unit/test_expression_tree/test_concatenations.py
+++ b/tests/unit/test_expression_tree/test_concatenations.py
@@ -40,6 +40,29 @@ class TestConcatenations(unittest.TestCase):
         with self.assertRaises(pybamm.DomainError):
             pybamm.Concatenation(a, b, d)
 
+    def test_concatenation_auxiliary_domains(self):
+        a = pybamm.Symbol(
+            "a",
+            domain=["negative electrode"],
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        b = pybamm.Symbol(
+            "b",
+            domain=["separator", "positive electrode"],
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        conc = pybamm.Concatenation(a, b)
+        self.assertEqual(conc.auxiliary_domains, {"secondary": ["current collector"]})
+
+        # Can't concatenate nodes with overlapping domains
+        c = pybamm.Symbol(
+            "c", domain=["test"], auxiliary_domains={"secondary": "something else"}
+        )
+        with self.assertRaisesRegex(
+            pybamm.DomainError, "children must have same or empty auxiliary domains"
+        ):
+            pybamm.Concatenation(a, b, c)
+
     def test_numpy_concatenation_vectors(self):
         # with entries
         y = np.linspace(0, 1, 15)[:, np.newaxis]
diff --git a/tests/unit/test_expression_tree/test_unary_operators.py b/tests/unit/test_expression_tree/test_unary_operators.py
index 38649e5ee..f62a7ef6b 100644
--- a/tests/unit/test_expression_tree/test_unary_operators.py
+++ b/tests/unit/test_expression_tree/test_unary_operators.py
@@ -139,8 +139,11 @@ class TestUnaryOperators(unittest.TestCase):
         # errors
         with self.assertRaisesRegex(TypeError, "index must be integer or slice"):
             pybamm.Index(vec, 0.0)
+        debug_mode = pybamm.settings.debug_mode
+        pybamm.settings.debug_mode = True
         with self.assertRaisesRegex(ValueError, "slice size exceeds child size"):
             pybamm.Index(vec, 5)
+        pybamm.settings.debug_mode = debug_mode
 
     def test_diff(self):
         a = pybamm.StateVector(slice(0, 1))
diff --git a/tests/unit/test_models/test_submodels/test_porosity/test_leading_reaction_driven_porosity.py b/tests/unit/test_models/test_submodels/test_porosity/test_leading_reaction_driven_porosity.py
index 6d9406543..ce90a031c 100644
--- a/tests/unit/test_models/test_submodels/test_porosity/test_leading_reaction_driven_porosity.py
+++ b/tests/unit/test_models/test_submodels/test_porosity/test_leading_reaction_driven_porosity.py
@@ -10,7 +10,7 @@ import unittest
 class TestLeadingOrder(unittest.TestCase):
     def test_public_functions(self):
         param = pybamm.standard_parameters_lead_acid
-        a = pybamm.Broadcast(pybamm.Scalar(0), "test")
+        a = pybamm.Broadcast(pybamm.Scalar(0), "current collector")
         variables = {
             "X-averaged negative electrode interfacial current density": a,
             "X-averaged positive electrode interfacial current density": a,
diff --git a/tests/unit/test_processed_variable.py b/tests/unit/test_processed_variable.py
index 2f360180b..bd9f9cd0a 100644
--- a/tests/unit/test_processed_variable.py
+++ b/tests/unit/test_processed_variable.py
@@ -80,8 +80,12 @@ class TestProcessedVariable(unittest.TestCase):
         c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"])
         pybamm.ProcessedVariable(c, solution.t, solution.y, mesh)
 
-    def test_processed_variable_3D(self):
-        var = pybamm.Variable("var", domain=["negative particle"])
+    def test_processed_variable_3D_x_r(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
         x = pybamm.SpatialVariable("x", domain=["negative electrode"])
         r = pybamm.SpatialVariable("r", domain=["negative particle"])
 
@@ -99,6 +103,41 @@ class TestProcessedVariable(unittest.TestCase):
             np.reshape(y_sol, [len(x_sol), len(r_sol), len(t_sol)]),
         )
 
+    def test_processed_variable_3D_x_z(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative electrode", "separator"],
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
+        z = pybamm.SpatialVariable("z", domain=["current collector"])
+
+        disc = tests.get_1p1d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        z_sol = disc.process_symbol(z).entries[:, 0]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = np.ones(len(x_sol) * len(z_sol))[:, np.newaxis] * np.linspace(0, 5)
+
+        processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
+        np.testing.assert_array_equal(
+            processed_var.entries,
+            np.reshape(y_sol, [len(x_sol), len(z_sol), len(t_sol)]),
+        )
+
+        # On edges
+        x_s_edge = pybamm.Matrix(
+            np.repeat(disc.mesh["separator"][0].edges, len(z_sol)),
+            domain="separator",
+            auxiliary_domains={"secondary": "current collector"},
+        )
+        processed_x_s_edge = pybamm.ProcessedVariable(x_s_edge, t_sol, y_sol, disc.mesh)
+        np.testing.assert_array_equal(
+            x_s_edge.entries[:, 0],
+            processed_x_s_edge.entries[:, :, 0].reshape(-1, 1)[:, 0],
+        )
+
     def test_processed_variable_3D_scikit(self):
         var = pybamm.Variable("var", domain=["current collector"])
         y = pybamm.SpatialVariable("y", domain=["current collector"])
@@ -206,7 +245,11 @@ class TestProcessedVariable(unittest.TestCase):
         )
 
     def test_processed_var_3D_interpolation(self):
-        var = pybamm.Variable("var", domain=["negative particle"])
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
         x = pybamm.SpatialVariable("x", domain=["negative electrode"])
         r = pybamm.SpatialVariable("r", domain=["negative particle"])
 
@@ -239,7 +282,11 @@ class TestProcessedVariable(unittest.TestCase):
         np.testing.assert_array_equal(processed_var(0.2, 0.2, 0.2).shape, ())
 
         # positive particle
-        var = pybamm.Variable("var", domain=["positive particle"])
+        var = pybamm.Variable(
+            "var",
+            domain=["positive particle"],
+            auxiliary_domains={"secondary": ["positive electrode"]},
+        )
         x = pybamm.SpatialVariable("x", domain=["positive electrode"])
         r = pybamm.SpatialVariable("r", domain=["positive particle"])
 
@@ -257,7 +304,11 @@ class TestProcessedVariable(unittest.TestCase):
         )
 
     def test_processed_var_3D_r_first_dimension(self):
-        var = pybamm.Variable("var", domain=["negative particle"])
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
         broad_var = pybamm.PrimaryBroadcast(var, "negative electrode")
         x = pybamm.SpatialVariable("x", domain=["negative electrode"])
         r = pybamm.SpatialVariable("r", domain=["negative particle"])
@@ -394,7 +445,11 @@ class TestProcessedVariable(unittest.TestCase):
         whole_cell = ["negative electrode", "separator", "positive electrode"]
         model = pybamm.BaseBatteryModel()
         c = pybamm.Variable("conc", domain=whole_cell)
-        c_s = pybamm.Variable("particle conc", domain="negative particle")
+        c_s = pybamm.Variable(
+            "particle conc",
+            domain="negative particle",
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
         model.rhs = {c: -c, c_s: 1 - c_s}
         model.initial_conditions = {c: 1, c_s: 0.5}
         model.boundary_conditions = {
@@ -423,16 +478,6 @@ class TestProcessedVariable(unittest.TestCase):
             np.ones_like(x_sol)[:, np.newaxis] * np.exp(-t_sol),
         )
 
-    def test_failure(self):
-        t = np.ones(25)
-        y = np.ones((120, 25))
-        mat = pybamm.Vector(np.ones(120), domain=["negative particle"])
-        disc = tests.get_p2d_discretisation_for_testing()
-        with self.assertRaisesRegex(
-            ValueError, "3D variable shape does not match domain shape"
-        ):
-            pybamm.ProcessedVariable(mat, t, y, disc.mesh)
-
     def test_call_failure(self):
         # x domain
         var = pybamm.Variable("var x", domain=["negative electrode", "separator"])
@@ -445,9 +490,7 @@ class TestProcessedVariable(unittest.TestCase):
         y_sol = x_sol[:, np.newaxis] * np.linspace(0, 5)
 
         processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
-        with self.assertRaisesRegex(
-            ValueError, "x cannot be None for macroscale variable"
-        ):
+        with self.assertRaisesRegex(ValueError, "x cannot be None"):
             processed_var(0)
 
         # r domain
@@ -460,13 +503,9 @@ class TestProcessedVariable(unittest.TestCase):
         y_sol = r_sol[:, np.newaxis] * np.linspace(0, 5)
 
         processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
-        with self.assertRaisesRegex(
-            ValueError, "r cannot be None for microscale variable"
-        ):
+        with self.assertRaisesRegex(ValueError, "r cannot be None"):
             processed_var(0)
-        with self.assertRaisesRegex(
-            ValueError, "r cannot be None for microscale variable"
-        ):
+        with self.assertRaisesRegex(ValueError, "r cannot be None"):
             processed_var(0, 1)
 
 
diff --git a/tests/unit/test_quick_plot.py b/tests/unit/test_quick_plot.py
index ca0907ce7..601f34892 100644
--- a/tests/unit/test_quick_plot.py
+++ b/tests/unit/test_quick_plot.py
@@ -102,7 +102,9 @@ class TestQuickPlot(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "mismatching variable domains"):
             pybamm.QuickPlot(model, mesh, solution, [["a", "b broadcasted"]])
         model.variables["3D variable"] = disc.process_symbol(
-            pybamm.Broadcast(1, ["negative particle"])
+            pybamm.FullBroadcast(
+                1, "negative particle", {"secondary": "negative electrode"}
+            )
         )
         with self.assertRaisesRegex(NotImplementedError, "cannot plot 3D variables"):
             pybamm.QuickPlot(model, mesh, solution, ["3D variable"])