Repository: pybamm-team/PyBaMM

Issue:
Issue exporting fluxes
**Describe the bug**
Shape size error when trying to output fluxes and currents.

**To Reproduce**
Steps to reproduce the behaviour:
1. Go to SPM.ipynb
2. Run all
3. Scroll to In[17] (below the list of output variables)
4. Add the line:
```python
pybamm.ProcessedVariable(
    model.variables['Positive particle flux'], solver.t, solver.y, mesh=mesh
)
```

Get the following error message:
```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-74cbf4a13b3d> in <module>
      7 )
      8 pybamm.ProcessedVariable(
----> 9     model.variables['Positive particle flux'], solver.t, solver.y, mesh=mesh
     10 )

~/PycharmProjects/PyBaMM/pybamm/processed_variable.py in __init__(self, base_variable, t_sol, y_sol, mesh, interp_kind)
     52         # Evaluate the base_variable index-by-index
     53         for idx in range(len(t_sol)):
---> 54             entries[:, idx] = base_variable.evaluate(t_sol[idx], y_sol[:, idx])
     55 
     56         # assign attributes for reference

ValueError: could not broadcast input array from shape (9) into shape (10)
```


Discussion:

Patch:
diff --git a/docs/source/processed_variable.rst b/docs/source/processed_variable.rst
index 9002601bf..b2cc38428 100644
--- a/docs/source/processed_variable.rst
+++ b/docs/source/processed_variable.rst
@@ -1,5 +1,7 @@
-Processed Variable
-==================
+Post-Process Variables
+======================
+
+.. autofunction:: pybamm.post_process_variables
 
 .. autoclass:: pybamm.ProcessedVariable
   :members:
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 2656bc73d..7e2ed6000 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -193,7 +193,7 @@ from .solvers.scikits_ode_solver import ScikitsOdeSolver
 #
 # other
 #
-from .processed_variable import ProcessedVariable
+from .processed_variable import post_process_variables, ProcessedVariable
 
 #
 # Remove any imported modules, so we don't expose them as part of pybamm
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
index c91e64757..c5fd67ea7 100644
--- a/pybamm/discretisations/discretisation.py
+++ b/pybamm/discretisations/discretisation.py
@@ -284,7 +284,6 @@ class Discretisation(object):
                             eqn, eqn_key.domain
                         )
 
-            # Process symbol (original or broadcasted)
             new_var_eqn_dict[eqn_key] = self.process_symbol(eqn)
 
             # note we are sending in the key.id here so we don't have to
diff --git a/pybamm/expression_tree/unary_operators.py b/pybamm/expression_tree/unary_operators.py
index 0688818a6..5ed0409e0 100644
--- a/pybamm/expression_tree/unary_operators.py
+++ b/pybamm/expression_tree/unary_operators.py
@@ -432,7 +432,7 @@ def div(expression):
 #
 
 
-def surf(variable):
+def surf(variable, set_domain=False):
     """convenience function for creating a right :class:`BoundaryValue`, usually in the
     spherical geometry
 
@@ -448,8 +448,13 @@ def surf(variable):
     :class:`GetSurfaceValue`
         the surface value of ``variable``
     """
-
-    return boundary_value(variable, "right")
+    out = boundary_value(variable, "right")
+    if set_domain:
+        if variable.domain == ["negative particle"]:
+            out.domain = ["negative electrode"]
+        elif variable.domain == ["positive particle"]:
+            out.domain = ["positive electrode"]
+    return out
 
 
 def average(symbol):
diff --git a/pybamm/models/lithium_ion/dfn.py b/pybamm/models/lithium_ion/dfn.py
index 456965bf9..526faa34a 100644
--- a/pybamm/models/lithium_ion/dfn.py
+++ b/pybamm/models/lithium_ion/dfn.py
@@ -32,8 +32,8 @@ class DFN(pybamm.LithiumIonBaseModel):
         "Submodels"
         # Exchange-current density
         c_e_n, _, c_e_p = c_e.orphans
-        c_s_n_surf = pybamm.surf(c_s_n)
-        c_s_p_surf = pybamm.surf(c_s_p)
+        c_s_n_surf = pybamm.surf(c_s_n, set_domain=True)
+        c_s_p_surf = pybamm.surf(c_s_p, set_domain=True)
         int_curr_model = pybamm.interface.LithiumIonReaction(param)
         j0_n = int_curr_model.get_exchange_current_densities(c_e_n, c_s_n_surf)
         j0_p = int_curr_model.get_exchange_current_densities(c_e_p, c_s_p_surf)
diff --git a/pybamm/models/submodels/electrolyte_diffusion.py b/pybamm/models/submodels/electrolyte_diffusion.py
index 1220ccc39..90728d33f 100644
--- a/pybamm/models/submodels/electrolyte_diffusion.py
+++ b/pybamm/models/submodels/electrolyte_diffusion.py
@@ -62,8 +62,9 @@ class StefanMaxwell(pybamm.SubModel):
         self.boundary_conditions = {N_e: {"left": 0, "right": 0}}
         self.variables = self.get_variables(c_e, N_e)
 
-        # Cut off if concentration goes negative
-        self.events = [pybamm.Function(np.min, c_e)]
+        # Cut off if concentration goes too small
+        # (open-circuit potential poorly defined)
+        self.events = [pybamm.Function(np.min, c_e) - 0.002]
 
     def set_leading_order_system(self, c_e, reactions, epsilon=None):
         """
@@ -114,8 +115,9 @@ class StefanMaxwell(pybamm.SubModel):
         )
         self.variables = self.get_variables(c_e_var, N_e)
 
-        # Cut off if concentration goes negative
-        self.events = [pybamm.Function(np.min, c_e)]
+        # Cut off if concentration goes too small
+        # (open-circuit potential poorly defined)
+        self.events = [c_e - 0.002]
 
     def get_variables(self, c_e, N_e):
         """
