Implement discharge capacity as an optional x-axis in QuickPlot

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- Added functionality to set `Discharge capacity` as the x-axis in `QuickPlot`.([#4775](https://github.com/pybamm-team/PyBaMM/pull/4775))
 - Revision of the hysteresis notebook to include the method implemented in the module `axen_ocp`. ([#4880](https://github.com/pybamm-team/PyBaMM/pull/4880))
 - Added `axen_ocp` module within submodel `interface.open_circuit_potential` to handle an OCP with hysteresis. ([#4816](https://github.com/pybamm-team/PyBaMM/pull/4816))
 - Creates a 'calc_esoh' property in battery models ([#4825](https://github.com/pybamm-team/PyBaMM/pull/4825))

src/pybamm/plotting/quick_plot.py

@@ -91,6 +91,10 @@ class QuickPlot:
         - "tight": make axes tight to plot at each time
         - dictionary: fine-grain control for each variable, can be either "fixed" or \
         "tight" or a specific tuple (lower, upper).
+    x_axis : str, optional
+        The variable to use for the x-axis. Options are:
+        - "Time" (default): Use time as the x-axis.
+        - "Discharge capacity [A.h]": Use discharge capacity as the x-axis.
 
     """
 
@@ -108,6 +112,7 @@ def __init__(
         spatial_unit="um",
         variable_limits="fixed",
         n_t_linear=100,
+        x_axis="Time",
     ):
         solutions = self.preprocess_solutions(solutions)
 
@@ -214,6 +219,27 @@ def t_sample(sol):
         self.min_t_unscaled = min_t
         self.max_t_unscaled = max_t
 
+        # set x_axis
+        self.x_axis = x_axis
+
+        if x_axis == "Discharge capacity [A.h]":
+            # Use discharge capacity as x-axis
+            discharge_capacities = [
+                solution["Discharge capacity [A.h]"].entries for solution in solutions
+            ]
+            self.x_values = discharge_capacities
+
+            self.x_scaling_factor = 1
+            self.x_label = "Discharge capacity [A.h]"
+
+        elif x_axis == "Time":
+            self.x_values = ts_seconds
+            self.x_scaling_factor = self.time_scaling_factor
+
+        else:
+            msg = "Invalid value for `x_axis`."
+            raise ValueError(msg)
+
         # Prepare dictionary of variables
         # output_variables is a list of strings or lists, e.g.
         # ["var 1", ["variable 2", "var 3"]]
@@ -436,7 +462,7 @@ def reset_axis(self):
 
             # Get min and max variable values
             if self.variable_limits[key] == "fixed":
-                # fixed variable limits: calculate "globlal" min and max
+                # fixed variable limits: calculate "global" min and max
                 spatial_vars = self.spatial_variable_dict[key]
                 var_min = np.min(
                     [
@@ -521,7 +547,11 @@ def plot(self, t, dynamic=False):
             # 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
-                ax.set_xlabel(f"Time [{self.time_unit}]")
+                if self.x_axis == "Time":
+                    ax.set_xlabel(f"Time [{self.time_unit}]")
+                if self.x_axis == "Discharge capacity [A.h]":
+                    ax.set_xlabel(f"{self.x_label}")
+
                 for i, variable_list in enumerate(variable_lists):
                     for j, variable in enumerate(variable_list):
                         if len(variable_list) == 1:
@@ -531,10 +561,10 @@ def plot(self, t, dynamic=False):
                             # multiple variables -> use linestyle to differentiate
                             # variables (color differentiates models)
                             linestyle = self.linestyles[j]
-                        full_t = self.ts_seconds[i]
+                        full_val = self.x_values[i]
                         (self.plots[key][i][j],) = ax.plot(
-                            full_t / self.time_scaling_factor,
-                            variable(full_t),
+                            full_val / self.x_scaling_factor,
+                            variable(full_val),
                             color=self.colors[i],
                             linestyle=linestyle,
                         )
@@ -665,13 +695,13 @@ def plot(self, t, dynamic=False):
 
     def dynamic_plot(self, show_plot=True, step=None):
         """
-        Generate a dynamic plot with a slider to control the time.
+        Generate a dynamic plot with a slider to control the x-axis.
 
         Parameters
         ----------
         step : float, optional
             For notebook mode, size of steps to allow in the slider. Defaults to 1/100th
-            of the total time.
+            of the total range (time or discharge capacity).
         show_plot : bool, optional
             Whether to show the plots. Default is True. Set to False if you want to
             only display the plot after plt.show() has been called.
@@ -692,17 +722,25 @@ def dynamic_plot(self, show_plot=True, step=None):
             plt = import_optional_dependency("matplotlib.pyplot")
             Slider = import_optional_dependency("matplotlib.widgets", "Slider")
 
-            # create an initial plot at time self.min_t
+            # Set initial x-axis values and slider
             self.plot(self.min_t, dynamic=True)
 
+            # Set x-axis label correctly
+            if self.x_axis == "Time":
+                ax_label = f"Time [{self.time_unit}]"
+            elif self.x_axis == "Discharge capacity [A.h]":
+                ax_label = "Discharge capacity [A.h]"
+
+            ax_min, ax_max, val_init = self.min_t, self.max_t, self.min_t
+
             axcolor = "lightgoldenrodyellow"
             ax_slider = plt.axes([0.315, 0.02, 0.37, 0.03], facecolor=axcolor)
             self.slider = Slider(
                 ax_slider,
-                f"Time [{self.time_unit}]",
-                self.min_t,
-                self.max_t,
-                valinit=self.min_t,
+                ax_label,
+                ax_min,
+                ax_max,
+                valinit=val_init,
                 color="#1f77b4",
             )
             self.slider.on_changed(self.slider_update)

tests/unit/test_plotting/test_quick_plot.py

@@ -281,6 +281,47 @@ def test_simple_ode_model(self, solver):
 
         pybamm.close_plots()
 
+    def test_invalid_x_axis(self):
+        model = pybamm.lithium_ion.SPM()
+        sim = pybamm.Simulation(model)
+        solution = sim.solve([0, 3600])
+
+        with pytest.raises(ValueError, match="Invalid value for `x_axis`."):
+            pybamm.QuickPlot([solution], x_axis="Invalid axis")
+
+    def test_plot_with_discharge_capacity(self):
+        model = pybamm.lithium_ion.BaseModel(name="Simple ODE Model")
+        a = pybamm.Variable("a", domain=[])
+        model.rhs = {a: pybamm.Scalar(0.2)}
+        model.initial_conditions = {a: pybamm.Scalar(0)}
+        model.variables = {"a": a, "Discharge capacity [A.h]": a * 2}
+
+        t_eval = np.linspace(0, 2, 100)
+        solution = pybamm.CasadiSolver().solve(model, t_eval)
+
+        quick_plot = pybamm.QuickPlot(
+            solution,
+            ["a"],
+            x_axis="Discharge capacity [A.h]",
+        )
+        quick_plot.plot(0)
+
+        # Test discharge capacity values
+        np.testing.assert_allclose(
+            quick_plot.plots[("a",)][0][0].get_xdata(),
+            solution["Discharge capacity [A.h]"].data,
+        )
+
+        # Test x-axis label
+        x_label = quick_plot.fig.axes[0].get_xlabel()
+        assert x_label == "Discharge capacity [A.h]", (
+            f"Expected 'Discharge capacity [A.h]', got '{x_label}'"
+        )
+
+        # check dynamic plot loads
+        quick_plot.dynamic_plot(show_plot=False)
+        quick_plot.slider_update(0.01)
+
     def test_plot_with_different_models(self):
         model = pybamm.BaseModel()
         a = pybamm.Variable("a")