Processed variables now get spatial variables automatically

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Features
 
+- Processed variables now get the spatial variables automatically, allowing plotting of more generic models ([#3234](https://github.com/pybamm-team/PyBaMM/pull/3234))
+
+## Breaking changes
+
 - Numpy functions now work with PyBaMM symbols (e.g. `np.exp(pybamm.Symbol("a"))` returns `pybamm.Exp(pybamm.Symbol("a"))`). This means that parameter functions can be specified using numpy functions instead of pybamm functions. Additionally, combining numpy arrays with pybamm objects now works (the numpy array is converted to a pybamm array) ([#3205](https://github.com/pybamm-team/PyBaMM/pull/3205))
 
 ## Bug fixes

examples/scripts/compare_comsol/compare_comsol_DFN.py

@@ -102,6 +102,7 @@ def get_interp_fun(variable_name, domain):
 
 # Create comsol model with dictionary of Matrix variables
 comsol_model = pybamm.lithium_ion.BaseModel()
+comsol_model._geometry = pybamm_model.default_geometry
 comsol_model.variables = {
     "Negative particle surface concentration [mol.m-3]": comsol_c_n_surf,
     "Electrolyte concentration [mol.m-3]": comsol_c_e,

pybamm/discretisations/discretisation.py

@@ -241,6 +241,10 @@ def process_model(
             model_disc
         )
 
+        # Save geometry
+        pybamm.logger.verbose("Save geometry for {}".format(model.name))
+        model_disc._geometry = getattr(self.mesh, "_geometry", None)
+
         # Check that resulting model makes sense
         if check_model:
             pybamm.logger.verbose("Performing model checks for {}".format(model.name))

pybamm/meshes/meshes.py

@@ -25,6 +25,9 @@ class Mesh(dict):
     def __init__(self, geometry, submesh_types, var_pts):
         super().__init__()
 
+        # Save geometry
+        self.geometry = geometry
+
         # Preprocess var_pts
         var_pts_input = var_pts
         var_pts = {}
@@ -205,6 +208,14 @@ def add_ghost_meshes(self):
                 rgs_edges, submesh.coord_sys
             )
 
+    @property
+    def geometry(self):
+        return self._geometry
+
+    @geometry.setter
+    def geometry(self, geometry):
+        self._geometry = geometry
+
 
 class SubMesh:
     """

pybamm/models/base_model.py

@@ -113,6 +113,7 @@ def __init__(self, name="Unnamed model"):
         self._input_parameters = None
         self._parameter_info = None
         self._variables_casadi = {}
+        self._geometry = pybamm.Geometry({})
 
         # Default behaviour is to use the jacobian
         self.use_jacobian = True
@@ -312,6 +313,10 @@ def length_scales(self, values):
             "length_scales has been removed since models are now dimensional"
         )
 
+    @property
+    def geometry(self):
+        return self._geometry
+
     @property
     def default_var_pts(self):
         return {}

pybamm/simulation.py

@@ -1027,7 +1027,7 @@ def geometry(self):
 
     @geometry.setter
     def geometry(self, geometry):
-        self._geometry = geometry.copy()
+        self._geometry = geometry
 
     @property
     def parameter_values(self):

pybamm/solvers/processed_variable.py

@@ -56,6 +56,15 @@ def __init__(
         self.warn = warn
         self.cumtrapz_ic = cumtrapz_ic
 
+        # Process spatial variables
+        geometry = solution.all_models[0].geometry
+        self.spatial_variables = {}
+        for domain_level, domain_names in self.domains.items():
+            variables = []
+            for domain in domain_names:
+                variables += list(geometry[domain].keys())
+            self.spatial_variables[domain_level] = variables
+
         # Sensitivity starts off uninitialized, only set when called
         self._sensitivities = None
         self.solution_sensitivities = solution.sensitivities
@@ -172,25 +181,11 @@ def initialise_1D(self, fixed_t=False):
         # assign attributes for reference (either x_sol or r_sol)
         self.entries = entries
         self.dimensions = 1
-        if self.domain[0].endswith("particle"):
-            self.first_dimension = "r"
-            self.r_sol = space
-        elif self.domain[0] in [
-            "negative electrode",
-            "separator",
-            "positive electrode",
-        ]:
-            self.first_dimension = "x"
-            self.x_sol = space
-        elif self.domain == ["current collector"]:
-            self.first_dimension = "z"
-            self.z_sol = space
-        elif self.domain[0].endswith("particle size"):
-            self.first_dimension = "R"
-            self.R_sol = space
-        else:
-            self.first_dimension = "x"
-            self.x_sol = space
+        self.spatial_variable_names = {
+            k: self._process_spatial_variable_names(v)
+            for k, v in self.spatial_variables.items()
+        }
+        self.first_dimension = self.spatial_variable_names["primary"]
 
         # assign attributes for reference
         pts_for_interp = space
@@ -285,48 +280,13 @@ def initialise_2D(self):
             axis=1,
         )
 
-        # Process r-x, x-z, r-R, R-x, or R-z
-        if self.domain[0].endswith("particle") and self.domains["secondary"][
-            0
-        ].endswith("electrode"):
-            self.first_dimension = "r"
-            self.second_dimension = "x"
-            self.r_sol = first_dim_pts
-            self.x_sol = second_dim_pts
-        elif self.domain[0] in [
-            "negative electrode",
-            "separator",
-            "positive electrode",
-        ] and self.domains["secondary"] == ["current collector"]:
-            self.first_dimension = "x"
-            self.second_dimension = "z"
-            self.x_sol = first_dim_pts
-            self.z_sol = second_dim_pts
-        elif self.domain[0].endswith("particle") and self.domains["secondary"][
-            0
-        ].endswith("particle size"):
-            self.first_dimension = "r"
-            self.second_dimension = "R"
-            self.r_sol = first_dim_pts
-            self.R_sol = second_dim_pts
-        elif self.domain[0].endswith("particle size") and self.domains["secondary"][
-            0
-        ].endswith("electrode"):
-            self.first_dimension = "R"
-            self.second_dimension = "x"
-            self.R_sol = first_dim_pts
-            self.x_sol = second_dim_pts
-        elif self.domain[0].endswith("particle size") and self.domains["secondary"] == [
-            "current collector"
-        ]:
-            self.first_dimension = "R"
-            self.second_dimension = "z"
-            self.R_sol = first_dim_pts
-            self.z_sol = second_dim_pts
-        else:  # pragma: no cover
-            raise pybamm.DomainError(
-                f"Cannot process 2D object with domains '{self.domains}'."
-            )
+        self.spatial_variable_names = {
+            k: self._process_spatial_variable_names(v)
+            for k, v in self.spatial_variables.items()
+        }
+
+        self.first_dimension = self.spatial_variable_names["primary"]
+        self.second_dimension = self.spatial_variable_names["secondary"]
 
         # assign attributes for reference
         self.entries = entries
@@ -386,6 +346,35 @@ def initialise_2D_scikit_fem(self):
             coords={"y": y_sol, "z": z_sol, "t": self.t_pts},
         )
 
+    def _process_spatial_variable_names(self, spatial_variable):
+        if len(spatial_variable) == 0:
+            return None
+
+        # Extract names
+        raw_names = []
+        for var in spatial_variable:
+            # Ignore tabs in domain names
+            if var == "tabs":
+                continue
+            if isinstance(var, str):
+                raw_names.append(var)
+            else:
+                raw_names.append(var.name)
+
+        # Rename battery variables to match PyBaMM convention
+        if all([var.startswith("r") for var in raw_names]):
+            return "r"
+        elif all([var.startswith("x") for var in raw_names]):
+            return "x"
+        elif all([var.startswith("R") for var in raw_names]):
+            return "R"
+        elif len(raw_names) == 1:
+            return raw_names[0]
+        else:
+            raise NotImplementedError(
+                "Spatial variable name not recognized for {}".format(spatial_variable)
+            )
+
     def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None, warn=True):
         """
         Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R),

tests/__init__.py

@@ -36,5 +36,6 @@
     get_2p1d_discretisation_for_testing,
     get_unit_2p1D_mesh_for_testing,
     get_cylindrical_discretisation_for_testing,
+    get_base_model_with_battery_geometry,
 )
 from .testcase import TestCase

tests/shared.py

@@ -268,3 +268,9 @@ def get_cylindrical_discretisation_for_testing(
         mesh=get_cylindrical_mesh_for_testing(xpts, rpts, rcellpts, include_particles),
         cc_method=pybamm.FiniteVolume,
     )
+
+
+def get_base_model_with_battery_geometry(**kwargs):
+    model = pybamm.BaseModel()
+    model._geometry = pybamm.battery_geometry(**kwargs)
+    return model

tests/unit/test_meshes/test_meshes.py

@@ -35,6 +35,9 @@ def test_mesh_creation_no_parameters(self):
         # create mesh
         mesh = pybamm.Mesh(geometry, submesh_types, var_pts)
 
+        # check geometry
+        self.assertEqual(mesh.geometry, geometry)
+
         # check boundary locations
         self.assertEqual(mesh["negative particle"].edges[0], 0)
         self.assertEqual(mesh["negative particle"].edges[-1], 1)
@@ -77,6 +80,9 @@ def test_mesh_creation(self):
         # create mesh
         mesh = pybamm.Mesh(geometry, submesh_types, var_pts)
 
+        # check geometry
+        self.assertEqual(mesh.geometry, geometry)
+
         # check boundary locations
         self.assertEqual(mesh["negative electrode"].edges[0], 0)
         self.assertAlmostEqual(mesh["positive electrode"].edges[-1], 0.6)