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
Comment thread
Gui-FernandesBR marked this conversation as resolved.
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.
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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
Comment thread
Bizo883 marked this conversation as resolved.
Outdated

# 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])
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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.
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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}'."
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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")
Comment thread
Bizo883 marked this conversation as resolved.
Outdated

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.
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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])
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
Comment thread
Bizo883 marked this conversation as resolved.
Outdated
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)
Comment thread
Bizo883 marked this conversation as resolved.
Loading