[WIP] Allow 3D process variables

src/pybamm/discretisations/discretisation.py

@@ -783,6 +783,13 @@ def process_symbol(self, symbol):
                 ]
             else:
                 discretised_symbol.secondary_mesh = None
+
+            # Assign tertiary mesh
+            if symbol.domains["tertiary"] != []:
+                discretised_symbol.tertiary_mesh = self.mesh[symbol.domains["tertiary"]]
+            else:
+                discretised_symbol.tertiary_mesh = None
+
             return discretised_symbol
 
     def _process_symbol(self, symbol):

src/pybamm/models/base_model.py

@@ -122,6 +122,11 @@
                 else:
                     var.secondary_mesh = None
 
+                if var.domains["tertiary"] != []:
+                    var.tertiary_mesh = properties["mesh"][var.domains["tertiary"]]
+                else:
+                    var.tertiary_mesh = None
+
             if properties["geometry"]:
                 instance._geometry = pybamm.Geometry(properties["geometry"])
         else:
@@ -854,51 +859,48 @@
         initial_conditions = {}
         if isinstance(solution, pybamm.Solution):
             solution = solution.last_state
+
+        def get_final_state_eval(final_state):
+            if isinstance(solution, pybamm.Solution):
+                final_state = final_state.data
+
+            if final_state.ndim == 0:
+                return np.array([final_state])
+            elif final_state.ndim == 1:
+                return final_state[-1:]
+            elif final_state.ndim == 2:
+                return final_state[:, -1]
+            elif final_state.ndim == 3:
+                return final_state[:, :, -1].flatten(order="F")
+            elif final_state.ndim == 4:
+                return final_state[:, :, :, -1].flatten(order="F")
+            else:
+                raise NotImplementedError("Variable must be 0D, 1D, 2D, or 3D")
+
+        def get_variable_state(var_name):
+            try:
+                return solution[var_name]
+            except KeyError as e:
+                raise pybamm.ModelError(
+                    "To update a model from a solution, each variable in "
+                    "model.initial_conditions must appear in the solution with "
+                    "the same key as the variable name. In the solution provided, "
+                    f"'{e.args[0]}' was not found."
+                ) from e
+
         for var in self.initial_conditions:
             if isinstance(var, pybamm.Variable):
-                try:
-                    final_state = solution[var.name]
-                except KeyError as e:
-                    raise pybamm.ModelError(
-                        "To update a model from a solution, each variable in "
-                        "model.initial_conditions must appear in the solution with "
-                        "the same key as the variable name. In the solution provided, "
-                        f"'{e.args[0]}' was not found."
-                    ) from e
-                if isinstance(solution, pybamm.Solution):
-                    final_state = final_state.data
-                if final_state.ndim == 0:
-                    final_state_eval = np.array([final_state])
-                elif final_state.ndim == 1:
-                    final_state_eval = final_state[-1:]
-                elif final_state.ndim == 2:
-                    final_state_eval = final_state[:, -1]
-                elif final_state.ndim == 3:
-                    final_state_eval = final_state[:, :, -1].flatten(order="F")
-                else:
-                    raise NotImplementedError("Variable must be 0D, 1D, or 2D")
+                final_state = get_variable_state(var.name)
+                final_state_eval = get_final_state_eval(final_state)
+
             elif isinstance(var, pybamm.Concatenation):
                 children = []
                 for child in var.orphans:
-                    try:
-                        final_state = solution[child.name]
-                    except KeyError as e:
-                        raise pybamm.ModelError(
-                            "To update a model from a solution, each variable in "
-                            "model.initial_conditions must appear in the solution with "
-                            "the same key as the variable name. In the solution "
-                            f"provided, {e.args[0]}"
-                        ) from e
-                    if isinstance(solution, pybamm.Solution):
-                        final_state = final_state.data
-                    if final_state.ndim == 2:
-                        final_state_eval = final_state[:, -1]
-                    else:
-                        raise NotImplementedError(
-                            "Variable in concatenation must be 1D"
-                        )
+                    final_state = get_variable_state(child.name)
+                    final_state_eval = get_final_state_eval(final_state)
                     children.append(final_state_eval)
                 final_state_eval = np.concatenate(children)
+
             else:
                 raise NotImplementedError(
                     "Variable must have type 'Variable' or 'Concatenation'"

src/pybamm/models/submodels/particle/base_particle.py

@@ -337,15 +337,7 @@ def _get_standard_concentration_distribution_variables(self, c_s):
             c_s_distribution = c_s
 
             # x-average the *tertiary* domain.
-            # NOTE: not yet implemented. Make 0.5 everywhere
-            c_s_xav_distribution = pybamm.FullBroadcast(
-                0.5,
-                [f"{domain} {phase_name}particle"],
-                {
-                    "secondary": f"{domain} {phase_name}particle size",
-                    "tertiary": "current collector",
-                },
-            )
+            c_s_xav_distribution = pybamm.x_average(c_s)
 
             # Surface concentration distribution variables
             c_s_surf_distribution = pybamm.surf(c_s)