diff --git a/pybamm/models/submodels/potential.py b/pybamm/models/submodels/potential.py
index 489fe81da..94d1a17cb 100644
--- a/pybamm/models/submodels/potential.py
+++ b/pybamm/models/submodels/potential.py
@@ -34,12 +34,6 @@ class Potential(pybamm.SubModel):
         # Load parameters and spatial variables
         param = self.set_of_parameters
 
-        # Broadcast if necessary
-        if ocp_n.domain == []:
-            ocp_n = pybamm.Broadcast(ocp_n, ["negative electrode"])
-        if ocp_p.domain == []:
-            ocp_p = pybamm.Broadcast(ocp_p, ["positive electrode"])
-
         # Dimensionless
         ocp_n_av = pybamm.average(ocp_n)
         ocp_p_av = pybamm.average(ocp_p)
@@ -88,12 +82,6 @@ class Potential(pybamm.SubModel):
         # Load parameters
         param = self.set_of_parameters
 
-        # Broadcast if necessary
-        if eta_r_n.domain == []:
-            eta_r_n = pybamm.Broadcast(eta_r_n, ["negative electrode"])
-        if eta_r_p.domain == []:
-            eta_r_p = pybamm.Broadcast(eta_r_p, ["positive electrode"])
-
         # Derived and dimensional reaction overpotentials
         eta_r_n_av = pybamm.average(eta_r_n)
         eta_r_p_av = pybamm.average(eta_r_p)
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index c38a5b014..9f80f4f19 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -257,6 +257,9 @@ class ParameterValues(dict):
         elif isinstance(symbol, pybamm.SpatialVariable):
             return pybamm.SpatialVariable(symbol.name, symbol.domain, symbol.coord_sys)
 
+        elif isinstance(symbol, pybamm.StateVector):
+            return symbol.__class__(symbol.y_slice, symbol.name, symbol.domain)
+
         elif isinstance(symbol, pybamm.Time):
             return pybamm.Time()
 
diff --git a/pybamm/processed_variable.py b/pybamm/processed_variable.py
index d5b14d1f2..c1b88dc22 100644
--- a/pybamm/processed_variable.py
+++ b/pybamm/processed_variable.py
@@ -1,10 +1,43 @@
 #
 # Processed Variable class
 #
+import numbers
 import numpy as np
 import scipy.interpolate as interp
 
 
+def post_process_variables(variables, t_sol, y_sol, mesh=None, interp_kind="linear"):
+    """
+    Post-process all variables in a model
+
+    Parameters
+    ----------
+    variables : dict
+        Dictionary of variables
+    t_sol : array_like, size (m,)
+        The time vector returned by the solver
+    y_sol : array_like, size (m, k)
+        The solution vector returned by the solver. Can include solution values that
+        other than those that get read by base_variable.evaluate() (i.e. k>=n)
+    mesh : :class:`pybamm.Mesh`
+        The mesh used to solve, used here to calculate the reference x values for
+        interpolation
+    interp_kind : str
+        The method to use for interpolation
+
+    Returns
+    -------
+    dict
+        Dictionary of processed variables
+    """
+    processed_variables = {}
+    for var, eqn in variables.items():
+        processed_variables[var] = ProcessedVariable(
+            eqn, t_sol, y_sol, mesh, interp_kind
+        )
+    return processed_variables
+
+
 class ProcessedVariable(object):
     """
     An object that can be evaluated at arbitrary (scalars or vectors) t and x, and
@@ -27,50 +60,143 @@ class ProcessedVariable(object):
         interpolation
     interp_kind : str
         The method to use for interpolation
-
     """
 
     def __init__(self, base_variable, t_sol, y_sol, mesh=None, interp_kind="linear"):
-        if base_variable.domain != []:
-            if mesh is not None:
-                # Process the discretisation to get x values
-                x_sol = np.concatenate(
-                    [mesh[dom][0].nodes for dom in base_variable.domain]
-                )
-                len_x = len(x_sol)
-            else:
-                # We must provide a mesh for reference x values  for interpolation
-                raise ValueError("mesh must be provided for intepolation")
+        self.base_variable = base_variable
+        self.t_sol = t_sol
+        self.y_sol = y_sol
+        self.mesh = mesh
+        self.interp_kind = interp_kind
+        self.domain = base_variable.domain
+
+        self.base_eval = base_variable.evaluate(t_sol[0], y_sol[:, 0])
+
+        if (
+            isinstance(self.base_eval, numbers.Number)
+            or len(self.base_eval.shape) == 0
+            or self.base_eval.shape[0] == 1
+        ):
+            self.initialise_1D()
         else:
-            # No discretisation provided, or variable has no domain (function of t only)
-            # We don't need x values for interpolation
-            x_sol = None
-            len_x = 1
+            if len(self.mesh.combine_submeshes(*self.domain)) == 1:
+                self.initialise_2D()
+            else:
+                self.initialise_3D()
 
+    def initialise_1D(self):
         # initialise empty array of the correct size
-        entries = np.empty((len_x, len(t_sol)))
+        entries = np.empty(len(self.t_sol))
         # Evaluate the base_variable index-by-index
