#1051 fix 2D plot bug

CHANGELOG.md

@@ -1,3 +1,15 @@
+# Unreleased
+
+## Features
+
+## Optimizations
+
+## Bug fixes
+
+-   Fix a bug where variables that depend on y and z were transposed in `QuickPlot` ([#1055](https://github.com/pybamm-team/PyBaMM/pull/1055))
+
+## Breaking changes
+
 # [v0.2.2](https://github.com/pybamm-team/PyBaMM/tree/v0.2.2) - 2020-06-01
 
 New SEI models, simplification of submodel structure, as well as optimisations and general bug fixes.

pybamm/__init__.py

@@ -52,8 +52,12 @@ def version(formatted=False):
 script_path = os.path.abspath(__file__)
 
 from .util import root_dir
+
 ABSOLUTE_PATH = root_dir()
-PARAMETER_PATH = [os.getcwd(), os.path.join(root_dir(), "pybamm", "input", "parameters")]
+PARAMETER_PATH = [
+    os.getcwd(),
+    os.path.join(root_dir(), "pybamm", "input", "parameters"),
+]
 
 #
 # Utility classes and methods
@@ -217,7 +221,7 @@ def version(formatted=False):
 #
 # Plotting
 #
-from .plotting.quick_plot import QuickPlot
+from .plotting.quick_plot import QuickPlot, close_plots
 from .plotting.plot import plot
 from .plotting.plot2D import plot2D
 from .plotting.dynamic_plot import dynamic_plot

pybamm/plotting/quick_plot.py

@@ -44,6 +44,13 @@ def split_long_string(title, max_words=4):
         return first_line + "\n" + second_line
 
 
+def close_plots():
+    "Close all open figures"
+    import matplotlib.pyplot as plt
+
+    plt.close("all")
+
+
 class QuickPlot(object):
     """
     Generates a quick plot of a subset of key outputs of the model so that the model
@@ -265,6 +272,7 @@ def set_output_variables(self, output_variables, solutions):
         self.first_spatial_scale = {}
         self.second_spatial_scale = {}
         self.is_x_r = {}
+        self.is_y_z = {}
 
         # Calculate subplot positions based on number of variables supplied
         self.subplot_positions = {}
@@ -354,8 +362,15 @@ def set_output_variables(self, output_variables, solutions):
                     )
                     if first_spatial_var_name == "r" and second_spatial_var_name == "x":
                         self.is_x_r[variable_tuple] = True
+                        self.is_y_z[variable_tuple] = False
+                    elif (
+                        first_spatial_var_name == "y" and second_spatial_var_name == "z"
+                    ):
+                        self.is_x_r[variable_tuple] = False
+                        self.is_y_z[variable_tuple] = True
                     else:
                         self.is_x_r[variable_tuple] = False
+                        self.is_y_z[variable_tuple] = False
 
             # Store variables and subplot position
             self.variables[variable_tuple] = variables
@@ -585,17 +600,24 @@ def plot(self, t):
                     y_name = list(spatial_vars.keys())[1][0]
                     x = self.first_dimensional_spatial_variable[key]
                     y = self.second_dimensional_spatial_variable[key]
-                    var = variable(t_in_seconds, **spatial_vars, warn=False).T
+                    # need to transpose if domain is x-z
+                    if self.is_y_z[key] is True:
+                        var = variable(t_in_seconds, **spatial_vars, warn=False)
+                    else:
+                        var = variable(t_in_seconds, **spatial_vars, warn=False).T
                 ax.set_xlabel(
                     "{} [{}]".format(x_name, self.spatial_unit), fontsize=fontsize
                 )
                 ax.set_ylabel(
                     "{} [{}]".format(y_name, self.spatial_unit), fontsize=fontsize
                 )
                 vmin, vmax = self.variable_limits[key]
-                ax.contourf(
+                # store the plot and the var data (for testing) as cant access
+                # z data from QuadContourSet object
+                self.plots[key][0][0] = ax.contourf(
                     x, y, var, levels=100, vmin=vmin, vmax=vmax, cmap="coolwarm"
                 )
+                self.plots[key][0][1] = var
                 if vmin is None and vmax is None:
                     vmin = ax_min(var)
                     vmax = ax_max(var)
@@ -717,10 +739,17 @@ def slider_update(self, t):
                 else:
                     x = self.first_dimensional_spatial_variable[key]
                     y = self.second_dimensional_spatial_variable[key]
-                    var = variable(time_in_seconds, **spatial_vars, warn=False).T
-                ax.contourf(
+                    # need to transpose if domain is x-z
+                    if self.is_y_z[key] is True:
+                        var = variable(time_in_seconds, **spatial_vars, warn=False)
+                    else:
+                        var = variable(time_in_seconds, **spatial_vars, warn=False).T
+                # store the plot and the updated var data (for testing) as cant
+                # access z data from QuadContourSet object
+                self.plots[key][0][0] = ax.contourf(
                     x, y, var, levels=100, vmin=vmin, vmax=vmax, cmap="coolwarm"
                 )
+                self.plots[key][0][1] = var
                 if (vmin, vmax) == (None, None):
                     vmin = ax_min(var)
                     vmax = ax_max(var)

tests/unit/test_quick_plot.py

@@ -269,6 +269,8 @@ def test_simple_ode_model(self):
         ):
             pybamm.QuickPlot(solution, ["NaN variable"])
 
+        pybamm.close_plots()
+
     def test_spm_simulation(self):
         # SPM
         model = pybamm.lithium_ion.SPM()
@@ -282,26 +284,51 @@ def test_spm_simulation(self):
         quick_plot = pybamm.QuickPlot([sim, sim.solution])
         quick_plot.plot(0)
 
-    def test_loqs_spm_base(self):
+        pybamm.close_plots()
+
+    def test_loqs_spme(self):
         t_eval = np.linspace(0, 10, 2)
 
-        # SPM
-        for model in [pybamm.lithium_ion.SPM(), pybamm.lead_acid.LOQS()]:
+        for model in [pybamm.lithium_ion.SPMe(), pybamm.lead_acid.LOQS()]:
             geometry = model.default_geometry
             param = model.default_parameter_values
             param.process_model(model)
             param.process_geometry(geometry)
-            mesh = pybamm.Mesh(
-                geometry, model.default_submesh_types, model.default_var_pts
-            )
+            var = pybamm.standard_spatial_vars
+            var_pts = {var.x_n: 5, var.x_s: 5, var.x_p: 5, var.r_n: 5, var.r_p: 5}
+            mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts)
             disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
             disc.process_model(model)
             solver = model.default_solver
             solution = solver.solve(model, t_eval)
             pybamm.QuickPlot(solution)
 
-            # test quick plot of particle for spm
-            if model.name == "Single Particle Model":
+            # check 1D (space) variables update properly for different time units
+            c_e = solution["Electrolyte concentration [mol.m-3]"].entries
+
+            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],
+                )[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],
+                )[0][:, 0]
+                np.testing.assert_array_almost_equal(qp_data, c_e[:, 1])
+
+            # test quick plot of particle for spme
+            if model.name == "Single Particle Model with electrolyte":
                 output_variables = [
                     "X-averaged negative particle concentration [mol.m-3]",
                     "X-averaged positive particle concentration [mol.m-3]",
@@ -310,6 +337,61 @@ def test_loqs_spm_base(self):
                 ]
                 pybamm.QuickPlot(solution, output_variables)
 
+                # check 2D (space) variables update properly for different time units
+                c_n = solution["Negative particle concentration [mol.m-3]"].entries
+
+                for unit, scale in zip(["seconds", "minutes", "hours"], [1, 60, 3600]):
+                    quick_plot = pybamm.QuickPlot(
+                        solution,
+                        ["Negative particle concentration [mol.m-3]"],
+                        time_unit=unit,
+                    )
+                    quick_plot.plot(0)
+                    qp_data = quick_plot.plots[
+                        ("Negative particle concentration [mol.m-3]",)
+                    ][0][1]
+                    np.testing.assert_array_almost_equal(qp_data, c_n[:, :, 0])
+                    quick_plot.slider_update(t_eval[-1] / scale)
+                    qp_data = quick_plot.plots[
+                        ("Negative particle concentration [mol.m-3]",)
+                    ][0][1]
+                    np.testing.assert_array_almost_equal(qp_data, c_n[:, :, 1])
+
+        pybamm.close_plots()
+
+    def test_plot_1plus1D_spme(self):
+        spm = pybamm.lithium_ion.SPMe(
+            {"current collector": "potential pair", "dimensionality": 1}
+        )
+        geometry = spm.default_geometry
+        param = spm.default_parameter_values
+        param.process_model(spm)
+        param.process_geometry(geometry)
+        var = pybamm.standard_spatial_vars
+        var_pts = {var.x_n: 5, var.x_s: 5, var.x_p: 5, var.r_n: 5, var.r_p: 5, var.z: 5}
+        mesh = pybamm.Mesh(geometry, spm.default_submesh_types, var_pts)
+        disc_spm = pybamm.Discretisation(mesh, spm.default_spatial_methods)
+        disc_spm.process_model(spm)
+        t_eval = np.linspace(0, 100, 10)
+        solution = spm.default_solver.solve(spm, t_eval)
+
+        # check 2D (x,z space) variables update properly for different time units
+        # Note: these should be the transpose of the entries in the processed variable
+        c_e = solution["Electrolyte concentration [mol.m-3]"].entries
+
+        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)
+            qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
+            np.testing.assert_array_almost_equal(qp_data.T, c_e[:, :, 0])
+            quick_plot.slider_update(t_eval[-1] / scale)
+            qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
+            np.testing.assert_array_almost_equal(qp_data.T, c_e[:, :, -1])
+
+        pybamm.close_plots()
+
     def test_plot_2plus1D_spm(self):
         spm = pybamm.lithium_ion.SPM(
             {"current collector": "potential pair", "dimensionality": 2}
@@ -331,11 +413,11 @@ def test_plot_2plus1D_spm(self):
         mesh = pybamm.Mesh(geometry, spm.default_submesh_types, var_pts)
         disc_spm = pybamm.Discretisation(mesh, spm.default_spatial_methods)
         disc_spm.process_model(spm)
-        t_eval = np.linspace(0, 3600, 100)
-        solution_spm = spm.default_solver.solve(spm, t_eval)
+        t_eval = np.linspace(0, 100, 10)
+        solution = spm.default_solver.solve(spm, t_eval)
 
         quick_plot = pybamm.QuickPlot(
-            solution_spm,
+            solution,
             [
                 "Negative current collector potential [V]",
                 "Positive current collector potential [V]",
@@ -345,10 +427,28 @@ def test_plot_2plus1D_spm(self):
         quick_plot.dynamic_plot(testing=True)
         quick_plot.slider_update(1)
 
-        with self.assertRaisesRegex(NotImplementedError, "Shape not recognized for"):
-            pybamm.QuickPlot(
-                solution_spm, ["Negative particle concentration [mol.m-3]"]
+        # check 2D (y,z space) variables update properly for different time units
+        phi_n = solution["Negative current collector potential [V]"].entries
+
+        for unit, scale in zip(["seconds", "minutes", "hours"], [1, 60, 3600]):
+            quick_plot = pybamm.QuickPlot(
+                solution, ["Negative current collector potential [V]"], time_unit=unit
             )
+            quick_plot.plot(0)
+            qp_data = quick_plot.plots[("Negative current collector potential [V]",)][
+                0
+            ][1]
+            np.testing.assert_array_almost_equal(qp_data, phi_n[:, :, 0])
+            quick_plot.slider_update(t_eval[-1] / scale)
+            qp_data = quick_plot.plots[("Negative current collector potential [V]",)][
+                0
+            ][1]
+            np.testing.assert_array_almost_equal(qp_data, phi_n[:, :, -1])
+
+        with self.assertRaisesRegex(NotImplementedError, "Shape not recognized for"):
+            pybamm.QuickPlot(solution, ["Negative particle concentration [mol.m-3]"])
+
+        pybamm.close_plots()
 
     def test_failure(self):
         with self.assertRaisesRegex(TypeError, "solutions must be"):