Skip to content

Add universal Bernstein-Khushalani engine (bk_native) to do_fit#326

Open
matthewholman wants to merge 8 commits into
mainfrom
feat/bk-everywhere
Open

Add universal Bernstein-Khushalani engine (bk_native) to do_fit#326
matthewholman wants to merge 8 commits into
mainfrom
feat/bk-everywhere

Conversation

@matthewholman
Copy link
Copy Markdown
Collaborator

This branch was developed with Claude Code, on top of an initial
design discussion captured in the project notes.

Adds a universal Bernstein-Khushalani (BK) fitter as a new engine
option for do_fit, living entirely inside layup's C++ on top of the
existing variational-integration + residual machinery. No external
dependencies (does not call into liborbfit). The new engine is
always a valid choice -- not distance-dispatched -- because the BK
basis with a fixed bound-orbit energy prior on gdot reduces to
ordinary least-squares for well-arced orbits while smoothly providing
ill-conditioning protection for distant short-arc cases.

What's new

src/lib/orbit_fit/bk_basis.{h,cpp}    BK parameterization primitives (Eigen-only)
src/lib/orbit_fit/bk_fit.cpp          LM driver run_bk_native_fit (reuses
                                      compute_residuals from orbit_fit.cpp)
src/main.cpp                          register bk_basis_bindings + bk_fit_bindings

src/layup/orbitfit.py                 engine='bk_native' dispatch in do_fit
                                      via a new _run_fit(...) helper

tools/bk_engine_sweep.py              full-sweep harness over a diagnostic
                                      dataset; produces a CSV + summary table

tests/layup/test_bk_basis.py          Layer 1: 25 pure-math invariants
tests/layup/test_bk_fit.py            Layer 2: 10 LM-driver smoke + convergence
tests/layup/test_bk_everywhere.py     Layer 3:  6 engine-sweep regressions

How it works (briefly)

BK parameterizes the orbit at epoch by (α, β, γ, α̇, β̇, γ̇) where
γ = 1/r_helio and (α, β) are gnomonic tangent-plane coordinates
of the line-of-sight at a fiducial direction n̂₀. The fit:

  1. Pick n̂₀ from the mean of the observations' ρ̂ vectors.
  2. Convert the Cartesian seed to BK via cartesian_to_bk.
  3. Compute a fixed bound-orbit energy prior on γ̇:
    σ_γ̇² = γ²·(2μγ³ − |α̇·ρ̂_α + β̇·ρ̂_β|²) (this is the exact
    off-fiducial form; the naive γ²(2μγ³ − α̇² − β̇²) is only valid
    at α=β=0 and was caught by an invariant test on first run).
  4. LM loop: at each iteration, convert the current BK state to
    Cartesian, call layup's compute_residuals to get tangent-plane
    residuals and the 6-element Cartesian partials, chain-rule through
    the 6×6 dcart_dbk to get BK partials, assemble
    C = BᵀWB + λI + P_prior and grad = BᵀWr + P_prior·p_bk,
    solve, Marquardt-accept/reject, repeat.
  5. On convergence, propagate the BK covariance to Cartesian via the
    same 6×6 chain rule and return a standard FitResult.

The Marquardt damping λI and the fixed energy prior P_prior
play complementary roles: λI provides LM's adaptive step control
(large early, decays to zero at convergence), P_prior provides
the persistent physical regularization that survives at convergence
in the line-of-sight direction whenever the data don't constrain it.

Validation on the diagnostic scan

tools/bk_engine_sweep.py runs both engines on every case in a
diagnostic scan (default: 98 cases across 7 populations × 14 arc
lengths, σ_obs = 0.1″, truth state as LM seed). Result:

Population n BK win Cart win Cart fail BK fail Both fail
centaur_15AU 14 13 0 0 0 1
centaur_25AU 14 9 2 2 0 1
classical_42AU 14 10 3 0 0 1
mainbelt_2.5AU 14 10 3 0 0 1
mainbelt_3.5AU 14 12 1 0 0 1
scattered_70AU 14 7 2 4 0 1
sednoid_80AU 14 6 1 6 0 1
TOTAL 98 67 12 12 0 7

Across the 79 cases where both engines succeed:

  • Drift ratio (BK / Cart): median 0.56, mean 187
  • Iteration ratio (BK / Cart): median 0.39, mean 0.52

