Skip to content
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
Attention: The newest changes should be on top -->

### Added

- ENH: Add animations for motor propellant mass and tank fluid volume [#656](https://github.com/RocketPy-Team/RocketPy/issues/656)
- ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848)
- ENH: Add persistent caching for ThrustCurve API [#881](https://github.com/RocketPy-Team/RocketPy/pull/881)
- ENH: Compatibility with MERRA-2 atmosphere reanalysis files [#825](https://github.com/RocketPy-Team/RocketPy/pull/825)
Expand Down
67 changes: 66 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,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.
Expand Down
35 changes: 35 additions & 0 deletions rocketpy/plots/plot_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
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.
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.
Expand Down
36 changes: 35 additions & 1 deletion tests/unit/test_plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Loading