-        for idx in range(len(t_sol)):
-            entries[:, idx] = base_variable.evaluate(t_sol[idx], y_sol[:, idx])
+        for idx in range(len(self.t_sol)):
+            entries[idx] = self.base_variable.evaluate(
+                self.t_sol[idx], self.y_sol[:, idx]
+            )
+
+        # No discretisation provided, or variable has no domain (function of t only)
+        self._interpolation_function = interp.interp1d(
+            self.t_sol, entries, kind=self.interp_kind
+        )
 
-        # assign attributes for reference
         self.entries = entries
-        self.x_sol = x_sol
-        self.t_sol = t_sol
+        self.t_x_r_sol = (self.t_sol, None, None)
+        self.dimensions = 1
 
-        # set up interpolation
-        if x_sol is None:
-            self._interpolation_function = interp.interp1d(
-                t_sol, entries, kind=interp_kind
+    def initialise_2D(self):
+        len_space = self.base_eval.shape[0]
+        entries = np.empty((len_space, len(self.t_sol)))
+
+        # Evaluate the base_variable index-by-index
+        for idx in range(len(self.t_sol)):
+            entries[:, idx] = self.base_variable.evaluate(
+                self.t_sol[idx], self.y_sol[:, idx]
             )
+
+        # Process the discretisation to get x values
+        nodes = self.mesh.combine_submeshes(*self.domain)[0].nodes
+        edges = self.mesh.combine_submeshes(*self.domain)[0].edges
+        if entries.shape[0] == len(nodes):
+            space = nodes
+        elif entries.shape[0] == len(edges) - 2:
+            space = edges[1:-1]
         else:
-            self._interpolation_function = interp.interp2d(
-                t_sol, x_sol, entries, kind=interp_kind
-            )
+            raise ValueError("variable shape does not match domain shape")
+
+        # assign attributes for reference (either x_sol or r_sol)
+        self.entries = entries
+        self.dimensions = 2
+        if any("particle" in dom for dom in self.domain):
+            self.scale = "micro"
+            self.r_sol = space
+            self.t_x_r_sol = (self.t_sol, None, self.r_sol)
+        else:
+            self.scale = "macro"
+            self.x_sol = space
+            self.t_x_r_sol = (self.t_sol, self.x_sol, None)
+
+        # set up interpolation
+        # note that the order of 't' and 'space' is the reverse of what you'd expect
+        self._interpolation_function = interp.interp2d(
+            self.t_sol, space, entries, kind=self.interp_kind
+        )
+
+    def initialise_3D(self):
+        len_x = len(self.mesh.combine_submeshes(*self.domain))
+        len_r = self.base_eval.shape[0] // len_x
+        entries = np.empty((len_r, len_x, len(self.t_sol)))
 
-    def __call__(self, t, x=None):
-        "Evaluate the variable at arbitrary t (and x), using interpolation"
-        if self.x_sol is None:
-            return self._interpolation_function(t)[0]
+        # Evaluate the base_variable index-by-index
+        for idx in range(len(self.t_sol)):
+            entries[:, :, idx] = np.reshape(
+                self.base_variable.evaluate(self.t_sol[idx], self.y_sol[:, idx]),
+                [len_r, len_x],
+            )
+        # Process the discretisation to get x values
+        nodes = self.mesh.combine_submeshes(*self.domain)[0].nodes
+        edges = self.mesh.combine_submeshes(*self.domain)[0].edges
+        if entries.shape[0] == len(nodes):
+            r_sol = nodes
+        elif entries.shape[0] == len(edges) - 2:
+            r_sol = edges[1:-1]
         else:
-            return self._interpolation_function(t, x)
+            raise ValueError("variable shape does not match domain shape")
+
+        # Get x values
+        if self.domain == ["negative particle"]:
+            x_sol = self.mesh["negative electrode"][0].nodes
+        elif self.domain == ["positive particle"]:
+            x_sol = self.mesh["positive electrode"][0].nodes
+
+        # assign attributes for reference
+        self.entries = entries
+        self.dimensions = 3
+        self.x_sol = x_sol
+        self.r_sol = r_sol
+        self.t_x_r_sol = (self.t_sol, x_sol, r_sol)
+
+        # set up interpolation
+        self._interpolation_function = interp.RegularGridInterpolator(
+            (r_sol, x_sol, self.t_sol), entries, method=self.interp_kind
+        )
+
+    def __call__(self, t, x=None, r=None):
+        "Evaluate the variable at arbitrary t (and x and/or r), using interpolation"
+        if self.dimensions == 1:
+            return self._interpolation_function(t)
+        elif self.dimensions == 2:
+            if self.scale == "micro":
+                return self._interpolation_function(t, r)
+            else:
+                return self._interpolation_function(t, x)
+        elif self.dimensions == 3:
+            if isinstance(r, np.ndarray):
+                if isinstance(x, np.ndarray) and isinstance(t, np.ndarray):
+                    r = r[:, np.newaxis, np.newaxis]
+                    x = x[:, np.newaxis]
+                elif isinstance(x, np.ndarray) or isinstance(t, np.ndarray):
+                    r = r[:, np.newaxis]
+            else:
+                if isinstance(x, np.ndarray) and isinstance(t, np.ndarray):
+                    x = x[:, np.newaxis]
+
+            return self._interpolation_function((r, x, t))
diff --git a/pybamm/spatial_methods/finite_volume.py b/pybamm/spatial_methods/finite_volume.py
index 4573ed997..3e2b07613 100644
--- a/pybamm/spatial_methods/finite_volume.py
+++ b/pybamm/spatial_methods/finite_volume.py
@@ -518,6 +518,9 @@ class FiniteVolume(pybamm.SpatialMethod):
             else:
                 raise ValueError("at least one boundary condition must be provided")
 
+        # Keep same domain
+        new_discretised_symbol.domain = discretised_symbol.domain
+
         return new_discretised_symbol
 
     def boundary_value(self, symbol, discretised_symbol, side):
@@ -552,7 +555,14 @@ class FiniteVolume(pybamm.SpatialMethod):
 
         # Return boundary value with domain removed
         boundary_value = pybamm.Matrix(matrix) @ discretised_symbol
-        boundary_value.domain = []
+
+        if sec_pts > 1:
+            if symbol.domain == ["negative particle"]:
+                boundary_value.domain = ["negative electrode"]
+            elif symbol.domain == ["positive particle"]:
+                boundary_value.domain = ["positive electrode"]
+        else:
+            boundary_value.domain = []
 
         return boundary_value
 


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index 7f5463e23..79dba49f6 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -11,4 +11,5 @@ from .shared import (
     get_mesh_for_testing,
     get_p2d_mesh_for_testing,
     get_discretisation_for_testing,
+    get_p2d_discretisation_for_testing,
 )
diff --git a/tests/shared.py b/tests/shared.py
index a4225955e..595f0d03f 100644
--- a/tests/shared.py
+++ b/tests/shared.py
@@ -104,3 +104,14 @@ def get_discretisation_for_testing(npts=None):
     }
 
     return pybamm.Discretisation(mesh, spatial_methods)