Headline observations:

  • BK never fails when Cartesian succeeds (0/98). In 12 cases
    Cartesian's flag=2 (chi²/dof out of range) while BK converges.
  • On the typical case, BK is ~2× closer to truth in ~40% the
    iterations. Mean drift ratio of 187 reflects the distant short-arc
    cases where Cartesian wanders 5-13 AU while BK stays within 0.1 AU.
  • The 7 both-failed cases are all 0.04-day (1-hour) arcs with 3
    observations -- a data limitation (0 effective d.o.f.), not an
    algorithm issue.

Standout extreme cases:

  • scattered_70AU_arc_014.00d: Cartesian 4.52 AU / 58 iter,
    BK 0.017 AU / 6 iter (≈265× tighter, ≈10× fewer iter).
  • sednoid_80AU_arc_007.00d: Cartesian 8.65 AU / 100 iter,
    BK 0.24 AU / 10 iter.

Test coverage (41 tests total, all passing)

  • Layer 1 (test_bk_basis.py, 25): Cartesian↔BK round-trip,
    finite-difference vs analytic Jacobian, mixed-partial symmetry,
    fiducial-direction gauge invariance, special-case forms at the
    fiducial, σ_γ̇² agreement with the Cartesian-side parabolic
    boundary. Pure math, no ASSIST needed.
  • Layer 2 (test_bk_fit.py, 10): empty-obs and few-obs guards,
    recovery of synthetic mainbelt and TNO orbits from perfect and
    0.1%-perturbed seeds, BK↔Cartesian agreement, engine dispatch via
    _run_fit. Needs the ASSIST ephemeris; skips cleanly if absent.
  • Layer 3 (test_bk_everywhere.py, 6): regression tests against
    the diagnostic-scan dataset for well-arced (both engines converge,
    agree, drift < 1% r_helio) and distant short-arc (BK drifts no
    more than Cartesian, uses fewer LM iterations).

Dependencies

Stacks on feat/pybind11-eigen-includes (PR #322 — adds the
pybind11/eigen.h headers required to expose BKState, BKFiducial,
dcart_dbk, etc. to Python). Once #322 lands, this PR collapses to
its 7 new commits.

