diff --git a/pyproject.toml b/pyproject.toml index 48aeb8bc6..a3dbcd5e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,7 +68,11 @@ monte-carlo = [ "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.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/simulation/flight.py b/rocketpy/simulation/flight.py index a7e0374b2..82f8c99f4 100644 --- a/rocketpy/simulation/flight.py +++ b/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 @@ -4039,7 +4040,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() diff --git a/tests/animation_verification/__init__.py b/tests/animation_verification/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/animation_verification/main.py b/tests/animation_verification/main.py new file mode 100644 index 000000000..5538bd7fe --- /dev/null +++ b/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() diff --git a/tests/animation_verification/rocket_setup.py b/tests/animation_verification/rocket_setup.py new file mode 100644 index 000000000..4f3c11ba1 --- /dev/null +++ b/tests/animation_verification/rocket_setup.py @@ -0,0 +1,94 @@ +import os +from rocketpy import SolidMotor, Rocket + +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +ROOT_DIR = os.path.dirname(os.path.dirname(CURRENT_DIR)) +DATA_DIR = os.path.join(ROOT_DIR, "data") + +OFF_DRAG_PATH = os.path.join(DATA_DIR, "rockets/calisto/powerOffDragCurve.csv") +ON_DRAG_PATH = os.path.join(DATA_DIR, "rockets/calisto/powerOnDragCurve.csv") +AIRFOIL_PATH = os.path.join(DATA_DIR, "airfoils/NACA0012-radians.txt") +MOTOR_PATH = os.path.join(DATA_DIR, "motors/cesaroni/Cesaroni_M1670.eng") + + +def get_motor(): + """Locates the motor file and returns a configured SolidMotor object.""" + # We can now point directly to the file without searching + if not os.path.exists(MOTOR_PATH): + raise FileNotFoundError(f"Could not find Cesaroni_M1670.eng at: {MOTOR_PATH}") + + return SolidMotor( + thrust_source=MOTOR_PATH, + dry_mass=1.815, + dry_inertia=(0.125, 0.125, 0.002), + nozzle_radius=33 / 1000, + grain_number=5, + grain_density=1815, + grain_outer_radius=33 / 1000, + grain_initial_inner_radius=15 / 1000, + grain_initial_height=120 / 1000, + grain_separation=5 / 1000, + grains_center_of_mass_position=0.397, + center_of_dry_mass_position=0.317, + nozzle_position=0, + burn_time=3.9, + throat_radius=11 / 1000, + coordinate_system_orientation="nozzle_to_combustion_chamber", + ) + + +def get_calisto_rocket(): + """Configures and returns the Calisto Rocket object.""" + motor = get_motor() + + calisto = Rocket( + radius=127 / 2000, + mass=14.426, + inertia=(6.321, 6.321, 0.034), + power_off_drag=OFF_DRAG_PATH, + power_on_drag=ON_DRAG_PATH, + center_of_mass_without_motor=0, + coordinate_system_orientation="tail_to_nose", + ) + + calisto.add_motor(motor, position=-1.255) + calisto.set_rail_buttons( + upper_button_position=0.0818, + lower_button_position=-0.618, + angular_position=45, + ) + + # Aerodynamic surfaces + calisto.add_nose(length=0.55829, kind="vonKarman", position=1.27) + calisto.add_trapezoidal_fins( + n=4, + root_chord=0.120, + tip_chord=0.060, + span=0.110, + position=-1.04956, + cant_angle=0, + airfoil=(AIRFOIL_PATH, "radians"), + ) + calisto.add_tail( + top_radius=0.0635, bottom_radius=0.0435, length=0.060, position=-1.194656 + ) + + # Parachutes + calisto.add_parachute( + name="Main", + cd_s=10.0, + trigger=800, + sampling_rate=105, + lag=1.5, + noise=(0, 8.3, 0.5), + ) + calisto.add_parachute( + name="Drogue", + cd_s=1.0, + trigger="apogee", + sampling_rate=105, + lag=1.5, + noise=(0, 8.3, 0.5), + ) + + return calisto diff --git a/tests/animation_verification/rocket_stl.py b/tests/animation_verification/rocket_stl.py new file mode 100644 index 000000000..700285d07 --- /dev/null +++ b/tests/animation_verification/rocket_stl.py @@ -0,0 +1,232 @@ +import math +import numpy as np +from stl import mesh # Requires numpy-stl package: pip install numpy-stl +import struct + + +def create_rocket_stl(filename="rocket_model.stl", length=100, radius=10): + """ + Creates a detailed rocket-shaped STL file with proper scale. + + Parameters + ---------- + filename : str + Output filename + length : float + Rocket length in meters + radius : float + Rocket base radius in meters + """ + + # More realistic proportions + nose_length = length * 0.25 + engine_length = length * 0.1 + + # Fin parameters + fin_height = radius * 3 + fin_width = radius * 2 + fin_thickness = radius * 0.2 + num_fins = 4 + + # Calculate number of vertices for smooth curves + num_segments = 36 + vertices = [] + faces = [] + + # Helper function to add triangle + def add_triangle(v1, v2, v3): + nonlocal faces + idx = len(vertices) + vertices.extend([v1, v2, v3]) + faces.append([idx, idx + 1, idx + 2]) + + # Helper function to add quadrilateral as two triangles + def add_quad(v1, v2, v3, v4): + add_triangle(v1, v2, v3) + add_triangle(v1, v3, v4) + + # 1. Create nose cone (pointed ogive) + nose_vertices = [] + for i in range(num_segments + 1): + angle = 2 * math.pi * i / num_segments + x = math.cos(angle) * radius + y = math.sin(angle) * radius + z = length # Tip at full length + nose_vertices.append([x, y, z]) + + # Create nose cone triangles + nose_tip = [0, 0, length] + for i in range(num_segments): + v1 = nose_tip + v2 = nose_vertices[i] + v3 = nose_vertices[(i + 1) % num_segments] + add_triangle(v1, v2, v3) + + # 2. Create main body cylinder + body_z_start = length - nose_length + body_z_end = engine_length + + # Create vertices for top and bottom circles of body + top_circle = [] + bottom_circle = [] + + for i in range(num_segments): + angle = 2 * math.pi * i / num_segments + x = math.cos(angle) * radius + y = math.sin(angle) * radius + + top_circle.append([x, y, body_z_start]) + bottom_circle.append([x, y, body_z_end]) + + # Create body cylinder triangles + for i in range(num_segments): + next_i = (i + 1) % num_segments + + # Side quad + v1 = top_circle[i] + v2 = top_circle[next_i] + v3 = bottom_circle[next_i] + v4 = bottom_circle[i] + add_quad(v1, v2, v3, v4) + + # Top circle (connecting to nose) + v1 = top_circle[i] + v2 = top_circle[next_i] + v3 = nose_vertices[i] + add_triangle(v1, v2, v3) + + # 3. Create engine nozzle (truncated cone) + engine_radius = radius * 0.7 + engine_z_end = 0 + + # Create vertices for engine circles + engine_top_circle = [] + engine_bottom_circle = [] + + for i in range(num_segments): + angle = 2 * math.pi * i / num_segments + x_top = math.cos(angle) * radius + y_top = math.sin(angle) * radius + x_bottom = math.cos(angle) * engine_radius + y_bottom = math.sin(angle) * engine_radius + + engine_top_circle.append([x_top, y_top, body_z_end]) + engine_bottom_circle.append([x_bottom, y_bottom, engine_z_end]) + + # Create engine nozzle triangles + for i in range(num_segments): + next_i = (i + 1) % num_segments + + # Side quad + v1 = engine_top_circle[i] + v2 = engine_top_circle[next_i] + v3 = engine_bottom_circle[next_i] + v4 = engine_bottom_circle[i] + add_quad(v1, v2, v3, v4) + + # Connect to body + v1 = engine_top_circle[i] + v2 = engine_top_circle[next_i] + v3 = bottom_circle[i] + add_triangle(v1, v2, v3) + + # 4. Create fins + for fin_num in range(num_fins): + fin_angle = 2 * math.pi * fin_num / num_fins + + # Inner fin edge (attached to rocket) + inner_x = math.cos(fin_angle) * radius + inner_y = math.sin(fin_angle) * radius + inner_top = [inner_x, inner_y, body_z_end + fin_height * 0.3] + inner_bottom = [inner_x, inner_y, body_z_end - fin_height * 0.7] + + # Outer fin edge + outer_x = math.cos(fin_angle) * (radius + fin_width) + outer_y = math.sin(fin_angle) * (radius + fin_width) + outer_top = [outer_x, outer_y, body_z_end + fin_height * 0.3] + outer_bottom = [outer_x, outer_y, body_z_end - fin_height * 0.7] + + # Fin side that attaches to rocket + add_quad( + inner_top, + inner_bottom, + bottom_circle[fin_num * num_segments // num_fins], + bottom_circle[(fin_num * num_segments // num_fins + 1) % num_segments], + ) + + # Fin surfaces + # Front face + add_quad(inner_top, outer_top, outer_bottom, inner_bottom) + + # Top face + add_triangle( + inner_top, + outer_top, + [ + inner_top[0] + math.cos(fin_angle + math.pi / 2) * fin_thickness * 0.5, + inner_top[1] + math.sin(fin_angle + math.pi / 2) * fin_thickness * 0.5, + inner_top[2], + ], + ) + + # Bottom face + add_triangle( + inner_bottom, + outer_bottom, + [ + inner_bottom[0] + + math.cos(fin_angle + math.pi / 2) * fin_thickness * 0.5, + inner_bottom[1] + + math.sin(fin_angle + math.pi / 2) * fin_thickness * 0.5, + inner_bottom[2], + ], + ) + + # 5. Create bottom cap + center_bottom = [0, 0, engine_z_end] + for i in range(num_segments): + next_i = (i + 1) % num_segments + add_triangle( + center_bottom, engine_bottom_circle[i], engine_bottom_circle[next_i] + ) + + # Convert to numpy arrays + vertices_array = np.array(vertices) + faces_array = np.array(faces) + + # Create the mesh + rocket_mesh = mesh.Mesh(np.zeros(faces_array.shape[0], dtype=mesh.Mesh.dtype)) + for i, face in enumerate(faces_array): + for j in range(3): + rocket_mesh.vectors[i][j] = vertices_array[face[j]] + + # Write the mesh to file + rocket_mesh.save(filename) + print(f"Rocket STL saved to {filename}") + print(f"Total triangles: {len(faces)}") + + +# Alternative version using pure Python STL generation +def create_rocket_stl_simple(filename="rocket_simple.stl", length=100, radius=10): + """Simpler version using pure Python without numpy-stl dependency""" + + def write_stl(faces, filename): + with open(filename, "wb") as f: + # Write 80 byte header + f.write(b"\x00" * 80) + # Write number of faces + f.write(struct.pack("