+
+
+def get_p2d_discretisation_for_testing(npts=None, mpts=None):
+    mesh = get_p2d_mesh_for_testing(npts, mpts)
+    spatial_methods = {
+        "macroscale": SpatialMethodForTesting,
+        "negative particle": SpatialMethodForTesting,
+        "positive particle": SpatialMethodForTesting,
+    }
+
+    return pybamm.Discretisation(mesh, spatial_methods)
diff --git a/tests/unit/test_expression_tree/test_simplify.py b/tests/unit/test_expression_tree/test_simplify.py
index dde29c087..68c6a2f2e 100644
--- a/tests/unit/test_expression_tree/test_simplify.py
+++ b/tests/unit/test_expression_tree/test_simplify.py
@@ -342,11 +342,11 @@ class TestSimplify(unittest.TestCase):
         # we expect simplified solution to be much faster
         timer = pybamm.Timer()
         start = timer.time()
-        for _ in range(100):
+        for _ in range(200):
             expr3.evaluate(y=np.ones(300))
         end = timer.time()
         start_simp = timer.time()
-        for _ in range(100):
+        for _ in range(200):
             expr3simp.evaluate(y=np.ones(300))
         end_simp = timer.time()
         self.assertLess(end_simp - start_simp, 1.5 * (end - start))
diff --git a/tests/unit/test_models/standard_model_tests.py b/tests/unit/test_models/standard_model_tests.py
index e51fa97b7..7acd2b8ae 100644
--- a/tests/unit/test_models/standard_model_tests.py
+++ b/tests/unit/test_models/standard_model_tests.py
@@ -85,33 +85,21 @@ class StandardModelTest(object):
         )
         std_out_test.test_all()
 
-    def test_all(self, param=None, disc=None, solver=None, t_eval=None):
+    def test_all(
+        self, param=None, disc=None, solver=None, t_eval=None, skip_output_tests=False
+    ):
         self.model.check_well_posedness()
         self.test_processing_parameters(param)
         self.test_processing_disc(disc)
         self.test_solving(solver, t_eval)
 
