From de4f9670043a99394ce09ac94bce21e5a5d268c1 Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Sun, 20 Oct 2024 19:18:59 +0530
Subject: [PATCH 1/8] Modified  plot method to take a list of times .

---
 src/pybamm/plotting/quick_plot.py | 74 ++++++++++++++++++-------------
 1 file changed, 44 insertions(+), 30 deletions(-)

diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py
index babfd2e761..fe50a76a8a 100644
--- a/src/pybamm/plotting/quick_plot.py
+++ b/src/pybamm/plotting/quick_plot.py
@@ -480,12 +480,12 @@ def reset_axis(self):
                 raise ValueError(f"Axis limits cannot be NaN for variables '{key}'")
 
     def plot(self, t, dynamic=False):
-        """Produces a quick plot with the internal states at time t.
+        """Produces a quick plot with the internal states at time t or a list of times.
 
         Parameters
         ----------
-        t : float
-            Dimensional time (in 'time_units') at which to plot.
+        t : float or list
+            Dimensional time (in 'time_units') at which to plot. Can be a single time or a list of times.
         dynamic : bool, optional
             Determine whether to allocate space for a slider at the bottom of the plot when generating a dynamic plot.
             If True, creates a dynamic plot with a slider.
@@ -496,9 +496,17 @@ def plot(self, t, dynamic=False):
         cm = import_optional_dependency("matplotlib", "cm")
         colors = import_optional_dependency("matplotlib", "colors")
 
-        t_in_seconds = t * self.time_scaling_factor
-        self.fig = plt.figure(figsize=self.figsize)
+        # Check if `t` is a float (single time) or a list (multiple times)
+        if isinstance(t, float):
+            t_list = [t]
+        elif isinstance(t, list):
+            t_list = t
+        else:
+            raise ValueError("t should be either a float or a list of floats")
 
+        t_in_seconds_list = [time * self.time_scaling_factor for time in t_list]
+
+        self.fig = plt.figure(figsize=self.figsize)
         self.gridspec = gridspec.GridSpec(self.n_rows, self.n_cols)
         self.plots = {}
         self.time_lines = {}
@@ -520,7 +528,7 @@ def plot(self, t, dynamic=False):
             variable_handles = []
             # Set labels for the first subplot only (avoid repetition)
             if variable_lists[0][0].dimensions == 0:
-                # 0D plot: plot as a function of time, indicating time t with a line
+                # 0D plot: plot as a function of time, indicating time t with a single line
                 ax.set_xlabel(f"Time [{self.time_unit}]")
                 for i, variable_list in enumerate(variable_lists):
                     for j, variable in enumerate(variable_list):
@@ -540,25 +548,26 @@ def plot(self, t, dynamic=False):
                         )
                         variable_handles.append(self.plots[key][0][j])
                     solution_handles.append(self.plots[key][i][0])
+
                 y_min, y_max = ax.get_ylim()
                 ax.set_ylim(y_min, y_max)
+
+                # Plot a single vertical time line for the first time `t`
                 (self.time_lines[key],) = ax.plot(
-                    [
-                        t_in_seconds / self.time_scaling_factor,
-                        t_in_seconds / self.time_scaling_factor,
-                    ],
+                    [t_in_seconds_list[0] / self.time_scaling_factor] * 2,
                     [y_min, y_max],
                     "k--",
                     lw=1.5,
                 )
+
             elif variable_lists[0][0].dimensions == 1:
                 # 1D plot: plot as a function of x at time t
                 # Read dictionary of spatial variables
                 spatial_vars = self.spatial_variable_dict[key]
                 spatial_var_name = next(iter(spatial_vars.keys()))
-                ax.set_xlabel(
-                    f"{spatial_var_name} [{self.spatial_unit}]",
-                )
+                ax.set_xlabel(f"{spatial_var_name} [{self.spatial_unit}]")
+
+                # Plot the variable at each time in the list
                 for i, variable_list in enumerate(variable_lists):
                     for j, variable in enumerate(variable_list):
                         if len(variable_list) == 1:
@@ -568,37 +577,44 @@ def plot(self, t, dynamic=False):
                             # multiple variables -> use linestyle to differentiate
                             # variables (color differentiates models)
                             linestyle = self.linestyles[j]
-                        (self.plots[key][i][j],) = ax.plot(
-                            self.first_spatial_variable[key],
-                            variable(t_in_seconds, **spatial_vars),
-                            color=self.colors[i],
-                            linestyle=linestyle,
-                            zorder=10,
-                        )
-                        variable_handles.append(self.plots[key][0][j])
+
+                        for t_in_seconds in t_in_seconds_list:
+                            (self.plots[key][i][j],) = ax.plot(
+                                self.first_spatial_variable[key],
+                                variable(t_in_seconds, **spatial_vars),
+                                color=self.colors[i],
+                                linestyle=linestyle,
+                                zorder=10,
+                                label=f"time = {t_in_seconds / self.time_scaling_factor} {self.time_unit}",
+                            )
+                            variable_handles.append(self.plots[key][0][j])
                     solution_handles.append(self.plots[key][i][0])
-                # add lines for boundaries between subdomains
+
+                # Add lines for boundaries between subdomains
                 for boundary in variable_lists[0][0].internal_boundaries:
                     boundary_scaled = boundary * self.spatial_factor
                     ax.axvline(boundary_scaled, color="0.5", lw=1, zorder=0)
+
             elif variable_lists[0][0].dimensions == 2:
