From 7407db314f2208e112079ccd8d278a4f3a298677 Mon Sep 17 00:00:00 2001 From: awilson110 Date: Fri, 28 Aug 2026 10:37:33 +0100 Subject: [PATCH 1/4] new branch for herget method to avoid conflicts --- src/layup/orbitfit.py | 1534 +++-------------------- src/layup/utilities/herget_iod.py | 306 +++++ src/layup/utilities/universal_kepler.py | 405 ++++++ tests/data/2000OK67_ephem.csv | 62 + tests/layup/test_herget_iod.py | 256 ++++ tests/layup/test_universal_kepler.py | 552 ++++++++ 6 files changed, 1753 insertions(+), 1362 deletions(-) create mode 100644 src/layup/utilities/herget_iod.py create mode 100644 src/layup/utilities/universal_kepler.py create mode 100644 tests/data/2000OK67_ephem.csv create mode 100644 tests/layup/test_herget_iod.py create mode 100644 tests/layup/test_universal_kepler.py diff --git a/src/layup/orbitfit.py b/src/layup/orbitfit.py index c35dbb5f..d482d6db 100644 --- a/src/layup/orbitfit.py +++ b/src/layup/orbitfit.py @@ -1,14 +1,13 @@ -import hashlib import logging import os -import re from argparse import Namespace -from dataclasses import dataclass from pathlib import Path from typing import Literal, Optional import numpy as np +import pooch import spiceypy as spice +from time import sleep from numpy.lib import recfunctions as rfn @@ -17,27 +16,11 @@ Observation, gauss, get_ephem, - run_bk_iod, - run_bk_native_fit, run_from_vector_with_initial_guess, - run_sequential_update, ) - -try: - from layup.routines import ( - get_ias15_adaptive_mode, - set_ias15_adaptive_mode, - ) -except ImportError: # extension not rebuilt yet - get_ias15_adaptive_mode = lambda: -1 - set_ias15_adaptive_mode = lambda m: None -# _MU_SUN (= heliocentric GM = k^2) is used by the BK-native fit for the -# bound-orbit energy prior on gdot; SPEED_OF_LIGHT (au/day) by the radar ingest. -from layup.constants import MU_SUN as _MU_SUN, SPEED_OF_LIGHT from layup.convert import convert -from layup.iod import filter_candidates_by_residual, get_iod, iod_methods -from layup.utilities.astrometric_uncertainty import astrometric_uncertainty_Veres2017 +from layup.utilities.astrometric_uncertainty import data_weight_Veres2017 from layup.utilities.data_processing_utilities import ( LayupObservatory, create_chunks, @@ -45,54 +28,21 @@ get_format, parse_fit_result, process_data_by_id, - resolve_num_workers, ) from layup.utilities.datetime_conversions import convert_tdb_date_to_julian_date from layup.utilities.debiasing import debias, generate_bias_dict -from layup.utilities.file_io import ( - ADESXMLDataReader, - CSVDataReader, - HDF5DataReader, - Obs80DataReader, -) -from layup.utilities.file_io.file_output import append_hdf5, write_csv, write_hdf5 -from layup.utilities.cache_location import default_cache_dir +from layup.utilities.file_io import CSVDataReader, HDF5DataReader, Obs80DataReader +from layup.utilities.file_io.file_output import write_csv, write_hdf5 +from layup.utilities.herget_iod import herget_with_assist logger = logging.getLogger(__name__) -# Observed sky-motion rates (ADES `raRate`/`decRate`) arrive in arcsec/hour and -# follow the great-circle convention: `raRate` is cos(Dec)*dRA/dt, NOT the bare -# coordinate rate dRA/dt. This matches Sorcha's `RARateCosDec` output (which -# projects drho_hat/dt onto A = (-sinRA, cosRA, 0)) and layup's own residual, -# which projects onto the same tangent vector `a_vec` -- so an observed rate is -# compared directly to omega.a_vec with no extra cos(Dec) factor. Internally the -# fitter works in radians and AU/day, so omega = d(rho_hat)/dt is in rad/day; -# convert the observed rates from arcsec/hour to rad/day at ingest. -ARCSEC_PER_HOUR_TO_RAD_PER_DAY = (np.pi / 180.0 / 3600.0) * 24.0 - -# Radar (delay/Doppler) observables arrive in JPL units: round-trip delay in -# microseconds and Doppler shift in Hz at a per-observation transmit frequency -# `freqTx` (Hz). The fitter models round-trip delay in days and round-trip -# range-rate in au/day (see RadarObservation in detection.cpp), so convert at -# ingest, mirroring the streak rate convention above: -# delay[days] = delay[us] * 1e-6 / 86400 -# doppler[au/day] = -c * doppler[Hz] / freqTx[Hz] -# The Doppler sign follows F = -(f_tx/c) * d(round-trip range)/dt, so a positive -# (receding) range-rate produces a negative frequency shift. -US_TO_DAYS = 1.0e-6 / 86400.0 -# Fallback 1-sigma weights when the JPL uncertainty columns are absent: ~1 us of -# round-trip delay and ~1 Hz of Doppler (converted per observation via freqTx). -_DEFAULT_DELAY_UNC_DAYS = US_TO_DAYS -_DEFAULT_DOPPLER_UNC_HZ = 1.0 - # The list of required input column names for the provided observations to be fit. # Note: This should not include the primary id column name. REQUIRED_INPUT_OBSERVATIONS_COLUMN_NAMES = [ ( set(["ra", "dec"]), # Either `ra` and `dec` must be in the file set(["raRate", "decRate"]), # Or `raRate` and `decRate` must be in the file - set(["delay"]), # Or a radar round-trip delay (us) - set(["doppler"]), # Or a radar Doppler shift (Hz; needs `freqTx`) ), "obsTime", "stn", @@ -109,168 +59,17 @@ "MPC80col": (Obs80DataReader, None), "ADES_csv": (CSVDataReader, "csv"), "ADES_psv": (CSVDataReader, "psv"), - "ADES_xml": (ADESXMLDataReader, None), + "ADES_xml": (None, None), "ADES_hdf5": (HDF5DataReader, None), } +GMtotal = 0.0002963092748799319 +AU_M = 149597870700 +SPEED_OF_LIGHT = 2.99792458e8 * 86400.0 / AU_M -def _run_fit(assist_ephem, initial_guess, observations, engine, iter_max=100): - """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. - - `iter_max` is the LM iteration budget used by the multi-root picker's - two-tier (cheap-screen then full) passes. The Cartesian engine honors - it; the BK-native engine uses its own internal cap (it takes `mu` for - the bound-orbit energy prior rather than an iteration budget), so - `iter_max` is ignored on that path. - """ - if engine == "cartesian": - return run_from_vector_with_initial_guess(assist_ephem, initial_guess, observations, iter_max) - 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'.") - - -# Non-gravitational Marsden parameters and their FitResult/C++ bitmask bits. -_NONGRAV_BITS = {"A1": 1, "A2": 2, "A3": 4} - - -def _parse_nongrav(fit_nongrav): - """Normalize the ``fit_nongrav`` argument to ``(mask, names)``. - - Accepts ``False``/``None`` (no non-grav fit), ``True`` (== ``["A2"]``, the - common asteroid Yarkovsky case), a string naming params (e.g. ``"A2"``, - ``"A1A2A3"``, ``"A1,A3"``), or an iterable of names. Returns the C++ bitmask - (bits 1/2/4 for A1/A2/A3) and the selected names ordered A1, A2, A3. - """ - if not fit_nongrav: - return 0, [] - if fit_nongrav is True: - sel = {"A2"} - elif isinstance(fit_nongrav, str): - sel = set(re.findall(r"A[123]", fit_nongrav.upper())) - else: - sel = {str(s).upper() for s in fit_nongrav} - names = [n for n in ("A1", "A2", "A3") if n in sel] - if not names: - raise ValueError(f"fit_nongrav={fit_nongrav!r}: expected some of 'A1', 'A2', 'A3'.") - return sum(_NONGRAV_BITS[n] for n in names), names - - -# fit_nongrav="auto" model ladder (issue #357): non-grav models are tried in -# increasing complexity, and the first that converges, is well-conditioned -# (flag 0), and is statistically warranted (see NongravAutoThresholds) is adopted; -# otherwise the gravity-only fit is kept. -_AUTO_NONGRAV_LADDER = (("A2",), ("A1", "A2"), ("A1", "A2", "A3")) - - -@dataclass(frozen=True) -class NongravAutoThresholds: - """Decision thresholds for adaptive non-grav selection (``fit_nongrav="auto"``, - issue #357). Pass a customized instance to ``orbitfit`` to tune how readily a - non-gravitational model is adopted; the defaults reproduce the standard - "introduce a non-grav only when gravity is unacceptable and the parameter is - statistically warranted" behavior. - - Parameters - ---------- - accept_reduced_chi2 : float - A gravity-only fit whose reduced chi-square is at or below this is kept - as-is; non-grav models are tried only above it. Default 1.5. - delta_chi2_per_param : float - Minimum chi-square drop required per added non-grav parameter to adopt a - model (9.0 ~ 3-sigma). Default 9.0. - nsigma : float - Each added non-grav parameter must exceed this many times its 1-sigma - uncertainty to be adopted. Default 3.0. - """ - - accept_reduced_chi2: float = 1.5 - delta_chi2_per_param: float = 9.0 - nsigma: float = 3.0 - - -_AUTO_DEFAULT_THRESHOLDS = NongravAutoThresholds() - -def _gravity_fit_acceptable(csq, ndof, thresholds=_AUTO_DEFAULT_THRESHOLDS): - """Whether a gravity-only fit is good enough that no non-grav is warranted. - - True when the reduced chi-square is at or below the acceptance threshold (or - there are no degrees of freedom to judge it by). - """ - return ndof <= 0 or csq / ndof <= thresholds.accept_reduced_chi2 - - -def _nongrav_warranted(csq_gravity, res_ng, names, thresholds=_AUTO_DEFAULT_THRESHOLDS): - """Whether adopting the (converged, well-conditioned) non-grav fit ``res_ng`` - for parameters ``names`` over the gravity-only fit is statistically warranted: - a significant chi-square drop AND every added parameter individually significant. - """ - if (csq_gravity - res_ng.csq) <= thresholds.delta_chi2_per_param * len(names): - return False # chi-square improvement not significant for the added parameter(s) - return all( - abs(getattr(res_ng, n.lower())) > thresholds.nsigma * getattr(res_ng, n.lower() + "_unc") - for n in names - ) - - -def _gofr_arg(nongrav_gr): - """Normalize a non-grav g(r) argument to the ``[alpha, nm, nn, nk, r0]`` list the - C++ fit expects (empty -> the default inverse-square law).""" - if nongrav_gr is None: - return [] - gr = list(nongrav_gr) - if len(gr) != 5: - raise ValueError("nongrav_gr must be [alpha, nm, nn, nk, r0] (5 values) or None") - return [float(v) for v in gr] - - -def _select_nongrav_auto( - assist_ephem, res_grav, observations, thresholds=_AUTO_DEFAULT_THRESHOLDS, gofr=None -): - """Adaptive non-grav selection for ``fit_nongrav="auto"`` (issue #357). - - ``res_grav`` is the converged gravity-only fit. Returns the most parsimonious - acceptable model: the gravity-only result unless a non-grav model is both - well-conditioned (flag 0) and statistically warranted per ``thresholds``. - ``gofr`` selects the g(r) sublimation law (see ``orbitfit``'s ``nongrav_gr``). - """ - if _gravity_fit_acceptable(res_grav.csq, res_grav.ndof, thresholds): - return res_grav # gravity-only fit is acceptable; no non-gravs needed - gr = _gofr_arg(gofr) - for names in _AUTO_NONGRAV_LADDER: - mask = sum(_NONGRAV_BITS[n] for n in names) - res_ng = run_from_vector_with_initial_guess( - assist_ephem, res_grav, observations, nongrav_mask=mask, gofr=gr - ) - if res_ng.flag == 0 and _nongrav_warranted(res_grav.csq, res_ng, names, thresholds): - return res_ng # parsimonious, well-determined non-grav model - return res_grav # no non-grav model is warranted - - -def _get_result_dtypes(primary_id_column_name: str, nongrav_names=(), per_arc=False): - """Helper function to create the result dtype with the correct primary ID column name. - - For each fitted non-gravitational parameter in ``nongrav_names`` (a subset of - ``A1``/``A2``/``A3``), two columns are appended -- e.g. ``a2`` and ``a2_unc`` - (the value and its 1-sigma uncertainty, au/day^2). With no non-grav params the - default 6-parameter output schema is unchanged. - - When ``per_arc`` is set (the two-apparition comet-linkage fit), the non-grav - columns above hold the *earlier* arc's amplitudes and a second block - ``a2_arc2``/``a2_arc2_unc`` is appended for the *later* arc. Both are opt-in, - so the ordinary fit schema is unaffected. - """ - per_arc_cols = [] - if per_arc: - per_arc_cols = [ - (col, "f8") for n in nongrav_names for col in (n.lower() + "_arc2", n.lower() + "_arc2_unc") - ] +def _get_result_dtypes(primary_id_column_name: str): + """Helper function to create the result dtype with the correct primary ID column name.""" # Define a structured dtype to match the OrbfitResult fields return np.dtype( [ @@ -290,166 +89,6 @@ def _get_result_dtypes(primary_id_column_name: str, nongrav_names=(), per_arc=Fa ("FORMAT", "O"), # Orbit format ] + [(col_name, "f8") for col_name in get_cov_columns()] # Flat covariance matrix (36 elements) - + [ # non-grav params (issue #351), value + 1-sigma per fitted param - (col, "f8") for n in nongrav_names for col in (n.lower(), n.lower() + "_unc") - ] - + per_arc_cols # later-arc non-grav amplitudes (comet linkage), when per_arc - # Observation provenance / incremental fingerprint (issue #419): a - # deterministic hash of the observation set this orbit was fit from, plus - # the number of fittable observations. Kept last so the positional output - # tuples above the conditional non-grav columns are unaffected. Lets a - # steady-state pipeline skip re-fitting objects whose observations are - # unchanged since the prior catalog (see ``_obs_fingerprint``). - + [("obs_hash", "O"), ("nobs_fit", "i4")] - ) - - -# Observation columns that determine the fit and therefore the fingerprint. Any -# change to these (new obs, corrected astrometry, changed uncertainties) yields a -# different hash and forces a re-fit; changes to purely cosmetic columns (e.g. a -# magnitude) do not. Only the columns actually present in the data are hashed. -_FINGERPRINT_COLUMNS = ( - "obsTime", # epoch - "ra", - "dec", # optical astrometry (raw, pre-debias) - "stn", # observatory code - "raStar", - "decStar", # occultation astrometry - "rmsRA", - "rmsDec", # reported astrometric uncertainties - "astCat", # star catalog (feeds debiasing + weighting) - "raRate", - "decRate", - "rmsRArate", - "rmsDecrate", # streak rates - "delay", - "doppler", - "freqTx", # radar observables -) - - -def _fmt_fingerprint_value(v): - """Canonical, round-trip-stable string for one observation field value. - - Uses Python's shortest round-trip ``repr`` for floats (deterministic across - runs and platforms) and decodes bytes so a numpy ``S``/``O`` string column - hashes identically however it was loaded. - """ - if isinstance(v, bytes): - return v.decode("utf-8", "replace") - if isinstance(v, (np.floating, float)): - return repr(float(v)) - return str(v) - - -def _obs_fingerprint(data, column_names): - """Deterministic fingerprint of an object's fittable observation set. - - Returns ``(nobs, hash16)`` where ``nobs`` is the row count and ``hash16`` is a - 16-hex-character SHA-1 digest over the fit-relevant columns (``_FINGERPRINT_COLUMNS``). - The per-row strings are sorted before hashing, so the fingerprint is - order-independent: the same physical observations produce the same hash - regardless of row order (an LM fit over a set is itself order-independent). - Computed on the raw observations before any in-place debiasing so it reflects - what was reported, not the current bias model. - """ - cols = [c for c in _FINGERPRINT_COLUMNS if c in column_names] - rows = ["\x1f".join(_fmt_fingerprint_value(d[c]) for c in cols) for d in data] - rows.sort() - payload = "\x1e".join(rows).encode("utf-8") - return len(rows), hashlib.sha1(payload).hexdigest()[:16] - - -def _is_radar(d, column_names): - """True if a row carries a populated radar observable (delay or Doppler). - - Radar rows have no ra/dec; they are dispatched to ``Observation.from_radar`` - rather than the astrometry/streak factories. - """ - has_delay = "delay" in column_names and not np.isnan(d["delay"]) - has_doppler = "doppler" in column_names and not np.isnan(d["doppler"]) - return has_delay or has_doppler - - -def _radar_observation(objID, d, epoch_jd, column_names): - """Build a radar ``Observation`` from a row, converting JPL units to the - fitter's internal units. - - delay (us, round-trip) -> days; Doppler (Hz) -> round-trip range-rate - (au/day) via the per-observation transmit frequency ``freqTx``. The - barycentric observer position/velocity columns (x,y,z,vx,vy,vz) must already - be present (added by ``orbitfit``); the observer acceleration columns - (ax,ay,az) drive the two-leg light-time model and default to zero when - absent. 1-sigma uncertainties come from ``rmsDelay``/``rmsDoppler`` when - present, else the module defaults. - """ - has_delay = "delay" in column_names and not np.isnan(d["delay"]) - has_doppler = "doppler" in column_names and not np.isnan(d["doppler"]) - - f_tx = d["freqTx"] if "freqTx" in column_names else np.nan - if has_doppler and (np.isnan(f_tx) or f_tx == 0.0): - raise ValueError(f"Radar Doppler observation for {objID} requires a nonzero 'freqTx' (Hz).") - - delay_days = (d["delay"] * US_TO_DAYS) if has_delay else 0.0 - doppler_audy = (-SPEED_OF_LIGHT * d["doppler"] / f_tx) if has_doppler else 0.0 - - if "rmsDelay" in column_names and not np.isnan(d["rmsDelay"]): - delay_unc = abs(d["rmsDelay"]) * US_TO_DAYS - else: - delay_unc = _DEFAULT_DELAY_UNC_DAYS - - if has_doppler: - rms_hz = ( - d["rmsDoppler"] - if ("rmsDoppler" in column_names and not np.isnan(d["rmsDoppler"])) - else _DEFAULT_DOPPLER_UNC_HZ - ) - doppler_unc = abs(SPEED_OF_LIGHT * rms_hz / f_tx) - else: - doppler_unc = abs(SPEED_OF_LIGHT * _DEFAULT_DOPPLER_UNC_HZ / f_tx) if not np.isnan(f_tx) else 1.0 - - if all(c in column_names for c in ("ax", "ay", "az")): - observer_acc = [float(d["ax"]), float(d["ay"]), float(d["az"])] - else: - observer_acc = [0.0, 0.0, 0.0] - - return Observation.from_radar_with_id( - str(objID), - delay_days, - doppler_audy, - has_delay, - has_doppler, - epoch_jd, - [d["x"], d["y"], d["z"]], # Barycentric observer position - [d["vx"], d["vy"], d["vz"]], # Barycentric observer velocity - delay_unc, - doppler_unc, - observer_acc, # Barycentric observer acceleration (au/day^2) - ) - - -def _append_observer_acceleration(data, observatory, dt_sec=2.0): - """Append barycentric observer acceleration columns (ax, ay, az; au/day^2). - - Finite-differences the barycentric station velocity from - ``obscodes_to_barycentric`` at the observation epoch +/- ``dt_sec``. Used only - by radar observations, whose two-leg light-time model extrapolates the station - state back to the signal transmit time. Requires the ``et`` column (seconds - past J2000, TDB) that ``orbitfit`` adds before computing the observer states. - """ - - def _vel(et_offset_sec): - shifted = data.copy() - shifted["et"] = data["et"] + et_offset_sec - pv = np.atleast_1d(observatory.obscodes_to_barycentric(shifted)) - return np.stack([pv["vx"], pv["vy"], pv["vz"]], axis=-1) - - # d(velocity[au/day]) / d(time[day]); 2*dt_sec seconds = 2*dt_sec/86400 days. - scale = 86400.0 / (2.0 * dt_sec) - acc = (_vel(dt_sec) - _vel(-dt_sec)) * scale - acc = np.atleast_2d(acc) - return rfn.append_fields( - data, ["ax", "ay", "az"], [acc[:, 0], acc[:, 1], acc[:, 2]], usemask=False, asrecarray=True ) @@ -682,320 +321,143 @@ def create_empty_result(id, dtypes): "NONE", # format ) + (np.nan,) * 36 # Flat covariance matrix - # non-grav columns (issue #351): NaN per a1/a2/a3 (+ _unc) that is present - + tuple(np.nan for n in ("a1", "a2", "a3") if n in dtypes.names for _ in (0, 1)) - # obs fingerprint (issue #419): empty hash never matches, so a failed - # or invalid fit is always retried next cycle rather than skipped. - + (("", 0) if "obs_hash" in dtypes.names else ()) ], dtype=dtypes, ) -def _carry_forward_result(prior_row, dtypes): - """Copy a prior fit result forward under the current output dtype (issue #419). +def do_gauss_iod(observations, seq): + """Calculate an initial orbit estimate using Gauss's method. - Used by the skip-unchanged path: when an object's observations are unchanged, - its stored fit is emitted verbatim rather than recomputed. Fields shared with - ``dtypes`` are copied by name (so the carried row is identical to the prior - fit); any field present in ``dtypes`` but absent from the prior (e.g. a - non-grav column the prior run did not fit) is left at its type default. - """ - prior_row = prior_row[0] if getattr(prior_row, "shape", None) == (1,) else prior_row - out = np.zeros(1, dtype=dtypes) - for name in dtypes.names: - if name in prior_row.dtype.names: - out[name][0] = prior_row[name] - return out - - -def _partition_unchanged(data, initial_guess, primary_id_column_name, fit_nongrav): - """Split raw observations into (to_fit, carried_forward) by obs fingerprint. - - The steady-state pre-filter for ``orbitfit(skip_unchanged=True)`` (issue #419): - an object whose fingerprint matches a converged prior fit over the same - observation set is carried forward verbatim (cast to the current output - dtype); every other object's rows are returned for fitting. Runs before any - ephemeris/observatory setup, so a skipped object costs only a fingerprint - hash. The fingerprint is computed on the raw (pre-debias) observations, so it - matches the one ``_orbitfit`` stores. + Parameters + ---------- + observations : list[Observation] + The list of Observations used for the orbit estimate + seq : list[list[int] + The list of lists of indexes of observations that are closely spaced in time. + + Returns + ------- + list[FitResult] + A collection of orbit fit results that can be used to perform a higher + quality fit estimate. """ - _, nongrav_names = _parse_nongrav(fit_nongrav) - out_dtype = _get_result_dtypes(primary_id_column_name, nongrav_names) - prior = np.atleast_1d(initial_guess) - prior_by_id = {row[primary_id_column_name]: row for row in prior} - ids = data[primary_id_column_name] - keep = np.ones(len(data), dtype=bool) - carried = [] - for oid in np.unique(ids): - obj_mask = ids == oid - nobs, obs_hash = _obs_fingerprint(data[obj_mask], data.dtype.names) - p = prior_by_id.get(oid) - if ( - p is not None - and int(p["flag"]) == 0 - and "nobs_fit" in prior.dtype.names - and int(p["nobs_fit"]) == nobs - and str(p["obs_hash"]) == obs_hash - ): - keep[obj_mask] = False - carried.append(_carry_forward_result(p, out_dtype)) - carried_arr = np.concatenate(carried) if carried else np.array([], dtype=out_dtype) - return data[keep], carried_arr + # Get gauss solution, using the first, middle, and last observation + # of the primary sequence + idx0, idx1, idx2 = seq[0][0], seq[0][int(len(seq[0]) / 2)], seq[0][-1] + logger.debug(f"Sequence indexs passed to gauss: {idx0}, {idx1}, {idx2}") + solns = gauss(GMtotal, observations[idx0], observations[idx1], observations[idx2], 0.0001, SPEED_OF_LIGHT) + return solns -def do_gauss_iod(observations, seq): - """Backward-compat wrapper for the Gauss IOD. - Prefer ``layup.iod.get_iod("gauss")`` for new code; this shim - exists so callers that imported ``do_gauss_iod`` directly continue - to work. - """ - return get_iod("gauss")(observations, seq) - - -# Multi-root picker tuning. Both knobs are exposed to do_fit() callers -# in case downstream code wants to override them, but the defaults are -# what worked best on the diagnostic/scan and neo_scan datasets. -_PICKER_MIN_R_HELIO_AU = 0.3 # reject roots with r < this as unphysical -_PICKER_SCREEN_ITER_MAX = 80 # cheap LM budget for the first pass -_PICKER_FULL_ITER_MAX = 100 # full LM budget for the fallback pass - -_PREFILTER_THRESHOLD_SIGMA = 1000.0 # held-out residual filter cutoff - -# IAS15 adaptive-step controller used during the multi-root picker. -# With the legacy controller (mode 1), LM grinds for minutes on phantom -# Gauss roots whose trajectories pass close to Earth (the integrator -# chases ever-smaller steps to resolve the close encounter — 100-1000× -# wallclock blowup observed on diagnostic/scan). The newer (Pham, Rein -# & Spiegel 2024) controller, mode 2, steps through those encounters -# efficiently: it brings the same pathological cases from >120 s to -# sub-second with the identical recovered orbit, and unlike a step-size -# floor it is a better controller rather than a truncation, so it costs -# no accuracy on genuine close-Earth encounters. Set to -1 to leave -# ASSIST's default (legacy mode 1). -_PICKER_IAS15_ADAPTIVE_MODE = 2 - - -# Cache of the Python-side assist.Ephem handle. The C-side get_ephem() -# from layup.routines returns the C struct; the Python residual filter -# needs the rebound/assist Python wrapper instead, so we cache one per -# cache_dir. -_assist_python_ephem_cache: dict = {} - - -def _get_python_ephem(cache_dir): - """Lazy-load and cache the Python-side assist.Ephem for the filter.""" - key = str(cache_dir) - if key in _assist_python_ephem_cache: - return _assist_python_ephem_cache[key] - try: - import assist - except ImportError: - return None - try: - eph = assist.Ephem(os.path.join(key, "linux_p1550p2650.440"), os.path.join(key, "sb441-n16.bsp")) - except Exception as e: - logger.warning(f"assist.Ephem load failed for {cache_dir}: {e}") - return None - _assist_python_ephem_cache[key] = eph - return eph - - -def _pick_best_root(candidates, min_r_au): - """Pick the best converged candidate from a list of LM results. - - "Best" means smallest χ² among candidates that - - 1. report ``flag == 0`` (LM converged), and - 2. have heliocentric distance > ``min_r_au`` (physical orbit). - - Returns None if no candidate satisfies (1); in that case the caller - typically retries at a larger LM budget. If (1) is met but (2) - isn't, the smallest-χ² convergent root is still returned (better - than nothing). +def do_herget_iod(observations, seq, args, aux): + """Calculate an initial orbit estimate using Herget's method. + + Parameters + ---------- + observations : list[Observation] + The list of Observations used for the orbit estimate + seq : list[list[int] + The list of lists of indexes of observations that are closely spaced in time. + + Returns + ------- + list[FitResult] + A collection of orbit fit results that can be used to perform a higher + quality fit estimate. """ - converged = [c for c in candidates if c.flag == 0] - if not converged: - return None - sane = [c for c in converged if (c.state[0] ** 2 + c.state[1] ** 2 + c.state[2] ** 2) > min_r_au**2] - pool = sane if sane else converged - return min(pool, key=lambda c: c.csq) - - -def do_fit( - observations, - seq, - cache_dir, - iod="auto", - engine="cartesian", - screen_iter_max: int = _PICKER_SCREEN_ITER_MAX, - full_iter_max: int = _PICKER_FULL_ITER_MAX, - min_r_helio_AU: float = _PICKER_MIN_R_HELIO_AU, - prefilter_threshold_sigma: float = _PREFILTER_THRESHOLD_SIGMA, - picker_ias15_adaptive_mode: int = _PICKER_IAS15_ADAPTIVE_MODE, -): - """Carry out an orbit fit to a list of observations. - - Pipeline: - 1. IOD: produce one or more candidate seed orbits via the - registered method named by `iod` (default: "auto"). The - registry lives in `layup.iod`; register new methods with - `iod.register_iod(name, callable)`. - 2. Multi-root picker: run LM from every IOD candidate on the - primary segment (`seq[0]`) at a cheap `screen_iter_max` - budget. Pick the smallest-χ² converged candidate with - heliocentric distance above `min_r_helio_AU`. If nothing - converges at the cheap budget, retry at `full_iter_max`. - Then refit on the full observation set. + # Get Herget solution, using first and last point + # of the primary sequence + solns = herget_with_assist(observations, seq, 0.001, args=args, aux=aux) + print(solns[0].niter) + return solns + + +def do_fit(observations, seq, cache_dir, iod="gauss", args=None, aux=None): + """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. + + A Gauss preliminary order is fit for the 0-th segment, + using the first, middle, and last observations in that + segment. + + Then an orbit fit is done on the 0-th segment, using the + initial orbit from Gauss. If that fails, any other preliminary + solutions are tried. + + Next, a fit to the full set of observations is attempted, given + the fit to the primary segment as an initial guess. If that + succeeds, the solution is returned. + + Otherwise, adjacent segments of observations are added and + the fit is updated, iteratively. Parameters ---------- observations : list - Time-ordered list of layup Observations. + A time-ordered list of observations seq : list of lists - Per-segment index lists; seq[0] is the primary segment. - cache_dir : str - Directory holding the ASSIST kernels. + A list of lists of observation indices. iod : str - Name of the registered IOD method "auto" (default) or "gauss". "auto" runs - Gauss and falls back to the BK 5-parameter linear IOD (run_bk_iod) on the - primary segment when every Gauss root fails to seed a converged fit. - 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. - screen_iter_max, full_iter_max : int - Two-tier LM iteration caps for the multi-root picker. - min_r_helio_AU : float - Lower bound on heliocentric distance for accepted IOD roots. + The IOD used to generate an initial guess orbit. Currently supports ['gauss']. + Default is 'gauss'. Returns ------- FitResult - Best converged fit (flag == 0) when one exists, else a - best-effort or sentinel FitResult with a non-zero flag. + The result of the orbit fit. """ - # 'auto' is a strategy, not a registered IOD: seed candidates with Gauss, - # then (after the picker, below) fall back to the BK 5-parameter linear IOD - # if every Gauss root fails to converge. Any other value is a registry name. - is_auto = isinstance(iod, str) and iod.lower() == "auto" - try: - iod_func = get_iod("gauss") if is_auto else (get_iod(iod) if isinstance(iod, str) else iod) - except ValueError as e: - raise ValueError(f"{e} Use iod.register_iod to add a new method.") - # Normalize to a list: an IOD may legitimately return None (e.g. Gauss - # finding no real roots), and the prefilter/picker below index and len() it. - solns = list(iod_func(observations, seq) or []) - - # If the iod produced no candidates, surface a sentinel -- unless we're in - # 'auto' mode, where the BK-IOD fallback below still has a shot. - if not solns and not is_auto: - logger.debug(f"IOD {iod!r} returned no candidates") + if iod.lower() == "gauss": + solns = do_gauss_iod(observations, seq) + elif iod.lower() == "herget": + solns = do_herget_iod(observations, seq, args, aux) + else: + 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: + logger.debug(f"The iod {iod} failed") x = FitResult() x.flag = 5 return x - # Pre-filter the IOD candidates by predicted-vs-observed residual - # on every observation. The right Gauss root predicts the full - # observation set within a few σ; phantom roots typically miss by - # 10⁵+ σ. Throwing those out before any LM iteration runs cuts - # the picker loop down to 1-2 LM fits per case in the common - # case (vs up to 8 brute-force LMs). Loose threshold (default - # 1000σ) so the right root is never rejected. - py_ephem = _get_python_ephem(cache_dir) - if py_ephem is not None and len(solns) > 1: - before = len(solns) - solns = filter_candidates_by_residual( - solns, observations, py_ephem, threshold_sigma=prefilter_threshold_sigma - ) - if len(solns) < before: - logger.debug( - f"IOD pre-filter: kept {len(solns)}/{before} " f"candidates at {prefilter_threshold_sigma}σ" - ) - assist_ephem = get_ephem(cache_dir) - # Multi-root picker. Fit every IOD candidate on the primary segment - # at the cheap screening budget, pick the best converged root, and - # only fall back to the full LM budget if nothing converged at the - # cheap tier. Gauss's polynomial gives up to 8 real roots; historic - # do_fit committed to solns[0] (largest r), which is often a - # phantom outer-SS solution for NEO-like targets. - # - # During this loop we select IAS15 adaptive_mode=2 so phantom roots - # whose trajectories pass close to Earth can't tie up the - # integrator for minutes (100-1000× wallclock blowup observed on - # diagnostic/scan with the legacy controller). The newer controller - # steps through close encounters efficiently with no accuracy cost. - # The setting is restored on every exit path. - # - # Each LM call dispatches through _run_fit so the picker honors the - # selected engine. The screen/full iteration budgets apply to the - # Cartesian engine; the BK-native engine uses its own internal cap. - saved_mode = get_ias15_adaptive_mode() - if picker_ias15_adaptive_mode >= 0: - set_ias15_adaptive_mode(picker_ias15_adaptive_mode) - + #! 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]] - try: - candidates = [_run_fit(assist_ephem, soln, obs, engine, screen_iter_max) for soln in solns] - x = _pick_best_root(candidates, min_r_helio_AU) - if x is None: - candidates = [_run_fit(assist_ephem, soln, obs, engine, full_iter_max) for soln in solns] - x = _pick_best_root(candidates, min_r_helio_AU) - finally: - set_ias15_adaptive_mode(saved_mode) - - if x is None and is_auto and len(obs) >= 3: - # Every Gauss root (if any) failed to seed a converged LM. Fall back to - # the BK 5-parameter linear IOD on the primary segment. BK-IOD shines on - # distant short arcs -- exactly where Gauss's three-point geometry is - # ill-conditioned (see bk_iod.cpp's regime-of-validity note); on the - # diagnostic scan Gauss+BK covers ~90% of cases vs ~84% for Gauss alone. - # Epoch convention matches do_gauss_iod's middle observation. - logger.debug(f"All {len(solns)} Gauss roots failed; trying BK-IOD fallback") - bk_seed = run_bk_iod(obs, float(obs[len(obs) // 2].epoch), _MU_SUN) - if bk_seed.flag == 0: - cand = _run_fit(assist_ephem, bk_seed, obs, engine, full_iter_max) - candidates.append(cand) - if cand.flag == 0: - x = cand - - if x is None: - # Still no convergence — surface the least-bad attempt so the caller has - # *something* to inspect, with a flag they can detect. - if not candidates: - # 'auto' with zero Gauss roots and no usable BK seed: nothing to - # surface, so return an explicit no-solution sentinel. - x = FitResult() - x.flag = 5 - return x - x = min(candidates, key=lambda c: c.csq) - logger.debug( - f"Primary interval: no root converged " f"(best csq={x.csq:.3g}, n_roots={len(candidates)})" - ) - x.flag = 3 + 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) + if x.flag != 0: + logger.debug(f"Primary interval failed. Total observations: {len(obs)}") + x.flag = 3 # caution return x # Attempt to fit all the data, given the fit of the primary interval - primary_x = x obs = observations - x = _run_fit(assist_ephem, x, obs, engine) + x = run_from_vector_with_initial_guess(assist_ephem, x, obs) # If that failed, build up the solution slowly if x.flag != 0: obs = [] - # Restart from the first IOD seed, or the converged primary fit when - # there were no IOD candidates (the iod='auto' BK-IOD fallback path). - x = solns[0] if solns else primary_x + x = solns[0] for i, sq in enumerate(seq): obs += [observations[i] for i in sq] - logger.debug(f"Incremental fit segment {i} of {len(seq)} " f"(n_obs={len(obs)})") - x = _run_fit(assist_ephem, x, obs, engine) + print(i, "of", len(seq), obs[0], sq) + x = run_from_vector_with_initial_guess(assist_ephem, x, obs) + print("flag:", x.flag) if x.flag != 0: x.flag = 4 break @@ -1014,31 +476,6 @@ def do_other_fit(iod: str): raise ValueError(f"The IOD, {iod} is not supported. Please use a supported IOD.") -# Minimum observational arc (in days) generally needed to constrain an orbit. -# Below this the fit is essentially unconstrained, so a failure is most likely a -# too-short baseline rather than anything wrong with the data. -_MIN_ARC_DAYS = 1.0 - - -def _warn_if_short_arc(jds, obj_id): - """Emit a helpful warning when a failed fit is likely caused by too short an - observational arc (less than ~24 hours / a single night). - - See issue #312: an orbit fit needs a baseline of more than 24 hours, so when - a fit fails on a sub-day arc we tell the user the likely cause rather than - leaving them with an opaque failure. - """ - if jds is None or len(jds) == 0: - return - arc_days = float(np.max(jds) - np.min(jds)) - if arc_days < _MIN_ARC_DAYS: - logger.warning( - f"Orbit fit failed for {obj_id}: the observations span only " - f"{arc_days * 24.0:.1f} hours. Constraining an orbit generally requires " - f"a baseline of more than ~24 hours (more than a single night of observations)." - ) - - def _orbitfit( data, cache_dir: str, @@ -1046,14 +483,10 @@ def _orbitfit( initial_guess=None, bias_dict: dict = None, sort_array: bool = True, - weight_data=False, # bool (Veres 2017) or "supplied" (rmsRA/rmsDec columns) - iod: str = "auto", - engine: str = "cartesian", - fit_nongrav: bool = False, - nongrav_auto_thresholds=None, - nongrav_gr=None, - per_arc: bool = False, - skip_unchanged: bool = False, + weight_data: bool = False, + iod: str = "gauss", + args=None, + aux=None, ): """This function will contain all of the calls to the c++ code that will calculate an orbit given a set of observations. Note that all observations @@ -1075,41 +508,14 @@ def _orbitfit( A dictionary containing bias corrections for different catalogs. sort_array : bool Whether to sort the observations by obstime before processing. Default is True. - weight_data : bool or str - Astrometric weighting. ``False`` (default) leaves the built-in default - uncertainty. ``True`` applies the Veres 2017 model (observation code, date, - catalog, program). ``"supplied"`` uses the per-observation ``rmsRA`` / - ``rmsDec`` columns (arcseconds) directly -- e.g. ADES-reported - uncertainties, or an external weighting model such as era-based historical - weighting for old comet apparitions (a row with a NaN/nonpositive value - falls back to the default). + weight_data : bool + Whether to apply data weighting based on the observation code, date, catalog + and program. Default is False. iod : str - The IOD used to generate an initial guess orbit. Supports 'gauss' - and 'auto' (Gauss with BK-IOD fallback). - Default is 'auto'. + The IOD used to generate an initial guess orbit. Currently supports ['gauss']. + Default is 'gauss'. """ - # Fitting non-gravitational params (issue #351) uses the joint state+nongrav - # LM, which only the Cartesian engine supports; the BK-native engine assumes a - # 6D state. Override with a warning, mirroring the radar path. - auto_nongrav = isinstance(fit_nongrav, str) and fit_nongrav.strip().lower() == "auto" - if auto_nongrav: - # 'auto' selects the non-grav model per object (issue #357); the schema - # carries all of A1/A2/A3 and each row reports only the adopted params. - nongrav_mask, nongrav_names = 0, ["A1", "A2", "A3"] - else: - nongrav_mask, nongrav_names = _parse_nongrav(fit_nongrav) - if (nongrav_mask or auto_nongrav) and engine != "cartesian": - logger.warning("Non-gravitational fitting requires engine='cartesian'; overriding %r.", engine) - engine = "cartesian" - - # Per-arc (piecewise-constant) non-grav amplitudes need an explicit non-grav - # mask (which params to split per apparition); it is meaningless for a - # gravity-only or 'auto'-selected fit. Ignore with a warning otherwise. - if per_arc and not nongrav_mask: - logger.warning("per_arc=True requires an explicit fit_nongrav mask (e.g. 'A1A2A3'); ignoring.") - per_arc = False - - _RESULT_DTYPES = _get_result_dtypes(primary_id_column_name, nongrav_names, per_arc=per_arc) + _RESULT_DTYPES = _get_result_dtypes(primary_id_column_name) if len(data) == 0: return np.array([], dtype=_RESULT_DTYPES) @@ -1139,47 +545,20 @@ def _orbitfit( # Check if certain columns are present in the data column_names = data.dtype.names - astcat_column_present = "astCat" in column_names + g_column_present = "astCat" in column_names program_column_present = "program" in column_names position_rates_columns_present = all(col in column_names for col in ["raRate", "decRate"]) - rate_unc_columns_present = all(col in column_names for col in ["rmsRArate", "rmsDecrate"]) - astrom_unc_columns_present = all(col in column_names for col in ["rmsRA", "rmsDec"]) - radar_columns_present = any(col in column_names for col in ["delay", "doppler"]) - - # Fingerprint the raw observation set (issue #419), before any in-place - # debiasing mutates ra/dec, so it reflects what was reported. - nobs_fit, obs_hash = _obs_fingerprint(data, column_names) - - # Lever 1 (skip-unchanged): if a prior converged fit for this object was - # built from the identical observation set, carry it forward verbatim - # instead of re-fitting. ``initial_guess`` has already been filtered to - # this object and reset to None when its flag != 0, so a match here is a - # successful prior fit over the same obs. Requires the prior catalog to - # carry the fingerprint columns; otherwise this never triggers and we fit. - if ( - skip_unchanged - and initial_guess is not None - and "obs_hash" in initial_guess.dtype.names - and "nobs_fit" in initial_guess.dtype.names - and int(initial_guess["nobs_fit"][0]) == nobs_fit - and str(initial_guess["obs_hash"][0]) == obs_hash - ): - return _carry_forward_result(initial_guess, _RESULT_DTYPES) # Accommodate occultation measurements. These measurements are implied when # the "ra" and "dec" columns are None. In this case, we will use the "starra" - # and "stardec" columns. Radar rows have no ra/dec and are skipped. + # and "stardec" columns. for d in data: - if radar_columns_present and _is_radar(d, column_names): - continue if _is_occultation(d): d = _use_star_astrometry(d) # bias_dict will be a dictionary when the debias flag is set to True. if bias_dict is not None: for d in data: - if radar_columns_present and _is_radar(d, column_names): - continue # debiasing is an astrometric (ra/dec) correction d["ra"], d["dec"] = debias( ra=d["ra"], dec=d["dec"], @@ -1194,38 +573,16 @@ def _orbitfit( # radians. observations = [] for d in data: - if radar_columns_present and _is_radar(d, column_names): - o = _radar_observation( - d[primary_id_column_name], - d, - convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), # JD TDB - column_names, - ) - observations.append(o) - continue if position_rates_columns_present and (not np.isnan(d["raRate"]) and not np.isnan(d["decRate"])): - # Rate uncertainties (rmsRArate/rmsDecrate) share raRate's - # arcsec/hour units; convert to rad/day. Absent -> C++ default. - streak_rate_unc = {} - if ( - rate_unc_columns_present - and not np.isnan(d["rmsRArate"]) - and not np.isnan(d["rmsDecrate"]) - ): - streak_rate_unc["ra_rate_unc"] = abs(d["rmsRArate"]) * ARCSEC_PER_HOUR_TO_RAD_PER_DAY - streak_rate_unc["dec_rate_unc"] = abs(d["rmsDecrate"]) * ARCSEC_PER_HOUR_TO_RAD_PER_DAY o = Observation.from_streak_with_id( str(d[primary_id_column_name]), d["ra"] * np.pi / 180.0, d["dec"] * np.pi / 180.0, - # arcsec/hour (great-circle) -> rad/day; raRate already - # carries the cos(Dec) factor (see module constant above). - d["raRate"] * ARCSEC_PER_HOUR_TO_RAD_PER_DAY, - d["decRate"] * ARCSEC_PER_HOUR_TO_RAD_PER_DAY, + d["raRate"], + d["decRate"], convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), # Convert obstime to JD TDB [d["x"], d["y"], d["z"]], # Barycentric position [d["vx"], d["vy"], d["vz"]], # Barycentric velocity - **streak_rate_unc, ) else: o = Observation.from_astrometry_with_id( @@ -1237,50 +594,22 @@ def _orbitfit( [d["vx"], d["vy"], d["vz"]], # Barycentric velocity ) - # Astrometric weighting. ``weight_data="supplied"`` uses the per-obs - # rmsRA/rmsDec columns (arcsec) directly -- e.g. ADES-reported - # uncertainties, or an external weighting model such as era-based - # historical weighting for old comet apparitions. ``weight_data=True`` - # uses the Veres 2017 model; ``False`` leaves the C++ default. Supplied - # takes precedence; a NaN/nonpositive supplied value on a row falls - # back to the C++ default for that row. - if isinstance(weight_data, str) and weight_data.lower() == "supplied": - if not astrom_unc_columns_present: - raise ValueError('weight_data="supplied" requires rmsRA and rmsDec columns (arcsec).') - if np.isfinite(d["rmsRA"]) and d["rmsRA"] > 0: - o.ra_unc = abs(d["rmsRA"]) * np.pi / (180.0 * 3600.0) - if np.isfinite(d["rmsDec"]) and d["rmsDec"] > 0: - o.dec_unc = abs(d["rmsDec"]) * np.pi / (180.0 * 3600.0) - elif weight_data: - # astrometric_uncertainty_Veres2017 returns the astrometric uncertainty in - # ARCSECONDS (per its docstring), but Observation.ra_unc / - # dec_unc are stored in RADIANS. Convert at the assignment. - sigma_arcsec = astrometric_uncertainty_Veres2017( + if weight_data: + data_weight = data_weight_Veres2017( obsCode=d["stn"], jd_tdb=convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), catalog=d["astCat"] if astcat_column_present else None, program=d["program"] if program_column_present else None, ) - sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) - o.ra_unc = sigma_rad - o.dec_unc = sigma_rad + o.ra_unc = data_weight + o.dec_unc = data_weight observations.append(o) - # Radar delay/Doppler rows use the variable-row packing, which only the - # Cartesian engine supports; the BK-native engine assumes 2 rows per - # observation and would silently drop them. - if radar_columns_present and engine != "cartesian": - logger.warning( - "Radar (delay/Doppler) observations require engine='cartesian'; overriding %r.", - engine, - ) - engine = "cartesian" - # if cache_dir is not provided, use the default os_cache if cache_dir is None: - kernels_loc = str(default_cache_dir()) + kernels_loc = str(pooch.os_cache("layup")) else: kernels_loc = str(cache_dir) @@ -1289,88 +618,23 @@ def _orbitfit( # Perform the orbit fitting if initial_guess is None or initial_guess["flag"] != 0: - if iod.lower() in ["gauss", "auto"]: + if iod.lower() in ["gauss", "herget"]: res = do_fit( observations=observations, seq=sequence, cache_dir=kernels_loc, iod=iod.lower(), - engine=engine, + args=args, + aux=aux, ) else: res = do_other_fit(iod=iod.lower()) else: guess_to_use = parse_fit_result(initial_guess) res = run_from_vector_with_initial_guess(get_ephem(kernels_loc), guess_to_use, observations) - - # Non-gravitational params (issue #351): once the 6-parameter orbit has - # converged, refine it jointly with the requested non-grav params, seeded - # from that solution. They are weakly constrained on short arcs, so if the - # joint fit is degenerate (flag 6) or fails to converge we keep the - # 6-parameter result and report the params as NaN (graceful guard). - if auto_nongrav and res.flag == 0: - # Adopt the most parsimonious statistically-warranted non-grav model, - # or keep the gravity-only fit (issue #357). - res = _select_nongrav_auto( - get_ephem(kernels_loc), - res, - observations, - nongrav_auto_thresholds or _AUTO_DEFAULT_THRESHOLDS, - gofr=nongrav_gr, - ) - elif nongrav_mask and res.flag == 0: - res_ng = run_from_vector_with_initial_guess( - get_ephem(kernels_loc), - res, - observations, - nongrav_mask=nongrav_mask, - gofr=_gofr_arg(nongrav_gr), - per_arc=per_arc, - ) - if res_ng.flag == 0: - res = res_ng - else: - logger.debug("Non-grav refinement did not converge; reporting non-grav params as NaN.") - # Populate our output structured array with the orbit fit results success = res.flag == 0 - if not success: - _warn_if_short_arc(jds, data[primary_id_column_name][0]) cov_matrix = tuple(res.cov[i] for i in range(36)) if success else (np.nan,) * 36 - nongrav_cols = () - if nongrav_names: - # Report each parameter only if it was actually adopted (its bit is set - # in the fit's nongrav_mask); for 'auto' that is the selected subset. - fitted_mask = getattr(res, "nongrav_mask", 0) if success else 0 - nongrav_cols = tuple( - v - for n in nongrav_names - for v in ( - (getattr(res, n.lower()) if (fitted_mask & _NONGRAV_BITS[n]) else np.nan), - (getattr(res, n.lower() + "_unc") if (fitted_mask & _NONGRAV_BITS[n]) else np.nan), - ) - ) - # Later-arc amplitudes (comet linkage): the C++ result reports them in the - # _arc2 fields only when per_arc fitting was on and converged. - per_arc_cols = () - if per_arc: - per_arc_on = success and getattr(res, "per_arc", False) - per_arc_cols = tuple( - v - for n in nongrav_names - for v in ( - ( - getattr(res, n.lower() + "_arc2") - if (per_arc_on and fitted_mask & _NONGRAV_BITS[n]) - else np.nan - ), - ( - getattr(res, n.lower() + "_arc2_unc") - if (per_arc_on and fitted_mask & _NONGRAV_BITS[n]) - else np.nan - ), - ) - ) output = np.array( [ ( @@ -1387,9 +651,6 @@ def _orbitfit( ("BCART_EQ" if success else "NONE"), # The base format returned by the C++ code ) + cov_matrix # Flat covariance matrix - + nongrav_cols # non-grav params + uncertainties (issue #351), when fit_nongrav - + per_arc_cols # later-arc amplitudes (comet linkage), when per_arc - + (obs_hash, nobs_fit) # obs fingerprint (issue #419) ], dtype=_RESULT_DTYPES, ) @@ -1405,13 +666,9 @@ def orbitfit( primary_id_column_name="provID", debias=False, weight_data=False, - iod="auto", - engine="cartesian", - fit_nongrav=False, - nongrav_auto_thresholds=None, - nongrav_gr=None, - per_arc=False, - skip_unchanged=False, + iod="gauss", + args=None, + aux=None, ): """This is the function that you would call interactively. i.e. from a notebook @@ -1429,79 +686,14 @@ def orbitfit( The name of the primary identifier column for the objects. Default is "provID". debias : bool Whether to apply debiasing corrections to the observations. Default is False. - weight_data : bool or str - Astrometric weighting. ``False`` (default) leaves the built-in default - uncertainty. ``True`` applies the Veres 2017 model (observation code, date, - catalog, program). ``"supplied"`` uses the per-observation ``rmsRA`` / - ``rmsDec`` columns (arcseconds) directly -- e.g. ADES-reported - uncertainties, or an external weighting model such as era-based historical - weighting for old comet apparitions (a row with a NaN/nonpositive value - falls back to the default). + weight_data : bool + Whether to apply data weighting based on the observation code, date, catalog + and program. Default is False. iod : str - The IOD used to generate an initial guess orbit. Supports 'gauss' - and 'auto' (Gauss with BK-IOD fallback). - Default is 'auto'. - fit_nongrav : bool | str | iterable of str - Which non-gravitational Marsden parameters to fit after the 6-parameter - orbit converges. ``False`` (default) fits none; ``True`` fits A2 (the - transverse Yarkovsky term, the common asteroid case); a string or iterable - naming params -- e.g. ``"A2"``, ``"A1A2A3"``, ``["A1", "A3"]`` -- selects a - subset of A1 (radial), A2 (transverse), A3 (normal). ``"auto"`` selects the - model adaptively per object (issue #357): the gravity-only fit is kept - unless its reduced chi-square is unacceptable, in which case the most - parsimonious non-grav model that is well-conditioned and statistically - significant is adopted (the A1/A2/A3 columns are all present, with only the - adopted params filled and the rest NaN). For each fitted param an ``a{n}`` - value and ``a{n}_unc`` 1-sigma column (au/day^2) are added to the result. - Cartesian engine only; params weakly constrained on short arcs are reported - as NaN (issue #351). - nongrav_auto_thresholds : NongravAutoThresholds, optional - Decision thresholds used when ``fit_nongrav="auto"`` -- how unacceptable the - gravity-only fit must be before a non-grav is tried, and how large the - chi-square drop and per-parameter significance must be to adopt one. Default - (``None``) uses the standard thresholds; pass a customized - ``NongravAutoThresholds`` to tune. Ignored unless ``fit_nongrav="auto"``. - nongrav_gr : sequence of float, optional - The non-gravitational g(r) sublimation law as ``[alpha, nm, nn, nk, r0]`` in - ASSIST's parameterization ``g(r) = alpha*(r/r0)^-nm*(1+(r/r0)^nn)^-nk``. - Default (``None``) is the asteroidal inverse-square law ``(r/r0)^-2`` used by - Yarkovsky A2 fits; pass a cometary law (e.g. Marsden water-ice) to fit a - comet's non-gravs. Applies to any non-grav fit (explicit or ``"auto"``). - per_arc : bool, optional - Fit piecewise-constant *per-apparition* non-grav amplitudes (comet - linkage). The state and ``g(r)`` are shared, but observations before the - fit epoch (the earlier arc) and after it (the later arc) each get their own - ``[A1,A2,A3]``. Requires an explicit ``fit_nongrav`` mask and an - ``initial_guess`` whose epoch sits between the two apparitions. The output - adds ``a{1,2,3}_arc2`` columns for the later arc; the base ``a{1,2,3}`` - columns then hold the earlier arc. Default False. - skip_unchanged : bool - Incremental / steady-state mode (issue #419). When True and ``initial_guess`` - is a prior result catalog carrying the ``obs_hash`` fingerprint columns, any - object whose observation set is byte-for-byte unchanged since that catalog is - carried forward verbatim instead of re-fitting. Objects with changed obs are - re-fit, warm-started from the prior state when available. Default False (every - object is fit). The output always carries the ``obs_hash``/``nobs_fit`` - columns so it can seed the next cycle. + The IOD used to generate an initial guess orbit. Currently supports ['gauss']. + Default is 'gauss'. """ - # Incremental / steady-state pre-filter (issue #419). Before any per-obs - # ephemeris or observatory setup, drop objects whose observation set is - # unchanged since the prior catalog and carry their stored fit forward - # verbatim. This is where the steady-state throughput win comes from -- a - # skipped object costs one fingerprint hash, not a fit. Changed objects fall - # through and are re-fit below, warm-started from the prior state via - # ``initial_guess`` (which is retained for exactly that purpose). - carried_forward = None - if ( - skip_unchanged - and initial_guess is not None - and "obs_hash" in getattr(initial_guess, "dtype", np.dtype([])).names - ): - data, carried_forward = _partition_unchanged(data, initial_guess, primary_id_column_name, fit_nongrav) - if len(data) == 0: # everything unchanged -> nothing to fit - return carried_forward - layup_observatory = LayupObservatory(cache_dir=cache_dir) # The units of et are seconds (from J2000). This new column is used by @@ -1512,17 +704,11 @@ def orbitfit( pos_vel = layup_observatory.obscodes_to_barycentric(data) data = rfn.merge_arrays([data, pos_vel], flatten=True, asrecarray=True, usemask=False) - # Radar (delay/Doppler) observations need the barycentric observer - # acceleration for the two-leg light-time model; compute it only when radar - # columns are present so optical/streak fits are unaffected. - if any(col in data.dtype.names for col in ("delay", "doppler")): - data = _append_observer_acceleration(data, layup_observatory) - bias_dict = None if debias: bias_dict = generate_bias_dict(cache_dir) - fitted = process_data_by_id( + return process_data_by_id( data, num_workers, _orbitfit, @@ -1532,370 +718,9 @@ def orbitfit( bias_dict=bias_dict, weight_data=weight_data, iod=iod, - engine=engine, - fit_nongrav=fit_nongrav, - nongrav_auto_thresholds=nongrav_auto_thresholds, - nongrav_gr=nongrav_gr, - per_arc=per_arc, - skip_unchanged=skip_unchanged, + args=args, + aux=aux, ) - # Re-attach objects carried forward unchanged by the #419 pre-filter. - if carried_forward is not None and len(carried_forward): - return np.concatenate([fitted, carried_forward]) - return fitted - - -def _observations_for_update(data, cache_dir, weight_data=False, bias_dict=None): - """Augment one object's observations with the observer barycentric state and - build the C++ ``Observation`` list, mirroring ``orbitfit()``'s preprocessing. - - Supports optical astrometry and streak (rate) rows -- the observation kinds a - steady-state catalog update sees. (Radar/occultation are not yet handled by - the sequential path; the driver's full-refit fallback covers them.) - """ - DEG = np.pi / 180.0 - kernels_loc = str(default_cache_dir()) if cache_dir is None else str(cache_dir) - observatory = LayupObservatory(cache_dir=cache_dir) - - et = np.array([spice.str2et(row["obsTime"]) for row in data], dtype=" the default layup os_cache). - all_data : numpy structured array, optional - The full observation set (old + new). Required for the nonlinearity - fallback: when the update is too large the driver refits over all_data. - weight_data, debias_data : bool - Apply Veres (2017) weighting / MPC debiasing to the new observations, - matching the corresponding ``orbitfit`` options. - max_update_sigma : float - Nonlinearity gate. If the update moves the state more than this many prior - standard deviations (Mahalanobis), fall back to a full refit over - ``all_data`` when provided, else flag the result (flag=8). - iter_max : int - LM iteration cap. - - Returns - ------- - FitResult - The updated fit. ``method`` is ``"sequential_update"`` for an accepted - information-filter update, or ``"orbit_fit"`` when the fallback refit ran. - """ - prior_fit = prior if isinstance(prior, FitResult) else parse_fit_result(prior) - kernels_loc = str(default_cache_dir()) if cache_dir is None else str(cache_dir) - ephem = get_ephem(kernels_loc) - bias_dict = generate_bias_dict(cache_dir) if debias_data else None - - new_obs = _observations_for_update(new_data, cache_dir, weight_data, bias_dict) - seq = run_sequential_update(ephem, prior_fit, new_obs, iter_max) - - def _full_refit(): - if all_data is None: - return None - all_obs = _observations_for_update(all_data, cache_dir, weight_data, bias_dict) - return run_from_vector_with_initial_guess(ephem, prior_fit, all_obs, iter_max) - - # The information update did not converge (e.g. a non-positive-definite prior, - # flag 7): fall back to a full refit if we can, else surface the failure. - if seq.flag != 0: - fallback = _full_refit() - return fallback if fallback is not None else seq - - # Nonlinearity gate: a large move means the linearization is untrustworthy. - if _update_mahalanobis(prior_fit, seq) > max_update_sigma: - fallback = _full_refit() - if fallback is not None: - return fallback - seq.flag = 8 # nonlinear update, no full-obs set supplied to refit - return seq - - -def _obs_row_keys(data, column_names): - """Per-observation identity keys over the fit-relevant columns (issue #419). - - Each key is the same per-row string that ``_obs_fingerprint`` hashes, so a - row's key equals its contribution to the object's fingerprint. Used to diff a - current observation set against the one a prior fit was built from. - """ - cols = [c for c in _FINGERPRINT_COLUMNS if c in column_names] - return ["\x1f".join(_fmt_fingerprint_value(d[c]) for c in cols) for d in data] - - -def _append_only_new_obs(current, prior_obs): - """Return the rows of ``current`` absent from ``prior_obs`` if the change is - append-only, else ``None``. - - Append-only means every observation the prior was fit from is still present - (none removed or modified) -- the case the sequential update handles exactly. - If any prior observation is gone, the summarised old-obs information no longer - matches the current set, so the object must be fully re-fit and this returns - ``None``. - """ - cur_keys = _obs_row_keys(current, current.dtype.names) - prior_keys = set(_obs_row_keys(prior_obs, prior_obs.dtype.names)) - if not prior_keys.issubset(cur_keys): - return None # an old observation was removed or changed -> not append-only - mask = np.array([k not in prior_keys for k in cur_keys], dtype=bool) - return current[mask] - - -def _group_by_id(data, primary_id_column_name): - if data is None: - return {} - ids = data[primary_id_column_name] - return {oid: data[ids == oid] for oid in np.unique(ids)} - - -def _fitresult_to_row(fit, obj_id, obs_hash, nobs_fit, dtypes): - """Pack a FitResult (from the sequential update or its refit fallback) into one - result-catalog row carrying the current-obs fingerprint, so the row is - identical in shape to ``orbitfit`` output and can seed the next cycle.""" - success = fit.flag == 0 - cov = tuple(fit.cov[i] for i in range(36)) if success else (np.nan,) * 36 - row = ( - (obj_id, (fit.csq if success else np.nan), fit.ndof) - + (tuple(fit.state[i] for i in range(6)) if success else (np.nan,) * 6) - + ( - (fit.epoch - 2400000.5) if success else np.nan, - fit.niter, - fit.method, - fit.flag, - "BCART_EQ" if success else "NONE", - ) - + cov - + (obs_hash, nobs_fit) - ) - return np.array([row], dtype=dtypes) - - -def incremental_orbitfit( - data, - cache_dir, - prior_catalog, - *, - prior_obs=None, - primary_id_column_name="provID", - weight_data=False, - debias=False, - max_update_sigma=4.0, - iod="auto", - engine="cartesian", - num_workers=1, -): - """Steady-state incremental fit over a batch of objects (issue #419 capstone). - - Ties the three levers into one operational maintenance pass. For each object - in ``data`` (the current observations), routes: - - * **skip** -- the observation set is unchanged since ``prior_catalog`` (matching - fingerprint): carry the prior fit forward verbatim, no fit. - * **sequential update** -- observations were only appended (``prior_obs`` given - and every prior observation is still present): update the prior with the new - observations only, via :func:`sequential_update` (integrating just the new - obs). Its nonlinearity gate falls back to a full refit when the update is - too large. - * **full refit** -- observations were removed or changed, or no per-object - ``prior_obs`` is available: refit over all current obs, warm-started from the - prior state when there is one, cold (IOD) when the object is new. - - Parameters - ---------- - data : numpy structured array - Current observations for all objects (grouped by ``primary_id_column_name``). - cache_dir : str or None - Kernel/ephemeris cache directory (None -> the default layup os_cache). - prior_catalog : numpy structured array or None - Prior fit results carrying state/cov/epoch and the ``obs_hash``/``nobs_fit`` - fingerprint columns (i.e. produced by ``orbitfit``/this driver). None -> every - object is cold-fit. - prior_obs : numpy structured array, optional - The observations the prior catalog was fit from, grouped by id. Enables the - sequential route (needs the per-object obs to diff). Without it, changed - objects are fully refit. - weight_data, debias : bool - Veres (2017) weighting / MPC debiasing, as in ``orbitfit``. - max_update_sigma : float - Nonlinearity gate passed to :func:`sequential_update`. - - Returns - ------- - (numpy structured array, dict) - The updated result catalog (one row per object, same schema as - ``orbitfit`` output) and a routing tally - ``{"skip", "sequential", "sequential_fallback", "full", "cold"}``. - """ - from collections import Counter - - pid = primary_id_column_name - out_dtype = _get_result_dtypes(pid) - prior_by_id = {row[pid]: row for row in np.atleast_1d(prior_catalog)} if prior_catalog is not None else {} - prior_obs_by_id = _group_by_id(prior_obs, pid) - - routing = Counter() - carried, seq_rows = [], [] - warm_ids, cold_ids = [], [] - - for oid in np.unique(data[pid]): - cur = data[data[pid] == oid] - nobs, obs_hash = _obs_fingerprint(cur, cur.dtype.names) - p = prior_by_id.get(oid) - has_prior = p is not None and int(p["flag"]) == 0 and "obs_hash" in np.atleast_1d(p).dtype.names - - # Route 1: unchanged -> skip (carry the prior row forward). - if has_prior and str(p["obs_hash"]) == obs_hash and int(p["nobs_fit"]) == nobs: - carried.append(_carry_forward_result(p, out_dtype)) - routing["skip"] += 1 - continue - - # Route 2: append-only change with the prior obs available -> sequential update. - if has_prior and oid in prior_obs_by_id: - new_rows = _append_only_new_obs(cur, prior_obs_by_id[oid]) - if new_rows is not None and len(new_rows) > 0: - seq = sequential_update( - p, - new_rows, - cache_dir, - all_data=cur, - weight_data=weight_data, - debias_data=debias, - max_update_sigma=max_update_sigma, - ) - routing["sequential" if seq.method == "sequential_update" else "sequential_fallback"] += 1 - seq_rows.append(_fitresult_to_row(seq, oid, obs_hash, nobs, out_dtype)) - continue - - # Route 3: full refit (warm if a prior exists, else cold IOD). - (warm_ids if has_prior else cold_ids).append(oid) - routing["full" if has_prior else "cold"] += 1 - - # Full/cold refits go through orbitfit (warm objects filtered to their priors; - # cold objects with no initial guess). Two calls keep _orbitfit's per-object - # initial-guess lookup happy (it errors on a guess with no row for the object). - fit_parts = [] - common = dict( - cache_dir=cache_dir, - primary_id_column_name=pid, - weight_data=weight_data, - debias=debias, - iod=iod, - engine=engine, - num_workers=num_workers, - ) - if warm_ids: - sub = data[np.isin(data[pid], warm_ids)] - fit_parts.append(orbitfit(sub, initial_guess=prior_catalog, **common)) - if cold_ids: - sub = data[np.isin(data[pid], cold_ids)] - fit_parts.append(orbitfit(sub, initial_guess=None, **common)) - - parts = ( - ([np.concatenate(carried)] if carried else []) - + ([np.concatenate(seq_rows)] if seq_rows else []) - + [p for p in fit_parts if len(p)] - ) - result = np.concatenate(parts) if parts else np.array([], dtype=out_dtype) - return result, dict(routing) def orbitfit_cli( @@ -1906,6 +731,7 @@ def orbitfit_cli( chunk_size: int = 10_000, num_workers: int = -1, cli_args: Optional[Namespace] = None, + aux: any = None, ): """This is the function that is called from the command line @@ -1934,15 +760,13 @@ def orbitfit_cli( weight_data = cli_args.weight_data output_orbit_format = cli_args.output_orbit_format iod = cli_args.iod - engine = getattr(cli_args, "engine", "cartesian") else: cache_dir = None debias = False guess_file = None weight_data = False output_orbit_format = "COM" # Default output orbit format. - iod = "auto" - engine = "cartesian" + iod = "gauss" _primary_id_column_name = cli_args.primary_id_column_name @@ -1977,7 +801,8 @@ def orbitfit_cli( else Path(f"{output_file_stem_flagged}.h5") ) - num_workers = resolve_num_workers(num_workers) + if num_workers < 0: + num_workers = os.cpu_count() # Check that input file exists if not input_file.exists(): @@ -2025,22 +850,6 @@ def orbitfit_cli( chunks = create_chunks(reader, chunk_size) - # Output is written one chunk at a time. The first write to each file - # overwrites any stale file left by a previous run (so re-runs are - # idempotent); later chunks append. Without this a re-run, or a re-merge onto - # an existing file, would duplicate every row. - _written_files = set() - - def _emit(arr, path): - first = path not in _written_files - _written_files.add(path) - if output_file_format == "hdf5": - (write_hdf5 if first else append_hdf5)(arr, path, key="data") - else: # csv: write_csv appends when the file already exists - if first and os.path.exists(path): - os.remove(path) - write_csv(arr, path) - for chunk in chunks: data = reader.read_objects(chunk) initial_guess = None @@ -2068,7 +877,8 @@ def _emit(arr, path): debias=debias, weight_data=weight_data, iod=iod, - engine=engine, + args=cli_args, + aux=aux, ) # Convert the fit_orbits to the preferred output format @@ -2087,14 +897,30 @@ def _emit(arr, path): fit_orbits_success = fit_orbits[success_mask] fit_orbits_failed = fit_orbits[~success_mask] - if len(fit_orbits_success) > 0: - _emit(fit_orbits_success, output_file) - - if len(fit_orbits_failed) > 0: - _emit(fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], output_file_flagged) + if output_file_format == "hdf5": + if len(fit_orbits_success) > 0: + write_hdf5(fit_orbits_success, output_file, key="data") + + if len(fit_orbits_failed) > 0: + write_hdf5( + fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], + output_file_flagged, + key="data", + ) + else: # csv output format + if len(fit_orbits_success) > 0: + write_csv(fit_orbits_success, output_file) + + if len(fit_orbits_failed) > 0: + write_csv( + fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], output_file_flagged + ) else: # All results go to a single output file - _emit(fit_orbits, output_file) + if output_file_format == "hdf5": + write_hdf5(fit_orbits, output_file, key="data") + else: + write_csv(fit_orbits, output_file) logger.info(f"Data has been written to {output_file}") @@ -2113,12 +939,13 @@ def _is_valid_data(data): bool True if the data is valid, False otherwise. """ - column_names = data.dtype.names valid_conditions = [ len(data) >= 3, np.all( data["et"] >= -6279962400.00 ), # excludes all datasets before 1801, data["et"] = 0 is j2000, 6279962400.00 is seconds between 1801 and j2000 + np.all(is_numeric(data["ra"])), + np.all(is_numeric(data["dec"])), np.all(is_numeric(data["x"])), np.all(is_numeric(data["y"])), np.all(is_numeric(data["z"])), @@ -2126,23 +953,6 @@ def _is_valid_data(data): np.all(is_numeric(data["vy"])), np.all(is_numeric(data["vz"])), ] - - # Each row must carry a usable observable: ra/dec for optical (and streak), - # or a delay/Doppler for radar. Validate per row so a radar file -- which has - # no ra/dec columns -- is not rejected, while optical rows still require - # numeric ra/dec. - if any(c in column_names for c in ["delay", "doppler"]): - have_radec = "ra" in column_names and "dec" in column_names - valid_conditions.append( - all( - _is_radar(d, column_names) or (have_radec and is_numeric(d["ra"]) and is_numeric(d["dec"])) - for d in data - ) - ) - else: - valid_conditions.append(np.all(is_numeric(data["ra"]))) - valid_conditions.append(np.all(is_numeric(data["dec"]))) - return all(valid_conditions) diff --git a/src/layup/utilities/herget_iod.py b/src/layup/utilities/herget_iod.py new file mode 100644 index 00000000..08e2efc4 --- /dev/null +++ b/src/layup/utilities/herget_iod.py @@ -0,0 +1,306 @@ +# Given an initial and final position, refine the velocity using +# Herget method as a way of creating a first guess for the inital orbit (IOD) +# import modules +import numpy as np +from sorcha.ephemeris.simulation_setup import create_assist_ephemeris +import assist +import rebound +from _layup_cpp._core import FitResult +from layup.utilities.universal_kepler import universal_step + +SPEED_OF_LIGHT_AU_DAY = 173.145 + + +def herget_with_assist(observations, seq, tolerance, args, aux, max_iterations=100): + """Runs the Herget method on a set of observations. + + Parameters + ---------- + observations : list + List of all the observations of the object + seq : list[list] + list of lists containing the indices of observations that are closely spaced in time + tolerance : float + the maximum delta_rho residuals allowed; will continue to converge until the residuals are below this value + args : argparse.Namespace + The argparse object that was created when running from the CLI. Needed to instantiate assist simulations + aux : LayupConfigs.auxiliary object + Auxiliary Layup configs; needed to instantiate assist simulations + max_iterations : int (optional, default: 100) + the maximum number of iterations before the fitting stops""" + seq_lengths = [len(i) for i in seq] + longest_i = np.argmax(seq_lengths) # finds the sequence index with the most observations contained in it + obs = np.array(observations)[seq[longest_i]] + + # Define our values + + obs_1 = obs[0] + r_e_1 = obs_1.observer_position + rho_hat_1 = np.array(obs_1.rho_hat) + rho_1 = 40 # this is the magnitude of rho, direction given by rho_hat, initial guess is 40au + t1 = obs_1.epoch + r1 = r_e_1 + rho_1 * rho_hat_1 + + obs_n = obs[-1] + r_e_n = obs_n.observer_position + rho_hat_n = np.array(obs_n.rho_hat) + rho_n = 40 # this is the magnitude of rho, direction given by rho_hat, initial guess is 40au + tn = obs_n.epoch + rn = r_e_n + rho_n * rho_hat_n + + iteration = 0 + delta_rho1 = tolerance + 1 + delta_rhon = tolerance + 1 + + # Get original epochs so we can light-time correct them each iteration + epochs = np.zeros(len(obs)) + for i, observation in enumerate(obs): + epochs[i] = observation.epoch + + while (abs(delta_rho1) + abs(delta_rhon)) / 2 > tolerance and iteration < max_iterations: + + # Light-time correct the observation times + for i, observation in enumerate(obs): + # print(observation.epoch) + observation.epoch = epochs[i] - ((rho_1) + (rho_n)) / (2 * SPEED_OF_LIGHT_AU_DAY) + # print(observation.epoch) + + delta_rho1, delta_rhon, state_1 = find_drho( + obs, t1, tn, r1, rn, tolerance, args, aux, rho_hat_1, rho_hat_n + ) + + # Update rho values + rho_1 -= delta_rho1 + r1 = r_e_1 + rho_1 * np.array(rho_hat_1) + rho_n -= delta_rhon + rn = r_e_n + rho_n * np.array(rho_hat_n) + print(delta_rho1, delta_rhon) + # print(rho_1, rho_n) + + iteration += 1 + + state = state_1 + solution = FitResult() + solution.state = state + solution.epoch = epochs[0] + solution.method = "herget" + solution.niter = iteration + solution.flag = 0 # Success flag + solution.ndof = len(observations) + solution.csq = 0.0 + solution.cov = [0.01] * 36 + + return [solution] + + +def find_drho(observations, t1, tn, r1, rn, tolerance, args, aux, rho_hat_1, rho_hat_n): + """Find the adjustment to make to rho_1 and rho_n to make in order to reduce the residuals of the observations + + Parameters + ---------- + observations : list + list of observation objects + t1 : float + light-time corrected time for position r1, in TDB MJD + tn : float + light-time corrected time for position rn, in TDB MJD + r1 : numpy array + position vector at time t1 + rn : numpy array + position vector at time tn + tolerance : float + the average value of delta_rho1 and delta_rhon at which the orbit is considered to have converged at + args : argparse.Namespace + The argparse object that was created when running from the CLI. Needed to instantiate assist simulations + aux : LayupConfigs.auxiliary object + Auxiliary Layup configs; needed to instantiate assist simulations + rho_hat_1 : numpy array + unit vector of rho at time t1 + rho_hat_n : numpy array + unit vector of rho at time tn + + Returns + ------- + delta_rho1 : float + the amount to adjust rho_1 by to return a more accurate orbit + delta_rhon : float + the amount to adjust rho_n by to return a more accurate orbit + state_1[x, y, z, vx, vy, vz] + the new guess for the state vector at t1 + """ + + # Find velocities at rho_1 and rho_n + [vx1, vy1, vz1], [vxn, vyn, vzn] = find_velocity(t1, tn, r1, rn, tolerance) + [var_vx1, var_vy1, var_vz1], _ = find_velocity(t1, tn, r1 + rho_hat_1, rn, tolerance) + + # Simulation setup + ephem, _, _ = create_assist_ephemeris(args, aux) + sim = rebound.Simulation() + + sim.add(x=r1[0], y=r1[1], z=r1[2], vx=vx1, vy=vy1, vz=vz1) + var = sim.add_variation(testparticle=0) + var.particles[0].xyz = rho_hat_1 + var.particles[0].vxyz = np.array([var_vx1 - vx1, var_vy1 - vy1, var_vz1 - vz1]) + + ex = assist.Extras(sim, ephem) + sim.t = t1 - ephem.jd_ref + a1, a2, b = np.zeros((3, 2 * len(observations))) + + # For each observation, integrate to that time and find the residuals + for i, observation in enumerate(observations): + + # For this observation, get A and D + A = observation.a_vec + D = observation.d_vec + + t = observation.epoch + sim.integrate(t - ephem.jd_ref) + + r_e = np.array(observation.observer_position) + r = sim.particles[0].xyz + r_var = var.particles[0].xyz + rho = r - r_e + + # Add these to the arrays + b[2 * i] = np.dot(rho / np.linalg.norm(rho), A) + b[2 * i + 1] = np.dot(rho / np.linalg.norm(rho), D) + a1[2 * i] = b[2 * i] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), A) + a1[2 * i + 1] = b[2 * i + 1] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), D) + + _, [var_vxn, var_vyn, var_vzn] = find_velocity(t1, tn, r1, rn + rho_hat_n, tolerance) + + # Do the same for rho_n, set up simulation again + vxn, vyn, vzn = sim.particles[0].vxyz + sim = rebound.Simulation() + sim.add(x=rn[0], y=rn[1], z=rn[2], vx=vxn, vy=vyn, vz=vzn) + var = sim.add_variation(testparticle=0) + var.particles[0].xyz = rho_hat_n + var.particles[0].vxyz = np.array([var_vxn - vxn, var_vyn - vyn, var_vzn - vzn]) + + ex = assist.Extras(sim, ephem) + sim.t = tn - ephem.jd_ref + + # Find residuals for each observation + for i, observation in enumerate(observations): + A = observation.a_vec + D = observation.d_vec + + t = observation.epoch + + sim.integrate(t - ephem.jd_ref) + r_e = np.array(observation.observer_position) + r = sim.particles[0].xyz + r_var = var.particles[0].xyz + rho = r - r_e + + # Add to array + a2[2 * i] = b[2 * i] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), A) + a2[2 * i + 1] = b[2 * i + 1] - np.dot((rho + r_var) / np.linalg.norm(rho + r_var), D) + + sigma_a1b = sum(a1 * b) + sigma_a2b = sum(a2 * b) + sigma_a1squared = sum(a1**2) + sigma_a2squared = sum(a2**2) + sigma_a1a2 = sum(a1 * a2) + + delta_rho1 = (sigma_a1b * sigma_a2squared - sigma_a2b * sigma_a1a2) / ( + sigma_a1a2**2 - sigma_a1squared * sigma_a2squared + ) + delta_rhon = (-delta_rho1 * sigma_a1squared - sigma_a1b) / sigma_a1a2 + + # Check this is the solution, should equal zero + # print(sigma_a1b + delta_rho1*sigma_a1squared + delta_rhon*sigma_a1a2) + # print(sigma_a2b + delta_rho1*sigma_a1a2 + delta_rhon*sigma_a2squared) + # print(sum(a1*(b + delta_rho1*a1 + delta_rhon*a2))) + + return delta_rho1, delta_rhon, [*r1, vx1, vy1, vz1] + + +def find_velocity(t1, tn, r1, rn, tolerance): + """Converge on a velocity which takes position r1 at time t1 to position rn at time tn. + Uses the universal kepler stepper to integrate over time. + + Parameters + ---------- + t1 : float + Light-time corrected time of first state, in TDB MJD + tn : float + Light-time corrected time of nth state, in TDB MJD + r1 : numpy array + Position vector at time t1 + rn : numpy array + Position vector at time tn + tolerance : float + how closely the calculated rn value must lie within the correct value + + Returns + ------- + state_1[vx, vy, vz] + The velocity components of state vector at t1 + state_n[vx, vy, vz] + The velocity components of state vector at tn + """ + + # Initialising data + delta_t = tn - t1 + state_1 = np.array([*r1, 0, 0, 0]) + + for i in range(3): + state_1[i + 3] = (rn[i] - r1[i]) / delta_t + + state_n = [*rn + abs(tolerance) + 1, 0, 0, 0] + + # Find new values for vx, vy and vz in turn + while np.linalg.norm(state_n[:3] - rn) > tolerance: + state_1[3:], state_n = find_new_vel_with_universal_kepler(t1, tn, state_1, rn) + + return state_1[3:], state_n[3:] + + +def find_mag_to_adjust(P, Q, R): + # This is the formula for the point on a line (defined by Q and R) + # that is closest to a point outside the line, P + # For our purpose, this is the scale factor to vary the velocity by so that it + # will be closest to rho_n next time + mag = np.dot(R - Q, P - Q) / np.dot(R - Q, R - Q) + return mag + + +def find_new_vel_with_universal_kepler(t1, tn, state_1, state_n): + """Adjust the velocity components of an input position-velocity state to land closer to a desired position + + Parameters + ---------- + t1 : float + Light-time corrected time of first state, in TDB MJD + tn : float + Light-time corrected time of nth state, in TDB MJD + state_1 : numpy array + the state vector at time t1; in AU and AU/day + state_n : numpy array + the state vector at time tn; in AU and AU/day + + Returns + ------- + state_1[vx, vy, vz] + Adjusted velocity components of state_1 + state_n[x, y, z, vx, vy, vz] + The new state vector at time tn when using state_1 as the input""" + + # Initialise variables + dt = tn - t1 + GMtotal = 0.0002963092748799319 + + for i in range(3): + variation = np.zeros(6) + variation[i + 3] = 1 # we are varying vx, vy and vz by 1, once at a time + var_results = universal_step(GMtotal, dt, state_1, variation=variation) + + diff = find_mag_to_adjust( + np.array(state_n[:3]), + np.array(var_results.state[:3]), + np.array(var_results.state[:3]) + np.array(var_results.variation[:3]), + ) # find the velocity in this cartesian direction that will get us closest to the desired position + state_1[i + 3] += diff + + return state_1[3:], np.array(var_results.state) diff --git a/src/layup/utilities/universal_kepler.py b/src/layup/utilities/universal_kepler.py new file mode 100644 index 00000000..9a1af744 --- /dev/null +++ b/src/layup/utilities/universal_kepler.py @@ -0,0 +1,405 @@ +"""Universal-variable Kepler propagator with variational partials. + +Python port of ``universal-kepler.c`` in this directory (Danby 1988, p.178), +intended as the numerical core of a Herget initial-orbit-determination +prototype. + +WHAT THIS IS +------------ +The *initial*-value problem: given a state (r0, v0) and an interval dt, +return the state at t0 + dt, optionally propagating a 6-vector deviation +alongside it. The deviation is the state-transition matrix applied to a +direction -- which is what Herget actually needs, since only two directions +matter (d/d rho_1 and d/d rho_3), not the full 6x6. + +WHAT THIS IS NOT +---------------- +A Lambert solver. It does not solve the two-point boundary-value problem +(r1, t1, r3, t3 -> v1). It is, however, the natural engine for a *shooting* +Lambert: guess v1, propagate here, compare to r3, and correct using the +variational output as the Jacobian d r(t3) / d v1. + +Handles elliptic, hyperbolic and near-parabolic motion without branching on +a bound-orbit assumption -- the sign of ``alpha = gm/a`` selects the regime. + +UNITS are whatever ``gm`` is expressed in; for AU and days use layup's +``constants.MU_SUN`` (heliocentric) or ``constants.GMtotal`` (barycentric), +matching the frame the state is expressed in. The two differ by 0.13% and +layup uses both deliberately, so pick consciously. + +DIFFERENCES FROM THE C +---------------------- +1. **Multi-revolution fix.** The C is silently wrong for ``|dt|`` longer + than one orbital period: its elliptic branch reduces the local ``dt`` by + whole revolutions to build the initial guess, but the Newton iteration + then solves against that *reduced* value while ``universal_step`` forms + ``g = dt - gm*g3`` from the *unreduced* one. The Lagrange g comes out + too large by exactly ``n_rev * P``, the state runs away along the + velocity direction, and the routine still reports success. + + Here the reduction is used for the guess *only*; ``n_rev`` revolutions + are added back as ``n_rev * 2*pi / sqrt(alpha)`` (one revolution advances + the universal anomaly by 2*pi) and the Newton solves the FULL dt. That + is deliberately not just a guard: reducing dt and keeping the result + would give the right *state* but the wrong *partials*, because + neighbouring orbits have different periods and the deviation grows + secularly with every revolution. A state-only fix would silently break + exactly the thing this module exists to provide. + +2. **Relative convergence tolerance.** The C tests ``|ds| > EPS`` with + ``EPS = 1e-13`` as an *absolute* tolerance on ``s``, which carries units + of time/length. Over one revolution ``s`` spans ``P/a``: about 577 for + a mainbelt orbit but **2395 for a TNO at 43 AU**, so the test demands + 4e-17 in relative terms -- below double precision, and unreachable. + Newton then exhausts its six iterations, Laguerre-Conway its fifteen, + and ``kepler()`` returns ``KEPLER_FLAG`` -- which ``universal_step`` + ignores, computing a state from the unconverged ``s`` regardless. The + failure therefore gets worse the more distant the object, which is + precisely backwards for this project. (The C's author evidently saw + it: both loops carry a commented-out alternative on ``f/dt``, a + scale-invariant relative residual.) Here the test is relative to + ``max(1, |s|)``. + +3. ``r`` is recomputed from the converged g-functions rather than reused + from the last Newton evaluation (the C's ``*rx = fp`` is one iteration + stale -- negligible once converged, but free to do properly). + +4. ``gdot`` falls back to ``1 - (gm/r)*g2`` when ``|f|`` is too small for + the Wronskian form ``(1 + g*fdot)/f`` the C uses unconditionally. + +5. Non-convergence raises rather than returning a flag the caller may + ignore -- as ``universal_step`` itself does in the C. + +The unused ``stumpff()`` in the C (dead code -- both call sites use +``cfun()``) is not ported. Note that layup's own +``utilities.orbit_conversion.stumpff`` returns c0..c3 only; the variational +path here needs c4 and c5 via g1a/g2a/g3a, which is why ``cfun`` is ported +rather than reused. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import numpy as np + +__all__ = [ + "KeplerConvergenceError", + "KeplerStep", + "stumpff_c", + "universal_step", + "state_transition_matrix", +] + +# Iteration budgets, matching the C. +_NEWTON_MAX = 6 +_LAGCON_MAX = 15 +_TWO_PI = 2.0 * math.pi + +# Convergence tolerance on the universal variable s, applied RELATIVE to |s| +# (see note 2 in the module docstring). The C uses this as an absolute +# tolerance, which is unreachable for distant orbits. +_TOL = 1e-13 + + +def _converged(ds: float, s: float) -> bool: + return abs(ds) <= _TOL * max(1.0, abs(s)) + + +# Below this |f| the Wronskian form of gdot loses precision; see note 3 above. +_F_FLOOR = 1e-8 + + +class KeplerConvergenceError(RuntimeError): + """Neither Newton nor Laguerre-Conway reached the tolerance.""" + + +@dataclass +class KeplerStep: + """Result of one propagation. + + Attributes + ---------- + state : (6,) ndarray + Propagated (x, y, z, vx, vy, vz). + variation : (6,) ndarray or None + The input deviation propagated to the same epoch, i.e. the + state-transition matrix applied to it. None if none was supplied. + s : float + Converged universal variable. + r : float + Heliocentric (or barycentric) distance at the end of the step. + n_rev : int + Whole revolutions spanned by dt. Nonzero means the multi-revolution + path was exercised -- the regime the C got wrong. + n_iter : int + Iterations used by whichever solver converged. + solver : str + "newton" or "laguerre-conway". + """ + + state: np.ndarray + variation: np.ndarray | None + s: float + r: float + n_rev: int + n_iter: int + solver: str + + +def stumpff_c(z: float) -> tuple[float, float, float, float, float, float]: + """Stumpff functions c0..c5 by Mikkola's argument four-folding. + + Port of the C ``cfun()``. The four-folding keeps the series arguments + small (|h| < 0.1) so the truncated rational approximations for c4 and c5 + stay accurate, then unfolds with the duplication identities. + """ + h = z + k = 0 + while abs(h) >= 0.1: + h *= 0.25 + k += 1 + + c4 = (1.0 - h * (1.0 - h * (1.0 - h / 90.0 / (1.0 + h / 132.0)) / 56.0) / 30.0) / 24.0 + c5 = (1.0 - h * (1.0 - h * (1.0 - h / 110.0 / (1.0 + h / 156.0)) / 72.0) / 42.0) / 120.0 + + for _ in range(k): + c3 = 1.0 / 6.0 - h * c5 + c2 = 0.5 - h * c4 + c5 = (c5 + c4 + c2 * c3) / 16.0 + c4 = c3 * (2.0 - h * c3) / 8.0 + h *= 4.0 + + c3 = 1.0 / 6.0 - z * c5 + c2 = 0.5 - z * c4 + c1 = 1.0 - z * c3 + c0 = 1.0 - z * c2 + return c0, c1, c2, c3, c4, c5 + + +def _initial_guess(gm: float, dt: float, r0: float, alpha: float, u: float) -> tuple[float, int]: + """Starting value for the universal variable, and the revolution count. + + Returns (s_guess, n_rev). n_rev is nonzero only on the elliptic branch, + and the returned guess already includes the whole revolutions -- callers + solve against the full dt. + """ + # Short step relative to the current distance: the series guess is good + # and there is no revolution structure to unwrap. + if abs(dt / r0) <= 0.2: + return dt / r0 - (dt * dt * u) / (2.0 * r0**3), 0 + + if alpha <= 0.0: + # Hyperbolic. a < 0 here, so the sqrt arguments below are positive. + a = gm / alpha + en = math.sqrt(-gm / (a * a * a)) + ch = 1.0 - r0 / a + sh = u / math.sqrt(-a * gm) + e = math.sqrt(ch * ch - sh * sh) + dm = en * dt + if dm < 0.0: + return -math.log((-2.0 * dm + 1.8 * e) / (ch - sh)) / math.sqrt(-alpha), 0 + return math.log((2.0 * dm + 1.8 * e) / (ch + sh)) / math.sqrt(-alpha), 0 + + # Elliptic. + a = gm / alpha + en = math.sqrt(gm / (a * a * a)) + ec = 1.0 - r0 / a + es = u / (en * a * a) + + # Whole revolutions are stripped to keep the RK4 guess inside one orbit, + # then added back below. Truncation toward zero (not floor) so the + # remainder keeps dt's sign, matching the C's (int) cast -- but in Python + # the int is arbitrary precision, so the C's overflow at ~2e9 revolutions + # does not apply. + n_rev = int(en * dt / _TWO_PI) + dt_red = dt - n_rev * _TWO_PI / en + + y = en * dt_red - es + + # One RK4 step of the (ec, es) rotation, as a cheap high-order guess for + # the eccentric-anomaly increment. (The C also computes the eccentricity + # here for Danby's alternative guess, which is commented out; omitted.) + xx, yy = ec, es + h = en * dt_red + omx = h / (1.0 - xx) + k0x, k0y = -yy * omx, xx * omx + xx1, yy1 = xx + k0x / 2.0, yy + k0y / 2.0 + omx = h / (1.0 - xx1) + k1x, k1y = -yy1 * omx, xx1 * omx + xx1, yy1 = xx + k1x / 2.0, yy + k1y / 2.0 + omx = h / (1.0 - xx1) + k2x, k2y = -yy1 * omx, xx1 * omx + xx1 = xx + k2x + omx = h / (1.0 - xx1) + k3y = xx1 * omx + yy += (k0y + 2.0 * (k1y + k2y) + k3y) / 6.0 + + root_alpha = math.sqrt(alpha) + # sqrt(alpha)*s is the eccentric-anomaly increment, so a revolution is + # worth 2*pi/sqrt(alpha) in s. + s = (y + yy) / root_alpha + n_rev * _TWO_PI / root_alpha + return s, n_rev + + +def _solve_kepler(gm, dt, r0, alpha, u, zeta): + """Solve the universal Kepler equation for s; return the g-functions. + + Returns (g0..g5, r, s, n_rev, n_iter, solver). + """ + s_guess, n_rev = _initial_guess(gm, dt, r0, alpha, u) + + def _fvals(s): + c0, c1, c2, c3, _, _ = stumpff_c(s * s * alpha) + c1 *= s + c2 *= s * s + c3 *= s * s * s + f = r0 * c1 + u * c2 + gm * c3 - dt + fp = r0 * c0 + u * c1 + gm * c2 # == r at this s, hence always > 0 + fpp = zeta * c1 + u * c0 + fppp = zeta * c0 - u * alpha * c1 + return f, fp, fpp, fppp + + # Newton with Danby's cubic correction: three nested refinements of the + # same step reuse one function evaluation. + s = s_guess + ds = math.inf + solver, n_iter = "newton", 0 + for n_iter in range(1, _NEWTON_MAX + 1): + f, fp, fpp, fppp = _fvals(s) + ds = -f / fp + ds = -f / (fp + ds * fpp / 2.0) + ds = -f / (fp + ds * fpp / 2.0 + ds * ds * fppp / 6.0) + s += ds + if _converged(ds, s): + break + + if not _converged(ds, s): + # Laguerre-Conway from the original guess: larger convergence basin, + # at the cost of a square root per iteration. + solver = "laguerre-conway" + s = s_guess + ln = 5.0 + for n_iter in range(1, _LAGCON_MAX + 1): + f, fp, fpp, _ = _fvals(s) + disc = (ln - 1.0) ** 2 * fp * fp - (ln - 1.0) * ln * f * fpp + ds = -ln * f / (fp + math.copysign(math.sqrt(abs(disc)), fp)) + s += ds + if _converged(ds, s): + break + if not _converged(ds, s): + raise KeplerConvergenceError( + f"Kepler equation did not converge: dt={dt!r} r0={r0!r} " + f"alpha={alpha!r} u={u!r} last |ds|={abs(ds):.3e} s={s!r}" + ) + + c0, c1, c2, c3, c4, c5 = stumpff_c(s * s * alpha) + g0 = c0 + g1 = c1 * s + g2 = c2 * s**2 + g3 = c3 * s**3 + g4 = c4 * s**4 + g5 = c5 * s**5 + r = r0 * g0 + u * g1 + gm * g2 + return g0, g1, g2, g3, g4, g5, r, s, n_rev, n_iter, solver + + +def universal_step(gm, dt, state, variation=None) -> KeplerStep: + """Propagate `state` by `dt`, optionally carrying a deviation along. + + Parameters + ---------- + gm : float + Gravitational parameter, in units consistent with `state` and `dt`. + dt : float + Interval. May be negative. + state : array_like, shape (6,) + (x, y, z, vx, vy, vz) at the start of the step. + variation : array_like, shape (6,), optional + A deviation in the *initial* state. Propagated exactly (to the + two-body model) to the end of the step, i.e. the state-transition + matrix applied to this vector. + + Returns + ------- + KeplerStep + """ + state = np.asarray(state, dtype=float) + if state.shape != (6,): + raise ValueError(f"state must have shape (6,), got {state.shape}") + + r0vec, v0vec = state[:3], state[3:] + r0 = float(np.linalg.norm(r0vec)) + if r0 == 0.0: + raise ValueError("state has zero position; the two-body problem is singular there") + + v0s = float(v0vec @ v0vec) + u = float(r0vec @ v0vec) + alpha = 2.0 * gm / r0 - v0s # = gm/a; sign selects ellipse vs hyperbola + zeta = gm - alpha * r0 + + g0, g1, g2, g3, g4, g5, r, s, n_rev, n_iter, solver = _solve_kepler(gm, dt, r0, alpha, u, zeta) + + f = 1.0 - (gm / r0) * g2 + g = dt - gm * g3 + fdot = -(gm / (r * r0)) * g1 + # The Wronskian form f*gdot - fdot*g = 1 is better conditioned than + # 1 - (gm/r)*g2 except where f itself is near zero. + gdot = (1.0 + g * fdot) / f if abs(f) > _F_FLOOR else 1.0 - (gm / r) * g2 + + out = np.empty(6) + out[:3] = f * r0vec + g * v0vec + out[3:] = fdot * r0vec + gdot * v0vec + + var_out = None + if variation is not None: + dvar = np.asarray(variation, dtype=float) + if dvar.shape != (6,): + raise ValueError(f"variation must have shape (6,), got {dvar.shape}") + dr, dv = dvar[:3], dvar[3:] + + # Derivatives of the scalars the solve depends on, along the deviation. + r0pr = float(r0vec @ dr) / r0 + alphapr = -(2.0 * gm / (r0 * r0)) * r0pr - 2.0 * float(v0vec @ dv) + upr = float(r0vec @ dv) + float(v0vec @ dr) + zetapr = -alpha * r0pr - r0 * alphapr + + # d g_k / d alpha at fixed s (Stumpff recurrences). + g1a = 0.5 * (g3 - s * g2) + g2a = 0.5 * (2.0 * g4 - s * g3) + g3a = 0.5 * (3.0 * g5 - s * g4) + + # d s / d(deviation) from differentiating the Kepler equation at fixed dt. + spr = -(s * r0pr + g3 * zetapr + g2 * upr + (g3a * zeta + u * g2a) * alphapr) / r + + g1pr = g0 * spr + g1a * alphapr + g2pr = g1 * spr + g2a * alphapr + g3pr = g2 * spr + g3a * alphapr + rpr = r0pr + g1 * upr + g2 * zetapr + u * g1pr + zeta * g2pr + + fpr = (gm * g2 / (r0 * r0)) * r0pr - (gm / r0) * g2pr + gpr = -gm * g3pr + fdotpr = (gm / (r * r * r0)) * g1 * rpr + (gm / (r * r0 * r0)) * g1 * r0pr - (gm / (r * r0)) * g1pr + gdotpr = (gm / (r * r)) * g2 * rpr - (gm / r) * g2pr + + var_out = np.empty(6) + var_out[:3] = f * dr + g * dv + fpr * r0vec + gpr * v0vec + var_out[3:] = fdot * dr + gdot * dv + fdotpr * r0vec + gdotpr * v0vec + + return KeplerStep(state=out, variation=var_out, s=s, r=r, n_rev=n_rev, n_iter=n_iter, solver=solver) + + +def state_transition_matrix(gm, dt, state) -> np.ndarray: + """Full 6x6 d state(t0+dt) / d state(t0), by six unit variations. + + Herget needs only two columns of this (d/d rho_1 and d/d rho_3), so + prefer calling `universal_step` directly with the deviation you care + about; this exists for testing and for callers that want the whole thing. + """ + stm = np.empty((6, 6)) + for j in range(6): + e = np.zeros(6) + e[j] = 1.0 + stm[:, j] = universal_step(gm, dt, state, variation=e).variation + return stm diff --git a/tests/data/2000OK67_ephem.csv b/tests/data/2000OK67_ephem.csv new file mode 100644 index 00000000..63530972 --- /dev/null +++ b/tests/data/2000OK67_ephem.csv @@ -0,0 +1,62 @@ +provID,ra,dec,mjd_utc,obsTime,stn,rmsRA,rmsDec,rmsTime,rmsCorr,uncTime +2000 OK67,9.32151,4.45406,61253.0,2026-Aug-01 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.31419,4.4515,61254.0,2026-Aug-02 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.30654,4.44879,61255.0,2026-Aug-03 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.29855,4.44595,61256.0,2026-Aug-04 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.29023,4.44295,61257.0,2026-Aug-05 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.28157,4.43982,61258.0,2026-Aug-06 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.27259,4.43654,61259.0,2026-Aug-07 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.26328,4.43312,61260.0,2026-Aug-08 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.25364,4.42956,61261.0,2026-Aug-09 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.24368,4.42587,61262.0,2026-Aug-10 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.23341,4.42203,61263.0,2026-Aug-11 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.22282,4.41806,61264.0,2026-Aug-12 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.21192,4.41396,61265.0,2026-Aug-13 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.20071,4.40972,61266.0,2026-Aug-14 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.18921,4.40536,61267.0,2026-Aug-15 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.17741,4.40087,61268.0,2026-Aug-16 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.16531,4.39625,61269.0,2026-Aug-17 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.15293,4.3915,61270.0,2026-Aug-18 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.14027,4.38664,61271.0,2026-Aug-19 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.12733,4.38165,61272.0,2026-Aug-20 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.11412,4.37655,61273.0,2026-Aug-21 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.10064,4.37133,61274.0,2026-Aug-22 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0869,4.36599,61275.0,2026-Aug-23 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0729,4.36055,61276.0,2026-Aug-24 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.05865,4.35499,61277.0,2026-Aug-25 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.04415,4.34932,61278.0,2026-Aug-26 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.0294,4.34355,61279.0,2026-Aug-27 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,9.01443,4.33768,61280.0,2026-Aug-28 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.99921,4.3317,61281.0,2026-Aug-29 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.98377,4.32563,61282.0,2026-Aug-30 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.96811,4.31945,61283.0,2026-Aug-31 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.95223,4.31318,61284.0,2026-Sep-01 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.93614,4.30682,61285.0,2026-Sep-02 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.91984,4.30037,61286.0,2026-Sep-03 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.90334,4.29382,61287.0,2026-Sep-04 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.88665,4.28719,61288.0,2026-Sep-05 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.86976,4.28048,61289.0,2026-Sep-06 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.85269,4.27368,61290.0,2026-Sep-07 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.83545,4.2668,61291.0,2026-Sep-08 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.81803,4.25985,61292.0,2026-Sep-09 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.80044,4.25282,61293.0,2026-Sep-10 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.7827,4.24572,61294.0,2026-Sep-11 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.7648,4.23855,61295.0,2026-Sep-12 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.74676,4.23132,61296.0,2026-Sep-13 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.72859,4.22403,61297.0,2026-Sep-14 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.71028,4.21667,61298.0,2026-Sep-15 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.69185,4.20926,61299.0,2026-Sep-16 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.6733,4.2018,61300.0,2026-Sep-17 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.65465,4.19429,61301.0,2026-Sep-18 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.63589,4.18673,61302.0,2026-Sep-19 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.61704,4.17912,61303.0,2026-Sep-20 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.5981,4.17148,61304.0,2026-Sep-21 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.57907,4.16379,61305.0,2026-Sep-22 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.55998,4.15608,61306.0,2026-Sep-23 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.54081,4.14832,61307.0,2026-Sep-24 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.52159,4.14054,61308.0,2026-Sep-25 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.50231,4.13273,61309.0,2026-Sep-26 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.48299,4.1249,61310.0,2026-Sep-27 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.46362,4.11704,61311.0,2026-Sep-28 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.44422,4.10917,61312.0,2026-Sep-29 00:00:00.000,X05,0.0001,0.0001,,0.0001, +2000 OK67,8.42479,4.10128,61313.0,2026-Sep-30 00:00:00.000,X05,0.0001,0.0001,,0.0001, diff --git a/tests/layup/test_herget_iod.py b/tests/layup/test_herget_iod.py new file mode 100644 index 00000000..5d1a0367 --- /dev/null +++ b/tests/layup/test_herget_iod.py @@ -0,0 +1,256 @@ +import os +import numpy as np +from numpy.lib import recfunctions as rfn +from numpy.testing import assert_allclose, assert_equal +import layup.utilities.herget_iod as herget +import spiceypy as spice + +from layup.utilities.layup_configs import LayupConfigs +from sorcha.ephemeris.simulation_setup import create_assist_ephemeris +from layup.utilities.data_utilities_for_tests import get_test_filepath +from layup.utilities.file_io.CSVReader import CSVDataReader +from layup.routines import Observation +from layup.utilities.data_processing_utilities import LayupObservatory +from layup.utilities.datetime_conversions import convert_tdb_date_to_julian_date +from layup.orbitfit import _build_sequence + +from sorcha.ephemeris.simulation_setup import create_assist_ephemeris +import assist +import rebound + +SPEED_OF_LIGHT_AU_DAY = 173.145 + + +def test_find_mag_to_adjust(): + """Testing that the formula for finding the closest point on a line works by using a trivial case""" + + point = np.array([1, 2]) + line_point1 = np.array([0, 0]) + line_point2 = np.array([1, 0]) + + answer = 1 + mag = herget.find_mag_to_adjust(point, line_point1, line_point2) + assert mag == answer + + # Trying again for the edge case that the point lies on the line + point = np.array([0, 0]) + answer = 0 + mag = herget.find_mag_to_adjust(point, line_point1, line_point2) + assert mag == answer + + +def test_find_new_vel_with_universal_kepler(): + """Start with an object moving away from target and check it changes direction of velocity + Also check end position ends up closer to the target""" + + pos = np.array([50, 0, 0]) + vel = np.array([-1e-5, -1e-5, -1e-5]) + target = np.array( + [55, 5, 5, 0, 0, 0] + ) # Function should attempt to add a positive velocity in any cartesian direction to get to this position + + t1 = 0 + tn = 300 + + [vx1, vy1, vz1], [*pos_new, vxn, vyn, vzn] = herget.find_new_vel_with_universal_kepler( + t1, tn, [*pos, *vel], target + ) + + assert ([vx1, vy1, vz1] > vel).all() + assert ( + np.sqrt(sum((target[:3] - pos_new) ** 2)) < np.sqrt(sum((target[:3] - pos) ** 2)) + ).all() # Check the new end-position is closer to the target + + # Run again to check it continues to converge + [vx1, vy1, vz1], [*pos_new_rerun, vxn, vyn, vzn] = herget.find_new_vel_with_universal_kepler( + t1, tn, [*pos, vx1, vy1, vz1], target + ) + + assert (np.sqrt(sum((target[:3] - pos_new_rerun) ** 2)) < np.sqrt(sum((target[:3] - pos_new) ** 2))).all() + + +def test_find_drho(tmpdir): + """Take an object with a known rho value (from JPL), see if this function approaches the right direction""" + + os.chdir(tmpdir) + data = CSVDataReader( + get_test_filepath("2000OK67_ephem.csv"), "csv", primary_id_column_name="provID" + ).read_rows() + + layup_observatory = LayupObservatory(cache_dir=None) + + # The units of et are seconds (from J2000). This new column is used by + # data_processing_utilities.obscodes_to_barycentric. + et_col = np.array([spice.str2et(row["obsTime"]) for row in data], dtype=" 0 the low-order Stumpff functions are trigonometric.""" + c0, c1, c2, c3, _, _ = stumpff_c(z) + sq = math.sqrt(z) + assert c0 == pytest.approx(math.cos(sq), rel=1e-13) + assert c1 == pytest.approx(math.sin(sq) / sq, rel=1e-13) + assert c2 == pytest.approx((1.0 - math.cos(sq)) / z, rel=1e-13) + assert c3 == pytest.approx((sq - math.sin(sq)) / z**1.5, rel=1e-13) + + +@pytest.mark.parametrize("z", [-0.3, -2.0, -25.0]) +def test_stumpff_closed_form_hyperbolic(z): + c0, c1, c2, c3, _, _ = stumpff_c(z) + sq = math.sqrt(-z) + assert c0 == pytest.approx(math.cosh(sq), rel=1e-13) + assert c1 == pytest.approx(math.sinh(sq) / sq, rel=1e-13) + assert c2 == pytest.approx((math.cosh(sq) - 1.0) / (-z), rel=1e-13) + assert c3 == pytest.approx((math.sinh(sq) - sq) / (-z) ** 1.5, rel=1e-13) + + +@pytest.mark.parametrize("z", [-5.0, -0.2, 0.0, 0.2, 5.0, 40.0]) +def test_stumpff_recurrence(z): + """c_k(z) = 1/k! - z*c_{k+2}(z) ties the high orders to the low ones.""" + c0, c1, c2, c3, c4, c5 = stumpff_c(z) + assert c0 == pytest.approx(1.0 - z * c2, rel=1e-14, abs=1e-16) + assert c1 == pytest.approx(1.0 - z * c3, rel=1e-14, abs=1e-16) + assert c2 == pytest.approx(0.5 - z * c4, rel=1e-14, abs=1e-16) + assert c3 == pytest.approx(1.0 / 6.0 - z * c5, rel=1e-14, abs=1e-16) + + +# -------------------------------------------------------------------------- +# 2. The propagated state +# -------------------------------------------------------------------------- + + +def _propagate_via_elements(gm, dt, state): + """Independent two-body propagation through classical elements. + + Deliberately a different formulation from the universal-variable solver + under test: elements -> Kepler's equation in E -> back to Cartesian. + Elliptic orbits only. + """ + r0v, v0v = state[:3], state[3:] + r0 = np.linalg.norm(r0v) + a = 1.0 / (2.0 / r0 - (v0v @ v0v) / gm) + assert a > 0, "elements reference handles ellipses only" + n = math.sqrt(gm / a**3) + + # Eccentricity vector and the eccentric anomaly at t0. + h = np.cross(r0v, v0v) + evec = np.cross(v0v, h) / gm - r0v / r0 + e = np.linalg.norm(evec) + + cosE0 = (1.0 - r0 / a) / e + sinE0 = (r0v @ v0v) / (e * math.sqrt(gm * a)) + E0 = math.atan2(sinE0, cosE0) + M = E0 - e * math.sin(E0) + n * dt + + # Newton on Kepler's equation. + E = M + for _ in range(200): + dE = (E - e * math.sin(E) - M) / (1.0 - e * math.cos(E)) + E -= dE + if abs(dE) < 1e-15: + break + + # Lagrange f/g in terms of the eccentric-anomaly increment. + dE_ = E - E0 + r = a * (1.0 - e * math.cos(E)) + f = 1.0 - a / r0 * (1.0 - math.cos(dE_)) + g = dt + (math.sin(dE_) - dE_) / n + fdot = -math.sqrt(gm * a) / (r * r0) * math.sin(dE_) + gdot = 1.0 - a / r * (1.0 - math.cos(dE_)) + + out = np.empty(6) + out[:3] = f * r0v + g * v0v + out[3:] = fdot * r0v + gdot * v0v + return out + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.4, -0.05, 0.001, 0.05, 0.3, 0.75, 0.99]) +def test_state_matches_classical_elements(name, frac): + state = STATES[name] + dt = frac * _period(state) + got = universal_step(GM, dt, state).state + want = _propagate_via_elements(GM, dt, state) + assert got == pytest.approx(want, rel=1e-11, abs=1e-13) + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_energy_and_angular_momentum_conserved(name): + state = STATES[name] + P = _period(state) + e0 = 0.5 * state[3:] @ state[3:] - GM / np.linalg.norm(state[:3]) + h0 = np.cross(state[:3], state[3:]) + for frac in (0.13, 0.4, 0.87, 2.6): + s = universal_step(GM, frac * P, state).state + e1 = 0.5 * s[3:] @ s[3:] - GM / np.linalg.norm(s[:3]) + assert e1 == pytest.approx(e0, rel=1e-12) + assert np.cross(s[:3], s[3:]) == pytest.approx(h0, rel=1e-12) + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_round_trip(name): + """Forward then back returns the state and the deviation.""" + state = STATES[name] + dt = 0.37 * _period(state) + var = np.array([0.3, -0.2, 0.11, 1e-3, 2e-3, -5e-4]) + fwd = universal_step(GM, dt, state, variation=var) + back = universal_step(GM, -dt, fwd.state, variation=fwd.variation) + assert back.state == pytest.approx(state, rel=1e-11, abs=1e-13) + assert back.variation == pytest.approx(var, rel=1e-9, abs=1e-13) + + +def test_hyperbolic_propagation(): + state = _hyperbolic_state() + r0 = np.linalg.norm(state[:3]) + alpha = 2.0 * GM / r0 - state[3:] @ state[3:] + assert alpha < 0, "test state should be unbound" + e0 = 0.5 * state[3:] @ state[3:] - GM / r0 + for dt in (-400.0, -20.0, 5.0, 200.0, 3000.0): + out = universal_step(GM, dt, state) + e1 = 0.5 * out.state[3:] @ out.state[3:] - GM / np.linalg.norm(out.state[:3]) + assert e1 == pytest.approx(e0, rel=1e-11) + # and it should come back + dt = 1500.0 + there = universal_step(GM, dt, state).state + back = universal_step(GM, -dt, there).state + assert back == pytest.approx(state, rel=1e-10, abs=1e-13) + + +def test_near_parabolic_branch_is_exercised(): + """|dt/r0| <= 0.2 takes the series guess; make sure it is right there.""" + state = STATES["mainbelt"] + r0 = np.linalg.norm(state[:3]) + dt = 0.1 * r0 # inside the branch + got = universal_step(GM, dt, state).state + want = _propagate_via_elements(GM, dt, state) + assert got == pytest.approx(want, rel=1e-11, abs=1e-13) + + +def test_zero_dt_is_identity(): + state = STATES["tno"] + var = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + out = universal_step(GM, 0.0, state, variation=var) + assert out.state == pytest.approx(state, abs=1e-14) + assert out.variation == pytest.approx(var, abs=1e-14) + + +# -------------------------------------------------------------------------- +# 3. The variational output -- the reason this module exists +# -------------------------------------------------------------------------- + + +def _fd_stm(gm, dt, state, rel=1e-7): + """Central-difference d state(t+dt) / d state(t).""" + stm = np.empty((6, 6)) + scale = np.array([np.linalg.norm(state[:3])] * 3 + [np.linalg.norm(state[3:])] * 3) + for j in range(6): + h = rel * scale[j] + plus, minus = state.copy(), state.copy() + plus[j] += h + minus[j] -= h + stm[:, j] = (universal_step(gm, dt, plus).state - universal_step(gm, dt, minus).state) / (2.0 * h) + return stm + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.3, 0.08, 0.45, 0.9]) +def test_variational_matches_finite_differences(name, frac): + state = STATES[name] + dt = frac * _period(state) + analytic = state_transition_matrix(GM, dt, state) + numeric = _fd_stm(GM, dt, state) + # Compare column-wise against that column's own magnitude: the position + # and velocity blocks differ by many orders of magnitude, so a single + # global tolerance would be meaningless. + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 2e-6 + + +def test_variational_matches_finite_differences_hyperbolic(): + state = _hyperbolic_state() + for dt in (-300.0, 50.0, 900.0): + analytic = state_transition_matrix(GM, dt, state) + numeric = _fd_stm(GM, dt, state) + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 2e-6 + + +@pytest.mark.parametrize("name", sorted(STATES)) +def test_stm_is_symplectic(name): + """M^T J M = J for a Hamiltonian flow. Stronger than det(M) = 1, and it + catches sign or index errors in the variational algebra that a + finite-difference check with loose tolerance could let through. + + The tolerance has to scale as ||M||^2 * eps: the position-vs-velocity + block grows like dt (order 1e5 for a TNO over a period), so forming + M^T J M cancels large numbers down to O(1). Measured worst case is + about 0.6 * ||M||^2 * eps, so 20x that is a real test rather than a + rubber stamp. + """ + state = STATES[name] + J = np.block([[np.zeros((3, 3)), np.eye(3)], [-np.eye(3), np.zeros((3, 3))]]) + for frac in (0.11, 0.5, 1.7): + M = state_transition_matrix(GM, frac * _period(state), state) + floor = 20.0 * np.linalg.norm(M) ** 2 * np.finfo(float).eps + assert M.T @ J @ M == pytest.approx(J, abs=floor) + assert np.linalg.det(M) == pytest.approx(1.0, abs=max(1e-13, floor)) + + +def test_variation_is_linear(): + """The propagated deviation must be linear in the input deviation.""" + state = STATES["tno"] + dt = 0.3 * _period(state) + a = np.array([0.5, -0.25, 0.1, 1e-4, -2e-4, 3e-5]) + b = np.array([-0.1, 0.4, 0.7, 5e-5, 1e-4, -1e-4]) + va = universal_step(GM, dt, state, variation=a).variation + vb = universal_step(GM, dt, state, variation=b).variation + vab = universal_step(GM, dt, state, variation=2.0 * a - 3.0 * b).variation + assert vab == pytest.approx(2.0 * va - 3.0 * vb, rel=1e-10, abs=1e-14) + + +# -------------------------------------------------------------------------- +# Multi-revolution: where the C is wrong +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("n", [0, 1, 2, 5, 40]) +def test_state_is_periodic_across_revolutions(n): + """Adding whole periods must not move the state. + + The C fails this from n=1: its g is too large by n*P*|v|, so |r| grows + without bound while the routine still reports success. + """ + state = STATES["mainbelt"] + P = _period(state) + ref = universal_step(GM, 0.3 * P, state) + got = universal_step(GM, (n + 0.3) * P, state) + assert got.n_rev == n + # Rounding in a single solve at large s accumulates with revolution + # count: measured 2.7e-14 AU at n=1 rising to 6.2e-10 AU at n=40. + # 1e-8 AU (about 1.5 m) is loose enough to be stable and still four + # orders tighter than any real effect. + assert got.state == pytest.approx(ref.state, abs=1e-8) + + +@pytest.mark.parametrize("n", [1, 3]) +def test_variational_still_matches_fd_across_revolutions(n): + """The partials must be right in the multi-revolution regime too. + + This is why the fix solves the full dt rather than simply reducing it: + reducing would give the correct state but the wrong partials, since + neighbouring orbits have different periods and the deviation grows + secularly with every revolution. + """ + state = STATES["mainbelt"] + dt = (n + 0.3) * _period(state) + analytic = state_transition_matrix(GM, dt, state) + # Larger FD step than the single-revolution case: the STM grows with + # revolution count (that is the point of this test), so the differencing + # cancels more and a smaller step is noisier, not better. + numeric = _fd_stm(GM, dt, state, rel=1e-6) + for j in range(6): + scale = max(np.max(np.abs(numeric[:, j])), 1e-12) + assert np.max(np.abs(analytic[:, j] - numeric[:, j])) / scale < 1e-5 + + +def test_partials_grow_secularly_with_revolutions(): + """The along-track deviation grows with revolution count -- the concrete + reason a state-only multi-rev fix would have been wrong.""" + state = STATES["mainbelt"] + P = _period(state) + norms = [np.linalg.norm(state_transition_matrix(GM, (n + 0.3) * P, state)) for n in (0, 1, 2, 5)] + assert norms == sorted(norms), f"STM norm should increase with revolutions: {norms}" + assert norms[-1] > 3.0 * norms[0] + + +# -------------------------------------------------------------------------- +# 4. Faithfulness to the original C +# -------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def c_universal_step(tmp_path_factory): + """Build universal-kepler.c as a shared library and expose universal_step. + + Skipped when there is no compiler, or when the C source has moved. + """ + src = HERE / "universal-kepler.c" + if not src.exists(): + pytest.skip(f"C source not found at {src}") + cc = shutil.which("gcc") or shutil.which("cc") + if cc is None: + pytest.skip("no C compiler available") + + build = tmp_path_factory.mktemp("uk_c") + # The C declares `extern double machine_epsilon` but never uses it; give + # it a definition so the link is clean. + shim = build / "shim.c" + shim.write_text("double machine_epsilon = 0.0;\n") + so = build / "libuk.so" + proc = subprocess.run( + [cc, "-std=gnu99", "-w", "-O2", "-fPIC", "-shared", str(src), str(shim), "-lm", "-o", str(so)], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + pytest.skip(f"could not build the C reference: {proc.stderr[-400:]}") + + lib = ctypes.CDLL(str(so)) + arr6 = ctypes.c_double * 6 + lib.universal_step.restype = ctypes.c_int + lib.universal_step.argtypes = [ + ctypes.c_double, + ctypes.c_double, + ctypes.POINTER(arr6), + ctypes.POINTER(arr6), + ctypes.POINTER(arr6), + ] + + def call(gm, dt, state, variation): + s0 = arr6(*[float(x) for x in state]) + out = arr6() + var = arr6(*[float(x) for x in variation]) + flag = lib.universal_step( + ctypes.c_double(gm), + ctypes.c_double(dt), + ctypes.byref(s0), + ctypes.byref(out), + ctypes.byref(var), + ) + return flag, np.array(out[:]), np.array(var[:]) + + return call + + +@pytest.mark.parametrize("name", sorted(STATES)) +@pytest.mark.parametrize("frac", [-0.4, 0.02, 0.25, 0.6, 0.95]) +def test_matches_c_within_one_period(c_universal_step, name, frac): + """Inside one revolution -- where the C is correct -- the port must + reproduce it, state and partials alike.""" + state = STATES[name] + dt = frac * _period(state) + var = np.array([0.3, -0.2, 0.11, 1e-3, 2e-3, -5e-4]) + + flag, c_state, c_var = c_universal_step(GM, dt, state, var) + if flag != 0: + # The C's absolute tolerance is unreachable for some orbits; that is + # pinned by test_c_tolerance_fails_on_distant_orbits below, and there + # is nothing to compare against here. + pytest.skip(f"C did not converge for {name} at {frac}P (its own defect)") + + py = universal_step(GM, dt, state, variation=var) + assert py.state == pytest.approx(c_state, rel=1e-11, abs=1e-13) + assert py.variation == pytest.approx(c_var, rel=1e-9, abs=1e-13) + + +def test_matches_c_hyperbolic(c_universal_step): + state = _hyperbolic_state() + var = np.array([0.2, 0.1, -0.3, 1e-3, -1e-3, 2e-4]) + for dt in (-250.0, 40.0, 800.0): + flag, c_state, c_var = c_universal_step(GM, dt, state, var) + assert flag == 0 + py = universal_step(GM, dt, state, variation=var) + assert py.state == pytest.approx(c_state, rel=1e-11, abs=1e-13) + assert py.variation == pytest.approx(c_var, rel=1e-9, abs=1e-13) + + +def test_c_tolerance_fails_on_distant_orbits(c_universal_step): + """Pins the C's scale-dependent convergence test. + + EPS = 1e-13 is applied as an absolute tolerance on the universal + variable s, which spans P/a over one revolution -- 577 for a mainbelt + orbit but 2395 for a TNO at 43 AU. The demand is then below double + precision, so convergence becomes a coin flip that the C loses more + often the more distant the object. The port's relative test converges + everywhere in this sweep. + """ + fracs = (-0.4, 0.02, 0.25, 0.5, 0.6, 0.75, 0.95) + c_failures, distant_failures = [], [] + for name in sorted(STATES): + state = STATES[name] + P = _period(state) + for frac in fracs: + dt = frac * P + flag, _, _ = c_universal_step(GM, dt, state, np.zeros(6)) + if flag != 0: + c_failures.append((name, frac)) + if name in ("tno", "eccentric"): + distant_failures.append((name, frac)) + # the port must converge regardless + universal_step(GM, dt, state) + + assert c_failures, "expected the C to fail somewhere in this sweep" + assert distant_failures, f"expected distant-orbit failures, got {c_failures}" + + +def test_c_is_wrong_past_one_period_and_the_port_is_not(c_universal_step): + """Pins the defect this port fixes, so nobody 'restores' the C behaviour. + + If this ever fails because the C now agrees, the C was fixed upstream and + this test should become an equality check. + """ + state = STATES["mainbelt"] + P = _period(state) + var = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + + ref = universal_step(GM, 0.3 * P, state).state + dt = 5.3 * P + flag, c_state, _ = c_universal_step(GM, dt, state, var) + + assert flag == 0, "the C reports success -- that is what makes it dangerous" + c_err = np.linalg.norm(c_state[:3] - ref[:3]) + assert c_err > 10.0, f"expected the C to be far off, got {c_err} AU" + + py = universal_step(GM, dt, state).state + assert py == pytest.approx(ref, rel=1e-9, abs=1e-11) + + +# -------------------------------------------------------------------------- +# Input handling +# -------------------------------------------------------------------------- + + +def test_rejects_bad_shapes(): + with pytest.raises(ValueError): + universal_step(GM, 1.0, np.zeros(5)) + with pytest.raises(ValueError): + universal_step(GM, 1.0, STATES["tno"], variation=np.zeros(3)) + + +def test_rejects_zero_position(): + with pytest.raises(ValueError): + universal_step(GM, 1.0, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0])) From c19134a0cc1422467b14076e49945da6d7665563 Mon Sep 17 00:00:00 2001 From: awilson110 Date: Fri, 28 Aug 2026 10:46:29 +0100 Subject: [PATCH 2/4] implement herget as iod --- src/layup/iod.py | 14 +- src/layup/orbitfit.py | 1534 ++++++++++++++++++++++++++++++++++++----- 2 files changed, 1375 insertions(+), 173 deletions(-) diff --git a/src/layup/iod.py b/src/layup/iod.py index 970ef505..806eb110 100644 --- a/src/layup/iod.py +++ b/src/layup/iod.py @@ -36,7 +36,10 @@ from typing import Callable, Optional, Sequence from layup.constants import GMtotal, SPEED_OF_LIGHT -from layup.routines import FitResult, Observation, gauss +from layup.routines import FitResult, Observation, gauss, get_ephem +from layup.utilities.herget_iod import herget_with_assist + +from layup.orbit_maths import build_ephem_and_mus logger = logging.getLogger(__name__) @@ -109,6 +112,15 @@ def gauss_iod(observations, seq): register_iod("gauss", gauss_iod) +def herget_iod(observations, seq): + '''''' + ephem, _, _ = build_ephem_and_mus() + solns = herget_with_assist(observations, seq, ephem, tolerance=0.0001, max_iterations=100) + return solns + + +register_iod("herget", herget_iod) + # ----------------------------------------------------------------------- # # Candidate filter (held-out angular residual). # diff --git a/src/layup/orbitfit.py b/src/layup/orbitfit.py index d482d6db..10c405c3 100644 --- a/src/layup/orbitfit.py +++ b/src/layup/orbitfit.py @@ -1,13 +1,14 @@ +import hashlib import logging import os +import re from argparse import Namespace +from dataclasses import dataclass from pathlib import Path from typing import Literal, Optional import numpy as np -import pooch import spiceypy as spice -from time import sleep from numpy.lib import recfunctions as rfn @@ -16,11 +17,27 @@ Observation, gauss, get_ephem, + run_bk_iod, + run_bk_native_fit, run_from_vector_with_initial_guess, + run_sequential_update, ) + +try: + from layup.routines import ( + get_ias15_adaptive_mode, + set_ias15_adaptive_mode, + ) +except ImportError: # extension not rebuilt yet + get_ias15_adaptive_mode = lambda: -1 + set_ias15_adaptive_mode = lambda m: None +# _MU_SUN (= heliocentric GM = k^2) is used by the BK-native fit for the +# bound-orbit energy prior on gdot; SPEED_OF_LIGHT (au/day) by the radar ingest. +from layup.constants import MU_SUN as _MU_SUN, SPEED_OF_LIGHT from layup.convert import convert +from layup.iod import filter_candidates_by_residual, get_iod, iod_methods -from layup.utilities.astrometric_uncertainty import data_weight_Veres2017 +from layup.utilities.astrometric_uncertainty import astrometric_uncertainty_Veres2017 from layup.utilities.data_processing_utilities import ( LayupObservatory, create_chunks, @@ -28,21 +45,54 @@ get_format, parse_fit_result, process_data_by_id, + resolve_num_workers, ) from layup.utilities.datetime_conversions import convert_tdb_date_to_julian_date from layup.utilities.debiasing import debias, generate_bias_dict -from layup.utilities.file_io import CSVDataReader, HDF5DataReader, Obs80DataReader -from layup.utilities.file_io.file_output import write_csv, write_hdf5 -from layup.utilities.herget_iod import herget_with_assist +from layup.utilities.file_io import ( + ADESXMLDataReader, + CSVDataReader, + HDF5DataReader, + Obs80DataReader, +) +from layup.utilities.file_io.file_output import append_hdf5, write_csv, write_hdf5 +from layup.utilities.cache_location import default_cache_dir logger = logging.getLogger(__name__) +# Observed sky-motion rates (ADES `raRate`/`decRate`) arrive in arcsec/hour and +# follow the great-circle convention: `raRate` is cos(Dec)*dRA/dt, NOT the bare +# coordinate rate dRA/dt. This matches Sorcha's `RARateCosDec` output (which +# projects drho_hat/dt onto A = (-sinRA, cosRA, 0)) and layup's own residual, +# which projects onto the same tangent vector `a_vec` -- so an observed rate is +# compared directly to omega.a_vec with no extra cos(Dec) factor. Internally the +# fitter works in radians and AU/day, so omega = d(rho_hat)/dt is in rad/day; +# convert the observed rates from arcsec/hour to rad/day at ingest. +ARCSEC_PER_HOUR_TO_RAD_PER_DAY = (np.pi / 180.0 / 3600.0) * 24.0 + +# Radar (delay/Doppler) observables arrive in JPL units: round-trip delay in +# microseconds and Doppler shift in Hz at a per-observation transmit frequency +# `freqTx` (Hz). The fitter models round-trip delay in days and round-trip +# range-rate in au/day (see RadarObservation in detection.cpp), so convert at +# ingest, mirroring the streak rate convention above: +# delay[days] = delay[us] * 1e-6 / 86400 +# doppler[au/day] = -c * doppler[Hz] / freqTx[Hz] +# The Doppler sign follows F = -(f_tx/c) * d(round-trip range)/dt, so a positive +# (receding) range-rate produces a negative frequency shift. +US_TO_DAYS = 1.0e-6 / 86400.0 +# Fallback 1-sigma weights when the JPL uncertainty columns are absent: ~1 us of +# round-trip delay and ~1 Hz of Doppler (converted per observation via freqTx). +_DEFAULT_DELAY_UNC_DAYS = US_TO_DAYS +_DEFAULT_DOPPLER_UNC_HZ = 1.0 + # The list of required input column names for the provided observations to be fit. # Note: This should not include the primary id column name. REQUIRED_INPUT_OBSERVATIONS_COLUMN_NAMES = [ ( set(["ra", "dec"]), # Either `ra` and `dec` must be in the file set(["raRate", "decRate"]), # Or `raRate` and `decRate` must be in the file + set(["delay"]), # Or a radar round-trip delay (us) + set(["doppler"]), # Or a radar Doppler shift (Hz; needs `freqTx`) ), "obsTime", "stn", @@ -59,17 +109,168 @@ "MPC80col": (Obs80DataReader, None), "ADES_csv": (CSVDataReader, "csv"), "ADES_psv": (CSVDataReader, "psv"), - "ADES_xml": (None, None), + "ADES_xml": (ADESXMLDataReader, None), "ADES_hdf5": (HDF5DataReader, None), } -GMtotal = 0.0002963092748799319 -AU_M = 149597870700 -SPEED_OF_LIGHT = 2.99792458e8 * 86400.0 / AU_M +def _run_fit(assist_ephem, initial_guess, observations, engine, iter_max=100): + """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. + + `iter_max` is the LM iteration budget used by the multi-root picker's + two-tier (cheap-screen then full) passes. The Cartesian engine honors + it; the BK-native engine uses its own internal cap (it takes `mu` for + the bound-orbit energy prior rather than an iteration budget), so + `iter_max` is ignored on that path. + """ + if engine == "cartesian": + return run_from_vector_with_initial_guess(assist_ephem, initial_guess, observations, iter_max) + 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'.") + + +# Non-gravitational Marsden parameters and their FitResult/C++ bitmask bits. +_NONGRAV_BITS = {"A1": 1, "A2": 2, "A3": 4} + + +def _parse_nongrav(fit_nongrav): + """Normalize the ``fit_nongrav`` argument to ``(mask, names)``. + + Accepts ``False``/``None`` (no non-grav fit), ``True`` (== ``["A2"]``, the + common asteroid Yarkovsky case), a string naming params (e.g. ``"A2"``, + ``"A1A2A3"``, ``"A1,A3"``), or an iterable of names. Returns the C++ bitmask + (bits 1/2/4 for A1/A2/A3) and the selected names ordered A1, A2, A3. + """ + if not fit_nongrav: + return 0, [] + if fit_nongrav is True: + sel = {"A2"} + elif isinstance(fit_nongrav, str): + sel = set(re.findall(r"A[123]", fit_nongrav.upper())) + else: + sel = {str(s).upper() for s in fit_nongrav} + names = [n for n in ("A1", "A2", "A3") if n in sel] + if not names: + raise ValueError(f"fit_nongrav={fit_nongrav!r}: expected some of 'A1', 'A2', 'A3'.") + return sum(_NONGRAV_BITS[n] for n in names), names + + +# fit_nongrav="auto" model ladder (issue #357): non-grav models are tried in +# increasing complexity, and the first that converges, is well-conditioned +# (flag 0), and is statistically warranted (see NongravAutoThresholds) is adopted; +# otherwise the gravity-only fit is kept. +_AUTO_NONGRAV_LADDER = (("A2",), ("A1", "A2"), ("A1", "A2", "A3")) + + +@dataclass(frozen=True) +class NongravAutoThresholds: + """Decision thresholds for adaptive non-grav selection (``fit_nongrav="auto"``, + issue #357). Pass a customized instance to ``orbitfit`` to tune how readily a + non-gravitational model is adopted; the defaults reproduce the standard + "introduce a non-grav only when gravity is unacceptable and the parameter is + statistically warranted" behavior. + + Parameters + ---------- + accept_reduced_chi2 : float + A gravity-only fit whose reduced chi-square is at or below this is kept + as-is; non-grav models are tried only above it. Default 1.5. + delta_chi2_per_param : float + Minimum chi-square drop required per added non-grav parameter to adopt a + model (9.0 ~ 3-sigma). Default 9.0. + nsigma : float + Each added non-grav parameter must exceed this many times its 1-sigma + uncertainty to be adopted. Default 3.0. + """ + + accept_reduced_chi2: float = 1.5 + delta_chi2_per_param: float = 9.0 + nsigma: float = 3.0 + + +_AUTO_DEFAULT_THRESHOLDS = NongravAutoThresholds() -def _get_result_dtypes(primary_id_column_name: str): - """Helper function to create the result dtype with the correct primary ID column name.""" + +def _gravity_fit_acceptable(csq, ndof, thresholds=_AUTO_DEFAULT_THRESHOLDS): + """Whether a gravity-only fit is good enough that no non-grav is warranted. + + True when the reduced chi-square is at or below the acceptance threshold (or + there are no degrees of freedom to judge it by). + """ + return ndof <= 0 or csq / ndof <= thresholds.accept_reduced_chi2 + + +def _nongrav_warranted(csq_gravity, res_ng, names, thresholds=_AUTO_DEFAULT_THRESHOLDS): + """Whether adopting the (converged, well-conditioned) non-grav fit ``res_ng`` + for parameters ``names`` over the gravity-only fit is statistically warranted: + a significant chi-square drop AND every added parameter individually significant. + """ + if (csq_gravity - res_ng.csq) <= thresholds.delta_chi2_per_param * len(names): + return False # chi-square improvement not significant for the added parameter(s) + return all( + abs(getattr(res_ng, n.lower())) > thresholds.nsigma * getattr(res_ng, n.lower() + "_unc") + for n in names + ) + + +def _gofr_arg(nongrav_gr): + """Normalize a non-grav g(r) argument to the ``[alpha, nm, nn, nk, r0]`` list the + C++ fit expects (empty -> the default inverse-square law).""" + if nongrav_gr is None: + return [] + gr = list(nongrav_gr) + if len(gr) != 5: + raise ValueError("nongrav_gr must be [alpha, nm, nn, nk, r0] (5 values) or None") + return [float(v) for v in gr] + + +def _select_nongrav_auto( + assist_ephem, res_grav, observations, thresholds=_AUTO_DEFAULT_THRESHOLDS, gofr=None +): + """Adaptive non-grav selection for ``fit_nongrav="auto"`` (issue #357). + + ``res_grav`` is the converged gravity-only fit. Returns the most parsimonious + acceptable model: the gravity-only result unless a non-grav model is both + well-conditioned (flag 0) and statistically warranted per ``thresholds``. + ``gofr`` selects the g(r) sublimation law (see ``orbitfit``'s ``nongrav_gr``). + """ + if _gravity_fit_acceptable(res_grav.csq, res_grav.ndof, thresholds): + return res_grav # gravity-only fit is acceptable; no non-gravs needed + gr = _gofr_arg(gofr) + for names in _AUTO_NONGRAV_LADDER: + mask = sum(_NONGRAV_BITS[n] for n in names) + res_ng = run_from_vector_with_initial_guess( + assist_ephem, res_grav, observations, nongrav_mask=mask, gofr=gr + ) + if res_ng.flag == 0 and _nongrav_warranted(res_grav.csq, res_ng, names, thresholds): + return res_ng # parsimonious, well-determined non-grav model + return res_grav # no non-grav model is warranted + + +def _get_result_dtypes(primary_id_column_name: str, nongrav_names=(), per_arc=False): + """Helper function to create the result dtype with the correct primary ID column name. + + For each fitted non-gravitational parameter in ``nongrav_names`` (a subset of + ``A1``/``A2``/``A3``), two columns are appended -- e.g. ``a2`` and ``a2_unc`` + (the value and its 1-sigma uncertainty, au/day^2). With no non-grav params the + default 6-parameter output schema is unchanged. + + When ``per_arc`` is set (the two-apparition comet-linkage fit), the non-grav + columns above hold the *earlier* arc's amplitudes and a second block + ``a2_arc2``/``a2_arc2_unc`` is appended for the *later* arc. Both are opt-in, + so the ordinary fit schema is unaffected. + """ + per_arc_cols = [] + if per_arc: + per_arc_cols = [ + (col, "f8") for n in nongrav_names for col in (n.lower() + "_arc2", n.lower() + "_arc2_unc") + ] # Define a structured dtype to match the OrbfitResult fields return np.dtype( [ @@ -89,6 +290,166 @@ def _get_result_dtypes(primary_id_column_name: str): ("FORMAT", "O"), # Orbit format ] + [(col_name, "f8") for col_name in get_cov_columns()] # Flat covariance matrix (36 elements) + + [ # non-grav params (issue #351), value + 1-sigma per fitted param + (col, "f8") for n in nongrav_names for col in (n.lower(), n.lower() + "_unc") + ] + + per_arc_cols # later-arc non-grav amplitudes (comet linkage), when per_arc + # Observation provenance / incremental fingerprint (issue #419): a + # deterministic hash of the observation set this orbit was fit from, plus + # the number of fittable observations. Kept last so the positional output + # tuples above the conditional non-grav columns are unaffected. Lets a + # steady-state pipeline skip re-fitting objects whose observations are + # unchanged since the prior catalog (see ``_obs_fingerprint``). + + [("obs_hash", "O"), ("nobs_fit", "i4")] + ) + + +# Observation columns that determine the fit and therefore the fingerprint. Any +# change to these (new obs, corrected astrometry, changed uncertainties) yields a +# different hash and forces a re-fit; changes to purely cosmetic columns (e.g. a +# magnitude) do not. Only the columns actually present in the data are hashed. +_FINGERPRINT_COLUMNS = ( + "obsTime", # epoch + "ra", + "dec", # optical astrometry (raw, pre-debias) + "stn", # observatory code + "raStar", + "decStar", # occultation astrometry + "rmsRA", + "rmsDec", # reported astrometric uncertainties + "astCat", # star catalog (feeds debiasing + weighting) + "raRate", + "decRate", + "rmsRArate", + "rmsDecrate", # streak rates + "delay", + "doppler", + "freqTx", # radar observables +) + + +def _fmt_fingerprint_value(v): + """Canonical, round-trip-stable string for one observation field value. + + Uses Python's shortest round-trip ``repr`` for floats (deterministic across + runs and platforms) and decodes bytes so a numpy ``S``/``O`` string column + hashes identically however it was loaded. + """ + if isinstance(v, bytes): + return v.decode("utf-8", "replace") + if isinstance(v, (np.floating, float)): + return repr(float(v)) + return str(v) + + +def _obs_fingerprint(data, column_names): + """Deterministic fingerprint of an object's fittable observation set. + + Returns ``(nobs, hash16)`` where ``nobs`` is the row count and ``hash16`` is a + 16-hex-character SHA-1 digest over the fit-relevant columns (``_FINGERPRINT_COLUMNS``). + The per-row strings are sorted before hashing, so the fingerprint is + order-independent: the same physical observations produce the same hash + regardless of row order (an LM fit over a set is itself order-independent). + Computed on the raw observations before any in-place debiasing so it reflects + what was reported, not the current bias model. + """ + cols = [c for c in _FINGERPRINT_COLUMNS if c in column_names] + rows = ["\x1f".join(_fmt_fingerprint_value(d[c]) for c in cols) for d in data] + rows.sort() + payload = "\x1e".join(rows).encode("utf-8") + return len(rows), hashlib.sha1(payload).hexdigest()[:16] + + +def _is_radar(d, column_names): + """True if a row carries a populated radar observable (delay or Doppler). + + Radar rows have no ra/dec; they are dispatched to ``Observation.from_radar`` + rather than the astrometry/streak factories. + """ + has_delay = "delay" in column_names and not np.isnan(d["delay"]) + has_doppler = "doppler" in column_names and not np.isnan(d["doppler"]) + return has_delay or has_doppler + + +def _radar_observation(objID, d, epoch_jd, column_names): + """Build a radar ``Observation`` from a row, converting JPL units to the + fitter's internal units. + + delay (us, round-trip) -> days; Doppler (Hz) -> round-trip range-rate + (au/day) via the per-observation transmit frequency ``freqTx``. The + barycentric observer position/velocity columns (x,y,z,vx,vy,vz) must already + be present (added by ``orbitfit``); the observer acceleration columns + (ax,ay,az) drive the two-leg light-time model and default to zero when + absent. 1-sigma uncertainties come from ``rmsDelay``/``rmsDoppler`` when + present, else the module defaults. + """ + has_delay = "delay" in column_names and not np.isnan(d["delay"]) + has_doppler = "doppler" in column_names and not np.isnan(d["doppler"]) + + f_tx = d["freqTx"] if "freqTx" in column_names else np.nan + if has_doppler and (np.isnan(f_tx) or f_tx == 0.0): + raise ValueError(f"Radar Doppler observation for {objID} requires a nonzero 'freqTx' (Hz).") + + delay_days = (d["delay"] * US_TO_DAYS) if has_delay else 0.0 + doppler_audy = (-SPEED_OF_LIGHT * d["doppler"] / f_tx) if has_doppler else 0.0 + + if "rmsDelay" in column_names and not np.isnan(d["rmsDelay"]): + delay_unc = abs(d["rmsDelay"]) * US_TO_DAYS + else: + delay_unc = _DEFAULT_DELAY_UNC_DAYS + + if has_doppler: + rms_hz = ( + d["rmsDoppler"] + if ("rmsDoppler" in column_names and not np.isnan(d["rmsDoppler"])) + else _DEFAULT_DOPPLER_UNC_HZ + ) + doppler_unc = abs(SPEED_OF_LIGHT * rms_hz / f_tx) + else: + doppler_unc = abs(SPEED_OF_LIGHT * _DEFAULT_DOPPLER_UNC_HZ / f_tx) if not np.isnan(f_tx) else 1.0 + + if all(c in column_names for c in ("ax", "ay", "az")): + observer_acc = [float(d["ax"]), float(d["ay"]), float(d["az"])] + else: + observer_acc = [0.0, 0.0, 0.0] + + return Observation.from_radar_with_id( + str(objID), + delay_days, + doppler_audy, + has_delay, + has_doppler, + epoch_jd, + [d["x"], d["y"], d["z"]], # Barycentric observer position + [d["vx"], d["vy"], d["vz"]], # Barycentric observer velocity + delay_unc, + doppler_unc, + observer_acc, # Barycentric observer acceleration (au/day^2) + ) + + +def _append_observer_acceleration(data, observatory, dt_sec=2.0): + """Append barycentric observer acceleration columns (ax, ay, az; au/day^2). + + Finite-differences the barycentric station velocity from + ``obscodes_to_barycentric`` at the observation epoch +/- ``dt_sec``. Used only + by radar observations, whose two-leg light-time model extrapolates the station + state back to the signal transmit time. Requires the ``et`` column (seconds + past J2000, TDB) that ``orbitfit`` adds before computing the observer states. + """ + + def _vel(et_offset_sec): + shifted = data.copy() + shifted["et"] = data["et"] + et_offset_sec + pv = np.atleast_1d(observatory.obscodes_to_barycentric(shifted)) + return np.stack([pv["vx"], pv["vy"], pv["vz"]], axis=-1) + + # d(velocity[au/day]) / d(time[day]); 2*dt_sec seconds = 2*dt_sec/86400 days. + scale = 86400.0 / (2.0 * dt_sec) + acc = (_vel(dt_sec) - _vel(-dt_sec)) * scale + acc = np.atleast_2d(acc) + return rfn.append_fields( + data, ["ax", "ay", "az"], [acc[:, 0], acc[:, 1], acc[:, 2]], usemask=False, asrecarray=True ) @@ -321,143 +682,320 @@ def create_empty_result(id, dtypes): "NONE", # format ) + (np.nan,) * 36 # Flat covariance matrix + # non-grav columns (issue #351): NaN per a1/a2/a3 (+ _unc) that is present + + tuple(np.nan for n in ("a1", "a2", "a3") if n in dtypes.names for _ in (0, 1)) + # obs fingerprint (issue #419): empty hash never matches, so a failed + # or invalid fit is always retried next cycle rather than skipped. + + (("", 0) if "obs_hash" in dtypes.names else ()) ], dtype=dtypes, ) -def do_gauss_iod(observations, seq): - """Calculate an initial orbit estimate using Gauss's method. - - Parameters - ---------- - observations : list[Observation] - The list of Observations used for the orbit estimate - seq : list[list[int] - The list of lists of indexes of observations that are closely spaced in time. +def _carry_forward_result(prior_row, dtypes): + """Copy a prior fit result forward under the current output dtype (issue #419). - Returns - ------- - list[FitResult] - A collection of orbit fit results that can be used to perform a higher - quality fit estimate. + Used by the skip-unchanged path: when an object's observations are unchanged, + its stored fit is emitted verbatim rather than recomputed. Fields shared with + ``dtypes`` are copied by name (so the carried row is identical to the prior + fit); any field present in ``dtypes`` but absent from the prior (e.g. a + non-grav column the prior run did not fit) is left at its type default. """ - # Get gauss solution, using the first, middle, and last observation - # of the primary sequence - idx0, idx1, idx2 = seq[0][0], seq[0][int(len(seq[0]) / 2)], seq[0][-1] - logger.debug(f"Sequence indexs passed to gauss: {idx0}, {idx1}, {idx2}") - solns = gauss(GMtotal, observations[idx0], observations[idx1], observations[idx2], 0.0001, SPEED_OF_LIGHT) - - return solns - - -def do_herget_iod(observations, seq, args, aux): - """Calculate an initial orbit estimate using Herget's method. - - Parameters - ---------- - observations : list[Observation] - The list of Observations used for the orbit estimate - seq : list[list[int] - The list of lists of indexes of observations that are closely spaced in time. - - Returns - ------- - list[FitResult] - A collection of orbit fit results that can be used to perform a higher - quality fit estimate. + prior_row = prior_row[0] if getattr(prior_row, "shape", None) == (1,) else prior_row + out = np.zeros(1, dtype=dtypes) + for name in dtypes.names: + if name in prior_row.dtype.names: + out[name][0] = prior_row[name] + return out + + +def _partition_unchanged(data, initial_guess, primary_id_column_name, fit_nongrav): + """Split raw observations into (to_fit, carried_forward) by obs fingerprint. + + The steady-state pre-filter for ``orbitfit(skip_unchanged=True)`` (issue #419): + an object whose fingerprint matches a converged prior fit over the same + observation set is carried forward verbatim (cast to the current output + dtype); every other object's rows are returned for fitting. Runs before any + ephemeris/observatory setup, so a skipped object costs only a fingerprint + hash. The fingerprint is computed on the raw (pre-debias) observations, so it + matches the one ``_orbitfit`` stores. """ - # Get Herget solution, using first and last point - # of the primary sequence - solns = herget_with_assist(observations, seq, 0.001, args=args, aux=aux) - print(solns[0].niter) - return solns - - -def do_fit(observations, seq, cache_dir, iod="gauss", args=None, aux=None): - """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. + _, nongrav_names = _parse_nongrav(fit_nongrav) + out_dtype = _get_result_dtypes(primary_id_column_name, nongrav_names) + prior = np.atleast_1d(initial_guess) + prior_by_id = {row[primary_id_column_name]: row for row in prior} + ids = data[primary_id_column_name] + keep = np.ones(len(data), dtype=bool) + carried = [] + for oid in np.unique(ids): + obj_mask = ids == oid + nobs, obs_hash = _obs_fingerprint(data[obj_mask], data.dtype.names) + p = prior_by_id.get(oid) + if ( + p is not None + and int(p["flag"]) == 0 + and "nobs_fit" in prior.dtype.names + and int(p["nobs_fit"]) == nobs + and str(p["obs_hash"]) == obs_hash + ): + keep[obj_mask] = False + carried.append(_carry_forward_result(p, out_dtype)) + carried_arr = np.concatenate(carried) if carried else np.array([], dtype=out_dtype) + return data[keep], carried_arr - A Gauss preliminary order is fit for the 0-th segment, - using the first, middle, and last observations in that - segment. - Then an orbit fit is done on the 0-th segment, using the - initial orbit from Gauss. If that fails, any other preliminary - solutions are tried. - - Next, a fit to the full set of observations is attempted, given - the fit to the primary segment as an initial guess. If that - succeeds, the solution is returned. +def do_gauss_iod(observations, seq): + """Backward-compat wrapper for the Gauss IOD. - Otherwise, adjacent segments of observations are added and - the fit is updated, iteratively. + Prefer ``layup.iod.get_iod("gauss")`` for new code; this shim + exists so callers that imported ``do_gauss_iod`` directly continue + to work. + """ + return get_iod("gauss")(observations, seq) + + +# Multi-root picker tuning. Both knobs are exposed to do_fit() callers +# in case downstream code wants to override them, but the defaults are +# what worked best on the diagnostic/scan and neo_scan datasets. +_PICKER_MIN_R_HELIO_AU = 0.3 # reject roots with r < this as unphysical +_PICKER_SCREEN_ITER_MAX = 80 # cheap LM budget for the first pass +_PICKER_FULL_ITER_MAX = 100 # full LM budget for the fallback pass + +_PREFILTER_THRESHOLD_SIGMA = 1000.0 # held-out residual filter cutoff + +# IAS15 adaptive-step controller used during the multi-root picker. +# With the legacy controller (mode 1), LM grinds for minutes on phantom +# Gauss roots whose trajectories pass close to Earth (the integrator +# chases ever-smaller steps to resolve the close encounter — 100-1000× +# wallclock blowup observed on diagnostic/scan). The newer (Pham, Rein +# & Spiegel 2024) controller, mode 2, steps through those encounters +# efficiently: it brings the same pathological cases from >120 s to +# sub-second with the identical recovered orbit, and unlike a step-size +# floor it is a better controller rather than a truncation, so it costs +# no accuracy on genuine close-Earth encounters. Set to -1 to leave +# ASSIST's default (legacy mode 1). +_PICKER_IAS15_ADAPTIVE_MODE = 2 + + +# Cache of the Python-side assist.Ephem handle. The C-side get_ephem() +# from layup.routines returns the C struct; the Python residual filter +# needs the rebound/assist Python wrapper instead, so we cache one per +# cache_dir. +_assist_python_ephem_cache: dict = {} + + +def _get_python_ephem(cache_dir): + """Lazy-load and cache the Python-side assist.Ephem for the filter.""" + key = str(cache_dir) + if key in _assist_python_ephem_cache: + return _assist_python_ephem_cache[key] + try: + import assist + except ImportError: + return None + try: + eph = assist.Ephem(os.path.join(key, "linux_p1550p2650.440"), os.path.join(key, "sb441-n16.bsp")) + except Exception as e: + logger.warning(f"assist.Ephem load failed for {cache_dir}: {e}") + return None + _assist_python_ephem_cache[key] = eph + return eph + + +def _pick_best_root(candidates, min_r_au): + """Pick the best converged candidate from a list of LM results. + + "Best" means smallest χ² among candidates that + + 1. report ``flag == 0`` (LM converged), and + 2. have heliocentric distance > ``min_r_au`` (physical orbit). + + Returns None if no candidate satisfies (1); in that case the caller + typically retries at a larger LM budget. If (1) is met but (2) + isn't, the smallest-χ² convergent root is still returned (better + than nothing). + """ + converged = [c for c in candidates if c.flag == 0] + if not converged: + return None + sane = [c for c in converged if (c.state[0] ** 2 + c.state[1] ** 2 + c.state[2] ** 2) > min_r_au**2] + pool = sane if sane else converged + return min(pool, key=lambda c: c.csq) + + +def do_fit( + observations, + seq, + cache_dir, + iod="auto", + engine="cartesian", + screen_iter_max: int = _PICKER_SCREEN_ITER_MAX, + full_iter_max: int = _PICKER_FULL_ITER_MAX, + min_r_helio_AU: float = _PICKER_MIN_R_HELIO_AU, + prefilter_threshold_sigma: float = _PREFILTER_THRESHOLD_SIGMA, + picker_ias15_adaptive_mode: int = _PICKER_IAS15_ADAPTIVE_MODE, +): + """Carry out an orbit fit to a list of observations. + + Pipeline: + 1. IOD: produce one or more candidate seed orbits via the + registered method named by `iod` (default: "auto"). The + registry lives in `layup.iod`; register new methods with + `iod.register_iod(name, callable)`. + 2. Multi-root picker: run LM from every IOD candidate on the + primary segment (`seq[0]`) at a cheap `screen_iter_max` + budget. Pick the smallest-χ² converged candidate with + heliocentric distance above `min_r_helio_AU`. If nothing + converges at the cheap budget, retry at `full_iter_max`. + Then refit on the full observation set. Parameters ---------- observations : list - A time-ordered list of observations + Time-ordered list of layup Observations. seq : list of lists - A list of lists of observation indices. + Per-segment index lists; seq[0] is the primary segment. + cache_dir : str + Directory holding the ASSIST kernels. iod : str - The IOD used to generate an initial guess orbit. Currently supports ['gauss']. - Default is 'gauss'. + Name of the registered IOD method "auto" (default) or "gauss". "auto" runs + Gauss and falls back to the BK 5-parameter linear IOD (run_bk_iod) on the + primary segment when every Gauss root fails to seed a converged fit. + 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. + screen_iter_max, full_iter_max : int + Two-tier LM iteration caps for the multi-root picker. + min_r_helio_AU : float + Lower bound on heliocentric distance for accepted IOD roots. Returns ------- FitResult - The result of the orbit fit. + Best converged fit (flag == 0) when one exists, else a + best-effort or sentinel FitResult with a non-zero flag. """ - if iod.lower() == "gauss": - solns = do_gauss_iod(observations, seq) - elif iod.lower() == "herget": - solns = do_herget_iod(observations, seq, args, aux) - else: - 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: - logger.debug(f"The iod {iod} failed") + # 'auto' is a strategy, not a registered IOD: seed candidates with Gauss, + # then (after the picker, below) fall back to the BK 5-parameter linear IOD + # if every Gauss root fails to converge. Any other value is a registry name. + is_auto = isinstance(iod, str) and iod.lower() == "auto" + try: + iod_func = get_iod("gauss") if is_auto else (get_iod(iod) if isinstance(iod, str) else iod) + except ValueError as e: + raise ValueError(f"{e} Use iod.register_iod to add a new method.") + # Normalize to a list: an IOD may legitimately return None (e.g. Gauss + # finding no real roots), and the prefilter/picker below index and len() it. + solns = list(iod_func(observations, seq) or []) + + # If the iod produced no candidates, surface a sentinel -- unless we're in + # 'auto' mode, where the BK-IOD fallback below still has a shot. + if not solns and not is_auto: + logger.debug(f"IOD {iod!r} returned no candidates") x = FitResult() x.flag = 5 return x + # Pre-filter the IOD candidates by predicted-vs-observed residual + # on every observation. The right Gauss root predicts the full + # observation set within a few σ; phantom roots typically miss by + # 10⁵+ σ. Throwing those out before any LM iteration runs cuts + # the picker loop down to 1-2 LM fits per case in the common + # case (vs up to 8 brute-force LMs). Loose threshold (default + # 1000σ) so the right root is never rejected. + py_ephem = _get_python_ephem(cache_dir) + if py_ephem is not None and len(solns) > 1: + before = len(solns) + solns = filter_candidates_by_residual( + solns, observations, py_ephem, threshold_sigma=prefilter_threshold_sigma + ) + if len(solns) < before: + logger.debug( + f"IOD pre-filter: kept {len(solns)}/{before} " f"candidates at {prefilter_threshold_sigma}σ" + ) + 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] + # Multi-root picker. Fit every IOD candidate on the primary segment + # at the cheap screening budget, pick the best converged root, and + # only fall back to the full LM budget if nothing converged at the + # cheap tier. Gauss's polynomial gives up to 8 real roots; historic + # do_fit committed to solns[0] (largest r), which is often a + # phantom outer-SS solution for NEO-like targets. + # + # During this loop we select IAS15 adaptive_mode=2 so phantom roots + # whose trajectories pass close to Earth can't tie up the + # integrator for minutes (100-1000× wallclock blowup observed on + # diagnostic/scan with the legacy controller). The newer controller + # steps through close encounters efficiently with no accuracy cost. + # The setting is restored on every exit path. + # + # Each LM call dispatches through _run_fit so the picker honors the + # selected engine. The screen/full iteration budgets apply to the + # Cartesian engine; the BK-native engine uses its own internal cap. + saved_mode = get_ias15_adaptive_mode() + if picker_ias15_adaptive_mode >= 0: + set_ias15_adaptive_mode(picker_ias15_adaptive_mode) + 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) - if x.flag != 0: - logger.debug(f"Primary interval failed. Total observations: {len(obs)}") - x.flag = 3 # caution + try: + candidates = [_run_fit(assist_ephem, soln, obs, engine, screen_iter_max) for soln in solns] + x = _pick_best_root(candidates, min_r_helio_AU) + if x is None: + candidates = [_run_fit(assist_ephem, soln, obs, engine, full_iter_max) for soln in solns] + x = _pick_best_root(candidates, min_r_helio_AU) + finally: + set_ias15_adaptive_mode(saved_mode) + + if x is None and is_auto and len(obs) >= 3: + # Every Gauss root (if any) failed to seed a converged LM. Fall back to + # the BK 5-parameter linear IOD on the primary segment. BK-IOD shines on + # distant short arcs -- exactly where Gauss's three-point geometry is + # ill-conditioned (see bk_iod.cpp's regime-of-validity note); on the + # diagnostic scan Gauss+BK covers ~90% of cases vs ~84% for Gauss alone. + # Epoch convention matches do_gauss_iod's middle observation. + logger.debug(f"All {len(solns)} Gauss roots failed; trying BK-IOD fallback") + bk_seed = run_bk_iod(obs, float(obs[len(obs) // 2].epoch), _MU_SUN) + if bk_seed.flag == 0: + cand = _run_fit(assist_ephem, bk_seed, obs, engine, full_iter_max) + candidates.append(cand) + if cand.flag == 0: + x = cand + + if x is None: + # Still no convergence — surface the least-bad attempt so the caller has + # *something* to inspect, with a flag they can detect. + if not candidates: + # 'auto' with zero Gauss roots and no usable BK seed: nothing to + # surface, so return an explicit no-solution sentinel. + x = FitResult() + x.flag = 5 + return x + x = min(candidates, key=lambda c: c.csq) + logger.debug( + f"Primary interval: no root converged " f"(best csq={x.csq:.3g}, n_roots={len(candidates)})" + ) + x.flag = 3 return x # Attempt to fit all the data, given the fit of the primary interval + primary_x = x 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: obs = [] - x = solns[0] + # Restart from the first IOD seed, or the converged primary fit when + # there were no IOD candidates (the iod='auto' BK-IOD fallback path). + x = solns[0] if solns else primary_x 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) - print("flag:", x.flag) + logger.debug(f"Incremental fit segment {i} of {len(seq)} " f"(n_obs={len(obs)})") + x = _run_fit(assist_ephem, x, obs, engine) if x.flag != 0: x.flag = 4 break @@ -476,6 +1014,31 @@ def do_other_fit(iod: str): raise ValueError(f"The IOD, {iod} is not supported. Please use a supported IOD.") +# Minimum observational arc (in days) generally needed to constrain an orbit. +# Below this the fit is essentially unconstrained, so a failure is most likely a +# too-short baseline rather than anything wrong with the data. +_MIN_ARC_DAYS = 1.0 + + +def _warn_if_short_arc(jds, obj_id): + """Emit a helpful warning when a failed fit is likely caused by too short an + observational arc (less than ~24 hours / a single night). + + See issue #312: an orbit fit needs a baseline of more than 24 hours, so when + a fit fails on a sub-day arc we tell the user the likely cause rather than + leaving them with an opaque failure. + """ + if jds is None or len(jds) == 0: + return + arc_days = float(np.max(jds) - np.min(jds)) + if arc_days < _MIN_ARC_DAYS: + logger.warning( + f"Orbit fit failed for {obj_id}: the observations span only " + f"{arc_days * 24.0:.1f} hours. Constraining an orbit generally requires " + f"a baseline of more than ~24 hours (more than a single night of observations)." + ) + + def _orbitfit( data, cache_dir: str, @@ -483,10 +1046,14 @@ def _orbitfit( initial_guess=None, bias_dict: dict = None, sort_array: bool = True, - weight_data: bool = False, - iod: str = "gauss", - args=None, - aux=None, + weight_data=False, # bool (Veres 2017) or "supplied" (rmsRA/rmsDec columns) + iod: str = "auto", + engine: str = "cartesian", + fit_nongrav: bool = False, + nongrav_auto_thresholds=None, + nongrav_gr=None, + per_arc: bool = False, + skip_unchanged: bool = False, ): """This function will contain all of the calls to the c++ code that will calculate an orbit given a set of observations. Note that all observations @@ -508,14 +1075,41 @@ def _orbitfit( A dictionary containing bias corrections for different catalogs. sort_array : bool Whether to sort the observations by obstime before processing. Default is True. - weight_data : bool - Whether to apply data weighting based on the observation code, date, catalog - and program. Default is False. + weight_data : bool or str + Astrometric weighting. ``False`` (default) leaves the built-in default + uncertainty. ``True`` applies the Veres 2017 model (observation code, date, + catalog, program). ``"supplied"`` uses the per-observation ``rmsRA`` / + ``rmsDec`` columns (arcseconds) directly -- e.g. ADES-reported + uncertainties, or an external weighting model such as era-based historical + weighting for old comet apparitions (a row with a NaN/nonpositive value + falls back to the default). iod : str - The IOD used to generate an initial guess orbit. Currently supports ['gauss']. - Default is 'gauss'. + The IOD used to generate an initial guess orbit. Supports 'gauss' + and 'auto' (Gauss with BK-IOD fallback). + Default is 'auto'. """ - _RESULT_DTYPES = _get_result_dtypes(primary_id_column_name) + # Fitting non-gravitational params (issue #351) uses the joint state+nongrav + # LM, which only the Cartesian engine supports; the BK-native engine assumes a + # 6D state. Override with a warning, mirroring the radar path. + auto_nongrav = isinstance(fit_nongrav, str) and fit_nongrav.strip().lower() == "auto" + if auto_nongrav: + # 'auto' selects the non-grav model per object (issue #357); the schema + # carries all of A1/A2/A3 and each row reports only the adopted params. + nongrav_mask, nongrav_names = 0, ["A1", "A2", "A3"] + else: + nongrav_mask, nongrav_names = _parse_nongrav(fit_nongrav) + if (nongrav_mask or auto_nongrav) and engine != "cartesian": + logger.warning("Non-gravitational fitting requires engine='cartesian'; overriding %r.", engine) + engine = "cartesian" + + # Per-arc (piecewise-constant) non-grav amplitudes need an explicit non-grav + # mask (which params to split per apparition); it is meaningless for a + # gravity-only or 'auto'-selected fit. Ignore with a warning otherwise. + if per_arc and not nongrav_mask: + logger.warning("per_arc=True requires an explicit fit_nongrav mask (e.g. 'A1A2A3'); ignoring.") + per_arc = False + + _RESULT_DTYPES = _get_result_dtypes(primary_id_column_name, nongrav_names, per_arc=per_arc) if len(data) == 0: return np.array([], dtype=_RESULT_DTYPES) @@ -545,20 +1139,47 @@ def _orbitfit( # Check if certain columns are present in the data column_names = data.dtype.names - g_column_present = "astCat" in column_names + astcat_column_present = "astCat" in column_names program_column_present = "program" in column_names position_rates_columns_present = all(col in column_names for col in ["raRate", "decRate"]) + rate_unc_columns_present = all(col in column_names for col in ["rmsRArate", "rmsDecrate"]) + astrom_unc_columns_present = all(col in column_names for col in ["rmsRA", "rmsDec"]) + radar_columns_present = any(col in column_names for col in ["delay", "doppler"]) + + # Fingerprint the raw observation set (issue #419), before any in-place + # debiasing mutates ra/dec, so it reflects what was reported. + nobs_fit, obs_hash = _obs_fingerprint(data, column_names) + + # Lever 1 (skip-unchanged): if a prior converged fit for this object was + # built from the identical observation set, carry it forward verbatim + # instead of re-fitting. ``initial_guess`` has already been filtered to + # this object and reset to None when its flag != 0, so a match here is a + # successful prior fit over the same obs. Requires the prior catalog to + # carry the fingerprint columns; otherwise this never triggers and we fit. + if ( + skip_unchanged + and initial_guess is not None + and "obs_hash" in initial_guess.dtype.names + and "nobs_fit" in initial_guess.dtype.names + and int(initial_guess["nobs_fit"][0]) == nobs_fit + and str(initial_guess["obs_hash"][0]) == obs_hash + ): + return _carry_forward_result(initial_guess, _RESULT_DTYPES) # Accommodate occultation measurements. These measurements are implied when # the "ra" and "dec" columns are None. In this case, we will use the "starra" - # and "stardec" columns. + # and "stardec" columns. Radar rows have no ra/dec and are skipped. for d in data: + if radar_columns_present and _is_radar(d, column_names): + continue if _is_occultation(d): d = _use_star_astrometry(d) # bias_dict will be a dictionary when the debias flag is set to True. if bias_dict is not None: for d in data: + if radar_columns_present and _is_radar(d, column_names): + continue # debiasing is an astrometric (ra/dec) correction d["ra"], d["dec"] = debias( ra=d["ra"], dec=d["dec"], @@ -573,16 +1194,38 @@ def _orbitfit( # radians. observations = [] for d in data: + if radar_columns_present and _is_radar(d, column_names): + o = _radar_observation( + d[primary_id_column_name], + d, + convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), # JD TDB + column_names, + ) + observations.append(o) + continue if position_rates_columns_present and (not np.isnan(d["raRate"]) and not np.isnan(d["decRate"])): + # Rate uncertainties (rmsRArate/rmsDecrate) share raRate's + # arcsec/hour units; convert to rad/day. Absent -> C++ default. + streak_rate_unc = {} + if ( + rate_unc_columns_present + and not np.isnan(d["rmsRArate"]) + and not np.isnan(d["rmsDecrate"]) + ): + streak_rate_unc["ra_rate_unc"] = abs(d["rmsRArate"]) * ARCSEC_PER_HOUR_TO_RAD_PER_DAY + streak_rate_unc["dec_rate_unc"] = abs(d["rmsDecrate"]) * ARCSEC_PER_HOUR_TO_RAD_PER_DAY o = Observation.from_streak_with_id( str(d[primary_id_column_name]), d["ra"] * np.pi / 180.0, d["dec"] * np.pi / 180.0, - d["raRate"], - d["decRate"], + # arcsec/hour (great-circle) -> rad/day; raRate already + # carries the cos(Dec) factor (see module constant above). + d["raRate"] * ARCSEC_PER_HOUR_TO_RAD_PER_DAY, + d["decRate"] * ARCSEC_PER_HOUR_TO_RAD_PER_DAY, convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), # Convert obstime to JD TDB [d["x"], d["y"], d["z"]], # Barycentric position [d["vx"], d["vy"], d["vz"]], # Barycentric velocity + **streak_rate_unc, ) else: o = Observation.from_astrometry_with_id( @@ -594,22 +1237,50 @@ def _orbitfit( [d["vx"], d["vy"], d["vz"]], # Barycentric velocity ) - if weight_data: - data_weight = data_weight_Veres2017( + # Astrometric weighting. ``weight_data="supplied"`` uses the per-obs + # rmsRA/rmsDec columns (arcsec) directly -- e.g. ADES-reported + # uncertainties, or an external weighting model such as era-based + # historical weighting for old comet apparitions. ``weight_data=True`` + # uses the Veres 2017 model; ``False`` leaves the C++ default. Supplied + # takes precedence; a NaN/nonpositive supplied value on a row falls + # back to the C++ default for that row. + if isinstance(weight_data, str) and weight_data.lower() == "supplied": + if not astrom_unc_columns_present: + raise ValueError('weight_data="supplied" requires rmsRA and rmsDec columns (arcsec).') + if np.isfinite(d["rmsRA"]) and d["rmsRA"] > 0: + o.ra_unc = abs(d["rmsRA"]) * np.pi / (180.0 * 3600.0) + if np.isfinite(d["rmsDec"]) and d["rmsDec"] > 0: + o.dec_unc = abs(d["rmsDec"]) * np.pi / (180.0 * 3600.0) + elif weight_data: + # astrometric_uncertainty_Veres2017 returns the astrometric uncertainty in + # ARCSECONDS (per its docstring), but Observation.ra_unc / + # dec_unc are stored in RADIANS. Convert at the assignment. + sigma_arcsec = astrometric_uncertainty_Veres2017( obsCode=d["stn"], jd_tdb=convert_tdb_date_to_julian_date(d["obsTime"], cache_dir), catalog=d["astCat"] if astcat_column_present else None, program=d["program"] if program_column_present else None, ) + sigma_rad = sigma_arcsec * np.pi / (180.0 * 3600.0) - o.ra_unc = data_weight - o.dec_unc = data_weight + o.ra_unc = sigma_rad + o.dec_unc = sigma_rad observations.append(o) + # Radar delay/Doppler rows use the variable-row packing, which only the + # Cartesian engine supports; the BK-native engine assumes 2 rows per + # observation and would silently drop them. + if radar_columns_present and engine != "cartesian": + logger.warning( + "Radar (delay/Doppler) observations require engine='cartesian'; overriding %r.", + engine, + ) + engine = "cartesian" + # if cache_dir is not provided, use the default os_cache if cache_dir is None: - kernels_loc = str(pooch.os_cache("layup")) + kernels_loc = str(default_cache_dir()) else: kernels_loc = str(cache_dir) @@ -618,23 +1289,88 @@ def _orbitfit( # Perform the orbit fitting if initial_guess is None or initial_guess["flag"] != 0: - if iod.lower() in ["gauss", "herget"]: + if iod.lower() in ["gauss", "auto", "herget"]: res = do_fit( observations=observations, seq=sequence, cache_dir=kernels_loc, iod=iod.lower(), - args=args, - aux=aux, + engine=engine, ) else: res = do_other_fit(iod=iod.lower()) else: guess_to_use = parse_fit_result(initial_guess) res = run_from_vector_with_initial_guess(get_ephem(kernels_loc), guess_to_use, observations) + + # Non-gravitational params (issue #351): once the 6-parameter orbit has + # converged, refine it jointly with the requested non-grav params, seeded + # from that solution. They are weakly constrained on short arcs, so if the + # joint fit is degenerate (flag 6) or fails to converge we keep the + # 6-parameter result and report the params as NaN (graceful guard). + if auto_nongrav and res.flag == 0: + # Adopt the most parsimonious statistically-warranted non-grav model, + # or keep the gravity-only fit (issue #357). + res = _select_nongrav_auto( + get_ephem(kernels_loc), + res, + observations, + nongrav_auto_thresholds or _AUTO_DEFAULT_THRESHOLDS, + gofr=nongrav_gr, + ) + elif nongrav_mask and res.flag == 0: + res_ng = run_from_vector_with_initial_guess( + get_ephem(kernels_loc), + res, + observations, + nongrav_mask=nongrav_mask, + gofr=_gofr_arg(nongrav_gr), + per_arc=per_arc, + ) + if res_ng.flag == 0: + res = res_ng + else: + logger.debug("Non-grav refinement did not converge; reporting non-grav params as NaN.") + # Populate our output structured array with the orbit fit results success = res.flag == 0 + if not success: + _warn_if_short_arc(jds, data[primary_id_column_name][0]) cov_matrix = tuple(res.cov[i] for i in range(36)) if success else (np.nan,) * 36 + nongrav_cols = () + if nongrav_names: + # Report each parameter only if it was actually adopted (its bit is set + # in the fit's nongrav_mask); for 'auto' that is the selected subset. + fitted_mask = getattr(res, "nongrav_mask", 0) if success else 0 + nongrav_cols = tuple( + v + for n in nongrav_names + for v in ( + (getattr(res, n.lower()) if (fitted_mask & _NONGRAV_BITS[n]) else np.nan), + (getattr(res, n.lower() + "_unc") if (fitted_mask & _NONGRAV_BITS[n]) else np.nan), + ) + ) + # Later-arc amplitudes (comet linkage): the C++ result reports them in the + # _arc2 fields only when per_arc fitting was on and converged. + per_arc_cols = () + if per_arc: + per_arc_on = success and getattr(res, "per_arc", False) + per_arc_cols = tuple( + v + for n in nongrav_names + for v in ( + ( + getattr(res, n.lower() + "_arc2") + if (per_arc_on and fitted_mask & _NONGRAV_BITS[n]) + else np.nan + ), + ( + getattr(res, n.lower() + "_arc2_unc") + if (per_arc_on and fitted_mask & _NONGRAV_BITS[n]) + else np.nan + ), + ) + ) output = np.array( [ ( @@ -651,6 +1387,9 @@ def _orbitfit( ("BCART_EQ" if success else "NONE"), # The base format returned by the C++ code ) + cov_matrix # Flat covariance matrix + + nongrav_cols # non-grav params + uncertainties (issue #351), when fit_nongrav + + per_arc_cols # later-arc amplitudes (comet linkage), when per_arc + + (obs_hash, nobs_fit) # obs fingerprint (issue #419) ], dtype=_RESULT_DTYPES, ) @@ -666,9 +1405,13 @@ def orbitfit( primary_id_column_name="provID", debias=False, weight_data=False, - iod="gauss", - args=None, - aux=None, + iod="auto", + engine="cartesian", + fit_nongrav=False, + nongrav_auto_thresholds=None, + nongrav_gr=None, + per_arc=False, + skip_unchanged=False, ): """This is the function that you would call interactively. i.e. from a notebook @@ -686,14 +1429,79 @@ def orbitfit( The name of the primary identifier column for the objects. Default is "provID". debias : bool Whether to apply debiasing corrections to the observations. Default is False. - weight_data : bool - Whether to apply data weighting based on the observation code, date, catalog - and program. Default is False. + weight_data : bool or str + Astrometric weighting. ``False`` (default) leaves the built-in default + uncertainty. ``True`` applies the Veres 2017 model (observation code, date, + catalog, program). ``"supplied"`` uses the per-observation ``rmsRA`` / + ``rmsDec`` columns (arcseconds) directly -- e.g. ADES-reported + uncertainties, or an external weighting model such as era-based historical + weighting for old comet apparitions (a row with a NaN/nonpositive value + falls back to the default). iod : str - The IOD used to generate an initial guess orbit. Currently supports ['gauss']. - Default is 'gauss'. + The IOD used to generate an initial guess orbit. Supports 'gauss', + 'herget' and 'auto' (Gauss with BK-IOD fallback). + Default is 'auto'. + fit_nongrav : bool | str | iterable of str + Which non-gravitational Marsden parameters to fit after the 6-parameter + orbit converges. ``False`` (default) fits none; ``True`` fits A2 (the + transverse Yarkovsky term, the common asteroid case); a string or iterable + naming params -- e.g. ``"A2"``, ``"A1A2A3"``, ``["A1", "A3"]`` -- selects a + subset of A1 (radial), A2 (transverse), A3 (normal). ``"auto"`` selects the + model adaptively per object (issue #357): the gravity-only fit is kept + unless its reduced chi-square is unacceptable, in which case the most + parsimonious non-grav model that is well-conditioned and statistically + significant is adopted (the A1/A2/A3 columns are all present, with only the + adopted params filled and the rest NaN). For each fitted param an ``a{n}`` + value and ``a{n}_unc`` 1-sigma column (au/day^2) are added to the result. + Cartesian engine only; params weakly constrained on short arcs are reported + as NaN (issue #351). + nongrav_auto_thresholds : NongravAutoThresholds, optional + Decision thresholds used when ``fit_nongrav="auto"`` -- how unacceptable the + gravity-only fit must be before a non-grav is tried, and how large the + chi-square drop and per-parameter significance must be to adopt one. Default + (``None``) uses the standard thresholds; pass a customized + ``NongravAutoThresholds`` to tune. Ignored unless ``fit_nongrav="auto"``. + nongrav_gr : sequence of float, optional + The non-gravitational g(r) sublimation law as ``[alpha, nm, nn, nk, r0]`` in + ASSIST's parameterization ``g(r) = alpha*(r/r0)^-nm*(1+(r/r0)^nn)^-nk``. + Default (``None``) is the asteroidal inverse-square law ``(r/r0)^-2`` used by + Yarkovsky A2 fits; pass a cometary law (e.g. Marsden water-ice) to fit a + comet's non-gravs. Applies to any non-grav fit (explicit or ``"auto"``). + per_arc : bool, optional + Fit piecewise-constant *per-apparition* non-grav amplitudes (comet + linkage). The state and ``g(r)`` are shared, but observations before the + fit epoch (the earlier arc) and after it (the later arc) each get their own + ``[A1,A2,A3]``. Requires an explicit ``fit_nongrav`` mask and an + ``initial_guess`` whose epoch sits between the two apparitions. The output + adds ``a{1,2,3}_arc2`` columns for the later arc; the base ``a{1,2,3}`` + columns then hold the earlier arc. Default False. + skip_unchanged : bool + Incremental / steady-state mode (issue #419). When True and ``initial_guess`` + is a prior result catalog carrying the ``obs_hash`` fingerprint columns, any + object whose observation set is byte-for-byte unchanged since that catalog is + carried forward verbatim instead of re-fitting. Objects with changed obs are + re-fit, warm-started from the prior state when available. Default False (every + object is fit). The output always carries the ``obs_hash``/``nobs_fit`` + columns so it can seed the next cycle. """ + # Incremental / steady-state pre-filter (issue #419). Before any per-obs + # ephemeris or observatory setup, drop objects whose observation set is + # unchanged since the prior catalog and carry their stored fit forward + # verbatim. This is where the steady-state throughput win comes from -- a + # skipped object costs one fingerprint hash, not a fit. Changed objects fall + # through and are re-fit below, warm-started from the prior state via + # ``initial_guess`` (which is retained for exactly that purpose). + carried_forward = None + if ( + skip_unchanged + and initial_guess is not None + and "obs_hash" in getattr(initial_guess, "dtype", np.dtype([])).names + ): + data, carried_forward = _partition_unchanged(data, initial_guess, primary_id_column_name, fit_nongrav) + if len(data) == 0: # everything unchanged -> nothing to fit + return carried_forward + layup_observatory = LayupObservatory(cache_dir=cache_dir) # The units of et are seconds (from J2000). This new column is used by @@ -704,11 +1512,17 @@ def orbitfit( pos_vel = layup_observatory.obscodes_to_barycentric(data) data = rfn.merge_arrays([data, pos_vel], flatten=True, asrecarray=True, usemask=False) + # Radar (delay/Doppler) observations need the barycentric observer + # acceleration for the two-leg light-time model; compute it only when radar + # columns are present so optical/streak fits are unaffected. + if any(col in data.dtype.names for col in ("delay", "doppler")): + data = _append_observer_acceleration(data, layup_observatory) + bias_dict = None if debias: bias_dict = generate_bias_dict(cache_dir) - return process_data_by_id( + fitted = process_data_by_id( data, num_workers, _orbitfit, @@ -718,9 +1532,370 @@ def orbitfit( bias_dict=bias_dict, weight_data=weight_data, iod=iod, - args=args, - aux=aux, + engine=engine, + fit_nongrav=fit_nongrav, + nongrav_auto_thresholds=nongrav_auto_thresholds, + nongrav_gr=nongrav_gr, + per_arc=per_arc, + skip_unchanged=skip_unchanged, ) + # Re-attach objects carried forward unchanged by the #419 pre-filter. + if carried_forward is not None and len(carried_forward): + return np.concatenate([fitted, carried_forward]) + return fitted + + +def _observations_for_update(data, cache_dir, weight_data=False, bias_dict=None): + """Augment one object's observations with the observer barycentric state and + build the C++ ``Observation`` list, mirroring ``orbitfit()``'s preprocessing. + + Supports optical astrometry and streak (rate) rows -- the observation kinds a + steady-state catalog update sees. (Radar/occultation are not yet handled by + the sequential path; the driver's full-refit fallback covers them.) + """ + DEG = np.pi / 180.0 + kernels_loc = str(default_cache_dir()) if cache_dir is None else str(cache_dir) + observatory = LayupObservatory(cache_dir=cache_dir) + + et = np.array([spice.str2et(row["obsTime"]) for row in data], dtype=" the default layup os_cache). + all_data : numpy structured array, optional + The full observation set (old + new). Required for the nonlinearity + fallback: when the update is too large the driver refits over all_data. + weight_data, debias_data : bool + Apply Veres (2017) weighting / MPC debiasing to the new observations, + matching the corresponding ``orbitfit`` options. + max_update_sigma : float + Nonlinearity gate. If the update moves the state more than this many prior + standard deviations (Mahalanobis), fall back to a full refit over + ``all_data`` when provided, else flag the result (flag=8). + iter_max : int + LM iteration cap. + + Returns + ------- + FitResult + The updated fit. ``method`` is ``"sequential_update"`` for an accepted + information-filter update, or ``"orbit_fit"`` when the fallback refit ran. + """ + prior_fit = prior if isinstance(prior, FitResult) else parse_fit_result(prior) + kernels_loc = str(default_cache_dir()) if cache_dir is None else str(cache_dir) + ephem = get_ephem(kernels_loc) + bias_dict = generate_bias_dict(cache_dir) if debias_data else None + + new_obs = _observations_for_update(new_data, cache_dir, weight_data, bias_dict) + seq = run_sequential_update(ephem, prior_fit, new_obs, iter_max) + + def _full_refit(): + if all_data is None: + return None + all_obs = _observations_for_update(all_data, cache_dir, weight_data, bias_dict) + return run_from_vector_with_initial_guess(ephem, prior_fit, all_obs, iter_max) + + # The information update did not converge (e.g. a non-positive-definite prior, + # flag 7): fall back to a full refit if we can, else surface the failure. + if seq.flag != 0: + fallback = _full_refit() + return fallback if fallback is not None else seq + + # Nonlinearity gate: a large move means the linearization is untrustworthy. + if _update_mahalanobis(prior_fit, seq) > max_update_sigma: + fallback = _full_refit() + if fallback is not None: + return fallback + seq.flag = 8 # nonlinear update, no full-obs set supplied to refit + return seq + + +def _obs_row_keys(data, column_names): + """Per-observation identity keys over the fit-relevant columns (issue #419). + + Each key is the same per-row string that ``_obs_fingerprint`` hashes, so a + row's key equals its contribution to the object's fingerprint. Used to diff a + current observation set against the one a prior fit was built from. + """ + cols = [c for c in _FINGERPRINT_COLUMNS if c in column_names] + return ["\x1f".join(_fmt_fingerprint_value(d[c]) for c in cols) for d in data] + + +def _append_only_new_obs(current, prior_obs): + """Return the rows of ``current`` absent from ``prior_obs`` if the change is + append-only, else ``None``. + + Append-only means every observation the prior was fit from is still present + (none removed or modified) -- the case the sequential update handles exactly. + If any prior observation is gone, the summarised old-obs information no longer + matches the current set, so the object must be fully re-fit and this returns + ``None``. + """ + cur_keys = _obs_row_keys(current, current.dtype.names) + prior_keys = set(_obs_row_keys(prior_obs, prior_obs.dtype.names)) + if not prior_keys.issubset(cur_keys): + return None # an old observation was removed or changed -> not append-only + mask = np.array([k not in prior_keys for k in cur_keys], dtype=bool) + return current[mask] + + +def _group_by_id(data, primary_id_column_name): + if data is None: + return {} + ids = data[primary_id_column_name] + return {oid: data[ids == oid] for oid in np.unique(ids)} + + +def _fitresult_to_row(fit, obj_id, obs_hash, nobs_fit, dtypes): + """Pack a FitResult (from the sequential update or its refit fallback) into one + result-catalog row carrying the current-obs fingerprint, so the row is + identical in shape to ``orbitfit`` output and can seed the next cycle.""" + success = fit.flag == 0 + cov = tuple(fit.cov[i] for i in range(36)) if success else (np.nan,) * 36 + row = ( + (obj_id, (fit.csq if success else np.nan), fit.ndof) + + (tuple(fit.state[i] for i in range(6)) if success else (np.nan,) * 6) + + ( + (fit.epoch - 2400000.5) if success else np.nan, + fit.niter, + fit.method, + fit.flag, + "BCART_EQ" if success else "NONE", + ) + + cov + + (obs_hash, nobs_fit) + ) + return np.array([row], dtype=dtypes) + + +def incremental_orbitfit( + data, + cache_dir, + prior_catalog, + *, + prior_obs=None, + primary_id_column_name="provID", + weight_data=False, + debias=False, + max_update_sigma=4.0, + iod="auto", + engine="cartesian", + num_workers=1, +): + """Steady-state incremental fit over a batch of objects (issue #419 capstone). + + Ties the three levers into one operational maintenance pass. For each object + in ``data`` (the current observations), routes: + + * **skip** -- the observation set is unchanged since ``prior_catalog`` (matching + fingerprint): carry the prior fit forward verbatim, no fit. + * **sequential update** -- observations were only appended (``prior_obs`` given + and every prior observation is still present): update the prior with the new + observations only, via :func:`sequential_update` (integrating just the new + obs). Its nonlinearity gate falls back to a full refit when the update is + too large. + * **full refit** -- observations were removed or changed, or no per-object + ``prior_obs`` is available: refit over all current obs, warm-started from the + prior state when there is one, cold (IOD) when the object is new. + + Parameters + ---------- + data : numpy structured array + Current observations for all objects (grouped by ``primary_id_column_name``). + cache_dir : str or None + Kernel/ephemeris cache directory (None -> the default layup os_cache). + prior_catalog : numpy structured array or None + Prior fit results carrying state/cov/epoch and the ``obs_hash``/``nobs_fit`` + fingerprint columns (i.e. produced by ``orbitfit``/this driver). None -> every + object is cold-fit. + prior_obs : numpy structured array, optional + The observations the prior catalog was fit from, grouped by id. Enables the + sequential route (needs the per-object obs to diff). Without it, changed + objects are fully refit. + weight_data, debias : bool + Veres (2017) weighting / MPC debiasing, as in ``orbitfit``. + max_update_sigma : float + Nonlinearity gate passed to :func:`sequential_update`. + + Returns + ------- + (numpy structured array, dict) + The updated result catalog (one row per object, same schema as + ``orbitfit`` output) and a routing tally + ``{"skip", "sequential", "sequential_fallback", "full", "cold"}``. + """ + from collections import Counter + + pid = primary_id_column_name + out_dtype = _get_result_dtypes(pid) + prior_by_id = {row[pid]: row for row in np.atleast_1d(prior_catalog)} if prior_catalog is not None else {} + prior_obs_by_id = _group_by_id(prior_obs, pid) + + routing = Counter() + carried, seq_rows = [], [] + warm_ids, cold_ids = [], [] + + for oid in np.unique(data[pid]): + cur = data[data[pid] == oid] + nobs, obs_hash = _obs_fingerprint(cur, cur.dtype.names) + p = prior_by_id.get(oid) + has_prior = p is not None and int(p["flag"]) == 0 and "obs_hash" in np.atleast_1d(p).dtype.names + + # Route 1: unchanged -> skip (carry the prior row forward). + if has_prior and str(p["obs_hash"]) == obs_hash and int(p["nobs_fit"]) == nobs: + carried.append(_carry_forward_result(p, out_dtype)) + routing["skip"] += 1 + continue + + # Route 2: append-only change with the prior obs available -> sequential update. + if has_prior and oid in prior_obs_by_id: + new_rows = _append_only_new_obs(cur, prior_obs_by_id[oid]) + if new_rows is not None and len(new_rows) > 0: + seq = sequential_update( + p, + new_rows, + cache_dir, + all_data=cur, + weight_data=weight_data, + debias_data=debias, + max_update_sigma=max_update_sigma, + ) + routing["sequential" if seq.method == "sequential_update" else "sequential_fallback"] += 1 + seq_rows.append(_fitresult_to_row(seq, oid, obs_hash, nobs, out_dtype)) + continue + + # Route 3: full refit (warm if a prior exists, else cold IOD). + (warm_ids if has_prior else cold_ids).append(oid) + routing["full" if has_prior else "cold"] += 1 + + # Full/cold refits go through orbitfit (warm objects filtered to their priors; + # cold objects with no initial guess). Two calls keep _orbitfit's per-object + # initial-guess lookup happy (it errors on a guess with no row for the object). + fit_parts = [] + common = dict( + cache_dir=cache_dir, + primary_id_column_name=pid, + weight_data=weight_data, + debias=debias, + iod=iod, + engine=engine, + num_workers=num_workers, + ) + if warm_ids: + sub = data[np.isin(data[pid], warm_ids)] + fit_parts.append(orbitfit(sub, initial_guess=prior_catalog, **common)) + if cold_ids: + sub = data[np.isin(data[pid], cold_ids)] + fit_parts.append(orbitfit(sub, initial_guess=None, **common)) + + parts = ( + ([np.concatenate(carried)] if carried else []) + + ([np.concatenate(seq_rows)] if seq_rows else []) + + [p for p in fit_parts if len(p)] + ) + result = np.concatenate(parts) if parts else np.array([], dtype=out_dtype) + return result, dict(routing) def orbitfit_cli( @@ -731,7 +1906,6 @@ def orbitfit_cli( chunk_size: int = 10_000, num_workers: int = -1, cli_args: Optional[Namespace] = None, - aux: any = None, ): """This is the function that is called from the command line @@ -760,13 +1934,15 @@ def orbitfit_cli( weight_data = cli_args.weight_data output_orbit_format = cli_args.output_orbit_format iod = cli_args.iod + engine = getattr(cli_args, "engine", "cartesian") else: cache_dir = None debias = False guess_file = None weight_data = False output_orbit_format = "COM" # Default output orbit format. - iod = "gauss" + iod = "auto" + engine = "cartesian" _primary_id_column_name = cli_args.primary_id_column_name @@ -801,8 +1977,7 @@ def orbitfit_cli( else Path(f"{output_file_stem_flagged}.h5") ) - if num_workers < 0: - num_workers = os.cpu_count() + num_workers = resolve_num_workers(num_workers) # Check that input file exists if not input_file.exists(): @@ -850,6 +2025,22 @@ def orbitfit_cli( chunks = create_chunks(reader, chunk_size) + # Output is written one chunk at a time. The first write to each file + # overwrites any stale file left by a previous run (so re-runs are + # idempotent); later chunks append. Without this a re-run, or a re-merge onto + # an existing file, would duplicate every row. + _written_files = set() + + def _emit(arr, path): + first = path not in _written_files + _written_files.add(path) + if output_file_format == "hdf5": + (write_hdf5 if first else append_hdf5)(arr, path, key="data") + else: # csv: write_csv appends when the file already exists + if first and os.path.exists(path): + os.remove(path) + write_csv(arr, path) + for chunk in chunks: data = reader.read_objects(chunk) initial_guess = None @@ -877,8 +2068,7 @@ def orbitfit_cli( debias=debias, weight_data=weight_data, iod=iod, - args=cli_args, - aux=aux, + engine=engine, ) # Convert the fit_orbits to the preferred output format @@ -897,30 +2087,14 @@ def orbitfit_cli( fit_orbits_success = fit_orbits[success_mask] fit_orbits_failed = fit_orbits[~success_mask] - if output_file_format == "hdf5": - if len(fit_orbits_success) > 0: - write_hdf5(fit_orbits_success, output_file, key="data") - - if len(fit_orbits_failed) > 0: - write_hdf5( - fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], - output_file_flagged, - key="data", - ) - else: # csv output format - if len(fit_orbits_success) > 0: - write_csv(fit_orbits_success, output_file) - - if len(fit_orbits_failed) > 0: - write_csv( - fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], output_file_flagged - ) + if len(fit_orbits_success) > 0: + _emit(fit_orbits_success, output_file) + + if len(fit_orbits_failed) > 0: + _emit(fit_orbits_failed[[_primary_id_column_name, "method", "flag"]], output_file_flagged) else: # All results go to a single output file - if output_file_format == "hdf5": - write_hdf5(fit_orbits, output_file, key="data") - else: - write_csv(fit_orbits, output_file) + _emit(fit_orbits, output_file) logger.info(f"Data has been written to {output_file}") @@ -939,13 +2113,12 @@ def _is_valid_data(data): bool True if the data is valid, False otherwise. """ + column_names = data.dtype.names valid_conditions = [ len(data) >= 3, np.all( data["et"] >= -6279962400.00 ), # excludes all datasets before 1801, data["et"] = 0 is j2000, 6279962400.00 is seconds between 1801 and j2000 - np.all(is_numeric(data["ra"])), - np.all(is_numeric(data["dec"])), np.all(is_numeric(data["x"])), np.all(is_numeric(data["y"])), np.all(is_numeric(data["z"])), @@ -953,6 +2126,23 @@ def _is_valid_data(data): np.all(is_numeric(data["vy"])), np.all(is_numeric(data["vz"])), ] + + # Each row must carry a usable observable: ra/dec for optical (and streak), + # or a delay/Doppler for radar. Validate per row so a radar file -- which has + # no ra/dec columns -- is not rejected, while optical rows still require + # numeric ra/dec. + if any(c in column_names for c in ["delay", "doppler"]): + have_radec = "ra" in column_names and "dec" in column_names + valid_conditions.append( + all( + _is_radar(d, column_names) or (have_radec and is_numeric(d["ra"]) and is_numeric(d["dec"])) + for d in data + ) + ) + else: + valid_conditions.append(np.all(is_numeric(data["ra"]))) + valid_conditions.append(np.all(is_numeric(data["dec"]))) + return all(valid_conditions) From fbc10afcc43705d04cf095c874b3947bf64e4f9f Mon Sep 17 00:00:00 2001 From: awilson110 Date: Fri, 28 Aug 2026 11:17:43 +0100 Subject: [PATCH 3/4] updated herget method --- src/layup/utilities/herget_iod.py | 7 +++---- tests/layup/test_herget_iod.py | 26 ++++++-------------------- 2 files changed, 9 insertions(+), 24 deletions(-) diff --git a/src/layup/utilities/herget_iod.py b/src/layup/utilities/herget_iod.py index 08e2efc4..d2977b6f 100644 --- a/src/layup/utilities/herget_iod.py +++ b/src/layup/utilities/herget_iod.py @@ -11,7 +11,7 @@ SPEED_OF_LIGHT_AU_DAY = 173.145 -def herget_with_assist(observations, seq, tolerance, args, aux, max_iterations=100): +def herget_with_assist(observations, seq, ephem, tolerance=0.001, max_iterations=100): """Runs the Herget method on a set of observations. Parameters @@ -66,7 +66,7 @@ def herget_with_assist(observations, seq, tolerance, args, aux, max_iterations=1 # print(observation.epoch) delta_rho1, delta_rhon, state_1 = find_drho( - obs, t1, tn, r1, rn, tolerance, args, aux, rho_hat_1, rho_hat_n + obs, t1, tn, r1, rn, tolerance, ephem, rho_hat_1, rho_hat_n ) # Update rho values @@ -93,7 +93,7 @@ def herget_with_assist(observations, seq, tolerance, args, aux, max_iterations=1 return [solution] -def find_drho(observations, t1, tn, r1, rn, tolerance, args, aux, rho_hat_1, rho_hat_n): +def find_drho(observations, t1, tn, r1, rn, tolerance, ephem, rho_hat_1, rho_hat_n): """Find the adjustment to make to rho_1 and rho_n to make in order to reduce the residuals of the observations Parameters @@ -134,7 +134,6 @@ def find_drho(observations, t1, tn, r1, rn, tolerance, args, aux, rho_hat_1, rho [var_vx1, var_vy1, var_vz1], _ = find_velocity(t1, tn, r1 + rho_hat_1, rn, tolerance) # Simulation setup - ephem, _, _ = create_assist_ephemeris(args, aux) sim = rebound.Simulation() sim.add(x=r1[0], y=r1[1], z=r1[2], vx=vx1, vy=vy1, vz=vz1) diff --git a/tests/layup/test_herget_iod.py b/tests/layup/test_herget_iod.py index 5d1a0367..e7677a24 100644 --- a/tests/layup/test_herget_iod.py +++ b/tests/layup/test_herget_iod.py @@ -5,18 +5,15 @@ import layup.utilities.herget_iod as herget import spiceypy as spice -from layup.utilities.layup_configs import LayupConfigs -from sorcha.ephemeris.simulation_setup import create_assist_ephemeris from layup.utilities.data_utilities_for_tests import get_test_filepath from layup.utilities.file_io.CSVReader import CSVDataReader -from layup.routines import Observation +from layup.routines import Observation, get_ephem from layup.utilities.data_processing_utilities import LayupObservatory from layup.utilities.datetime_conversions import convert_tdb_date_to_julian_date from layup.orbitfit import _build_sequence - -from sorcha.ephemeris.simulation_setup import create_assist_ephemeris import assist import rebound +from layup.orbit_maths import build_ephem_and_mus SPEED_OF_LIGHT_AU_DAY = 173.145 @@ -102,17 +99,6 @@ def test_find_drho(tmpdir): jds = convert_tdb_date_to_julian_date(data["obsTime"]) sequence = _build_sequence(jds, sep_dt=90.0) - class FakeCliArgs: - def __init__(self, g=None): - self.primary_id_column_name = "ObjID" - self.n = 1 - self.chunk = 10000 - self.ar_data_file_path = None - self.force = True - self.code_format = True - - args = FakeCliArgs() - aux = LayupConfigs().auxiliary obs_1 = observations[0] r_e_1 = obs_1.observer_position @@ -137,8 +123,9 @@ def __init__(self, g=None): observation.epoch = epochs[i] - ((rho_1) + (rho_n)) / (2 * SPEED_OF_LIGHT_AU_DAY) state_1[3:], _ = herget.find_velocity(t_1, t_n, state_1[:3], r_n, 0.001) - - ephem, _, _ = create_assist_ephemeris(args, aux) + + ephem, _, _ = build_ephem_and_mus() + print(dir(ephem)) sim = rebound.Simulation() ex = assist.Extras(sim, ephem) sim.t = t_1 - ephem.jd_ref @@ -166,7 +153,7 @@ def __init__(self, g=None): # call find_drho, check if it reduces the sum of the residuals delta_rho1, delta_rhon, state_1 = herget.find_drho( - observations, t_1, t_n, state_1[:3], r_n, 0.001, args, aux, rho_hat_1, rho_hat_n + observations, t_1, t_n, state_1[:3], r_n, 0.001, ephem, rho_hat_1, rho_hat_n ) # Update rho values @@ -180,7 +167,6 @@ def __init__(self, g=None): state_1[3:], _ = herget.find_velocity(t_1, t_n, state_1[:3], r_n, 0.001) - ephem, _, _ = create_assist_ephemeris(args, aux) sim = rebound.Simulation() ex = assist.Extras(sim, ephem) sim.t = t_1 - ephem.jd_ref From 734ca5dddfb8cbfb544a1684e9f41e048b3ccf3e Mon Sep 17 00:00:00 2001 From: awilson110 Date: Fri, 28 Aug 2026 11:26:59 +0100 Subject: [PATCH 4/4] used linter --- src/layup/iod.py | 3 ++- tests/layup/test_herget_iod.py | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/layup/iod.py b/src/layup/iod.py index 806eb110..0ff5677d 100644 --- a/src/layup/iod.py +++ b/src/layup/iod.py @@ -112,8 +112,9 @@ def gauss_iod(observations, seq): register_iod("gauss", gauss_iod) + def herget_iod(observations, seq): - '''''' + """""" ephem, _, _ = build_ephem_and_mus() solns = herget_with_assist(observations, seq, ephem, tolerance=0.0001, max_iterations=100) return solns diff --git a/tests/layup/test_herget_iod.py b/tests/layup/test_herget_iod.py index e7677a24..79b91d9c 100644 --- a/tests/layup/test_herget_iod.py +++ b/tests/layup/test_herget_iod.py @@ -99,7 +99,6 @@ def test_find_drho(tmpdir): jds = convert_tdb_date_to_julian_date(data["obsTime"]) sequence = _build_sequence(jds, sep_dt=90.0) - obs_1 = observations[0] r_e_1 = obs_1.observer_position rho_hat_1 = np.array(obs_1.rho_hat) @@ -123,7 +122,7 @@ def test_find_drho(tmpdir): observation.epoch = epochs[i] - ((rho_1) + (rho_n)) / (2 * SPEED_OF_LIGHT_AU_DAY) state_1[3:], _ = herget.find_velocity(t_1, t_n, state_1[:3], r_n, 0.001) - + ephem, _, _ = build_ephem_and_mus() print(dir(ephem)) sim = rebound.Simulation()