feat: Implement 3D flight animation methods using Vedo (Issue #523)

pyproject.toml

@@ -67,7 +67,11 @@ monte-carlo = [
     "prettytable",
 ]
 
-all = ["rocketpy[env-analysis]", "rocketpy[monte-carlo]"]
+animation = [
+    "vedo>=2024.5.1"
+]
+
+all = ["rocketpy[env-analysis]", "rocketpy[monte-carlo]", "rocketpy[animation]"]
 
 
 [tool.coverage.report]

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

rocketpy/simulation/flight.py

@@ -1,5 +1,6 @@
 # pylint: disable=too-many-lines
 import math
+import time
 import warnings
 from copy import deepcopy
 from functools import cached_property
@@ -3743,7 +3744,202 @@ def export_kml(
             color=color,
             altitude_mode=altitude_mode,
         )
+    def animate_trajectory(self, file_name, start=0, stop=None, time_step=0.1, **kwargs):
+        """
+        6-DOF Animation of the flight trajectory using Vedo.
+
+        Parameters
+        ----------
+        file_name : str
+            3D object file representing your rocket, usually in .stl format.
+            Example: "rocket.stl"
+        start : int, float, optional
+            Time for starting animation, in seconds. Default is 0.
+        stop : int, float, optional
+            Time for ending animation, in seconds. If None, uses self.t_final.
+            Default is None.
+        time_step : float, optional
+            Time step for data interpolation in the animation. Default is 0.1.
+        **kwargs : dict, optional
+            Additional keyword arguments to be passed to vedo.Plotter.show().
+            Common arguments:
+            - azimuth (float): Rotation in degrees around the vertical axis.
+            - elevation (float): Rotation in degrees above the horizon.
+            - roll (float): Rotation in degrees around the view axis.
+            - zoom (float): Zoom level (default 1).
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        ImportError
+            If the 'vedo' package is not installed.
+
+        Notes
+        -----
+        This feature requires the 'vedo' package. Install it with:
+        pip install rocketpy[animation]
+        """
+        try:
+            from vedo import Box, Line, Mesh, Plotter, settings
+        except ImportError as e:
+            raise ImportError(
+                "The animation feature requires the 'vedo' package. "
+                "Install it with:\n"
+                "    pip install rocketpy[animation]\n"
+                "Or directly:\n"
+                "    pip install vedo>=2024.5.1"
+            ) from e
+
+        # Enable interaction if needed
+        try:
+            settings.allow_interaction = True
+        except AttributeError:
+            pass  # Not available in newer versions of vedo
+
+        # Handle stop time
+        if stop is None:
+            stop = self.t_final
+
+        # Define the world bounds based on trajectory
+        max_x = max(self.x[:, 1])
+        max_y = max(self.y[:, 1])
+        # Use simple logic for bounds
+        world = Box(
+            pos=[self.x(start), self.y(start), self.apogee],
+            length=max_x * 2 if max_x != 0 else 1000,
+            width=max_y * 2 if max_y != 0 else 1000,
+            height=self.apogee,
+        ).wireframe()
+
+        # Load rocket mesh
+        rocket = Mesh(file_name).c("green")
+        rocket.pos(self.x(start), self.y(start), 0).add_trail(n=len(self.x[:, 1]))
+        # Create trail
+        trail_points = [[self.x(t), self.y(t), self.z(t) - self.env.elevation] 
+                        for t in np.arange(start, stop, time_step)]
+        trail = Line(trail_points, c="k", alpha=0.5)
+        # Setup Plotter
+        plt = Plotter(axes=1, interactive=False)
+        plt.show(world, rocket, __doc__, viewup="z", **kwargs)
+
+        # Animation Loop
+        for t in np.arange(start, stop, time_step):
+            # Calculate rotation angle and vector from quaternions
+            # Note: This simple rotation logic mimics the old branch. 
+            # Ideally, vedo handles orientation via matrix, but we stick 
+            # to the provided logic for now.
+
+            # e0 is the scalar part of the quaternion
+            angle = np.arccos(2 * self.e0(t)**2 - 1) 
+            k = np.sin(angle / 2) if np.sin(angle / 2) != 0 else 1
+
+            # Update position and rotation
+            # Adjusting for ground elevation
+            rocket.pos(self.x(t), self.y(t), self.z(t) - self.env.elevation)
+            rocket.rotate_x(self.e1(t) / k)
+            rocket.rotate_y(self.e2(t) / k)
+            rocket.rotate_z(self.e3(t) / k)
+
+            # update the scene
+            plt.show(world, rocket, trail)
+
+            # slow down to make animation visible
+            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, start=0, stop=None, time_step=0.1, **kwargs):
+        """
+        Animation of rocket attitude (rotation) during the flight.
 
+        Parameters
+        ----------
+        file_name : str
+            3D object file representing your rocket, usually in .stl format.
+        start : int, float, optional
+            Time for starting animation, in seconds. Default is 0.
+        stop : int, float, optional
+            Time for ending animation, in seconds. If None, uses self.t_final.
+            Default is None.
+        time_step : float, optional
+            Time step for data interpolation. Default is 0.1.
+        **kwargs : dict, optional
+            Additional keyword arguments to be passed to vedo.Plotter.show().
+            Common arguments:
+            - azimuth (float): Rotation in degrees around the vertical axis.
+            - elevation (float): Rotation in degrees above the horizon.
+            - roll (float): Rotation in degrees around the view axis.
+            - zoom (float): Zoom level (default 1).
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        ImportError
+            If the 'vedo' package is not installed.
+        """
+        try:
+            from vedo import Box, Mesh, Plotter, settings
+        except ImportError as e:
+            raise ImportError(
+                "The animation feature requires the 'vedo' package. "
+                "Install it with:\n"
+                "    pip install rocketpy[animation]\n"
+            ) from e
+
+        # Enable interaction if needed
+        try:
+            settings.allow_interaction = True
+        except AttributeError:
+            pass  # Not available in newer versions of vedo
+
+        if stop is None:
+            stop = self.t_final
+
+        # Smaller box for rotation view
+        world = Box(
+            pos=[self.x(start), self.y(start), self.apogee],
+            length=max(self.x[:, 1]) * 0.2,
+            width=max(self.y[:, 1]) * 0.2,
+            height=self.apogee * 0.1,
+        ).wireframe()
+
+        rocket = Mesh(file_name).c("green")
+        # Initialize at origin relative to view
+        rocket.pos(self.x(start), self.y(start), 0).add_trail(n=len(self.x[:, 1]))
+
+        plt = Plotter(axes=1, interactive=False)
+        plt.show(world, rocket, __doc__, viewup="z", **kwargs)
+
+        for t in np.arange(start, stop, time_step):
+            angle = np.arccos(2 * self.e0(t)**2 - 1)
+            k = np.sin(angle / 2) if np.sin(angle / 2) != 0 else 1
+
+            # Keep position static (relative start) to observe only rotation
+            rocket.pos(self.x(start), self.y(start), 0)
+            rocket.rotate_x(self.e1(t) / k)
+            rocket.rotate_y(self.e2(t) / k)
+            rocket.rotate_z(self.e3(t) / k)
+
+            plt.show(world, rocket)
+
+            if getattr(plt, 'escaped', False):
+                break
+
+        plt.interactive().close()
+        return None
+
     def info(self):
         """Prints out a summary of the data available about the Flight."""
         self.prints.all()