-                # Read dictionary of spatial variables
+                # 2D variables do not support plotting for multiple times
+                if len(t_list) > 1:
+                    raise ValueError("Cannot plot 2D variables for multiple times")
+
+                # Continue with 2D plotting as before for a single time
                 spatial_vars = self.spatial_variable_dict[key]
-                # there can only be one entry in the variable list
                 variable = variable_lists[0][0]
-                # different order based on whether the domains are x-r, x-z or y-z, etc
                 if self.x_first_and_y_second[key] is False:
                     x_name = list(spatial_vars.keys())[1][0]
                     y_name = next(iter(spatial_vars.keys()))[0]
                     x = self.second_spatial_variable[key]
                     y = self.first_spatial_variable[key]
-                    var = variable(t_in_seconds, **spatial_vars)
+                    var = variable(t_in_seconds_list[0], **spatial_vars)
                 else:
                     x_name = next(iter(spatial_vars.keys()))[0]
                     y_name = list(spatial_vars.keys())[1][0]
                     x = self.first_spatial_variable[key]
                     y = self.second_spatial_variable[key]
-                    var = variable(t_in_seconds, **spatial_vars).T
+                    var = variable(t_in_seconds_list[0], **spatial_vars).T
                 ax.set_xlabel(f"{x_name} [{self.spatial_unit}]")
                 ax.set_ylabel(f"{y_name} [{self.spatial_unit}]")
                 vmin, vmax = self.variable_limits[key]
@@ -618,13 +634,11 @@ def plot(self, t, dynamic=False):
                         x, y, var, levels=100, vmin=vmin, vmax=vmax
                     )
                 self.plots[key][0][1] = var
-                if vmin is None and vmax is None:
-                    vmin = ax_min(var)
-                    vmax = ax_max(var)
                 self.colorbars[key] = self.fig.colorbar(
                     cm.ScalarMappable(colors.Normalize(vmin=vmin, vmax=vmax)),
                     ax=ax,
                 )
+
             # Set either y label or legend entries
             if len(key) == 1:
                 title = split_long_string(key[0])

From 8f52938b95fdf1e9e0be93799fc4cf4d0b5cd25c Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Sun, 20 Oct 2024 22:08:46 +0530
Subject: [PATCH 2/8] fixed failing tests

---
 tests/unit/test_plotting/test_quick_plot.py | 36 ++++++++++-----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py
index 188a725680..f7de702b29 100644
--- a/tests/unit/test_plotting/test_quick_plot.py
+++ b/tests/unit/test_plotting/test_quick_plot.py
@@ -73,7 +73,7 @@ def test_simple_ode_model(self, solver):
                 "c broadcasted positive electrode",
             ],
         )
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         # update the axis
         new_axis = [0, 0.5, 0, 1]
@@ -92,7 +92,7 @@ def test_simple_ode_model(self, solver):
         # Test with different output variables
         quick_plot = pybamm.QuickPlot(solution, ["b broadcasted"])
         assert len(quick_plot.axis_limits) == 1
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         quick_plot = pybamm.QuickPlot(
             solution,
@@ -105,7 +105,7 @@ def test_simple_ode_model(self, solver):
             ],
         )
         assert len(quick_plot.axis_limits) == 5
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         # update the axis
         new_axis = [0, 0.5, 0, 1]
@@ -124,7 +124,7 @@ def test_simple_ode_model(self, solver):
 
         # Test longer name
         quick_plot = pybamm.QuickPlot(solution, ["Variable with a very long name"])
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         # Test different inputs
         quick_plot = pybamm.QuickPlot(
@@ -154,7 +154,7 @@ def test_simple_ode_model(self, solver):
         quick_plot = pybamm.QuickPlot(solution, ["a"])
         assert quick_plot.time_scaling_factor == 1
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="seconds")
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         assert quick_plot.time_scaling_factor == 1
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot
@@ -163,7 +163,7 @@ def test_simple_ode_model(self, solver):
             quick_plot.plots[("a",)][0][0].get_ydata(), 0.2 * t_plot
         )
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="minutes")
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         assert quick_plot.time_scaling_factor == 60
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot / 60
@@ -172,7 +172,7 @@ def test_simple_ode_model(self, solver):
             quick_plot.plots[("a",)][0][0].get_ydata(), 0.2 * t_plot
         )
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="hours")
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         assert quick_plot.time_scaling_factor == 3600
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot / 3600
@@ -201,7 +201,7 @@ def test_simple_ode_model(self, solver):
 
         # Test 2D variables
         quick_plot = pybamm.QuickPlot(solution, ["2D variable"])
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         quick_plot.dynamic_plot(show_plot=False)
         quick_plot.slider_update(0.01)
 
@@ -217,14 +217,14 @@ def test_simple_ode_model(self, solver):
             None,
             None,
         ]
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
             solution, ["2D variable"], variable_limits="tight"
         )
         assert quick_plot.variable_limits[("2D variable",)] == (None, None)
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
@@ -234,7 +234,7 @@ def test_simple_ode_model(self, solver):
         )
         assert quick_plot.axis_limits[("a",)][2:] == [1, 2]
         assert quick_plot.axis_limits[("c broadcasted", "c broadcasted")][2:] == [3, 4]
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
@@ -242,7 +242,7 @@ def test_simple_ode_model(self, solver):
         )
         assert quick_plot.axis_limits[("a",)][2:] == [None, None]
         assert quick_plot.axis_limits[("b broadcasted",)][2:] != [None, None]
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
         quick_plot.slider_update(1)
 
         with pytest.raises(
@@ -288,12 +288,12 @@ def test_spm_simulation(self):
         # pass only a simulation object
         # it should be converted to a list of corresponding solution
         quick_plot = pybamm.QuickPlot(sim)
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         # mixed simulation and solution input
         # solution should be extracted from the simulation
         quick_plot = pybamm.QuickPlot([sim, sim.solution])
-        quick_plot.plot(0)
+        quick_plot.plot(0.0)
 
         # test creating a GIF
         with TemporaryDirectory() as dir_name:
@@ -336,7 +336,7 @@ def test_loqs_spme(self):
                 quick_plot = pybamm.QuickPlot(
                     solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
                 )
-                quick_plot.plot(0)
+                quick_plot.plot(0.0)
 
                 qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
                     0
@@ -371,7 +371,7 @@ def test_loqs_spme(self):
                         ["Negative particle concentration [mol.m-3]"],
                         time_unit=unit,
                     )
