ENH: add animations for motor propellant mass and tank fluid volumes

Author: Bizo883

rocketpy/plots/motor_plots.py

@@ -1,8 +1,9 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.patches import Polygon
+from matplotlib.animation import FuncAnimation
 
-from ..plots.plot_helpers import show_or_save_plot
+from ..plots.plot_helpers import show_or_save_plot, show_or_save_animation
 
 
 class _MotorPlots:
@@ -520,6 +521,70 @@ def _set_plot_properties(self, ax):
         plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
         plt.tight_layout()
 
+    def animate_propellant_mass(self, filename=None, fps=30):
+        """Animates the propellant mass of the motor as a function of time.
+
+        Parameters
+        ----------
+        filename : str | None, optional
+            The path the animation should be saved to. By default None, in which
+            case the animation will be shown instead of saved.
+        fps : int, optional
+            Frames per second for the animation. Default is 30.
+
+        Returns
+        -------
+        matplotlib.animation.FuncAnimation
+            The created animation object.
+        """
+
+        # Extract time and mass data
+        times = self.motor.propellant_mass.times
+        values = self.motor.propellant_mass.values
+
+        # Create figure and axis
+        fig, ax = plt.subplots()
+
+        # Configure axis
+        ax.set_xlim(times[0], times[-1])
+        ax.set_ylim(min(values), max(values))
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Propellant Mass (kg)")
+        ax.set_title("Propellant Mass Evolution")
+
+        # Create line and current point marker
+        (line,) = ax.plot([], [], lw=2, color="blue", label="Propellant Mass")
+        (point,) = ax.plot([], [], "ko")
+
+        ax.legend()
+
+        # Initialization
+        def init():
+            line.set_data([], [])
+            point.set_data([], [])
+            return line, point
+
+        # Update per frame
+        def update(frame_index):
+            line.set_data(times[:frame_index], values[:frame_index])
+            point.set_data(times[frame_index], values[frame_index])
+            return line, point
+
+        # Build animation
+        animation = FuncAnimation(
+            fig,
+            update,
+            frames=len(times),
+            init_func=init,
+            interval=1000 / fps,
+            blit=True,
+        )
+
+        # Show or save animation
+        show_or_save_animation(animation, filename, fps=fps)
+
+        return animation
+
     def all(self):
         """Prints out all graphs available about the Motor. It simply calls
         all the other plotter methods in this class.

rocketpy/plots/plot_helpers.py

@@ -65,3 +65,38 @@ def show_or_save_fig(fig: Figure, filename=None):
         Path(filename).parent.mkdir(parents=True, exist_ok=True)
 
         fig.savefig(filename, dpi=SAVEFIG_DPI)
+
+
+def show_or_save_animation(animation, filename=None, fps=30):
+    """Shows or saves the given matplotlib animation. If a filename is given,
+    the animation will be saved. Otherwise, it will be shown.
+
+    Parameters
+    ----------
+    animation : matplotlib.animation.FuncAnimation
+        The animation object to be saved or shown.
+    filename : str | None, optional
+        The path the animation should be saved to, by default None. Supported
+        file endings is: gif.
+    fps : int, optional
+        Frames per second when saving the animation. Default is 30.
+    """
+    if filename is None:
+        plt.show()
+    else:
+        file_ending = Path(filename).suffix
+        supported_endings = [".gif"]
+
+        if file_ending not in supported_endings:
+            raise ValueError(
+                f"Unsupported file ending '{file_ending}'."
+                f"Supported file endings are: {supported_endings}."
+            )
+
+        # Before export, ensure the folder the file should go into exists
+        Path(filename).parent.mkdir(parents=True, exist_ok=True)
+
+        if file_ending == ".gif":
+            animation.save(filename, fps=fps, writer="pillow")
+
+        plt.close()

rocketpy/plots/tank_plots.py

@@ -1,10 +1,11 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.patches import Polygon
+from matplotlib.animation import FuncAnimation
 
 from rocketpy.mathutils.function import Function
 
-from .plot_helpers import show_or_save_plot
+from .plot_helpers import show_or_save_plot, show_or_save_animation
 
 
 class _TankPlots:
@@ -180,6 +181,77 @@ def fluid_center_of_mass(self, filename=None):
         ax.legend(["Liquid", "Gas", "Total"])
         show_or_save_plot(filename)
 
+    def animate_fluid_volume(self, filename=None, fps=30):
+        """Animates the liquid and gas volumes inside the tank as a function of time.
+
+        Parameters
+        ----------
+        filename : str | None, optional
+            The path the animation should be saved to. By default None, in which
+            case the animation will be shown instead of saved.
+        fps : int, optional
+            Frames per second for the animation. Default is 30.
+
+        Returns
+        -------
+        matplotlib.animation.FuncAnimation
+            The created animation object.
+        """
+
+        t_start, t_end = self.flux_time
+        times = np.linspace(t_start, t_end, 200)
+
+        liquid_values = self.tank.liquid_volume.get_value(times)
+        gas_values = self.tank.gas_volume.get_value(times)
+
+        fig, ax = plt.subplots()
+
+        ax.set_xlim(times[0], times[-1])
+        max_val = max(liquid_values.max(), gas_values.max())
+        ax.set_ylim(0, max_val * 1.1)
+
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Volume (m³)")
+        ax.set_title("Liquid/Gas Volume Evolution")
+        (line_liquid,) = ax.plot([], [], lw=2, color="blue", label="Liquid Volume")
+        (line_gas,) = ax.plot([], [], lw=2, color="red", label="Gas Volume")
+
+        (point_liquid,) = ax.plot([], [], "ko")
+        (point_gas,) = ax.plot([], [], "ko")
+
+        ax.legend()
+
+        def init():
+            line_liquid.set_data([], [])
+            line_gas.set_data([], [])
+            point_liquid.set_data([], [])
+            point_gas.set_data([], [])
+            return line_liquid, line_gas, point_liquid, point_gas
+
+        def update(frame_index):
+            # Liquid part
+            line_liquid.set_data(times[:frame_index], liquid_values[:frame_index])
+            point_liquid.set_data(times[frame_index], liquid_values[frame_index])
+
+            # Gas part
+            line_gas.set_data(times[:frame_index], gas_values[:frame_index])
+            point_gas.set_data(times[frame_index], gas_values[frame_index])
+
+            return line_liquid, line_gas, point_liquid, point_gas
+
+        animation = FuncAnimation(
+            fig,
+            update,
+            frames=len(times),
+            init_func=init,
+            interval=1000 / fps,
+            blit=True,
+        )
+
+        show_or_save_animation(animation, filename, fps=fps)
+
+        return animation
+
     def all(self):
         """Prints out all graphs available about the Tank. It simply calls
         all the other plotter methods in this class.

