201 lines
6.9 KiB
Python
201 lines
6.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.animation import FuncAnimation
|
|
from collections import defaultdict
|
|
import argparse
|
|
import re
|
|
|
|
def parse_line(line):
|
|
"""Parse a line from the simulation output."""
|
|
# Skip comments and empty lines
|
|
if line.startswith('#') or not line.strip():
|
|
return None
|
|
|
|
# Parse the line format: body_name: X = x1m x2m x3m, V = v1m/s v2m/s v3m/s
|
|
try:
|
|
# Split on colon to separate name and data
|
|
name_part, data_part = line.strip().split(':')
|
|
name = name_part.strip()
|
|
|
|
# Split data part into position and velocity
|
|
pos_part, vel_part = data_part.split(',')
|
|
|
|
# Extract position values
|
|
pos_str = pos_part.split('=')[1].strip()
|
|
pos_values = [float(x.replace('m', '').strip()) for x in pos_str.split() if x.strip()]
|
|
|
|
if len(pos_values) != 3:
|
|
return None
|
|
|
|
return {
|
|
'name': name,
|
|
'position': tuple(pos_values)
|
|
}
|
|
except (ValueError, IndexError):
|
|
return None
|
|
|
|
def read_output_file(filename):
|
|
"""Read the simulation output file and organize data by body and time."""
|
|
positions = defaultdict(list)
|
|
times = [] # We'll use frame numbers as times since actual time isn't in output
|
|
|
|
with open(filename, 'r') as f:
|
|
frame = 0
|
|
for line in f:
|
|
data = parse_line(line)
|
|
if data is None:
|
|
continue
|
|
|
|
name = data['name']
|
|
pos = data['position']
|
|
|
|
# Store position
|
|
positions[name].append((pos[0], pos[1]))
|
|
frame += 1
|
|
times.append(frame)
|
|
|
|
return positions, times
|
|
|
|
def calculate_l2_point(sun_pos, earth_pos):
|
|
"""Calculate L2 Lagrange point position."""
|
|
# Calculate distance between bodies
|
|
dx = earth_pos[0] - sun_pos[0]
|
|
dy = earth_pos[1] - sun_pos[1]
|
|
r = np.sqrt(dx*dx + dy*dy)
|
|
|
|
# Calculate mass ratio (using approximate values)
|
|
mu = 5.972e24 / (1.989e30 + 5.972e24) # Earth mass / (Sun mass + Earth mass)
|
|
|
|
# Calculate L2 position (beyond Earth)
|
|
L2_x = sun_pos[0] + (1 + mu**(1/3)) * dx
|
|
L2_y = sun_pos[1] + (1 + mu**(1/3)) * dy
|
|
|
|
return (L2_x, L2_y)
|
|
|
|
def center_positions(positions, center_body):
|
|
"""Center all positions relative to the specified body."""
|
|
if center_body not in positions:
|
|
print(f"Warning: Center body '{center_body}' not found in simulation")
|
|
return positions
|
|
|
|
centered_positions = defaultdict(list)
|
|
for frame in range(len(next(iter(positions.values())))):
|
|
# Get center body position for this frame
|
|
center_pos = positions[center_body][frame]
|
|
|
|
# Shift all positions relative to center body
|
|
for name, pos_list in positions.items():
|
|
if frame < len(pos_list):
|
|
pos = pos_list[frame]
|
|
centered_pos = (pos[0] - center_pos[0], pos[1] - center_pos[1])
|
|
centered_positions[name].append(centered_pos)
|
|
|
|
return centered_positions
|
|
|
|
def create_animation(positions, times, output_file=None, center_body=None):
|
|
"""Create an animation of the bodies' orbits."""
|
|
# Check if we have any data
|
|
if not positions or not times:
|
|
print("Error: No valid data found in the input file")
|
|
return
|
|
|
|
# Center positions if requested
|
|
if center_body:
|
|
positions = center_positions(positions, center_body)
|
|
|
|
# Set up the figure and axis
|
|
fig, ax = plt.subplots(figsize=(10, 10))
|
|
|
|
# Set equal aspect ratio
|
|
ax.set_aspect('equal')
|
|
|
|
# Create scatter plots for each body
|
|
scatters = {}
|
|
for name in positions.keys():
|
|
scatters[name], = ax.plot([], [], 'o-', label=name, alpha=0.7)
|
|
|
|
# Set up the plot
|
|
ax.set_xlabel('X (m)')
|
|
ax.set_ylabel('Y (m)')
|
|
title = 'Orbital Simulation (L2 centered)'
|
|
ax.set_title(title)
|
|
ax.legend()
|
|
|
|
# Find the bounds of the plot
|
|
all_x = []
|
|
all_y = []
|
|
for pos_list in positions.values():
|
|
if pos_list: # Only process if we have positions
|
|
x, y = zip(*pos_list)
|
|
all_x.extend(x)
|
|
all_y.extend(y)
|
|
|
|
if not all_x or not all_y:
|
|
print("Error: No valid position data found")
|
|
return
|
|
|
|
# Calculate initial L2 point for reference
|
|
if 'Sun' in positions and 'Earth' in positions:
|
|
initial_l2 = calculate_l2_point(positions['Sun'][0], positions['Earth'][0])
|
|
# Set a fixed zoom level around L2 (adjust the factor to change zoom)
|
|
zoom_factor = 2e9 # 2 million km view
|
|
ax.set_xlim(-zoom_factor, zoom_factor)
|
|
ax.set_ylim(-zoom_factor, zoom_factor)
|
|
|
|
# Create L2 point scatter plot at origin
|
|
l2_scatter, = ax.plot([0], [0], 'gx', markersize=10, label='L2')
|
|
scatters['L2'] = l2_scatter
|
|
else:
|
|
max_range = max(max(abs(min(all_x)), abs(max(all_x))),
|
|
max(abs(min(all_y)), abs(max(all_y))))
|
|
ax.set_xlim(-max_range, max_range)
|
|
ax.set_ylim(-max_range, max_range)
|
|
|
|
def init():
|
|
"""Initialize the animation."""
|
|
for scatter in scatters.values():
|
|
scatter.set_data([], [])
|
|
return list(scatters.values())
|
|
|
|
def update(frame):
|
|
"""Update the animation for each frame."""
|
|
if 'Sun' in positions and 'Earth' in positions:
|
|
# Calculate L2 point for current frame
|
|
l2_point = calculate_l2_point(positions['Sun'][frame], positions['Earth'][frame])
|
|
|
|
# Update positions relative to L2
|
|
for name, scatter in scatters.items():
|
|
if name == 'L2':
|
|
scatter.set_data([0], [0]) # Keep L2 at origin
|
|
elif positions[name]: # Only update if we have positions
|
|
x, y = zip(*positions[name][:frame+1])
|
|
# Shift positions relative to L2
|
|
x = [pos - l2_point[0] for pos in x]
|
|
y = [pos - l2_point[1] for pos in y]
|
|
scatter.set_data(x, y)
|
|
return list(scatters.values())
|
|
|
|
# Create the animation
|
|
anim = FuncAnimation(fig, update, frames=len(times),
|
|
init_func=init, blit=True,
|
|
interval=50) # 50ms between frames
|
|
|
|
if output_file:
|
|
anim.save(output_file, writer='ffmpeg', fps=30)
|
|
else:
|
|
plt.show()
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Animate orbital simulation output')
|
|
parser.add_argument('input_file', help='Input file from simulation')
|
|
parser.add_argument('--output', '-o', help='Output video file (optional)')
|
|
parser.add_argument('--center', '-c', help='Center the animation on this body')
|
|
args = parser.parse_args()
|
|
|
|
positions, times = read_output_file(args.input_file)
|
|
create_animation(positions, times, args.output, args.center)
|
|
|
|
if __name__ == '__main__':
|
|
main() |