571 lines
21 KiB
Python
Executable File
571 lines
21 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Plot orbital trajectories from binary output file generated by the orbital simulator.
|
|
|
|
The binary file contains snapshots serialized with bincode, where each snapshot has:
|
|
- time: f64
|
|
- bodies: array of Body structs
|
|
- name: String
|
|
- mass: f64
|
|
- position: [f64; 3] (x, y, z)
|
|
- velocity: [f64; 3] (vx, vy, vz)
|
|
- acceleration: [f64; 3] (ax, ay, az)
|
|
|
|
Usage:
|
|
python plot_trajectories.py <trajectory_file.bin>
|
|
"""
|
|
|
|
import sys
|
|
import struct
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits.mplot3d import Axes3D
|
|
import matplotlib.animation as animation
|
|
import numpy as np
|
|
from collections import defaultdict
|
|
import argparse
|
|
|
|
class BinaryReader:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
self.pos = 0
|
|
|
|
def read_u64(self):
|
|
result = struct.unpack('<Q', self.data[self.pos:self.pos+8])[0]
|
|
self.pos += 8
|
|
return result
|
|
|
|
def read_f64(self):
|
|
result = struct.unpack('<d', self.data[self.pos:self.pos+8])[0]
|
|
self.pos += 8
|
|
return result
|
|
|
|
def read_string(self):
|
|
# Read length (u64) then string bytes
|
|
length = self.read_u64()
|
|
result = self.data[self.pos:self.pos+length].decode('utf-8')
|
|
self.pos += length
|
|
return result
|
|
|
|
def read_vec3(self):
|
|
# Read 3 f64 values for position/velocity/acceleration
|
|
x = self.read_f64()
|
|
y = self.read_f64()
|
|
z = self.read_f64()
|
|
return np.array([x, y, z])
|
|
|
|
def read_trajectory_file(filename):
|
|
"""Read the binary trajectory file and return parsed data."""
|
|
with open(filename, 'rb') as f:
|
|
data = f.read()
|
|
|
|
reader = BinaryReader(data)
|
|
snapshots = []
|
|
|
|
try:
|
|
while reader.pos < len(data):
|
|
# Read snapshot
|
|
time = reader.read_f64()
|
|
|
|
# Read number of bodies (u64)
|
|
num_bodies = reader.read_u64()
|
|
|
|
bodies = []
|
|
for _ in range(num_bodies):
|
|
# Read Body struct
|
|
name = reader.read_string()
|
|
mass = reader.read_f64()
|
|
position = reader.read_vec3()
|
|
velocity = reader.read_vec3()
|
|
acceleration = reader.read_vec3()
|
|
|
|
bodies.append({
|
|
'name': name,
|
|
'mass': mass,
|
|
'position': position,
|
|
'velocity': velocity,
|
|
'acceleration': acceleration
|
|
})
|
|
|
|
snapshots.append({
|
|
'time': time,
|
|
'bodies': bodies
|
|
})
|
|
|
|
except struct.error:
|
|
# End of file or corrupted data
|
|
pass
|
|
|
|
return snapshots
|
|
|
|
def organize_trajectories(snapshots):
|
|
"""Organize snapshots into trajectories by body name."""
|
|
trajectories = defaultdict(list)
|
|
times = []
|
|
|
|
for snapshot in snapshots:
|
|
times.append(snapshot['time'])
|
|
for body in snapshot['bodies']:
|
|
trajectories[body['name']].append(body['position'])
|
|
|
|
# Convert lists to numpy arrays
|
|
for name in trajectories:
|
|
trajectories[name] = np.array(trajectories[name])
|
|
|
|
return dict(trajectories), np.array(times)
|
|
|
|
def plot_trajectories_2d(trajectories, times, center_body=None):
|
|
"""Plot 2D trajectories (X-Y plane)."""
|
|
plt.figure(figsize=(12, 10))
|
|
|
|
colors = plt.cm.tab10(np.linspace(0, 1, len(trajectories)))
|
|
|
|
for i, (name, positions) in enumerate(trajectories.items()):
|
|
x = positions[:, 0]
|
|
y = positions[:, 1]
|
|
|
|
plt.plot(x, y, color=colors[i], alpha=0.7, linewidth=1.5, label=name)
|
|
|
|
# Mark starting position
|
|
plt.plot(x[0], y[0], 'o', color=colors[i], markersize=8, alpha=0.8)
|
|
|
|
# Mark current position
|
|
plt.plot(x[-1], y[-1], 's', color=colors[i], markersize=6, alpha=0.8)
|
|
|
|
plt.xlabel('X Position (m)')
|
|
plt.ylabel('Y Position (m)')
|
|
|
|
title = 'Orbital Trajectories (X-Y Plane)'
|
|
if center_body:
|
|
title += f' - Centered on {center_body}'
|
|
plt.title(title)
|
|
|
|
plt.legend()
|
|
plt.grid(True, alpha=0.3)
|
|
plt.axis('equal')
|
|
plt.tight_layout()
|
|
|
|
def plot_trajectories_3d(trajectories, times, center_body=None):
|
|
"""Plot 3D trajectories."""
|
|
fig = plt.figure(figsize=(12, 10))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
colors = plt.cm.tab10(np.linspace(0, 1, len(trajectories)))
|
|
|
|
for i, (name, positions) in enumerate(trajectories.items()):
|
|
x = positions[:, 0]
|
|
y = positions[:, 1]
|
|
z = positions[:, 2]
|
|
|
|
ax.plot(x, y, z, color=colors[i], alpha=0.7, linewidth=1.5, label=name)
|
|
|
|
# Mark starting position
|
|
ax.scatter(x[0], y[0], z[0], color=colors[i], s=100, alpha=0.8, marker='o')
|
|
|
|
# Mark current position
|
|
ax.scatter(x[-1], y[-1], z[-1], color=colors[i], s=60, alpha=0.8, marker='s')
|
|
|
|
ax.set_xlabel('X Position (m)')
|
|
ax.set_ylabel('Y Position (m)')
|
|
ax.set_zlabel('Z Position (m)')
|
|
|
|
title = 'Orbital Trajectories (3D)'
|
|
if center_body:
|
|
title += f' - Centered on {center_body}'
|
|
ax.set_title(title)
|
|
|
|
ax.legend()
|
|
|
|
# Make axes equal
|
|
max_range = 0
|
|
for positions in trajectories.values():
|
|
range_val = np.max(np.abs(positions))
|
|
max_range = max(max_range, range_val)
|
|
|
|
ax.set_xlim(-max_range, max_range)
|
|
ax.set_ylim(-max_range, max_range)
|
|
ax.set_zlim(-max_range, max_range)
|
|
|
|
def plot_energy_over_time(snapshots, times):
|
|
"""Plot energy evolution over time."""
|
|
plt.figure(figsize=(12, 6))
|
|
|
|
total_energies = []
|
|
kinetic_energies = []
|
|
|
|
for snapshot in snapshots:
|
|
ke = 0
|
|
pe = 0
|
|
bodies = snapshot['bodies']
|
|
|
|
# Calculate kinetic energy
|
|
for body in bodies:
|
|
v_squared = np.sum(body['velocity']**2)
|
|
ke += 0.5 * body['mass'] * v_squared
|
|
|
|
# Calculate potential energy (simplified, assuming G=1 in normalized units)
|
|
G = 6.67430e-11 # You might need to adjust this based on your normalization
|
|
for i in range(len(bodies)):
|
|
for j in range(i+1, len(bodies)):
|
|
r = np.linalg.norm(bodies[i]['position'] - bodies[j]['position'])
|
|
pe -= G * bodies[i]['mass'] * bodies[j]['mass'] / r
|
|
|
|
kinetic_energies.append(ke)
|
|
total_energies.append(ke + pe)
|
|
|
|
plt.plot(times, kinetic_energies, label='Kinetic Energy', alpha=0.8)
|
|
plt.plot(times, total_energies, label='Total Energy', alpha=0.8)
|
|
|
|
plt.xlabel('Time (s)')
|
|
plt.ylabel('Energy (J)')
|
|
plt.title('Energy Conservation Over Time')
|
|
plt.legend()
|
|
plt.grid(True, alpha=0.3)
|
|
plt.tight_layout()
|
|
|
|
def plot_animated_2d(trajectories, times, interval=50, center_body=None):
|
|
"""Create animated 2D trajectory plot."""
|
|
fig, ax = plt.subplots(figsize=(12, 10))
|
|
|
|
colors = plt.cm.tab10(np.linspace(0, 1, len(trajectories)))
|
|
body_names = list(trajectories.keys())
|
|
|
|
# Set up the plot
|
|
ax.set_xlabel('X Position (m)')
|
|
ax.set_ylabel('Y Position (m)')
|
|
|
|
title = 'Animated Orbital Trajectories (X-Y Plane)'
|
|
if center_body:
|
|
title += f' - Centered on {center_body}'
|
|
ax.set_title(title)
|
|
ax.grid(True, alpha=0.3)
|
|
|
|
# Calculate plot limits
|
|
all_positions = np.concatenate(list(trajectories.values()))
|
|
margin = 0.1
|
|
x_range = np.max(all_positions[:, 0]) - np.min(all_positions[:, 0])
|
|
y_range = np.max(all_positions[:, 1]) - np.min(all_positions[:, 1])
|
|
x_center = (np.max(all_positions[:, 0]) + np.min(all_positions[:, 0])) / 2
|
|
y_center = (np.max(all_positions[:, 1]) + np.min(all_positions[:, 1])) / 2
|
|
|
|
max_range = max(x_range, y_range) * (1 + margin)
|
|
ax.set_xlim(x_center - max_range/2, x_center + max_range/2)
|
|
ax.set_ylim(y_center - max_range/2, y_center + max_range/2)
|
|
ax.set_aspect('equal')
|
|
|
|
# Initialize plot elements
|
|
trajectory_lines = []
|
|
body_points = []
|
|
body_trails = []
|
|
|
|
for i, name in enumerate(body_names):
|
|
# Trajectory line (will grow over time)
|
|
line, = ax.plot([], [], color=colors[i], alpha=0.7, linewidth=1.5, label=name)
|
|
trajectory_lines.append(line)
|
|
|
|
# Current body position
|
|
point, = ax.plot([], [], 'o', color=colors[i], markersize=8, alpha=0.9)
|
|
body_points.append(point)
|
|
|
|
# Trail of recent positions
|
|
trail, = ax.plot([], [], 'o', color=colors[i], markersize=3, alpha=0.3)
|
|
body_trails.append(trail)
|
|
|
|
# Time display
|
|
time_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, fontsize=12,
|
|
verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
|
|
|
|
ax.legend(loc='upper right')
|
|
|
|
def animate(frame):
|
|
current_time = times[frame]
|
|
|
|
# Format time for display
|
|
if current_time >= 86400: # More than a day
|
|
time_str = f"Time: {current_time/86400:.1f} days"
|
|
elif current_time >= 3600: # More than an hour
|
|
time_str = f"Time: {current_time/3600:.1f} hours"
|
|
else:
|
|
time_str = f"Time: {current_time:.1f} seconds"
|
|
|
|
time_text.set_text(time_str)
|
|
|
|
# Update each body
|
|
for i, name in enumerate(body_names):
|
|
positions = trajectories[name]
|
|
|
|
# Update trajectory line (show path up to current time)
|
|
x_data = positions[:frame+1, 0]
|
|
y_data = positions[:frame+1, 1]
|
|
trajectory_lines[i].set_data(x_data, y_data)
|
|
|
|
# Update current position
|
|
if frame < len(positions):
|
|
current_pos = positions[frame]
|
|
body_points[i].set_data([current_pos[0]], [current_pos[1]])
|
|
|
|
# Update trail (last 20 positions)
|
|
trail_start = max(0, frame - 20)
|
|
trail_x = positions[trail_start:frame, 0]
|
|
trail_y = positions[trail_start:frame, 1]
|
|
body_trails[i].set_data(trail_x, trail_y)
|
|
|
|
return trajectory_lines + body_points + body_trails + [time_text]
|
|
|
|
num_frames = len(times)
|
|
anim = animation.FuncAnimation(fig, animate, frames=num_frames,
|
|
interval=interval, blit=True, repeat=True)
|
|
|
|
return fig, anim
|
|
|
|
def plot_animated_3d(trajectories, times, interval=50, center_body=None):
|
|
"""Create animated 3D trajectory plot."""
|
|
fig = plt.figure(figsize=(12, 10))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
colors = plt.cm.tab10(np.linspace(0, 1, len(trajectories)))
|
|
body_names = list(trajectories.keys())
|
|
|
|
# Set up the plot
|
|
ax.set_xlabel('X Position (m)')
|
|
ax.set_ylabel('Y Position (m)')
|
|
ax.set_zlabel('Z Position (m)')
|
|
|
|
title = 'Animated Orbital Trajectories (3D)'
|
|
if center_body:
|
|
title += f' - Centered on {center_body}'
|
|
ax.set_title(title)
|
|
|
|
# Calculate plot limits
|
|
all_positions = np.concatenate(list(trajectories.values()))
|
|
max_range = np.max(np.abs(all_positions)) * 1.1
|
|
ax.set_xlim(-max_range, max_range)
|
|
ax.set_ylim(-max_range, max_range)
|
|
ax.set_zlim(-max_range, max_range)
|
|
|
|
# Initialize plot elements
|
|
trajectory_lines = []
|
|
body_points = []
|
|
|
|
for i, name in enumerate(body_names):
|
|
# Trajectory line
|
|
line, = ax.plot([], [], [], color=colors[i], alpha=0.7, linewidth=1.5, label=name)
|
|
trajectory_lines.append(line)
|
|
|
|
# Current body position
|
|
point = ax.scatter([], [], [], color=colors[i], s=100, alpha=0.9)
|
|
body_points.append(point)
|
|
|
|
# Time display
|
|
time_text = ax.text2D(0.02, 0.98, '', transform=ax.transAxes, fontsize=12,
|
|
verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
|
|
|
|
ax.legend(loc='upper right')
|
|
|
|
def animate(frame):
|
|
current_time = times[frame]
|
|
|
|
# Format time for display
|
|
if current_time >= 86400: # More than a day
|
|
time_str = f"Time: {current_time/86400:.1f} days"
|
|
elif current_time >= 3600: # More than an hour
|
|
time_str = f"Time: {current_time/3600:.1f} hours"
|
|
else:
|
|
time_str = f"Time: {current_time:.1f} seconds"
|
|
|
|
time_text.set_text(time_str)
|
|
|
|
# Update each body
|
|
for i, name in enumerate(body_names):
|
|
positions = trajectories[name]
|
|
|
|
# Update trajectory line
|
|
x_data = positions[:frame+1, 0]
|
|
y_data = positions[:frame+1, 1]
|
|
z_data = positions[:frame+1, 2]
|
|
trajectory_lines[i].set_data(x_data, y_data)
|
|
trajectory_lines[i].set_3d_properties(z_data)
|
|
|
|
# Update current position
|
|
if frame < len(positions):
|
|
current_pos = positions[frame]
|
|
# Remove old scatter point and create new one
|
|
body_points[i].remove()
|
|
body_points[i] = ax.scatter([current_pos[0]], [current_pos[1]], [current_pos[2]],
|
|
color=colors[i], s=100, alpha=0.9)
|
|
|
|
return trajectory_lines + body_points + [time_text]
|
|
|
|
num_frames = len(times)
|
|
anim = animation.FuncAnimation(fig, animate, frames=num_frames,
|
|
interval=interval, blit=False, repeat=True)
|
|
|
|
return fig, anim
|
|
|
|
def calculate_animation_params(num_frames, target_duration_sec=60, manual_interval=None):
|
|
"""Calculate animation parameters for optimal viewing experience."""
|
|
if manual_interval is not None:
|
|
# User specified manual interval
|
|
interval_ms = manual_interval
|
|
total_duration_sec = num_frames * interval_ms / 1000.0
|
|
time_scale_factor = target_duration_sec / total_duration_sec
|
|
return interval_ms, total_duration_sec, time_scale_factor, True
|
|
|
|
# Auto-calculate interval for target duration
|
|
target_duration_ms = target_duration_sec * 1000
|
|
optimal_interval = max(10, target_duration_ms // num_frames) # Minimum 10ms for smooth animation
|
|
actual_duration_sec = num_frames * optimal_interval / 1000.0
|
|
time_scale_factor = target_duration_sec / actual_duration_sec
|
|
|
|
return optimal_interval, actual_duration_sec, time_scale_factor, False
|
|
|
|
def center_trajectories_on_body(trajectories, center_body_name):
|
|
"""Center all trajectories relative to the specified body."""
|
|
if center_body_name not in trajectories:
|
|
available_bodies = list(trajectories.keys())
|
|
raise ValueError(f"Body '{center_body_name}' not found. Available bodies: {available_bodies}")
|
|
|
|
center_trajectory = trajectories[center_body_name]
|
|
centered_trajectories = {}
|
|
|
|
for body_name, trajectory in trajectories.items():
|
|
# Subtract the center body's position from each body's trajectory
|
|
centered_trajectories[body_name] = trajectory - center_trajectory
|
|
|
|
return centered_trajectories
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Plot orbital trajectories from binary file')
|
|
parser.add_argument('trajectory_file', help='Binary trajectory file to plot')
|
|
parser.add_argument('--2d-only', action='store_true', dest='two_d_only', help='Only show 2D plot')
|
|
parser.add_argument('--3d-only', action='store_true', dest='three_d_only', help='Only show 3D plot')
|
|
parser.add_argument('--energy', action='store_true', help='Show energy plot')
|
|
parser.add_argument('--animate', action='store_true', help='Show animated trajectories')
|
|
parser.add_argument('--save-animation', type=str, help='Save animation as MP4 file')
|
|
parser.add_argument('--static', action='store_true', help='Show static plots (default if no --animate)')
|
|
parser.add_argument('--interval', type=int, help='Animation interval in milliseconds (default: auto-scaled to ~60s total)')
|
|
parser.add_argument('--target-duration', type=int, default=60, help='Target animation duration in seconds (default: 60)')
|
|
parser.add_argument('--center', type=str, help='Center animation on specified body (e.g., "Sun", "Earth")')
|
|
parser.add_argument('--list-bodies', action='store_true', help='List available bodies and exit')
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"Reading trajectory file: {args.trajectory_file}")
|
|
|
|
try:
|
|
snapshots = read_trajectory_file(args.trajectory_file)
|
|
print(f"Loaded {len(snapshots)} snapshots")
|
|
|
|
if not snapshots:
|
|
print("No data found in file!")
|
|
return
|
|
|
|
trajectories, times = organize_trajectories(snapshots)
|
|
|
|
# Handle list-bodies option
|
|
if args.list_bodies:
|
|
print(f"Available bodies in trajectory file:")
|
|
for body_name in sorted(trajectories.keys()):
|
|
print(f" - {body_name}")
|
|
return
|
|
|
|
print(f"Bodies found: {list(trajectories.keys())}")
|
|
print(f"Time range: {times[0]:.2e} - {times[-1]:.2e} seconds")
|
|
print(f"Number of time steps: {len(times)}")
|
|
|
|
# Center trajectories on specified body if requested
|
|
original_trajectories = trajectories.copy()
|
|
if args.center:
|
|
try:
|
|
trajectories = center_trajectories_on_body(trajectories, args.center)
|
|
print(f"🎯 Centering animation on: {args.center}")
|
|
except ValueError as e:
|
|
print(f"❌ Error: {e}")
|
|
return
|
|
|
|
# Check if we should animate or show static plots
|
|
show_animation = args.animate or args.save_animation
|
|
|
|
if show_animation:
|
|
# Calculate animation parameters
|
|
interval_ms, anim_duration_sec, time_scale, is_manual = calculate_animation_params(
|
|
len(times), args.target_duration, args.interval
|
|
)
|
|
|
|
print(f"\n🎬 Animation Settings:")
|
|
print(f" Total frames: {len(times)}")
|
|
print(f" Animation duration: {anim_duration_sec:.1f} seconds")
|
|
print(f" Frame interval: {interval_ms}ms")
|
|
if is_manual:
|
|
print(f" ⚙️ Using manual interval (--interval {args.interval})")
|
|
if time_scale != 1.0:
|
|
print(f" ⏱️ Time scale: {time_scale:.2f}x (animation {'faster' if time_scale > 1 else 'slower'} than target)")
|
|
else:
|
|
print(f" 🤖 Auto-scaled for {args.target_duration}s target duration")
|
|
print(f" ⏱️ Time scale: 1.0x (optimized)")
|
|
|
|
simulation_duration = times[-1] - times[0]
|
|
if simulation_duration > 0:
|
|
compression_ratio = simulation_duration / anim_duration_sec
|
|
if compression_ratio >= 86400:
|
|
print(f" 📈 Compression: {compression_ratio/86400:.1f} days of simulation per second of animation")
|
|
elif compression_ratio >= 3600:
|
|
print(f" 📈 Compression: {compression_ratio/3600:.1f} hours of simulation per second of animation")
|
|
else:
|
|
print(f" 📈 Compression: {compression_ratio:.1f}x real-time")
|
|
print()
|
|
|
|
print("Creating animated plots...")
|
|
animations = []
|
|
|
|
# Create animated plots
|
|
if not args.three_d_only:
|
|
print("Creating 2D animation...")
|
|
fig_2d, anim_2d = plot_animated_2d(trajectories, times, interval_ms, args.center)
|
|
animations.append((fig_2d, anim_2d, '2d'))
|
|
|
|
if not args.two_d_only:
|
|
print("Creating 3D animation...")
|
|
fig_3d, anim_3d = plot_animated_3d(trajectories, times, interval_ms, args.center)
|
|
animations.append((fig_3d, anim_3d, '3d'))
|
|
|
|
# Save animations if requested
|
|
if args.save_animation:
|
|
for fig, anim, plot_type in animations:
|
|
filename = f"{args.save_animation}_{plot_type}.mp4"
|
|
print(f"Saving {plot_type.upper()} animation to {filename}...")
|
|
try:
|
|
anim.save(filename, writer='ffmpeg', fps=20)
|
|
print(f"Animation saved to {filename}")
|
|
except Exception as e:
|
|
print(f"Error saving animation: {e}")
|
|
print("Note: You may need to install ffmpeg for video export")
|
|
|
|
plt.show()
|
|
|
|
else:
|
|
print("Creating static plots...")
|
|
# Use original trajectories for static plots unless centering is requested
|
|
plot_trajectories = trajectories if args.center else original_trajectories
|
|
|
|
# Plot static trajectories
|
|
if not args.three_d_only:
|
|
plot_trajectories_2d(plot_trajectories, times, args.center)
|
|
|
|
if not args.two_d_only:
|
|
plot_trajectories_3d(plot_trajectories, times, args.center)
|
|
|
|
if args.energy:
|
|
plot_energy_over_time(snapshots, times)
|
|
|
|
plt.show()
|
|
|
|
except FileNotFoundError:
|
|
print(f"Error: File '{args.trajectory_file}' not found!")
|
|
except Exception as e:
|
|
print(f"Error reading file: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|