-        # cannot test dfn yet, and lead acid composite voltage
-        # only test the full models
-        if isinstance(
-            self.model, (pybamm.LithiumIonBaseModel, pybamm.LeadAcidBaseModel)
+        if (
+            isinstance(
+                self.model, (pybamm.LithiumIonBaseModel, pybamm.LeadAcidBaseModel)
+            )
+            and not skip_output_tests
         ):
-            # cannot test dfn at moment and Composite fails the voltage test
-            # annoyingly reaction-diffusion is an instance of LeadAcidBaseModel
-            # so need to exclude it here
-            if not (
-                isinstance(
-                    self.model,
-                    (
-                        pybamm.lithium_ion.DFN,
-                        pybamm.lead_acid.Composite,
-                        pybamm.lead_acid.CompositeCapacitance,
-                        pybamm.lead_acid.NewmanTiedemannCapacitance,
-                        pybamm.ReactionDiffusionModel,
-                    ),
-                )
-            ):
-                self.test_outputs()
+            self.test_outputs()
 
     def test_update_parameters(self, param):
         # check if geometry has changed, throw error if so (need to re-discretise)
diff --git a/tests/unit/test_models/standard_output_tests.py b/tests/unit/test_models/standard_output_tests.py
index 801845914..60858a59c 100644
--- a/tests/unit/test_models/standard_output_tests.py
+++ b/tests/unit/test_models/standard_output_tests.py
@@ -122,9 +122,8 @@ class VoltageTests(BaseOutputTest):
             - charge: ocp_n decreases, ocp_p increases
             - off: ocp_n, ocp_p constant
         """
-
-        neg_end_vs_start = self.ocp_n_av.entries[:, -1] - self.ocp_n_av.entries[:, 1]
-        pos_end_vs_start = self.ocp_p_av.entries[:, -1] - self.ocp_p_av.entries[:, 1]
+        neg_end_vs_start = self.ocp_n_av.entries[-1] - self.ocp_n_av.entries[1]
+        pos_end_vs_start = self.ocp_p_av.entries[-1] - self.ocp_p_av.entries[1]
         if self.operating_condition == "discharge":
             np.testing.assert_array_less(-neg_end_vs_start, 0)
             np.testing.assert_array_less(pos_end_vs_start, 0)
@@ -142,7 +141,7 @@ class VoltageTests(BaseOutputTest):
             - off: ocv constant
         """
 
-        end_vs_start = self.ocv_av.entries[:, -1] - self.ocv_av.entries[:, 1]
+        end_vs_start = self.ocv_av.entries[-1] - self.ocv_av.entries[1]
 
         if self.operating_condition == "discharge":
             np.testing.assert_array_less(end_vs_start, 0)
@@ -157,7 +156,7 @@ class VoltageTests(BaseOutputTest):
             - charge: voltage increases
             - off: voltage constant
         """
-        end_vs_start = self.voltage.entries[:, -1] - self.voltage.entries[:, 1]
+        end_vs_start = self.voltage.entries[-1] - self.voltage.entries[1]
 
         if self.operating_condition == "discharge":
             np.testing.assert_array_less(end_vs_start, 0)
@@ -180,6 +179,7 @@ class VoltageTests(BaseOutputTest):
             + self.eta_r_av.entries
             + self.eta_e_av.entries
             + self.Delta_Phi_s_av.entries,
+            decimal=3,
         )
 
     def test_all(self):
@@ -201,16 +201,18 @@ class ParticleConcentrationTests(BaseOutputTest):
         self.c_s_n_surf = self.get_var("Negative particle surface concentration")
         self.c_s_p_surf = self.get_var("Positive particle surface concentration")
 
-        # TODO: fix flux bug
-        # self.N_s_n = self.get_var("Negative particle flux")
-        # self.N_s_p = self.get_var("Positive particle flux")
+        self.N_s_n = self.get_var("Negative particle flux")
+        self.N_s_p = self.get_var("Positive particle flux")
 
     def test_concentration_increase_decrease(self):
         """Test all concentrations in negative particles decrease and all
         concentrations in positive particles increase over a discharge."""
 
-        neg_end_vs_start = self.c_s_n.entries[:, -1] - self.c_s_n.entries[:, 1]
-        pos_end_vs_start = self.c_s_p.entries[:, -1] - self.c_s_p.entries[:, 1]
+        t, x_n, r_n = self.c_s_n.t_x_r_sol
+        t, x_p, r_p = self.c_s_p.t_x_r_sol
+
+        neg_end_vs_start = self.c_s_n(t[1:], x_n, r_n) - self.c_s_n(t[:-1], x_n, r_n)
+        pos_end_vs_start = self.c_s_p(t[1:], x_p, r_p) - self.c_s_p(t[:-1], x_p, r_p)
 
         if self.operating_condition == "discharge":
             np.testing.assert_array_less(neg_end_vs_start, 0)
@@ -245,9 +247,20 @@ class ParticleConcentrationTests(BaseOutputTest):
 
     def test_fluxes(self):
         """Test that no flux holds in the centre of the particle. Test that surface
-        flux in the negative particles is less than zero and that the flux on the
-        surface of the positive particles is greater than zeros during a discharge."""
-        # TODO: implement after flux bug fix
+        flux in the negative particles is greater than zero and that the flux in the
+        positive particles is less than zero during a discharge."""
+        t, x_n, r_n = self.N_s_n.t_x_r_sol
+        t, x_p, r_p = self.N_s_p.t_x_r_sol
+
+        if self.operating_condition == "discharge":
+            np.testing.assert_array_less(0, self.N_s_n(t[1:], x_n, r_n))
+            np.testing.assert_array_less(self.N_s_p(t[1:], x_p, r_p), 0)
+        if self.operating_condition == "charge":
+            np.testing.assert_array_less(self.N_s_n(t[1:], x_n, r_n), 0)
+            np.testing.assert_array_less(0, self.N_s_p(t[1:], x_p, r_p))
+        if self.operating_condition == "off":
+            np.testing.assert_array_almost_equal(self.N_s_n(t, x_n, r_n), 0)
+            np.testing.assert_array_almost_equal(self.N_s_p(t, x_p, r_p), 0)
 
     def test_all(self):
         self.test_concentration_increase_decrease()
@@ -273,14 +286,8 @@ class ElectrolyteConcentrationTests(BaseOutputTest):
         # self.c_e_s_av = self.get_var("Average separator electrolyte concentration")
         # self.c_e_p_av = self.get_var("Average positive electrolyte concentration")
 
-        # TODO: fix flux bug
         # self.N_e = self.get_var("Electrolyte flux")
-
-        # self.N_e_n = self.get_var("Negative electrolyte flux")
-        # self.N_e_s = self.get_var("Separator electrolyte flux")
-        # self.N_e_p = self.get_var("Positive electrolyte flux")
-
-        # self.N_e_hat = self.get_var("Reduced cation flux")
+        self.N_e_hat = self.get_var("Reduced cation flux")
 
     def test_concentration_limit(self):
         "Test that the electrolyte concentration is always greater than zero."
@@ -326,7 +333,7 @@ class ElectrolyteConcentrationTests(BaseOutputTest):
         """Test that the internal boundary fluxes are continuous. Test current
         collector fluxes are zero."""
 
-        # TODO: fix flux bug
+        # TODO: fix extrapolation to test boundary values
 
     def test_splitting(self):
         """Test that when splitting the concentrations and fluxes by negative electrode,
@@ -363,8 +370,9 @@ class PotentialTests(BaseOutputTest):
     def test_negative_electrode_potential_profile(self):
         """Test that negative electrode potential is zero on left boundary. Test
         average negative electrode potential is less than or equal to zero."""
+        t, x, _ = self.phi_s_n.t_x_r_sol
 
-        np.testing.assert_array_almost_equal(self.phi_s_n.entries[0], 0)
+        np.testing.assert_array_almost_equal(self.phi_s_n(t, x=0), 0, decimal=5)
 
     def test_positive_electrode_potential_profile(self):
         """Test average positive electrode potential is less than the positive electrode
diff --git a/tests/unit/test_models/test_lead_acid/test_composite.py b/tests/unit/test_models/test_lead_acid/test_composite.py
index 3cfc579ba..c3a8be15b 100644
--- a/tests/unit/test_models/test_lead_acid/test_composite.py
+++ b/tests/unit/test_models/test_lead_acid/test_composite.py
@@ -13,9 +13,8 @@ import numpy as np
 class TestLeadAcidComposite(unittest.TestCase):
     def test_basic_processing(self):
         model = pybamm.lead_acid.Composite()
-
         modeltest = tests.StandardModelTest(model)
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     def test_optimisations(self):
         model = pybamm.lead_acid.Composite()
diff --git a/tests/unit/test_models/test_lead_acid/test_composite_capacitance.py b/tests/unit/test_models/test_lead_acid/test_composite_capacitance.py
index d67da055d..334358ae1 100644
--- a/tests/unit/test_models/test_lead_acid/test_composite_capacitance.py
+++ b/tests/unit/test_models/test_lead_acid/test_composite_capacitance.py
@@ -15,14 +15,14 @@ class TestLeadAcidCompositeCapacitance(unittest.TestCase):
     def test_basic_processing(self):
         model = pybamm.lead_acid.CompositeCapacitance()
         modeltest = tests.StandardModelTest(model)
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     @unittest.skipIf(scikits_odes_spec is None, "scikits.odes not installed")
     def test_basic_processing_no_capacitance(self):
         model = pybamm.lead_acid.CompositeCapacitance(use_capacitance=False)
         modeltest = tests.StandardModelTest(model)
 
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     def test_optimisations(self):
         model = pybamm.lead_acid.CompositeCapacitance()
diff --git a/tests/unit/test_models/test_lead_acid/test_loqs.py b/tests/unit/test_models/test_lead_acid/test_loqs.py
index 312d1bc9a..290e36a96 100644
--- a/tests/unit/test_models/test_lead_acid/test_loqs.py
+++ b/tests/unit/test_models/test_lead_acid/test_loqs.py
@@ -14,7 +14,6 @@ class TestLeadAcidLOQS(unittest.TestCase):
     def test_basic_processing(self):
         model = pybamm.lead_acid.LOQS()
         modeltest = tests.StandardModelTest(model)
-
         modeltest.test_all()
 
     def test_optimisations(self):
@@ -36,22 +35,6 @@ class TestLeadAcidLOQS(unittest.TestCase):
         modeltest = tests.StandardModelTest(model, parameter_values=parameter_values)
         modeltest.test_all()
 
-        t_sol, y_sol = modeltest.solver.t, modeltest.solver.y
-        # check surface concentration increases in negative particle and
-        # decreases in positive particle for charge
-        c_e = pybamm.ProcessedVariable(
-            model.variables["Electrolyte concentration"],
-            t_sol,
-            y_sol,
-            mesh=modeltest.disc.mesh,
-        )
-        voltage = pybamm.ProcessedVariable(
-            model.variables["Terminal voltage"], t_sol, y_sol
-        )
-        # neg surf concentration should be monotonically increasing for a charge
-        np.testing.assert_array_less(c_e.entries[:, :-1], c_e.entries[:, 1:])
-        np.testing.assert_array_less(voltage.entries[:-1], voltage.entries[1:])
-
     def test_zero_current(self):
         model = pybamm.lead_acid.LOQS()
         parameter_values = model.default_parameter_values
@@ -59,22 +42,6 @@ class TestLeadAcidLOQS(unittest.TestCase):
         modeltest = tests.StandardModelTest(model, parameter_values=parameter_values)
         modeltest.test_all()
 
-        t_sol, y_sol = modeltest.solver.t, modeltest.solver.y
-        # check surface concentration increases in negative particle and
-        # decreases in positive particle for charge
-        c_e = pybamm.ProcessedVariable(
-            model.variables["Electrolyte concentration"],
-            t_sol,
-            y_sol,
-            mesh=modeltest.disc.mesh,
-        )
-        voltage = pybamm.ProcessedVariable(
-            model.variables["Terminal voltage"], t_sol, y_sol
-        )
-        # variables should remain unchanged
-        np.testing.assert_almost_equal(c_e.entries - c_e.entries[0], 0)
-        np.testing.assert_almost_equal(voltage.entries - voltage.entries[0], 0)
-
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_models/test_lead_acid/test_newman_tiedemann_capacitance.py b/tests/unit/test_models/test_lead_acid/test_newman_tiedemann_capacitance.py
index 5ae0df5b7..6082648ee 100644
--- a/tests/unit/test_models/test_lead_acid/test_newman_tiedemann_capacitance.py
+++ b/tests/unit/test_models/test_lead_acid/test_newman_tiedemann_capacitance.py
@@ -13,14 +13,14 @@ class TestLeadAcidNewmanTiedemannCapacitance(unittest.TestCase):
     def test_basic_processing(self):
         model = pybamm.lead_acid.NewmanTiedemannCapacitance()
         modeltest = tests.StandardModelTest(model)
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     @unittest.skipIf(scikits_odes_spec is None, "scikits.odes not installed")
     def test_basic_processing_no_capacitance(self):
         model = pybamm.lead_acid.NewmanTiedemannCapacitance(use_capacitance=False)
         modeltest = tests.StandardModelTest(model)
 
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     def test_optimisations(self):
         model = pybamm.lead_acid.NewmanTiedemannCapacitance()
diff --git a/tests/unit/test_models/test_reaction_diffusion.py b/tests/unit/test_models/test_reaction_diffusion.py
index 931c11a8b..3d2eb5f52 100644
--- a/tests/unit/test_models/test_reaction_diffusion.py
+++ b/tests/unit/test_models/test_reaction_diffusion.py
@@ -15,7 +15,7 @@ class TestReactionDiffusionModel(unittest.TestCase):
         model = pybamm.ReactionDiffusionModel()
 
         modeltest = tests.StandardModelTest(model)
-        modeltest.test_all()
+        modeltest.test_all(skip_output_tests=True)
 
     def test_optimisations(self):
         model = pybamm.ReactionDiffusionModel()
diff --git a/tests/unit/test_parameters/test_parameter_values.py b/tests/unit/test_parameters/test_parameter_values.py
index aa28e4c76..f4f3e6712 100644
--- a/tests/unit/test_parameters/test_parameter_values.py
+++ b/tests/unit/test_parameters/test_parameter_values.py
@@ -128,6 +128,12 @@ class TestParameterValues(unittest.TestCase):
         self.assertIsInstance(processed_f, pybamm.Matrix)
         np.testing.assert_array_equal(processed_f.evaluate(), np.ones((5, 6)))
 
+        # process statevector
+        g = pybamm.StateVector(slice(0, 10))
+        processed_g = parameter_values.process_symbol(g)
+        self.assertIsInstance(processed_g, pybamm.StateVector)
+        np.testing.assert_array_equal(processed_g.evaluate(y=np.ones(10)), np.ones(10))
+
         # not implemented
         sym = pybamm.Symbol("sym")
         with self.assertRaises(NotImplementedError):
diff --git a/tests/unit/test_parameters/test_update_parameters.py b/tests/unit/test_parameters/test_update_parameters.py
index 6f80b5222..311fa98e4 100644
--- a/tests/unit/test_parameters/test_update_parameters.py
+++ b/tests/unit/test_parameters/test_update_parameters.py
@@ -39,14 +39,14 @@ class TestUpdateParameters(unittest.TestCase):
         model1 = pybamm.ReactionDiffusionModel()
         modeltest1 = tests.StandardModelTest(model1)
         t_eval = np.linspace(0, 0.9)
-        modeltest1.test_all(t_eval=t_eval)
+        modeltest1.test_all(t_eval=t_eval, skip_output_tests=True)
         Y1 = modeltest1.solver.y
 
         # double initial conditions
         model2 = pybamm.ReactionDiffusionModel()
         # process and solve the model a first time
         modeltest2 = tests.StandardModelTest(model2)
-        modeltest2.test_all()
+        modeltest2.test_all(skip_output_tests=True)
         # process and solve with updated parameter values
         parameter_values_update = pybamm.ParameterValues(
             base_parameters=model2.default_parameter_values,
@@ -64,7 +64,7 @@ class TestUpdateParameters(unittest.TestCase):
         model1 = pybamm.ReactionDiffusionModel()
         modeltest1 = tests.StandardModelTest(model1)
         t_eval = np.linspace(0, 0.9)
-        modeltest1.test_all(t_eval=t_eval)
+        modeltest1.test_all(t_eval=t_eval, skip_output_tests=True)
 
         T1, Y1 = modeltest1.solver.t, modeltest1.solver.y
 
@@ -92,7 +92,9 @@ class TestUpdateParameters(unittest.TestCase):
         )
         # nb: need to be careful make parameters a reasonable size
         modeltest2 = tests.StandardModelTest(model2)
-        modeltest2.test_all(param=parameter_values_update, t_eval=t_eval)
+        modeltest2.test_all(
+            param=parameter_values_update, t_eval=t_eval, skip_output_tests=True
+        )
         T2, Y2 = modeltest2.solver.t, modeltest2.solver.y
         # results should be different
         for idx in range(len(T1)):
diff --git a/tests/unit/test_processed_variable.py b/tests/unit/test_processed_variable.py
index cf1cabf2e..70c596b94 100644
--- a/tests/unit/test_processed_variable.py
+++ b/tests/unit/test_processed_variable.py
@@ -11,7 +11,7 @@ import unittest
 
 
 class TestProcessedVariable(unittest.TestCase):
-    def test_simple_processed_variable(self):
+    def test_processed_variable_1D(self):
         # without space
         t = pybamm.t
         y = pybamm.StateVector(slice(0, 1))
@@ -19,9 +19,9 @@ class TestProcessedVariable(unittest.TestCase):
         t_sol = np.linspace(0, 1)
         y_sol = np.array([np.linspace(0, 5)])
         processed_var = pybamm.ProcessedVariable(var, t_sol, y_sol)
-        np.testing.assert_array_equal(processed_var.entries, t_sol * y_sol)
+        np.testing.assert_array_equal(processed_var.entries, t_sol * y_sol[0])
 
-    def test_processed_var_space(self):
+    def test_processed_variable_2D(self):
         t = pybamm.t
         var = pybamm.Variable("var", domain=["negative electrode", "separator"])
         x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
@@ -42,12 +42,26 @@ class TestProcessedVariable(unittest.TestCase):
             processed_eqn.entries, t_sol * y_sol + x_sol[:, np.newaxis]
         )
 