-                    quick_plot.plot(0)
+                    quick_plot.plot(0.0)
                     qp_data = quick_plot.plots[
                         ("Negative particle concentration [mol.m-3]",)
                     ][0][1]
@@ -411,7 +411,7 @@ def test_plot_1plus1D_spme(self):
             quick_plot = pybamm.QuickPlot(
                 solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
             )
-            quick_plot.plot(0)
+            quick_plot.plot(0.0)
             qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
             c_e_eval = c_e(t_eval[0], x=c_e.first_dim_pts, z=c_e.second_dim_pts)
             np.testing.assert_array_almost_equal(qp_data.T, c_e_eval)
@@ -456,7 +456,7 @@ def test_plot_2plus1D_spm(self):
             quick_plot = pybamm.QuickPlot(
                 solution, ["Negative current collector potential [V]"], time_unit=unit
             )
-            quick_plot.plot(0)
+            quick_plot.plot(0.0)
             qp_data = quick_plot.plots[("Negative current collector potential [V]",)][
                 0
             ][1]

From 173d514901d480879877baa58a20a0904b124cd8 Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Mon, 21 Oct 2024 18:48:37 +0530
Subject: [PATCH 3/8] added different colours to different time plots

---
 src/pybamm/plotting/quick_plot.py | 198 +++++++++++++-----------------
 1 file changed, 82 insertions(+), 116 deletions(-)

diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py
index fe50a76a8a..66a23dd41d 100644
--- a/src/pybamm/plotting/quick_plot.py
+++ b/src/pybamm/plotting/quick_plot.py
@@ -313,11 +313,11 @@ def set_output_variables(self, output_variables, solutions):
             first_variable = first_solution[0]
             domain = first_variable.domain
             # check all other solutions against the first solution
-            for idx, variable in enumerate(first_solution):
+            for _idx, variable in enumerate(first_solution):
                 if variable.domain != domain:
                     raise ValueError(
                         "Mismatching variable domains. "
-                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[idx]}' has domain '{variable.domain}'"
+                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[_idx]}' has domain '{variable.domain}'"
                     )
                 self.spatial_variable_dict[variable_tuple] = {}
 
@@ -480,7 +480,7 @@ def reset_axis(self):
                 raise ValueError(f"Axis limits cannot be NaN for variables '{key}'")
 
     def plot(self, t, dynamic=False):
-        """Produces a quick plot with the internal states at time t or a list of times.
+        """Produces a quick plot with the internal states at time t.
 
         Parameters
         ----------
@@ -493,20 +493,12 @@ def plot(self, t, dynamic=False):
 
         plt = import_optional_dependency("matplotlib.pyplot")
         gridspec = import_optional_dependency("matplotlib.gridspec")
-        cm = import_optional_dependency("matplotlib", "cm")
-        colors = import_optional_dependency("matplotlib", "colors")
 
-        # Check if `t` is a float (single time) or a list (multiple times)
         if isinstance(t, float):
-            t_list = [t]
-        elif isinstance(t, list):
-            t_list = t
-        else:
-            raise ValueError("t should be either a float or a list of floats")
-
-        t_in_seconds_list = [time * self.time_scaling_factor for time in t_list]
+            t = [t]
 
         self.fig = plt.figure(figsize=self.figsize)
+
         self.gridspec = gridspec.GridSpec(self.n_rows, self.n_cols)
         self.plots = {}
         self.time_lines = {}
@@ -516,6 +508,11 @@ def plot(self, t, dynamic=False):
         # initialize empty handles, to be created only if the appropriate plots are made
         solution_handles = []
 
+        # Generate distinct colors for each time point
+        time_colors = plt.cm.coolwarm(
+            np.linspace(0, 1, len(t))
+        )  # Use a colormap for distinct colors
+
         for k, (key, variable_lists) in enumerate(self.variables.items()):
             ax = self.fig.add_subplot(self.gridspec[k])
             self.axes.add(key, ax)
@@ -526,158 +523,127 @@ def plot(self, t, dynamic=False):
             ax.xaxis.set_major_locator(plt.MaxNLocator(3))
             self.plots[key] = defaultdict(dict)
             variable_handles = []
-            # Set labels for the first subplot only (avoid repetition)
+
             if variable_lists[0][0].dimensions == 0:
-                # 0D plot: plot as a function of time, indicating time t with a single line
+                # 0D plot: plot as a function of time, indicating multiple times with lines
                 ax.set_xlabel(f"Time [{self.time_unit}]")
                 for i, variable_list in enumerate(variable_lists):
                     for j, variable in enumerate(variable_list):
