""" 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()