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
347 lines
15 KiB
Python
347 lines
15 KiB
Python
"""
|
|
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)
|