-                        if len(variable_list) == 1:
-                            # single variable -> use linestyle to differentiate model
-                            linestyle = self.linestyles[i]
-                        else:
-                            # multiple variables -> use linestyle to differentiate
-                            # variables (color differentiates models)
-                            linestyle = self.linestyles[j]
+                        linestyle = (
+                            self.linestyles[i]
+                            if len(variable_list) == 1
+                            else self.linestyles[j]
+                        )
                         full_t = self.ts_seconds[i]
                         (self.plots[key][i][j],) = ax.plot(
                             full_t / self.time_scaling_factor,
                             variable(full_t),
-                            color=self.colors[i],
+                            color=self.colors[i],  # Different colors for plot lines
                             linestyle=linestyle,
                         )
                         variable_handles.append(self.plots[key][0][j])
                     solution_handles.append(self.plots[key][i][0])
-
                 y_min, y_max = ax.get_ylim()
                 ax.set_ylim(y_min, y_max)
 
-                # Plot a single vertical time line for the first time `t`
-                (self.time_lines[key],) = ax.plot(
-                    [t_in_seconds_list[0] / self.time_scaling_factor] * 2,
-                    [y_min, y_max],
-                    "k--",
-                    lw=1.5,
-                )
+                # Add vertical lines for each time in the list, using different colors for each time
+                for _idx, t_single in enumerate(t):
+                    t_in_seconds = t_single * self.time_scaling_factor
+                    (self.time_lines[key],) = ax.plot(
+                        [
+                            t_in_seconds / self.time_scaling_factor,
+                            t_in_seconds / self.time_scaling_factor,
+                        ],
+                        [y_min, y_max],
+                        "--",  # Dashed lines
+                        lw=1.5,
+                        color=time_colors[_idx],  # Different color for each time
+                        label=f"t = {t_single:.2f} {self.time_unit}",
+                    )
+                ax.legend()
 
             elif variable_lists[0][0].dimensions == 1:
-                # 1D plot: plot as a function of x at time t
-                # Read dictionary of spatial variables
+                # 1D plot: plot as a function of x at different times
                 spatial_vars = self.spatial_variable_dict[key]
                 spatial_var_name = next(iter(spatial_vars.keys()))
                 ax.set_xlabel(f"{spatial_var_name} [{self.spatial_unit}]")
 
-                # Plot the variable at each time in the list
-                for i, variable_list in enumerate(variable_lists):
-                    for j, variable in enumerate(variable_list):
-                        if len(variable_list) == 1:
-                            # single variable -> use linestyle to differentiate model
-                            linestyle = self.linestyles[i]
-                        else:
-                            # multiple variables -> use linestyle to differentiate
-                            # variables (color differentiates models)
-                            linestyle = self.linestyles[j]
-
-                        for t_in_seconds in t_in_seconds_list:
+                for _idx, t_single in enumerate(t):
+                    t_in_seconds = t_single * self.time_scaling_factor
+
+                    for i, variable_list in enumerate(variable_lists):
+                        for j, variable in enumerate(variable_list):
+                            linestyle = (
+                                self.linestyles[i]
+                                if len(variable_list) == 1
+                                else self.linestyles[j]
+                            )
                             (self.plots[key][i][j],) = ax.plot(
                                 self.first_spatial_variable[key],
                                 variable(t_in_seconds, **spatial_vars),
-                                color=self.colors[i],
+                                color=time_colors[
+                                    _idx
+                                ],  # Different color for each time
                                 linestyle=linestyle,
+                                label=f"t = {t_single:.2f} {self.time_unit}",  # Add time label
                                 zorder=10,
-                                label=f"time = {t_in_seconds / self.time_scaling_factor} {self.time_unit}",
                             )
                             variable_handles.append(self.plots[key][0][j])
-                    solution_handles.append(self.plots[key][i][0])
+                        solution_handles.append(self.plots[key][i][0])
 
-                # Add lines for boundaries between subdomains
-                for boundary in variable_lists[0][0].internal_boundaries:
-                    boundary_scaled = boundary * self.spatial_factor
-                    ax.axvline(boundary_scaled, color="0.5", lw=1, zorder=0)
+                # Add a legend to indicate which plot corresponds to which time
+                ax.legend()
 
             elif variable_lists[0][0].dimensions == 2:
-                # 2D variables do not support plotting for multiple times
-                if len(t_list) > 1:
-                    raise ValueError("Cannot plot 2D variables for multiple times")
-
-                # Continue with 2D plotting as before for a single time
+                # 2D plot: superimpose plots at different times
                 spatial_vars = self.spatial_variable_dict[key]
                 variable = variable_lists[0][0]
-                if self.x_first_and_y_second[key] is False:
-                    x_name = list(spatial_vars.keys())[1][0]
-                    y_name = next(iter(spatial_vars.keys()))[0]
-                    x = self.second_spatial_variable[key]
-                    y = self.first_spatial_variable[key]
-                    var = variable(t_in_seconds_list[0], **spatial_vars)
-                else:
-                    x_name = next(iter(spatial_vars.keys()))[0]
-                    y_name = list(spatial_vars.keys())[1][0]
+
+                for _idx, t_single in enumerate(t):
+                    t_in_seconds = t_single * self.time_scaling_factor
                     x = self.first_spatial_variable[key]
                     y = self.second_spatial_variable[key]
