test: add modular verification suite for 3D animation methods

GuilhermeAsura · Gui-FernandesBR · commit 62912d79e1bf · 2025-12-08T21:25:01.000-03:00
diff --git a/tests/animation_verification/__init__.py b/tests/animation_verification/__init__.py
diff --git a/tests/animation_verification/main.py b/tests/animation_verification/main.py
@@ -0,0 +1,71 @@
+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
@@ -0,0 +1,82 @@
+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
@@ -0,0 +1,220 @@
+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
+    body_length = length * 0.65
+    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
+        
+        # Fin vertices
+        fin_base_center = [0, 0, body_z_end * 0.5]
+        
+        # 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('<I', len(faces)))
+            
+            for face in faces:
+                # Write normal (dummy, will be recalculated)
+                f.write(struct.pack('<fff', 0.0, 0.0, 0.0))
+                # Write three vertices
+                for vertex in face:
+                    f.write(struct.pack('<fff', vertex[0], vertex[1], vertex[2]))
+                # Write attribute byte count
+                f.write(struct.pack('<H', 0))
+    
+    # Simplified geometry - create a more detailed rocket than original but simpler than full version
+    num_segments = 16
+    faces = []
+    
+    # Add more triangles for better shape...
+    # (You can combine the face generation logic from above with this write_stl function)
+    
+    write_stl(faces, filename)