-    def test_failure(self):
-        var_with_domain = pybamm.Symbol("var", domain=["test"])
-        with self.assertRaisesRegex(ValueError, "mesh must be provided"):
-            pybamm.ProcessedVariable(var_with_domain, None, None)
+    def test_processed_variable_3D(self):
+        var = pybamm.Variable("var", domain=["negative particle"])
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable("r", domain=["negative particle"])
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries
+        r_sol = disc.process_symbol(r).entries
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5)
 
-    def test_processed_var_interpolation(self):
+        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(r_sol), len(x_sol), len(t_sol)]),
+        )
+
+    def test_processed_var_1D_interpolation(self):
         # without spatial dependence
         t = pybamm.t
         y = pybamm.StateVector(slice(0, 1))
@@ -67,7 +81,8 @@ class TestProcessedVariable(unittest.TestCase):
         np.testing.assert_array_equal(processed_eqn(t_sol), t_sol * y_sol[0])
         np.testing.assert_array_almost_equal(processed_eqn(0.5), 0.5 * 2.5)
 
-        # with spatial dependence
+    def test_processed_var_2D_interpolation(self):
+        t = pybamm.t
         var = pybamm.Variable("var", domain=["negative electrode", "separator"])
         x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
         eqn = t * var + x
@@ -105,6 +120,39 @@ class TestProcessedVariable(unittest.TestCase):
         # 2 scalars
         self.assertEqual(processed_eqn(0.5, x_sol[-1]).shape, (1,))
 