-                    var = variable(t_in_seconds_list[0], **spatial_vars).T
-                ax.set_xlabel(f"{x_name} [{self.spatial_unit}]")
-                ax.set_ylabel(f"{y_name} [{self.spatial_unit}]")
-                vmin, vmax = self.variable_limits[key]
-                # store the plot and the var data (for testing) as cant access
-                # z data from QuadMesh or QuadContourSet object
-                if self.is_y_z[key] is True:
-                    self.plots[key][0][0] = ax.pcolormesh(
-                        x,
-                        y,
-                        var,
-                        vmin=vmin,
-                        vmax=vmax,
-                        shading=self.shading,
+                    var = variable(t_in_seconds, **spatial_vars).T
+
+                    ax.set_xlabel(
+                        f"{next(iter(spatial_vars.keys()))[0]} [{self.spatial_unit}]"
                     )
-                else:
-                    self.plots[key][0][0] = ax.contourf(
-                        x, y, var, levels=100, vmin=vmin, vmax=vmax
+                    ax.set_ylabel(
+                        f"{list(spatial_vars.keys())[1][0]} [{self.spatial_unit}]"
                     )
-                self.plots[key][0][1] = var
-                self.colorbars[key] = self.fig.colorbar(
-                    cm.ScalarMappable(colors.Normalize(vmin=vmin, vmax=vmax)),
-                    ax=ax,
-                )
+                    vmin, vmax = self.variable_limits[key]
 
-            # Set either y label or legend entries
-            if len(key) == 1:
-                title = split_long_string(key[0])
-                ax.set_title(title, fontsize="medium")
-            else:
-                ax.legend(
-                    variable_handles,
-                    [split_long_string(s, 6) for s in key],
-                    bbox_to_anchor=(0.5, 1),
-                    loc="lower center",
-                )
+                    # Use contourf and colorbars to represent the values
+                    contour_plot = ax.contourf(
+                        x, y, var, levels=100, vmin=vmin, vmax=vmax, cmap="coolwarm"
+                    )
+                    self.plots[key][0][0] = contour_plot
+                    self.colorbars[key] = self.fig.colorbar(contour_plot, ax=ax)
 
-        # Set global legend
+                    # Add time to the colorbar title or axis title to indicate which time it represents
+                    ax.set_title(f"t = {t_single:.2f} {self.time_unit}")
+
+        # Set global legend if there are multiple models
         if len(self.labels) > 1:
             fig_legend = self.fig.legend(
                 solution_handles, self.labels, loc="lower right"
             )
-            # Get the position of the top of the legend in relative figure units
-            # There may be a better way ...
-            try:
-                legend_top_inches = fig_legend.get_window_extent(
-                    renderer=self.fig.canvas.get_renderer()
-                ).get_points()[1, 1]
-                fig_height_inches = (self.fig.get_size_inches() * self.fig.dpi)[1]
-                legend_top = legend_top_inches / fig_height_inches
-            except AttributeError:  # pragma: no cover
-                # When testing the examples we set the matplotlib backend to "Template"
-                # which means that the above code doesn't work. Since this is just for
-                # that particular test we can just skip it
-                legend_top = 0
         else:
-            legend_top = 0
+            fig_legend = None
 
-        # Fix layout
+        # Fix layout for sliders if dynamic
         if dynamic:
             slider_top = 0.05
         else:
             slider_top = 0
-        bottom = max(legend_top, slider_top)
+        bottom = max(
+            fig_legend.get_window_extent(
+                renderer=self.fig.canvas.get_renderer()
+            ).get_points()[1, 1]
+            if fig_legend
+            else 0,
+            slider_top,
+        )
         self.gridspec.tight_layout(self.fig, rect=[0, bottom, 1, 1])
 
     def dynamic_plot(self, show_plot=True, step=None):

From 8313c81be584bc3d046f56f8b4b407286697e41d Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Sun, 27 Oct 2024 01:05:04 +0530
Subject: [PATCH 4/8] fixed failing tests

---
 src/pybamm/plotting/quick_plot.py           | 21 ++++++------
 tests/unit/test_plotting/test_quick_plot.py | 36 ++++++++++-----------
 2 files changed, 28 insertions(+), 29 deletions(-)

diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py
index 66a23dd41d..c5bcd759da 100644
--- a/src/pybamm/plotting/quick_plot.py
+++ b/src/pybamm/plotting/quick_plot.py
@@ -313,11 +313,11 @@ def set_output_variables(self, output_variables, solutions):
             first_variable = first_solution[0]
             domain = first_variable.domain
             # check all other solutions against the first solution
-            for _idx, variable in enumerate(first_solution):
+            for idx, variable in enumerate(first_solution):
                 if variable.domain != domain:
                     raise ValueError(
                         "Mismatching variable domains. "
-                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[_idx]}' has domain '{variable.domain}'"
+                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[ idx]}' has domain '{variable.domain}'"
                     )
                 self.spatial_variable_dict[variable_tuple] = {}
 
@@ -494,7 +494,7 @@ def plot(self, t, dynamic=False):
         plt = import_optional_dependency("matplotlib.pyplot")
         gridspec = import_optional_dependency("matplotlib.gridspec")
 
-        if isinstance(t, float):
+        if isinstance(t, (int, float)):
             t = [t]
 
         self.fig = plt.figure(figsize=self.figsize)
@@ -538,7 +538,7 @@ def plot(self, t, dynamic=False):
                         (self.plots[key][i][j],) = ax.plot(
                             full_t / self.time_scaling_factor,
                             variable(full_t),
-                            color=self.colors[i],  # Different colors for plot lines
+                            color=self.colors[i],
                             linestyle=linestyle,
                         )
                         variable_handles.append(self.plots[key][0][j])
@@ -547,7 +547,7 @@ def plot(self, t, dynamic=False):
                 ax.set_ylim(y_min, y_max)
 
                 # Add vertical lines for each time in the list, using different colors for each time
