"""Orbit propagation."""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, TypeVar
import numpy as np
import pyarrow as pa
import quivr as qv
if TYPE_CHECKING:
from empyrean.ephemeris.sensitivity import StateSensitivities
from empyrean.orbits.thrust import ThrustParams
from empyrean.propagation.tagged_covariance import TaggedCovariances
from empyrean._convert import (
_COORD_TYPE_MAP,
AnyOrbits,
coordinates_to_arrays,
extract_joint,
extract_non_grav_covariance,
extract_photometry,
extract_srp,
int_to_frame,
joint_columns_from_result,
naif_to_origin,
non_grav_border_only,
validate_non_grav_marsden_only,
)
from empyrean.coordinates.coordinates import CartesianCoordinates
from empyrean.coordinates.covariance import (
CartesianCovariance as _CartesianCovariance,
)
from empyrean.coordinates.covariance import _CovarianceTable
from empyrean.coordinates.epoch import Epochs, _epochs_mjd_tdb
from empyrean.orbits.orbits import CartesianOrbits
from empyrean.propagation.config import (
_DATACLASS_TO_INT,
_FORCE_MODEL_TO_INT,
_UNCERTAINTY_METHOD_TO_INT,
Auto,
ForceModelTier,
GaussianMixture,
MonteCarlo,
PropagationConfig,
SigmaPoint,
UncertaintyMethod,
_uncertainty_method_params,
)
from empyrean.propagation.events import (
AtmosphericEntries,
AtmosphericExits,
CaptureEnds,
CaptureStarts,
CloseApproachEnds,
CloseApproachStarts,
CovarianceRegimeChanges,
EventConfig,
Events,
EventSummary,
Impacts,
Periapses,
PossibleImpacts,
ShadowEntries,
ShadowExits,
)
from empyrean.propagation.mixtures import build_mixture_chains
from empyrean.propagation.result import PropagationResult
from empyrean.propagation.tagged_covariance import _KIND_BY_CODE
# Quivr event sub-tables share the build helpers below; the helpers
# return the same concrete table subclass they are handed.
_EventTableT = TypeVar("_EventTableT", bound=qv.Table)
# ``CartesianCovariance`` is built dynamically by ``_make_covariance_class``,
# whose declared return type is the bare ``type``; that loses the injected
# ``from_matrix`` / ``from_kwargs`` constructors. Re-bind it through
# ``type[_CovarianceTable]`` (the Protocol describing exactly that
# dynamically-injected surface) so those constructors type-check at the call
# sites below. The runtime object is unchanged.
CartesianCovariance: type[_CovarianceTable] = _CartesianCovariance
UncertaintyMethodLike = UncertaintyMethod | SigmaPoint | MonteCarlo | GaussianMixture | Auto | str
[docs]
def propagate(
orbits: AnyOrbits,
epochs: Epochs,
config: PropagationConfig | None = None,
*,
# ── Sugar for quick, inline overrides ─────────────────────
# Any of these will populate a fresh PropagationConfig when `config`
# isn't supplied. Ignored when `config` is passed.
#
# Sugar mirrors top-level `PropagationConfig` fields only — knobs
# nested under `config.advanced` (`epsilon`, step bounds, loop
# guards) are deliberately not surfaced here. Reaching for those
# means you're tuning the integrator itself, which is a structured-
# config conversation, not an inline override.
force_model: ForceModelTier | str | None = None,
uncertainty_method: UncertaintyMethodLike | None = None,
num_threads: int | None = None,
events: EventConfig | None = None,
tagged_covariance: bool = False,
thrust_arcs: "Sequence[ThrustParams | None] | None" = None,
# Internal: a pre-built force-model handle
# (``empyrean._empyrean_rs.BuiltSystem``). When supplied, the forward
# model runs through the frozen handle (identity-guarded, never a
# silent rebuild) instead of assembling the force model one-shot.
# Set by :meth:`empyrean.BuiltSystem.propagate`; not part of the
# public call surface. The result is bit-identical to the one-shot on
# the matching key.
_builtsystem: Any = None,
) -> PropagationResult:
"""Propagate orbits to target epochs.
Parameters
----------
orbits : CartesianOrbits | CometaryOrbits | KeplerianOrbits | SphericalOrbits
Input orbits with optional covariance and non-gravitational
parameters.
epochs : Epochs
Target epochs as an :class:`~empyrean.types.Epochs` table,
converted to TDB internally. Build one with
``Epochs.from_mjd(values, scale="tdb")`` (or ``scale="utc"``); a
bare array is refused, as it carries no time scale.
config : PropagationConfig, optional
Full propagation configuration. Construct with
``PropagationConfig(force_model=..., uncertainty_method=...)``
etc. If omitted, one is built from the sugar kwargs below (or
defaults).
Other Parameters
----------------
force_model : ForceModelTier or str, optional
Quick override for ``config.force_model``. Ignored if ``config``
is given.
uncertainty_method : UncertaintyMethod | SigmaPoint | MonteCarlo | GaussianMixture | Auto | str
Optional quick override for ``config.uncertainty_method``. Accepts either
an enum / string (default parameters) or a parameterized
dataclass (:class:`SigmaPoint`, :class:`MonteCarlo`,
:class:`GaussianMixture`, :class:`Auto`). Ignored if ``config`` is
given. Pass :class:`Auto` (rather than the bare ``AUTO`` enum) to
tune the per-close-approach κ band edges and AGM knobs.
All six methods run in :func:`propagate`:
* ``FIRST_ORDER`` / ``SECOND_ORDER`` / ``AUTO`` attach the
STM-based state covariance (``AUTO`` escalating to a
second-order ellipsoid over close-approach windows).
* ``SIGMA_POINT`` reconstructs a genuine sample-based state
covariance (the second moment of the propagated canonical
2N+1 sigma-point set), read back tagged ``sigma_point``.
* ``MONTE_CARLO`` draws ``n_samples`` from the input covariance
(reproducibly, for a fixed ``seed``) and reports the
Monte-Carlo impact probability on the
:attr:`~empyrean.PossibleImpacts.ip_mc` column of any
possible-impact event. It does **not** reconstruct a per-epoch
state covariance — the propagated states carry no state
covariance under ``MONTE_CARLO`` (use ``SIGMA_POINT`` for a
sampled state covariance, or
:func:`~empyrean.compute_impact_probabilities` for the
full Monte-Carlo impact-probability workflow).
* ``GAUSSIAN_MIXTURE`` splits the input Gaussian into an adaptive
mixture at close approaches; its distinctive product is the
mixture-corrected impact probability. Away from encounters the
output-state covariance is the linear ``Φ·Σ·Φᵀ`` mapping (like
``SECOND_ORDER``), so for a well-determined object it reads back
very close to ``FIRST_ORDER`` (tagged ``linear``) — that is
expected, not a bug.
The mixture components themselves come back on
:attr:`PropagationResult.mixtures` (a
:class:`~empyrean.propagation.mixtures.MixtureChains` table),
so a consumer can evaluate ``sum_k w_k * N(x | mu_k, Sigma_k)``
directly at a retained close-approach epoch. Four limits apply
to what is retained, each a real property of the engine's
retention rather than a marshaling shortfall:
- **Depth-0 only.** Only the initial split is retained;
recursive AGM calls (depth > 0) are not captured.
- **Only CA epochs where AGM fired.** An orbit that never
triggered a split contributes no rows, not a one-component
chain.
- **Component covariance is the linear map.** Each component's
covariance is ``Phi Sigma_k Phi^T``; the second-order mean
correction is intentionally omitted.
- **Retained weights may sum to less than 1.** A sub-Gaussian
whose own sub-propagation missed the close approach (or
failed to integrate) contributes no component, and the
deficit is not recorded anywhere. Do not assume
``sum_k w_k == 1``; sum ``weight`` and check.
num_threads : int, optional
Threads for multi-orbit propagation. ``None`` (default) and
``0`` both use all available cores; ``n`` > 0 pins exactly
``n`` threads. Each orbit is integrated on a single thread;
parallelism is across orbits, not within a single trajectory.
events : EventConfig, optional
Event-detection toggles + body filter + dense-output cadence.
Override individual flags here without rebuilding a full
:class:`PropagationConfig`. See
:class:`~empyrean.EventConfig`.
tagged_covariance : bool, default False
When ``True``, also read back the provenance-tagged,
resolved-kind covariance at every output epoch (the honest
covariance that distinguishes a second-order close-approach
ellipsoid from the bare linear ``Φ Σ₀ Φᵀ`` mapping on the
states). The result's
:attr:`~empyrean.PropagationResult.tagged_covariance` table is
populated and
:meth:`~empyrean.PropagationResult.tagged_covariance_series`
becomes usable. Off by default — the readback recomputes the
resolved kind per orbit, so it isn't free.
thrust_arcs : sequence of ThrustParams or None, optional
Structured continuous-thrust / finite-burn input, one entry per
orbit and positionally aligned with ``orbits`` (pass ``None`` for
the gravity / non-grav-only orbits, or the whole argument
``None`` for a fully ballistic batch). Build each entry from
:class:`~empyrean.ThrustParams` /
:class:`~empyrean.ThrustArc` / a
:class:`~empyrean.orbits.thrust.SteeringLaw` variant. A non-empty
:attr:`~empyrean.ThrustParams.correction_covariances` triggers the
burn-sensitivity propagation whose solved segments surface in the
tagged-covariance
:attr:`~empyrean.TaggedCovariance.thrust_segments` (requires
``tagged_covariance=True``). Length or arc/correction mismatches
raise, never silently degrade.
Returns
-------
PropagationResult
Propagated states, detected events, and per-orbit state
sensitivity chains.
Notes
-----
Within each orbit, states come back in **ascending epoch order,
always**, regardless of the order the epochs were requested in.
Positional pairing against an ascending, duplicate-free request grid
is therefore exact; for any other request shape, join on the
result's ``epoch_mjd_tdb`` column.
The joint covariance travels in both directions. An input orbit
carrying cross terms — on ``non_grav.non_grav_cross`` (the 6x3
state-Marsden border) and ``wide_cross`` (everything else) — is
propagated against them rather than against its ``6x6`` alone, and
every output row carries the propagated cross terms in the same two
columns. This is what makes a chained propagation agree with the
single-leg answer: the state-parameter columns are non-zero even
when the input was block-diagonal, because propagation itself
generates that correlation, so a second leg handed only the ``6x6``
reports a tighter uncertainty than the first leg supports.
Cross terms are refused without the parameter blocks they are
conditioned on (the non-grav ``3x3``, the DT and AMRAT prior
variances). Propagation passes those through unchanged rather than
restating them on every output row, so when chaining legs by hand
take them from the orbit that started the chain. A row with no cross
terms is null, never zero — an absent correlation and a measured
zero correlation are different claims.
Raises
------
TypeError
``epochs`` is not an :class:`Epochs` table.
Examples
--------
Defaults (Standard force model, FirstOrder uncertainty):
>>> times = Epochs.from_mjd([60500.0, 60501.0], scale="tdb")
>>> result = empyrean.propagate(orbits, times)
With a config object:
>>> cfg = PropagationConfig(
... force_model=ForceModelTier.STANDARD,
... uncertainty_method=SigmaPoint(),
... num_threads=8,
... )
>>> result = empyrean.propagate(orbits, times, cfg)
With inline kwargs (sugar):
>>> result = empyrean.propagate(orbits, times, force_model="standard")
"""
from empyrean._empyrean_rs import _propagate
# Refused here, at the entry point, before any orbit is marshaled —
# a bare array is a caller error, not a modelling result.
times_mjd_tdb = _epochs_mjd_tdb(epochs, "propagate() epochs")
# ── Assemble PropagationConfig ────────────────────────────
if config is None:
# PropagationConfig.force_model is typed as ForceModelTier, while the
# `force_model` sugar additionally accepts a str. Resolve a str tier
# to its ForceModelTier member here so the constructed config carries
# the precise enum type; the case-insensitive lookup mirrors the
# downstream `_FORCE_MODEL_TO_INT` mapping.
force_model_tier: ForceModelTier
if force_model is None:
force_model_tier = ForceModelTier.STANDARD
elif isinstance(force_model, str):
force_model_tier = ForceModelTier(force_model.lower())
else:
force_model_tier = force_model
config = PropagationConfig(
force_model=force_model_tier,
uncertainty_method=(
uncertainty_method
if uncertainty_method is not None
else UncertaintyMethod.FIRST_ORDER
),
num_threads=num_threads,
events=events if events is not None else EventConfig(),
)
elif any(v is not None for v in (force_model, uncertainty_method, num_threads, events)):
raise TypeError(
"propagate(): pass either `config` or the sugar kwargs "
"(force_model / uncertainty_method / num_threads / events), "
"not both"
)
# Pull fields off the config from here on
force_model = config.force_model
uncertainty_method = config.uncertainty_method
num_threads = config.num_threads
epsilon = config.epsilon
events = config.events
if events is None:
events = EventConfig()
# ── Extract coordinate arrays from orbits ────────────────
coords = orbits.coordinates
coord_type = type(coords)
if coord_type not in _COORD_TYPE_MAP:
raise TypeError(f"unsupported coordinate type: {coord_type}")
(
epochs_arr,
elements_arr,
covariances_arr,
has_cov_arr,
representations_arr,
frames_arr,
origins_arr,
) = coordinates_to_arrays(coords)
# IDs
orbit_ids = orbits.orbit_id.to_pylist()
if orbits.object_id is not None:
object_ids = [s if s else "" for s in orbits.object_id.to_pylist()]
else:
object_ids = [""] * len(orbits)
# Non-grav parameters
n = len(orbits)
# NonGravParams is Marsden-only; reject a stray model='srp' / cr before
# marshaling (SRP rides its own slot, extracted below).
validate_non_grav_marsden_only(orbits)
has_srp, srp_amrat, srp_cr, srp_amrat_variance = extract_srp(orbits)
non_grav_dts: np.ndarray | None = None
non_grav_dt_variances: np.ndarray | None = None
# g(r) Marsden–Sekanina exponents. Passed only when a non-default g(r)
# is present (any non-zero α/r0/m/n/k); all-zero is the inverse-square
# asteroid default that the engine applies without a marshal.
ng_alphas: np.ndarray | None = None
ng_r0s: np.ndarray | None = None
ng_ms: np.ndarray | None = None
ng_ns: np.ndarray | None = None
ng_ks: np.ndarray | None = None
if orbits.non_grav is not None:
ng = orbits.non_grav
# Handle nullable columns: fill None with 0.0
a1s = np.asarray(ng.a1.to_numpy(zero_copy_only=False), dtype=np.float64)
a2s = np.asarray(ng.a2.to_numpy(zero_copy_only=False), dtype=np.float64)
a3s = np.asarray(ng.a3.to_numpy(zero_copy_only=False), dtype=np.float64)
a1s = np.nan_to_num(a1s, nan=0.0)
a2s = np.nan_to_num(a2s, nan=0.0)
a3s = np.nan_to_num(a3s, nan=0.0)
# SBDB non-grav DT (days). NaN entries → no delay; pass the
# whole array only when at least one row populated, so the
# asteroid-only case avoids an FFI marshal.
dt_col = np.asarray(ng.dt.to_numpy(zero_copy_only=False), dtype=np.float64)
if np.isfinite(dt_col).any():
non_grav_dts = dt_col
# DT prior variance — opens the DT column in a StateAndNonGravAndDT
# solve. Gated like non_grav_dts (finite positive) so the no-prior
# asteroid-only case avoids an FFI marshal.
dtv_col = np.asarray(ng.dt_variance.to_numpy(zero_copy_only=False), dtype=np.float64)
if (np.isfinite(dtv_col) & (dtv_col > 0.0)).any():
non_grav_dt_variances = dtv_col
# g(r) exponents — carry the comet Marsden–Sekanina g(r) so a fitted
# or SBDB comet orbit isn't silently propagated with inverse-square.
alpha_col = np.nan_to_num(
np.asarray(ng.alpha.to_numpy(zero_copy_only=False), dtype=np.float64), nan=0.0
)
r0_col = np.nan_to_num(
np.asarray(ng.r0.to_numpy(zero_copy_only=False), dtype=np.float64), nan=0.0
)
m_col = np.nan_to_num(
np.asarray(ng.m.to_numpy(zero_copy_only=False), dtype=np.float64), nan=0.0
)
n_col = np.nan_to_num(
np.asarray(ng.n.to_numpy(zero_copy_only=False), dtype=np.float64), nan=0.0
)
k_col = np.nan_to_num(
np.asarray(ng.k.to_numpy(zero_copy_only=False), dtype=np.float64), nan=0.0
)
if (
(alpha_col != 0).any()
or (r0_col != 0).any()
or (m_col != 0).any()
or (n_col != 0).any()
or (k_col != 0).any()
):
ng_alphas = alpha_col
ng_r0s = r0_col
ng_ms = m_col
ng_ns = n_col
ng_ks = k_col
else:
a1s = np.zeros(n, dtype=np.float64)
a2s = np.zeros(n, dtype=np.float64)
a3s = np.zeros(n, dtype=np.float64)
# Photometric parameters, including the optional 3x3 covariance
# over (H, slope1, slope2). Gated like the non-grav covariance
# below so the common no-covariance case skips the FFI marshal.
(
phot_h,
phot_slope1,
phot_slope2,
phot_model,
has_phot_cov_arr,
phot_cov_arr,
) = extract_photometry(orbits)
has_phot_cov: np.ndarray | None = has_phot_cov_arr if has_phot_cov_arr.any() else None
phot_cov: np.ndarray | None = phot_cov_arr if has_phot_cov_arr.any() else None
# Fitted non-grav covariance — passed through only when a row carries one
# (mirrors the OD output path) so a StateAndNonGrav-fitted orbit re-fed
# into propagate keeps its prior. Gated like the other optional non-grav
# arrays so the common no-cov case skips the FFI marshal.
has_ng_cov_arr, ng_cov_arr = extract_non_grav_covariance(orbits)
has_non_grav_cov: np.ndarray | None = has_ng_cov_arr if has_ng_cov_arr.any() else None
non_grav_cov: np.ndarray | None = ng_cov_arr if has_ng_cov_arr.any() else None
# ── Map force model to int ───────────────────────────────
if isinstance(force_model, str):
fm_int = _FORCE_MODEL_TO_INT.get(force_model.lower())
if fm_int is None:
raise ValueError(f"unknown force model: {force_model}")
elif isinstance(force_model, ForceModelTier):
fm_int = _FORCE_MODEL_TO_INT[force_model]
elif isinstance(force_model, int):
fm_int = force_model
else:
raise TypeError(f"force_model must be ForceModelTier, str, or int, got {type(force_model)}")
# ── Map uncertainty method to int + extract params ─────────
#
# Three input shapes:
# 1. str / UncertaintyMethod enum → default parameters
# 2. SigmaPoint / MonteCarlo / GaussianMixture dataclass → method + params
# 3. int (legacy) → default parameters
#
# The flat parameter slots come from the one shared lowering helper, which
# the ephemeris and impact entry points also use, so the three paths cannot
# drift apart. The tag dispatch below still owns type validation.
(
sigma_n_sigma,
sigma_samples_per_plane,
mc_n_samples,
mc_seed,
gm_threshold,
gm_max_depth,
gm_components_per_split,
auto_threshold_first,
auto_threshold_mixture,
auto_threshold_ip_skip,
auto_gmm_max_depth,
auto_gmm_components_per_split,
) = _uncertainty_method_params(uncertainty_method)
if isinstance(uncertainty_method, (SigmaPoint, MonteCarlo, GaussianMixture, Auto)):
um_int = _DATACLASS_TO_INT[type(uncertainty_method)]
elif isinstance(uncertainty_method, str):
um_int_opt = _UNCERTAINTY_METHOD_TO_INT.get(uncertainty_method.lower())
if um_int_opt is None:
raise ValueError(f"unknown uncertainty method: {uncertainty_method}")
um_int = um_int_opt
elif isinstance(uncertainty_method, UncertaintyMethod):
um_int = _UNCERTAINTY_METHOD_TO_INT[uncertainty_method]
elif isinstance(uncertainty_method, int):
um_int = uncertainty_method
else:
raise TypeError(
"uncertainty_method must be UncertaintyMethod, a SigmaPoint / "
"MonteCarlo / GaussianMixture / Auto dataclass, str, or int; got "
f"{type(uncertainty_method).__name__}"
)
# ── Structured thrust input ──────────────────────────────
# One ThrustParams (or None) per orbit, positionally aligned with the
# batch. The binding reconstructs each into a wrapper ThrustParams and
# attaches it per orbit; None entries stay gravity / non-grav only.
thrust_arg: list[ThrustParams | None] | None = None
if thrust_arcs is not None:
thrust_arg = list(thrust_arcs)
if len(thrust_arg) != n:
raise ValueError(
f"thrust_arcs must have one entry per orbit (got {len(thrust_arg)} for {n} orbits)"
)
# ── Call Rust ─────────────────────────────────────────────
# Thread the full nested PropagationConfig as a single dict so that
# advanced fields (events.dense_output, diagnostics.*, advanced.*,
# excluded_perturbers, max_propagation_time_days, etc.) are honored
# without growing _propagate's flat-arg signature.
result = _propagate(
orbit_ids,
object_ids,
epochs_arr,
elements_arr,
covariances_arr,
has_cov_arr,
representations_arr,
frames_arr,
origins_arr,
times_mjd_tdb,
fm_int,
um_int,
a1s,
a2s,
a3s,
phot_h,
phot_slope1,
phot_model,
num_threads=num_threads,
epsilon=epsilon,
thrust_arcs=thrust_arg,
non_grav_dts=non_grav_dts,
non_grav_dt_variances=non_grav_dt_variances,
has_srp=has_srp,
srp_amrat=srp_amrat,
srp_cr=srp_cr,
srp_amrat_variance=srp_amrat_variance,
has_non_grav_cov=has_non_grav_cov,
non_grav_cov=non_grav_cov,
phot_slope2=phot_slope2,
has_phot_cov=has_phot_cov,
phot_cov=phot_cov,
# The joint's off-diagonal terms, when the input carries them.
# A leg fed its predecessor's joint propagates the covariance
# that leg actually computed; fed the 6x6 alone it reports a
# tighter uncertainty than the first leg supports. No keys at
# all when no orbit carries cross terms.
**extract_joint(orbits),
ng_alphas=ng_alphas,
ng_r0s=ng_r0s,
ng_ms=ng_ms,
ng_ns=ng_ns,
ng_ks=ng_ks,
# GaussianMixture (AGM) parameters — honored by the binding when
# uncertainty_method resolves to GAUSSIAN_MIXTURE (tag 5); benign
# defaults otherwise.
gm_threshold=gm_threshold,
gm_max_depth=gm_max_depth,
gm_components_per_split=gm_components_per_split,
sigma_n_sigma=sigma_n_sigma,
sigma_samples_per_plane=sigma_samples_per_plane,
mc_n_samples=mc_n_samples,
mc_seed=mc_seed,
# AUTO's caller-tunable κ band edges + AGM knobs. A
# default-constructed Auto() / "auto" lowers to the engine
# defaults, so the common path is unchanged; a parameterized
# Auto(...) takes effect instead of silently using defaults.
auto_threshold_first=auto_threshold_first,
auto_threshold_mixture=auto_threshold_mixture,
auto_threshold_ip_skip=auto_threshold_ip_skip,
auto_gmm_max_depth=auto_gmm_max_depth,
auto_gmm_components_per_split=auto_gmm_components_per_split,
propagation_config_dict=config._to_wire_dict(),
with_tagged_covariance=tagged_covariance,
builtsystem=_builtsystem,
)
# ── Build CartesianOrbits from result ─────────────────────
states = _build_cartesian_orbits(result)
detected_events = _build_events(result)
# ── Build per-orbit StateSensitivity chains ───────────────
sensitivity = _build_state_sensitivity(result)
# ── Build provenance-tagged covariance table (opt-in) ─────
tagged = _build_tagged_covariance(result) if tagged_covariance else None
# ── Retained AGM mixture components (always on) ───────────
# Free: the components already crossed the boundary with the result.
# Empty table when nothing split.
mixtures = build_mixture_chains(result)
return PropagationResult(
states=states,
events=detected_events,
sensitivity=sensitivity,
tagged_covariance=tagged,
mixtures=mixtures,
)
def _build_tagged_covariance(
result: dict[str, Any],
) -> "TaggedCovariances | None":
"""Build a :class:`TaggedCovariances` table from the pyo3 result.
The per-``(orbit, epoch)`` tagged arrays in ``result`` are aligned
1:1 with the states, so the orbit / object ids and epochs are reused
straight from the flat state columns. Returns ``None`` when the
extension emitted no tagged sub-dict.
"""
from empyrean.propagation.tagged_covariance import build_tagged_covariances
if "tagged_covariance" not in result:
return None
orbit_ids: list[str] = list(result["orbit_ids"])
object_ids: list[str | None] = [s if s else None for s in result["object_ids"]]
epochs_arr = np.asarray(result["epochs"], dtype=np.float64)
return build_tagged_covariances(result, orbit_ids, object_ids, epochs_arr)
def _build_state_sensitivity(result: dict[str, Any]) -> "StateSensitivities | None":
"""Build a :class:`StateSensitivities` table from the pyo3 result.
Flattens the per-row (6, 6) STM and (6, 6, 6) STT arrays into the
row-major lists the table expects (length 36 / 216 per row), with
``None`` per row when ``has_stm`` / ``has_stt`` is false on that row.
Returns ``None`` if no row has an STM (sample-based methods or
FirstOrder without input covariance).
"""
from empyrean.ephemeris.sensitivity import StateSensitivities
stms = np.asarray(result["stms"]) if "stms" in result else None
has_stm = np.asarray(result["has_stm"]) if "has_stm" in result else None
stts = np.asarray(result["stts"]) if "stts" in result else None
has_stt = np.asarray(result["has_stt"]) if "has_stt" in result else None
if has_stm is None or not bool(has_stm.any()):
return None
orbit_ids = list(result["orbit_ids"])
object_ids = [s if s else None for s in result["object_ids"]]
epochs_arr = np.asarray(result["epochs"], dtype=np.float64)
n = len(orbit_ids)
# Flatten STMs to length-36 lists; null per-row where has_stm is false.
if stms is not None:
stms_flat = stms.reshape(n, 36).tolist()
stm_col = [
row if (has_stm is None or bool(has_stm[i])) else None
for i, row in enumerate(stms_flat)
]
else:
stm_col = [None] * n
# STTs: only populate rows where has_stt is true; null elsewhere.
if stts is not None and has_stt is not None and bool(has_stt.any()):
stts_flat = stts.reshape(n, 216).tolist()
stt_col = [row if bool(has_stt[i]) else None for i, row in enumerate(stts_flat)]
else:
stt_col = [None] * n
# Per-row resolved covariance kind (linear / second_order / …).
rk_codes = result.get("resolved_kind")
if rk_codes is not None:
resolved_kind_col: list[str | None] = [
_KIND_BY_CODE[int(c)].value if int(c) >= 0 else None for c in np.asarray(rk_codes)
]
else:
resolved_kind_col = [None] * n
return StateSensitivities.from_kwargs(
orbit_id=orbit_ids,
object_id=object_ids,
epoch_mjd_tdb=epochs_arr,
stm=stm_col,
stt=stt_col,
resolved_kind=resolved_kind_col,
)
def _build_cartesian_orbits(result: dict[str, Any]) -> CartesianOrbits:
"""Build CartesianOrbits from the Rust result dict."""
m = len(result["epochs"])
out_epochs = np.asarray(result["epochs"])
out_x = np.asarray(result["x"])
out_y = np.asarray(result["y"])
out_z = np.asarray(result["z"])
out_vx = np.asarray(result["vx"])
out_vy = np.asarray(result["vy"])
out_vz = np.asarray(result["vz"])
out_frames = np.asarray(result["frames"])
out_origins = np.asarray(result["origins"])
out_covariances = np.asarray(result["covariances"])
out_has_cov = np.asarray(result["has_covariance"])
out_orbit_ids = result["orbit_ids"]
out_object_ids = result["object_ids"]
# Build frame attribute (all rows should have the same frame)
from empyrean.coordinates.enums import Frame as FrameEnum
frame = int_to_frame(int(out_frames[0])) if m > 0 else FrameEnum.ICRF
# Build origin column
origin_strs = [naif_to_origin(int(o)) for o in out_origins]
# Build covariance
if out_has_cov.any():
cov = CartesianCovariance.from_matrix(out_covariances)
else:
cov = None
frame_value = frame.value if hasattr(frame, "value") else frame
if cov is not None:
cart_coords = CartesianCoordinates.from_kwargs(
epoch=out_epochs,
x=out_x,
y=out_y,
z=out_z,
vx=out_vx,
vy=out_vy,
vz=out_vz,
frame=frame_value,
origin=origin_strs,
covariance=cov,
)
else:
cart_coords = CartesianCoordinates.from_kwargs(
epoch=out_epochs,
x=out_x,
y=out_y,
z=out_z,
vx=out_vx,
vy=out_vy,
vz=out_vz,
frame=frame_value,
origin=origin_strs,
)
# Convert empty strings to None for nullable object_id
object_id_list = [s if s else None for s in out_object_ids]
orbits_kwargs: dict[str, Any] = {
"orbit_id": out_orbit_ids,
"object_id": object_id_list,
"coordinates": cart_coords,
}
# The propagated joint's off-diagonal blocks. `coordinates.covariance`
# above is only the state block: propagation generates state-parameter
# correlation even from a block-diagonal input, so a chained run built
# on the 6x6 alone is a tighter claim than this leg supports. Attached
# only when the engine produced cross terms — a gravity-only or
# parameter-free orbit leaves both columns unset rather than carrying
# rows of zeros.
border, wide = joint_columns_from_result(result, m)
if border is not None:
orbits_kwargs["non_grav"] = non_grav_border_only(border)
if wide is not None:
orbits_kwargs["wide_cross"] = wide
return CartesianOrbits.from_kwargs(**orbits_kwargs)
def _nullable_float(values: np.ndarray) -> pa.Array | np.ndarray:
"""Convert a list of floats to a pyarrow array with NaN -> null."""
arr = np.asarray(values, dtype=np.float64)
mask = np.isnan(arr)
if mask.any():
return pa.array(arr.tolist(), type=pa.float64(), mask=mask)
return arr
def _nullable_str_list(values: Sequence[str | None]) -> list[str | None]:
"""Convert a list of strings to a list with empty string -> None."""
return [s if s else None for s in values]
def _build_events(result: dict[str, Any]) -> Events:
"""Build Events container from the Rust result dict.
The Rust extension returns a single flat events sub-dict with an
``event_types`` discriminator column carrying ``distance_au`` /
``distance_km`` / ``relative_velocity_au_day``. We dispatch each row
into the appropriate per-subtype quivr table. Subtype-specific fields
that the flat schema *does* carry are read across (e.g. the atmospheric
entry altitude rides in ``distance_km``); fields it does not carry
(latitude / longitude on impacts, jacobi constants on captures,
illumination on shadow events) are filled with NaN / null.
"""
ev = result.get("events")
if ev is None or len(ev["event_types"]) == 0:
return Events(
summary=EventSummary.empty(),
close_approach_starts=CloseApproachStarts.empty(),
close_approach_ends=CloseApproachEnds.empty(),
periapses=Periapses.empty(),
impacts=Impacts.empty(),
possible_impacts=PossibleImpacts.empty(),
atmospheric_entries=AtmosphericEntries.empty(),
atmospheric_exits=AtmosphericExits.empty(),
capture_starts=CaptureStarts.empty(),
capture_ends=CaptureEnds.empty(),
shadow_entries=ShadowEntries.empty(),
shadow_exits=ShadowExits.empty(),
covariance_regime_changes=CovarianceRegimeChanges.empty(),
)
orbit_ids = list(ev["orbit_ids"])
object_ids = list(ev["object_ids"])
event_types = list(ev["event_types"])
bodies = list(ev["bodies"])
epochs = np.asarray(ev["epochs"], dtype=np.float64)
distance_au = np.asarray(ev["distance_au"], dtype=np.float64)
distance_km = np.asarray(ev["distance_km"], dtype=np.float64)
rel_v = np.asarray(ev["relative_velocity_au_day"], dtype=np.float64)
# Subtype payload columns the C ABI now carries (NaN / -1 sentinels on
# rows the field doesn't apply to). `.get` with a sentinel fallback so
# older result dicts (pre-extension) still load.
n_all = len(event_types)
def _ev_f(key: str) -> np.ndarray:
return np.asarray(ev.get(key, np.full(n_all, np.nan)), dtype=np.float64)
two_body_energy = _ev_f("two_body_energy")
jacobi = _ev_f("jacobi_constant")
jacobi_sigma = _ev_f("jacobi_constant_sigma")
jacobi_l1 = _ev_f("jacobi_constant_l1")
jacobi_l2 = _ev_f("jacobi_constant_l2")
n_periapses = np.asarray(ev.get("n_periapses", np.full(n_all, -1)), dtype=np.int32)
impact_lat = _ev_f("impact_latitude_deg")
impact_lon = _ev_f("impact_longitude_deg")
impact_alt = _ev_f("impact_altitude_km")
shadow_fraction = _ev_f("shadow_fraction")
illumination = _ev_f("illumination")
relative_x = _ev_f("relative_x")
relative_y = _ev_f("relative_y")
relative_z = _ev_f("relative_z")
relative_vx = _ev_f("relative_vx")
relative_vy = _ev_f("relative_vy")
relative_vz = _ev_f("relative_vz")
pi_effective_radius_au = _ev_f("effective_radius_au")
pi_effective_radius_km = _ev_f("effective_radius_km")
pi_sigma_distance_au = _ev_f("sigma_distance_au")
pi_ip_linear = _ev_f("ip_linear")
pi_ip_second_order = _ev_f("ip_second_order")
pi_nonlinearity = _ev_f("nonlinearity")
pi_ip_agm = _ev_f("ip_agm")
pi_ip_mc = _ev_f("ip_mc")
previous_kind = np.asarray(ev.get("previous_kind", np.full(n_all, -1)), dtype=np.int64)
regime_resolved_kind = np.asarray(
ev.get("regime_resolved_kind", np.full(n_all, -1)), dtype=np.int64
)
kappa = _ev_f("kappa")
threshold_below = _ev_f("threshold_below")
threshold_above = _ev_f("threshold_above")
# Cross-cutting summary table — every event lands here.
summary = EventSummary.from_kwargs(
orbit_id=orbit_ids,
object_id=_nullable_str_list(object_ids),
event_type=event_types,
body=bodies,
epoch=epochs,
)
# Filter helpers ---------------------------------------------------
def _idx(tag: str) -> list[int]:
return [i for i, t in enumerate(event_types) if t == tag]
def _str(values: Sequence[str], idx: list[int]) -> list[str]:
return [values[i] for i in idx]
def _str_opt(values: Sequence[str], idx: list[int]) -> list[str | None]:
return [values[i] if values[i] else None for i in idx]
def _arr(values: np.ndarray, idx: list[int]) -> np.ndarray:
return values[idx] if len(idx) > 0 else np.zeros(0, dtype=values.dtype)
def _common(tag: str, cls: type[_EventTableT]) -> _EventTableT:
idx = _idx(tag)
if not idx:
return cls.empty()
return cls.from_kwargs(
orbit_id=_str(orbit_ids, idx),
object_id=_str_opt(object_ids, idx),
body=_str(bodies, idx),
epoch=_arr(epochs, idx),
distance_au=_arr(distance_au, idx),
distance_km=_arr(distance_km, idx),
)
close_approach_starts = _common("close_approach_start", CloseApproachStarts)
close_approach_ends = _common("close_approach_end", CloseApproachEnds)
# Periapses carry relative state vectors wired through the C ABI.
per_idx = _idx("periapsis")
if per_idx:
periapses = Periapses.from_kwargs(
orbit_id=_str(orbit_ids, per_idx),
object_id=_str_opt(object_ids, per_idx),
body=_str(bodies, per_idx),
epoch=_arr(epochs, per_idx),
distance_au=_arr(distance_au, per_idx),
distance_km=_arr(distance_km, per_idx),
relative_velocity_au_day=_arr(rel_v, per_idx),
relative_x=_arr(relative_x, per_idx),
relative_y=_arr(relative_y, per_idx),
relative_z=_arr(relative_z, per_idx),
relative_vx=_arr(relative_vx, per_idx),
relative_vy=_arr(relative_vy, per_idx),
relative_vz=_arr(relative_vz, per_idx),
)
else:
periapses = Periapses.empty()
# Impacts: planetodetic surface-intercept lat/lon/alt now carried by
# the flat schema (NaN -> null where the impact geometry was
# unresolved).
imp_idx = _idx("impact")
if imp_idx:
impacts = Impacts.from_kwargs(
orbit_id=_str(orbit_ids, imp_idx),
object_id=_str_opt(object_ids, imp_idx),
body=_str(bodies, imp_idx),
epoch=_arr(epochs, imp_idx),
latitude_deg=_nullable_float(_arr(impact_lat, imp_idx)),
longitude_deg=_nullable_float(_arr(impact_lon, imp_idx)),
altitude_km=_nullable_float(_arr(impact_alt, imp_idx)),
# Body-relative impact speed carried by the flat schema
# (same column the periapses / atmospheric entries read);
# NaN -> null where the ABI did not resolve a velocity.
relative_velocity_au_day=_nullable_float(_arr(rel_v, imp_idx)),
)
else:
impacts = Impacts.empty()
# Possible impacts: probabilistic fields not in the flat schema.
pi_idx = _idx("possible_impact")
if pi_idx:
# PossibleImpact probability payload is wired through the C ABI.
# ip_second_order / nonlinearity / ip_agm / ip_mc are absent
# (NaN on the flat row) unless the matching uncertainty method
# ran; they are the nullable columns and marshal to Arrow **null**
# — matching compute_impact_probabilities' encoding, so
# `row.ip_agm is None` is a true test for "AGM did not run"
# rather than a false positive on every row. ip_linear /
# effective_radius / sigma_distance are always populated and stay
# dense.
possible_impacts = PossibleImpacts.from_kwargs(
orbit_id=_str(orbit_ids, pi_idx),
object_id=_str_opt(object_ids, pi_idx),
body=_str(bodies, pi_idx),
epoch=_arr(epochs, pi_idx),
miss_distance_au=_arr(distance_au, pi_idx),
miss_distance_km=_arr(distance_km, pi_idx),
effective_radius_au=_arr(pi_effective_radius_au, pi_idx),
effective_radius_km=_arr(pi_effective_radius_km, pi_idx),
sigma_distance_au=_arr(pi_sigma_distance_au, pi_idx),
ip_linear=_arr(pi_ip_linear, pi_idx),
relative_velocity_au_day=_arr(rel_v, pi_idx),
ip_second_order=_nullable_float(_arr(pi_ip_second_order, pi_idx)),
nonlinearity=_nullable_float(_arr(pi_nonlinearity, pi_idx)),
ip_agm=_nullable_float(_arr(pi_ip_agm, pi_idx)),
ip_mc=_nullable_float(_arr(pi_ip_mc, pi_idx)),
)
else:
possible_impacts = PossibleImpacts.empty()
# Atmospheric entries: distance_au is the body-CENTER crossing
# distance (the Karman radius), NOT an altitude. The true altitude
# above the reference ellipsoid and the surface lat/lon come from the
# planetodetic ground track (impact_altitude_km / impact_*_deg on the
# flat event), NaN -> null when the ground track is unresolved. The
# entry speed rides in relative_velocity_au_day.
ae_idx = _idx("atmospheric_entry")
if ae_idx:
atmospheric_entries = AtmosphericEntries.from_kwargs(
orbit_id=_str(orbit_ids, ae_idx),
object_id=_str_opt(object_ids, ae_idx),
body=_str(bodies, ae_idx),
epoch=_arr(epochs, ae_idx),
distance_au=_arr(distance_au, ae_idx),
altitude_km=_nullable_float(_arr(impact_alt, ae_idx)),
relative_velocity_au_day=_nullable_float(_arr(rel_v, ae_idx)),
latitude_deg=_nullable_float(_arr(impact_lat, ae_idx)),
longitude_deg=_nullable_float(_arr(impact_lon, ae_idx)),
)
else:
atmospheric_entries = AtmosphericEntries.empty()
def _simple_entry_exit(tag: str, cls: type[_EventTableT]) -> _EventTableT:
idx = _idx(tag)
if not idx:
return cls.empty()
return cls.from_kwargs(
orbit_id=_str(orbit_ids, idx),
object_id=_str_opt(object_ids, idx),
body=_str(bodies, idx),
epoch=_arr(epochs, idx),
distance_au=_arr(distance_au, idx),
)
atmospheric_exits = _simple_entry_exit("atmospheric_exit", AtmosphericExits)
# Capture starts/ends: two-body energy + CR3BP Jacobi constants (and
# the escape periapsis count) now carried by the flat schema.
def _capture(tag: str, cls: type[_EventTableT], with_n_periapses: bool = False) -> _EventTableT:
idx = _idx(tag)
if not idx:
return cls.empty()
kwargs: dict[str, Any] = {
"orbit_id": _str(orbit_ids, idx),
"object_id": _str_opt(object_ids, idx),
"body": _str(bodies, idx),
"epoch": _arr(epochs, idx),
"distance_au": _arr(distance_au, idx),
"distance_km": _arr(distance_km, idx),
"relative_velocity_au_day": _arr(rel_v, idx),
"two_body_energy": _arr(two_body_energy, idx),
"jacobi_constant": _nullable_float(_arr(jacobi, idx)),
"jacobi_constant_sigma": _nullable_float(_arr(jacobi_sigma, idx)),
"jacobi_constant_l1": _nullable_float(_arr(jacobi_l1, idx)),
"jacobi_constant_l2": _nullable_float(_arr(jacobi_l2, idx)),
}
if with_n_periapses:
kwargs["n_periapses"] = _arr(n_periapses, idx)
return cls.from_kwargs(**kwargs)
capture_starts = _capture("capture_start", CaptureStarts)
capture_ends = _capture("capture_end", CaptureEnds, with_n_periapses=True)
# Shadow events: shadow_fraction / illumination are carried through the
# C ABI.
def _shadow(tag: str, cls: type[_EventTableT]) -> _EventTableT:
idx = _idx(tag)
if not idx:
return cls.empty()
return cls.from_kwargs(
orbit_id=_str(orbit_ids, idx),
object_id=_str_opt(object_ids, idx),
body=_str(bodies, idx),
epoch=_arr(epochs, idx),
shadow_fraction=_arr(shadow_fraction, idx),
illumination=_arr(illumination, idx),
)
shadow_entries = _shadow("shadow_entry", ShadowEntries)
shadow_exits = _shadow("shadow_exit", ShadowExits)
# Covariance-regime-change events: the UncertaintyMethod.AUTO audit
# trail (linear <-> second-order transitions at CA-window boundaries).
# Kind codes: -1 = not applicable, else EMPYREAN_COVARIANCE_KIND_*.
def _kind_label(code: int) -> str | None:
return _KIND_BY_CODE[int(code)].value if int(code) >= 0 else None
crc_idx = _idx("covariance_regime_change")
if crc_idx:
covariance_regime_changes = CovarianceRegimeChanges.from_kwargs(
orbit_id=_str(orbit_ids, crc_idx),
object_id=_str_opt(object_ids, crc_idx),
body=_str_opt(bodies, crc_idx),
epoch=_arr(epochs, crc_idx),
previous_kind=[_kind_label(c) for c in _arr(previous_kind, crc_idx)],
resolved_kind=[_kind_label(c) for c in _arr(regime_resolved_kind, crc_idx)],
kappa=_nullable_float(_arr(kappa, crc_idx)),
threshold_below=_nullable_float(_arr(threshold_below, crc_idx)),
threshold_above=_nullable_float(_arr(threshold_above, crc_idx)),
)
else:
covariance_regime_changes = CovarianceRegimeChanges.empty()
return Events(
summary=summary,
close_approach_starts=close_approach_starts,
close_approach_ends=close_approach_ends,
periapses=periapses,
impacts=impacts,
possible_impacts=possible_impacts,
atmospheric_entries=atmospheric_entries,
atmospheric_exits=atmospheric_exits,
capture_starts=capture_starts,
capture_ends=capture_ends,
shadow_entries=shadow_entries,
shadow_exits=shadow_exits,
covariance_regime_changes=covariance_regime_changes,
)