diff --git a/CHANGELOG.md b/CHANGELOG.md index e46ee3faa..ae43d86d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ Attention: The newest changes should be on top --> ### Added +- ENH: Add 3D flight trajectory and attitude animations in Flight plots layer [#909](https://github.com/RocketPy-Team/RocketPy/pull/909) - ENH: Air brakes controller functions now support 8-parameter signature [#854](https://github.com/RocketPy-Team/RocketPy/pull/854) - TST: Add acceptance tests for 3DOF flight simulation based on Bella Lui rocket [#914] (https://github.com/RocketPy-Team/RocketPy/pull/914_ - ENH: Add background map auto download functionality to Monte Carlo plots [#896](https://github.com/RocketPy-Team/RocketPy/pull/896) diff --git a/docs/user/flight.rst b/docs/user/flight.rst index 31e7ab588..b005477db 100644 --- a/docs/user/flight.rst +++ b/docs/user/flight.rst @@ -274,7 +274,7 @@ During the rail launch phase, RocketPy calculates reaction forces and internal b **Rail Button Forces (N):** - ``rail_button1_normal_force`` : Normal reaction force at upper rail button -- ``rail_button1_shear_force`` : Shear (tangential) reaction force at upper rail button +- ``rail_button1_shear_force`` : Shear (tangential) reaction force at upper rail button - ``rail_button2_normal_force`` : Normal reaction force at lower rail button - ``rail_button2_shear_force`` : Shear (tangential) reaction force at lower rail button @@ -282,7 +282,7 @@ During the rail launch phase, RocketPy calculates reaction forces and internal b - ``rail_button1_bending_moment`` : Time-dependent bending moment at upper rail button attachment - ``max_rail_button1_bending_moment`` : Maximum absolute bending moment at upper rail button -- ``rail_button2_bending_moment`` : Time-dependent bending moment at lower rail button attachment +- ``rail_button2_bending_moment`` : Time-dependent bending moment at lower rail button attachment - ``max_rail_button2_bending_moment`` : Maximum absolute bending moment at lower rail button **Calculation Method:** @@ -454,6 +454,51 @@ Flight Data Plots # Flight path and orientation flight.plots.flight_path_angle_data() +3D Flight Animation +~~~~~~~~~~~~~~~~~~~ + +RocketPy can animate the simulated flight trajectory and attitude through the +Flight plots layer. + +.. note:: + + Install optional animation dependencies first: + + .. code-block:: bash + + pip install rocketpy[animation] + +.. code-block:: python + + # Fast start using RocketPy's built-in default STL model + flight.plots.animate_trajectory( + start=0.0, + stop=min(flight.t_final, 20.0), + time_step=0.05, + ) + + # Or provide your own STL model file + flight.plots.animate_trajectory( + file_name="rocket.stl", + start=0.0, + stop=flight.t_final, + time_step=0.05, + azimuth=45, + elevation=20, + ) + + # Keep rocket centered and animate only attitude changes + flight.plots.animate_rotate( + file_name="rocket.stl", + start=0.0, + stop=min(flight.t_final, 20.0), + time_step=0.05, + ) + +Both methods validate the selected time interval and STL path before rendering. +If ``vedo`` is not installed, RocketPy raises an informative ``ImportError`` +with installation instructions. + Forces and Moments ~~~~~~~~~~~~~~~~~~ diff --git a/pyproject.toml b/pyproject.toml index 35ea34382..01b5d502b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,9 @@ build-backend = "setuptools.build_meta" [tool.setuptools] packages = { find = { where = ["."], include = ["rocketpy*"] } } +[tool.setuptools.package-data] +"rocketpy.plots" = ["assets/*.stl"] + [tool.setuptools.dynamic] dependencies = { file = ["requirements.txt"] } @@ -61,14 +64,18 @@ env-analysis = [ ] monte-carlo = [ - "imageio", + "imageio", "multiprocess>=0.70", "statsmodels", "prettytable", "contextily>=1.0.0; python_version < '3.14'", ] -all = ["rocketpy[env-analysis]", "rocketpy[monte-carlo]"] +animation = [ + "vedo>=2024.5.1" +] + +all = ["rocketpy[env-analysis]", "rocketpy[monte-carlo]", "rocketpy[animation]"] [tool.coverage.report] diff --git a/requirements-optional.txt b/requirements-optional.txt index 58ed1030b..2961946ca 100644 --- a/requirements-optional.txt +++ b/requirements-optional.txt @@ -6,4 +6,5 @@ timezonefinder imageio multiprocess>=0.70 statsmodels -prettytable \ No newline at end of file +prettytable +vedo>=2024.5.1 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 61a594320..e08c9a81d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,7 @@ numpy>=1.13 scipy>=1.0 matplotlib>=3.9.0 # Released May 15th 2024 netCDF4>=1.6.4 -requests -pytz +requests>=2.25.0 +pytz>=2020.1 simplekml dill diff --git a/rocketpy/plots/assets/default_rocket.stl b/rocketpy/plots/assets/default_rocket.stl new file mode 100644 index 000000000..e3889fe36 --- /dev/null +++ b/rocketpy/plots/assets/default_rocket.stl @@ -0,0 +1,16 @@ +solid default_rocket + facet normal 0 0 1 + outer loop + vertex 0 0 0 + vertex 1 0 0 + vertex 0 1 0 + endloop + endfacet + facet normal 0 0 -1 + outer loop + vertex 0 0 0 + vertex 0 1 0 + vertex 1 0 0 + endloop + endfacet +endsolid default_rocket diff --git a/rocketpy/plots/flight_plots.py b/rocketpy/plots/flight_plots.py index 7eb41a8b2..86118037f 100644 --- a/rocketpy/plots/flight_plots.py +++ b/rocketpy/plots/flight_plots.py @@ -1,8 +1,12 @@ +import os +import time from functools import cached_property +from importlib import resources import matplotlib.pyplot as plt import numpy as np +from ..tools import import_optional_dependency from .plot_helpers import show_or_save_plot @@ -133,6 +137,256 @@ def trajectory_3d(self, *, filename=None): # pylint: disable=too-many-statement ax1.set_box_aspect(None, zoom=0.95) # 95% for label adjustment show_or_save_plot(filename) + def _resolve_animation_model_path(self, file_name): + """Resolve model path, defaulting to the built-in STL when omitted.""" + if file_name is not None: + return file_name + + return str( + resources.files("rocketpy.plots").joinpath("assets/default_rocket.stl") + ) + + def _validate_animation_inputs(self, file_name, start, stop, time_step): + """Validate shared input parameters for 3D animation methods.""" + if time_step <= 0: + raise ValueError( + f"Invalid time_step: {time_step}. It must be greater than 0." + ) + + if stop is None: + stop = self.flight.t_final + + if ( + start < 0 + or stop < 0 + or start > self.flight.t_final + or stop > self.flight.t_final + or start >= stop + ): + raise ValueError( + f"Invalid animation time range: start={start}, stop={stop}. " + f"Both must be within [0, {self.flight.t_final}] and start < stop." + ) + + if not os.path.isfile(file_name): + raise FileNotFoundError( + f"Could not find the 3D model file: '{file_name}'. " + "Provide a valid .stl file path." + ) + + return stop + + @staticmethod + def _rotation_from_quaternion(q0, q1, q2, q3): + """Convert unit quaternion to axis-angle representation in degrees.""" + norm = np.sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3) + if norm == 0: + return 0.0, (1.0, 0.0, 0.0) + + q0 = q0 / norm + q1 = q1 / norm + q2 = q2 / norm + q3 = q3 / norm + + # q and -q represent the same orientation. Keep q0 non-negative to + # reduce discontinuities in axis-angle interpolation across frames. + if q0 < 0: + q0 = -q0 + q1 = -q1 + q2 = -q2 + q3 = -q3 + + q0 = np.clip(q0, -1.0, 1.0) + angle = 2 * np.arccos(q0) + sin_half = np.sqrt(max(1 - q0 * q0, 0.0)) + + if sin_half < 1e-12: + return 0.0, (1.0, 0.0, 0.0) + + axis = (q1 / sin_half, q2 / sin_half, q3 / sin_half) + return np.degrees(angle), axis + + def _create_animation_box(self, start, scale=1.0): + """Create a world box with minimum visible dimensions.""" + min_box_dim = 10.0 + x_values = self.flight.x[:, 1] + y_values = self.flight.y[:, 1] + z_values = self.flight.z[:, 1] - self.flight.env.elevation + + center_x = 0.5 * (np.max(x_values) + np.min(x_values)) + center_y = 0.5 * (np.max(y_values) + np.min(y_values)) + center_z = max(self.flight.z(start) - self.flight.env.elevation, 0.0) + + length = max(np.ptp(x_values) * scale, min_box_dim) + width = max(np.ptp(y_values) * scale, min_box_dim) + height = max(np.ptp(z_values) * scale, min_box_dim) + + # Keep z center inside visible space while preserving minimum box size. + center_z = max(center_z, 0.5 * min_box_dim) + + vedo = import_optional_dependency("vedo") + Box = vedo.Box + + return Box( + pos=[center_x, center_y, center_z], + length=length, + width=width, + height=height, + ).wireframe() + + def animate_trajectory( + self, file_name=None, start=0, stop=None, time_step=0.1, **kwargs + ): + """Animate 6-DOF trajectory and attitude using vedo. + + Parameters + ---------- + file_name : str | None, optional + Path to a 3D model file representing the rocket, usually ``.stl``. + If None, RocketPy uses a built-in default STL model. + Default is None. + start : int, float, optional + Animation start time in seconds. Default is 0. + stop : int, float | None, optional + Animation end time in seconds. If None, uses ``flight.t_final``. + Default is None. + time_step : float, optional + Animation frame step in seconds. Must be greater than 0. + Default is 0.1. + **kwargs : dict, optional + Additional keyword arguments passed to ``vedo.Plotter.show``. + """ + + vedo = import_optional_dependency("vedo") + + Line = vedo.Line + Mesh = vedo.Mesh + Plotter = vedo.Plotter + settings = vedo.settings + + file_name = self._resolve_animation_model_path(file_name) + stop = self._validate_animation_inputs(file_name, start, stop, time_step) + + try: + settings.allow_interaction = True + except AttributeError: + pass + + world = self._create_animation_box(start, scale=1.2) + base_rocket = Mesh(file_name).c("green") + time_steps = np.arange(start, stop, time_step) + trajectory_points = [] + + plt = Plotter(axes=1, interactive=False) + plt.show(world, "Rocket Trajectory Animation", viewup="z", **kwargs) + + for t in time_steps: + rocket = base_rocket.clone() + x_position = self.flight.x(t) + y_position = self.flight.y(t) + z_position = self.flight.z(t) - self.flight.env.elevation + + angle_deg, axis = self._rotation_from_quaternion( + self.flight.e0(t), + self.flight.e1(t), + self.flight.e2(t), + self.flight.e3(t), + ) + + rocket.pos(x_position, y_position, z_position) + if angle_deg != 0.0: + rocket.rotate(angle_deg, axis=axis) + + trajectory_points.append([x_position, y_position, z_position]) + actors = [world, rocket] + if len(trajectory_points) > 1: + actors.append(Line(trajectory_points, c="k", alpha=0.5)) + + plt.show(*actors, resetcam=False) + + start_pause = time.time() + while time.time() - start_pause < time_step: + plt.render() + + if getattr(plt, "escaped", False): + break + + plt.interactive().close() + return None + + def animate_rotate( + self, file_name=None, start=0, stop=None, time_step=0.1, **kwargs + ): + """Animate rocket attitude (rotation-only view) using vedo. + + Parameters + ---------- + file_name : str | None, optional + Path to a 3D model file representing the rocket, usually ``.stl``. + If None, RocketPy uses a built-in default STL model. + Default is None. + start : int, float, optional + Animation start time in seconds. Default is 0. + stop : int, float | None, optional + Animation end time in seconds. If None, uses ``flight.t_final``. + Default is None. + time_step : float, optional + Animation frame step in seconds. Must be greater than 0. + Default is 0.1. + **kwargs : dict, optional + Additional keyword arguments passed to ``vedo.Plotter.show``. + """ + + vedo = import_optional_dependency("vedo") + + Mesh = vedo.Mesh + Plotter = vedo.Plotter + settings = vedo.settings + + file_name = self._resolve_animation_model_path(file_name) + stop = self._validate_animation_inputs(file_name, start, stop, time_step) + + try: + settings.allow_interaction = True + except AttributeError: + pass + + world = self._create_animation_box(start, scale=0.3) + base_rocket = Mesh(file_name).c("green") + time_steps = np.arange(start, stop, time_step) + + x_start = self.flight.x(start) + y_start = self.flight.y(start) + z_start = self.flight.z(start) - self.flight.env.elevation + + plt = Plotter(axes=1, interactive=False) + plt.show(world, "Rocket Rotation Animation", viewup="z", **kwargs) + + for t in time_steps: + rocket = base_rocket.clone() + angle_deg, axis = self._rotation_from_quaternion( + self.flight.e0(t), + self.flight.e1(t), + self.flight.e2(t), + self.flight.e3(t), + ) + + rocket.pos(x_start, y_start, z_start) + if angle_deg != 0.0: + rocket.rotate(angle_deg, axis=axis) + + plt.show(world, rocket, resetcam=False) + + start_pause = time.time() + while time.time() - start_pause < time_step: + plt.render() + + if getattr(plt, "escaped", False): + break + + plt.interactive().close() + return None + def linear_kinematics_data(self, *, filename=None): # pylint: disable=too-many-statements """Prints out all Kinematics graphs available about the Flight diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index fc412d3b9..f4a218995 100644 --- a/tests/unit/test_plots.py +++ b/tests/unit/test_plots.py @@ -1,4 +1,7 @@ +import builtins import os +import sys +import types from unittest.mock import MagicMock, patch import matplotlib.pyplot as plt @@ -171,3 +174,200 @@ def test_animate_fluid_volume(example_mass_flow_rate_based_tank_seblm): assert os.path.exists("test_fluid_volume.gif") os.remove("test_fluid_volume.gif") + + +class _DummyVedoActor: + """Minimal actor mock that supports the methods used by animation plots.""" + + def __init__(self): + self.rotations = [] + + def c(self, *_args, **_kwargs): + return self + + def pos(self, *_args, **_kwargs): + return self + + def wireframe(self): + return self + + def rotate(self, angle, axis=None): + self.rotations.append((angle, axis)) + return self + + def clone(self): + return _DummyVedoActor() + + +class _DummyPlotter: + """Minimal plotter mock for non-interactive animation tests.""" + + def __init__(self, *_args, **_kwargs): + self.escaped = False + + def show(self, *_args, **_kwargs): + return self + + def render(self): + return None + + def interactive(self): + return self + + def close(self): + return None + + +def _mock_vedo_module(monkeypatch): + """Install a minimal vedo module in sys.modules for tests.""" + + vedo_module = types.ModuleType("vedo") + vedo_module.Mesh = lambda *_args, **_kwargs: _DummyVedoActor() + vedo_module.Box = lambda *_args, **_kwargs: _DummyVedoActor() + vedo_module.Line = lambda *_args, **_kwargs: _DummyVedoActor() + vedo_module.Plotter = _DummyPlotter + vedo_module.settings = types.SimpleNamespace() + monkeypatch.setitem(sys.modules, "vedo", vedo_module) + + +def test_animate_trajectory_runs_with_mocked_vedo(flight_calisto, monkeypatch): + """Test flight trajectory animation entry point through the plots layer.""" + + # Arrange + _mock_vedo_module(monkeypatch) + + # Act + result = flight_calisto.plots.animate_trajectory( + start=0.0, + stop=0.001, + time_step=0.001, + ) + + # Assert + assert result is None + + +def test_animate_rotate_runs_with_mocked_vedo(flight_calisto, monkeypatch): + """Test flight rotation animation entry point through the plots layer.""" + + # Arrange + _mock_vedo_module(monkeypatch) + + # Act + result = flight_calisto.plots.animate_rotate( + start=0.0, + stop=0.001, + time_step=0.001, + ) + + # Assert + assert result is None + + +def test_animate_trajectory_raises_when_vedo_is_missing(flight_calisto, monkeypatch): + """Test that an informative ImportError is raised when vedo is unavailable.""" + + # Arrange + real_import = builtins.__import__ + + def import_without_vedo(name, *args, **kwargs): + if name == "vedo" or name.startswith("vedo."): + raise ImportError("No module named 'vedo'") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_vedo) + + # Act / Assert + with pytest.raises(ImportError, match="optional dependency"): + flight_calisto.plots.animate_trajectory( + start=0.0, + stop=0.001, + time_step=0.001, + ) + + +def test_animate_rotate_raises_when_time_range_is_invalid(flight_calisto, monkeypatch): + """Test validation error for invalid animation time range.""" + + # Arrange + _mock_vedo_module(monkeypatch) + # Act / Assert + with pytest.raises(ValueError, match="Invalid animation time range"): + flight_calisto.plots.animate_rotate( + start=1.0, + stop=0.5, + time_step=0.1, + ) + + +def test_animate_trajectory_raises_when_stl_file_is_missing( + flight_calisto, monkeypatch +): + """Test file validation when STL path does not exist.""" + + # Arrange + _mock_vedo_module(monkeypatch) + + # Act / Assert + with pytest.raises(FileNotFoundError, match="Could not find the 3D model file"): + flight_calisto.plots.animate_trajectory( + "missing_model.stl", + start=0.0, + stop=0.1, + time_step=0.1, + ) + + +@pytest.mark.parametrize("invalid_time_step", [0, -0.1]) +def test_animate_trajectory_raises_when_time_step_is_non_positive( + flight_calisto, monkeypatch, invalid_time_step +): + """Test validation error when animation time_step is not strictly positive.""" + + # Arrange + _mock_vedo_module(monkeypatch) + # Act / Assert + with pytest.raises(ValueError, match="Invalid time_step"): + flight_calisto.plots.animate_trajectory( + start=0.0, + stop=0.1, + time_step=invalid_time_step, + ) + + +def test_animate_rotate_raises_when_stop_exceeds_flight_end( + flight_calisto, monkeypatch +): + """Test validation error when stop time exceeds available simulation range.""" + + # Arrange + _mock_vedo_module(monkeypatch) + # Act / Assert + with pytest.raises(ValueError, match="Invalid animation time range"): + flight_calisto.plots.animate_rotate( + start=0.0, + stop=flight_calisto.t_final + 0.1, + time_step=0.1, + ) + + +def test_animate_trajectory_raises_when_default_model_is_missing( + flight_calisto, monkeypatch +): + """Test failure path when default packaged STL model is unavailable.""" + + # Arrange + _mock_vedo_module(monkeypatch) + monkeypatch.setattr( + flight_calisto.plots, + "_resolve_animation_model_path", + lambda _file_name: "missing_default_model.stl", + ) + + # Act / Assert + with pytest.raises(FileNotFoundError, match="Could not find the 3D model file"): + flight_calisto.plots.animate_trajectory( + start=0.0, + stop=0.1, + time_step=0.1, + )