-                for _idx, t_single in enumerate(t):
+                for idx, t_single in enumerate(t):
                     t_in_seconds = t_single * self.time_scaling_factor
                     (self.time_lines[key],) = ax.plot(
                         [
@@ -557,7 +557,7 @@ def plot(self, t, dynamic=False):
                         [y_min, y_max],
                         "--",  # Dashed lines
                         lw=1.5,
-                        color=time_colors[_idx],  # Different color for each time
+                        color=time_colors[idx],  # Different color for each time
                         label=f"t = {t_single:.2f} {self.time_unit}",
                     )
                 ax.legend()
@@ -568,7 +568,7 @@ def plot(self, t, dynamic=False):
                 spatial_var_name = next(iter(spatial_vars.keys()))
                 ax.set_xlabel(f"{spatial_var_name} [{self.spatial_unit}]")
 
-                for _idx, t_single in enumerate(t):
+                for idx, t_single in enumerate(t):
                     t_in_seconds = t_single * self.time_scaling_factor
 
                     for i, variable_list in enumerate(variable_lists):
@@ -581,9 +581,7 @@ def plot(self, t, dynamic=False):
                             (self.plots[key][i][j],) = ax.plot(
                                 self.first_spatial_variable[key],
                                 variable(t_in_seconds, **spatial_vars),
-                                color=time_colors[
-                                    _idx
-                                ],  # Different color for each time
+                                color=time_colors[idx],  # Different color for each time
                                 linestyle=linestyle,
                                 label=f"t = {t_single:.2f} {self.time_unit}",  # Add time label
                                 zorder=10,
@@ -620,7 +618,8 @@ def plot(self, t, dynamic=False):
                     self.plots[key][0][0] = contour_plot
                     self.colorbars[key] = self.fig.colorbar(contour_plot, ax=ax)
 
-                    # Add time to the colorbar title or axis title to indicate which time it represents
+                    self.plots[key][0][1] = var
+
                     ax.set_title(f"t = {t_single:.2f} {self.time_unit}")
 
         # Set global legend if there are multiple models
diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py
index f7de702b29..188a725680 100644
--- a/tests/unit/test_plotting/test_quick_plot.py
+++ b/tests/unit/test_plotting/test_quick_plot.py
@@ -73,7 +73,7 @@ def test_simple_ode_model(self, solver):
                 "c broadcasted positive electrode",
             ],
         )
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         # update the axis
         new_axis = [0, 0.5, 0, 1]
@@ -92,7 +92,7 @@ def test_simple_ode_model(self, solver):
         # Test with different output variables
         quick_plot = pybamm.QuickPlot(solution, ["b broadcasted"])
         assert len(quick_plot.axis_limits) == 1
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         quick_plot = pybamm.QuickPlot(
             solution,
@@ -105,7 +105,7 @@ def test_simple_ode_model(self, solver):
             ],
         )
         assert len(quick_plot.axis_limits) == 5
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         # update the axis
         new_axis = [0, 0.5, 0, 1]
@@ -124,7 +124,7 @@ def test_simple_ode_model(self, solver):
 
         # Test longer name
         quick_plot = pybamm.QuickPlot(solution, ["Variable with a very long name"])
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         # Test different inputs
         quick_plot = pybamm.QuickPlot(
@@ -154,7 +154,7 @@ def test_simple_ode_model(self, solver):
         quick_plot = pybamm.QuickPlot(solution, ["a"])
         assert quick_plot.time_scaling_factor == 1
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="seconds")
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         assert quick_plot.time_scaling_factor == 1
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot
@@ -163,7 +163,7 @@ def test_simple_ode_model(self, solver):
             quick_plot.plots[("a",)][0][0].get_ydata(), 0.2 * t_plot
         )
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="minutes")
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         assert quick_plot.time_scaling_factor == 60
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot / 60
@@ -172,7 +172,7 @@ def test_simple_ode_model(self, solver):
             quick_plot.plots[("a",)][0][0].get_ydata(), 0.2 * t_plot
         )
         quick_plot = pybamm.QuickPlot(solution, ["a"], time_unit="hours")
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         assert quick_plot.time_scaling_factor == 3600
         np.testing.assert_array_almost_equal(
             quick_plot.plots[("a",)][0][0].get_xdata(), t_plot / 3600
@@ -201,7 +201,7 @@ def test_simple_ode_model(self, solver):
 
         # Test 2D variables
         quick_plot = pybamm.QuickPlot(solution, ["2D variable"])
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         quick_plot.dynamic_plot(show_plot=False)
         quick_plot.slider_update(0.01)
 
@@ -217,14 +217,14 @@ def test_simple_ode_model(self, solver):
             None,
             None,
         ]
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
             solution, ["2D variable"], variable_limits="tight"
         )
         assert quick_plot.variable_limits[("2D variable",)] == (None, None)
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
@@ -234,7 +234,7 @@ def test_simple_ode_model(self, solver):
         )
         assert quick_plot.axis_limits[("a",)][2:] == [1, 2]
         assert quick_plot.axis_limits[("c broadcasted", "c broadcasted")][2:] == [3, 4]
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         quick_plot.slider_update(1)
 
         quick_plot = pybamm.QuickPlot(
@@ -242,7 +242,7 @@ def test_simple_ode_model(self, solver):
         )
         assert quick_plot.axis_limits[("a",)][2:] == [None, None]
         assert quick_plot.axis_limits[("b broadcasted",)][2:] != [None, None]
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
         quick_plot.slider_update(1)
 
         with pytest.raises(
@@ -288,12 +288,12 @@ def test_spm_simulation(self):
         # pass only a simulation object
         # it should be converted to a list of corresponding solution
         quick_plot = pybamm.QuickPlot(sim)
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         # mixed simulation and solution input
         # solution should be extracted from the simulation
         quick_plot = pybamm.QuickPlot([sim, sim.solution])
-        quick_plot.plot(0.0)
+        quick_plot.plot(0)
 
         # test creating a GIF
         with TemporaryDirectory() as dir_name:
@@ -336,7 +336,7 @@ def test_loqs_spme(self):
                 quick_plot = pybamm.QuickPlot(
                     solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
                 )
-                quick_plot.plot(0.0)
+                quick_plot.plot(0)
 
                 qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][
                     0
@@ -371,7 +371,7 @@ def test_loqs_spme(self):
                         ["Negative particle concentration [mol.m-3]"],
                         time_unit=unit,
                     )