+    def test_processed_var_3D_interpolation(self):
+        var = pybamm.Variable("var", domain=["negative particle"])
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable("r", domain=["negative particle"])
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries
+        r_sol = disc.process_symbol(r).entries
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5)
+
+        processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
+        # 3 vectors
+        np.testing.assert_array_equal(
+            processed_var(t_sol, x_sol, r_sol).shape, (10, 40, 50)
+        )
+        np.testing.assert_array_equal(
+            processed_var(t_sol, x_sol, r_sol),
+            np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]),
+        )
+        # 2 vectors, 1 scalar
+        np.testing.assert_array_equal(processed_var(0.5, x_sol, r_sol).shape, (10, 40))
+        np.testing.assert_array_equal(processed_var(t_sol, 0.2, r_sol).shape, (10, 50))
+        np.testing.assert_array_equal(processed_var(t_sol, x_sol, 0.5).shape, (40, 50))
+        # 1 vectors, 2 scalar
+        np.testing.assert_array_equal(processed_var(0.5, 0.2, r_sol).shape, (10,))
+        np.testing.assert_array_equal(processed_var(0.5, x_sol, 0.5).shape, (40,))
+        np.testing.assert_array_equal(processed_var(t_sol, 0.2, 0.5).shape, (50,))
+        # 3 scalars
+        np.testing.assert_array_equal(processed_var(0.2, 0.2, 0.2).shape, ())
+
     def test_processed_variable_ode_pde_solution(self):
         # without space
         model = pybamm.StandardBatteryBaseModel()