This branch is intentionally independent of feat/bk-engine-dispatch
(PR #323), which adds engine='auto' for distance-based dispatch
between Cartesian and liborbfit's BK fit. The two PRs can land in
either order; once both are in, engine={'cartesian', 'bk', 'auto', 'bk_native'} are all available and the diagnostic harness can sweep
all four for comparison.

Review Checklist for Source Code Changes

  • Does pip install still work?
  • Have you written a unit test for any new functions? — 41 tests across test_bk_basis.py, test_bk_fit.py, test_bk_everywhere.py
  • Do all the units tests run successfully? — full suite green
  • Does Layup run successfully on a test set of input files/databases? — yes, plus full diagnostic-scan sweep (98 cases)
  • Have you used black on the files you have updated?

matthewholman and others added 8 commits May 15, 2026 18:03
… Python

Without these, accessing rho_hat / a_vec / d_vec from Python raises
"Unable to convert function return value to a Python type" because the
binding can't see the Eigen and std::array conversions.

This is a minimal enabler (no behavior change); the A/D-vector
correctness fix on fix/tangent-basis-vectors is independent and lives
on its own branch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pure C++/Eigen math layer for the Bernstein-Khushalani parameterization,
with the barycenter as the BK coordinate origin and gnomonic projection
defining the tangent plane at a fiducial direction n0.  No ASSIST,
REBOUND, or pybind11 dependencies, so this translation unit can be
tested in isolation.

New files in src/lib/orbit_fit/:
  bk_basis.h    -- types (BKState, BKFiducial) and function declarations
  bk_basis.cpp  -- implementations: choose_fiducial, bk_to_cartesian,
                   cartesian_to_bk, dcart_dbk (full 6x6 including the
                   bottom-left cross-term block), sigma_gdot_sq

The 6x6 Jacobian dcart_dbk has the block structure expected from the
math:
  [ d r / d (alpha,beta,gamma)     0                          ]
  [ d v / d (alpha,beta,gamma)     d v / d (adot,bdot,gdot)   ]
with the top-left and bottom-right 3x3 blocks identical (both built
from the (1/gamma)-scaled tangent vectors), and the bottom-left block
holding the cross-term contributions through the second derivatives
d^2 rho_hat / d (alpha, beta)^2.

sigma_gdot_sq returns the bound-orbit energy-prior variance,
gamma^2 (2 mu gamma^3 - adot^2 - bdot^2), or +infinity when the
right-hand side is non-positive (tangential rates already exceed
escape).  Returning +infinity yields zero precision in the prior
matrix used by the LM step, which is the desired "no prior" behavior.

orbit_fit.cpp adds a single line to its unity-build chain:
  #include "bk_basis.cpp"

so the new translation unit compiles into the existing _core module.
No pybind11 binding yet -- that comes in a follow-up commit alongside
Python-side unit tests for the math primitives.

The math derivation, design decisions (barycenter origin, fixed
gdot prior, eigendecomp+energy-prior solver, file layout, layered
test plan) live in the project memory file bk_everywhere_design.md.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds a bk_basis_bindings(py::module&) entry point that exposes BKState,
BKFiducial, bk_choose_fiducial, bk_to_cartesian, cartesian_to_bk,
dcart_dbk, and sigma_gdot_sq to Python via pybind11 / pybind11/eigen.h,
and wires the binding into main.cpp's _core module alongside the
existing detection_bindings etc.

tests/layup/test_bk_basis.py covers the pure-math invariants
(25 tests, all passing):

  * Round-trip Cartesian <-> BK across mainbelt / NEO / TNO regimes
    (rtol 1e-12).
  * Analytic dcart_dbk vs central-difference, per-element relative
    error < 1e-5 with parameter-scaled epsilon.
  * Mixed-partial symmetry of the second-derivative cross-terms
    appearing in the bottom-left block of dcart_dbk
    (d^2 r / d alpha d beta == d^2 r / d beta d alpha to FD tolerance).
  * Fiducial-direction gauge invariance: two valid n0 choices recover
    the same Cartesian orbit through round-trip.
  * Special-case forms at the fiducial direction alpha = beta = 0:
    position is (1/gamma) n0, top-left and bottom-right Jacobian
    blocks are [(1/gamma) a, (1/gamma) b, -(1/gamma^2) n0] as columns,
    bottom-left block vanishes when rates are zero.
  * sigma_gdot_sq agreement with the Cartesian energy bound at the
    parabolic boundary, and +inf return when tangential rates already
    exceed escape.

The energy-bound test caught a real bug in the first cut of
sigma_gdot_sq: the formula gamma^2 (2 mu gamma^3 - adot^2 - bdot^2)
is only exact at the fiducial direction.  Off-fiducial, the gnomonic
tangent vectors rho_hat_alpha, rho_hat_beta have magnitudes
sqrt((1+beta^2))/s^2 and sqrt((1+alpha^2))/s^2 respectively, and an
inner product -alpha*beta/s^4, so the true tangential-velocity term is

  |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

which reduces to adot^2 + bdot^2 only at alpha = beta = 0.  Fixed in
sigma_gdot_sq (and bk_basis.h documentation) to use the exact form,
which reproduces the parabolic-boundary condition |v|^2 = 2 mu / |r|
to machine precision in the test.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
bk_fit.cpp contains the LM driver that performs an orbit fit in the
universal Bernstein-Khushalani parameterization on top of layup's
existing Cartesian variational machinery.  Included from orbit_fit.cpp
inside the `namespace orbit_fit` block, after the Cartesian helpers
(compute_residuals, create_sequences, get_weight_matrix, converged)
and the Observation/FitResult types are in scope, so no forward
declarations are needed.

The driver structure mirrors run_from_vector_with_initial_guess but
operates in BK basis throughout:

  1. Pick a fiducial direction from the observations' rho_hat vectors
     (mean direction, Gram-Schmidt for the orthonormal a, b).
  2. Convert the Cartesian seed to BK via cartesian_to_bk.
  3. Compute a fixed bound-orbit energy prior precision on gdot from
     the BK seed (1 / sigma_gdot_sq), zero precision otherwise.
  4. LM loop:
       - convert current BK state to Cartesian -> reb_particle
       - call compute_residuals to get tangent-plane residuals and
         Cartesian 6-element partials per observation
       - chain-rule: B_bk = B_cart * dcart_dbk(current BK, fiducial)
       - assemble C = B_bk^T W B_bk + lambda I + P_prior,
                 grad = B_bk^T W r + P_prior * p_bk
       - solve, Marquardt rho-ratio accept/reject, update BK state,
         check convergence (using the existing `converged` predicate).
  5. On convergence:
       - cov_bk = (B^T W B + P_prior)^-1 (Hessian without lambda)
       - cov_cart = J cov_bk J^T (J = dcart_dbk at converged BK state)
       - return FitResult with state = bk_to_cartesian(BK_final) and
         cov flattened from cov_cart.  method = "bk_native".

Initial lambda and Marquardt accept threshold match the Cartesian
fit at orbit_fit.cpp:553.  Early-exit guard: returns a non-success
FitResult (flag = 1) without crashing when detections.size() < 3.

main.cpp gains an orbit_fit::bk_fit_bindings(m) call alongside the
existing orbit_fit bindings, exposing run_bk_native_fit to Python.

tests/layup/test_bk_fit.py covers the Layer 2 smoke tests:
  * binding loads and run_bk_native_fit returns a FitResult
  * empty-obs path returns flag != 0 without crashing
  * <3 obs path triggers the early-exit guard

Layer 2 convergence tests against synthetic observations from a
known orbit (and the Cartesian/BK agreement test on well-arced
mainbelt) are next steps -- they need either the predict-path
output piped back in or the diagnostic/scan dataset wired up to
this branch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Generate synthetic observations from a known Cartesian state via
layup's predict_sequence path (a fixed barycenter observer, so the
only dynamical content is the orbit itself), then feed those
observations back into both run_bk_native_fit and the existing
Cartesian fit.