-                    quick_plot.plot(0.0)
+                    quick_plot.plot(0)
                     qp_data = quick_plot.plots[
                         ("Negative particle concentration [mol.m-3]",)
                     ][0][1]
@@ -411,7 +411,7 @@ def test_plot_1plus1D_spme(self):
             quick_plot = pybamm.QuickPlot(
                 solution, ["Electrolyte concentration [mol.m-3]"], time_unit=unit
             )
-            quick_plot.plot(0.0)
+            quick_plot.plot(0)
             qp_data = quick_plot.plots[("Electrolyte concentration [mol.m-3]",)][0][1]
             c_e_eval = c_e(t_eval[0], x=c_e.first_dim_pts, z=c_e.second_dim_pts)
             np.testing.assert_array_almost_equal(qp_data.T, c_e_eval)
@@ -456,7 +456,7 @@ def test_plot_2plus1D_spm(self):
             quick_plot = pybamm.QuickPlot(
                 solution, ["Negative current collector potential [V]"], time_unit=unit
             )
-            quick_plot.plot(0.0)
+            quick_plot.plot(0)
             qp_data = quick_plot.plots[("Negative current collector potential [V]",)][
                 0
             ][1]

From 650b156b0555bf5204450ef6abad0086a815624a Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Sun, 27 Oct 2024 16:19:08 +0530
Subject: [PATCH 5/8] added tests for using list of times

---
 tests/unit/test_plotting/test_quick_plot.py | 35 +++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py
index 188a725680..092f3784f2 100644
--- a/tests/unit/test_plotting/test_quick_plot.py
+++ b/tests/unit/test_plotting/test_quick_plot.py
@@ -252,6 +252,41 @@ def test_simple_ode_model(self, solver):
                 solution, ["a", "b broadcasted"], variable_limits="bad variable limits"
             )
 
+        # Test with a list of times
+        # Test for a 0D variable
+        quick_plot = pybamm.QuickPlot(solution, ["a"])
+        quick_plot.plot(t=[0, 1, 2])
+        assert len(quick_plot.plots[("a",)]) == 1
+
+        # Test for a 1D variable
+        quick_plot = pybamm.QuickPlot(solution, ["c broadcasted"])
+
+        time_list = [0.5, 1.5, 2.5]
+        quick_plot.plot(time_list)
+
+        ax = quick_plot.fig.axes[0]
+        lines = ax.get_lines()
+
+        variable_key = ("c broadcasted",)
+        if variable_key in quick_plot.variables:
+            num_variables = len(quick_plot.variables[variable_key][0])
+        else:
+            raise KeyError(
+                f"'{variable_key}' is not in quick_plot.variables. Available keys: "
+                + str(quick_plot.variables.keys())
+            )
+
+        expected_lines = len(time_list) * num_variables
+
+        assert (
+            len(lines) == expected_lines
+        ), f"Expected {expected_lines} superimposed lines, but got {len(lines)}"
+
+        # Test for a 2D variable
+        quick_plot = pybamm.QuickPlot(solution, ["2D variable"])
+        quick_plot.plot(t=[0, 1, 2])
+        assert len(quick_plot.plots[("2D variable",)]) == 1
+
         # Test errors
         with pytest.raises(ValueError, match="Mismatching variable domains"):
             pybamm.QuickPlot(solution, [["a", "b broadcasted"]])

From e8b9d26c205c408429d9a056eae8d1cc28198664 Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Thu, 21 Nov 2024 10:34:50 +0530
Subject: [PATCH 6/8] added changelog enrty

