From 30b8127a55570e4c1e0866f20106844de54e1356 Mon Sep 17 00:00:00 2001 From: cinnaboot Date: Tue, 28 Apr 2026 13:04:37 -0400 Subject: [PATCH] feat: add sim_engine.py and test_orbital_period.py for precalculation - sim_engine.py: generic orbital mechanics simulator with RK4 + analytical propagation, used to verify C++ test expected values - test_orbital_period.py: precalculation script with safety timeout for measuring Earth/Mars orbital periods and direction test values - Remove SOI transition logic (buggy in both Python and C++ versions) - Use dataclasses.replace() for immutable state updates - Support single-line TOML inline tables via tomllib --- scripts/sim_engine.py | 555 +++++++++++++++++++++++++++++++++ scripts/test_orbital_period.py | 117 +++++++ 2 files changed, 672 insertions(+) create mode 100644 scripts/sim_engine.py create mode 100644 scripts/test_orbital_period.py diff --git a/scripts/sim_engine.py b/scripts/sim_engine.py new file mode 100644 index 0000000..6c1704a --- /dev/null +++ b/scripts/sim_engine.py @@ -0,0 +1,555 @@ +#!/usr/bin/env python3 +""" +Generic orbital mechanics simulation engine. +Replicates the exact physics from src/orbital_mechanics.cpp and src/simulation.cpp. + +Usage: + from sim_engine import Simulator + sim = Simulator("path/to/config.toml", dt=60.0) + sim.run(steps=1000) + for event in sim.events: + print(event) +""" + +import math +import tomllib +from dataclasses import dataclass, field, replace +from typing import Dict, Tuple, Any + + +# ============================================================================= +# Constants +# ============================================================================= + +G = 6.67430e-11 +PARABOLIC_TOLERANCE = 1e-3 +KEPLER_TOLERANCE = 1e-10 +KEPLER_MAX_ITER = 50 +VEL_DRIFT_THRESHOLD = 1e-6 # m/s + + +# ============================================================================= +# Vector operations +# ============================================================================= + +def vadd(a, b): + return (a[0]+b[0], a[1]+b[1], a[2]+b[2]) + +def vsub(a, b): + return (a[0]-b[0], a[1]-b[1], a[2]-b[2]) + +def vscale(v, s): + return (v[0]*s, v[1]*s, v[2]*s) + +def vmag(v): + return math.sqrt(v[0]**2 + v[1]**2 + v[2]**2) + +def vdot(a, b): + return a[0]*b[0] + a[1]*b[1] + a[2]*b[2] + +def vcross(a, b): + return ( + a[1]*b[2] - a[2]*b[1], + a[2]*b[0] - a[0]*b[2], + a[0]*b[1] - a[1]*b[0] + ) + +def vnorm(v): + m = vmag(v) + if m < 1e-15: + return (0.0, 0.0, 0.0) + return (v[0]/m, v[1]/m, v[2]/m) + +def normalize_angle(angle): + while angle < 0.0: + angle += 2.0 * math.pi + while angle >= 2.0 * math.pi: + angle -= 2.0 * math.pi + return angle + + +# ============================================================================= +# Data structures +# ============================================================================= + +@dataclass +class OrbitalElements: + a: float = 0.0 # semi-major axis (elliptical) / semi-latus rectum (parabolic) + e: float = 0.0 # eccentricity + nu: float = 0.0 # true anomaly + inc: float = 0.0 # inclination + Omega: float = 0.0 # longitude of ascending node + omega: float = 0.0 # argument of periapsis + p: float = 0.0 # semi-latus rectum (parabolic only) + + +@dataclass +class Body: + name: str = "" + mass: float = 0.0 + radius: float = 0.0 + parent_index: int = -1 + orbit: OrbitalElements = field(default_factory=OrbitalElements) + local_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0) + local_vel: Tuple[float, float, float] = (0.0, 0.0, 0.0) + global_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0) + global_vel: Tuple[float, float, float] = (0.0, 0.0, 0.0) + + +@dataclass +class Event: + """Recorded simulation event.""" + kind: str = "state" + time: float = 0.0 + data: Dict[str, Any] = field(default_factory=dict) + + +# ============================================================================= +# Kepler equation solvers (exact C++ logic) +# ============================================================================= + +def get_initial_trial_value(mean_anomaly, eccentricity): + """Initial guess for Kepler solver (C++ get_initial_trial_value).""" + return (mean_anomaly + eccentricity * math.sin(mean_anomaly) + + ((eccentricity ** 2 / 2.0) * math.sin(2.0 * mean_anomaly))) + + +def solve_kepler_elliptical(mean_anomaly, eccentricity): + E = get_initial_trial_value(mean_anomaly, eccentricity) + E_prev = E + 2.0 * KEPLER_TOLERANCE + for _ in range(KEPLER_MAX_ITER): + if abs(E - E_prev) < KEPLER_TOLERANCE: + break + E_prev = E + sin_E = math.sin(E) + E = E - (E - eccentricity * sin_E - mean_anomaly) / (1.0 - eccentricity * math.cos(E)) + return E + + +# ============================================================================= +# Coordinate transforms +# ============================================================================= + +def orbital_to_cartesian(elements, parent_mass): + """Convert orbital elements to local position/velocity vectors.""" + mu = G * parent_mass + a = elements.a + e = elements.e + nu = elements.nu + + if abs(e - 1.0) < PARABOLIC_TOLERANCE: + p = elements.p + else: + p = a * (1.0 - e * e) + + r = p / (1.0 + e * math.cos(nu)) + + x_orb = r * math.cos(nu) + y_orb = r * math.sin(nu) + + vx_orb = -math.sqrt(mu / p) * math.sin(nu) + vy_orb = math.sqrt(mu / p) * (e + math.cos(nu)) + + # z-x-z rotation: Rz(Omega) * Rx(inc) * Rz(omega) + cos_w, sin_w = math.cos(elements.omega), math.sin(elements.omega) + x1 = x_orb * cos_w - y_orb * sin_w + y1 = x_orb * sin_w + y_orb * cos_w + + cos_i, sin_i = math.cos(elements.inc), math.sin(elements.inc) + x2 = x1 + y2 = y1 * cos_i + z2 = y1 * sin_i + + cos_O, sin_O = math.cos(elements.Omega), math.sin(elements.Omega) + pos = (x2 * cos_O - y2 * sin_O, + x2 * sin_O + y2 * cos_O, + z2) + + vx1 = vx_orb * cos_w - vy_orb * sin_w + vy1 = vx_orb * sin_w + vy_orb * cos_w + vx2 = vx1 + vy2 = vy1 * cos_i + vz2 = vy1 * sin_i + + vel = (vx2 * cos_O - vy2 * sin_O, + vx2 * sin_O + vy2 * cos_O, + vz2) + + return pos, vel + + +def cartesian_to_orbital_elements(pos, vel, parent_mass): + """Convert local position/velocity to orbital elements.""" + mu = G * parent_mass + r = vmag(pos) + v = vmag(vel) + v_sq = v * v + + specific_energy = -mu / r + v_sq / 2.0 + h_vec = vcross(pos, vel) + h = vmag(h_vec) + + # Eccentricity vector: e_vec = (v² - μ/r)r - (r·v)v all divided by μ + r_dot_v = vdot(pos, vel) + e_vec = ((v_sq - mu / r) * pos[0] - r_dot_v * vel[0]) / mu, \ + ((v_sq - mu / r) * pos[1] - r_dot_v * vel[1]) / mu, \ + ((v_sq - mu / r) * pos[2] - r_dot_v * vel[2]) / mu + e = vmag(e_vec) + + # Semi-major axis + if abs(specific_energy) < 1e-10: + a = 1e10 + else: + a = -mu / (2.0 * specific_energy) + + # True anomaly + if e < 1e-10: + # Nearly circular: use argument of latitude + n_vec = vcross((0.0, 0.0, 1.0), h_vec) + n_mag = vmag(n_vec) + sin_i = (n_mag / h) if h > 1e-10 else 1.0 + if sin_i > 1e-6 and n_mag > 1e-10: + # Well-defined ascending node: compute argument of latitude + x_AN = n_vec[0] / n_mag + y_AN = n_vec[1] / n_mag + hcn = vcross(h_vec, n_vec) + hcn_mag = vmag(hcn) + if hcn_mag > 1e-10: + hcn = vscale(hcn, 1.0 / hcn_mag) + r_xAN = pos[0] * x_AN + pos[1] * y_AN + r_yAN = pos[0] * hcn[0] + pos[1] * hcn[1] + pos[2] * hcn[2] + nu = math.atan2(r_yAN, r_xAN) + else: + # Nearly coplanar: use atan2(y, x) as argument of latitude + nu = math.atan2(pos[1], pos[0]) + nu = normalize_angle(nu) + else: + cos_nu = vdot(pos, e_vec) / (r * e) + cos_nu = max(-1.0, min(1.0, cos_nu)) + sin_nu = None + + if abs(cos_nu) > 1.0 - 1e-10: + h_cross_e = vcross(h_vec, e_vec) + denom = r * e * h + sin_nu = vdot(pos, h_cross_e) / denom if denom > 1e-10 else 0.0 + else: + r_cross_h = vcross(pos, h_vec) + denom = r * e * h + sin_nu = vdot(r_cross_h, e_vec) / denom if denom > 1e-10 else 0.0 + + nu = math.atan2(sin_nu, cos_nu) + if nu == -math.pi: + nu = math.pi + nu = normalize_angle(nu) + + # Inclination + i = math.acos(h_vec[2] / h) if h > 1e-10 else 0.0 + + # RAAN + n_vec = vcross((0.0, 0.0, 1.0), h_vec) + n_mag = vmag(n_vec) + if n_mag > 1e-10: + Omega = math.acos(n_vec[0] / n_mag) + if n_vec[1] < 0.0: + Omega = 2.0 * math.pi - Omega + else: + Omega = 0.0 + + # Argument of periapsis + inclination_threshold = 0.01 + if e > 1e-10 and n_mag > 1e-10 and i > inclination_threshold: + cos_omega = vdot(e_vec, n_vec) / (e * n_mag) + n_cross_e = vcross(n_vec, e_vec) + sin_omega = vdot(n_cross_e, h_vec) / (e * n_mag * h) + omega = math.atan2(sin_omega, cos_omega) + if omega < 0.0: + omega += 2.0 * math.pi + elif e > 1e-10: + omega = math.atan2(e_vec[1], e_vec[0]) + if omega < 0.0: + omega += 2.0 * math.pi + else: + omega = 0.0 + + elements = OrbitalElements() + if abs(e - 1.0) < 1e-3: + elements.p = (h * h) / mu + else: + elements.a = a + elements.e = e + elements.nu = nu + elements.inc = i + elements.Omega = Omega + elements.omega = omega + + return elements + + +# ============================================================================= +# Propagation +# ============================================================================= + +def propagate(elements, dt, parent_mass): + """Propagate orbital elements forward by dt. Returns new elements.""" + mu = G * parent_mass + a = elements.a + e = elements.e + nu = elements.nu + + if abs(e - 1.0) < PARABOLIC_TOLERANCE: + # Parabolic (Barker's equation) + p = elements.p + D = math.tan(nu / 2.0) + M = D + (D * D * D) / 3.0 + n = math.sqrt(mu / (p ** 3.0)) + M = M + n * dt + # Solve Barker's: D + D^3/3 = M + c = 1.5 * M + disc = c * c + 1.0 + sqrt_disc = math.sqrt(disc) + D_new = math.cbrt(c + sqrt_disc) + math.cbrt(c - sqrt_disc) + return replace(elements, nu=2.0 * math.atan(D_new)) + + elif e < 1.0: + # Elliptical + n = math.sqrt(mu / (a ** 3.0)) + E = 2.0 * math.atan(math.sqrt((1.0 - e) / (1.0 + e)) * math.tan(nu / 2.0)) + M = E - e * math.sin(E) + M = M + n * dt + E_new = get_initial_trial_value(M, e) + E_prev = E_new + 2.0 * KEPLER_TOLERANCE + for _ in range(KEPLER_MAX_ITER): + if abs(E_new - E_prev) < KEPLER_TOLERANCE: + break + E_prev = E_new + sin_E = math.sin(E_new) + E_new = E_new - (E_new - e * sin_E - M) / (1.0 - e * math.cos(E_new)) + nu_new = 2.0 * math.atan(math.sqrt((1.0 + e) / (1.0 - e)) * math.tan(E_new / 2.0)) + return replace(elements, nu=nu_new) + + else: + # Hyperbolic + raise NotImplementedError("hyperbolic propagation not yet implemented") + + +# ============================================================================= +# Global coordinate computation +# ============================================================================= + +def compute_global_coordinates(bodies): + """ + Compute global position/velocity for all bodies. + Matches C++ compute_global_coordinates() exactly. + """ + for body in bodies: + if body.parent_index == -1: + body.global_pos = body.local_pos + body.global_vel = body.local_vel + elif 0 <= body.parent_index < len(bodies): + parent = bodies[body.parent_index] + body.global_pos = vadd(body.local_pos, parent.global_pos) + body.global_vel = vadd(body.local_vel, parent.global_vel) + + +# ============================================================================= +# Velocity drift check +# ============================================================================= + +def check_velocity_drift(body, parent, parent_mass): + """ + Check if local velocity has drifted from expected Keplerian velocity. + If so, reconstruct orbital elements from current state. + Matches C++ update_bodies_physics() drift check. + """ + if parent is None: + return + + _, expected_vel = orbital_to_cartesian(body.orbit, parent_mass) + vel_diff = vmag(vsub(body.local_vel, expected_vel)) + if vel_diff > VEL_DRIFT_THRESHOLD: + body.orbit = cartesian_to_orbital_elements(body.local_pos, body.local_vel, parent_mass) + + +# ============================================================================= +# Body physics update +# ============================================================================= + +def update_body(bodies, body_index, dt): + """ + Update a single body: drift check, propagation. + Matches C++ update_bodies_physics() per-body logic (without SOI). + """ + body = bodies[body_index] + + if body.parent_index == -1: + return # Root body doesn't propagate + + if 0 <= body.parent_index < len(bodies): + parent = bodies[body.parent_index] + check_velocity_drift(body, parent, parent.mass) + body.orbit = propagate(body.orbit, dt, parent.mass) + body.local_pos, body.local_vel = orbital_to_cartesian(body.orbit, parent.mass) + + +# ============================================================================= +# TOML config loader +# ============================================================================= + +def load_config(config_path): + """Load a TOML 1.0 config file and return parsed data.""" + with open(config_path, "rb") as f: + return tomllib.load(f) + + +def bodies_from_config(config): + """ + Create Body objects from TOML config. + Parent references are resolved by name, then by index. + """ + bodies = [] + name_to_idx = {} + + # First pass: create bodies without positions + for body_cfg in config.get("bodies", []): + orbit_cfg = body_cfg.get("orbit", {}) + elements = OrbitalElements( + a=orbit_cfg.get("semi_major_axis", 0.0), + e=orbit_cfg.get("eccentricity", 0.0), + nu=orbit_cfg.get("true_anomaly", 0.0), + inc=orbit_cfg.get("inclination", 0.0), + Omega=orbit_cfg.get("longitude_of_ascending_node", 0.0), + omega=orbit_cfg.get("argument_of_periapsis", 0.0), + ) + + parent_ref = body_cfg.get("parent_index", -1) + if isinstance(parent_ref, str): + # Resolve by name + if parent_ref in name_to_idx: + parent_index = name_to_idx[parent_ref] + elif parent_ref == "Sun" or parent_ref == "root" or parent_ref == "-1": + parent_index = -1 + else: + raise ValueError(f"Unknown parent name: {parent_ref}") + else: + parent_index = int(parent_ref) + + body = Body( + name=body_cfg.get("name", f"Body_{len(bodies)}"), + mass=body_cfg.get("mass", 0.0), + radius=body_cfg.get("radius", 0.0), + parent_index=parent_index, + orbit=elements, + ) + bodies.append(body) + name_to_idx[body.name] = len(bodies) - 1 + + return bodies + + +# ============================================================================= +# Initialization +# ============================================================================= + +def initialize_bodies(bodies): + """ + Initialize orbital objects from orbital elements. + Matches C++ initialize_orbital_objects() exactly (without SOI). + """ + for i, body in enumerate(bodies): + if body.parent_index >= 0 and body.parent_index < len(bodies): + parent = bodies[body.parent_index] + local_pos, local_vel = orbital_to_cartesian(body.orbit, parent.mass) + body.local_pos = local_pos + body.local_vel = local_vel + body.global_pos = vadd(parent.global_pos, local_pos) + body.global_vel = vadd(parent.global_vel, local_vel) + else: + body.local_pos = (0.0, 0.0, 0.0) + body.local_vel = (0.0, 0.0, 0.0) + body.global_pos = (0.0, 0.0, 0.0) + body.global_vel = (0.0, 0.0, 0.0) + + +# ============================================================================= +# Simulator — public API +# ============================================================================= + +class Simulator: + """ + Generic orbital mechanics simulator. + + Usage: + sim = Simulator("config.toml", dt=60.0) + sim.run(steps=1000) + + # Access results + for event in sim.events: + print(event) + + # Access final state + for body in sim.bodies: + print(f"{body.name}: r={vmag(body.global_pos):.0f} m") + """ + + def __init__(self, config_path, dt=60.0): + self.dt = dt + self.time = 0.0 + self.events = [] + self._body_count = 0 + + config = load_config(config_path) + self.bodies = bodies_from_config(config) + initialize_bodies(self.bodies) + self._body_count = len(self.bodies) + + def run(self, steps): + """Run simulation for the given number of timesteps.""" + for _ in range(steps): + self._step() + + def _step(self): + """Single simulation step. Matches C++ update_simulation() order.""" + # 1. Update body physics (drift, propagation) + for i in range(self._body_count): + update_body(self.bodies, i, self.dt) + + # 2. Compute global coordinates + compute_global_coordinates(self.bodies) + + self.time += self.dt + + def record_state(self, label=""): + """Record current simulation state as an event.""" + state = {} + for body in self.bodies: + r = vmag(body.global_pos) + state[body.name] = { + "r": r, + "nu": body.orbit.nu, + "a": body.orbit.a, + "e": body.orbit.e, + "parent": body.parent_index, + "parent_name": self.bodies[body.parent_index].name if body.parent_index >= 0 else "root", + } + self.events.append(Event(kind="state", time=self.time, data={"label": label, "state": state})) + + def get_body(self, name_or_index): + """Get a body by name or index.""" + if isinstance(name_or_index, int): + return self.bodies[name_or_index] + for body in self.bodies: + if body.name == name_or_index: + return body + raise KeyError(f"Body not found: {name_or_index}") + + def print_summary(self): + """Print a summary of all recorded state events.""" + for event in self.events: + label = event.data.get("label", "") + if label: + print(f"\n*** {label} (t={event.time:.1f}s) ***") + for name, info in event.data.get("state", {}).items(): + print(f" {name}: r={info['r']:.0f} m, " + f"nu={math.degrees(info['nu']):.1f}°, " + f"a={info['a']:.0f}, e={info['e']:.6f}, " + f"parent={info['parent_name']}") diff --git a/scripts/test_orbital_period.py b/scripts/test_orbital_period.py new file mode 100644 index 0000000..f099a63 --- /dev/null +++ b/scripts/test_orbital_period.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +""" +Precalculate expected values for test_orbital_period.cpp. + +Measures: +1. Earth orbital period (seconds, days) — track global angle for circular orbit +2. Mars orbital period (seconds, days) +3. Direction test: prograde check over 1 day +""" + +import sys +import math + +sys.path.insert(0, "scripts") +from sim_engine import Simulator, vmag, G, OrbitalElements, propagate + +MAX_STEPS = 1_100_000 # safety limit (687 days × 1440 steps/day) +DT = 60.0 + + +def measure_period(sim, body_name, parent_mass, analytical_days): + """ + Measure period by tracking global angle for one full revolution. + For circular orbits, nu stays at 0 so we track atan2(y, x) instead. + """ + body = sim.get_body(body_name) + parent = sim.get_body(body.parent_index) if body.parent_index >= 0 else None + + # Track global angle + if parent: + angle_start = math.atan2( + body.global_pos[1] - parent.global_pos[1], + body.global_pos[0] - parent.global_pos[0] + ) + else: + angle_start = math.atan2(body.global_pos[1], body.global_pos[0]) + + total_angle = 0.0 + prev_angle = angle_start + + for step in range(1, MAX_STEPS + 1): + sim._step() + + if parent: + angle = math.atan2( + body.global_pos[1] - parent.global_pos[1], + body.global_pos[0] - parent.global_pos[0] + ) + else: + angle = math.atan2(body.global_pos[1], body.global_pos[0]) + + # Accumulate angle (handle wrap) + delta = angle - prev_angle + if delta > math.pi: + delta -= 2 * math.pi + elif delta < -math.pi: + delta += 2 * math.pi + total_angle += delta + prev_angle = angle + + if total_angle >= 2 * math.pi: + break + + if step >= MAX_STEPS: + print(f" TIMEOUT after {MAX_STEPS} steps ({sim.time/86400:.1f} days)") + return None + + period_s = sim.time + period_days = period_s / 86400.0 + print(f" Measured: {period_s:.1f}s = {period_days:.4f} days") + print(f" Analytical: {analytical_days:.4f} days") + print(f" Error: {abs(period_days - analytical_days):.4f} days ({abs(period_days - analytical_days)/analytical_days*100:.4f}%)") + print(f" e after: {body.orbit.e:.15f}") + + return period_days + + +def main(): + print("=== Earth Period ===") + sim = Simulator("tests/test_orbital_period.toml", dt=DT) + earth_a = 1.496e11 + earth_mu = G * 1.989e30 # Sun mass + earth_analytical = 2.0 * math.pi * math.sqrt(earth_a**3 / earth_mu) / 86400.0 + measure_period(sim, "Earth", 1.989e30, earth_analytical) + + print("\n=== Mars Period ===") + sim = Simulator("tests/test_orbital_period.toml", dt=DT) + mars_a = 2.244e11 + mars_mu = G * 1.989e30 # Sun mass + mars_analytical = 2.0 * math.pi * math.sqrt(mars_a**3 / mars_mu) / 86400.0 + measure_period(sim, "Mars", 1.989e30, mars_analytical) + + print("\n=== Direction Test (1 day) ===") + sim = Simulator("tests/test_orbital_period.toml", dt=DT) + earth = sim.get_body("Earth") + sun = sim.get_body("Sun") + theta_start = math.atan2(earth.global_pos[1] - sun.global_pos[1], + earth.global_pos[0] - sun.global_pos[0]) + + sim.run(steps=1440) # 1 day = 86400s / 60s + + theta_end = math.atan2(earth.global_pos[1] - sun.global_pos[1], + earth.global_pos[0] - sun.global_pos[0]) + delta = theta_end - theta_start + print(f" theta_start: {theta_start:.10f} rad") + print(f" theta_end: {theta_end:.10f} rad") + print(f" delta: {delta:.10f} rad") + print(f" prograde: {delta > 0}") + + # Expected delta for 1 day of Earth orbit + expected_delta = math.sqrt(earth_mu / earth_a**3) * 86400.0 + print(f" expected: {expected_delta:.10f} rad") + print(f" error: {abs(delta - expected_delta):.10f} rad") + + +if __name__ == "__main__": + main()