Skip to content
Merged
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
15 changes: 15 additions & 0 deletions docs/user/motors/liquidmotor.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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::
Expand Down
15 changes: 15 additions & 0 deletions docs/user/motors/tanks.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
-----------------
Expand Down
68 changes: 67 additions & 1 deletion rocketpy/plots/motor_plots.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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.
Expand Down
34 changes: 34 additions & 0 deletions rocketpy/plots/plot_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
74 changes: 73 additions & 1 deletion rocketpy/plots/tank_plots.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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.
Expand Down
84 changes: 83 additions & 1 deletion tests/unit/test_plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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")