---
 CHANGELOG.md                      |  1 +
 src/pybamm/plotting/quick_plot.py | 10 +++++-----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f5ef9bee27..237fce6dcc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 - Added phase-dependent particle options to LAM ([#4369](https://github.com/pybamm-team/PyBaMM/pull/4369))
 - Added a lithium ion equivalent circuit model with split open circuit voltages for each electrode (`SplitOCVR`). ([#4330](https://github.com/pybamm-team/PyBaMM/pull/4330))
 - Added the `pybamm.DiscreteTimeSum` expression node to sum an expression over a sequence of data times, and accompanying `pybamm.DiscreteTimeData` class to store the data times and values ([#4501](https://github.com/pybamm-team/PyBaMM/pull/4501))
+- Modified `quick_plot.plot` to accept a list of times and generate superimposed graphs for specified time points.([#4529](https://github.com/pybamm-team/PyBaMM/pull/4529))
 
 ## Optimizations
 
diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py
index c5bcd759da..50ec0c6605 100644
--- a/src/pybamm/plotting/quick_plot.py
+++ b/src/pybamm/plotting/quick_plot.py
@@ -479,13 +479,13 @@ def reset_axis(self):
             ):  # pragma: no cover
                 raise ValueError(f"Axis limits cannot be NaN for variables '{key}'")
 
-    def plot(self, t, dynamic=False):
+    def plot(self, t: float | list[float], dynamic: bool = False):
         """Produces a quick plot with the internal states at time t.
 
         Parameters
         ----------
-        t : float or list
-            Dimensional time (in 'time_units') at which to plot. Can be a single time or a list of times.
+        t : float or list of float
+        Dimensional time (in 'time_units') at which to plot. Can be a single time or a list of times.
         dynamic : bool, optional
             Determine whether to allocate space for a slider at the bottom of the plot when generating a dynamic plot.
             If True, creates a dynamic plot with a slider.
@@ -494,7 +494,7 @@ def plot(self, t, dynamic=False):
         plt = import_optional_dependency("matplotlib.pyplot")
         gridspec = import_optional_dependency("matplotlib.gridspec")
 
-        if isinstance(t, (int, float)):
+        if not isinstance(t, list):
             t = [t]
 
         self.fig = plt.figure(figsize=self.figsize)
@@ -597,7 +597,7 @@ def plot(self, t, dynamic=False):
                 spatial_vars = self.spatial_variable_dict[key]
                 variable = variable_lists[0][0]
 
-                for _idx, t_single in enumerate(t):
+                for t_single in t:
                     t_in_seconds = t_single * self.time_scaling_factor
                     x = self.first_spatial_variable[key]
                     y = self.second_spatial_variable[key]

From e9e0963b7e0963a65e1b728643c0c7668cfa96ef Mon Sep 17 00:00:00 2001
From: medha-14 <medha140904@gmail.com>
Date: Fri, 22 Nov 2024 18:34:32 +0530
Subject: [PATCH 7/8] made small changes

---
 CHANGELOG.md                      | 5 ++++-
 src/pybamm/plotting/quick_plot.py | 5 +++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 45ca7ea1cc..1ddeaab55a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Features
+
+- Modified `quick_plot.plot` to accept a list of times and generate superimposed graphs for specified time points.([#4529](https://github.com/pybamm-team/PyBaMM/pull/4529))
+
 ## Breaking changes
 
 - Double-layer SEI models have been removed (with the corresponding parameters). All models assume now a single SEI layer. ([#4470](https://github.com/pybamm-team/PyBaMM/pull/4470))
@@ -20,7 +24,6 @@
 - Added phase-dependent particle options to LAM ([#4369](https://github.com/pybamm-team/PyBaMM/pull/4369))
 - Added a lithium ion equivalent circuit model with split open circuit voltages for each electrode (`SplitOCVR`). ([#4330](https://github.com/pybamm-team/PyBaMM/pull/4330))
 - Added the `pybamm.DiscreteTimeSum` expression node to sum an expression over a sequence of data times, and accompanying `pybamm.DiscreteTimeData` class to store the data times and values ([#4501](https://github.com/pybamm-team/PyBaMM/pull/4501))
-- Modified `quick_plot.plot` to accept a list of times and generate superimposed graphs for specified time points.([#4529](https://github.com/pybamm-team/PyBaMM/pull/4529))
 
 ## Optimizations
 
diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py
index 50ec0c6605..e3ad40855e 100644
--- a/src/pybamm/plotting/quick_plot.py
+++ b/src/pybamm/plotting/quick_plot.py
@@ -1,6 +1,7 @@
 #
 # Class for quick plotting of variables from models
 #
+from __future__ import annotations
 import os
 import numpy as np
 import pybamm
@@ -317,7 +318,7 @@ def set_output_variables(self, output_variables, solutions):
                 if variable.domain != domain:
                     raise ValueError(
                         "Mismatching variable domains. "
-                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[ idx]}' has domain '{variable.domain}'"
+                        f"'{variable_tuple[0]}' has domain '{domain}', but '{variable_tuple[idx]}' has domain '{variable.domain}'"
                     )
                 self.spatial_variable_dict[variable_tuple] = {}
 
@@ -485,7 +486,7 @@ def plot(self, t: float | list[float], dynamic: bool = False):
         Parameters
         ----------
         t : float or list of float
-        Dimensional time (in 'time_units') at which to plot. Can be a single time or a list of times.
+            Dimensional time (in 'time_units') at which to plot. Can be a single time or a list of times.
         dynamic : bool, optional
             Determine whether to allocate space for a slider at the bottom of the plot when generating a dynamic plot.
             If True, creates a dynamic plot with a slider.

From 4bea86c52f08c0b7c1897c0fa7f42f583b10e658 Mon Sep 17 00:00:00 2001
From: "Eric G. Kratz" <kratman@users.noreply.github.com>
Date: Fri, 22 Nov 2024 10:32:00 -0500
Subject: [PATCH 8/8] Update CHANGELOG.md

Co-authored-by: Agriya Khetarpal <74401230+agriyakhetarpal@users.noreply.github.com>
---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ddeaab55a..3005a6c16e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 ## Features
 
-- Modified `quick_plot.plot` to accept a list of times and generate superimposed graphs for specified time points.([#4529](https://github.com/pybamm-team/PyBaMM/pull/4529))
+- Modified `quick_plot.plot` to accept a list of times and generate superimposed graphs for specified time points. ([#4529](https://github.com/pybamm-team/PyBaMM/pull/4529))
 
 ## Breaking changes