diff --git a/CHANGELOG.md b/CHANGELOG.md index 17b654414..b5a0edf74 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ Attention: The newest changes should be on top --> ### Added +- ENH: add animations for motor propellant mass and tank fluid volumes [#894](https://github.com/RocketPy-Team/RocketPy/pull/894) - ENH: Add axial_acceleration attribute to the Flight class [#876](https://github.com/RocketPy-Team/RocketPy/pull/876) - ENH: Rail button bending moments calculation in Flight class [#893](https://github.com/RocketPy-Team/RocketPy/pull/893) - ENH: Built-in flight comparison tool (`FlightComparator`) to validate simulations against external data [#888](https://github.com/RocketPy-Team/RocketPy/pull/888) diff --git a/docs/user/motors/liquidmotor.rst b/docs/user/motors/liquidmotor.rst index 401ab62c6..6a9a22635 100644 --- a/docs/user/motors/liquidmotor.rst +++ b/docs/user/motors/liquidmotor.rst @@ -160,6 +160,21 @@ For example: example_liquid.exhaust_velocity.plot(0, 5) +The tanks added to a ``LiquidMotor`` can now be animated to visualize +how the liquid and gas volumes evolve during the burn. + +To animate the tanks, we can use the ``animate_fluid_volume()`` method: + +.. jupyter-execute:: + + example_liquid.animate_fluid_volume(fps=30) + +Optionally, the animation can be saved to a GIF file: + +.. jupyter-execute:: + + example_liquid.animate_fluid_volume(fps=30, save_as="liquid_motor.gif") + Alternatively, you can plot all the information at once: .. jupyter-execute:: diff --git a/docs/user/motors/tanks.rst b/docs/user/motors/tanks.rst index f1a3342b8..f57383476 100644 --- a/docs/user/motors/tanks.rst +++ b/docs/user/motors/tanks.rst @@ -263,6 +263,21 @@ We can see some outputs with: # Evolution of the Propellant center of mass position N2O_mass_tank.center_of_mass.plot() +All tank types now include a built-in method for animating the evolution +of liquid and gas volumes over time. This visualization aids in understanding the dynamic behavior +of the tank's contents. To animate the tanks, we can use the +``animate_fluid_volume()`` method: + +.. jupyter-execute:: + + N2O_mass_tank.animate_fluid_volume(fps=30) + +Optionally, the animation can be saved to a GIF file: + +.. jupyter-execute:: + + N2O_mass_tank.animate_fluid_volume(fps=30, save_as="mass_based_tank.gif") + Ullage Based Tank ----------------- diff --git a/rocketpy/plots/motor_plots.py b/rocketpy/plots/motor_plots.py index 8081371a4..7e8a072da 100644 --- a/rocketpy/plots/motor_plots.py +++ b/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,71 @@ 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.Supported file + ending is: .gif + 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.x_array + values = self.motor.propellant_mass.y_array + + # 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 + 1], values[: frame_index + 1]) + 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. diff --git a/rocketpy/plots/plot_helpers.py b/rocketpy/plots/plot_helpers.py index f795d46e9..629ebd22c 100644 --- a/rocketpy/plots/plot_helpers.py +++ b/rocketpy/plots/plot_helpers.py @@ -65,3 +65,37 @@ 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 ending 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) + + animation.save(filename, fps=fps, writer="pillow") + + plt.close() diff --git a/rocketpy/plots/tank_plots.py b/rocketpy/plots/tank_plots.py index 9c6ebb78a..68df643e7 100644 --- a/rocketpy/plots/tank_plots.py +++ b/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. Supported file + ending is: .gif + 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]) + ax.set_ylim(0, max(liquid_values.max(), gas_values.max()) * 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(): + for item in (line_liquid, line_gas, point_liquid, point_gas): + item.set_data([], []) + return line_liquid, line_gas, point_liquid, point_gas + + def update(frame_index): + # Liquid part + line_liquid.set_data( + times[: frame_index + 1], liquid_values[: frame_index + 1] + ) + point_liquid.set_data([times[frame_index]], [liquid_values[frame_index]]) + + # Gas part + line_gas.set_data(times[: frame_index + 1], gas_values[: frame_index + 1]) + 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. diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index 6d7c0cce2..3883c9459 100644 --- a/tests/unit/test_plots.py +++ b/tests/unit/test_plots.py @@ -2,10 +2,15 @@ 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 +94,80 @@ 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) + + +def test_show_or_save_animation_unsupported_format(): + # Test that show_or_save_animation raises ValueError for unsupported formats. + fig, ax = plt.subplots() + + def update(frame): + ax.plot([0, frame], [0, frame]) + return ax + + animation = FuncAnimation(fig, update, frames=5) + + with pytest.raises(ValueError, match="Unsupported file ending"): + show_or_save_animation(animation, "test.mp4") + + +def test_animate_propellant_mass(cesaroni_m1670): + """Test that animate_propellant_mass saves a .gif file correctly.""" + + motor = cesaroni_m1670 + animation = motor.plots.animate_propellant_mass(filename="cesaroni_m1670.gif") + + # Check animation type + assert isinstance(animation, FuncAnimation) + + # check if file exists + assert os.path.exists("cesaroni_m1670.gif") + + os.remove("cesaroni_m1670.gif") + + +def test_animate_fluid_volume(example_mass_flow_rate_based_tank_seblm): + """Test that animate_fluid_volume saves a .gif file correctly.""" + + tank = example_mass_flow_rate_based_tank_seblm + animation = tank.plots.animate_fluid_volume(filename="test_fluid_volume.gif") + + # Check animation type + assert isinstance(animation, FuncAnimation) + + # Check if file exists + assert os.path.exists("test_fluid_volume.gif") + + os.remove("test_fluid_volume.gif")