Lucy
655f1c9af6
- orbital_elements.py: elliptical + hyperbolic orbit support - orbit_drawer.py: orbit point generation with SOI truncation - soi_calculator.py: SOI crossing time calculator - frame_transition.py: reference frame switching - test_orbital.py: 147 assertions, all passing - visual_test.py: pygame flyby visualization
364 lines
14 KiB
Python
364 lines
14 KiB
Python
"""
|
||
Pygame visualisation: Keplerian flyby with SOI frame transition.
|
||
|
||
Shows:
|
||
- Star at centre (yellow)
|
||
- Planet (blue) in near-circular orbit, SOI bubble
|
||
- Asteroid (red) on crossing orbit
|
||
- When asteroid enters SOI: switches to planet-relative hyperbolic
|
||
trajectory, orbit trail shows truncated hyperbola inside SOI
|
||
- On exit: switches back to sun frame
|
||
- Countdown in window title
|
||
|
||
Controls:
|
||
SPACE — pause
|
||
LEFT/RIGHT — jump ±10 s
|
||
UP/DOWN — speed ×2 / ÷2
|
||
ESC — quit
|
||
"""
|
||
|
||
import math
|
||
import sys
|
||
from typing import List, Optional, Tuple
|
||
|
||
import pygame
|
||
|
||
from orbital_elements import OrbitalElements, orbital_elements_from_cartesian
|
||
from orbit_drawer import generate_orbit_points
|
||
from soi_calculator import find_soi_crossings
|
||
from frame_transition import (
|
||
sun_state_to_planet_orbit,
|
||
planet_orbit_to_sun_orbit,
|
||
)
|
||
|
||
# ═══════════════════════════════════════════════════════════════
|
||
# Simulation parameters
|
||
# ═══════════════════════════════════════════════════════════════
|
||
|
||
MU_SUN = 1_500_000.0
|
||
MU_PLANET = 10_000.0
|
||
SOI_RADIUS = 70.0
|
||
|
||
PLANET = OrbitalElements(
|
||
a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=MU_SUN,
|
||
)
|
||
|
||
ASTEROID_SUN = OrbitalElements(
|
||
a=380.0, e=0.35, omega=2.5, M0=0.8, epoch=0.0, mu=MU_SUN,
|
||
)
|
||
|
||
SEARCH_ORBITS = 5
|
||
|
||
# ═══════════════════════════════════════════════════════════════
|
||
# Display
|
||
# ═══════════════════════════════════════════════════════════════
|
||
|
||
WIDTH, HEIGHT = 1000, 1000
|
||
CENTER = (WIDTH // 2, HEIGHT // 2)
|
||
FPS = 60
|
||
|
||
STAR_COLOUR = (255, 220, 50)
|
||
STAR_RADIUS = 14
|
||
PLANET_COLOUR = (70, 130, 230)
|
||
PLANET_RADIUS = 8
|
||
SOI_COLOUR = (70, 130, 230, 35)
|
||
ASTEROID_COLOUR = (230, 80, 50)
|
||
ASTEROID_INSIDE_COLOUR = (50, 230, 80)
|
||
ASTEROID_RADIUS = 5
|
||
ORBIT_PLANET_COLOUR = (50, 80, 150)
|
||
ORBIT_ASTEROID_COLOUR = (150, 50, 30)
|
||
ORBIT_FLYBY_COLOUR = (80, 200, 100)
|
||
UI_COLOUR = (200, 200, 200)
|
||
BACKGROUND = (12, 12, 24)
|
||
GRID_COLOUR = (28, 28, 44)
|
||
|
||
|
||
def to_screen(x: float, y: float) -> Tuple[float, float]:
|
||
return CENTER[0] + x, CENTER[1] - y
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════
|
||
# SOI exit-time finder
|
||
# ═══════════════════════════════════════════════════════════════
|
||
|
||
def find_soi_exit_time(
|
||
orbit_planet: OrbitalElements,
|
||
soi_r: float,
|
||
from_time: float,
|
||
max_search: float = 500.0,
|
||
) -> Optional[float]:
|
||
"""
|
||
Given a planet-relative orbit, find the next time *after* from_time
|
||
when distance = soi_r (the exit point).
|
||
|
||
Uses coarse stepping + Newton refinement, similar to soi_calculator.
|
||
"""
|
||
dt = 0.1 # fine step
|
||
t = from_time + dt
|
||
|
||
# Walk forward until we cross outside SOI
|
||
for _ in range(int(max_search / dt)):
|
||
bx, by = orbit_planet.position_at(t)
|
||
r = math.hypot(bx, by)
|
||
if r >= soi_r:
|
||
# Refine with Newton
|
||
return _refine_soi(orbit_planet, soi_r, t - dt, t)
|
||
t += dt
|
||
return None
|
||
|
||
|
||
def _refine_soi(
|
||
orbit: OrbitalElements,
|
||
soi_r: float,
|
||
t_low: float,
|
||
t_high: float,
|
||
) -> float:
|
||
"""Newton-Raphson for r(t) = soi_r (exiting side)."""
|
||
t = t_high # start from outside
|
||
tol = 1e-8
|
||
for _ in range(30):
|
||
bx, by, bvx, bvy = orbit.compute_state_at(t)
|
||
r = math.hypot(bx, by)
|
||
f_val = r - soi_r
|
||
if r < 1e-14:
|
||
break
|
||
fprime = (bx * bvx + by * bvy) / r
|
||
if abs(fprime) < 1e-14:
|
||
break
|
||
t -= f_val / fprime
|
||
t = max(t_low, min(t_high, t))
|
||
if abs(f_val / fprime) < tol:
|
||
break
|
||
return t
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════
|
||
# Main
|
||
# ═══════════════════════════════════════════════════════════════
|
||
|
||
def main() -> None:
|
||
pygame.init()
|
||
screen = pygame.display.set_mode((WIDTH, HEIGHT))
|
||
pygame.display.set_caption("Orbital Flyby — initialising…")
|
||
clock = pygame.time.Clock()
|
||
font = pygame.font.Font(None, 18)
|
||
|
||
# ── Pre-compute sun-frame orbit trails ─────────────────
|
||
trail_planet = [to_screen(x, y) for x, y in
|
||
generate_orbit_points(PLANET)]
|
||
trail_asteroid_sun = [to_screen(x, y) for x, y in
|
||
generate_orbit_points(ASTEROID_SUN)]
|
||
|
||
# ── Pre-compute SOI entry times ────────────────────────
|
||
window = SEARCH_ORBITS * max(ASTEROID_SUN.period, PLANET.period)
|
||
print(f"Scanning {window:.0f}s for SOI crossings…")
|
||
sun_crossings = find_soi_crossings(
|
||
ASTEROID_SUN, PLANET, SOI_RADIUS, 0.0, window, n_steps=1000,
|
||
)
|
||
print(f"Found {len(sun_crossings)} SOI crossing(s):")
|
||
for i, (enter, exit_) in enumerate(sun_crossings):
|
||
print(f" [{i}] enter={enter:.1f}s (sun-frame)")
|
||
|
||
# ── Simulation state ───────────────────────────────────
|
||
# Start 10 s before the first SOI entry so the user sees
|
||
# the countdown tick down without waiting long.
|
||
_first_entry = sun_crossings[0][0] if sun_crossings else 0.0
|
||
sim_time: float = max(0.0, _first_entry - 10.0)
|
||
sim_speed: float = 1.0
|
||
paused: bool = False
|
||
|
||
in_soi: bool = False
|
||
soi_enter_time: Optional[float] = sun_crossings[0][0] if sun_crossings else None
|
||
soi_exit_time: Optional[float] = None
|
||
|
||
# Asteroid orbits (one per frame)
|
||
asteroid_sun: OrbitalElements = ASTEROID_SUN
|
||
asteroid_planet: Optional[OrbitalElements] = None
|
||
|
||
|
||
|
||
cross_idx: int = 0 # index into sun_crossings
|
||
|
||
running = True
|
||
while running:
|
||
dt = clock.tick(FPS) / 1000.0
|
||
|
||
# ── Input ─────────────────────────────────────────
|
||
for event in pygame.event.get():
|
||
if event.type == pygame.QUIT:
|
||
running = False
|
||
elif event.type == pygame.KEYDOWN:
|
||
if event.key == pygame.K_ESCAPE:
|
||
running = False
|
||
elif event.key == pygame.K_SPACE:
|
||
paused = not paused
|
||
elif event.key == pygame.K_RIGHT:
|
||
sim_time += 10.0
|
||
elif event.key == pygame.K_LEFT:
|
||
sim_time = max(0.0, sim_time - 10.0)
|
||
elif event.key == pygame.K_UP:
|
||
sim_speed = min(16.0, sim_speed * 2.0)
|
||
elif event.key == pygame.K_DOWN:
|
||
sim_speed = max(1 / 16, sim_speed / 2.0)
|
||
|
||
if not paused:
|
||
sim_time += dt * sim_speed
|
||
|
||
# ── SOI state machine ─────────────────────────────
|
||
# Check entry
|
||
if (not in_soi and soi_enter_time is not None
|
||
and sim_time >= soi_enter_time):
|
||
# ── ENTER SOI ──
|
||
in_soi = True
|
||
print(f"\n>>> SOI ENTRY at t={sim_time:.1f}s")
|
||
|
||
px, py, pvx, pvy = PLANET.compute_state_at(sim_time)
|
||
ax, ay, avx, avy = asteroid_sun.compute_state_at(sim_time)
|
||
asteroid_planet = sun_state_to_planet_orbit(
|
||
ax, ay, avx, avy, px, py, pvx, pvy, MU_PLANET, sim_time,
|
||
)
|
||
print(f" Planet-relative orbit: a={asteroid_planet.a:.1f}, "
|
||
f"e={asteroid_planet.e:.3f}, "
|
||
f"hyperbolic={asteroid_planet.is_hyperbolic}")
|
||
|
||
# Compute exit time
|
||
soi_exit_time = find_soi_exit_time(asteroid_planet, SOI_RADIUS, sim_time)
|
||
if soi_exit_time:
|
||
print(f" Computed SOI exit: t={soi_exit_time:.1f}s "
|
||
f"(dur={soi_exit_time - sim_time:.1f}s)")
|
||
else:
|
||
print(" WARNING: could not find SOI exit time!")
|
||
|
||
|
||
|
||
# Check exit
|
||
if in_soi and soi_exit_time is not None and sim_time >= soi_exit_time:
|
||
# ── EXIT SOI ──
|
||
in_soi = False
|
||
print(f">>> SOI EXIT at t={sim_time:.1f}s")
|
||
|
||
px, py, pvx, pvy = PLANET.compute_state_at(sim_time)
|
||
asteroid_sun = planet_orbit_to_sun_orbit(
|
||
asteroid_planet, px, py, pvx, pvy, MU_SUN, sim_time,
|
||
)
|
||
print(f" New sun orbit: a={asteroid_sun.a:.1f}, "
|
||
f"e={asteroid_sun.e:.3f}")
|
||
|
||
asteroid_planet = None
|
||
soi_exit_time = None
|
||
|
||
# Advance to next crossing
|
||
cross_idx += 1
|
||
while cross_idx < len(sun_crossings):
|
||
if sun_crossings[cross_idx][0] > sim_time + 0.1:
|
||
soi_enter_time = sun_crossings[cross_idx][0]
|
||
break
|
||
cross_idx += 1
|
||
else:
|
||
soi_enter_time = None
|
||
|
||
# ── Compute current positions ──────────────────────
|
||
px, py = PLANET.position_at(sim_time)
|
||
|
||
if in_soi and asteroid_planet is not None:
|
||
# Asteroid in planet frame → convert to sun frame
|
||
ax_rel, ay_rel = asteroid_planet.position_at(sim_time)
|
||
ax = ax_rel + px
|
||
ay = ay_rel + py
|
||
else:
|
||
ax, ay = asteroid_sun.position_at(sim_time)
|
||
|
||
sp = to_screen(px, py)
|
||
sa = to_screen(ax, ay)
|
||
|
||
# ── Draw ──────────────────────────────────────────
|
||
screen.fill(BACKGROUND)
|
||
|
||
# Grid
|
||
for g in range(-400, 401, 100):
|
||
sx1, sy1 = to_screen(g, -400)
|
||
sx2, sy2 = to_screen(g, 400)
|
||
pygame.draw.line(screen, GRID_COLOUR, (sx1, sy1), (sx2, sy2), 1)
|
||
sx1, sy1 = to_screen(-400, g)
|
||
sx2, sy2 = to_screen(400, g)
|
||
pygame.draw.line(screen, GRID_COLOUR, (sx1, sy1), (sx2, sy2), 1)
|
||
|
||
# Sun-frame asteroid trail (always shown)
|
||
if len(trail_asteroid_sun) > 1:
|
||
pygame.draw.lines(screen, ORBIT_ASTEROID_COLOUR, True,
|
||
trail_asteroid_sun, 1)
|
||
|
||
# Planet orbit trail
|
||
if len(trail_planet) > 1:
|
||
pygame.draw.lines(screen, ORBIT_PLANET_COLOUR, True,
|
||
trail_planet, 1)
|
||
|
||
# Flyby trail (planet-relative hyperbolic arc, inside SOI)
|
||
if in_soi and asteroid_planet is not None:
|
||
pts = generate_orbit_points(
|
||
asteroid_planet, num_points=200, max_radius=SOI_RADIUS,
|
||
)
|
||
if len(pts) > 1:
|
||
# pts are in planet frame — offset by planet world position
|
||
offset_trail = [to_screen(x + px, y + py) for x, y in pts]
|
||
pygame.draw.lines(screen, ORBIT_FLYBY_COLOUR, False,
|
||
offset_trail, 2)
|
||
|
||
# Planet SOI bubble
|
||
soi_surf = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
|
||
pygame.draw.circle(
|
||
soi_surf, SOI_COLOUR,
|
||
(int(sp[0]), int(sp[1])), int(SOI_RADIUS),
|
||
)
|
||
screen.blit(soi_surf, (0, 0))
|
||
|
||
# Star
|
||
pygame.draw.circle(screen, STAR_COLOUR, CENTER, STAR_RADIUS)
|
||
pygame.draw.circle(screen, (255, 240, 160), CENTER, STAR_RADIUS + 4, 2)
|
||
|
||
# Planet
|
||
pygame.draw.circle(screen, PLANET_COLOUR,
|
||
(int(sp[0]), int(sp[1])), PLANET_RADIUS)
|
||
|
||
# Asteroid (green inside SOI)
|
||
acol = ASTEROID_INSIDE_COLOUR if in_soi else ASTEROID_COLOUR
|
||
arad = ASTEROID_RADIUS + (2 if in_soi else 0)
|
||
pygame.draw.circle(screen, acol, (int(sa[0]), int(sa[1])), arad)
|
||
|
||
# ── Title / countdown ─────────────────────────────
|
||
if in_soi and soi_exit_time is not None:
|
||
remaining = soi_exit_time - sim_time
|
||
title = f"⚫ INSIDE SOI — flyby exit in {remaining:.1f}s"
|
||
elif soi_enter_time is not None:
|
||
countdown = soi_enter_time - sim_time
|
||
title = (f"SOI entry in {countdown:.1f}s "
|
||
f"(crossing {cross_idx + 1}/{len(sun_crossings)})")
|
||
else:
|
||
title = "No upcoming SOI crossing"
|
||
pygame.display.set_caption(title)
|
||
|
||
# ── HUD ───────────────────────────────────────────
|
||
lines = [
|
||
f"t={sim_time:.1f}s speed={sim_speed:.1f}× "
|
||
f"{'PAUSED' if paused else 'RUNNING'} "
|
||
f"{'[SOI]' if in_soi else '[SUN]'}",
|
||
f"Planet T={PLANET.period:.1f}s "
|
||
f"r=[{PLANET.periapsis:.0f}, {PLANET.apoapsis:.0f}] px",
|
||
f"SOI={SOI_RADIUS:.0f} px μ_planet={MU_PLANET}",
|
||
"SPACE=pause ←→=skip ↑↓=speed ESC=quit",
|
||
]
|
||
y = 12
|
||
for line in lines:
|
||
surf = font.render(line, True, UI_COLOUR)
|
||
screen.blit(surf, (12, y))
|
||
y += 22
|
||
|
||
pygame.display.flip()
|
||
|
||
pygame.quit()
|
||
sys.exit()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|