135 lines
4.2 KiB
Python
135 lines
4.2 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']
|
|
positions[name].append((pos[0], pos[1]))
|
|
frame += 1
|
|
times.append(frame)
|
|
|
|
return positions, times
|
|
|
|
def create_animation(positions, times, output_file=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
|
|
|
|
# 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)')
|
|
ax.set_title('Orbital Simulation')
|
|
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
|
|
|
|
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."""
|
|
for name, scatter in scatters.items():
|
|
if positions[name]: # Only update if we have positions
|
|
x, y = zip(*positions[name][:frame+1])
|
|
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)')
|
|
args = parser.parse_args()
|
|
|
|
positions, times = read_output_file(args.input_file)
|
|
create_animation(positions, times, args.output)
|
|
|
|
if __name__ == '__main__':
|
|
main() |