Enhance plot_voltage_components.py for Simulation and Solution classes

pybamm/plotting/plot_voltage_components.py

@@ -4,10 +4,12 @@
 import numpy as np
 
 from pybamm.util import have_optional_dependency
+from pybamm.simulation import Simulation
+from pybamm.solvers.solution import Solution
 
 
 def plot_voltage_components(
-    solution,
+    input_data,
     ax=None,
     show_legend=True,
     split_by_electrode=False,
@@ -19,8 +21,8 @@ def plot_voltage_components(
 
     Parameters
     ----------
-    solution : :class:`pybamm.Solution`
-        Solution object from which to extract voltage components
+    input_data : :class:`pybamm.Solution` or :class:`pybamm.Simulation`
+        Solution or Simulation object from which to extract voltage components.
     ax : matplotlib Axis, optional
         The axis on which to put the plot. If None, a new figure and axis is created.
     show_legend : bool, optional
@@ -34,6 +36,13 @@ def plot_voltage_components(
         Keyword arguments, passed to ax.fill_between
 
     """
+    # Check if the input is a Simulation and extract Solution
+    if isinstance(input_data, Simulation):
+        solution = input_data.solution
+    elif isinstance(input_data, Solution):
+        solution = input_data
+    else:
+        raise ValueError("Input must be a Solution or Simulation object")
     plt = have_optional_dependency("matplotlib.pyplot")
 
     # Set a default value for alpha, the opacity

pybamm/simulation.py

@@ -1194,6 +1194,44 @@ def save_model(
                 """
             )
 
+    def plot_voltage_components(
+        self,
+        ax=None,
+        show_legend=True,
+        split_by_electrode=False,
+        testing=False,
+        **kwargs_fill,
+    ):
+        """
+        Generate a plot showing the component overpotentials that make up the voltage
+
+        Parameters
+        ----------
+        ax : matplotlib Axis, optional
+            The axis on which to put the plot. If None, a new figure and axis is created.
+        show_legend : bool, optional
+            Whether to display the legend. Default is True.
+        split_by_electrode : bool, optional
+            Whether to show the overpotentials for the negative and positive electrodes
+            separately. Default is False.
+        testing : bool, optional
+            Whether to actually make the plot (turned off for unit tests).
+        kwargs_fill
+            Keyword arguments, passed to ax.fill_between.
+
+        """
+        if self.solution is None:
+            raise ValueError("The simulation has not been solved yet.")
+
+        return pybamm.plot_voltage_components(
+            self.solution,
+            ax=ax,
+            show_legend=show_legend,
+            split_by_electrode=split_by_electrode,
+            testing=testing,
+            **kwargs_fill,
+        )
+
 
 def load_sim(filename):
     """Load a saved simulation"""

pybamm/solvers/solution.py

@@ -801,6 +801,42 @@ def copy(self):
 
         return new_sol
 
+    def plot_voltage_components(
+        self,
+        ax=None,
+        show_legend=True,
+        split_by_electrode=False,
+        testing=False,
+        **kwargs_fill,
+    ):
+        """
+        Generate a plot showing the component overpotentials that make up the voltage
+
+        Parameters
+        ----------
+        ax : matplotlib Axis, optional
+            The axis on which to put the plot. If None, a new figure and axis is created.
+        show_legend : bool, optional
+            Whether to display the legend. Default is True.
+        split_by_electrode : bool, optional
+            Whether to show the overpotentials for the negative and positive electrodes
+            separately. Default is False.
+        testing : bool, optional
+            Whether to actually make the plot (turned off for unit tests).
+        kwargs_fill
+            Keyword arguments, passed to ax.fill_between.
+
+        """
+        # Use 'self' here as the solution object
+        return pybamm.plot_voltage_components(
+            self,
+            ax=ax,
+            show_legend=show_legend,
+            split_by_electrode=split_by_electrode,
+            testing=testing,
+            **kwargs_fill,
+        )
+
 
 class EmptySolution:
     def __init__(self, termination=None, t=None):

tests/unit/test_plotting/test_plot_voltage_components.py

@@ -6,20 +6,34 @@
 
 
 class TestPlotVoltageComponents(TestCase):
-    def test_plot(self):
+    def test_plot_with_solution(self):
         model = pybamm.lithium_ion.SPM()
         sim = pybamm.Simulation(model)
         sol = sim.solve([0, 3600])
         for split in [True, False]:
-            _, ax = pybamm.plot_voltage_components(
-                sol, testing=True, split_by_electrode=split
-            )
+            _, ax = sol.plot_voltage_components(testing=True, split_by_electrode=split)
             t, V = ax.get_lines()[0].get_data()
             np.testing.assert_array_equal(t, sol["Time [h]"].data)
             np.testing.assert_array_equal(V, sol["Battery voltage [V]"].data)
 
         _, ax = plt.subplots()
-        _, ax_out = pybamm.plot_voltage_components(sol, ax=ax, show_legend=True)
+        _, ax_out = sol.plot_voltage_components(ax=ax, show_legend=True)
+        self.assertEqual(ax_out, ax)
+
+    def test_plot_with_simulation(self):
+        model = pybamm.lithium_ion.SPM()
+        sim = pybamm.Simulation(model)
+        sim.solve([0, 3600])
+
+        for split in [True, False]:
+            # Use the method from the Simulation instance
+            _, ax = sim.plot_voltage_components(testing=True, split_by_electrode=split)
+            t, V = ax.get_lines()[0].get_data()
+            np.testing.assert_array_equal(t, sim.solution["Time [h]"].data)
+            np.testing.assert_array_equal(V, sim.solution["Battery voltage [V]"].data)
+
+        _, ax = plt.subplots()
+        _, ax_out = sim.plot_voltage_components(ax=ax, show_legend=True)
         self.assertEqual(ax_out, ax)