commit 655f1c9af633cabf717d8844f9e7fe9a03c1ed25 Author: Lucy Date: Thu May 21 20:31:17 2026 +0200 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 diff --git a/frame_transition.py b/frame_transition.py new file mode 100644 index 0000000..b4d1817 --- /dev/null +++ b/frame_transition.py @@ -0,0 +1,90 @@ +""" +Reference-frame transitions for patched-conics SOI handling. + +When a body enters a planet's sphere of influence its orbital +elements must be re-expressed relative to the planet. On exit +the reverse transition occurs. +""" + +from typing import Tuple + +from orbital_elements import OrbitalElements, orbital_elements_from_cartesian + + +def sun_state_to_planet_orbit( + body_x: float, + body_y: float, + body_vx: float, + body_vy: float, + planet_x: float, + planet_y: float, + planet_vx: float, + planet_vy: float, + mu_planet: float, + epoch: float, +) -> OrbitalElements: + """ + Convert a body's sun-frame state to a planet-relative orbit. + + Parameters + ---------- + body_* — body position & velocity in sun frame + planet_* — planet position & velocity in sun frame + mu_planet — planet's gravitational parameter + epoch — time of transition + + Returns + ------- + OrbitalElements describing the body's trajectory around the planet. + """ + x_rel = body_x - planet_x + y_rel = body_y - planet_y + vx_rel = body_vx - planet_vx + vy_rel = body_vy - planet_vy + return orbital_elements_from_cartesian( + x_rel, y_rel, vx_rel, vy_rel, mu_planet, epoch, + ) + + +def planet_orbit_to_sun_state( + orbit_planet: OrbitalElements, + planet_x: float, + planet_y: float, + planet_vx: float, + planet_vy: float, + t: float, +) -> Tuple[float, float, float, float]: + """ + Convert a planet-relative orbit back to sun-frame state at time t. + + Returns (x, y, vx, vy) in sun frame. + """ + x_rel, y_rel, vx_rel, vy_rel = orbit_planet.compute_state_at(t) + return ( + x_rel + planet_x, + y_rel + planet_y, + vx_rel + planet_vx, + vy_rel + planet_vy, + ) + + +def planet_orbit_to_sun_orbit( + orbit_planet: OrbitalElements, + planet_x: float, + planet_y: float, + planet_vx: float, + planet_vy: float, + mu_sun: float, + t: float, +) -> OrbitalElements: + """ + Full round-trip: planet-relative orbit → sun-frame orbital elements. + + Used after SOI exit to put the body back into a heliocentric orbit. + """ + x_sun, y_sun, vx_sun, vy_sun = planet_orbit_to_sun_state( + orbit_planet, planet_x, planet_y, planet_vx, planet_vy, t, + ) + return orbital_elements_from_cartesian( + x_sun, y_sun, vx_sun, vy_sun, mu_sun, t, + ) diff --git a/orbit_drawer.py b/orbit_drawer.py new file mode 100644 index 0000000..323c398 --- /dev/null +++ b/orbit_drawer.py @@ -0,0 +1,77 @@ +""" +Orbit trajectory point generation for rendering. + +generate_orbit_points() — full orbit (elliptical) or truncated + hyperbolic path within a given cut-off radius. +""" + +import math +from typing import List, Optional, Tuple + +from orbital_elements import OrbitalElements, TAU + + +def generate_orbit_points( + orbit: OrbitalElements, + num_points: int = 200, + max_radius: Optional[float] = None, +) -> List[Tuple[float, float]]: + """ + Generate *num_points* world-space positions along the orbit. + + Elliptical orbits → full closed loop, equally spaced in E. + Hyperbolic orbits → samples H ∈ [−H_max, +H_max], truncated + at *max_radius* (e.g. the planet's SOI edge). + + Parameters + ---------- + max_radius : float or None + Only used for hyperbolic orbits. Points with r > max_radius + are omitted. If None, samples out to r ≈ 5·|a| from focus. + """ + if orbit.is_hyperbolic: + return _hyperbolic_points(orbit, num_points, max_radius) + return _elliptical_points(orbit, num_points) + + +def _elliptical_points( + orbit: OrbitalElements, num_points: int, +) -> List[Tuple[float, float]]: + points = [] + for i in range(num_points): + E = TAU * i / num_points + points.append(orbit.position_from_E(E)) + return points + + +def _hyperbolic_points( + orbit: OrbitalElements, + num_points: int, + max_radius: Optional[float], +) -> List[Tuple[float, float]]: + """ + Sample the hyperbolic branch that passes through periapsis. + + Truncates at *max_radius* by solving for the H values where + r = max_radius, then sampling uniformly between those limits. + """ + r_limit = max_radius if max_radius is not None else 5.0 * abs(orbit.a) + + # Find H where r = r_limit. + # r = |a|(e·cosh(H) − 1) → cosh(H) = (r/|a| + 1) / e + abs_a = abs(orbit.a) + cosh_Hmax = (r_limit / abs_a + 1.0) / orbit.e + if cosh_Hmax < 1.0: + return [] # entire orbit is outside the radius + H_max = math.acosh(cosh_Hmax) + + points = [] + for i in range(num_points): + frac = -1.0 + 2.0 * i / (num_points - 1) if num_points > 1 else 0.0 + H = frac * H_max + x, y = orbit.position_from_H(H) + r = math.hypot(x, y) + if r <= r_limit + 1e-9: + points.append((x, y)) + + return points diff --git a/orbital_elements.py b/orbital_elements.py new file mode 100644 index 0000000..0334295 --- /dev/null +++ b/orbital_elements.py @@ -0,0 +1,303 @@ +""" +2D Keplerian orbital elements — elliptical and hyperbolic. + +Provides OrbitalElements dataclass with cached derived quantities for +efficient position/velocity computation, plus cartesian ↔ orbital +conversion. + +Supports both elliptical (e < 1, a > 0) and hyperbolic (e > 1, a < 0) +orbits using the standard Keplerian formulation. +""" + +import math +from dataclasses import dataclass +from typing import Tuple + +TAU = 2.0 * math.pi +KEPLER_TOL = 1e-12 +KEPLER_MAX_ITER = 50 + + +@dataclass +class OrbitalElements: + """ + 2D Keplerian orbital elements. + + a — semi-major axis. a > 0 for ellipses, a < 0 for hyperbolas. + e — eccentricity. e ≥ 0. e < 1 = ellipse, e > 1 = hyperbola. + omega — argument of periapsis (rad), +x → periapsis. + M0 — mean anomaly at epoch (rad). Wrapped to [0, 2π) for ellipses; + unbounded (can be any real) for hyperbolas. + epoch — reference time. + mu — gravitational parameter of central body. + """ + a: float + e: float + omega: float + M0: float + epoch: float + mu: float + + def __post_init__(self) -> None: + if self.e < 0.0: + raise ValueError("Eccentricity must be ≥ 0") + + self._hyperbolic: bool = self.e > 1.0 + + if self._hyperbolic: + if self.a > 0: + raise ValueError("Hyperbolic orbits require a < 0") + self._abs_a = abs(self.a) + self._p = self._abs_a * (self.e * self.e - 1.0) # semi-latus rectum + self._sqrt_mu_over_p = math.sqrt(self.mu / self._p) + self._n = math.sqrt(self.mu / self._abs_a ** 3) # mean motion + self._sqrt_em1 = math.sqrt(self.e * self.e - 1.0) + else: + if self.a <= 0: + raise ValueError("Elliptical orbits require a > 0") + self._p = self.a * (1.0 - self.e * self.e) # semi-latus rectum + self._sqrt_mu_over_p = math.sqrt(self.mu / self._p) + self._n = math.sqrt(self.mu / self.a ** 3) # mean motion + self._sqrt_1me2 = math.sqrt(1.0 - self.e * self.e) + self._sqrt_mu_a = math.sqrt(self.mu * self.a) + + # ── Convenience properties ────────────────────────────────── + + @property + def period(self) -> float: + """Orbital period (s). +inf for parabolas, meaningless for hyperbolas.""" + if self._hyperbolic: + return float("inf") + return TAU / self._n + + @property + def periapsis(self) -> float: + """Distance at closest approach.""" + if self._hyperbolic: + return self._abs_a * (self.e - 1.0) + return self.a * (1.0 - self.e) + + @property + def apoapsis(self) -> float: + """Distance at farthest point (ellipses only).""" + if self._hyperbolic: + return float("inf") + return self.a * (1.0 + self.e) + + @property + def is_hyperbolic(self) -> bool: + return self._hyperbolic + + def radius_range(self) -> Tuple[float, float]: + """(minimum, maximum) distance. Max is inf for hyperbolas.""" + return self.periapsis, self.apoapsis + + # ── Kepler solvers ────────────────────────────────────────── + + def solve_kepler(self, M: float) -> float: + """Solve Kepler's equation. Returns E (ellipse) or H (hyperbola).""" + if self._hyperbolic: + return self._solve_kepler_hyperbolic(M) + return self._solve_kepler_elliptic(M) + + def _solve_kepler_elliptic(self, M: float) -> float: + """M = E - e·sin(E) → E.""" + e = self.e + if e < 1e-14: + return M + if e < 0.8: + E = M + else: + E = M + 0.85 * e * math.sin(M) + for _ in range(KEPLER_MAX_ITER): + dE = (E - e * math.sin(E) - M) / (1.0 - e * math.cos(E)) + E -= dE + if abs(dE) < KEPLER_TOL: + break + return E + + def _solve_kepler_hyperbolic(self, M: float) -> float: + """M = e·sinh(H) − H → H.""" + e = self.e + if abs(M) < 1e-14: + return 0.0 + # Robust initial guess + H = math.asinh(M / e) + for _ in range(KEPLER_MAX_ITER): + dH = (e * math.sinh(H) - H - M) / (e * math.cosh(H) - 1.0) + H -= dH + if abs(dH) < KEPLER_TOL: + break + return H + + # ── Mean anomaly at time ──────────────────────────────────── + + def mean_anomaly_at(self, t: float) -> float: + """Mean anomaly M at time t.""" + M = self.M0 + self._n * (t - self.epoch) + if self._hyperbolic: + return M # unbounded — can be any real + return M % TAU # wrap to [0, 2π) + + # ── Position / state computation ──────────────────────────── + + def compute_state_at(self, t: float) -> Tuple[float, float, float, float]: + """World-space (x, y, vx, vy) at time t.""" + M = self.mean_anomaly_at(t) + anomaly = self.solve_kepler(M) # E for ellipse, H for hyperbola + if self._hyperbolic: + return self._state_from_H(anomaly) + return self._state_from_E(anomaly) + + def position_at(self, t: float) -> Tuple[float, float]: + """World-space (x, y) at time t.""" + x, y, _, _ = self.compute_state_at(t) + return x, y + + # ── Elliptical: state from eccentric anomaly E ───────────── + + def _state_from_E(self, E: float) -> Tuple[float, float, float, float]: + cosE = math.cos(E) + sinE = math.sin(E) + + # Perifocal position + x_pf = self.a * (cosE - self.e) + y_pf = self.a * self._sqrt_1me2 * sinE + + # Perifocal velocity + inv_denom = 1.0 / (1.0 - self.e * cosE) + sqrt_mu_over_a = math.sqrt(self.mu / self.a) + vx_pf = -sqrt_mu_over_a * sinE * inv_denom + vy_pf = sqrt_mu_over_a * self._sqrt_1me2 * cosE * inv_denom + + return self._rotate_to_world(x_pf, y_pf, vx_pf, vy_pf) + + def position_from_E(self, E: float) -> Tuple[float, float]: + cosE = math.cos(E) + sinE = math.sin(E) + x_pf = self.a * (cosE - self.e) + y_pf = self.a * self._sqrt_1me2 * sinE + cos_om = math.cos(self.omega) + sin_om = math.sin(self.omega) + return ( + x_pf * cos_om - y_pf * sin_om, + x_pf * sin_om + y_pf * cos_om, + ) + + # ── Hyperbolic: state from hyperbolic anomaly H ──────────── + + def _state_from_H(self, H: float) -> Tuple[float, float, float, float]: + coshH = math.cosh(H) + sinhH = math.sinh(H) + + # Perifocal position x = |a|(e − coshH), y = |a|√(e²−1)·sinhH + x_pf = self._abs_a * (self.e - coshH) + y_pf = self._abs_a * self._sqrt_em1 * sinhH + + # Distance r = |a|(e·coshH − 1) + r = self._abs_a * (self.e * coshH - 1.0) + + # True anomaly trig + cos_nu = x_pf / r + sin_nu = y_pf / r + + # Perifocal velocity (unified ν-formula, works for both orbit types) + vx_pf = -self._sqrt_mu_over_p * sin_nu + vy_pf = self._sqrt_mu_over_p * (cos_nu + self.e) + + return self._rotate_to_world(x_pf, y_pf, vx_pf, vy_pf) + + def position_from_H(self, H: float) -> Tuple[float, float]: + coshH = math.cosh(H) + sinhH = math.sinh(H) + x_pf = self._abs_a * (self.e - coshH) + y_pf = self._abs_a * self._sqrt_em1 * sinhH + cos_om = math.cos(self.omega) + sin_om = math.sin(self.omega) + return ( + x_pf * cos_om - y_pf * sin_om, + x_pf * sin_om + y_pf * cos_om, + ) + + # ── Rotate perifocal → world frame ───────────────────────── + + def _rotate_to_world( + self, x_pf: float, y_pf: float, vx_pf: float, vy_pf: float, + ) -> Tuple[float, float, float, float]: + cos_om = math.cos(self.omega) + sin_om = math.sin(self.omega) + return ( + x_pf * cos_om - y_pf * sin_om, + x_pf * sin_om + y_pf * cos_om, + vx_pf * cos_om - vy_pf * sin_om, + vx_pf * sin_om + vy_pf * cos_om, + ) + + +# ── Cartesian → Orbital Elements ───────────────────────────────── + + +def orbital_elements_from_cartesian( + x: float, + y: float, + vx: float, + vy: float, + mu: float, + epoch: float = 0.0, +) -> OrbitalElements: + """ + Convert (x, y, vx, vy) to Keplerian orbital elements. + + Handles both elliptical (energy < 0) and hyperbolic (energy > 0) + trajectories. + """ + r = math.hypot(x, y) + v_sq = vx * vx + vy * vy + energy = 0.5 * v_sq - mu / r + r_dot_v = x * vx + y * vy + + # Semi-major axis (negative for hyperbolas) + if abs(energy) < 1e-15: + raise ValueError("Parabolic trajectory not supported (energy ≈ 0)") + a = -mu / (2.0 * energy) # negative for hyperbolic, positive for elliptical + + # Eccentricity vector + factor = v_sq / mu - 1.0 / r + ex = factor * x - r_dot_v * vx / mu + ey = factor * y - r_dot_v * vy / mu + e = math.hypot(ex, ey) + + # Argument of periapsis + omega = math.atan2(ey, ex) % TAU + + # True anomaly + if e > 1e-14: + cos_nu = (ex * x + ey * y) / (e * r) + cos_nu = max(-1.0, min(1.0, cos_nu)) + nu = math.acos(cos_nu) + if r_dot_v < 0: + nu = -nu + else: + nu = 0.0 + + # Mean anomaly + if e < 1.0: # elliptical + cos_E = (e + math.cos(nu)) / (1.0 + e * math.cos(nu)) + cos_E = max(-1.0, min(1.0, cos_E)) + sin_E = math.sqrt(max(0.0, 1.0 - cos_E * cos_E)) + if nu < 0: + sin_E = -sin_E + E = math.atan2(sin_E, cos_E) + M0 = (E - e * math.sin(E)) % TAU + else: # hyperbolic + # cosh(H) = (e + cos(ν)) / (1 + e·cos(ν)) + cosh_H = (e + math.cos(nu)) / (1.0 + e * math.cos(nu)) + cosh_H = max(1.0, cosh_H) + H = math.acosh(cosh_H) + if nu < 0: + H = -H + M0 = e * math.sinh(H) - H # unbounded, no wrap + + return OrbitalElements( + a=a, e=e, omega=omega, M0=M0, epoch=epoch, mu=mu, + ) diff --git a/soi_calculator.py b/soi_calculator.py new file mode 100644 index 0000000..580c9b2 --- /dev/null +++ b/soi_calculator.py @@ -0,0 +1,197 @@ +""" +SOI (Sphere of Influence) crossing-time calculator. + +Given two bodies in Keplerian orbits around the same parent, +find all times when the distance between them equals a given +SOI radius — i.e., when one body enters or exits the other's +gravitational sphere of influence. + +Algorithm: + 1. Quick radial-overlap check (rejects impossible cases). + 2. Coarse time-scan across the search window. + 3. Newton-Raphson refinement of each sign-change bracket. +""" + +import math +from typing import List, Optional, Tuple + +from orbital_elements import OrbitalElements + + +def find_soi_crossings( + body: OrbitalElements, + planet: OrbitalElements, + soi_radius: float, + window_start: float, + window_end: float, + n_steps: int = 500, +) -> List[Tuple[float, float]]: + """ + Find all (enter_time, exit_time) SOI crossing pairs. + + Parameters + ---------- + body, planet : OrbitalElements + Two bodies orbiting the same central body (same mu). + soi_radius : float + Radius of the planet's sphere of influence. + window_start, window_end : float + Search window in simulation time. + n_steps : int + Number of coarse-scan steps (default 500). + + Returns + ------- + List of (enter, exit) time pairs, sorted chronologically. + Empty list if no crossings found. + """ + # ── Quick rejection: orbital radius ranges must overlap ── + b_min, b_max = body.radius_range() + p_min, p_max = planet.radius_range() + + if b_max + soi_radius < p_min or p_max + soi_radius < b_min: + return [] # orbits never come close enough + + # ── Coarse time scan ──────────────────────────────────── + duration = window_end - window_start + dt = duration / n_steps + + crossings: List[Tuple[float, float]] = [] + prev_dist: Optional[float] = None + enter_bracket: Optional[Tuple[float, float]] = None + + for i in range(n_steps + 1): + t = window_start + i * dt + + bx, by = body.position_at(t) + px, py = planet.position_at(t) + dist = math.hypot(bx - px, by - py) + + if prev_dist is not None: + prev_t = t - dt + was_outside = prev_dist > soi_radius + is_inside = dist <= soi_radius + + if was_outside and is_inside: + # ── Entering SOI — save narrow bracket ── + enter_bracket = (prev_t, t) + + elif not was_outside and not is_inside and enter_bracket is not None: + # ── Exiting SOI — refine both crossings ── + exit_bracket = (prev_t, t) + enter_t = _refine( + body=body, planet=planet, soi_radius=soi_radius, + bracket=enter_bracket, entering=True, + ) + exit_t = _refine( + body=body, planet=planet, soi_radius=soi_radius, + bracket=exit_bracket, entering=False, + ) + crossings.append((enter_t, exit_t)) + enter_bracket = None + + prev_dist = dist + + # ── Still inside SOI at window end? ───────────────────── + if enter_bracket is not None: + t_end = window_end + exit_bracket = (window_end - dt, window_end) + enter_t = _refine( + body=body, planet=planet, soi_radius=soi_radius, + bracket=enter_bracket, entering=True, + ) + exit_t = _refine( + body=body, planet=planet, soi_radius=soi_radius, + bracket=exit_bracket, entering=False, + ) + crossings.append((enter_t, exit_t)) + + return crossings + + +# ── Newton-Raphson refinement ─────────────────────────────────── + + +def _refine( + body: OrbitalElements, + planet: OrbitalElements, + soi_radius: float, + bracket: Tuple[float, float], + entering: bool, +) -> float: + """ + Newton-Raphson refinement of an SOI crossing time. + + Solves |r_body(t) − r_planet(t)| − soi_radius = 0. + + Parameters + ---------- + bracket : (t_low, t_high) + Narrow bracket around the crossing (ideally one scan-step wide). + entering : bool + True → distance *drops* to SOI (outer → inner). + False → distance *rises* to SOI (inner → outer). + """ + t_low, t_high = bracket + + # Start from the side closest to the expected root. + # For entering: outside edge (t_low, where dist > soi). + # For exiting: outside edge (t_high, where dist > soi). + t = t_low if entering else t_high + + newton_tol = 1e-8 + max_iter = 30 + + for _ in range(max_iter): + bx, by, bvx, bvy = body.compute_state_at(t) + px, py, pvx, pvy = planet.compute_state_at(t) + + dx = bx - px + dy = by - py + dvx = bvx - pvx + dvy = bvy - pvy + + dist = math.hypot(dx, dy) + f_val = dist - soi_radius + + if dist < 1e-14: + break # co-located — degenerate case + + fprime = (dx * dvx + dy * dvy) / dist + + if abs(fprime) < 1e-14: + break # stationary w.r.t. SOI boundary + + dt_correction = -f_val / fprime + t += dt_correction + + # Stay inside bracket + t = max(t_low, min(t_high, t)) + + if abs(dt_correction) < newton_tol: + break + + return t + + +# ── Convenience: next SOI event ────────────────────────────────── + + +def find_next_soi_event( + body: OrbitalElements, + planet: OrbitalElements, + soi_radius: float, + from_time: float, + max_orbits: int = 3, +) -> Optional[Tuple[float, float]]: + """ + Find the first (enter, exit) SOI crossing after *from_time*. + + Searches *max_orbits* × the longer of the two orbital periods. + Returns None if no crossing found in that window. + """ + window = max_orbits * max(body.period, planet.period) + crossings = find_soi_crossings( + body, planet, soi_radius, from_time, from_time + window, + ) + return crossings[0] if crossings else None diff --git a/test_orbital.py b/test_orbital.py new file mode 100644 index 0000000..71ee241 --- /dev/null +++ b/test_orbital.py @@ -0,0 +1,346 @@ +""" +Tests for orbital mechanics, SOI detection, and frame transitions. + +Run with: python3 test_orbital.py +or: python3 -m pytest test_orbital.py -v +""" + +import math +import sys +from orbital_elements import OrbitalElements, orbital_elements_from_cartesian, TAU +from soi_calculator import find_soi_crossings +from orbit_drawer import generate_orbit_points +from frame_transition import ( + sun_state_to_planet_orbit, + planet_orbit_to_sun_state, + planet_orbit_to_sun_orbit, +) + +_assertions = 0 +_failures = 0 + + +def check(condition: bool, msg: str = "") -> None: + global _assertions, _failures + _assertions += 1 + if not condition: + _failures += 1 + print(f" FAIL: {msg}", file=sys.stderr) + else: + print(f" ok: {msg}") + + +def approx(a: float, b: float, tol: float = 1e-9) -> bool: + return abs(a - b) < tol + + +# ═══════════════════════════════════════════════════════════════ +# Elliptical orbit tests (unchanged) +# ═══════════════════════════════════════════════════════════════ + +def test_circular_orbit_position() -> None: + print("test_circular_orbit_position") + orbit = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=1_000_000.0) + for t in [0.0, 10.0, 50.0, orbit.period / 4, orbit.period / 2]: + x, y = orbit.position_at(t) + check(approx(math.hypot(x, y), 100.0), f"t={t:.1f} dist={math.hypot(x,y):.6f}") + + +def test_elliptical_periapsis_apoapsis() -> None: + print("test_elliptical_periapsis_apoapsis") + a, e = 200.0, 0.3 + mu = 500_000.0 + orbit = OrbitalElements(a=a, e=e, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + x, y = orbit.position_at(0.0) + check(approx(x, a*(1-e)), f"periapsis x={x} == {a*(1-e)}") + check(approx(y, 0.0), f"periapsis y={y}") + half_T = orbit.period / 2 + x2, y2 = orbit.position_at(half_T) + check(approx(x2, -a*(1+e)), f"apoapsis x={x2} == {-a*(1+e)}") + check(approx(y2, 0.0, 1e-8), f"apoapsis y={y2}") + + +def test_omega_rotation() -> None: + print("test_omega_rotation") + orbit = OrbitalElements(a=100.0, e=0.2, omega=math.pi/2, M0=0.0, epoch=0.0, mu=1_000_000.0) + x, y = orbit.position_at(0.0) + check(approx(x, 0.0, 1e-8), f"omega=pi/2: x={x} ≈ 0") + check(approx(y, 80.0), f"omega=pi/2: y={y} == 80") + + +def test_period() -> None: + print("test_period") + a, mu = 200.0, 1_000_000.0 + orbit = OrbitalElements(a=a, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + expected = 2*math.pi*math.sqrt(a**3/mu) + check(approx(orbit.period, expected), f"period {orbit.period} == {expected}") + x0, y0 = orbit.position_at(0.0) + x1, y1 = orbit.position_at(orbit.period) + check(approx(x0, x1) and approx(y0, y1), "closed after 1 period") + + +def test_cartesian_roundtrip() -> None: + print("test_cartesian_roundtrip") + orbit1 = OrbitalElements(a=300.0, e=0.25, omega=1.2, M0=2.5, epoch=1000.0, mu=1_500_000.0) + x, y, vx, vy = orbit1.compute_state_at(2000.0) + orbit2 = orbital_elements_from_cartesian(x, y, vx, vy, orbit1.mu, epoch=2000.0) + check(approx(orbit2.a, orbit1.a, 1e-6), f"a: {orbit2.a} == {orbit1.a}") + check(approx(orbit2.e, orbit1.e, 1e-6), f"e: {orbit2.e} == {orbit1.e}") + omega_diff = abs(orbit2.omega - orbit1.omega) % TAU + check(omega_diff < 1e-6 or abs(omega_diff-TAU) < 1e-6, f"ω: {orbit2.omega} ≈ {orbit1.omega}") + x2, y2, vx2, vy2 = orbit2.compute_state_at(2000.0) + check(approx(x, x2, 1e-6) and approx(y, y2, 1e-6), "position at epoch") + check(approx(vx, vx2, 1e-6) and approx(vy, vy2, 1e-6), "velocity at epoch") + + +def test_orbit_points() -> None: + print("test_orbit_points") + orbit = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=100_000.0) + points = generate_orbit_points(orbit, 128) + check(len(points) == 128, f"{len(points)} == 128") + for x, y in points: + if not approx(math.hypot(x, y), 100.0, 1e-8): + check(False, f"point at {math.hypot(x,y):.9f} ≠ 100.0") + break + else: + check(True, "all 128 points at correct distance") + + +def test_radial_overlap_quick_reject() -> None: + print("test_radial_overlap_quick_reject") + mu = 1_000_000.0 + inner = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + outer = OrbitalElements(a=500.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + crossings = find_soi_crossings(inner, outer, 30.0, 0.0, 200.0) + check(len(crossings) == 0, f"non-overlapping: {len(crossings)} crossings") + + +# ═══════════════════════════════════════════════════════════════ +# SOI crossing tests +# ═══════════════════════════════════════════════════════════════ + +def test_soi_crossing_detection() -> None: + print("test_soi_crossing_detection") + mu = 1_500_000.0 + planet = OrbitalElements(a=350.0, e=0.02, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + body = OrbitalElements(a=320.0, e=0.25, omega=math.pi*0.7, M0=1.0, epoch=0.0, mu=mu) + soi_r = 50.0 + window = 3*max(body.period, planet.period) + crossings = find_soi_crossings(body, planet, soi_r, 0.0, window) + check(len(crossings) > 0, f"found {len(crossings)} crossing(s)") + for i, (enter_t, exit_t) in enumerate(crossings): + dur = exit_t - enter_t + check(enter_t <= exit_t, f"crossing {i}: {enter_t:.3f} ≤ {exit_t:.3f}") + for label, t in [("enter", enter_t), ("exit", exit_t)]: + bx, by = body.position_at(t) + px, py = planet.position_at(t) + d = math.hypot(bx-px, by-py) + check(approx(d, soi_r, 0.5), f"{label} d={d:.3f} ≈ {soi_r}") + if dur > 0.2: + mid_t = (enter_t+exit_t)/2 + bx, by = body.position_at(mid_t) + px, py = planet.position_at(mid_t) + d_mid = math.hypot(bx-px, by-py) + check(d_mid < soi_r-1.0, f"deep crossing midpoint d={d_mid:.1f}") + else: + print(f" (crossing {i} is a graze, dur={dur:.3f}s)") + + +def test_deep_soi_crossing() -> None: + print("test_deep_soi_crossing") + mu = 1_500_000.0 + planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu) + body = OrbitalElements(a=345.0, e=0.02, omega=0.5, M0=-0.35, epoch=0.0, mu=mu) + soi = 80.0 + window = 5*max(body.period, planet.period) + crossings = find_soi_crossings(body, planet, soi, 0.0, window, n_steps=800) + check(len(crossings) >= 1, f"found {len(crossings)} crossing(s)") + for i, (enter_t, exit_t) in enumerate(crossings): + dur = exit_t-enter_t + check(enter_t < exit_t, f"crossing {i}: {enter_t:.2f} < {exit_t:.2f} dur={dur:.2f}s") + for label, t in [("enter", enter_t), ("exit", exit_t)]: + bx, by = body.position_at(t) + px, py = planet.position_at(t) + d = math.hypot(bx-px, by-py) + check(approx(d, soi, 1.0), f"{label} d={d:.2f} ≈ {soi}") + if dur > 0.5: + mid_t = (enter_t+exit_t)/2 + bx, by = body.position_at(mid_t) + px, py = planet.position_at(mid_t) + d_mid = math.hypot(bx-px, by-py) + check(d_mid < soi-2.0, f"midpoint d={d_mid:.1f} < {soi}") + + +# ═══════════════════════════════════════════════════════════════ +# Hyperbolic orbit tests +# ═══════════════════════════════════════════════════════════════ + +def test_hyperbolic_orbit_position() -> None: + """Body at periapsis at t=0 on a hyperbolic trajectory.""" + print("test_hyperbolic_orbit_position") + mu = 50_000.0 + # a < 0 for hyperbolas, e > 1 + orbit = OrbitalElements(a=-40.0, e=2.5, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + + # At t=0 (H=0), should be at periapsis + x, y = orbit.position_at(0.0) + r_peri = abs(orbit.a)*(orbit.e - 1.0) # = 40*1.5 = 60 + check(approx(x, r_peri), f"hyperbolic periapsis x={x:.3f} == {r_peri}") + check(approx(y, 0.0), f"y={y:.3f} ≈ 0") + + +def test_hyperbolic_velocity() -> None: + """Velocity at periapsis should match the vis-viva equation.""" + print("test_hyperbolic_velocity") + mu = 50_000.0 + orbit = OrbitalElements(a=-40.0, e=2.5, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + _, _, vx, vy = orbit.compute_state_at(0.0) + + r_p = abs(orbit.a)*(orbit.e - 1.0) # 60 + # v_p = sqrt(mu * (2/r_p + 1/|a|)) + v_expected = math.sqrt(mu * (2.0/r_p + 1.0/abs(orbit.a))) + v_actual = math.hypot(vx, vy) + check(approx(v_actual, v_expected, 1e-6), + f"v_peri: {v_actual:.3f} == {v_expected:.3f}") + + +def test_hyperbolic_cartesian_roundtrip() -> None: + """Hyperbolic orbit → cartesian → same hyperbolic orbit.""" + print("test_hyperbolic_cartesian_roundtrip") + mu = 50_000.0 + orbit1 = OrbitalElements(a=-30.0, e=3.0, omega=1.0, M0=0.5, epoch=0.0, mu=mu) + x, y, vx, vy = orbit1.compute_state_at(10.0) + orbit2 = orbital_elements_from_cartesian(x, y, vx, vy, mu, epoch=10.0) + check(approx(orbit2.a, orbit1.a, 1e-6), f"a: {orbit2.a} == {orbit1.a}") + check(approx(orbit2.e, orbit1.e, 1e-6), f"e: {orbit2.e} == {orbit1.e}") + omega_diff = abs(orbit2.omega - orbit1.omega) % TAU + check(omega_diff < 1e-6 or abs(omega_diff-TAU) < 1e-6, f"ω: {orbit2.omega} ≈ {orbit1.omega}") + x2, y2 = orbit2.position_at(10.0) + check(approx(x, x2, 1e-6) and approx(y, y2, 1e-6), "position at epoch") + + +def test_hyperbolic_orbit_points() -> None: + """generate_orbit_points should truncate at max_radius.""" + print("test_hyperbolic_orbit_points") + mu = 50_000.0 + orbit = OrbitalElements(a=-20.0, e=2.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu) + points = generate_orbit_points(orbit, num_points=100, max_radius=80.0) + check(len(points) > 0, f"generated {len(points)} hyperbolic points") + for x, y in points: + d = math.hypot(x, y) + check(d <= 80.0 + 1e-6, f"point at r={d:.3f} ≤ 80.0") + + +# ═══════════════════════════════════════════════════════════════ +# Frame transition tests +# ═══════════════════════════════════════════════════════════════ + +def test_frame_transition_roundtrip() -> None: + """ + Body in sun frame → planet frame → back → should match original. + """ + print("test_frame_transition_roundtrip") + mu_sun = 1_500_000.0 + mu_planet = 30_000.0 + + planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu_sun) + body_sun = OrbitalElements(a=380.0, e=0.25, omega=2.0, M0=1.0, epoch=0.0, mu=mu_sun) + + t_transition = 50.0 + + # Original body position in sun frame + bx, by, bvx, bvy = body_sun.compute_state_at(t_transition) + px, py, pvx, pvy = planet.compute_state_at(t_transition) + + # Transition to planet frame + body_planet = sun_state_to_planet_orbit( + bx, by, bvx, bvy, px, py, pvx, pvy, mu_planet, t_transition, + ) + + # Body's planet-relative position should be the vector difference + bx_rel, by_rel = body_planet.position_at(t_transition) + check(approx(bx_rel, bx-px, 1e-6), f"rel x: {bx_rel:.3f} == {bx-px:.3f}") + check(approx(by_rel, by-py, 1e-6), f"rel y: {by_rel:.3f} == {by-py:.3f}") + + # Transition back to sun frame + x_back, y_back, vx_back, vy_back = planet_orbit_to_sun_state( + body_planet, px, py, pvx, pvy, t_transition, + ) + check(approx(x_back, bx, 1e-6) and approx(y_back, by, 1e-6), "position roundtrip") + check(approx(vx_back, bvx, 1e-6) and approx(vy_back, bvy, 1e-6), "velocity roundtrip") + + +def test_frame_transition_hyperbolic_flyby() -> None: + """ + Body on planet-crossing orbit → planet frame at SOI entry. + Verify the planet-relative orbit is hyperbolic (flyby, not capture). + """ + print("test_frame_transition_hyperbolic_flyby") + mu_sun = 1_500_000.0 + mu_planet = 10_000.0 # v_rel ≈ 20.4, critical ≈ 16.9 → hyperbolic ✓ + + # Planet on near-circular orbit + planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu_sun) + + # Asteroid: crossing orbit (verified to intersect planet's SOI) + asteroid = OrbitalElements(a=380.0, e=0.35, omega=2.5, M0=0.8, epoch=0.0, mu=mu_sun) + + # Find first SOI crossing + soi = 70.0 + window = 4 * max(asteroid.period, planet.period) + crossings = find_soi_crossings(asteroid, planet, soi, 0.0, window, n_steps=1000) + + check(len(crossings) > 0, f"found {len(crossings)} SOI crossing(s)") + + if crossings: + enter_t = crossings[0][0] + + # Get states at SOI entry + ax, ay, avx, avy = asteroid.compute_state_at(enter_t) + px, py, pvx, pvy = planet.compute_state_at(enter_t) + + # Transition to planet frame + body_planet = sun_state_to_planet_orbit( + ax, ay, avx, avy, px, py, pvx, pvy, mu_planet, enter_t, + ) + + # This SHOULD be hyperbolic (flyby, energy > 0) + check(body_planet.is_hyperbolic, + f"planet-relative orbit is hyperbolic (e={body_planet.e:.3f})") + + # Verify position matches relative to planet + bx_rel, by_rel = body_planet.position_at(enter_t) + d_rel = math.hypot(bx_rel, by_rel) + check(approx(d_rel, soi, 2.0), + f"rel distance at SOI entry: {d_rel:.1f} ≈ {soi}") + + # Verify the orbit goes out past the SOI (will exit) + check(body_planet.apoapsis > soi, + f"hyperbolic: path extends past SOI (max r={body_planet.apoapsis})") + + +# ═══════════════════════════════════════════════════════════════ +# Runner +# ═══════════════════════════════════════════════════════════════ + +if __name__ == "__main__": + print("=== Orbital Mechanics Tests ===\n") + + test_circular_orbit_position() + test_elliptical_periapsis_apoapsis() + test_omega_rotation() + test_period() + test_cartesian_roundtrip() + test_orbit_points() + test_radial_overlap_quick_reject() + test_soi_crossing_detection() + test_deep_soi_crossing() + test_hyperbolic_orbit_position() + test_hyperbolic_velocity() + test_hyperbolic_cartesian_roundtrip() + test_hyperbolic_orbit_points() + test_frame_transition_roundtrip() + test_frame_transition_hyperbolic_flyby() + + print(f"\n---\n{_assertions} assertions, {_failures} failures") + sys.exit(0 if _failures == 0 else 1) diff --git a/visual_test.py b/visual_test.py new file mode 100644 index 0000000..8f64164 --- /dev/null +++ b/visual_test.py @@ -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()