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
304 lines
10 KiB
Python
304 lines
10 KiB
Python
"""
|
||
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,
|
||
)
|