Repository: pybamm-team/PyBaMM (python)

Issue:
plotting of 2D current collector variables
Doing `sim.plot` with a 2D current collector variable doesn't give the correct plot, but also doesn't raise an error. We should either fix it or raise an error.

e.g. doing 
```
options = {
    "current collector": "potential pair",
    "dimensionality": 2,  # 2D current collectors
    "thermal": "x-lumped",
}
model = pybamm.lithium_ion.SPM(options)
sim = pybamm.Simulation(model)
solution = sim.solve()

sim.plot(["Negative current collector potential [V]"])
```

gives the wrong plot. But

```
phi_s_cn = solution["Negative current collector potential [V]"]

param = model.default_parameter_values
L_z = param.evaluate(pybamm.standard_parameters_lithium_ion.L_z)
l_y = phi_s_cp.y_sol[-1]  # dimensionless width
l_z = phi_s_cp.z_sol[-1]  # dimensionless height
y_plot = np.linspace(0, l_y, 21) * L_z
z_plot = np.linspace(0, l_z, 21) * L_z
t_plot = 1000

phi_s_cn_plot = plt.pcolormesh(
    y_plot, z_plot, phi_s_cn(y=y_plot, z=z_plot, t=t_plot), shading="gouraud"
)
```

gives the correct result. The result from `sim.plot` is transposed, so we just need to catch/fix this in quickplot.

Pull request description:
#1051 fix 2D plot bug
# Description

Fixes a bug where 2D variables that depend on y and z are transposed in `QuickPlot`. 

This was introduced in PR #1020 where the way 2D (y,z) variables are processed was changed, and I missed updating in quickplot as the scripts I had been using for 2+1D models all did the plotting manually... sorry!

Fixes #1051 

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/master/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [x] Bug fix (non-breaking change which fixes an issue)


# Key checklist:

- [x] No style issues: `$ flake8`
- [x] All tests pass: `$ python run-tests.py --unit`
- [x] The documentation builds: `$ cd docs` and then `$ make clean; make html`

You can run all three at once, using `$ python run-tests.py --quick`.

## Further checks:

- [ ] Code is commented, particularly in hard-to-understand areas
- [ ] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6ee541bda..c6f57ddeb 100644
--- a/CHANGELOG.md
+++ b/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.
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 066e6ef8d..9c16b3b55 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -52,8 +52,12 @@ FLOAT_FORMAT = "{: .17e}"
 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 @@ from . import experiments
 #
 # 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
diff --git a/pybamm/plotting/quick_plot.py b/pybamm/plotting/quick_plot.py
index a14b8d435..ed5649dd4 100644
--- a/pybamm/plotting/quick_plot.py
+++ b/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 @@ class QuickPlot(object):
         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 @@ class QuickPlot(object):
                     )
                     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,7 +600,11 @@ class QuickPlot(object):
                     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
                 )
@@ -593,9 +612,12 @@ class QuickPlot(object):
                     "{} [{}]".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 @@ class QuickPlot(object):
                 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)


Test patch:
diff --git a/tests/unit/test_quick_plot.py b/tests/unit/test_quick_plot.py
index a5dee5379..dfda89566 100644
--- a/tests/unit/test_quick_plot.py
+++ b/tests/unit/test_quick_plot.py
@@ -269,6 +269,8 @@ class TestQuickPlot(unittest.TestCase):
         ):
             pybamm.QuickPlot(solution, ["NaN variable"])
 
+        pybamm.close_plots()
+
     def test_spm_simulation(self):
         # SPM
         model = pybamm.lithium_ion.SPM()
@@ -282,26 +284,51 @@ class TestQuickPlot(unittest.TestCase):
         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 @@ class TestQuickPlot(unittest.TestCase):
                 ]
                 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 @@ class TestQuickPlot(unittest.TestCase):
         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 @@ class TestQuickPlot(unittest.TestCase):
         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"):