diff --git a/src/layup/iod.py b/src/layup/iod.py index 970ef505..0ff5677d 100644 --- a/src/layup/iod.py +++ b/src/layup/iod.py @@ -36,7 +36,10 @@ from typing import Callable, Optional, Sequence from layup.constants import GMtotal, SPEED_OF_LIGHT -from layup.routines import FitResult, Observation, gauss +from layup.routines import FitResult, Observation, gauss, get_ephem +from layup.utilities.herget_iod import herget_with_assist + +from layup.orbit_maths import build_ephem_and_mus logger = logging.getLogger(__name__) @@ -110,6 +113,16 @@ def gauss_iod(observations, seq): register_iod("gauss", gauss_iod) +def herget_iod(observations, seq): + """""" + ephem, _, _ = build_ephem_and_mus() + solns = herget_with_assist(observations, seq, ephem, tolerance=0.0001, max_iterations=100) + return solns + + +register_iod("herget", herget_iod) + + # ----------------------------------------------------------------------- # # Candidate filter (held-out angular residual). # # ----------------------------------------------------------------------- # diff --git a/src/layup/orbitfit.py b/src/layup/orbitfit.py index c35dbb5f..10c405c3 100644 --- a/src/layup/orbitfit.py +++ b/src/layup/orbitfit.py @@ -1289,7 +1289,7 @@ def _orbitfit( # Perform the orbit fitting if initial_guess is None or initial_guess["flag"] != 0: - if iod.lower() in ["gauss", "auto"]: + if iod.lower() in ["gauss", "auto", "herget"]: res = do_fit( observations=observations, seq=sequence, @@ -1438,8 +1438,8 @@ def orbitfit( weighting for old comet apparitions (a row with a NaN/nonpositive value falls back to the default). iod : str - The IOD used to generate an initial guess orbit. Supports 'gauss' - and 'auto' (Gauss with BK-IOD fallback). + The IOD used to generate an initial guess orbit. Supports 'gauss', + 'herget' and 'auto' (Gauss with BK-IOD fallback). Default is 'auto'. fit_nongrav : bool | str | iterable of str Which non-gravitational Marsden parameters to fit after the 6-parameter diff --git a/src/layup/utilities/herget_iod.py b/src/layup/utilities/herget_iod.py new file mode 100644 index 00000000..d2977b6f --- /dev/null +++ b/src/layup/utilities/herget_iod.py @@ -0,0 +1,305 @@ +# Given an initial and final position, refine the velocity using +# Herget method as a way of creating a first guess for the inital orbit (IOD) +# import modules +import numpy as np +from sorcha.ephemeris.simulation_setup import create_assist_ephemeris +import assist +import rebound +from _layup_cpp._core import FitResult +from layup.utilities.universal_kepler import universal_step + +SPEED_OF_LIGHT_AU_DAY = 173.145 + + +def herget_with_assist(observations, seq, ephem, tolerance=0.001, max_iterations=100): + """Runs the Herget method on a set of observations. + + Parameters + ---------- + observations : list + List of all the observations of the object + seq : list[list] + list of lists containing the indices of observations that are closely spaced in time + tolerance : float + the maximum delta_rho residuals allowed; will continue to converge until the residuals are below this value + args : argparse.Namespace + The argparse object that was created when running from the CLI. Needed to instantiate assist simulations + aux : LayupConfigs.auxiliary object + Auxiliary Layup configs; needed to instantiate assist simulations + max_iterations : int (optional, default: 100) + the maximum number of iterations before the fitting stops""" + seq_lengths = [len(i) for i in seq] + longest_i = np.argmax(seq_lengths) # finds the sequence index with the most observations contained in it + obs = np.array(observations)[seq[longest_i]] + + # Define our values + + obs_1 = obs[0] + r_e_1 = obs_1.observer_position + rho_hat_1 = np.array(obs_1.rho_hat) + rho_1 = 40 # this is the magnitude of rho, direction given by rho_hat, initial guess is 40au + t1 = obs_1.epoch + r1 = r_e_1 + rho_1 * rho_hat_1 + + obs_n = obs[-1] + r_e_n = obs_n.observer_position + rho_hat_n = np.array(obs_n.rho_hat) + rho_n = 40 # this is the magnitude of rho, direction given by rho_hat, initial guess is 40au + tn = obs_n.epoch + rn = r_e_n + rho_n * rho_hat_n + + iteration = 0 + delta_rho1 = tolerance + 1 + delta_rhon = tolerance + 1 + + # Get original epochs so we can light-time correct them each iteration + epochs = np.zeros(len(obs)) + for i, observation in enumerate(obs): + epochs[i] = observation.epoch + + while (abs(delta_rho1) + abs(delta_rhon)) / 2 > tolerance and iteration < max_iterations: + + # Light-time correct the observation times + for i, observation in enumerate(obs): + # print(observation.epoch) + observation.epoch = epochs[i] - ((rho_1) + (rho_n)) / (2 * SPEED_OF_LIGHT_AU_DAY) + # print(observation.epoch) + + delta_rho1, delta_rhon, state_1 = find_drho( + obs, t1, tn, r1, rn, tolerance, ephem, rho_hat_1, rho_hat_n + ) + + # Update rho values + rho_1 -= delta_rho1 + r1 = r_e_1 + rho_1 * np.array(rho_hat_1) + rho_n -= delta_rhon + rn = r_e_n + rho_n * np.array(rho_hat_n) + print(delta_rho1, delta_rhon) + # print(rho_1, rho_n) + + iteration += 1 + + state = state_1 + solution = FitResult() + solution.state = state + solution.epoch = epochs[0] + solution.method = "herget" + solution.niter = iteration + solution.flag = 0 # Success flag + solution.ndof = len(observations) + solution.csq = 0.0 + solution.cov = [0.01] * 36 + + return [solution] + + +def find_drho(observations, t1, tn, r1, rn, tolerance, ephem, rho_hat_1, rho_hat_n): + """Find the adjustment to make to rho_1 and rho_n to make in order to reduce the residuals of the observations + + Parameters + ---------- + observations : list + list of observation objects + t1 : float + light-time corrected time for position r1, in TDB MJD + tn : float + light-time corrected time for position rn, in TDB MJD + r1 : numpy array + position vector at time t1 + rn : numpy array + position vector at time tn + tolerance : float + the average value of delta_rho1 and delta_rhon at which the orbit is considered to have converged at + args : argparse.Namespace + The argparse object that was created when running from the CLI. Needed to instantiate assist simulations + aux : LayupConfigs.auxiliary object + Auxiliary Layup configs; needed to instantiate assist simulations + rho_hat_1 : numpy array + unit vector of rho at time t1 + rho_hat_n : numpy array + unit vector of rho at time tn + + Returns + ------- + delta_rho1 : float + the amount to adjust rho_1 by to return a more accurate orbit + delta_rhon : float + the amount to adjust rho_n by to return a more accurate orbit + state_1[x, y, z, vx, vy, vz] + the new guess for the state vector at t1 + """ + + # Find velocities at rho_1 and rho_n + [vx1, vy1, vz1], [vxn, vyn, vzn] = find_velocity(t1, tn, r1, rn, tolerance) + [var_vx1, var_vy1, var_vz1], _ = find_velocity(t1, tn, r1 + rho_hat_1, rn, tolerance) + + # Simulation setup + sim = rebound.Simulation() + + sim.add(x=r1[0], y=r1[1], z=r1[2], vx=vx1, vy=vy1, vz=vz1) + var = sim.add_variation(testparticle=0) + var.particles[0].xyz = rho_hat_1 + var.particles[0].vxyz = np.array([var_vx1 - vx1, var_vy1 - vy1, var_vz1 - vz1]) + + ex = assist.Extras(sim, ephem) + sim.t = t1 - ephem.jd_ref + a1, a2, b = np.zeros((3, 2 * len(observations))) + + # For each observation, integrate to that time and find the residuals + for i, observation in enumerate(observations): + + # For this observation, get A and D + A = observation.a_vec + D = observation.d_vec + + t = observation.epoch + sim.integrate(t - ephem.jd_ref) + + r_e = np.array(observation.observer_position) + r = sim.particles[0].xyz + r_var = var.particles[0].xyz + rho = r - r_e + + # Add these to the arrays + b[2 * i] = np.dot(rho / np.linalg.norm(rho), A) + b[2 * i + 1] = np.dot(rho / np.linalg.norm(rho), D) + a1[2 * i] = b[2 * i] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), A) + a1[2 * i + 1] = b[2 * i + 1] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), D) + + _, [var_vxn, var_vyn, var_vzn] = find_velocity(t1, tn, r1, rn + rho_hat_n, tolerance) + + # Do the same for rho_n, set up simulation again + vxn, vyn, vzn = sim.particles[0].vxyz + sim = rebound.Simulation() + sim.add(x=rn[0], y=rn[1], z=rn[2], vx=vxn, vy=vyn, vz=vzn) + var = sim.add_variation(testparticle=0) + var.particles[0].xyz = rho_hat_n + var.particles[0].vxyz = np.array([var_vxn - vxn, var_vyn - vyn, var_vzn - vzn]) + + ex = assist.Extras(sim, ephem) + sim.t = tn - ephem.jd_ref + + # Find residuals for each observation + for i, observation in enumerate(observations): + A = observation.a_vec + D = observation.d_vec + + t = observation.epoch + + sim.integrate(t - ephem.jd_ref) + r_e = np.array(observation.observer_position) + r = sim.particles[0].xyz + r_var = var.particles[0].xyz + rho = r - r_e + + # Add to array + a2[2 * i] = b[2 * i] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), A) + a2[2 * i + 1] = b[2 * i + 1] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), D) + + sigma_a1b = sum(a1 * b) + sigma_a2b = sum(a2 * b) + sigma_a1squared = sum(a1**2) + sigma_a2squared = sum(a2**2) + sigma_a1a2 = sum(a1 * a2) + + delta_rho1 = (sigma_a1b * sigma_a2squared - sigma_a2b * sigma_a1a2) / ( + sigma_a1a2**2 - sigma_a1squared * sigma_a2squared + ) + delta_rhon = (-delta_rho1 * sigma_a1squared - sigma_a1b) / sigma_a1a2 + + # Check this is the solution, should equal zero + # print(sigma_a1b + delta_rho1*sigma_a1squared + delta_rhon*sigma_a1a2) + # print(sigma_a2b + delta_rho1*sigma_a1a2 + delta_rhon*sigma_a2squared) + # print(sum(a1*(b + delta_rho1*a1 + delta_rhon*a2))) + + return delta_rho1, delta_rhon, [*r1, vx1, vy1, vz1] + + +def find_velocity(t1, tn, r1, rn, tolerance): + """Converge on a velocity which takes position r1 at time t1 to position rn at time tn. + Uses the universal kepler stepper to integrate over time. + + Parameters + ---------- + t1 : float + Light-time corrected time of first state, in TDB MJD + tn : float + Light-time corrected time of nth state, in TDB MJD + r1 : numpy array + Position vector at time t1 + rn : numpy array + Position vector at time tn + tolerance : float + how closely the calculated rn value must lie within the correct value + + Returns + ------- + state_1[vx, vy, vz] + The velocity components of state vector at t1 + state_n[vx, vy, vz] + The velocity components of state vector at tn + """ + + # Initialising data + delta_t = tn - t1 + state_1 = np.array([*r1, 0, 0, 0]) + + for i in range(3): + state_1[i + 3] = (rn[i] - r1[i]) / delta_t + + state_n = [*rn + abs(tolerance) + 1, 0, 0, 0] + + # Find new values for vx, vy and vz in turn + while np.linalg.norm(state_n[:3] - rn) > tolerance: + state_1[3:], state_n = find_new_vel_with_universal_kepler(t1, tn, state_1, rn) + + return state_1[3:], state_n[3:] + + +def find_mag_to_adjust(P, Q, R): + # This is the formula for the point on a line (defined by Q and R) + # that is closest to a point outside the line, P + # For our purpose, this is the scale factor to vary the velocity by so that it + # will be closest to rho_n next time + mag = np.dot(R - Q, P - Q) / np.dot(R - Q, R - Q) + return mag + + +def find_new_vel_with_universal_kepler(t1, tn, state_1, state_n): + """Adjust the velocity components of an input position-velocity state to land closer to a desired position + + Parameters + ---------- + t1 : float + Light-time corrected time of first state, in TDB MJD + tn : float + Light-time corrected time of nth state, in TDB MJD + state_1 : numpy array + the state vector at time t1; in AU and AU/day + state_n : numpy array + the state vector at time tn; in AU and AU/day + + Returns + ------- + state_1[vx, vy, vz] + Adjusted velocity components of state_1 + state_n[x, y, z, vx, vy, vz] + The new state vector at time tn when using state_1 as the input""" + + # Initialise variables + dt = tn - t1 + GMtotal = 0.0002963092748799319 + + for i in range(3): + variation = np.zeros(6) + variation[i + 3] = 1 # we are varying vx, vy and vz by 1, once at a time + var_results = universal_step(GMtotal, dt, state_1, variation=variation) + + diff = find_mag_to_adjust( + np.array(state_n[:3]), + np.array(var_results.state[:3]), + np.array(var_results.state[:3]) + np.array(var_results.variation[:3]), + ) # find the velocity in this cartesian direction that will get us closest to the desired position + state_1[i + 3] += diff + + return state_1[3:], np.array(var_results.state) diff --git a/src/layup/utilities/universal_kepler.py b/src/layup/utilities/universal_kepler.py new file mode 100644 index 00000000..9a1af744 --- /dev/null +++ b/src/layup/utilities/universal_kepler.py @@ -0,0 +1,405 @@ +"""Universal-variable Kepler propagator with variational partials. + +Python port of ``universal-kepler.c`` in this directory (Danby 1988, p.178), +intended as the numerical core of a Herget initial-orbit-determination +prototype. + +WHAT THIS IS +------------ +The *initial*-value problem: given a state (r0, v0) and an interval dt, +return the state at t0 + dt, optionally propagating a 6-vector deviation +alongside it. The deviation is the state-transition matrix applied to a +direction -- which is what Herget actually needs, since only two directions +matter (d/d rho_1 and d/d rho_3), not the full 6x6. + +WHAT THIS IS NOT +---------------- +A Lambert solver. It does not solve the two-point boundary-value problem +(r1, t1, r3, t3 -> v1). It is, however, the natural engine for a *shooting* +Lambert: guess v1, propagate here, compare to r3, and correct using the +variational output as the Jacobian d r(t3) / d v1. + +Handles elliptic, hyperbolic and near-parabolic motion without branching on +a bound-orbit assumption -- the sign of ``alpha = gm/a`` selects the regime. + +UNITS are whatever ``gm`` is expressed in; for AU and days use layup's +``constants.MU_SUN`` (heliocentric) or ``constants.GMtotal`` (barycentric), +matching the frame the state is expressed in. The two differ by 0.13% and +layup uses both deliberately, so pick consciously. + +DIFFERENCES FROM THE C +---------------------- +1. **Multi-revolution fix.** The C is silently wrong for ``|dt|`` longer + than one orbital period: its elliptic branch reduces the local ``dt`` by + whole revolutions to build the initial guess, but the Newton iteration + then solves against that *reduced* value while ``universal_step`` forms + ``g = dt - gm*g3`` from the *unreduced* one. The Lagrange g comes out + too large by exactly ``n_rev * P``, the state runs away along the + velocity direction, and the routine still reports success. + + Here the reduction is used for the guess *only*; ``n_rev`` revolutions + are added back as ``n_rev * 2*pi / sqrt(alpha)`` (one revolution advances + the universal anomaly by 2*pi) and the Newton solves the FULL dt. That + is deliberately not just a guard: reducing dt and keeping the result + would give the right *state* but the wrong *partials*, because + neighbouring orbits have different periods and the deviation grows + secularly with every revolution. A state-only fix would silently break + exactly the thing this module exists to provide. + +2. **Relative convergence tolerance.** The C tests ``|ds| > EPS`` with + ``EPS = 1e-13`` as an *absolute* tolerance on ``s``, which carries units + of time/length. Over one revolution ``s`` spans ``P/a``: about 577 for + a mainbelt orbit but **2395 for a TNO at 43 AU**, so the test demands + 4e-17 in relative terms -- below double precision, and unreachable. + Newton then exhausts its six iterations, Laguerre-Conway its fifteen, + and ``kepler()`` returns ``KEPLER_FLAG`` -- which ``universal_step`` + ignores, computing a state from the unconverged ``s`` regardless. The + failure therefore gets worse the more distant the object, which is + precisely backwards for this project. (The C's author evidently saw + it: both loops carry a commented-out alternative on ``f/dt``, a + scale-invariant relative residual.) Here the test is relative to + ``max(1, |s|)``. + +3. ``r`` is recomputed from the converged g-functions rather than reused + from the last Newton evaluation (the C's ``*rx = fp`` is one iteration + stale -- negligible once converged, but free to do properly). + +4. ``gdot`` falls back to ``1 - (gm/r)*g2`` when ``|f|`` is too small for + the Wronskian form ``(1 + g*fdot)/f`` the C uses unconditionally. + +5. Non-convergence raises rather than returning a flag the caller may + ignore -- as ``universal_step`` itself does in the C. + +The unused ``stumpff()`` in the C (dead code -- both call sites use +``cfun()``) is not ported. Note that layup's own +``utilities.orbit_conversion.stumpff`` returns c0..c3 only; the variational +path here needs c4 and c5 via g1a/g2a/g3a, which is why ``cfun`` is ported +rather than reused. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import numpy as np + +__all__ = [ + "KeplerConvergenceError", + "KeplerStep", + "stumpff_c", + "universal_step", + "state_transition_matrix", +] + +# Iteration budgets, matching the C. +_NEWTON_MAX = 6 +_LAGCON_MAX = 15 +_TWO_PI = 2.0 * math.pi + +# Convergence tolerance on the universal variable s, applied RELATIVE to |s| +# (see note 2 in the module docstring). The C uses this as an absolute +# tolerance, which is unreachable for distant orbits. +_TOL = 1e-13 + + +def _converged(ds: float, s: float) -> bool: + return abs(ds) <= _TOL * max(1.0, abs(s)) + + +# Below this |f| the Wronskian form of gdot loses precision; see note 3 above. +_F_FLOOR = 1e-8 + + +class KeplerConvergenceError(RuntimeError): + """Neither Newton nor Laguerre-Conway reached the tolerance.""" + + +@dataclass +class KeplerStep: + """Result of one propagation. + + Attributes + ---------- + state : (6,) ndarray + Propagated (x, y, z, vx, vy, vz). + variation : (6,) ndarray or None + The input deviation propagated to the same epoch, i.e. the + state-transition matrix applied to it. None if none was supplied. + s : float + Converged universal variable. + r : float + Heliocentric (or barycentric) distance at the end of the step. + n_rev : int + Whole revolutions spanned by dt. Nonzero means the multi-revolution + path was exercised -- the regime the C got wrong. + n_iter : int + Iterations used by whichever solver converged. + solver : str + "newton" or "laguerre-conway". + """ + + state: np.ndarray + variation: np.ndarray | None + s: float + r: float + n_rev: int + n_iter: int + solver: str + + +def stumpff_c(z: float) -> tuple[float, float, float, float, float, float]: + """Stumpff functions c0..c5 by Mikkola's argument four-folding. + + Port of the C ``cfun()``. The four-folding keeps the series arguments + small (|h| < 0.1) so the truncated rational approximations for c4 and c5 + stay accurate, then unfolds with the duplication identities. + """ + h = z + k = 0 + while abs(h) >= 0.1: + h *= 0.25 + k += 1 + + c4 = (1.0 - h * (1.0 - h * (1.0 - h / 90.0 / (1.0 + h / 132.0)) / 56.0) / 30.0) / 24.0 + c5 = (1.0 - h * (1.0 - h * (1.0 - h / 110.0 / (1.0 + h / 156.0)) / 72.0) / 42.0) / 120.0 + + for _ in range(k): + c3 = 1.0 / 6.0 - h * c5 + c2 = 0.5 - h * c4 + c5 = (c5 + c4 + c2 * c3) / 16.0 + c4 = c3 * (2.0 - h * c3) / 8.0 + h *= 4.0 + + c3 = 1.0 / 6.0 - z * c5 + c2 = 0.5 - z * c4 + c1 = 1.0 - z * c3 + c0 = 1.0 - z * c2 + return c0, c1, c2, c3, c4, c5 + + +def _initial_guess(gm: float, dt: float, r0: float, alpha: float, u: float) -> tuple[float, int]: + """Starting value for the universal variable, and the revolution count. + + Returns (s_guess, n_rev). n_rev is nonzero only on the elliptic branch, + and the returned guess already includes the whole revolutions -- callers + solve against the full dt. + """ + # Short step relative to the current distance: the series guess is good + # and there is no revolution structure to unwrap. + if abs(dt / r0) <= 0.2: + return dt / r0 - (dt * dt * u) / (2.0 * r0**3), 0 + + if alpha <= 0.0: + # Hyperbolic. a < 0 here, so the sqrt arguments below are positive. + a = gm / alpha + en = math.sqrt(-gm / (a * a * a)) + ch = 1.0 - r0 / a + sh = u / math.sqrt(-a * gm) + e = math.sqrt(ch * ch - sh * sh) + dm = en * dt + if dm < 0.0: + return -math.log((-2.0 * dm + 1.8 * e) / (ch - sh)) / math.sqrt(-alpha), 0 + return math.log((2.0 * dm + 1.8 * e) / (ch + sh)) / math.sqrt(-alpha), 0 + + # Elliptic. + a = gm / alpha + en = math.sqrt(gm / (a * a * a)) + ec = 1.0 - r0 / a + es = u / (en * a * a) + + # Whole revolutions are stripped to keep the RK4 guess inside one orbit, + # then added back below. Truncation toward zero (not floor) so the + # remainder keeps dt's sign, matching the C's (int) cast -- but in Python + # the int is arbitrary precision, so the C's overflow at ~2e9 revolutions + # does not apply. + n_rev = int(en * dt / _TWO_PI) + dt_red = dt - n_rev * _TWO_PI / en + + y = en * dt_red - es + + # One RK4 step of the (ec, es) rotation, as a cheap high-order guess for + # the eccentric-anomaly increment. (The C also computes the eccentricity + # here for Danby's alternative guess, which is commented out; omitted.) + xx, yy = ec, es + h = en * dt_red + omx = h / (1.0 - xx) + k0x, k0y = -yy * omx, xx * omx + xx1, yy1 = xx + k0x / 2.0, yy + k0y / 2.0 + omx = h / (1.0 - xx1) + k1x, k1y = -yy1 * omx, xx1 * omx + xx1, yy1 = xx + k1x / 2.0, yy + k1y / 2.0 + omx = h / (1.0 - xx1) + k2x, k2y = -yy1 * omx, xx1 * omx + xx1 = xx + k2x + omx = h / (1.0 - xx1) + k3y = xx1 * omx + yy += (k0y + 2.0 * (k1y + k2y) + k3y) / 6.0 + + root_alpha = math.sqrt(alpha) + # sqrt(alpha)*s is the eccentric-anomaly increment, so a revolution is + # worth 2*pi/sqrt(alpha) in s. + s = (y + yy) / root_alpha + n_rev * _TWO_PI / root_alpha + return s, n_rev + + +def _solve_kepler(gm, dt, r0, alpha, u, zeta): + """Solve the universal Kepler equation for s; return the g-functions. + + Returns (g0..g5, r, s, n_rev, n_iter, solver). + """ + s_guess, n_rev = _initial_guess(gm, dt, r0, alpha, u) + + def _fvals(s): + c0, c1, c2, c3, _, _ = stumpff_c(s * s * alpha) + c1 *= s + c2 *= s * s + c3 *= s * s * s + f = r0 * c1 + u * c2 + gm * c3 - dt + fp = r0 * c0 + u * c1 + gm * c2 # == r at this s, hence always > 0 + fpp = zeta * c1 + u * c0 + fppp = zeta * c0 - u * alpha * c1 + return f, fp, fpp, fppp + + # Newton with Danby's cubic correction: three nested refinements of the + # same step reuse one function evaluation. + s = s_guess + ds = math.inf + solver, n_iter = "newton", 0 + for n_iter in range(1, _NEWTON_MAX + 1): + f, fp, fpp, fppp = _fvals(s) + ds = -f / fp + ds = -f / (fp + ds * fpp / 2.0) + ds = -f / (fp + ds * fpp / 2.0 + ds * ds * fppp / 6.0) + s += ds + if _converged(ds, s): + break + + if not _converged(ds, s): + # Laguerre-Conway from the original guess: larger convergence basin, + # at the cost of a square root per iteration. + solver = "laguerre-conway" + s = s_guess + ln = 5.0 + for n_iter in range(1, _LAGCON_MAX + 1): + f, fp, fpp, _ = _fvals(s) + disc = (ln - 1.0) ** 2 * fp * fp - (ln - 1.0) * ln * f * fpp + ds = -ln * f / (fp + math.copysign(math.sqrt(abs(disc)), fp)) + s += ds + if _converged(ds, s): + break + if not _converged(ds, s): + raise KeplerConvergenceError( + f"Kepler equation did not converge: dt={dt!r} r0={r0!r} " + f"alpha={alpha!r} u={u!r} last |ds|={abs(ds):.3e} s={s!r}" + ) + + c0, c1, c2, c3, c4, c5 = stumpff_c(s * s * alpha) + g0 = c0 + g1 = c1 * s + g2 = c2 * s**2 + g3 = c3 * s**3 + g4 = c4 * s**4 + g5 = c5 * s**5 + r = r0 * g0 + u * g1 + gm * g2 + return g0, g1, g2, g3, g4, g5, r, s, n_rev, n_iter, solver + + +def universal_step(gm, dt, state, variation=None) -> KeplerStep: + """Propagate `state` by `dt`, optionally carrying a deviation along. + + Parameters + ---------- + gm : float + Gravitational parameter, in units consistent with `state` and `dt`. + dt : float + Interval. May be negative. + state : array_like, shape (6,) + (x, y, z, vx, vy, vz) at the start of the step. + variation : array_like, shape (6,), optional + A deviation in the *initial* state. Propagated exactly (to the + two-body model) to the end of the step, i.e. the state-transition + matrix applied to this vector. + + Returns + ------- + KeplerStep + """ + state = np.asarray(state, dtype=float) + if state.shape != (6,): + raise ValueError(f"state must have shape (6,), got {state.shape}") + + r0vec, v0vec = state[:3], state[3:] + r0 = float(np.linalg.norm(r0vec)) + if r0 == 0.0: + raise ValueError("state has zero position; the two-body problem is singular there") + + v0s = float(v0vec @ v0vec) + u = float(r0vec @ v0vec) + alpha = 2.0 * gm / r0 - v0s # = gm/a; sign selects ellipse vs hyperbola + zeta = gm - alpha * r0 + + g0, g1, g2, g3, g4, g5, r, s, n_rev, n_iter, solver = _solve_kepler(gm, dt, r0, alpha, u, zeta) + + f = 1.0 - (gm / r0) * g2 + g = dt - gm * g3 + fdot = -(gm / (r * r0)) * g1 + # The Wronskian form f*gdot - fdot*g = 1 is better conditioned than + # 1 - (gm/r)*g2 except where f itself is near zero. + gdot = (1.0 + g * fdot) / f if abs(f) > _F_FLOOR else 1.0 - (gm / r) * g2 + + out = np.empty(6) + out[:3] = f * r0vec + g * v0vec + out[3:] = fdot * r0vec + gdot * v0vec + + var_out = None + if variation is not None: + dvar = np.asarray(variation, dtype=float) + if dvar.shape != (6,): + raise ValueError(f"variation must have shape (6,), got {dvar.shape}") + dr, dv = dvar[:3], dvar[3:] + + # Derivatives of the scalars the solve depends on, along the deviation. + r0pr = float(r0vec @ dr) / r0 + alphapr = -(2.0 * gm / (r0 * r0)) * r0pr - 2.0 * float(v0vec @ dv) + upr = float(r0vec @ dv) + float(v0vec @ dr) + zetapr = -alpha * r0pr - r0 * alphapr + + # d g_k / d alpha at fixed s (Stumpff recurrences). + g1a = 0.5 * (g3 - s * g2) + g2a = 0.5 * (2.0 * g4 - s * g3) + g3a = 0.5 * (3.0 * g5 - s * g4) + + # d s / d(deviation) from differentiating the Kepler equation at fixed dt. + spr = -(s * r0pr + g3 * zetapr + g2 * upr + (g3a * zeta + u * g2a) * alphapr) / r + + g1pr = g0 * spr + g1a * alphapr + g2pr = g1 * spr + g2a * alphapr + g3pr = g2 * spr + g3a * alphapr + rpr = r0pr + g1 * upr + g2 * zetapr + u * g1pr + zeta * g2pr + + fpr = (gm * g2 / (r0 * r0)) * r0pr - (gm / r0) * g2pr + gpr = -gm * g3pr + fdotpr = (gm / (r * r * r0)) * g1 * rpr + (gm / (r * r0 * r0)) * g1 * r0pr - (gm / (r * r0)) * g1pr + gdotpr = (gm / (r * r)) * g2 * rpr - (gm / r) * g2pr + + var_out = np.empty(6) + var_out[:3] = f * dr + g * dv + fpr * r0vec + gpr * v0vec + var_out[3:] = fdot * dr + gdot * dv + fdotpr * r0vec + gdotpr * v0vec + + return KeplerStep(state=out, variation=var_out, s=s, r=r, n_rev=n_rev, n_iter=n_iter, solver=solver) + + +def state_transition_matrix(gm, dt, state) -> np.ndarray: + """Full 6x6 d state(t0+dt) / d state(t0), by six unit variations. + + Herget needs only two columns of this (d/d rho_1 and d/d rho_3), so + prefer calling `universal_step` directly with the deviation you care + about; this exists for testing and for callers that want the whole thing. + """ + stm = np.empty((6, 6)) + for j in range(6): + e = np.zeros(6) + e[j] = 1.0 + stm[:, j] = universal_step(gm, dt, state, variation=e).variation + return stm diff --git a/tests/data/2000OK67_ephem.csv b/tests/data/2000OK67_ephem.csv new file mode 100644 index 00000000..63530972 --- /dev/null +++ b/tests/data/2000OK67_ephem.csv @@ -0,0 +1,62 @@ +provID,ra,dec,mjd_utc,obsTime,stn,rmsRA,rmsDec,rmsTime,rmsCorr,uncTime +2000 OK67,9.32151,4.45406,61253.0,2026-Aug-01 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.31419,4.4515,61254.0,2026-Aug-02 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.30654,4.44879,61255.0,2026-Aug-03 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.29855,4.44595,61256.0,2026-Aug-04 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.29023,4.44295,61257.0,2026-Aug-05 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.28157,4.43982,61258.0,2026-Aug-06 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.27259,4.43654,61259.0,2026-Aug-07 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.26328,4.43312,61260.0,2026-Aug-08 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.25364,4.42956,61261.0,2026-Aug-09 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.24368,4.42587,61262.0,2026-Aug-10 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.23341,4.42203,61263.0,2026-Aug-11 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.22282,4.41806,61264.0,2026-Aug-12 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.21192,4.41396,61265.0,2026-Aug-13 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.20071,4.40972,61266.0,2026-Aug-14 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.18921,4.40536,61267.0,2026-Aug-15 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.17741,4.40087,61268.0,2026-Aug-16 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.16531,4.39625,61269.0,2026-Aug-17 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.15293,4.3915,61270.0,2026-Aug-18 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.14027,4.38664,61271.0,2026-Aug-19 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.12733,4.38165,61272.0,2026-Aug-20 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.11412,4.37655,61273.0,2026-Aug-21 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.10064,4.37133,61274.0,2026-Aug-22 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0869,4.36599,61275.0,2026-Aug-23 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0729,4.36055,61276.0,2026-Aug-24 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.05865,4.35499,61277.0,2026-Aug-25 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.04415,4.34932,61278.0,2026-Aug-26 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0294,4.34355,61279.0,2026-Aug-27 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.01443,4.33768,61280.0,2026-Aug-28 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.99921,4.3317,61281.0,2026-Aug-29 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.98377,4.32563,61282.0,2026-Aug-30 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.96811,4.31945,61283.0,2026-Aug-31 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.95223,4.31318,61284.0,2026-Sep-01 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.93614,4.30682,61285.0,2026-Sep-02 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.91984,4.30037,61286.0,2026-Sep-03 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.90334,4.29382,61287.0,2026-Sep-04 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.88665,4.28719,61288.0,2026-Sep-05 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.86976,4.28048,61289.0,2026-Sep-06 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.85269,4.27368,61290.0,2026-Sep-07 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.83545,4.2668,61291.0,2026-Sep-08 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.81803,4.25985,61292.0,2026-Sep-09 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.80044,4.25282,61293.0,2026-Sep-10 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.7827,4.24572,61294.0,2026-Sep-11 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.7648,4.23855,61295.0,2026-Sep-12 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.74676,4.23132,61296.0,2026-Sep-13 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.72859,4.22403,61297.0,2026-Sep-14 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.71028,4.21667,61298.0,2026-Sep-15 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.69185,4.20926,61299.0,2026-Sep-16 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.6733,4.2018,61300.0,2026-Sep-17 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.65465,4.19429,61301.0,2026-Sep-18 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.63589,4.18673,61302.0,2026-Sep-19 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.61704,4.17912,61303.0,2026-Sep-20 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.5981,4.17148,61304.0,2026-Sep-21 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.57907,4.16379,61305.0,2026-Sep-22 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.55998,4.15608,61306.0,2026-Sep-23 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.54081,4.14832,61307.0,2026-Sep-24 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.52159,4.14054,61308.0,2026-Sep-25 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.50231,4.13273,61309.0,2026-Sep-26 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.48299,4.1249,61310.0,2026-Sep-27 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.46362,4.11704,61311.0,2026-Sep-28 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.44422,4.10917,61312.0,2026-Sep-29 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.42479,4.10128,61313.0,2026-Sep-30 00:00:00.000,X05,0.0001,0.0001,,0.0001, diff --git a/tests/layup/test_herget_iod.py b/tests/layup/test_herget_iod.py new file mode 100644 index 00000000..79b91d9c --- /dev/null +++ b/tests/layup/test_herget_iod.py @@ -0,0 +1,241 @@ +import os +import numpy as np +from numpy.lib import recfunctions as rfn +from numpy.testing import assert_allclose, assert_equal +import layup.utilities.herget_iod as herget +import spiceypy as spice + +from layup.utilities.data_utilities_for_tests import get_test_filepath +from layup.utilities.file_io.CSVReader import CSVDataReader +from layup.routines import Observation, get_ephem +from layup.utilities.data_processing_utilities import LayupObservatory +from layup.utilities.datetime_conversions import convert_tdb_date_to_julian_date +from layup.orbitfit import _build_sequence +import assist +import rebound +from layup.orbit_maths import build_ephem_and_mus + +SPEED_OF_LIGHT_AU_DAY = 173.145 + + +def test_find_mag_to_adjust(): + """Testing that the formula for finding the closest point on a line works by using a trivial case""" + + point = np.array([1, 2]) + line_point1 = np.array([0, 0]) + line_point2 = np.array([1, 0]) + + answer = 1 + mag = herget.find_mag_to_adjust(point, line_point1, line_point2) + assert mag == answer + + # Trying again for the edge case that the point lies on the line + point = np.array([0, 0]) + answer = 0 + mag = herget.find_mag_to_adjust(point, line_point1, line_point2) + assert mag == answer + + +def test_find_new_vel_with_universal_kepler(): + """Start with an object moving away from target and check it changes direction of velocity + Also check end position ends up closer to the target""" + + pos = np.array([50, 0, 0]) + vel = np.array([-1e-5, -1e-5, -1e-5]) + target = np.array( + [55, 5, 5, 0, 0, 0] + ) # Function should attempt to add a positive velocity in any cartesian direction to get to this position + + t1 = 0 + tn = 300 + + [vx1, vy1, vz1], [*pos_new, vxn, vyn, vzn] = herget.find_new_vel_with_universal_kepler( + t1, tn, [*pos, *vel], target + ) + + assert ([vx1, vy1, vz1] > vel).all() + assert ( + np.sqrt(sum((target[:3] - pos_new) ** 2)) < np.sqrt(sum((target[:3] - pos) ** 2)) + ).all() # Check the new end-position is closer to the target + + # Run again to check it continues to converge + [vx1, vy1, vz1], [*pos_new_rerun, vxn, vyn, vzn] = herget.find_new_vel_with_universal_kepler( + t1, tn, [*pos, vx1, vy1, vz1], target + ) + + assert (np.sqrt(sum((target[:3] - pos_new_rerun) ** 2)) < np.sqrt(sum((target[:3] - pos_new) ** 2))).all() + + +def test_find_drho(tmpdir): + """Take an object with a known rho value (from JPL), see if this function approaches the right direction""" + + os.chdir(tmpdir) + data = CSVDataReader( + get_test_filepath("2000OK67_ephem.csv"), "csv", primary_id_column_name="provID" + ).read_rows() + + layup_observatory = LayupObservatory(cache_dir=None) + + # The units of et are seconds (from J2000). This new column is used by + # data_processing_utilities.obscodes_to_barycentric. + et_col = np.array([spice.str2et(row["obsTime"]) for row in data], dtype=" 0 the low-order Stumpff functions are trigonometric.""" + c0, c1, c2, c3, _, _ = stumpff_c(z) + sq = math.sqrt(z) + assert c0 == pytest.approx(math.cos(sq), rel=1e-13) + assert c1 == pytest.approx(math.sin(sq) / sq, rel=1e-13) + assert c2 == pytest.approx((1.0 - math.cos(sq)) / z, rel=1e-13) + assert c3 == pytest.approx((sq - math.sin(sq)) / z**1.5, rel=1e-13) + + +@pytest.mark.parametrize("z", [-0.3, -2.0, -25.0]) +def test_stumpff_closed_form_hyperbolic(z): + c0, c1, c2, c3, _, _ = stumpff_c(z) + sq = math.sqrt(-z) + assert c0 == pytest.approx(math.cosh(sq), rel=1e-13) + assert c1 == pytest.approx(math.sinh(sq) / sq, rel=1e-13) + assert c2 == pytest.approx((math.cosh(sq) - 1.0) / (-z), rel=1e-13) + assert c3 == pytest.approx((math.sinh(sq) - sq) / (-z) ** 1.5, rel=1e-13) + + +@pytest.mark.parametrize("z", [-5.0, -0.2, 0.0, 0.2, 5.0, 40.0]) +def test_stumpff_recurrence(z): + """c_k(z) = 1/k! - z*c_{k+2}(z) ties the high orders to the low ones.""" + c0, c1, c2, c3, c4, c5 = stumpff_c(z) + assert c0 == pytest.approx(1.0 - z * c2, rel=1e-14, abs=1e-16) + assert c1 == pytest.approx(1.0 - z * c3, rel=1e-14, abs=1e-16) + assert c2 == pytest.approx(0.5 - z * c4, rel=1e-14, abs=1e-16) + assert c3 == pytest.approx(1.0 / 6.0 - z * c5, rel=1e-14, abs=1e-16) + + +# -------------------------------------------------------------------------- +# 2. The propagated state +# -------------------------------------------------------------------------- + + +def _propagate_via_elements(gm, dt, state): + """Independent two-body propagation through classical elements. + + Deliberately a different formulation from the universal-variable solver + under test: elements -> Kepler's equation in E -> back to Cartesian. + Elliptic orbits only. + """ + r0v, v0v = state[:3], state[3:] + r0 = np.linalg.norm(r0v) + a = 1.0 / (2.0 / r0 - (v0v @ v0v) / gm) + assert a > 0, "elements reference handles ellipses only" + n = math.sqrt(gm / a**3) + + # Eccentricity vector and the eccentric anomaly at t0. + h = np.cross(r0v, v0v) + evec = np.cross(v0v, h) / gm - r0v / r0 + e = np.linalg.norm(evec) + + cosE0 = (1.0 - r0 / a) / e + sinE0 = (r0v @ v0v) / (e * math.sqrt(gm * a)) + E0 = math.atan2(sinE0, cosE0) + M = E0 - e * math.sin(E0) + n * dt + + # Newton on Kepler's equation. + E = M + for _ in range(200): + dE = (E - e * math.sin(E) - M) / (1.0 - e * math.cos(E)) + E -= dE + if abs(dE) < 1e-15: + break + + # Lagrange f/g in terms of the eccentric-anomaly increment. + dE_ = E - E0 + r = a * (1.0 - e * math.cos(E)) + f = 1.0 - a / r0 * (1.0 - math.cos(dE_)) + g = dt + (math.sin(dE_) - dE_) / n + fdot = -math.sqrt(gm * a) / (r * r0) * math.sin(dE_) + gdot = 1.0 - a / r * (1.0 - math.cos(dE_)) + + out = np.empty(6) + out[:3] = f * r0v + g * v0v + out[3:] = fdot * r0v + gdot * v0v + return out + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.4, -0.05, 0.001, 0.05, 0.3, 0.75, 0.99]) +def test_state_matches_classical_elements(name, frac): + state = STATES[name] + dt = frac * _period(state) + got = universal_step(GM, dt, state).state + want = _propagate_via_elements(GM, dt, state) + assert got == pytest.approx(want, rel=1e-11, abs=1e-13) + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_energy_and_angular_momentum_conserved(name): + state = STATES[name] + P = _period(state) + e0 = 0.5 * state[3:] @ state[3:] - GM / np.linalg.norm(state[:3]) + h0 = np.cross(state[:3], state[3:]) + for frac in (0.13, 0.4, 0.87, 2.6): + s = universal_step(GM, frac * P, state).state + e1 = 0.5 * s[3:] @ s[3:] - GM / np.linalg.norm(s[:3]) + assert e1 == pytest.approx(e0, rel=1e-12) + assert np.cross(s[:3], s[3:]) == pytest.approx(h0, rel=1e-12) + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_round_trip(name): + """Forward then back returns the state and the deviation.""" + state = STATES[name] + dt = 0.37 * _period(state) + var = np.array([0.3, -0.2, 0.11, 1e-3, 2e-3, -5e-4]) + fwd = universal_step(GM, dt, state, variation=var) + back = universal_step(GM, -dt, fwd.state, variation=fwd.variation) + assert back.state == pytest.approx(state, rel=1e-11, abs=1e-13) + assert back.variation == pytest.approx(var, rel=1e-9, abs=1e-13) + + +def test_hyperbolic_propagation(): + state = _hyperbolic_state() + r0 = np.linalg.norm(state[:3]) + alpha = 2.0 * GM / r0 - state[3:] @ state[3:] + assert alpha < 0, "test state should be unbound" + e0 = 0.5 * state[3:] @ state[3:] - GM / r0 + for dt in (-400.0, -20.0, 5.0, 200.0, 3000.0): + out = universal_step(GM, dt, state) + e1 = 0.5 * out.state[3:] @ out.state[3:] - GM / np.linalg.norm(out.state[:3]) + assert e1 == pytest.approx(e0, rel=1e-11) + # and it should come back + dt = 1500.0 + there = universal_step(GM, dt, state).state + back = universal_step(GM, -dt, there).state + assert back == pytest.approx(state, rel=1e-10, abs=1e-13) + + +def test_near_parabolic_branch_is_exercised(): + """|dt/r0| <= 0.2 takes the series guess; make sure it is right there.""" + state = STATES["mainbelt"] + r0 = np.linalg.norm(state[:3]) + dt = 0.1 * r0 # inside the branch + got = universal_step(GM, dt, state).state + want = _propagate_via_elements(GM, dt, state) + assert got == pytest.approx(want, rel=1e-11, abs=1e-13) + + +def test_zero_dt_is_identity(): + state = STATES["tno"] + var = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + out = universal_step(GM, 0.0, state, variation=var) + assert out.state == pytest.approx(state, abs=1e-14) + assert out.variation == pytest.approx(var, abs=1e-14) + + +# -------------------------------------------------------------------------- +# 3. The variational output -- the reason this module exists +# -------------------------------------------------------------------------- + + +def _fd_stm(gm, dt, state, rel=1e-7): + """Central-difference d state(t+dt) / d state(t).""" + stm = np.empty((6, 6)) + scale = np.array([np.linalg.norm(state[:3])] * 3 + [np.linalg.norm(state[3:])] * 3) + for j in range(6): + h = rel * scale[j] + plus, minus = state.copy(), state.copy() + plus[j] += h + minus[j] -= h + stm[:, j] = (universal_step(gm, dt, plus).state - universal_step(gm, dt, minus).state) / (2.0 * h) + return stm + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.3, 0.08, 0.45, 0.9]) +def test_variational_matches_finite_differences(name, frac): + state = STATES[name] + dt = frac * _period(state) + analytic = state_transition_matrix(GM, dt, state) + numeric = _fd_stm(GM, dt, state) + # Compare column-wise against that column's own magnitude: the position + # and velocity blocks differ by many orders of magnitude, so a single + # global tolerance would be meaningless. + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 2e-6 + + +def test_variational_matches_finite_differences_hyperbolic(): + state = _hyperbolic_state() + for dt in (-300.0, 50.0, 900.0): + analytic = state_transition_matrix(GM, dt, state) + numeric = _fd_stm(GM, dt, state) + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 2e-6 + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_stm_is_symplectic(name): + """M^T J M = J for a Hamiltonian flow. Stronger than det(M) = 1, and it + catches sign or index errors in the variational algebra that a + finite-difference check with loose tolerance could let through. + + The tolerance has to scale as ||M||^2 * eps: the position-vs-velocity + block grows like dt (order 1e5 for a TNO over a period), so forming + M^T J M cancels large numbers down to O(1). Measured worst case is + about 0.6 * ||M||^2 * eps, so 20x that is a real test rather than a + rubber stamp. + """ + state = STATES[name] + J = np.block([[np.zeros((3, 3)), np.eye(3)], [-np.eye(3), np.zeros((3, 3))]]) + for frac in (0.11, 0.5, 1.7): + M = state_transition_matrix(GM, frac * _period(state), state) + floor = 20.0 * np.linalg.norm(M) ** 2 * np.finfo(float).eps + assert M.T @ J @ M == pytest.approx(J, abs=floor) + assert np.linalg.det(M) == pytest.approx(1.0, abs=max(1e-13, floor)) + + +def test_variation_is_linear(): + """The propagated deviation must be linear in the input deviation.""" + state = STATES["tno"] + dt = 0.3 * _period(state) + a = np.array([0.5, -0.25, 0.1, 1e-4, -2e-4, 3e-5]) + b = np.array([-0.1, 0.4, 0.7, 5e-5, 1e-4, -1e-4]) + va = universal_step(GM, dt, state, variation=a).variation + vb = universal_step(GM, dt, state, variation=b).variation + vab = universal_step(GM, dt, state, variation=2.0 * a - 3.0 * b).variation + assert vab == pytest.approx(2.0 * va - 3.0 * vb, rel=1e-10, abs=1e-14) + + +# -------------------------------------------------------------------------- +# Multi-revolution: where the C is wrong +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("n", [0, 1, 2, 5, 40]) +def test_state_is_periodic_across_revolutions(n): + """Adding whole periods must not move the state. + + The C fails this from n=1: its g is too large by n*P*|v|, so |r| grows + without bound while the routine still reports success. + """ + state = STATES["mainbelt"] + P = _period(state) + ref = universal_step(GM, 0.3 * P, state) + got = universal_step(GM, (n + 0.3) * P, state) + assert got.n_rev == n + # Rounding in a single solve at large s accumulates with revolution + # count: measured 2.7e-14 AU at n=1 rising to 6.2e-10 AU at n=40. + # 1e-8 AU (about 1.5 m) is loose enough to be stable and still four + # orders tighter than any real effect. + assert got.state == pytest.approx(ref.state, abs=1e-8) + + +@pytest.mark.parametrize("n", [1, 3]) +def test_variational_still_matches_fd_across_revolutions(n): + """The partials must be right in the multi-revolution regime too. + + This is why the fix solves the full dt rather than simply reducing it: + reducing would give the correct state but the wrong partials, since + neighbouring orbits have different periods and the deviation grows + secularly with every revolution. + """ + state = STATES["mainbelt"] + dt = (n + 0.3) * _period(state) + analytic = state_transition_matrix(GM, dt, state) + # Larger FD step than the single-revolution case: the STM grows with + # revolution count (that is the point of this test), so the differencing + # cancels more and a smaller step is noisier, not better. + numeric = _fd_stm(GM, dt, state, rel=1e-6) + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 1e-5 + + +def test_partials_grow_secularly_with_revolutions(): + """The along-track deviation grows with revolution count -- the concrete + reason a state-only multi-rev fix would have been wrong.""" + state = STATES["mainbelt"] + P = _period(state) + norms = [np.linalg.norm(state_transition_matrix(GM, (n + 0.3) * P, state)) for n in (0, 1, 2, 5)] + assert norms == sorted(norms), f"STM norm should increase with revolutions: {norms}" + assert norms[-1] > 3.0 * norms[0] + + +# -------------------------------------------------------------------------- +# 4. Faithfulness to the original C +# -------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def c_universal_step(tmp_path_factory): + """Build universal-kepler.c as a shared library and expose universal_step. + + Skipped when there is no compiler, or when the C source has moved. + """ + src = HERE / "universal-kepler.c" + if not src.exists(): + pytest.skip(f"C source not found at {src}") + cc = shutil.which("gcc") or shutil.which("cc") + if cc is None: + pytest.skip("no C compiler available") + + build = tmp_path_factory.mktemp("uk_c") + # The C declares `extern double machine_epsilon` but never uses it; give + # it a definition so the link is clean. + shim = build / "shim.c" + shim.write_text("double machine_epsilon = 0.0;\n") + so = build / "libuk.so" + proc = subprocess.run( + [cc, "-std=gnu99", "-w", "-O2", "-fPIC", "-shared", str(src), str(shim), "-lm", "-o", str(so)], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + pytest.skip(f"could not build the C reference: {proc.stderr[-400:]}") + + lib = ctypes.CDLL(str(so)) + arr6 = ctypes.c_double * 6 + lib.universal_step.restype = ctypes.c_int + lib.universal_step.argtypes = [ + ctypes.c_double, + ctypes.c_double, + ctypes.POINTER(arr6), + ctypes.POINTER(arr6), + ctypes.POINTER(arr6), + ] + + def call(gm, dt, state, variation): + s0 = arr6(*[float(x) for x in state]) + out = arr6() + var = arr6(*[float(x) for x in variation]) + flag = lib.universal_step( + ctypes.c_double(gm), + ctypes.c_double(dt), + ctypes.byref(s0), + ctypes.byref(out), + ctypes.byref(var), + ) + return flag, np.array(out[:]), np.array(var[:]) + + return call + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.4, 0.02, 0.25, 0.6, 0.95]) +def test_matches_c_within_one_period(c_universal_step, name, frac): + """Inside one revolution -- where the C is correct -- the port must + reproduce it, state and partials alike.""" + state = STATES[name] + dt = frac * _period(state) + var = np.array([0.3, -0.2, 0.11, 1e-3, 2e-3, -5e-4]) + + flag, c_state, c_var = c_universal_step(GM, dt, state, var) + if flag != 0: + # The C's absolute tolerance is unreachable for some orbits; that is + # pinned by test_c_tolerance_fails_on_distant_orbits below, and there + # is nothing to compare against here. + pytest.skip(f"C did not converge for {name} at {frac}P (its own defect)") + + py = universal_step(GM, dt, state, variation=var) + assert py.state == pytest.approx(c_state, rel=1e-11, abs=1e-13) + assert py.variation == pytest.approx(c_var, rel=1e-9, abs=1e-13) + + +def test_matches_c_hyperbolic(c_universal_step): + state = _hyperbolic_state() + var = np.array([0.2, 0.1, -0.3, 1e-3, -1e-3, 2e-4]) + for dt in (-250.0, 40.0, 800.0): + flag, c_state, c_var = c_universal_step(GM, dt, state, var) + assert flag == 0 + py = universal_step(GM, dt, state, variation=var) + assert py.state == pytest.approx(c_state, rel=1e-11, abs=1e-13) + assert py.variation == pytest.approx(c_var, rel=1e-9, abs=1e-13) + + +def test_c_tolerance_fails_on_distant_orbits(c_universal_step): + """Pins the C's scale-dependent convergence test. + + EPS = 1e-13 is applied as an absolute tolerance on the universal + variable s, which spans P/a over one revolution -- 577 for a mainbelt + orbit but 2395 for a TNO at 43 AU. The demand is then below double + precision, so convergence becomes a coin flip that the C loses more + often the more distant the object. The port's relative test converges + everywhere in this sweep. + """ + fracs = (-0.4, 0.02, 0.25, 0.5, 0.6, 0.75, 0.95) + c_failures, distant_failures = [], [] + for name in sorted(STATES): + state = STATES[name] + P = _period(state) + for frac in fracs: + dt = frac * P + flag, _, _ = c_universal_step(GM, dt, state, np.zeros(6)) + if flag != 0: + c_failures.append((name, frac)) + if name in ("tno", "eccentric"): + distant_failures.append((name, frac)) + # the port must converge regardless + universal_step(GM, dt, state) + + assert c_failures, "expected the C to fail somewhere in this sweep" + assert distant_failures, f"expected distant-orbit failures, got {c_failures}" + + +def test_c_is_wrong_past_one_period_and_the_port_is_not(c_universal_step): + """Pins the defect this port fixes, so nobody 'restores' the C behaviour. + + If this ever fails because the C now agrees, the C was fixed upstream and + this test should become an equality check. + """ + state = STATES["mainbelt"] + P = _period(state) + var = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + + ref = universal_step(GM, 0.3 * P, state).state + dt = 5.3 * P + flag, c_state, _ = c_universal_step(GM, dt, state, var) + + assert flag == 0, "the C reports success -- that is what makes it dangerous" + c_err = np.linalg.norm(c_state[:3] - ref[:3]) + assert c_err > 10.0, f"expected the C to be far off, got {c_err} AU" + + py = universal_step(GM, dt, state).state + assert py == pytest.approx(ref, rel=1e-9, abs=1e-11) + + +# -------------------------------------------------------------------------- +# Input handling +# -------------------------------------------------------------------------- + + +def test_rejects_bad_shapes(): + with pytest.raises(ValueError): + universal_step(GM, 1.0, np.zeros(5)) + with pytest.raises(ValueError): + universal_step(GM, 1.0, STATES["tno"], variation=np.zeros(3)) + + +def test_rejects_zero_position(): + with pytest.raises(ValueError): + universal_step(GM, 1.0, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]))