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:
303
orbital_elements.py
Normal file
303
orbital_elements.py
Normal file
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user