@@ -115,8 +163,10 @@ class TestProcessedVariable(unittest.TestCase):
         modeltest = tests.StandardModelTest(model)
         modeltest.test_all()
         t_sol, y_sol = modeltest.solver.t, modeltest.solver.y
-        processed_var = pybamm.ProcessedVariable(model.variables["c"], t_sol, y_sol)
-        np.testing.assert_array_almost_equal(processed_var.entries[0], np.exp(-t_sol))
+        processed_vars = pybamm.post_process_variables(model.variables, t_sol, y_sol)
+        np.testing.assert_array_almost_equal(
+            processed_vars["c"].entries, np.exp(-t_sol)
+        )
 
         # with space
         # set up and solve model
@@ -132,14 +182,34 @@ class TestProcessedVariable(unittest.TestCase):
         t_sol, y_sol = modeltest.solver.t, modeltest.solver.y
         x = pybamm.SpatialVariable("x", domain=whole_cell)
         x_sol = modeltest.disc.process_symbol(x).entries
-        processed_var = pybamm.ProcessedVariable(
-            model.variables["c"], t_sol, y_sol, mesh=modeltest.disc.mesh
+        processed_vars = pybamm.post_process_variables(
+            model.variables, t_sol, y_sol, modeltest.disc.mesh
         )
+
         # test
         np.testing.assert_array_almost_equal(
-            processed_var.entries, np.ones_like(x_sol)[:, np.newaxis] * np.exp(-t_sol)
+            processed_vars["c"].entries,
+            np.ones_like(x_sol)[:, np.newaxis] * np.exp(-t_sol),
         )
 
+    def test_failure(self):
+        t = np.ones(25)
+        y = np.ones((15, 25))
+        mat = pybamm.Vector(np.ones(15), domain=["negative electrode"])
+        disc = tests.get_discretisation_for_testing()
+        with self.assertRaisesRegex(
+            ValueError, "variable shape does not match domain shape"
+        ):
+            pybamm.ProcessedVariable(mat, t, y, disc.mesh)
+
+        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, "variable shape does not match domain shape"
+        ):
+            pybamm.ProcessedVariable(mat, t, y, disc.mesh)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")