diff --git a/src/layup/orbitfit.py b/src/layup/orbitfit.py index afab1dbe..83caf1cf 100644 --- a/src/layup/orbitfit.py +++ b/src/layup/orbitfit.py @@ -15,8 +15,11 @@ Observation, gauss, get_ephem, + run_bk_iod, + run_bk_native_fit, run_from_vector_with_initial_guess, ) + from layup.convert import convert from layup.utilities.astrometric_uncertainty import data_weight_Veres2017 @@ -65,6 +68,25 @@ AU_M = 149597870700 SPEED_OF_LIGHT = 2.99792458e8 * 86400.0 / AU_M +# Heliocentric GM in AU^3 / day^2 (k^2, k = Gaussian gravitational constant). +# Used by the BK-native fit for the bound-orbit energy prior on gdot. +_MU_SUN = 0.00029591220828559104 + + +def _run_fit(assist_ephem, initial_guess, observations, engine): + """Dispatch a single LM fit step to the configured engine. + + Centralizing the dispatch here keeps do_fit's IOD-then-fit pipeline + parameterization-agnostic and lets us add new engines (e.g., a + future distance-dispatched 'auto') with a single edit instead of + threading the choice through every call site. + """ + if engine == "cartesian": + return run_from_vector_with_initial_guess(assist_ephem, initial_guess, observations) + if engine == "bk_native": + return run_bk_native_fit(assist_ephem, initial_guess, observations, _MU_SUN) + raise ValueError(f"Unknown engine {engine!r}; expected one of 'cartesian', 'bk_native'.") + def _get_result_dtypes(primary_id_column_name: str): """Helper function to create the result dtype with the correct primary ID column name.""" @@ -349,7 +371,7 @@ def do_gauss_iod(observations, seq): return solns -def do_fit(observations, seq, cache_dir, iod="gauss"): +def do_fit(observations, seq, cache_dir, iod="gauss", engine="cartesian"): """Carry out an orbit fit to the observations in a series of steps. A list of lists of observation indices specifies the order in which the fit proceeds. @@ -376,8 +398,20 @@ def do_fit(observations, seq, cache_dir, iod="gauss"): seq : list of lists A list of lists of observation indices. iod : str - The IOD used to generate an initial guess orbit. Currently supports ['gauss']. - Default is 'gauss'. + The IOD used to generate an initial guess orbit. Supported values: + - 'gauss' (default): Gauss's classic three-observation IOD. + - 'auto': Gauss first, falling back to the universal-BK + 5-parameter linear IOD (`run_bk_iod`) if every + Gauss root fails to seed a successful LM fit. + Empirically the two are complementary -- on the + diagnostic-scan dataset Gauss + BK fallback + covers ~90% of cases vs ~84% for Gauss alone. + engine : str + Which LM fitter to dispatch to. Supported: + - 'cartesian' (default): the existing 6D Cartesian-state fit. + - 'bk_native': the universal Bernstein-Khushalani fit + (run_bk_native_fit), with a fixed bound-orbit energy prior + on gdot. Recovers the Cartesian state at the same epoch. Returns ------- @@ -385,34 +419,45 @@ def do_fit(observations, seq, cache_dir, iod="gauss"): The result of the orbit fit. """ - if iod.lower() == "gauss": - solns = do_gauss_iod(observations, seq) - else: + iod_choice = iod.lower() + if iod_choice not in ("gauss", "auto"): raise ValueError(f"The IOD: {iod} is not supported. Please use a supported IOD.") - # If the selected iod fails, try something else. - if not solns: + solns = do_gauss_iod(observations, seq) + + if not solns and iod_choice == "gauss": + # Pure Gauss path with no roots: nothing else to try. logger.debug(f"The iod {iod} failed") x = FitResult() x.flag = 5 return x assist_ephem = get_ephem(cache_dir) - - #! I think this can be a `for/else loop...` - # Fit primary interval, starting with gauss solution - x = solns[0] obs = [observations[i] for i in seq[0]] - x = run_from_vector_with_initial_guess(assist_ephem, x, obs) - - if (x.flag != 0) and len(solns) > 1: - x = solns[1] - obs = [observations[i] for i in seq[0]] - x = run_from_vector_with_initial_guess(assist_ephem, x, obs) - elif (x.flag != 0) and len(solns) > 2: - x = solns[2] - obs = [observations[i] for i in seq[0]] - x = run_from_vector_with_initial_guess(assist_ephem, x, obs) + + # Fit primary interval: try each Gauss root in turn until one converges. + x = FitResult() + x.flag = 5 # No roots tried yet. + for soln in solns: + x = _run_fit(assist_ephem, soln, obs, engine) + if x.flag == 0: + break + + # If Gauss didn't get us a converged primary fit and iod='auto', fall + # back to the BK 5-parameter linear IOD on the primary segment. See + # bk_iod.cpp's regime-of-validity block: BK-IOD shines on distant short + # arcs, exactly the regime where Gauss's three-point geometry is + # ill-conditioned. + if x.flag != 0 and iod_choice == "auto" and len(obs) >= 3: + logger.debug(f"All {len(solns)} Gauss roots failed; trying BK-IOD fallback") + # Use the middle observation's epoch -- same convention as Gauss's + # idx1 in do_gauss_iod -- so the seed's frame matches what the LM + # would expect from a Gauss seed. + epoch = float(obs[len(obs) // 2].epoch) + bk_seed = run_bk_iod(obs, epoch, _MU_SUN) + if bk_seed.flag == 0: + x = _run_fit(assist_ephem, bk_seed, obs, engine) + if x.flag != 0: logger.debug(f"Primary interval failed. Total observations: {len(obs)}") x.flag = 3 # caution @@ -420,7 +465,7 @@ def do_fit(observations, seq, cache_dir, iod="gauss"): # Attempt to fit all the data, given the fit of the primary interval obs = observations - x = run_from_vector_with_initial_guess(assist_ephem, x, obs) + x = _run_fit(assist_ephem, x, obs, engine) # If that failed, build up the solution slowly if x.flag != 0: @@ -429,7 +474,7 @@ def do_fit(observations, seq, cache_dir, iod="gauss"): for i, sq in enumerate(seq): obs += [observations[i] for i in sq] print(i, "of", len(seq), obs[0], sq) - x = run_from_vector_with_initial_guess(assist_ephem, x, obs) + x = _run_fit(assist_ephem, x, obs, engine) print("flag:", x.flag) if x.flag != 0: x.flag = 4 diff --git a/src/lib/detection.cpp b/src/lib/detection.cpp index 80e5f2af..df0c0484 100644 --- a/src/lib/detection.cpp +++ b/src/lib/detection.cpp @@ -5,6 +5,8 @@ #include #include +#include +#include namespace py = pybind11; // --- Observation Variant Types --- diff --git a/src/lib/orbit_fit/bk_basis.cpp b/src/lib/orbit_fit/bk_basis.cpp new file mode 100644 index 00000000..b0febb05 --- /dev/null +++ b/src/lib/orbit_fit/bk_basis.cpp @@ -0,0 +1,294 @@ +// Bernstein-Khushalani parameterization primitives for the layup-internal +// universal BK fitter (feat/bk-everywhere). +// +// The math layer is pure C++/Eigen -- no ASSIST or REBOUND dependencies -- +// so this translation unit can be reasoned about and tested in isolation +// of the dynamics path. pybind11 bindings at the bottom expose the +// primitives to Python so Layer 1 tests (round-trip, finite-difference +// Jacobian, mixed-partial symmetry, etc.) can run via pytest. The design +// and math derivation live in the project memory file +// bk_everywhere_design.md. + +#include "bk_basis.h" + +#include +#include +#include + +#include +#include +#include + +namespace orbit_fit +{ + + namespace + { + // Internal cached quantities at the BK position (alpha, beta). + // + // p = n0 + alpha*a + beta*b + // s_sq = 1 + alpha^2 + beta^2 = |p|^2 + // rho_hat = p / sqrt(s_sq) + // rho_hat_alpha = (a - (a . rho_hat) * rho_hat) / sqrt(s_sq) + // rho_hat_beta = (b - (b . rho_hat) * rho_hat) / sqrt(s_sq) + // + // rho_hat_alpha and rho_hat_beta are the gnomonic-projection tangent + // vectors at rho_hat; they are NOT unit length in general (they + // scale as 1/sqrt(s_sq) times the projection of a/b onto T_{rho_hat}). + struct RhoFrame + { + double s_sq; + double s; + Eigen::Vector3d rho_hat; + Eigen::Vector3d rho_hat_alpha; + Eigen::Vector3d rho_hat_beta; + }; + + RhoFrame compute_rho_frame(double alpha, double beta, const BKFiducial &fid) + { + RhoFrame f; + f.s_sq = 1.0 + alpha * alpha + beta * beta; + f.s = std::sqrt(f.s_sq); + const Eigen::Vector3d p = fid.n0 + alpha * fid.a + beta * fid.b; + f.rho_hat = p / f.s; + const double rho_dot_a = f.rho_hat.dot(fid.a); + const double rho_dot_b = f.rho_hat.dot(fid.b); + f.rho_hat_alpha = (fid.a - rho_dot_a * f.rho_hat) / f.s; + f.rho_hat_beta = (fid.b - rho_dot_b * f.rho_hat) / f.s; + return f; + } + } // namespace + + BKFiducial choose_fiducial(const std::vector &rho_hats) + { + BKFiducial fid; + Eigen::Vector3d mean = Eigen::Vector3d::Zero(); + for (const auto &r : rho_hats) + mean += r; + if (mean.norm() < 1e-12) + { + // Pathological: observation directions cancel out. Pick ICRS x + // as a fallback; the fit is gauge-invariant under fiducial choice + // anyway, so any nonzero direction works. + mean = Eigen::Vector3d::UnitX(); + } + fid.n0 = mean.normalized(); + + // Gram-Schmidt against the ICRS axis least parallel to n0 so we + // don't divide by something tiny. + const Eigen::Vector3d seed = std::abs(fid.n0.z()) < 0.9 + ? Eigen::Vector3d::UnitZ() + : Eigen::Vector3d::UnitX(); + fid.a = (seed - seed.dot(fid.n0) * fid.n0).normalized(); + fid.b = fid.n0.cross(fid.a); + return fid; + } + + Eigen::Matrix bk_to_cartesian( + const BKState &bk, const BKFiducial &fid) + { + const RhoFrame f = compute_rho_frame(bk.alpha, bk.beta, fid); + const double inv_g = 1.0 / bk.gamma; + const double inv_g2 = inv_g * inv_g; + + const Eigen::Vector3d r = inv_g * f.rho_hat; + const Eigen::Vector3d v = inv_g * (bk.adot * f.rho_hat_alpha + bk.bdot * f.rho_hat_beta) + - bk.gdot * inv_g2 * f.rho_hat; + + Eigen::Matrix cart; + cart << r, v; + return cart; + } + + BKState cartesian_to_bk( + const Eigen::Matrix &cart, const BKFiducial &fid) + { + const Eigen::Vector3d r = cart.head<3>(); + const Eigen::Vector3d v = cart.tail<3>(); + + const double r_norm = r.norm(); + const double gamma = 1.0 / r_norm; + const Eigen::Vector3d rho_hat = gamma * r; + + // Gnomonic tangent-plane coordinates of rho_hat at n0. + const double u = rho_hat.dot(fid.n0); + const double alpha = rho_hat.dot(fid.a) / u; + const double beta = rho_hat.dot(fid.b) / u; + + // gdot = d/dt (1/|r|) = -(r . v) / |r|^3 = -gamma^2 * (rho_hat . v) + const double gdot = -gamma * gamma * rho_hat.dot(v); + + // d(rho_hat)/dt = gamma * (v - (v . rho_hat) * rho_hat) -- v's component perp to rho_hat, + // scaled to a sphere tangent vector. alpha-dot, beta-dot come from + // applying the quotient rule to alpha = (rho_hat . a) / (rho_hat . n0) + // and similarly for beta. + const Eigen::Vector3d rho_dot = gamma * (v - v.dot(rho_hat) * rho_hat); + const double rho_dot_n0 = rho_dot.dot(fid.n0); + const double adot = (rho_dot.dot(fid.a) - alpha * rho_dot_n0) / u; + const double bdot = (rho_dot.dot(fid.b) - beta * rho_dot_n0) / u; + + BKState bk; + bk.alpha = alpha; + bk.beta = beta; + bk.gamma = gamma; + bk.adot = adot; + bk.bdot = bdot; + bk.gdot = gdot; + return bk; + } + + Eigen::Matrix dcart_dbk( + const BKState &bk, const BKFiducial &fid) + { + const double alpha = bk.alpha; + const double beta = bk.beta; + const double gamma = bk.gamma; + const double adot = bk.adot; + const double bdot = bk.bdot; + const double gdot = bk.gdot; + + const RhoFrame f = compute_rho_frame(alpha, beta, fid); + const double inv_g = 1.0 / gamma; + const double inv_g2 = inv_g * inv_g; + const double inv_g3 = inv_g2 * inv_g; + const double inv_s2 = 1.0 / f.s_sq; + const double inv_s4 = inv_s2 * inv_s2; + + Eigen::Matrix J = Eigen::Matrix::Zero(); + + // Top-left and bottom-right 3x3 blocks: d(r)/d(alpha,beta,gamma) and + // d(v)/d(adot,bdot,gdot) -- identical shape. + const Eigen::Vector3d dr_dalpha = inv_g * f.rho_hat_alpha; + const Eigen::Vector3d dr_dbeta = inv_g * f.rho_hat_beta; + const Eigen::Vector3d dr_dgamma = -inv_g2 * f.rho_hat; + + J.block<3, 1>(0, 0) = dr_dalpha; + J.block<3, 1>(0, 1) = dr_dbeta; + J.block<3, 1>(0, 2) = dr_dgamma; + + J.block<3, 1>(3, 3) = dr_dalpha; + J.block<3, 1>(3, 4) = dr_dbeta; + J.block<3, 1>(3, 5) = dr_dgamma; + + // Bottom-left 3x3 block: d(v)/d(alpha,beta,gamma). Needs second + // derivatives of rho_hat with respect to (alpha, beta): + // d rho_hat_alpha / d alpha = -(1+beta^2)/s^4 * rho_hat + // - 2 alpha / s^2 * rho_hat_alpha + // d rho_hat_alpha / d beta = (alpha*beta)/s^4 * rho_hat + // - alpha/s^2 * rho_hat_beta + // - beta/s^2 * rho_hat_alpha + // d rho_hat_beta / d beta = -(1+alpha^2)/s^4 * rho_hat + // - 2 beta / s^2 * rho_hat_beta + // and d rho_hat_beta / d alpha == d rho_hat_alpha / d beta by mixed-partial symmetry. + const Eigen::Vector3d d_rha_dalpha = -(1.0 + beta * beta) * inv_s4 * f.rho_hat + - 2.0 * alpha * inv_s2 * f.rho_hat_alpha; + const Eigen::Vector3d d_rha_dbeta = (alpha * beta) * inv_s4 * f.rho_hat + - alpha * inv_s2 * f.rho_hat_beta + - beta * inv_s2 * f.rho_hat_alpha; + const Eigen::Vector3d d_rhb_dbeta = -(1.0 + alpha * alpha) * inv_s4 * f.rho_hat + - 2.0 * beta * inv_s2 * f.rho_hat_beta; + + // d v / d alpha + const Eigen::Vector3d dv_dalpha = inv_g * (adot * d_rha_dalpha + bdot * d_rha_dbeta) + - gdot * inv_g2 * f.rho_hat_alpha; + // d v / d beta + const Eigen::Vector3d dv_dbeta = inv_g * (adot * d_rha_dbeta + bdot * d_rhb_dbeta) + - gdot * inv_g2 * f.rho_hat_beta; + // d v / d gamma + const Eigen::Vector3d dv_dgamma = -inv_g2 * (adot * f.rho_hat_alpha + bdot * f.rho_hat_beta) + + 2.0 * gdot * inv_g3 * f.rho_hat; + + J.block<3, 1>(3, 0) = dv_dalpha; + J.block<3, 1>(3, 1) = dv_dbeta; + J.block<3, 1>(3, 2) = dv_dgamma; + + return J; + } + + double sigma_gdot_sq(const BKState &bk, double mu) + { + // Bound-orbit constraint: 0.5 |v|^2 <= mu / |r| = mu * gamma. + // Substituting |v|^2 = gdot^2 / gamma^4 + |adot rho_hat_alpha + bdot rho_hat_beta|^2 / gamma^2 + // (cross-terms with rho_hat vanish since rho_hat_alpha, rho_hat_beta are tangent to rho_hat), + // + // gdot^2 <= gamma^2 * (2 mu gamma^3 - |adot rho_hat_alpha + bdot rho_hat_beta|^2) + // + // The tangential term expands via + // |rho_hat_alpha|^2 = (1+beta^2)/s^4, |rho_hat_beta|^2 = (1+alpha^2)/s^4, + // rho_hat_alpha . rho_hat_beta = -alpha*beta/s^4, + // so + // |adot rho_hat_alpha + bdot rho_hat_beta|^2 = + // [adot^2 (1+beta^2) - 2 adot bdot alpha beta + bdot^2 (1+alpha^2)] / s^4 . + // At alpha = beta = 0 (the fiducial direction) this reduces to adot^2 + bdot^2. + const double alpha = bk.alpha; + const double beta = bk.beta; + const double gamma = bk.gamma; + const double adot = bk.adot; + const double bdot = bk.bdot; + const double s_sq = 1.0 + alpha * alpha + beta * beta; + const double s4 = s_sq * s_sq; + const double v_tan_sq = (adot * adot * (1.0 + beta * beta) + - 2.0 * adot * bdot * alpha * beta + + bdot * bdot * (1.0 + alpha * alpha)) / s4; + const double rhs = 2.0 * mu * gamma * gamma * gamma - v_tan_sq; + if (rhs <= 0.0) + { + // Tangential rates already exceed escape velocity for this + // gamma; the energy bound provides no constraint on gdot. + return std::numeric_limits::infinity(); + } + return gamma * gamma * rhs; + } + + static void bk_basis_bindings(pybind11::module &m) + { + namespace py = pybind11; + + py::class_(m, "BKState") + .def(py::init<>()) + .def(py::init([](double alpha, double beta, double gamma, + double adot, double bdot, double gdot) + { + BKState bk; + bk.alpha = alpha; bk.beta = beta; bk.gamma = gamma; + bk.adot = adot; bk.bdot = bdot; bk.gdot = gdot; + return bk; }), + py::arg("alpha") = 0.0, py::arg("beta") = 0.0, py::arg("gamma") = 0.0, + py::arg("adot") = 0.0, py::arg("bdot") = 0.0, py::arg("gdot") = 0.0) + .def_readwrite("alpha", &BKState::alpha) + .def_readwrite("beta", &BKState::beta) + .def_readwrite("gamma", &BKState::gamma) + .def_readwrite("adot", &BKState::adot) + .def_readwrite("bdot", &BKState::bdot) + .def_readwrite("gdot", &BKState::gdot) + .def("__repr__", [](const BKState &b) + { + std::ostringstream s; + s << ""; + return s.str(); }); + + py::class_(m, "BKFiducial") + .def(py::init<>()) + .def_readwrite("n0", &BKFiducial::n0) + .def_readwrite("a", &BKFiducial::a) + .def_readwrite("b", &BKFiducial::b); + + m.def("bk_choose_fiducial", &choose_fiducial, py::arg("rho_hats"), + "Construct a BKFiducial frame from a list of unit line-of-sight vectors."); + m.def("bk_to_cartesian", &bk_to_cartesian, + py::arg("bk"), py::arg("fid"), + "Forward transform: BK state -> 6-vector of barycentric Cartesian (r, v)."); + m.def("cartesian_to_bk", &cartesian_to_bk, + py::arg("cart"), py::arg("fid"), + "Inverse transform: 6-vector barycentric Cartesian -> BK state."); + m.def("dcart_dbk", &dcart_dbk, + py::arg("bk"), py::arg("fid"), + "6x6 Jacobian d(r, v) / d(alpha, beta, gamma, adot, bdot, gdot)."); + m.def("sigma_gdot_sq", &sigma_gdot_sq, + py::arg("bk"), py::arg("mu"), + "Variance of the bound-orbit energy prior on gdot."); + } + +} // namespace orbit_fit diff --git a/src/lib/orbit_fit/bk_basis.h b/src/lib/orbit_fit/bk_basis.h new file mode 100644 index 00000000..c0cb263d --- /dev/null +++ b/src/lib/orbit_fit/bk_basis.h @@ -0,0 +1,81 @@ +#pragma once + +#include +#include + +namespace orbit_fit +{ + + // Bernstein-Khushalani parameters at epoch. Origin: barycenter. + // (alpha, beta) are gnomonic tangent-plane coordinates of the + // line-of-sight direction rho_hat at a fiducial direction n0, + // gamma = 1/|r_helio|, and (adot, bdot, gdot) are their time + // derivatives at the epoch. + struct BKState + { + double alpha = 0.0; + double beta = 0.0; + double gamma = 0.0; + double adot = 0.0; + double bdot = 0.0; + double gdot = 0.0; + }; + + // Orthonormal frame defining the BK gnomonic tangent plane. + // {a, b, n0} form a right-handed orthonormal basis; n0 is the + // fiducial line-of-sight, (a, b) span its tangent plane. + struct BKFiducial + { + Eigen::Vector3d n0; + Eigen::Vector3d a; + Eigen::Vector3d b; + }; + + // Choose a fiducial frame from a list of line-of-sight unit vectors. + // n0 := normalize(sum(rho_hats)); (a, b) constructed by Gram-Schmidt + // against ICRS z (or ICRS x if n0 is near the z-axis). This is one + // of many valid choices -- the BK fit is gauge-invariant under any + // rotation of (a, b) about n0. + BKFiducial choose_fiducial(const std::vector &rho_hats); + + // Forward transform: BK -> barycentric Cartesian (position + velocity). + // r_vec = (1/gamma) * rho_hat(alpha, beta) + // v_vec = (1/gamma) [adot * rho_hat_alpha + bdot * rho_hat_beta] + // - (gdot/gamma^2) * rho_hat(alpha, beta) + Eigen::Matrix bk_to_cartesian( + const BKState &bk, const BKFiducial &fid); + + // Inverse transform: barycentric Cartesian -> BK. Well-defined for + // any state with r_vec . n0 > 0 (object on the n0-facing hemisphere) + // and gamma > 0. + BKState cartesian_to_bk( + const Eigen::Matrix &cart, const BKFiducial &fid); + + // 6x6 Jacobian d(r_vec, v_vec) / d(alpha, beta, gamma, adot, bdot, gdot). + // Block structure (each block is 3x3): + // [ d r / d (alpha,beta,gamma) 0 ] + // [ d v / d (alpha,beta,gamma) d v / d (adot,bdot,gdot) ] + // Top-left and bottom-right blocks are identical (both arise from the + // (1/gamma) * tangent-vector structure). + Eigen::Matrix dcart_dbk( + const BKState &bk, const BKFiducial &fid); + + // Variance of the bound-orbit gdot prior: + // sigma_gdot^2 = gamma^2 * (2 mu gamma^3 - |adot rho_hat_alpha + bdot rho_hat_beta|^2) + // + // The tangential-velocity term in the energy bound depends on the gnomonic + // tangent-vector norms at (alpha, beta), so the exact form expands to + // sigma_gdot^2 = gamma^2 * (2 mu gamma^3 - + // [adot^2 (1+beta^2) + // - 2 adot bdot alpha beta + // + bdot^2 (1+alpha^2)] / s^4) + // where s^2 = 1 + alpha^2 + beta^2. At the fiducial direction (alpha = beta = 0) + // this reduces to the familiar gamma^2 (2 mu gamma^3 - adot^2 - bdot^2). + // + // Returns +infinity when the tangential rates already exceed escape (the + // right-hand side would be non-positive), signalling "no prior." The + // caller's precision is 1 / sigma_gdot_sq, so +inf -> 0 precision -> + // no contribution, which is the correct behavior. + double sigma_gdot_sq(const BKState &bk, double mu); + +} // namespace orbit_fit diff --git a/src/lib/orbit_fit/bk_fit.cpp b/src/lib/orbit_fit/bk_fit.cpp new file mode 100644 index 00000000..ef325dbd --- /dev/null +++ b/src/lib/orbit_fit/bk_fit.cpp @@ -0,0 +1,230 @@ +// Levenberg-Marquardt driver for the universal-BK fitter +// (feat/bk-everywhere). Included from orbit_fit.cpp at the bottom of +// its `namespace orbit_fit { ... }` block, so all of layup's existing +// Cartesian-side machinery is in scope without forward declarations: +// +// Observation, FitResult (from detection.cpp / fit_result.cpp) +// residuals, partials (from orbit_fit.h) +// compute_residuals, create_sequences, get_weight_matrix, converged +// (from orbit_fit.cpp) +// bk_basis::* (from bk_basis.cpp -- included earlier in TU) +// +// The driver mirrors `orbit_fit()` and `run_from_vector_with_initial_guess()` +// but works in BK basis throughout: chain-rules the per-observation +// Cartesian partials produced by `compute_residuals` through the 6x6 +// dcart_dbk Jacobian into BK basis, then assembles and solves +// C = B^T W B + lambda I + P_prior in BK coordinates. + +// FitResult run_bk_native_fit +// +// See bk_everywhere_design.md for the algorithmic design. + +FitResult run_bk_native_fit( + struct assist_ephem *ephem, + FitResult initial_guess, + std::vector &detections, + double mu) +{ + FitResult result; + result.method = "bk_native"; + result.flag = 1; + result.epoch = initial_guess.epoch; + result.csq = HUGE_VAL; + result.niter = 0; + result.ndof = (int)(2 * detections.size() - 6); + result.state = initial_guess.state; + result.cov.fill(0.0); + + if (detections.size() < 3) + { + // Not enough observations to constrain a 6-parameter fit. + return result; + } + + // ----- Pick a fiducial direction from the observations ----- + std::vector rho_hats; + rho_hats.reserve(detections.size()); + for (const auto &obs : detections) + { + rho_hats.push_back(obs.rho_hat); + } + const BKFiducial fid = choose_fiducial(rho_hats); + + // ----- Convert the seed to BK ----- + Eigen::Matrix cart_seed; + for (int i = 0; i < 6; i++) cart_seed(i) = initial_guess.state[i]; + BKState bk = cartesian_to_bk(cart_seed, fid); + const BKState bk_seed = bk; + + // ----- Fixed bound-orbit energy prior on gdot ----- + const double sgsq = sigma_gdot_sq(bk_seed, mu); + const double gdot_precision = (std::isfinite(sgsq) && sgsq > 0.0) ? (1.0 / sgsq) : 0.0; + Eigen::Matrix P_prior = Eigen::Matrix::Zero(); + P_prior(5, 5) = gdot_precision; + + // ----- LM workspace ----- + std::vector reverse_in_seq, reverse_out_seq; + std::vector forward_in_seq, forward_out_seq; + std::vector times(detections.size()); + for (size_t i = 0; i < detections.size(); i++) + { + times[i] = detections[i].epoch; + } + create_sequences(times, initial_guess.epoch, + reverse_in_seq, reverse_out_seq, + forward_in_seq, forward_out_seq); + + Eigen::SparseMatrix W = get_weight_matrix(detections); + + std::vector resid_vec(detections.size()); + std::vector partials_vec(detections.size()); + + // ----- LM loop ----- + // Same initial lambda and accept threshold as the Cartesian fit + // (orbit_fit at orbit_fit.cpp L553). + double lambda = (206265.0 * 206265.0) / 1000.0; + const double rho_accept = 0.1; + double chi2_prev = HUGE_VAL; + double chi2_cur = HUGE_VAL; + Eigen::Matrix C_no_lambda; // last accepted Hessian sans Marquardt damping + bool have_accepted_step = false; + + const size_t iter_max = 100; + const double eps = 1e-12; + size_t iters; + for (iters = 0; iters < iter_max; iters++) + { + // --- Build current Cartesian state from BK --- + const Eigen::Matrix cart_now = bk_to_cartesian(bk, fid); + struct reb_particle p0; + p0.x = cart_now(0); p0.y = cart_now(1); p0.z = cart_now(2); + p0.vx = cart_now(3); p0.vy = cart_now(4); p0.vz = cart_now(5); + p0.m = 0.0; p0.r = 0.0; + p0.ax = 0.0; p0.ay = 0.0; p0.az = 0.0; + p0.hash = 0; + + // --- Residuals + Cartesian partials via the existing variational pipeline --- + compute_residuals(ephem, p0, initial_guess.epoch, + detections, + resid_vec, partials_vec, + forward_in_seq, forward_out_seq, + reverse_in_seq, reverse_out_seq); + + // --- Assemble B_cart (2N x 6) and the residual vector --- + const int N = (int)detections.size(); + Eigen::MatrixXd B_cart(2 * N, 6); + Eigen::VectorXd r_vec(2 * N); + for (int i = 0; i < N; i++) + { + for (int j = 0; j < 6; j++) + { + B_cart(2 * i, j) = partials_vec[i].x_partials[j]; + B_cart(2 * i + 1, j) = partials_vec[i].y_partials[j]; + } + r_vec(2 * i) = resid_vec[i].x_resid; + r_vec(2 * i + 1) = resid_vec[i].y_resid; + } + + // --- Chain rule: B_bk = B_cart * dcart_dbk(current BK state) --- + const Eigen::Matrix J = dcart_dbk(bk, fid); + const Eigen::MatrixXd B_bk = B_cart * J; + + // --- Normal equations in BK basis with Marquardt damping + fixed prior --- + const Eigen::MatrixXd Bt = B_bk.transpose(); + const Eigen::MatrixXd BtW = Bt * W; + Eigen::Matrix C_data = BtW * B_bk; // pure data Hessian + Eigen::Matrix C = C_data + + lambda * Eigen::Matrix::Identity() + + P_prior; + + // bk as a 6-vector for the prior-gradient term. The prior mean is zero + // (only gdot is constrained), so grad_prior = P_prior * bk_vec. + Eigen::Matrix bk_vec; + bk_vec << bk.alpha, bk.beta, bk.gamma, bk.adot, bk.bdot, bk.gdot; + Eigen::Matrix grad = BtW * r_vec + P_prior * bk_vec; + + // chi-square including the prior contribution + const double chi2_data = (r_vec.transpose() * W * r_vec)(0); + const double chi2_prior = bk_vec.transpose() * P_prior * bk_vec; + chi2_cur = chi2_data + chi2_prior; + + // --- Solve and Marquardt accept/reject --- + Eigen::Matrix dX = C.colPivHouseholderQr().solve(-grad); + const double rho_num = chi2_prev - chi2_cur; + const double rho_den = (dX.transpose() * (lambda * dX - grad)).norm(); + const double rho = rho_den > 0.0 ? rho_num / rho_den : -1.0; + + if (rho > rho_accept) + { + lambda *= 0.5; + bk.alpha += dX(0); + bk.beta += dX(1); + bk.gamma += dX(2); + bk.adot += dX(3); + bk.bdot += dX(4); + bk.gdot += dX(5); + chi2_prev = chi2_cur; + C_no_lambda = C_data + P_prior; // Hessian for the final covariance + have_accepted_step = true; + } + else + { + lambda *= 2.0; + } + + // --- Convergence (same predicate as the Cartesian fit) --- + const size_t ndof = detections.size() * 2 - 6; + const double thresh = 10.0; + Eigen::MatrixXd dX_mat = dX; + if (converged(dX_mat, eps, chi2_cur, ndof, thresh)) + { + result.flag = 0; + result.csq = chi2_cur; + break; + } + } + + result.niter = (int)iters; + + // ----- Covariance in BK, then transform to Cartesian ----- + if (have_accepted_step) + { + const Eigen::Matrix cov_bk = C_no_lambda.inverse(); + const Eigen::Matrix J = dcart_dbk(bk, fid); + const Eigen::Matrix cov_cart = J * cov_bk * J.transpose(); + for (int i = 0; i < 6; i++) + for (int j = 0; j < 6; j++) + result.cov[i * 6 + j] = cov_cart(i, j); + } + + const size_t ndof = detections.size() * 2 - 6; + const double thresh = 10.0; + if ((result.csq / (double)ndof) > thresh) + { + result.flag = 2; // "converged" but chi2/dof is too large + } + + // ----- Cartesian state of the converged BK fit ----- + const Eigen::Matrix cart_final = bk_to_cartesian(bk, fid); + for (int i = 0; i < 6; i++) + { + result.state[i] = cart_final(i); + } + + return result; +} + +#ifdef Py_PYTHON_H +static void bk_fit_bindings(pybind11::module &m) +{ + namespace py = pybind11; + m.def("run_bk_native_fit", &run_bk_native_fit, + py::arg("ephem"), + py::arg("initial_guess"), + py::arg("detections"), + py::arg("mu"), + "Universal-BK Levenberg-Marquardt orbit fit. Reuses layup's " + "Cartesian variational machinery, with chain-rule + bound-orbit " + "energy prior applied in BK basis."); +} +#endif /* Py_PYTHON_H */ diff --git a/src/lib/orbit_fit/bk_iod.cpp b/src/lib/orbit_fit/bk_iod.cpp new file mode 100644 index 00000000..f044cec9 --- /dev/null +++ b/src/lib/orbit_fit/bk_iod.cpp @@ -0,0 +1,237 @@ +// Universal-BK 5-parameter linear initial orbit determination. +// +// Included from orbit_fit.cpp at the bottom of `namespace orbit_fit`, +// after bk_fit.cpp, so all of layup's existing types and the BK +// primitives from bk_basis.cpp are in scope: +// Observation, FitResult (from detection.cpp / fit_result.cpp) +// SPEED_OF_LIGHT (from predict.cpp; AU/day) +// BKState, BKFiducial, +// choose_fiducial, bk_to_cartesian, +// dcart_dbk, sigma_gdot_sq (from bk_basis.cpp) +// +// The model is the small-angle / no-gravity linear projection used by +// Bernstein & Khushalani's `prelim_fit` (liborbfit/orbfit1.c), adapted +// to barycenter-origin BK. Per observation, the predicted tangent- +// plane components are linear in (alpha, beta, gamma, adot, bdot): +// +// t_i = (obs_jd_i - (r_obs_i . n0) / c) - epoch // light-time corrected dt +// X_i = r_obs_i . a, Y_i = r_obs_i . b // observer in BK tangent plane +// x_obs_i = (rho_hat_i . a) / (rho_hat_i . n0) // observation, gnomonic +// y_obs_i = (rho_hat_i . b) / (rho_hat_i . n0) +// +// x_obs_i = alpha + adot * t_i - gamma * X_i +// y_obs_i = beta + bdot * t_i - gamma * Y_i +// +// gdot is pinned to 0 with a nominal prior variance from the bound- +// orbit energy constraint (sigma_gdot_sq, also matching prelim_fit's +// covar[5][5] = 0.1 * (2 pi)^2 * gamma^3 convention up to a small +// numerical factor). +// +// --------------------------------------------------------------------------- +// REGIME OF VALIDITY +// --------------------------------------------------------------------------- +// +// The model omits the perspective denominator 1 / (1 - gamma * ze), where +// ze = r_obs . n0 is the observer's projection onto the line of sight. +// For a fixed observer (e.g., Earth at ~1 AU) the omitted factor scales as +// |gamma * ze| ~ (1 / r_helio) * 1 AU +// which becomes O(1) for inner-solar-system objects and is small only for +// distant ones. Empirically, BK-IOD's heliocentric-position accuracy on +// the diagnostic-scan dataset is (best window across arc lengths): +// +// population best arc best drift / r_helio +// -------------- -------- -------------------- +// mainbelt 2.5AU 1 d 1.9% <- BUT 2d+ catastrophically fails +// mainbelt 3.5AU 1 d 0.9% <- BUT 2d+ catastrophically fails +// centaur 15 AU 0.5 d 2.5% +// centaur 25 AU 0.5 d 0.4% +// classical 42 AU 10 d 2.1% +// scattered 70 AU 7 d 0.5% +// sednoid 80 AU 10 d 1.0% +// +// Practical guidance: +// * Distant objects (centaurs and TNOs): BK-IOD is excellent. +// Sub-percent on heliocentric distance at well-chosen arc lengths. +// Sweet spot moves to longer arcs (5-30 d) as distance increases. +// * Mainbelt and closer: BK is essentially outside its regime of +// validity; only the narrow 1-day window gives a usable answer. +// Use Gauss IOD for these objects. +// +// In all cases, BK-IOD is intended as a SEED for the full BK or +// Cartesian LM fit, not as a final orbit. The LM cleans up the +// O(gamma*ze) perspective bias as part of its convergence. + +FitResult run_bk_iod( + std::vector &observations, + double epoch, + double mu) +{ + FitResult result; + result.method = "bk_iod"; + result.flag = 1; + result.epoch = epoch; + result.csq = HUGE_VAL; + result.niter = 0; + result.ndof = 0; + for (int i = 0; i < 6; i++) result.state[i] = 0.0; + result.cov.fill(0.0); + + const int N = (int)observations.size(); + // Need 2N >= 5 equations to constrain 5 parameters. + if (N < 3) + { + return result; + } + + // ----- Fiducial direction from the observations ----- + std::vector rho_hats; + rho_hats.reserve(observations.size()); + for (const auto &obs : observations) + { + rho_hats.push_back(obs.rho_hat); + } + const BKFiducial fid = choose_fiducial(rho_hats); + + // ----- Build the (2N x 5) weighted least-squares system ----- + // Parameter ordering matches BKState's positional fields: + // p = (alpha, beta, gamma, adot, bdot) + Eigen::MatrixXd A(2 * N, 5); + Eigen::VectorXd b(2 * N); + Eigen::VectorXd w(2 * N); // diagonal weights (precision = 1/variance) + + for (int i = 0; i < N; i++) + { + const auto &obs = observations[i]; + const Eigen::Vector3d r_obs( + obs.observer_position[0], + obs.observer_position[1], + obs.observer_position[2]); + + // Light-time corrected dt: when the observer is closer to the + // target along the line of sight (ze > 0), the topocentric + // distance is smaller, light-time is smaller, and emission time + // is *later* relative to the baseline -- so we ADD ze/c. Any + // constant offset is absorbed by the rates and doesn't affect + // the linear fit. + const double ze = r_obs.dot(fid.n0); + const double t_i = (obs.epoch + ze / SPEED_OF_LIGHT) - epoch; + + // Observer projected onto the BK tangent plane. + const double X_i = r_obs.dot(fid.a); + const double Y_i = r_obs.dot(fid.b); + + // Observed line-of-sight projected onto the BK tangent plane + // via gnomonic projection. Matches the predicted x, y in the + // linear model: alpha, beta are themselves gnomonic coordinates, + // so the observed tangent-plane coords must also be gnomonic. + // (Small-angle approximation rho_hat . a would introduce + // O(alpha^2) error -- noticeable at TNO arcs of several degrees.) + const double rho_n0 = obs.rho_hat.dot(fid.n0); + const double x_obs = obs.rho_hat.dot(fid.a) / rho_n0; + const double y_obs = obs.rho_hat.dot(fid.b) / rho_n0; + + // Per-observation uncertainties. Default to a 1" guess if the + // Observation didn't carry an explicit uncertainty; the LM + // pipeline applies the same fallback elsewhere. + const double sigma_ra = obs.ra_unc.value_or(1.0 / 206265.0); + const double sigma_dec = obs.dec_unc.value_or(1.0 / 206265.0); + const double w_x = 1.0 / (sigma_ra * sigma_ra); + const double w_y = 1.0 / (sigma_dec * sigma_dec); + + // Row for the x (alpha-axis) equation. + A(2 * i, 0) = 1.0; + A(2 * i, 1) = 0.0; + A(2 * i, 2) = -X_i; + A(2 * i, 3) = t_i; + A(2 * i, 4) = 0.0; + b(2 * i) = x_obs; + w(2 * i) = w_x; + + // Row for the y (beta-axis) equation. + A(2 * i + 1, 0) = 0.0; + A(2 * i + 1, 1) = 1.0; + A(2 * i + 1, 2) = -Y_i; + A(2 * i + 1, 3) = 0.0; + A(2 * i + 1, 4) = t_i; + b(2 * i + 1) = y_obs; + w(2 * i + 1) = w_y; + } + + // Weighted normal equations: H p = g with H = A^T W A, g = A^T W b. + const Eigen::MatrixXd AtW = A.transpose() * w.asDiagonal(); + const Eigen::Matrix H = AtW * A; + const Eigen::Matrix g = AtW * b; + const Eigen::Matrix p = H.colPivHouseholderQr().solve(g); + + // Pack into a BKState (gdot pinned to 0). + BKState bk; + bk.alpha = p(0); + bk.beta = p(1); + bk.gamma = p(2); + bk.adot = p(3); + bk.bdot = p(4); + bk.gdot = 0.0; + + // Reject pathological solutions (gamma must be positive for the + // forward transform to make sense). + if (!std::isfinite(bk.alpha) || !std::isfinite(bk.beta) || !std::isfinite(bk.gamma) + || !std::isfinite(bk.adot) || !std::isfinite(bk.bdot) || bk.gamma <= 0.0) + { + return result; + } + + // 6x6 BK covariance: top-left 5x5 from the fit's inverse-normal + // matrix; gdot row/column from the bound-orbit energy prior (gdot + // is treated as uncorrelated with the data-constrained params). + const Eigen::Matrix cov_5 = H.inverse(); + Eigen::Matrix cov_bk = Eigen::Matrix::Zero(); + cov_bk.block<5, 5>(0, 0) = cov_5; + const double sgsq = sigma_gdot_sq(bk, mu); + cov_bk(5, 5) = std::isfinite(sgsq) ? sgsq : 0.0; + + // Transform to Cartesian via the analytic 6x6 chain rule. + const Eigen::Matrix cart = bk_to_cartesian(bk, fid); + const Eigen::Matrix J = dcart_dbk(bk, fid); + const Eigen::Matrix cov_cart = J * cov_bk * J.transpose(); + + for (int i = 0; i < 6; i++) + { + result.state[i] = cart(i); + } + for (int i = 0; i < 6; i++) + { + for (int j = 0; j < 6; j++) + { + result.cov[i * 6 + j] = cov_cart(i, j); + } + } + + // Residual chi-square at the linear solution (for diagnostic only). + const Eigen::VectorXd resid = A * p - b; + double csq = 0.0; + for (int k = 0; k < 2 * N; k++) + { + csq += w(k) * resid(k) * resid(k); + } + result.csq = csq; + result.ndof = (2 * N >= 5) ? (2 * N - 5) : 0; + result.niter = 1; // closed-form -- one "iteration" + result.flag = 0; + return result; +} + +#ifdef Py_PYTHON_H +static void bk_iod_bindings(pybind11::module &m) +{ + namespace py = pybind11; + m.def("run_bk_iod", &run_bk_iod, + py::arg("observations"), + py::arg("epoch"), + py::arg("mu"), + "Universal-BK 5-parameter linear initial orbit determination. " + "Closed-form weighted least-squares over (alpha, beta, gamma, " + "adot, bdot); gdot pinned to 0 with covariance from the bound-" + "orbit energy prior. Returns a FitResult with barycentric " + "Cartesian state at the supplied epoch."); +} +#endif /* Py_PYTHON_H */ diff --git a/src/lib/orbit_fit/orbit_fit.cpp b/src/lib/orbit_fit/orbit_fit.cpp index 69d682ac..2f2049a2 100644 --- a/src/lib/orbit_fit/orbit_fit.cpp +++ b/src/lib/orbit_fit/orbit_fit.cpp @@ -53,6 +53,7 @@ #include #include "orbit_fit.h" +#include "bk_basis.cpp" #include "../gauss/gauss.cpp" #include "predict.cpp" @@ -782,6 +783,20 @@ namespace orbit_fit } +// Universal-BK fit (LM driver in BK basis). Inlined here, inside +// `namespace orbit_fit`, so all of layup's existing Cartesian-fit +// helpers (compute_residuals, create_sequences, get_weight_matrix, +// converged) and the Observation/FitResult types are in scope without +// forward declarations. Math primitives come from bk_basis.cpp, +// included at the top of orbit_fit.cpp. +#include "bk_fit.cpp" + +// Universal-BK 5-parameter linear IOD. Same inlining rationale as +// bk_fit.cpp above -- inlined inside `namespace orbit_fit` so that +// Observation, FitResult, SPEED_OF_LIGHT, and the bk_basis primitives +// are all in scope. +#include "bk_iod.cpp" + #ifdef Py_PYTHON_H static void orbit_fit_bindings(py::module &m) { diff --git a/src/main.cpp b/src/main.cpp index c5a06424..875f28e2 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -71,6 +71,9 @@ PYBIND11_MODULE(_core, m) orbit_fit::orbit_fit_result_bindings(m); orbit_fit::predict_bindings(m); orbit_fit::predict_result_bindings(m); + orbit_fit::bk_basis_bindings(m); + orbit_fit::bk_fit_bindings(m); + orbit_fit::bk_iod_bindings(m); #ifdef VERSION_INFO m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); diff --git a/tests/layup/test_bk_basis.py b/tests/layup/test_bk_basis.py new file mode 100644 index 00000000..991678db --- /dev/null +++ b/tests/layup/test_bk_basis.py @@ -0,0 +1,281 @@ +"""Layer 1 tests for the BK basis primitives. + +These tests exercise the pure-math layer of the universal BK fitter +(`src/lib/orbit_fit/bk_basis.cpp`) via the pybind11 bindings exposed +through `layup.routines`. No ASSIST/REBOUND/ephemeris setup is +required -- the math primitives operate on barycentric Cartesian +states and BK parameters directly. + +Coverage: + * round-trip Cartesian <-> BK transforms + * analytic dcart_dbk vs central-difference finite-differences + * mixed-partial symmetry of the second derivatives that show up + in the bottom-left cross-term block + * fiducial-direction gauge invariance (different choices of n0 + must give the same Cartesian orbit) + * special-case forms at alpha = beta = 0 (the fiducial direction) + * sigma_gdot_sq agreement with the bound-orbit energy bound + computed independently in Cartesian +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from layup.routines import ( + BKState, + BKFiducial, + bk_choose_fiducial, + bk_to_cartesian, + cartesian_to_bk, + dcart_dbk, + sigma_gdot_sq, +) + +# GM_sun in AU^3 / day^2 (Gaussian gravitational constant squared). +MU_SUN = 0.00029591220828559104 + + +# A spread of test BK states covering mainbelt / NEO / TNO regimes. +# Format: (alpha, beta, gamma, adot, bdot, gdot) with gamma = 1/r_helio in 1/AU. +_BK_CASES = [ + # mainbelt ~3 AU, near-circular + (0.1, -0.05, 1.0 / 3.0, 1e-3, -8e-4, 3e-6), + # NEO ~1.2 AU, modest rates + (-0.2, 0.15, 1.0 / 1.2, 5e-3, 4e-3, -1e-5), + # TNO ~42 AU, slow + (0.02, 0.01, 1.0 / 42.0, 4e-5, -3e-5, 1e-8), + # near the fiducial direction (small alpha, beta) -- a common case + (1e-4, 2e-4, 0.025, 6e-5, 5e-5, -2e-7), + # off the fiducial direction + (0.5, -0.4, 0.05, 2e-4, 1e-4, -3e-8), +] + + +def _make_fiducial(rng: np.random.Generator) -> BKFiducial: + """Pick a reproducible fiducial direction not aligned with an ICRS axis.""" + n0 = rng.normal(size=3) + n0 /= np.linalg.norm(n0) + return bk_choose_fiducial([n0]) + + +def _bk_from_tuple(t): + return BKState(*t) + + +# --------------------------------------------------------------------------- +# Round-trip Cartesian <-> BK +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("case", _BK_CASES) +def test_round_trip_bk_to_cart_to_bk(case): + """BK -> Cartesian -> BK recovers the input to machine precision.""" + rng = np.random.default_rng(seed=12345) + fid = _make_fiducial(rng) + bk = _bk_from_tuple(case) + cart = bk_to_cartesian(bk, fid) + bk_back = cartesian_to_bk(cart, fid) + for name in ("alpha", "beta", "gamma", "adot", "bdot", "gdot"): + original = getattr(bk, name) + recovered = getattr(bk_back, name) + np.testing.assert_allclose( + recovered, original, rtol=1e-12, atol=1e-15, err_msg=f"BK.{name} not recovered through round-trip" + ) + + +@pytest.mark.parametrize("case", _BK_CASES) +def test_round_trip_cart_to_bk_to_cart(case): + """Cartesian -> BK -> Cartesian recovers the input to machine precision.""" + rng = np.random.default_rng(seed=67890) + fid = _make_fiducial(rng) + bk = _bk_from_tuple(case) + cart_in = np.asarray(bk_to_cartesian(bk, fid)).flatten() + bk_round = cartesian_to_bk(cart_in, fid) + cart_out = np.asarray(bk_to_cartesian(bk_round, fid)).flatten() + np.testing.assert_allclose(cart_out, cart_in, rtol=1e-12, atol=1e-15) + + +# --------------------------------------------------------------------------- +# Analytic dcart_dbk vs finite-difference +# --------------------------------------------------------------------------- + + +def _bk_perturb(bk: BKState, idx: int, delta: float) -> BKState: + names = ("alpha", "beta", "gamma", "adot", "bdot", "gdot") + vals = [getattr(bk, n) for n in names] + vals[idx] += delta + return BKState(*vals) + + +@pytest.mark.parametrize("case", _BK_CASES) +def test_dcart_dbk_matches_finite_difference(case): + """Analytic 6x6 Jacobian agrees with central-difference per element.""" + rng = np.random.default_rng(seed=11111) + fid = _make_fiducial(rng) + bk = _bk_from_tuple(case) + J_analytic = np.asarray(dcart_dbk(bk, fid)) + + # Step sizes scaled by parameter magnitude so we get a sensible FD + # for both the O(1) (alpha, beta) and O(1e-5) (rates) axes. + param_vals = (bk.alpha, bk.beta, bk.gamma, bk.adot, bk.bdot, bk.gdot) + eps = np.array([max(abs(v), 1.0) * 1e-6 for v in param_vals]) + + J_fd = np.zeros((6, 6)) + for i in range(6): + cart_plus = np.asarray(bk_to_cartesian(_bk_perturb(bk, i, eps[i]), fid)).flatten() + cart_minus = np.asarray(bk_to_cartesian(_bk_perturb(bk, i, -eps[i]), fid)).flatten() + J_fd[:, i] = (cart_plus - cart_minus) / (2.0 * eps[i]) + + # Relative tolerance covering the dynamic range of J entries. + scale = np.maximum(np.abs(J_analytic), np.abs(J_fd)) + scale = np.where(scale > 0, scale, 1.0) + np.testing.assert_array_less( + np.abs(J_analytic - J_fd) / scale, + 1e-5, + err_msg="Analytic dcart_dbk disagrees with finite-difference", + ) + + +# --------------------------------------------------------------------------- +# Mixed-partial symmetry of the second-derivative cross-terms +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("case", _BK_CASES) +def test_mixed_partial_symmetry_via_finite_difference(case): + """d(r,v)/(d alpha d beta) -- approached via FD of dcart_dbk -- is symmetric.""" + rng = np.random.default_rng(seed=22222) + fid = _make_fiducial(rng) + bk = _bk_from_tuple(case) + eps_a = max(abs(bk.alpha), 1.0) * 1e-6 + eps_b = max(abs(bk.beta), 1.0) * 1e-6 + + # FD of dcart_dbk's alpha column with respect to beta + J_plus_b = np.asarray(dcart_dbk(_bk_perturb(bk, 1, eps_b), fid)) + J_minus_b = np.asarray(dcart_dbk(_bk_perturb(bk, 1, -eps_b), fid)) + d2r_dadb = (J_plus_b[:, 0] - J_minus_b[:, 0]) / (2.0 * eps_b) + + # FD of dcart_dbk's beta column with respect to alpha + J_plus_a = np.asarray(dcart_dbk(_bk_perturb(bk, 0, eps_a), fid)) + J_minus_a = np.asarray(dcart_dbk(_bk_perturb(bk, 0, -eps_a), fid)) + d2r_dbda = (J_plus_a[:, 1] - J_minus_a[:, 1]) / (2.0 * eps_a) + + np.testing.assert_allclose(d2r_dadb, d2r_dbda, atol=1e-5, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# Fiducial-direction gauge invariance +# --------------------------------------------------------------------------- + + +def test_fiducial_gauge_invariance(): + """Two valid n0 choices give the same physical Cartesian orbit.""" + rng = np.random.default_rng(seed=33333) + # An arbitrary Cartesian state (40 AU object, small velocity) + r = np.array([20.0, 30.0, 5.0]) + v = np.array([-2e-4, 1e-4, 5e-5]) + cart = np.concatenate([r, v]) + + # Two different fiducial frames: one nearly aligned with r_hat, one tilted. + r_hat = r / np.linalg.norm(r) + fid_aligned = bk_choose_fiducial([r_hat]) + + tilt = np.array([0.0, 0.0, 1.0]) + fid_tilted = bk_choose_fiducial([r_hat + 0.3 * tilt]) + + bk_aligned = cartesian_to_bk(cart, fid_aligned) + bk_tilted = cartesian_to_bk(cart, fid_tilted) + + cart_back_aligned = np.asarray(bk_to_cartesian(bk_aligned, fid_aligned)).flatten() + cart_back_tilted = np.asarray(bk_to_cartesian(bk_tilted, fid_tilted)).flatten() + + np.testing.assert_allclose(cart_back_aligned, cart, rtol=1e-12, atol=1e-13) + np.testing.assert_allclose(cart_back_tilted, cart, rtol=1e-12, atol=1e-13) + + +# --------------------------------------------------------------------------- +# Special-case forms at alpha = beta = 0 (the fiducial direction) +# --------------------------------------------------------------------------- + + +def test_special_case_at_fiducial(): + """At alpha = beta = 0, rho_hat = n0, rho_hat_alpha = a, rho_hat_beta = b.""" + rng = np.random.default_rng(seed=44444) + fid = _make_fiducial(rng) + bk = BKState(alpha=0.0, beta=0.0, gamma=0.05, adot=0.0, bdot=0.0, gdot=0.0) + cart = np.asarray(bk_to_cartesian(bk, fid)).flatten() + + # Position at (alpha, beta) = (0, 0) is exactly (1/gamma) * n0. + expected_r = (1.0 / bk.gamma) * np.asarray(fid.n0) + np.testing.assert_allclose(cart[:3], expected_r, rtol=1e-13, atol=1e-13) + + # With all rates zero, velocity should be zero. + np.testing.assert_allclose(cart[3:], np.zeros(3), atol=1e-15) + + +def test_jacobian_at_fiducial(): + """At the fiducial direction, the Jacobian's top-left 3x3 block is + [a | b | -(1/gamma) * n0] (each as a column).""" + rng = np.random.default_rng(seed=55555) + fid = _make_fiducial(rng) + gamma = 0.025 + bk = BKState(alpha=0.0, beta=0.0, gamma=gamma, adot=0.0, bdot=0.0, gdot=0.0) + J = np.asarray(dcart_dbk(bk, fid)) + + expected_col_alpha = (1.0 / gamma) * np.asarray(fid.a) + expected_col_beta = (1.0 / gamma) * np.asarray(fid.b) + expected_col_gamma = -(1.0 / gamma**2) * np.asarray(fid.n0) + + np.testing.assert_allclose(J[:3, 0], expected_col_alpha, rtol=1e-13, atol=1e-15) + np.testing.assert_allclose(J[:3, 1], expected_col_beta, rtol=1e-13, atol=1e-15) + np.testing.assert_allclose(J[:3, 2], expected_col_gamma, rtol=1e-13, atol=1e-15) + + # Bottom-right block should equal the top-left block (same shape). + np.testing.assert_allclose(J[3:, 3:], J[:3, :3], rtol=1e-13, atol=1e-15) + + # Bottom-left block should be zero when adot = bdot = gdot = 0. + np.testing.assert_allclose(J[3:, :3], np.zeros((3, 3)), atol=1e-15) + + +# --------------------------------------------------------------------------- +# sigma_gdot_sq vs Cartesian-side energy-bound calculation +# --------------------------------------------------------------------------- + + +def test_sigma_gdot_sq_consistent_with_energy_bound(): + """sigma_gdot_sq matches the Cartesian-side bound on |gdot|^2 for a bound orbit. + + Derivation: the bound-orbit energy constraint 0.5 |v|^2 <= mu / |r| + rearranges, in BK at fixed (alpha, beta, gamma, adot, bdot), to + gdot^2 <= gamma^2 * (2 * mu * gamma^3 - adot^2 - bdot^2), + which is exactly the formula sigma_gdot_sq returns. + """ + rng = np.random.default_rng(seed=66666) + fid = _make_fiducial(rng) + # Pick (alpha, beta, gamma, adot, bdot) such that the orbit is bound for + # at least some gdot range. + bk = BKState(alpha=0.05, beta=-0.03, gamma=1.0 / 40.0, adot=4e-5, bdot=-3e-5, gdot=0.0) + sigma_sq = sigma_gdot_sq(bk, MU_SUN) + + # The bound-orbit constraint -> |v|^2 <= 2 * mu * gamma. At the BK state + # with gdot pinned to +sqrt(sigma_sq), the orbit should be exactly at the + # boundary (|v|^2 == 2 * mu * gamma). + assert sigma_sq > 0.0 and np.isfinite(sigma_sq) + bk_at_boundary = BKState( + alpha=bk.alpha, beta=bk.beta, gamma=bk.gamma, adot=bk.adot, bdot=bk.bdot, gdot=np.sqrt(sigma_sq) + ) + cart = np.asarray(bk_to_cartesian(bk_at_boundary, fid)).flatten() + r_norm = np.linalg.norm(cart[:3]) + v_norm_sq = np.dot(cart[3:], cart[3:]) + energy = 0.5 * v_norm_sq - MU_SUN / r_norm + # At the boundary, total energy = 0 (parabolic orbit). + np.testing.assert_allclose(energy, 0.0, atol=1e-12) + + +def test_sigma_gdot_sq_returns_inf_for_hyperbolic_tangentials(): + """When tangential rates already exceed escape velocity, sigma_gdot_sq + signals 'no prior' by returning +infinity.""" + bk = BKState(alpha=0.0, beta=0.0, gamma=1.0 / 50.0, adot=1e-3, bdot=1e-3, gdot=0.0) # huge rates at 50 AU + assert np.isinf(sigma_gdot_sq(bk, MU_SUN)) diff --git a/tests/layup/test_bk_everywhere.py b/tests/layup/test_bk_everywhere.py new file mode 100644 index 00000000..e1d20f1c --- /dev/null +++ b/tests/layup/test_bk_everywhere.py @@ -0,0 +1,235 @@ +"""Layer 3 engine-sweep tests for the universal BK fitter. + +Drives both engine='cartesian' and engine='bk_native' against the +diagnostic/scan dataset (outside the repo, at +``~/Dropbox/claude_layup/diagnostic/scan/truth/``) so the design +memory's prediction -- ``bk_native`` matches Cartesian across regimes +and shines on distant short arcs -- can be validated against real +ASSIST-integrated truth. + +These tests skip cleanly when either the ASSIST ephemeris or the +diagnostic scan data is unavailable, so machines without either +setup are unaffected. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import pytest + +from layup.orbitfit import _MU_SUN, _run_fit +from layup.routines import ( + FitResult, + Observation, + get_ephem, + run_bk_native_fit, + run_from_vector_with_initial_guess, +) + +# --------------------------------------------------------------------------- +# Environment guards +# --------------------------------------------------------------------------- + +CACHE = os.path.expanduser("~/Library/Caches/layup") +EPHEM_PLANETS = os.path.join(CACHE, "linux_p1550p2650.440") +EPHEM_SMALLBODIES = os.path.join(CACHE, "sb441-n16.bsp") +EPHEM_AVAILABLE = os.path.exists(EPHEM_PLANETS) and os.path.exists(EPHEM_SMALLBODIES) + +DIAGNOSTIC_SCAN = Path("~/Dropbox/claude_layup/diagnostic/scan/truth").expanduser() +DIAGNOSTIC_AVAILABLE = DIAGNOSTIC_SCAN.is_dir() + +pytestmark = pytest.mark.skipif( + not (EPHEM_AVAILABLE and DIAGNOSTIC_AVAILABLE), + reason=( + f"Skipping Layer 3 BK-everywhere tests: " + f"ephem at {CACHE} = {EPHEM_AVAILABLE}, " + f"diagnostic scan at {DIAGNOSTIC_SCAN} = {DIAGNOSTIC_AVAILABLE}." + ), +) + + +# --------------------------------------------------------------------------- +# Helpers for loading and converting diagnostic-scan cases +# --------------------------------------------------------------------------- + + +def _load_case(name: str) -> dict: + """Load a diagnostic-scan case by stem (e.g., 'classical_42AU_arc_007.00d').""" + with open(DIAGNOSTIC_SCAN / f"{name}.json") as f: + return json.load(f) + + +def _build_observations(case: dict) -> list: + """Convert a case's observation list into layup Observation objects.""" + obs_list = [] + sigma_arcsec = float(case["sigma_arcsec"]) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) + for o in case["observations"]: + # observer_state_AU is position-only; we fill velocity with zeros. + # Velocity is only used for second-order corrections (parallax, light- + # time second derivative) and the design here doesn't depend on it. + pos = list(o["observer_state_AU"]) + vel = [0.0, 0.0, 0.0] + obs_list.append( + Observation.from_astrometry( + ra=np.deg2rad(o["ra"]), + dec=np.deg2rad(o["dec"]), + epoch=float(o["jd_tdb"]), + observer_position=pos, + observer_velocity=vel, + ) + ) + # Per-observation astrometric uncertainty (in radians, matching sigma_arcsec). + obs_list[-1].ra_unc = sigma_rad + obs_list[-1].dec_unc = sigma_rad + return obs_list + + +def _truth_seed(case: dict) -> FitResult: + """Return a FitResult populated with the case's truth state at epoch.""" + fit = FitResult() + fit.state = list(map(float, case["truth_state_at_epoch"])) + fit.epoch = float(case["epoch_jd_tdb"]) + return fit + + +def _r_helio_AU(state) -> float: + return float(np.linalg.norm(state[:3])) + + +# --------------------------------------------------------------------------- +# Engine-sweep tests +# --------------------------------------------------------------------------- + + +# Well-arced cases. With ~30-60 day arcs of mainbelt or TNO objects, the +# data alone is enough to constrain the orbit; both engines should reach a +# state within sub-AU of truth and agree with each other. +WELL_ARCED_CASES = [ + "mainbelt_2.5AU_arc_030.00d", + "mainbelt_3.5AU_arc_060.00d", + "classical_42AU_arc_060.00d", +] + + +@pytest.mark.parametrize("case_name", WELL_ARCED_CASES) +def test_engine_sweep_well_arced_cases(case_name): + """With the truth state as the LM seed on a well-constrained arc, both + engines should converge near truth and agree with each other.""" + ephem = get_ephem(CACHE) + case = _load_case(case_name) + obs = _build_observations(case) + seed = _truth_seed(case) + truth = np.asarray(case["truth_state_at_epoch"]) + r_helio = _r_helio_AU(truth) + tol = 0.01 * r_helio # ~1% of heliocentric distance + + cart_res = _run_fit(ephem, seed, obs, "cartesian") + bk_res = _run_fit(ephem, seed, obs, "bk_native") + + assert cart_res.flag == 0, f"[{case_name}] Cartesian flag={cart_res.flag}" + assert bk_res.flag == 0, f"[{case_name}] BK flag={bk_res.flag}" + + cart_drift = np.linalg.norm(np.asarray(cart_res.state)[:3] - truth[:3]) + bk_drift = np.linalg.norm(np.asarray(bk_res.state)[:3] - truth[:3]) + assert cart_drift < tol, f"[{case_name}] Cartesian drift {cart_drift:.4f} > tol {tol:.4f}" + assert bk_drift < tol, f"[{case_name}] BK drift {bk_drift:.4f} > tol {tol:.4f}" + + # Engine agreement: BK and Cartesian should converge to nearly the same point. + state_disagreement = np.linalg.norm(np.asarray(bk_res.state)[:3] - np.asarray(cart_res.state)[:3]) + assert state_disagreement < tol, f"[{case_name}] BK and Cartesian disagree by {state_disagreement:.4f} AU" + + +# Short-arc / distant cases. These are the cases that motivated BK in the +# first place: the line-of-sight direction is poorly constrained, so the +# Cartesian fit's LM step can walk significantly along that direction. We +# test that BK does at least as well as Cartesian in this regime, AND that +# BK uses substantially fewer LM iterations (the BK basis is better +# conditioned, so the Marquardt damping doesn't need to fight as hard). +SHORT_ARC_DISTANT_CASES = [ + "scattered_70AU_arc_014.00d", + "classical_42AU_arc_010.00d", +] + + +@pytest.mark.parametrize("case_name", SHORT_ARC_DISTANT_CASES) +def test_bk_beats_cartesian_on_short_arc_distant(case_name): + """In the distant short-arc regime where the line-of-sight is poorly + constrained, BK should drift no more than Cartesian from truth and use + substantially fewer LM iterations.""" + ephem = get_ephem(CACHE) + case = _load_case(case_name) + obs = _build_observations(case) + seed = _truth_seed(case) + truth = np.asarray(case["truth_state_at_epoch"]) + + cart_res = _run_fit(ephem, seed, obs, "cartesian") + bk_res = _run_fit(ephem, seed, obs, "bk_native") + + assert cart_res.flag == 0, f"[{case_name}] Cartesian flag={cart_res.flag}" + assert bk_res.flag == 0, f"[{case_name}] BK flag={bk_res.flag}" + + cart_drift = np.linalg.norm(np.asarray(cart_res.state)[:3] - truth[:3]) + bk_drift = np.linalg.norm(np.asarray(bk_res.state)[:3] - truth[:3]) + + # BK should drift no more than Cartesian. (In practice it's often + # *much* less -- e.g. on scattered_70AU_arc_014.00d BK stays ~0.02 AU + # from truth while Cartesian wanders ~4.5 AU.) + assert ( + bk_drift <= cart_drift + 1e-6 + ), f"[{case_name}] BK drift {bk_drift:.4f} > Cartesian drift {cart_drift:.4f}" + + # BK should use fewer LM iterations than Cartesian: the BK basis is + # naturally better-conditioned than the Cartesian state at epoch for + # short-arc distant targets, so the LM step direction is healthier. + assert ( + bk_res.niter < cart_res.niter + ), f"[{case_name}] BK niter={bk_res.niter} not < Cartesian niter={cart_res.niter}" + + +def test_engine_sweep_produces_method_strings(): + """Sanity: each engine populates FitResult.method with its tag, so a + downstream sweep harness can tell which engine produced each fit.""" + ephem = get_ephem(CACHE) + case = _load_case("classical_42AU_arc_060.00d") + obs = _build_observations(case) + seed = _truth_seed(case) + + cart_res = _run_fit(ephem, seed, obs, "cartesian") + bk_res = _run_fit(ephem, seed, obs, "bk_native") + assert cart_res.method == "orbit_fit" + assert bk_res.method == "bk_native" + + +# --------------------------------------------------------------------------- +# Diagnostic helper (not a test) -- used by sweep harness scripts. +# --------------------------------------------------------------------------- + + +def sweep_cases_from_diagnostic(case_names=None) -> list: + """Return a list of (case_name, cartesian FitResult, bk_native FitResult) + tuples for the requested case names (or all 98 cases if None). + + Intended for ad-hoc use from a sweep script that produces tables; not + invoked by pytest collection. + """ + if case_names is None: + case_names = sorted(p.stem for p in DIAGNOSTIC_SCAN.glob("*.json")) + ephem = get_ephem(CACHE) + rows = [] + for name in case_names: + case = _load_case(name) + obs = _build_observations(case) + seed = _truth_seed(case) + rows.append( + ( + name, + _run_fit(ephem, seed, obs, "cartesian"), + _run_fit(ephem, seed, obs, "bk_native"), + ) + ) + return rows diff --git a/tests/layup/test_bk_fit.py b/tests/layup/test_bk_fit.py new file mode 100644 index 00000000..f22eb4bb --- /dev/null +++ b/tests/layup/test_bk_fit.py @@ -0,0 +1,300 @@ +"""Layer 2 tests for the universal BK fitter (`run_bk_native_fit`). + +These tests cover the LM driver itself. They reuse the same Gauss IOD + +observation setup as the existing Cartesian fit so the only difference +between the two engines is the parameterization + the energy prior on +gdot, isolating any disagreement to the BK-specific code path. + +Tests skip when the ASSIST ephemeris files aren't available, so CI on +machines without `~/Library/Caches/layup/{linux_p1550p2650.440, +sb441-n16.bsp}` is unaffected. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from layup.routines import ( + FitResult, + Observation, + get_ephem, + predict_sequence, + run_bk_native_fit, + run_from_vector_with_initial_guess, +) + +CACHE = os.path.expanduser("~/Library/Caches/layup") +EPHEM_PLANETS = os.path.join(CACHE, "linux_p1550p2650.440") +EPHEM_SMALLBODIES = os.path.join(CACHE, "sb441-n16.bsp") +EPHEM_AVAILABLE = os.path.exists(EPHEM_PLANETS) and os.path.exists(EPHEM_SMALLBODIES) + +# GM_sun in AU^3 / day^2. +MU_SUN = 0.00029591220828559104 + +pytestmark = pytest.mark.skipif( + not EPHEM_AVAILABLE, + reason=f"ASSIST ephemeris missing at {CACHE}; skipping BK-fit Layer 2 tests.", +) + + +# --------------------------------------------------------------------------- +# Tests that don't need real observations -- exercise the API + early-exit +# guards. +# --------------------------------------------------------------------------- + + +def test_run_bk_native_fit_returns_fitresult_for_empty_obs(): + """With zero observations, the fit returns a FitResult with flag != 0 and + does not crash.""" + ephem = get_ephem(CACHE) + ig = FitResult() + ig.state = [40.0, 10.0, 5.0, -8e-4, 9e-4, 1e-4] + ig.epoch = 2460000.5 + result = run_bk_native_fit(ephem, ig, [], MU_SUN) + assert result.method == "bk_native" + assert result.flag != 0 + + +def test_run_bk_native_fit_returns_fitresult_for_too_few_obs(): + """With <3 observations the early-exit guard fires; no crash, flag != 0.""" + ephem = get_ephem(CACHE) + ig = FitResult() + ig.state = [40.0, 10.0, 5.0, -8e-4, 9e-4, 1e-4] + ig.epoch = 2460000.5 + obs = [ + Observation.from_astrometry( + ra=1.57, + dec=0.1, + epoch=2459995.5, + observer_position=[-0.5, 0.8, 0.0], + observer_velocity=[-0.018, -0.009, 0.0], + ), + Observation.from_astrometry( + ra=1.57, + dec=0.1, + epoch=2460005.5, + observer_position=[-0.5, 0.8, 0.0], + observer_velocity=[-0.018, -0.009, 0.0], + ), + ] + result = run_bk_native_fit(ephem, ig, obs, MU_SUN) + assert result.method == "bk_native" + assert result.flag != 0 + + +# --------------------------------------------------------------------------- +# Synthetic-orbit convergence tests. +# +# We pick a known Cartesian state, generate synthetic observations from it +# via the layup C++ predict path, then feed those observations back into +# both run_bk_native_fit and run_from_vector_with_initial_guess and check +# that (a) BK converges, (b) BK recovers the input state, and (c) the BK +# and Cartesian fits agree at convergence. +# --------------------------------------------------------------------------- + + +def _generate_synthetic_observations(ephem, truth_state, truth_epoch, obs_times): + """Generate synthetic Observation objects consistent with `truth_state` + at `truth_epoch`, observed from a fixed point (Sun in barycentric coords) + at each of `obs_times`. + + Returns a list of Observation objects whose epochs and rho_hat directions + match what the truth orbit predicts. + """ + # Use a fixed observer at the solar system barycenter so the only + # dynamical content in the synthetic data is the orbit itself. The + # observer-velocity is zero -- consistent with a barycenter "observer." + observer_position = [0.0, 0.0, 0.0] + observer_velocity = [0.0, 0.0, 0.0] + + # Template observations at the desired times with dummy (ra, dec) + template = [ + Observation.from_astrometry( + ra=0.0, + dec=0.0, + epoch=float(t), + observer_position=observer_position, + observer_velocity=observer_velocity, + ) + for t in obs_times + ] + + # Run predict on each template; the FitResult holds the truth state at epoch. + truth_fit = FitResult() + truth_fit.state = list(map(float, truth_state)) + truth_fit.epoch = float(truth_epoch) + cov = np.zeros((6, 6)) + preds = predict_sequence(ephem, truth_fit, template, cov) + + # Build real Observations from the predicted rho unit vectors. + synth = [] + for t, pr in zip(obs_times, preds): + rho = np.asarray(pr.rho) + # Defensive normalization (predict returns a unit vector already). + rho = rho / np.linalg.norm(rho) + ra = np.arctan2(rho[1], rho[0]) + dec = np.arcsin(np.clip(rho[2], -1.0, 1.0)) + synth.append( + Observation.from_astrometry( + ra=float(ra), + dec=float(dec), + epoch=float(t), + observer_position=observer_position, + observer_velocity=observer_velocity, + ) + ) + return synth + + +def _seed_from_state(state, epoch): + fit = FitResult() + fit.state = list(map(float, state)) + fit.epoch = float(epoch) + return fit + + +@pytest.mark.parametrize( + "name, state, arc_days, nobs", + [ + # ~3 AU mainbelt-ish (well-constrained) + ("mainbelt_3au_60d", [3.0, 0.0, 0.0, 0.0, 0.0102, 0.001], 60.0, 12), + # ~40 AU TNO, longer arc + ("tno_40au_300d", [40.0, 0.0, 5.0, 0.0, 0.00125, 0.0], 300.0, 12), + ], +) +def test_bk_native_fit_recovers_known_state(name, state, arc_days, nobs): + """Synthetic obs from a known state, fitted with BK from a perfect seed, + should recover the input state and produce a tiny chi-square.""" + ephem = get_ephem(CACHE) + truth_epoch = 2460000.5 + obs_times = np.linspace(truth_epoch - 0.5 * arc_days, truth_epoch + 0.5 * arc_days, nobs) + + obs = _generate_synthetic_observations(ephem, state, truth_epoch, obs_times) + seed = _seed_from_state(state, truth_epoch) + + result = run_bk_native_fit(ephem, seed, obs, MU_SUN) + assert result.flag == 0, f"[{name}] BK fit did not converge (flag={result.flag})" + np.testing.assert_allclose( + np.asarray(result.state), + np.asarray(state), + rtol=1e-6, + atol=1e-9, + err_msg=f"[{name}] BK fit did not recover the truth state", + ) + # 2N residuals, 6 free params, noise-free obs -> chi2 essentially zero. + assert result.csq < 1e-12, f"[{name}] BK fit chi-square unexpectedly large: {result.csq}" + + +@pytest.mark.parametrize( + "name, state, arc_days, nobs, rel_perturb", + [ + # Modest perturbation -- exercises the LM loop without falling out of basin. + ("mainbelt_3au_60d_pert", [3.0, 0.0, 0.0, 0.0, 0.0102, 0.001], 60.0, 12, 1e-3), + ("tno_40au_300d_pert", [40.0, 0.0, 5.0, 0.0, 0.00125, 0.0], 300.0, 12, 1e-3), + ], +) +def test_bk_native_fit_recovers_from_perturbed_seed(name, state, arc_days, nobs, rel_perturb): + """With a 0.1%-perturbed seed, the LM loop still converges to the truth state.""" + ephem = get_ephem(CACHE) + truth_epoch = 2460000.5 + obs_times = np.linspace(truth_epoch - 0.5 * arc_days, truth_epoch + 0.5 * arc_days, nobs) + + obs = _generate_synthetic_observations(ephem, state, truth_epoch, obs_times) + + # Perturb each component by rel_perturb * |component| (deterministic, no RNG). + perturbed = np.asarray(state) * (1.0 + rel_perturb) + seed = _seed_from_state(perturbed.tolist(), truth_epoch) + + result = run_bk_native_fit(ephem, seed, obs, MU_SUN) + assert result.flag == 0, f"[{name}] BK fit did not converge (flag={result.flag})" + np.testing.assert_allclose( + np.asarray(result.state), + np.asarray(state), + rtol=1e-6, + atol=1e-9, + err_msg=f"[{name}] BK fit did not recover truth from perturbed seed", + ) + # niter should be > 1 since we actually had to iterate. + assert result.niter >= 1, f"[{name}] niter={result.niter} -- expected at least 1" + + +@pytest.mark.parametrize( + "name, state, arc_days, nobs", + [ + ("mainbelt_3au_60d", [3.0, 0.0, 0.0, 0.0, 0.0102, 0.001], 60.0, 12), + ], +) +def test_bk_and_cartesian_fits_agree(name, state, arc_days, nobs): + """For well-constrained synthetic observations, the BK and Cartesian + engines should converge to states that match to within numerical noise.""" + ephem = get_ephem(CACHE) + truth_epoch = 2460000.5 + obs_times = np.linspace(truth_epoch - 0.5 * arc_days, truth_epoch + 0.5 * arc_days, nobs) + + obs = _generate_synthetic_observations(ephem, state, truth_epoch, obs_times) + seed = _seed_from_state(state, truth_epoch) + + bk_result = run_bk_native_fit(ephem, seed, obs, MU_SUN) + cart_result = run_from_vector_with_initial_guess(ephem, seed, obs) + + assert bk_result.flag == 0, f"[{name}] BK fit failed: {bk_result.flag}" + assert cart_result.flag == 0, f"[{name}] Cartesian fit failed: {cart_result.flag}" + np.testing.assert_allclose( + np.asarray(bk_result.state), + np.asarray(cart_result.state), + rtol=1e-6, + atol=1e-9, + err_msg=f"[{name}] BK and Cartesian fits disagree at convergence", + ) + + +# --------------------------------------------------------------------------- +# Engine dispatch through orbitfit._run_fit +# --------------------------------------------------------------------------- + + +def test_run_fit_dispatch_cartesian(): + """orbitfit._run_fit(engine='cartesian') matches direct + run_from_vector_with_initial_guess.""" + from layup.orbitfit import _run_fit + + ephem = get_ephem(CACHE) + state = [3.0, 0.0, 0.0, 0.0, 0.0102, 0.001] + epoch = 2460000.5 + obs = _generate_synthetic_observations(ephem, state, epoch, np.linspace(epoch - 30, epoch + 30, 12)) + seed = _seed_from_state(state, epoch) + + via_dispatch = _run_fit(ephem, seed, obs, "cartesian") + direct = run_from_vector_with_initial_guess(ephem, seed, obs) + np.testing.assert_array_equal(via_dispatch.state, direct.state) + assert via_dispatch.method == direct.method + + +def test_run_fit_dispatch_bk_native(): + """orbitfit._run_fit(engine='bk_native') matches direct + run_bk_native_fit with MU_SUN.""" + from layup.orbitfit import _MU_SUN, _run_fit + + ephem = get_ephem(CACHE) + state = [3.0, 0.0, 0.0, 0.0, 0.0102, 0.001] + epoch = 2460000.5 + obs = _generate_synthetic_observations(ephem, state, epoch, np.linspace(epoch - 30, epoch + 30, 12)) + seed = _seed_from_state(state, epoch) + + via_dispatch = _run_fit(ephem, seed, obs, "bk_native") + direct = run_bk_native_fit(ephem, seed, obs, _MU_SUN) + np.testing.assert_array_equal(via_dispatch.state, direct.state) + assert via_dispatch.method == "bk_native" + + +def test_run_fit_dispatch_unknown_engine_raises(): + """An unrecognized engine name raises ValueError.""" + from layup.orbitfit import _run_fit + + ephem = get_ephem(CACHE) + seed = _seed_from_state([3.0, 0.0, 0.0, 0.0, 0.01, 0.0], 2460000.5) + with pytest.raises(ValueError, match="Unknown engine"): + _run_fit(ephem, seed, [], "not_an_engine") diff --git a/tests/layup/test_bk_iod.py b/tests/layup/test_bk_iod.py new file mode 100644 index 00000000..48468f54 --- /dev/null +++ b/tests/layup/test_bk_iod.py @@ -0,0 +1,176 @@ +"""Tests for the universal-BK 5-parameter linear IOD (`run_bk_iod`). + +The IOD is the layup-side analog of liborbfit's `prelim_fit`: a single +closed-form weighted least-squares solve over (alpha, beta, gamma, +adot, bdot) with gdot pinned to 0. See `bk_iod.cpp` for the model and +the documented regime of validity (works best for distant objects, +single-percent on heliocentric distance for TNOs at sweet-spot arc +lengths; not intended for mainbelt or as a final orbit). + +Test layers: + * Smoke: empty / few-obs guards return without crashing. + * Sweet-spot diagnostic: on a representative distant case, BK-IOD + recovers truth to within a few percent of heliocentric distance. + * Seeds the LM to truth: the BK-IOD output, fed into + run_bk_native_fit, converges to truth at rtol=1e-6 -- the actual + intended use of BK-IOD. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import pytest + +from layup.routines import ( + FitResult, + Observation, + get_ephem, + run_bk_iod, + run_bk_native_fit, +) + +CACHE = os.path.expanduser("~/Library/Caches/layup") +EPHEM_PLANETS = os.path.join(CACHE, "linux_p1550p2650.440") +EPHEM_SMALLBODIES = os.path.join(CACHE, "sb441-n16.bsp") +EPHEM_AVAILABLE = os.path.exists(EPHEM_PLANETS) and os.path.exists(EPHEM_SMALLBODIES) + +DIAGNOSTIC_SCAN = Path("~/Dropbox/claude_layup/diagnostic/scan/truth").expanduser() +DIAGNOSTIC_AVAILABLE = DIAGNOSTIC_SCAN.is_dir() + +MU_SUN = 0.00029591220828559104 + + +# --------------------------------------------------------------------------- +# Smoke tests -- no ephemeris or diagnostic data needed +# --------------------------------------------------------------------------- + + +def test_run_bk_iod_empty_obs(): + """No observations -> flag != 0, no crash.""" + result = run_bk_iod([], 2460000.5, MU_SUN) + assert result.method == "bk_iod" + assert result.flag != 0 + + +def test_run_bk_iod_too_few_obs(): + """<3 observations -> flag != 0, no crash.""" + obs = [ + Observation.from_astrometry(1.57, 0.1, 2460000.5, [-0.5, 0.8, 0.0], [0.0, 0.0, 0.0]), + Observation.from_astrometry(1.57, 0.1, 2460010.5, [-0.5, 0.8, 0.0], [0.0, 0.0, 0.0]), + ] + result = run_bk_iod(obs, 2460000.5, MU_SUN) + assert result.flag != 0 + + +# --------------------------------------------------------------------------- +# Diagnostic-data tests -- skip if scan + ephem not available +# --------------------------------------------------------------------------- + + +pytestmark_diagnostic = pytest.mark.skipif( + not (EPHEM_AVAILABLE and DIAGNOSTIC_AVAILABLE), + reason="Need both ephemeris and diagnostic scan data.", +) + + +def _load_diagnostic_case(name): + with open(DIAGNOSTIC_SCAN / f"{name}.json") as f: + return json.load(f) + + +def _build_observations_from_case(case): + sigma_arcsec = float(case["sigma_arcsec"]) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) + obs_list = [] + for o in case["observations"]: + pos = list(o["observer_state_AU"]) + vel = [0.0, 0.0, 0.0] + obs = Observation.from_astrometry( + ra=np.deg2rad(o["ra"]), + dec=np.deg2rad(o["dec"]), + epoch=float(o["jd_tdb"]), + observer_position=pos, + observer_velocity=vel, + ) + obs.ra_unc = sigma_rad + obs.dec_unc = sigma_rad + obs_list.append(obs) + return obs_list + + +@pytestmark_diagnostic +@pytest.mark.parametrize( + "case_name, max_drift_frac", + [ + # Tolerances chosen by the empirical sweep in bk_iod.cpp's docstring: + # distant objects in their sweet-spot arc length recover r_helio to + # within a few percent. + ("classical_42AU_arc_010.00d", 0.05), + ("scattered_70AU_arc_007.00d", 0.02), + ("sednoid_80AU_arc_010.00d", 0.03), + # Within-regime longer arcs -- still acceptable, looser bound. + ("classical_42AU_arc_060.00d", 0.08), + ], +) +def test_bk_iod_distant_objects(case_name, max_drift_frac): + """BK-IOD on a distant case should recover the truth heliocentric + position to within a few percent (regime-of-validity expectation).""" + case = _load_diagnostic_case(case_name) + obs = _build_observations_from_case(case) + truth = np.asarray(case["truth_state_at_epoch"]) + epoch = float(case["epoch_jd_tdb"]) + r_helio = float(np.linalg.norm(truth[:3])) + + result = run_bk_iod(obs, epoch, MU_SUN) + assert result.flag == 0, f"[{case_name}] BK-IOD did not converge (flag={result.flag})" + drift = np.linalg.norm(np.asarray(result.state)[:3] - truth[:3]) + assert drift < max_drift_frac * r_helio, ( + f"[{case_name}] BK-IOD drifted {drift:.3f} AU " f"> {max_drift_frac:.0%} of r_helio={r_helio:.1f} AU" + ) + + +# --------------------------------------------------------------------------- +# IOD's intended use: seeding the full BK LM fit +# --------------------------------------------------------------------------- + + +@pytestmark_diagnostic +@pytest.mark.parametrize( + "case_name", + [ + "classical_42AU_arc_010.00d", + "scattered_70AU_arc_007.00d", + "sednoid_80AU_arc_010.00d", + "classical_42AU_arc_060.00d", + ], +) +def test_bk_iod_seeds_lm_to_truth(case_name): + """The actual purpose of BK-IOD: produce a seed that, fed into + run_bk_native_fit, converges to the truth state. This is the + end-to-end test of "is BK-IOD useful?" -- and the answer should + be yes even on cases where the IOD itself sits a few percent off + the truth, because LM convergence basins are wider than that.""" + ephem = get_ephem(CACHE) + case = _load_diagnostic_case(case_name) + obs = _build_observations_from_case(case) + truth = np.asarray(case["truth_state_at_epoch"]) + epoch = float(case["epoch_jd_tdb"]) + r_helio = float(np.linalg.norm(truth[:3])) + + iod = run_bk_iod(obs, epoch, MU_SUN) + assert iod.flag == 0, f"[{case_name}] IOD failed (flag={iod.flag})" + + # Seed the LM with the IOD result and let it converge. + fit = run_bk_native_fit(ephem, iod, obs, MU_SUN) + assert fit.flag == 0, f"[{case_name}] LM (seeded by IOD) did not converge (flag={fit.flag})" + + # LM should land near truth (sub-AU on a sub-AU-noise dataset). + drift = np.linalg.norm(np.asarray(fit.state)[:3] - truth[:3]) + assert drift < 0.01 * r_helio, ( + f"[{case_name}] LM (IOD seed) drifted {drift:.3f} AU " + f"= {100 * drift / r_helio:.2f}% of r_helio={r_helio:.1f} AU" + ) diff --git a/tests/layup/test_iod_auto.py b/tests/layup/test_iod_auto.py new file mode 100644 index 00000000..7d6a353a --- /dev/null +++ b/tests/layup/test_iod_auto.py @@ -0,0 +1,130 @@ +"""Tests for `do_fit(iod='auto')`: Gauss first, falling back to the BK +5-parameter linear IOD if every Gauss root fails to seed the LM. + +The empirical motivation is the sweep in `tools/bk_iod_sweep.py`, +which showed Gauss + BK-IOD fallback covers 90/98 cases on the +diagnostic-scan dataset vs Gauss alone (82) or BK-IOD alone (72). +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import pytest + +from layup.orbitfit import do_fit +from layup.routines import Observation + +CACHE = os.path.expanduser("~/Library/Caches/layup") +EPHEM_PLANETS = os.path.join(CACHE, "linux_p1550p2650.440") +EPHEM_SMALLBODIES = os.path.join(CACHE, "sb441-n16.bsp") +EPHEM_AVAILABLE = os.path.exists(EPHEM_PLANETS) and os.path.exists(EPHEM_SMALLBODIES) + +DIAGNOSTIC_SCAN = Path("~/Dropbox/claude_layup/diagnostic/scan/truth").expanduser() +DIAGNOSTIC_AVAILABLE = DIAGNOSTIC_SCAN.is_dir() + + +def _build_observations(case_dict): + sigma_arcsec = float(case_dict["sigma_arcsec"]) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) + obs_list = [] + for o in case_dict["observations"]: + pos = list(o["observer_state_AU"]) + vel = [0.0, 0.0, 0.0] + obs = Observation.from_astrometry( + ra=np.deg2rad(o["ra"]), + dec=np.deg2rad(o["dec"]), + epoch=float(o["jd_tdb"]), + observer_position=pos, + observer_velocity=vel, + ) + obs.ra_unc = sigma_rad + obs.dec_unc = sigma_rad + obs_list.append(obs) + return obs_list + + +def _load(name): + return json.load(open(DIAGNOSTIC_SCAN / f"{name}.json")) + + +# --------------------------------------------------------------------------- +# Dispatch validation -- no ephemeris needed +# --------------------------------------------------------------------------- + + +def test_do_fit_unknown_iod_raises(): + """A typo'd iod value raises ValueError before any fitting starts.""" + with pytest.raises(ValueError, match="not supported"): + do_fit(observations=[], seq=[[]], cache_dir=CACHE, iod="not_a_real_iod") + + +# --------------------------------------------------------------------------- +# Fallback behavior -- needs ephemeris + diagnostic data +# --------------------------------------------------------------------------- + + +pytestmark_diagnostic = pytest.mark.skipif( + not (EPHEM_AVAILABLE and DIAGNOSTIC_AVAILABLE), + reason="Need both ephemeris and diagnostic scan data.", +) + + +@pytestmark_diagnostic +def test_iod_auto_matches_gauss_on_easy_case(): + """On a well-conditioned mainbelt case, iod='auto' should produce the + same result as iod='gauss' (the fallback path never fires).""" + case = _load("mainbelt_2.5AU_arc_030.00d") + obs = _build_observations(case) + seq = [list(range(len(obs)))] + + res_gauss = do_fit(obs, seq, CACHE, iod="gauss") + res_auto = do_fit(obs, seq, CACHE, iod="auto") + assert res_gauss.flag == 0 + assert res_auto.flag == 0 + np.testing.assert_allclose(np.asarray(res_auto.state), np.asarray(res_gauss.state), rtol=1e-9, atol=1e-12) + + +@pytestmark_diagnostic +def test_iod_auto_recovers_when_gauss_fails(): + """On a case where every Gauss root fails to seed the LM, iod='auto' + falls back to BK-IOD and produces a converged fit. + + The case `sednoid_80AU_arc_001.00d` is one we know `do_fit` with + iod='gauss' fails on (flag=3, 'primary interval failed') under its + default Cartesian LM -- the 80 AU object with a 1-day arc has a + degenerate Gauss geometry. iod='auto' picks up the BK-IOD fallback + and recovers (flag=0).""" + case = _load("sednoid_80AU_arc_001.00d") + obs = _build_observations(case) + seq = [list(range(len(obs)))] + + res_gauss = do_fit(obs, seq, CACHE, iod="gauss") + res_auto = do_fit(obs, seq, CACHE, iod="auto") + + # The gauss path fails (flag != 0). + assert res_gauss.flag != 0, ( + f"Test premise broken: this case was supposed to be 'gauss fails'." + f" Got gauss flag={res_gauss.flag}. Maybe pick a different case." + ) + # The auto path recovers via BK-IOD fallback. + assert res_auto.flag == 0, ( + f"iod='auto' did not recover: flag={res_auto.flag}, " f"method={res_auto.method}" + ) + + +@pytestmark_diagnostic +def test_iod_auto_engine_choice_propagates(): + """The engine parameter is independent of iod choice -- iod='auto' with + engine='bk_native' should work just like the Cartesian engine.""" + case = _load("classical_42AU_arc_010.00d") + obs = _build_observations(case) + seq = [list(range(len(obs)))] + + res_cart = do_fit(obs, seq, CACHE, iod="auto", engine="cartesian") + res_bk = do_fit(obs, seq, CACHE, iod="auto", engine="bk_native") + assert res_cart.flag == 0 + assert res_bk.flag == 0 diff --git a/tools/bk_engine_sweep.py b/tools/bk_engine_sweep.py new file mode 100644 index 00000000..e94ee26c --- /dev/null +++ b/tools/bk_engine_sweep.py @@ -0,0 +1,350 @@ +"""Engine-sweep harness: run both engine='cartesian' and engine='bk_native' +on every case in a diagnostic scan dataset, write a CSV with per-case +metrics, and print a population-level summary table. + +Each case is seeded with the truth state at epoch, so this is a "given +a perfect starting point, how does each LM behave on the data" comparison +rather than an IOD-recovery test. The data themselves carry the diagnostic +scan's pre-baked Gaussian noise (sigma_arcsec is read from each JSON). + +Expected input format: a directory containing one .json file per case, +each with at least the following keys:: + + { + "population": "...", + "arc_length_days": ..., + "epoch_jd_tdb": ..., + "sigma_arcsec": ..., + "truth_state_at_epoch": [x, y, z, vx, vy, vz], + "observations": [ + {"ra": , "dec": , "jd_tdb": ..., + "observer_state_AU": [x, y, z], ...}, + ... + ] + } + +The diagnostic-scan dataset shipping with the project lives at +``~/Dropbox/claude_layup/diagnostic/scan/truth/`` (98 cases, 7 +populations x 14 arc lengths). + +Usage:: + + python tools/bk_engine_sweep.py --scan-dir ~/path/to/truth/ \\ + --cache-dir ~/Library/Caches/layup \\ + --output bk_engine_sweep.csv + +Defaults: + --scan-dir ~/Dropbox/claude_layup/diagnostic/scan/truth + --cache-dir ~/Library/Caches/layup + --output bk_engine_sweep.csv (in the cwd) +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +import sys +from pathlib import Path +from typing import List + +import numpy as np + +from layup.orbitfit import _MU_SUN +from layup.routines import ( + FitResult, + Observation, + get_ephem, + run_bk_native_fit, + run_from_vector_with_initial_guess, +) + +DEFAULT_SCAN_DIR = "~/Dropbox/claude_layup/diagnostic/scan/truth" +DEFAULT_CACHE_DIR = "~/Library/Caches/layup" +DEFAULT_OUTPUT = "bk_engine_sweep.csv" + + +# ---------------------------------------------------------------------- +# Case loading / observation construction +# ---------------------------------------------------------------------- + + +def load_case(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +def build_observations(case: dict) -> list: + """Construct layup Observations from a case dict's observation list. + + `observer_state_AU` is treated as position-only (velocity zero); the + layup fit pipeline only uses observer velocity for second-order + corrections that don't affect the chain-rule comparison here. + """ + sigma_arcsec = float(case["sigma_arcsec"]) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) + obs_list = [] + for o in case["observations"]: + pos = list(o["observer_state_AU"]) + vel = [0.0, 0.0, 0.0] + obs = Observation.from_astrometry( + ra=np.deg2rad(o["ra"]), + dec=np.deg2rad(o["dec"]), + epoch=float(o["jd_tdb"]), + observer_position=pos, + observer_velocity=vel, + ) + obs.ra_unc = sigma_rad + obs.dec_unc = sigma_rad + obs_list.append(obs) + return obs_list + + +def truth_seed(case: dict) -> FitResult: + fit = FitResult() + fit.state = list(map(float, case["truth_state_at_epoch"])) + fit.epoch = float(case["epoch_jd_tdb"]) + return fit + + +# ---------------------------------------------------------------------- +# Per-case sweep +# ---------------------------------------------------------------------- + + +def sweep_one(ephem, case_path: Path) -> dict: + case = load_case(case_path) + obs = build_observations(case) + seed = truth_seed(case) + truth = np.asarray(case["truth_state_at_epoch"]) + + cart = run_from_vector_with_initial_guess(ephem, seed, obs) + bk = run_bk_native_fit(ephem, seed, obs, _MU_SUN) + + r_helio = float(np.linalg.norm(truth[:3])) + cart_drift = float(np.linalg.norm(np.asarray(cart.state)[:3] - truth[:3])) + bk_drift = float(np.linalg.norm(np.asarray(bk.state)[:3] - truth[:3])) + + return { + "case": case_path.stem, + "population": case["population"], + "arc_days": float(case["arc_length_days"]), + "n_obs": len(case["observations"]), + "r_helio_AU": r_helio, + "cart_flag": int(cart.flag), + "cart_niter": int(cart.niter), + "cart_csq": float(cart.csq), + "cart_drift_AU": cart_drift, + "bk_flag": int(bk.flag), + "bk_niter": int(bk.niter), + "bk_csq": float(bk.csq), + "bk_drift_AU": bk_drift, + } + + +# ---------------------------------------------------------------------- +# Reporting +# ---------------------------------------------------------------------- + + +def write_csv(rows: List[dict], path: Path) -> None: + if not rows: + return + fieldnames = list(rows[0].keys()) + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in rows: + writer.writerow(row) + print(f"Wrote {len(rows)} rows to {path}") + + +def _fmt_drift(d: float) -> str: + if d < 1e-6: + return f"{d:.2e}" + return f"{d:.4f}" + + +def print_summary(rows: List[dict]) -> None: + print() + print("=" * 96) + print( + f"{'Case':<40s} {'r_helio':>8s} {'arc':>6s} " + f"{'cart_drift':>12s} {'bk_drift':>12s} {'cart_it':>7s} {'bk_it':>5s}" + ) + print("=" * 96) + pop_summary: dict = {} + for row in rows: + pop = row["population"] + s = pop_summary.setdefault( + pop, + { + "n": 0, + "bk_better": 0, + "cart_better": 0, + "tie": 0, + "both_failed": 0, + "only_cart_failed": 0, + "only_bk_failed": 0, + "cart_total_iter": 0, + "bk_total_iter": 0, + }, + ) + s["n"] += 1 + cf = row["cart_flag"] + bf = row["bk_flag"] + if cf != 0 and bf != 0: + s["both_failed"] += 1 + elif cf != 0: + s["only_cart_failed"] += 1 + elif bf != 0: + s["only_bk_failed"] += 1 + elif row["bk_drift_AU"] < row["cart_drift_AU"]: + s["bk_better"] += 1 + elif row["cart_drift_AU"] < row["bk_drift_AU"]: + s["cart_better"] += 1 + else: + s["tie"] += 1 + if cf == 0: + s["cart_total_iter"] += row["cart_niter"] + if bf == 0: + s["bk_total_iter"] += row["bk_niter"] + + print( + f"{row['case']:<40s} {row['r_helio_AU']:>8.2f} " + f"{row['arc_days']:>6.2f} " + f"{_fmt_drift(row['cart_drift_AU']):>12s} " + f"{_fmt_drift(row['bk_drift_AU']):>12s} " + f"{row['cart_niter']:>7d} {row['bk_niter']:>5d}" + ) + + print() + print("=" * 102) + print(f"Per-population summary ({len(rows)} cases total)") + print("=" * 102) + header = ( + f"{'Population':<20s} {'n':>3s} {'BK win':>7s} {'Cart win':>9s} " + f"{'cart fail':>10s} {'bk fail':>8s} {'both fail':>10s} " + f"{'mean cart it':>13s} {'mean bk it':>11s}" + ) + print(header) + print("-" * len(header)) + total = { + "n": 0, + "bk_better": 0, + "cart_better": 0, + "tie": 0, + "only_cart_failed": 0, + "only_bk_failed": 0, + "both_failed": 0, + "cart_total_iter": 0, + "bk_total_iter": 0, + } + for pop in sorted(pop_summary): + s = pop_summary[pop] + for k in total: + total[k] += s[k] + cart_succ = max(s["n"] - s["only_cart_failed"] - s["both_failed"], 0) + bk_succ = max(s["n"] - s["only_bk_failed"] - s["both_failed"], 0) + cart_mean_it = s["cart_total_iter"] / cart_succ if cart_succ else float("nan") + bk_mean_it = s["bk_total_iter"] / bk_succ if bk_succ else float("nan") + print( + f"{pop:<20s} {s['n']:>3d} {s['bk_better']:>7d} " + f"{s['cart_better']:>9d} {s['only_cart_failed']:>10d} " + f"{s['only_bk_failed']:>8d} {s['both_failed']:>10d} " + f"{cart_mean_it:>13.1f} {bk_mean_it:>11.1f}" + ) + print("-" * len(header)) + cart_succ = max(total["n"] - total["only_cart_failed"] - total["both_failed"], 0) + bk_succ = max(total["n"] - total["only_bk_failed"] - total["both_failed"], 0) + cart_mean = total["cart_total_iter"] / cart_succ if cart_succ else float("nan") + bk_mean = total["bk_total_iter"] / bk_succ if bk_succ else float("nan") + print( + f"{'TOTAL':<20s} {total['n']:>3d} {total['bk_better']:>7d} " + f"{total['cart_better']:>9d} {total['only_cart_failed']:>10d} " + f"{total['only_bk_failed']:>8d} {total['both_failed']:>10d} " + f"{cart_mean:>13.1f} {bk_mean:>11.1f}" + ) + + # Drift / iter ratios across cases where both engines succeed. + both_ok = [r for r in rows if r["cart_flag"] == 0 and r["bk_flag"] == 0] + if both_ok: + ratios = [r["bk_drift_AU"] / r["cart_drift_AU"] for r in both_ok if r["cart_drift_AU"] > 1e-9] + iter_ratios = [r["bk_niter"] / max(r["cart_niter"], 1) for r in both_ok] + print() + print(f"Across {len(both_ok)} cases where both engines succeed:") + if ratios: + print( + f" drift ratio (BK / Cart): median={statistics.median(ratios):.3f}, " + f"mean={statistics.mean(ratios):.3f}" + ) + print( + f" iter ratio (BK / Cart): median={statistics.median(iter_ratios):.3f}, " + f"mean={statistics.mean(iter_ratios):.3f}" + ) + + +# ---------------------------------------------------------------------- +# Main +# ---------------------------------------------------------------------- + + +def parse_args(argv: List[str]) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument( + "--scan-dir", + default=DEFAULT_SCAN_DIR, + help=f"Directory of diagnostic-scan .json cases (default: {DEFAULT_SCAN_DIR})", + ) + p.add_argument( + "--cache-dir", + default=DEFAULT_CACHE_DIR, + help=f"layup ephemeris cache directory (default: {DEFAULT_CACHE_DIR})", + ) + p.add_argument( + "--output", + default=DEFAULT_OUTPUT, + help=f"Output CSV path (default: {DEFAULT_OUTPUT})", + ) + return p.parse_args(argv) + + +def main(argv: List[str] | None = None) -> int: + args = parse_args(argv or sys.argv[1:]) + scan_dir = Path(args.scan_dir).expanduser() + cache_dir = Path(args.cache_dir).expanduser() + output = Path(args.output).expanduser() + + if not scan_dir.is_dir(): + print(f"ERROR: diagnostic scan not found at {scan_dir}", file=sys.stderr) + return 1 + if not cache_dir.is_dir(): + print(f"ERROR: ephem cache not found at {cache_dir}", file=sys.stderr) + return 1 + + ephem = get_ephem(str(cache_dir)) + case_paths = sorted(scan_dir.glob("*.json")) + print(f"Running engine sweep on {len(case_paths)} cases from {scan_dir}...") + + rows = [] + for i, path in enumerate(case_paths, start=1): + try: + row = sweep_one(ephem, path) + except Exception as exc: # noqa: BLE001 -- want full coverage + print( + f" [{i}/{len(case_paths)}] {path.stem}: raised " f"{type(exc).__name__}: {exc}", + file=sys.stderr, + ) + continue + rows.append(row) + if i % 10 == 0: + print(f" ...{i}/{len(case_paths)} done") + + write_csv(rows, output) + print_summary(rows) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/bk_iod_sweep.py b/tools/bk_iod_sweep.py new file mode 100644 index 00000000..6b84a761 --- /dev/null +++ b/tools/bk_iod_sweep.py @@ -0,0 +1,415 @@ +"""IOD-sweep harness: compare BK-IOD vs Gauss-IOD on a diagnostic +scan, both as raw IOD output and as a seed for the BK LM fit. + +For each case we record: + * IOD output drift from truth (Gauss picks its first viable root; + BK-IOD's single output) + * Final LM-converged drift from truth and iteration count, starting + from the IOD output (Gauss uses the same "try each root in turn" + fallback as orbitfit.do_fit) + * Success/failure flags + +The point is to see whether BK-IOD is at least as good a seed as +Gauss for the BK LM, and whether one or the other expands the regime +where the full fit converges to truth. + +Usage:: + + python tools/bk_iod_sweep.py --scan-dir --output + +Defaults: + --scan-dir ~/Dropbox/claude_layup/diagnostic/scan/truth + --cache-dir ~/Library/Caches/layup + --output bk_iod_sweep.csv (in the cwd) +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +import sys +from pathlib import Path +from typing import List + +import numpy as np + +from layup.orbitfit import _MU_SUN +from layup.routines import ( + FitResult, + Observation, + gauss, + get_ephem, + run_bk_iod, + run_bk_native_fit, +) + +DEFAULT_SCAN_DIR = "~/Dropbox/claude_layup/diagnostic/scan/truth" +DEFAULT_CACHE_DIR = "~/Library/Caches/layup" +DEFAULT_OUTPUT = "bk_iod_sweep.csv" + +# Constants for the Gauss call -- match orbitfit.do_gauss_iod's call site. +_GMTOTAL = 0.0002963092748799319 +_AU_M = 149597870700 +_SPEED_OF_LIGHT = 2.99792458e8 * 86400.0 / _AU_M + + +# ---------------------------------------------------------------------- +# Case loading and observation construction +# ---------------------------------------------------------------------- + + +def load_case(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +def build_observations(case: dict) -> list: + sigma_arcsec = float(case["sigma_arcsec"]) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) + obs_list = [] + for o in case["observations"]: + pos = list(o["observer_state_AU"]) + vel = [0.0, 0.0, 0.0] + obs = Observation.from_astrometry( + ra=np.deg2rad(o["ra"]), + dec=np.deg2rad(o["dec"]), + epoch=float(o["jd_tdb"]), + observer_position=pos, + observer_velocity=vel, + ) + obs.ra_unc = sigma_rad + obs.dec_unc = sigma_rad + obs_list.append(obs) + return obs_list + + +# ---------------------------------------------------------------------- +# IOD wrappers +# ---------------------------------------------------------------------- + + +def gauss_iod(observations: list) -> list: + """Run Gauss IOD using the first, middle, last observations. + Returns a list of FitResult candidates (Gauss roots); empty if no + viable solutions.""" + if len(observations) < 3: + return [] + idx0 = 0 + idx1 = len(observations) // 2 + idx2 = len(observations) - 1 + return gauss( + _GMTOTAL, + observations[idx0], + observations[idx1], + observations[idx2], + 0.0001, + _SPEED_OF_LIGHT, + ) + + +def bk_iod(observations: list, epoch: float) -> FitResult: + return run_bk_iod(observations, epoch, _MU_SUN) + + +# ---------------------------------------------------------------------- +# Per-case sweep +# ---------------------------------------------------------------------- + + +def _drift_AU(state, truth): + return float(np.linalg.norm(np.asarray(state)[:3] - np.asarray(truth)[:3])) + + +def _try_lm(ephem, seed: FitResult, observations: list, mu: float) -> FitResult: + """Run BK LM from a seed. Returns the FitResult.""" + return run_bk_native_fit(ephem, seed, observations, mu) + + +def sweep_one(ephem, case_path: Path) -> dict: + case = load_case(case_path) + obs = build_observations(case) + truth = np.asarray(case["truth_state_at_epoch"]) + epoch = float(case["epoch_jd_tdb"]) + r_helio = float(np.linalg.norm(truth[:3])) + + # --- Gauss IOD: take its candidate roots --- + gauss_solns = gauss_iod(obs) + # Default values when Gauss returns nothing usable. + g_iod_drift = float("nan") + g_lm_drift = float("nan") + g_lm_niter = -1 + g_lm_flag = -1 + g_n_roots = len(gauss_solns) + + # Use only the FIRST Gauss root. orbitfit.do_fit also starts with + # solns[0]; the fall-back to roots 1, 2 only fires when the first + # root's LM fails. Running LM on all 3 roots for every case made + # the sweep effectively unusable (pathological seeds from spurious + # Gauss roots burn iter_max=100 with no convergence). This way we + # measure the typical do_fit experience. Note: Gauss returns flag + # in an uninitialized state, so we ignore soln.flag and only check + # the LM's flag afterwards. + if gauss_solns: + soln = gauss_solns[0] + g_iod_drift = _drift_AU(soln.state, truth) + lm_res = _try_lm(ephem, soln, obs, _MU_SUN) + if lm_res.flag == 0: + g_lm_drift = _drift_AU(lm_res.state, truth) + g_lm_niter = lm_res.niter + g_lm_flag = lm_res.flag + + # --- BK IOD --- + bk_iod_result = bk_iod(obs, epoch) + b_iod_drift = float("nan") + b_lm_drift = float("nan") + b_lm_niter = -1 + b_lm_flag = -1 + if bk_iod_result.flag == 0: + b_iod_drift = _drift_AU(bk_iod_result.state, truth) + bk_lm_res = _try_lm(ephem, bk_iod_result, obs, _MU_SUN) + b_lm_niter = bk_lm_res.niter + b_lm_flag = bk_lm_res.flag + if bk_lm_res.flag == 0: + b_lm_drift = _drift_AU(bk_lm_res.state, truth) + + return { + "case": case_path.stem, + "population": case["population"], + "arc_days": float(case["arc_length_days"]), + "n_obs": len(obs), + "r_helio_AU": r_helio, + # Gauss + "gauss_n_roots": g_n_roots, + "gauss_iod_drift_AU": g_iod_drift, + "gauss_lm_drift_AU": g_lm_drift, + "gauss_lm_niter": g_lm_niter, + "gauss_lm_flag": g_lm_flag, + # BK + "bk_iod_flag": int(bk_iod_result.flag), + "bk_iod_drift_AU": b_iod_drift, + "bk_lm_drift_AU": b_lm_drift, + "bk_lm_niter": b_lm_niter, + "bk_lm_flag": b_lm_flag, + } + + +# ---------------------------------------------------------------------- +# Reporting +# ---------------------------------------------------------------------- + + +def write_csv(rows: List[dict], path: Path) -> None: + if not rows: + return + fieldnames = list(rows[0].keys()) + with open(path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for row in rows: + writer.writerow(row) + print(f"Wrote {len(rows)} rows to {path}") + + +def _fmt_drift(d): + if d is None or (isinstance(d, float) and (np.isnan(d) or not np.isfinite(d))): + return "--" + if d < 1e-6: + return f"{d:.2e}" + return f"{d:.4f}" + + +def _converged_drift(row, key): + """Returns the LM-converged drift, or NaN if LM didn't converge.""" + flag = row[f"{key}_lm_flag"] + if flag != 0: + return float("nan") + return row[f"{key}_lm_drift_AU"] + + +def print_summary(rows: List[dict]) -> None: + print() + print("=" * 110) + print( + f"{'Case':<40s} {'r/AU':>6s} {'arc':>5s} " + f"{'g_iod':>10s} {'g_lm':>10s} {'g_it':>4s} " + f"{'b_iod':>10s} {'b_lm':>10s} {'b_it':>4s}" + ) + print("=" * 110) + pop_summary: dict = {} + for row in rows: + pop = row["population"] + s = pop_summary.setdefault( + pop, + { + "n": 0, + "gauss_lm_ok": 0, + "bk_lm_ok": 0, + "both_ok": 0, + "bk_better": 0, + "gauss_better": 0, + "bk_only": 0, + "gauss_only": 0, + "neither": 0, + "gauss_total_iter": 0, + "bk_total_iter": 0, + }, + ) + s["n"] += 1 + g_ok = row["gauss_lm_flag"] == 0 + b_ok = row["bk_lm_flag"] == 0 + if g_ok: + s["gauss_lm_ok"] += 1 + s["gauss_total_iter"] += row["gauss_lm_niter"] + if b_ok: + s["bk_lm_ok"] += 1 + s["bk_total_iter"] += row["bk_lm_niter"] + if g_ok and b_ok: + s["both_ok"] += 1 + gd = row["gauss_lm_drift_AU"] + bd = row["bk_lm_drift_AU"] + if bd < gd: + s["bk_better"] += 1 + elif gd < bd: + s["gauss_better"] += 1 + elif b_ok and not g_ok: + s["bk_only"] += 1 + elif g_ok and not b_ok: + s["gauss_only"] += 1 + else: + s["neither"] += 1 + + print( + f"{row['case']:<40s} {row['r_helio_AU']:>6.1f} " + f"{row['arc_days']:>5.2f} " + f"{_fmt_drift(row['gauss_iod_drift_AU']):>10s} " + f"{_fmt_drift(row['gauss_lm_drift_AU']):>10s} " + f"{row['gauss_lm_niter']:>4d} " + f"{_fmt_drift(row['bk_iod_drift_AU']):>10s} " + f"{_fmt_drift(row['bk_lm_drift_AU']):>10s} " + f"{row['bk_lm_niter']:>4d}" + ) + + print() + print("=" * 115) + print(f"Per-population summary ({len(rows)} cases total)") + print("=" * 115) + header = ( + f"{'Population':<20s} {'n':>3s} {'G->LM ok':>8s} {'B->LM ok':>8s} " + f"{'BK win':>7s} {'G win':>5s} {'BK only':>7s} {'G only':>6s} {'neither':>7s} " + f"{'mean g_it':>10s} {'mean b_it':>10s}" + ) + print(header) + print("-" * len(header)) + total = { + k: 0 + for k in ( + "n", + "gauss_lm_ok", + "bk_lm_ok", + "bk_better", + "gauss_better", + "bk_only", + "gauss_only", + "neither", + "gauss_total_iter", + "bk_total_iter", + ) + } + for pop in sorted(pop_summary): + s = pop_summary[pop] + for k in total: + total[k] += s[k] + g_mean_it = s["gauss_total_iter"] / s["gauss_lm_ok"] if s["gauss_lm_ok"] else float("nan") + b_mean_it = s["bk_total_iter"] / s["bk_lm_ok"] if s["bk_lm_ok"] else float("nan") + print( + f"{pop:<20s} {s['n']:>3d} {s['gauss_lm_ok']:>8d} {s['bk_lm_ok']:>8d} " + f"{s['bk_better']:>7d} {s['gauss_better']:>5d} {s['bk_only']:>7d} " + f"{s['gauss_only']:>6d} {s['neither']:>7d} " + f"{g_mean_it:>10.1f} {b_mean_it:>10.1f}" + ) + print("-" * len(header)) + g_mean = total["gauss_total_iter"] / total["gauss_lm_ok"] if total["gauss_lm_ok"] else float("nan") + b_mean = total["bk_total_iter"] / total["bk_lm_ok"] if total["bk_lm_ok"] else float("nan") + print( + f"{'TOTAL':<20s} {total['n']:>3d} {total['gauss_lm_ok']:>8d} {total['bk_lm_ok']:>8d} " + f"{total['bk_better']:>7d} {total['gauss_better']:>5d} {total['bk_only']:>7d} " + f"{total['gauss_only']:>6d} {total['neither']:>7d} " + f"{g_mean:>10.1f} {b_mean:>10.1f}" + ) + + # Drift-ratio + iter-ratio across cases where both LMs succeed. + both_ok = [r for r in rows if r["gauss_lm_flag"] == 0 and r["bk_lm_flag"] == 0] + if both_ok: + ratios = [] + for r in both_ok: + gd = r["gauss_lm_drift_AU"] + bd = r["bk_lm_drift_AU"] + if gd > 1e-9: + ratios.append(bd / gd) + iter_ratios = [r["bk_lm_niter"] / max(r["gauss_lm_niter"], 1) for r in both_ok] + print() + print(f"Across {len(both_ok)} cases where BOTH IOD->LM pipelines converge:") + if ratios: + print( + f" drift ratio (BK_LM / Gauss_LM): " + f"median={statistics.median(ratios):.3f}, mean={statistics.mean(ratios):.3f}" + ) + print( + f" iter ratio (BK_LM / Gauss_LM): " + f"median={statistics.median(iter_ratios):.3f}, mean={statistics.mean(iter_ratios):.3f}" + ) + + +# ---------------------------------------------------------------------- +# Main +# ---------------------------------------------------------------------- + + +def parse_args(argv: List[str]) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + p.add_argument("--scan-dir", default=DEFAULT_SCAN_DIR) + p.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR) + p.add_argument("--output", default=DEFAULT_OUTPUT) + return p.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv or sys.argv[1:]) + scan_dir = Path(args.scan_dir).expanduser() + cache_dir = Path(args.cache_dir).expanduser() + output = Path(args.output).expanduser() + + if not scan_dir.is_dir(): + print(f"ERROR: diagnostic scan not found at {scan_dir}", file=sys.stderr) + return 1 + if not cache_dir.is_dir(): + print(f"ERROR: ephem cache not found at {cache_dir}", file=sys.stderr) + return 1 + + ephem = get_ephem(str(cache_dir)) + case_paths = sorted(scan_dir.glob("*.json")) + print(f"Running IOD sweep on {len(case_paths)} cases from {scan_dir}...") + + rows = [] + for i, path in enumerate(case_paths, start=1): + try: + row = sweep_one(ephem, path) + except Exception as exc: # noqa: BLE001 + print( + f" [{i}/{len(case_paths)}] {path.stem}: raised " f"{type(exc).__name__}: {exc}", + file=sys.stderr, + ) + continue + rows.append(row) + if i % 10 == 0: + print(f" ...{i}/{len(case_paths)} done") + + write_csv(rows, output) + print_summary(rows) + return 0 + + +if __name__ == "__main__": + sys.exit(main())