tests/unit/test_plots.py

@@ -2,10 +2,11 @@
 from unittest.mock import MagicMock, patch
 
 import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
 import pytest
 
 from rocketpy.plots.compare import Compare
-from rocketpy.plots.plot_helpers import show_or_save_fig, show_or_save_plot
+from rocketpy.plots.plot_helpers import show_or_save_fig, show_or_save_plot, show_or_save_animation
 
 
 @patch("matplotlib.pyplot.show")
@@ -89,3 +90,36 @@ def test_show_or_save_fig(filename):
     else:
         assert os.path.exists(filename)
         os.remove(filename)
+
+
+@pytest.mark.parametrize("filename", [None, "test.gif"])
+@patch("matplotlib.pyplot.show")
+def test_show_or_save_animation(mock_show, filename):
+    """This test is to check if the show_or_save_animation function is
+    working properly.
+
+    Parameters
+    ----------
+    mock_show :
+        Mocks the matplotlib.pyplot.show() function to avoid showing the animation.
+    filename : str
+        Name of the file to save the animation. If None, the animation will be
+        shown instead.
+    """
+
+    # Create a simple animation object
+    fig, ax = plt.subplots()
+
+    def update(frame):
+        ax.plot([0, frame], [0, frame])
+        return ax
+
+    animation = FuncAnimation(fig, update, frames=5)
+
+    show_or_save_animation(animation, filename)
+
+    if filename is None:
+        mock_show.assert_called_once()
+    else:
+        assert os.path.exists(filename)
+        os.remove(filename)