tests/animation_verification/main.py

@@ -0,0 +1,73 @@
+import os
+import traceback
+from rocketpy import Environment, Flight
+from rocket_stl import create_rocket_stl
+from rocket_setup import get_calisto_rocket
+
+
+def run_simulation_and_test_animation():
+    print("🚀 Setting up simulation (Calisto Example)...")
+
+    # 1. Setup Environment
+    env = Environment(latitude=32.990254, longitude=-106.974998, elevation=1400)
+    env.set_date((2025, 12, 5, 12))
+    env.set_atmospheric_model(type="standard_atmosphere")
+
+    # 2. Get Rocket
+    try:
+        calisto = get_calisto_rocket()
+    except Exception as e:
+        print(f"❌ Failed to configure rocket: {e}")
+        return
+
+    # 3. Simulate Flight
+    test_flight = Flight(
+        rocket=calisto, environment=env, rail_length=5.2, inclination=85, heading=0
+    )
+
+    print(f"✅ Flight simulated successfully! Apogee: {test_flight.apogee:.2f} m")
+
+    # 4. Test Animation Methods
+    stl_file = "rocket_model.stl"
+    # Note: Depending on where you run this, you might need to adjust imports
+    # or ensure create_rocket_stl is available in scope.
+    create_rocket_stl(stl_file, length=300, radius=50)
+
+    print("\n🎥 Testing animate_trajectory()...")
+
+    try:
+        test_flight.animate_trajectory(
+            file_name=stl_file,
+            stop=15.0,
+            time_step=0.05,
+            azimuth=-45,  # Rotates view 45 degrees left
+            elevation=30,  # Tilts view 30 degrees up
+            zoom=1.2,
+        )
+        print("✅ animate_trajectory() executed successfully.")
+    except Exception as e:
+        print(f"❌ animate_trajectory() Failed: {e}")
+        traceback.print_exc()
+
+    print("\n🔄 Testing animate_rotate()...")
+
+    try:
+        test_flight.animate_rotate(
+            file_name=stl_file,
+            time_step=1.0,
+            azimuth=-45,  # Rotates view 45 degrees left
+            elevation=30,  # Tilts view 30 degrees up
+            zoom=1.2,
+        )
+        print("✅ animate_rotate() executed successfully.")
+    except Exception as e:
+        print(f"❌ animate_rotate() Failed: {e}")
+        traceback.print_exc()
+
+    # Cleanup
+    if os.path.exists(stl_file):
+        os.remove(stl_file)
+
+
+if __name__ == "__main__":
+    run_simulation_and_test_animation()