Initial commit: Keplerian orbital mechanics with patched-conics SOI transitions
- 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
This commit is contained in:
363
visual_test.py
Normal file
363
visual_test.py
Normal file
@@ -0,0 +1,363 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user