Three new test categories on top of the smoke tests:

  * test_bk_native_fit_recovers_known_state: with the truth state as
    the seed, BK converges in essentially one iteration to a fit
    state matching the truth to rtol=1e-6 and chi2 < 1e-12.
    Parameterized over a 3 AU mainbelt 60-day arc and a 40 AU TNO
    300-day arc.

  * test_bk_native_fit_recovers_from_perturbed_seed: with the seed
    perturbed by 0.1% in each component, the LM loop still converges
    to the truth state (rtol=1e-6) -- exercising the chain-rule
    Jacobian + Marquardt damping on a non-trivial number of
    iterations.  Same two orbital regimes.

  * test_bk_and_cartesian_fits_agree: for the well-constrained
    mainbelt case, run_bk_native_fit and run_from_vector_with_initial_guess
    converge to states that agree at rtol=1e-6.  Establishes the
    "no regression on the easy case" baseline.

All seven tests in tests/layup/test_bk_fit.py pass with ASSIST
ephemeris available; skip cleanly without it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wires the universal BK fitter (run_bk_native_fit) into layup's Python
do_fit pipeline alongside the existing Cartesian fit, so callers can
choose the engine per call rather than via the C++ entry point.

Changes:

  - src/layup/orbitfit.py
      * Import run_bk_native_fit from layup.routines.
      * Add a module-level _MU_SUN constant (heliocentric GM in
        AU^3 / day^2) used to construct the BK fixed energy prior.
      * Add _run_fit(assist_ephem, initial_guess, obs, engine) helper
        that dispatches to:
          - run_from_vector_with_initial_guess  for engine='cartesian'
          - run_bk_native_fit (with _MU_SUN)    for engine='bk_native'
          - ValueError otherwise
      * do_fit gains an `engine='cartesian'` parameter (default
        preserves the existing behavior).  All five
        run_from_vector_with_initial_guess call sites inside do_fit
        are now routed through _run_fit so the engine choice
        propagates uniformly.

  - tests/layup/test_bk_fit.py
      * test_run_fit_dispatch_cartesian: _run_fit(..., 'cartesian')
        matches direct run_from_vector_with_initial_guess on
        synthetic mainbelt observations.
      * test_run_fit_dispatch_bk_native: _run_fit(..., 'bk_native')
        matches direct run_bk_native_fit(ephem, ig, obs, MU_SUN).
      * test_run_fit_dispatch_unknown_engine_raises: ValueError on
        an unknown engine name.

The 'auto' (distance-dispatched) engine from PR 323 is intentionally
not wired up here; when 323 lands first this branch rebases and
gains both options.  Likewise the 'bk' (liborbfit-backed) engine
from PR 323 is independent of this work.

All 10 tests in test_bk_fit.py and 25 tests in test_bk_basis.py
continue to pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
tests/layup/test_bk_everywhere.py drives both engine='cartesian' and
engine='bk_native' against the diagnostic/scan dataset at
~/Dropbox/claude_layup/diagnostic/scan/truth/ -- the same 7-population,
14-arc-length scan that PR 323's auto-dispatch was validated against.
Skips when either the ASSIST ephemeris or the diagnostic scan is
unavailable, so CI is unaffected.

