Repository: pybamm-team/PyBaMM

Issue:
Add length scales to models
Add `x_scale`, `r_scale`, etc. to models similar to how they have a timescale. This would clean up `ProcessedVariable` and `QuickPlot` and make it easy for users to see how lengths are scaled in a given model

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 852cadcc0..62242ee1d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,8 @@
 
 ## Features
 
--   Add averaging in secondary dimensions ([#1057](https://github.com/pybamm-team/PyBaMM/pull/1057))
+-   Added `length_scales` attribute to models ([#1058](https://github.com/pybamm-team/PyBaMM/pull/1058))
+-   Added averaging in secondary dimensions ([#1057](https://github.com/pybamm-team/PyBaMM/pull/1057))
 
 ## Optimizations
 
diff --git a/examples/scripts/compare_comsol/compare_comsol_DFN.py b/examples/scripts/compare_comsol/compare_comsol_DFN.py
index 876c59f3a..a9cafb0f1 100644
--- a/examples/scripts/compare_comsol/compare_comsol_DFN.py
+++ b/examples/scripts/compare_comsol/compare_comsol_DFN.py
@@ -119,7 +119,7 @@ comsol_voltage.mesh = None
 comsol_voltage.secondary_mesh = None
 
 # Create comsol model with dictionary of Matrix variables
-comsol_model = pybamm.BaseModel()
+comsol_model = pybamm.lithium_ion.BaseModel()
 comsol_model.variables = {
     "Negative particle surface concentration [mol.m-3]": comsol_c_n_surf,
     "Electrolyte concentration [mol.m-3]": comsol_c_e,
@@ -129,15 +129,13 @@ comsol_model.variables = {
     "Electrolyte potential [V]": comsol_phi_e,
     "Positive electrode potential [V]": comsol_phi_p,
     "Terminal voltage [V]": comsol_voltage,
-    # Add spatial variables for use in QuickPlot
-    "x": pybamm_model.variables["x"],
-    "x [m]": pybamm_model.variables["x [m]"],
 }
 
 # Make new solution with same t and y
 comsol_solution = pybamm.Solution(pybamm_solution.t, pybamm_solution.y)
-# Update model timescale to match the pybamm model
+# Update model scales to match the pybamm model
 comsol_model.timescale = pybamm_model.timescale
+comsol_model.length_scales = pybamm_model.length_scales
 comsol_solution.model = comsol_model
 
 # plot
diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py
index 0776e40ba..bb158382b 100644
--- a/pybamm/models/base_model.py
+++ b/pybamm/models/base_model.py
@@ -128,6 +128,7 @@ class BaseModel(object):
 
         # Default timescale is 1 second
         self.timescale = pybamm.Scalar(1)
+        self.length_scales = {}
 
     @property
     def name(self):
@@ -335,6 +336,7 @@ class BaseModel(object):
         new_model.use_simplify = self.use_simplify
         new_model.convert_to_format = self.convert_to_format
         new_model.timescale = self.timescale
+        new_model.length_scales = self.length_scales
         return new_model
 
     def update(self, *submodels):
diff --git a/pybamm/models/full_battery_models/lead_acid/base_lead_acid_model.py b/pybamm/models/full_battery_models/lead_acid/base_lead_acid_model.py
index e961af923..463a2c8b6 100644
--- a/pybamm/models/full_battery_models/lead_acid/base_lead_acid_model.py
+++ b/pybamm/models/full_battery_models/lead_acid/base_lead_acid_model.py
@@ -21,6 +21,16 @@ class BaseModel(pybamm.BaseBatteryModel):
 
         # Default timescale is discharge timescale
         self.timescale = self.param.tau_discharge
+
+        # Set default length scales
+        self.length_scales = {
+            "negative electrode": self.param.L_x,
+            "separator": self.param.L_x,
+            "positive electrode": self.param.L_x,
+            "current collector y": self.param.L_y,
+            "current collector z": self.param.L_z,
+        }
+
         self.set_standard_output_variables()
 
     @property
diff --git a/pybamm/models/full_battery_models/lead_acid/basic_full.py b/pybamm/models/full_battery_models/lead_acid/basic_full.py
index 9b6408752..60addb80f 100644
--- a/pybamm/models/full_battery_models/lead_acid/basic_full.py
+++ b/pybamm/models/full_battery_models/lead_acid/basic_full.py
@@ -288,8 +288,6 @@ class BasicFull(BaseModel):
             - param.U_n_ref
             + pot * phi_s_p,
             "Terminal voltage [V]": param.U_p_ref - param.U_n_ref + pot * voltage,
-            "x [m]": pybamm.standard_spatial_vars.x * param.L_x,
-            "x": pybamm.standard_spatial_vars.x,
             "Porosity": eps,
             "Volume-averaged velocity": v,
             "X-averaged separator transverse volume-averaged velocity": div_V_s,
diff --git a/pybamm/models/full_battery_models/lithium_ion/base_lithium_ion_model.py b/pybamm/models/full_battery_models/lithium_ion/base_lithium_ion_model.py
index 5176da217..bfa6d963c 100644
--- a/pybamm/models/full_battery_models/lithium_ion/base_lithium_ion_model.py
+++ b/pybamm/models/full_battery_models/lithium_ion/base_lithium_ion_model.py
@@ -19,6 +19,17 @@ class BaseModel(pybamm.BaseBatteryModel):
 
         # Default timescale is discharge timescale
         self.timescale = self.param.tau_discharge
+
+        # Set default length scales
+        self.length_scales = {
+            "negative electrode": self.param.L_x,
+            "separator": self.param.L_x,
+            "positive electrode": self.param.L_x,
+            "negative particle": self.param.R_n,
+            "positive particle": self.param.R_p,
+            "current collector y": self.param.L_y,
+            "current collector z": self.param.L_z,
+        }
         self.set_standard_output_variables()
 
     def set_standard_output_variables(self):
diff --git a/pybamm/parameters/parameter_values.py b/pybamm/parameters/parameter_values.py
index d9afc3084..48361f8bf 100644
--- a/pybamm/parameters/parameter_values.py
+++ b/pybamm/parameters/parameter_values.py
@@ -386,6 +386,10 @@ class ParameterValues:
         # Process timescale
         model.timescale = self.process_symbol(model.timescale)
 
+        # Process length scales
+        for domain, scale in model.length_scales.items():
+            model.length_scales[domain] = self.process_symbol(scale)
+
         pybamm.logger.info("Finish setting parameters for {}".format(model.name))
 
         return model
diff --git a/pybamm/plotting/quick_plot.py b/pybamm/plotting/quick_plot.py
index ed5649dd4..a965a4a0b 100644
--- a/pybamm/plotting/quick_plot.py
+++ b/pybamm/plotting/quick_plot.py
@@ -146,30 +146,11 @@ class QuickPlot(object):
         else:
             raise ValueError("spatial unit '{}' not recognized".format(spatial_unit))
 
-        variables = models[0].variables
-        # empty spatial scales, will raise error later if can't find a particular one
-        self.spatial_scales = {}
-        if "x [m]" in variables and "x" in variables:
-            x_scale = (variables["x [m]"] / variables["x"]).evaluate()[
-                -1
-            ] * self.spatial_factor
-            self.spatial_scales.update({dom: x_scale for dom in variables["x"].domain})
-        if "y [m]" in variables and "y" in variables:
-            self.spatial_scales["current collector y"] = (
-                variables["y [m]"] / variables["y"]
-            ).evaluate()[-1] * self.spatial_factor
-        if "z [m]" in variables and "z" in variables:
-            self.spatial_scales["current collector z"] = (
-                variables["z [m]"] / variables["z"]
-            ).evaluate()[-1] * self.spatial_factor
-        if "r_n [m]" in variables and "r_n" in variables:
-            self.spatial_scales["negative particle"] = (
-                variables["r_n [m]"] / variables["r_n"]
-            ).evaluate()[-1] * self.spatial_factor
-        if "r_p [m]" in variables and "r_p" in variables:
-            self.spatial_scales["positive particle"] = (
-                variables["r_p [m]"] / variables["r_p"]
-            ).evaluate()[-1] * self.spatial_factor
+        # Set length scales
+        self.length_scales = {
+            domain: scale.evaluate() * self.spatial_factor
+            for domain, scale in models[0].length_scales.items()
+        }
 
         # Time parameters
         model_timescale_in_seconds = models[0].timescale_eval
@@ -269,8 +250,6 @@ class QuickPlot(object):
         self.spatial_variable_dict = {}
         self.first_dimensional_spatial_variable = {}
         self.second_dimensional_spatial_variable = {}
-        self.first_spatial_scale = {}
-        self.second_spatial_scale = {}
         self.is_x_r = {}
         self.is_y_z = {}
 
@@ -317,18 +296,15 @@ class QuickPlot(object):
 
             # Set the x variable (i.e. "x" or "r" for any one-dimensional variables)
             if first_variable.dimensions == 1:
-                (
-                    spatial_var_name,
-                    spatial_var_value,
-                    spatial_scale,
-                ) = self.get_spatial_var(variable_tuple, first_variable, "first")
+                (spatial_var_name, spatial_var_value,) = self.get_spatial_var(
+                    variable_tuple, first_variable, "first"
+                )
                 self.spatial_variable_dict[variable_tuple] = {
                     spatial_var_name: spatial_var_value
                 }
                 self.first_dimensional_spatial_variable[variable_tuple] = (
                     spatial_var_value * self.spatial_factor
                 )
-                self.first_spatial_scale[variable_tuple] = spatial_scale
 
             elif first_variable.dimensions == 2:
                 # Don't allow 2D variables if there are multiple solutions
@@ -343,12 +319,10 @@ class QuickPlot(object):
                     (
                         first_spatial_var_name,
                         first_spatial_var_value,
-                        first_spatial_scale,
                     ) = self.get_spatial_var(variable_tuple, first_variable, "first")
                     (
                         second_spatial_var_name,
                         second_spatial_var_value,
-                        second_spatial_scale,
                     ) = self.get_spatial_var(variable_tuple, first_variable, "second")
                     self.spatial_variable_dict[variable_tuple] = {
                         first_spatial_var_name: first_spatial_var_value,
@@ -397,19 +371,7 @@ class QuickPlot(object):
         if domain == "current collector":
             domain += " {}".format(spatial_var_name)
 
-        # Get scale to go from dimensionless to dimensional in the units
-        # specified by spatial_unit
-        try:
-            spatial_scale = self.spatial_scales[domain]
-        except KeyError:
-            raise KeyError(
-                (
-                    "Can't find spatial scale for '{}', make sure both '{} [m]' "
-                    + "and '{}' are defined in the model variables"
-                ).format(domain, *[spatial_var_name] * 2)
-            )
-
-        return spatial_var_name, spatial_var_value, spatial_scale
+        return spatial_var_name, spatial_var_value
 
     def reset_axis(self):
         """
@@ -578,10 +540,14 @@ class QuickPlot(object):
                 # add dashed lines for boundaries between subdomains
                 y_min, y_max = ax.get_ylim()
                 ax.set_ylim(y_min, y_max)
-                for bnd in variable_lists[0][0].internal_boundaries:
-                    bnd_dim = bnd * self.first_spatial_scale[key]
+                for boundary in variable_lists[0][0].internal_boundaries:
+                    boundary_scaled = boundary * self.spatial_factor
                     ax.plot(
-                        [bnd_dim, bnd_dim], [y_min, y_max], color="0.5", lw=1, zorder=0
+                        [boundary_scaled, boundary_scaled],
+                        [y_min, y_max],
+                        color="0.5",
+                        lw=1,
+                        zorder=0,
                     )
             elif variable_lists[0][0].dimensions == 2:
                 # Read dictionary of spatial variables
@@ -716,10 +682,10 @@ class QuickPlot(object):
                 if y_min is None and y_max is None:
                     y_min, y_max = ax_min(var_min), ax_max(var_max)
                     ax.set_ylim(y_min, y_max)
-                    for bnd in self.variables[key][0][0].internal_boundaries:
-                        bnd_dim = bnd * self.first_spatial_scale[key]
+                    for boundary in self.variables[key][0][0].internal_boundaries:
+                        boundary_scaled = boundary * self.spatial_factor
                         ax.plot(
-                            [bnd_dim, bnd_dim],
+                            [boundary_scaled, boundary_scaled],
                             [y_min, y_max],
                             color="0.5",
                             lw=1,
diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
index bec894688..8e0096f41 100644
--- a/pybamm/solvers/processed_variable.py
+++ b/pybamm/solvers/processed_variable.py
@@ -63,18 +63,12 @@ class ProcessedVariable(object):
         self.timescale = solution.model.timescale.evaluate()
         self.t_pts = self.t_sol * self.timescale
 
-        # Store spatial variables to get scales
-        self.spatial_vars = {}
+        # Store length scales
         if solution.model:
-            for var in ["x", "y", "z", "r_n", "r_p"]:
-                if (
-                    var in solution.model.variables
-                    and var + " [m]" in solution.model.variables
-                ):
-                    self.spatial_vars[var] = solution.model.variables[var]
-                    self.spatial_vars[var + " [m]"] = solution.model.variables[
-                        var + " [m]"
-                    ]
+            self.length_scales = {
+                domain: scale.evaluate()
+                for domain, scale in solution.model.length_scales.items()
+            }
 
         # Evaluate base variable at initial time
         if self.known_evals:
@@ -222,16 +216,20 @@ class ProcessedVariable(object):
             self.x_sol = space
 
         # assign attributes for reference
-        self.first_dim_pts = space * self.get_spatial_scale(
-            self.first_dimension, self.domain[0]
-        )
-        self.internal_boundaries = self.mesh.internal_boundaries
+        length_scale = self.get_spatial_scale(self.first_dimension, self.domain[0])
+        pts_for_interp = space * length_scale
+        self.internal_boundaries = [
+            bnd * length_scale for bnd in self.mesh.internal_boundaries
+        ]
+
+        # Set first_dim_pts to edges for nicer plotting
+        self.first_dim_pts = edges * length_scale
 
         # set up interpolation
         if len(self.t_sol) == 1:
             # function of space only
             interpolant = interp.interp1d(
-                self.first_dim_pts,
+                pts_for_interp,
                 entries_for_interp[:, 0],
                 kind="linear",
                 fill_value=np.nan,
@@ -250,7 +248,7 @@ class ProcessedVariable(object):
             # is the reverse of what you'd expect
             self._interpolation_function = interp.interp2d(
                 self.t_pts,
-                self.first_dim_pts,
+                pts_for_interp,
                 entries_for_interp,
                 kind="linear",
                 fill_value=np.nan,
@@ -325,12 +323,15 @@ class ProcessedVariable(object):
         # assign attributes for reference
         self.entries = entries
         self.dimensions = 2
-        self.first_dim_pts = first_dim_pts * self.get_spatial_scale(
+        first_length_scale = self.get_spatial_scale(
             self.first_dimension, self.domain[0]
         )
-        self.second_dim_pts = second_dim_pts * self.get_spatial_scale(
-            self.second_dimension
+        self.first_dim_pts = first_dim_pts * first_length_scale
+
+        second_length_scale = self.get_spatial_scale(
+            self.second_dimension, self.auxiliary_domains["secondary"][0]
         )
+        self.second_dim_pts = second_dim_pts * second_length_scale
 
         # set up interpolation
         if len(self.t_sol) == 1:
@@ -394,8 +395,8 @@ class ProcessedVariable(object):
         self.z_sol = z_sol
         self.first_dimension = "y"
         self.second_dimension = "z"
-        self.first_dim_pts = y_sol * self.get_spatial_scale("y")
-        self.second_dim_pts = z_sol * self.get_spatial_scale("z")
+        self.first_dim_pts = y_sol * self.get_spatial_scale("y", "current collector")
+        self.second_dim_pts = z_sol * self.get_spatial_scale("z", "current collector")
 
         # set up interpolation
         if len(self.t_sol) == 1:
@@ -477,27 +478,22 @@ class ProcessedVariable(object):
                 second_dim = second_dim[:, np.newaxis]
         return self._interpolation_function((first_dim, second_dim, t))
 
-    def get_spatial_scale(self, name, domain=None):
+    def get_spatial_scale(self, name, domain):
         "Returns the spatial scale for a named spatial variable"
-        # Different scale in negative and positive particles
-        if domain == "negative particle":
-            name = "r_n"
-        elif domain == "positive particle":
-            name = "r_p"
-
-        # Try to get length scale
-        if name + " [m]" in self.spatial_vars and name in self.spatial_vars:
-            scale = (
-                self.spatial_vars[name + " [m]"] / self.spatial_vars[name]
-            ).evaluate()[-1]
-        else:
+        try:
+            if name == "y" and domain == "current collector":
+                return self.length_scales["current collector y"]
+            elif name == "z" and domain == "current collector":
+                return self.length_scales["current collector z"]
+            else:
+                return self.length_scales[domain]
+        except KeyError:
             if self.warn:
                 pybamm.logger.warning(
-                    "No scale set for spatial variable {}. "
-                    "Using default of 1 [m].".format(name)
+                    "No length scale set for {}. "
+                    "Using default of 1 [m].".format(domain)
                 )
-            scale = 1
-        return scale
+            return 1
 
     @property
     def data(self):


Test patch:
diff --git a/tests/integration/test_models/standard_output_tests.py b/tests/integration/test_models/standard_output_tests.py
index a20d5f4e5..4a48d2791 100644
--- a/tests/integration/test_models/standard_output_tests.py
+++ b/tests/integration/test_models/standard_output_tests.py
@@ -629,13 +629,13 @@ class VelocityTests(BaseOutputTest):
 
     def test_velocity_boundaries(self):
         """Test the boundary values of the current densities"""
-        L_x = self.v_box.first_dim_pts[-1] / self.v_box.x_sol[-1]
+        L_x = self.x_edge[-1]
         np.testing.assert_array_almost_equal(self.v_box(self.t, 0), 0, decimal=4)
         np.testing.assert_array_almost_equal(self.v_box(self.t, L_x), 0, decimal=4)
 
     def test_vertical_velocity(self):
         """Test the boundary values of the current densities"""
-        L_x = self.v_box.first_dim_pts[-1] / self.v_box.x_sol[-1]
+        L_x = self.x_edge[-1]
         np.testing.assert_array_equal(self.dVbox_dz(self.t, 0), 0)
         np.testing.assert_array_less(self.dVbox_dz(self.t, 0.5 * L_x), 0)
         np.testing.assert_array_equal(self.dVbox_dz(self.t, L_x), 0)
diff --git a/tests/unit/test_parameters/test_parameter_values.py b/tests/unit/test_parameters/test_parameter_values.py
index 94204b6a3..bd85c292d 100644
--- a/tests/unit/test_parameters/test_parameter_values.py
+++ b/tests/unit/test_parameters/test_parameter_values.py
@@ -511,6 +511,9 @@ class TestParameterValues(unittest.TestCase):
             "grad_var1": pybamm.grad(var1),
             "d_var1": d * var1,
         }
+        model.timescale = b
+        model.length_scales = {"test": c}
+
         parameter_values = pybamm.ParameterValues({"a": 1, "b": 2, "c": 3, "d": 42})
         parameter_values.process_model(model)
         # rhs
@@ -547,6 +550,9 @@ class TestParameterValues(unittest.TestCase):
         self.assertTrue(
             isinstance(model.variables["d_var1"].children[1], pybamm.Variable)
         )
+        # timescale and length scales
+        self.assertEqual(model.timescale.evaluate(), 2)
+        self.assertEqual(model.length_scales["test"].evaluate(), 3)
 
         # bad boundary conditions
         model = pybamm.BaseModel()
diff --git a/tests/unit/test_quick_plot.py b/tests/unit/test_quick_plot.py
index dfda89566..ed69408f1 100644
--- a/tests/unit/test_quick_plot.py
+++ b/tests/unit/test_quick_plot.py
@@ -5,7 +5,7 @@ import numpy as np
 
 class TestQuickPlot(unittest.TestCase):
     def test_simple_ode_model(self):
-        model = pybamm.BaseBatteryModel(name="Simple ODE Model")
+        model = pybamm.lithium_ion.BaseModel(name="Simple ODE Model")
 
         whole_cell = ["negative electrode", "separator", "positive electrode"]
         # Create variables: domain is explicitly empty since these variables are only
@@ -37,13 +37,8 @@ class TestQuickPlot(unittest.TestCase):
             "c broadcasted positive electrode": pybamm.PrimaryBroadcast(
                 c, "positive particle"
             ),
-            "x [m]": pybamm.standard_spatial_vars.x,
-            "x": pybamm.standard_spatial_vars.x,
-            "r_n [m]": pybamm.standard_spatial_vars.r_n,
-            "r_n": pybamm.standard_spatial_vars.r_n,
-            "r_p [m]": pybamm.standard_spatial_vars.r_p,
-            "r_p": pybamm.standard_spatial_vars.r_p,
         }
+        model.timescale = pybamm.Scalar(1)
 
         # ODEs only (don't use jacobian)
         model.use_jacobian = False
@@ -255,13 +250,6 @@ class TestQuickPlot(unittest.TestCase):
                 [solution, solution], ["a"], labels=["sol 1", "sol 2", "sol 3"]
             )
 
-        # Remove 'x [m]' from the variables and make sure a key error is raise
-        del solution.model.variables["x [m]"]
-        with self.assertRaisesRegex(
-            KeyError, "Can't find spatial scale for 'negative electrode'"
-        ):
-            pybamm.QuickPlot(solution, ["b broadcasted"])
-
         # No variable can be NaN
         model.variables["NaN variable"] = disc.process_symbol(pybamm.Scalar(np.nan))
         with self.assertRaisesRegex(
@@ -304,26 +292,30 @@ class TestQuickPlot(unittest.TestCase):
             pybamm.QuickPlot(solution)
 
             # check 1D (space) variables update properly for different time units
-            c_e = solution["Electrolyte concentration [mol.m-3]"].entries
+            t = solution["Time [s]"].entries
+            c_e_var = solution["Electrolyte concentration [mol.m-3]"]
+            # 1D variables should be evaluated on edges
+            L_x = param.evaluate(pybamm.geometric_parameters.L_x)
+            c_e = c_e_var(t=t, x=mesh.combine_submeshes(*c_e_var.domain).edges * L_x)
 
             for unit, scale in zip(["seconds", "minutes", "hours"], [1, 60, 3600]):
                 quick_plot = pybamm.QuickPlot(
                     solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
                 )
                 quick_plot.plot(0)
-                # take off extrapolated points
+
                 qp_data = (
                     quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
                         0
-                    ].get_ydata()[1:-1],
+                    ].get_ydata(),
                 )[0]
                 np.testing.assert_array_almost_equal(qp_data, c_e[:, 0])
                 quick_plot.slider_update(t_eval[-1] / scale)
-                # take off extrapolated points
+
                 qp_data = (
                     quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
                         0
-                    ].get_ydata()[1:-1],
+                    ].get_ydata(),
                 )[0][:, 0]
                 np.testing.assert_array_almost_equal(qp_data, c_e[:, 1])