Two test groups:

  * test_engine_sweep_well_arced_cases: on long-arc cases (30-60d
    mainbelt + 60d classical TNO), both engines converge near truth
    (drift < 1% of heliocentric distance) and agree with each other.

  * test_bk_beats_cartesian_on_short_arc_distant: on distant short-arc
    cases (70 AU scattered / 42 AU classical at 10-14 day arcs), BK
    drifts no more than Cartesian from truth AND uses fewer LM
    iterations.  This is the regime BK was designed for, and the
    diagnostic data shows it strongly: on scattered_70AU_arc_014.00d
    BK stays 0.02 AU from truth in 6 iterations while Cartesian
    wanders 4.5 AU over 58 iterations.

The module also exposes a sweep_cases_from_diagnostic() helper for
ad-hoc engine-sweep harness scripts.

All 6 Layer 3 tests pass (in addition to the 25 Layer 1 + 10 Layer 2
tests, for 41 total BK tests on this branch).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A runnable CLI script that drives both engine='cartesian' and
engine='bk_native' across an entire diagnostic-scan directory, writes
per-case metrics to CSV, and prints a population-level summary
(BK wins / Cartesian wins / per-engine failures / mean iteration
counts, plus median+mean drift and iteration ratios).

Usage:
    python tools/bk_engine_sweep.py --scan-dir <dir> --output <csv>

Defaults discover the project's diagnostic scan at
~/Dropbox/claude_layup/diagnostic/scan/truth and the layup ephemeris
cache at ~/Library/Caches/layup.  Both are overrideable via flags,
so anyone with a compatible truth dataset can reproduce.

Running on the 98-case scan (truth state as LM seed, sigma_arcsec=0.1):

  Population              n   BK win   Cart win   cart fail   bk fail   both fail
  -------------------------------------------------------------------------------
  centaur_15AU           14       13          0           0         0           1
  centaur_25AU           14        9          2           2         0           1
  classical_42AU         14       10          3           0         0           1
  mainbelt_2.5AU         14       10          3           0         0           1
  mainbelt_3.5AU         14       12          1           0         0           1
  scattered_70AU         14        7          2           4         0           1
  sednoid_80AU           14        6          1           6         0           1
  -------------------------------------------------------------------------------
  TOTAL                  98       67         12          12         0           7

  Across 79 cases where both engines succeed:
    drift ratio  (BK / Cart):  median=0.560, mean=187.199
    iter ratio   (BK / Cart):  median=0.386, mean=0.524

Headline: BK never fails when Cartesian succeeds (0 / 98), succeeds in
12 cases where Cartesian flag=2's out, and on the typical case is ~2x
closer to truth in ~40% the iterations.  The mean drift ratio of 187 is
inflated by the 70-80 AU short-arc cases where Cartesian wanders 5-13 AU
while BK stays put.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@matthewholman
Copy link
Copy Markdown
Collaborator Author

Heads up: CI is failing here at the "Install dependencies" step, but the failure is an upstream dependency issue, not anything in this PR.

Root cause: ABI mismatch between the latest assist and rebound wheels on PyPI. The build-isolated environment that pip creates for layup pulls the latest of each, and libassist now references a rebound symbol that the latest librebound no longer exports:

OSError: dlopen(libassist.cpython-312-darwin.so):
  Symbol not found: _reb_simulation_create_from_simulationarchive_with_messages
  Referenced from: libassist.cpython-312-darwin.so
  Expected in:     librebound.cpython-312-darwin.so
CMake Error at CMakeLists.txt:34:
  Could not find assist package. Is assist installed in the build environment?

Reproduces locally on a clean clone today (2026-05-15). Affects every branch, not just this one — main's last successful build was 2026-05-11, and the matrix runs since then have stayed queued or quietly failed. PR #325 passes only because its successful CI ran before the upstream break.

The 41 BK tests in this PR all pass locally against the previously-installed assist 1.2.0 / rebound 4.6.0 combo (i.e., this PR doesn't cause any regression independent of the upstream issue).

Suggested fix (separate from this PR, since it's a repo-wide problem):

[build-system]
requires = [
    "scikit-build-core>=0.10",
    "pybind11",
    "assist>=1.2,<1.3",   # or whichever range is known-compatible
    "rebound>=4.6,<4.7",
]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant