orblet 0.1.0__py3-none-any.whl

orblet/__init__.py

@@ -0,0 +1,207 @@
+"""orblet — a standalone orbit fitter: arrays in, orbital elements out.
+
+No mission, no data files, no network. orblet is a CONSUMER of data and never
+a producer: you hand it epochs, measurements and uncertainties, and it hands
+back a model, a likelihood, a solution. Where those arrays came from, and
+whether you are allowed to have them, is not orblet's business.
+
+Work the open way
+-----------------
+orblet ships ATOMS, not a black box. A Kepler solver, forward models, design
+matrices, the two linear cores, proposed likelihoods, default priors, a
+frequency search, the element conversions, the plots. You assemble the
+optimiser or sampler you actually want and plug your own target function into
+it; the default recipes are written out in the manuals and the example
+notebooks rather than hidden inside a wrapper.
+
+That is a deliberate trade. A one-line `fit()` is quicker to call and harder
+to trust: when it disagrees with your expectation you cannot see which step
+disagreed. Here every step is a function you can call, print and plot on its
+own.
+
+The shape of the thing
+----------------------
+    model       the forward models: elements and epochs -> observables
+    design      the design matrices, one column per linear parameter
+    solve       the linear cores — fix the geometry, solve the rest exactly
+    search      the frequency scan and the periodograms
+    likelihood  proposed likelihoods (Gaussian with jitter), not mandates
+    priors      prior classes, defaults, and the spec parser
+    prepare     epoch and time-scale handling
+    sampling    MCMC atoms: chains, stacking, R-hat, ESS, dead chains
+    elements    conversions between parameter conventions
+    residuals   the 5-parameter residual routine
+    parallax    observer position -> on-sky parallax factors
+    interpret   what a fitted orbit implies: masses, limits, flux ratios
+    plotting    the orbit, the corner, the sky plane
+    simulate    synthetic data with known truth, for testing and teaching
+
+Why the split matters: the geometry of a Keplerian orbit is non-linear, but
+once it is fixed the amplitudes enter LINEARLY and have a closed-form best
+fit. `solve` is that closed form. A sampler over a handful of non-linear
+parameters, with `solve` as its inner step, beats one over all of them.
+
+Import cost
+-----------
+Nothing heavy is imported when you `import orblet`. Every name below is
+resolved on first attribute access (PEP 562), so scipy, emcee and matplotlib
+arrive only if you touch something that needs them.
+
+**This package is under construction.** It is being assembled in place inside
+the gaia-orbits monorepo, one dependency layer at a time, before being
+extracted to its own repository. The modules arrive by relocation, unchanged.
+
+State of the move:  `docs/plans/2026-09-23-orblet-b1-chunks.md`
+Plan of record:     `docs/plans/2026-09-23-orblet-extraction-plan.md`
+Machine-checked:    `tests/_orblet_move_map.py`
+"""
+
+from __future__ import annotations
+
+import importlib
+
+# public name -> (module, attribute). One mechanism, one table, grouped by
+# what the name is FOR rather than by which file it happens to live in.
+#: The package version. The ONLY place it is written: the build reads it from
+#: here (``[tool.setuptools.dynamic]`` in orblet's pyproject), and the packaging
+#: test pins that the installed metadata agrees. A plain literal on purpose —
+#: setuptools parses it from the source without importing the package.
+__version__ = "0.1.0"
+
+_LAZY_EXPORTS: dict[str, tuple[str, str]] = {
+    # ── Design matrices: one column per linear parameter ───────────────
+    "rv_design_matrix": ("orblet.design.columns", "rv_design_matrix"),
+    "ti_design_matrix": ("orblet.design.columns", "ti_design_matrix"),
+    "acceleration_columns": ("orblet.design.columns", "acceleration_columns"),
+    "astrometric_5param_design_matrix": (
+        "orblet.solve.astrometry",
+        "astrometric_5param_design_matrix",
+    ),
+
+    # ── Linear cores: geometry fixed, amplitudes solved exactly ────────
+    "linear_solve_rv": ("orblet.solve.rv", "linear_solve_rv"),
+    "linear_solve_ti": ("orblet.solve.astrometry", "linear_solve_ti"),
+    "best_linear_params_rv": ("orblet.solve.rv", "best_linear_params_rv"),
+    "best_linear_params_ti": (
+        "orblet.solve.astrometry",
+        "best_linear_params_ti",
+    ),
+    "fit_astrometric_5param": (
+        "orblet.solve.astrometry",
+        "fit_astrometric_5param",
+    ),
+    "recover_K": ("orblet.solve.rv", "recover_K"),
+    "recover_omega": ("orblet.solve.rv", "recover_omega"),
+    "semi_amplitude_kms": ("orblet.model", "semi_amplitude_kms"),
+
+    # ── Searching for the period ───────────────────────────────────────
+    "scan_ti_frequency": ("orblet.search", "scan_ti_frequency"),
+
+    # ── Seeds: turning one solution into a starting point for the next ─
+    "compose_rv_seed": ("orblet.solve.rv", "compose_rv_seed"),
+    "compose_ti_seed": ("orblet.search", "compose_ti_seed"),
+    "ti_amplitude_chains": ("orblet.solve.astrometry", "ti_amplitude_chains"),
+    "mirror_inclination": ("orblet.elements", "mirror_inclination"),
+
+    # ── Epochs and time scales ─────────────────────────────────────────
+    "prepare_rv_for_orbit": ("orblet.prepare", "prepare_rv_for_orbit"),
+    "resolve_epochs_mjd": ("orblet.prepare", "resolve_epochs_mjd"),
+
+    # ── Where the observer is, as on-sky parallax factors ──────────────
+    "per_direction_parallax_factors": (
+        "orblet.parallax",
+        "per_direction_parallax_factors",
+    ),
+    "OBSERVER_GEOCENTRE": ("orblet.parallax", "OBSERVER_GEOCENTRE"),
+    "OBSERVER_L2": ("orblet.parallax", "OBSERVER_L2"),
+    "L2_OFFSET_AU": ("orblet.parallax", "L2_OFFSET_AU"),
+
+    # ── Element conversions and derived curves ─────────────────────────
+    "extract_orbital_elements": (
+        "orblet.elements",
+        "extract_orbital_elements",
+    ),
+    "ti_to_kepler": ("orblet.elements", "ti_to_kepler"),
+    "to_nss_convention": ("orblet.elements", "to_nss_convention"),
+    "compute_rv_model_curve": ("orblet.elements", "compute_rv_model_curve"),
+    "compute_residuals": ("orblet.elements", "compute_residuals"),
+
+    # ── Carrying one channel's chain into the other's prior ────────────
+    "a1_sini_from_rv_chain": ("orblet.rv_chain", "a1_sini_from_rv_chain"),
+    "astrometric_priors_from_rv_chain": (
+        "orblet.rv_chain",
+        "astrometric_priors_from_rv_chain",
+    ),
+    "predicted_k_kms_from_astrometric_chain": (
+        "orblet.rv_chain",
+        "predicted_k_kms_from_astrometric_chain",
+    ),
+
+    # ── Plots (matplotlib arrives here, not before) ────────────────────
+    "plot_orbit_fit": ("orblet.plotting", "plot_orbit_fit"),
+    "plot_orbit_corner": ("orblet.plotting", "plot_orbit_corner"),
+    "plot_astrometric_orbit": ("orblet.plotting", "plot_astrometric_orbit"),
+    "plot_astrometric_sky_overlay": (
+        "orblet.plotting",
+        "plot_astrometric_sky_overlay",
+    ),
+}
+
+
+def __getattr__(name: str):
+    """Resolve a public name on first access (PEP 562).
+
+    Keeps ``import orblet`` cheap: nothing below the front door is executed
+    until something is actually asked for.
+    """
+    target = _LAZY_EXPORTS.get(name)
+    if target is None:
+        raise AttributeError(
+            f"module {__name__!r} has no attribute {name!r}"
+        )
+    module_name, attr = target
+    return getattr(importlib.import_module(module_name), attr)
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)
+
+
+# Kept as a literal, same names as the table above.
+__all__ = [
+    "rv_design_matrix",
+    "ti_design_matrix",
+    "acceleration_columns",
+    "astrometric_5param_design_matrix",
+    "linear_solve_rv",
+    "linear_solve_ti",
+    "best_linear_params_rv",
+    "best_linear_params_ti",
+    "fit_astrometric_5param",
+    "recover_K",
+    "recover_omega",
+    "semi_amplitude_kms",
+    "scan_ti_frequency",
+    "compose_rv_seed",
+    "compose_ti_seed",
+    "ti_amplitude_chains",
+    "mirror_inclination",
+    "prepare_rv_for_orbit",
+    "resolve_epochs_mjd",
+    "per_direction_parallax_factors",
+    "OBSERVER_GEOCENTRE",
+    "OBSERVER_L2",
+    "L2_OFFSET_AU",
+    "extract_orbital_elements",
+    "ti_to_kepler",
+    "to_nss_convention",
+    "compute_rv_model_curve",
+    "compute_residuals",
+    "a1_sini_from_rv_chain",
+    "astrometric_priors_from_rv_chain",
+    "predicted_k_kms_from_astrometric_chain",
+    "plot_orbit_fit",
+    "plot_orbit_corner",
+    "plot_astrometric_orbit",
+    "plot_astrometric_sky_overlay",
+]

orblet/chain_stats.py

@@ -0,0 +1,194 @@
+"""
+Chain-summary helpers for posterior samples.
+
+Linear-quantity helpers (:func:`chain_quantiles`,
+:func:`chain_credible_interval`, :func:`chain_summary_table`) are
+straightforward wrappers over :func:`numpy.quantile` with strict
+input validation.
+
+Angular-quantity helper (:func:`chain_circular_summary`) computes the
+circular mean and circular standard deviation of an angle chain.  Two
+implementation notes worth preserving:
+
+- Inputs are wrapped to ``[0, 2π)`` before the circular-std call
+  because :func:`scipy.stats.circstd` does NOT internally fold values
+  outside ``[low, high]``.
+- For known-multimodal posteriors (e.g., the ``(ω + π, Ω + π)`` 180°
+  degeneracy in retrograde-inclination astrometric fits), the circular
+  mean collapses to a meaningless midpoint between modes; the docstring
+  flags this so the helper is not used naively for such cases.
+
+*Dependencies.*  Three of the four helpers need only numpy and scipy.
+:func:`chain_summary_table` returns a :class:`pandas.DataFrame`, and pandas is
+an OPTIONAL dependency (``pip install orblet[tables]``) imported inside that
+function rather than at module scope — otherwise one unused table helper would
+make the whole module unimportable for anyone without pandas.  The numbers are
+available without it: :func:`chain_quantiles` returns the same values as a
+plain dict, and :func:`chain_summary_table` only arranges them.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from scipy.stats import circstd
+
+if TYPE_CHECKING:  # pragma: no cover - for type checkers only
+    import pandas as pd
+
+
+def _quantile_key(level: float) -> str:
+    """Encode a quantile level as a 3-digit-millis key (e.g., ``q025``)."""
+    return f"q{int(round(level * 1000)):03d}"
+
+
+def chain_quantiles(
+    samples: np.ndarray,
+    levels: tuple[float, ...] = (0.025, 0.5, 0.975),
+) -> dict[str, float]:
+    """
+    Compute quantiles of a 1-D chain.
+
+    Parameters
+    ----------
+    samples : np.ndarray, shape (n,)
+        1-D array of posterior samples.  Empty arrays raise
+        :class:`ValueError`.  NaN samples raise :class:`ValueError`.
+    levels : tuple of floats in (0, 1), default ``(0.025, 0.5, 0.975)``
+        Quantile levels.  Returned dict keys are integer-encoded:
+        ``f"q{int(round(lvl*1000)):03d}"`` → ``"q025"``, ``"q500"``,
+        ``"q975"``.
+
+    Returns
+    -------
+    dict[str, float]
+        Mapping of quantile-key to value.
+    """
+    samples = np.asarray(samples, dtype=np.float64)
+    if samples.size == 0:
+        raise ValueError("samples is empty")
+    if np.any(np.isnan(samples)):
+        raise ValueError("samples contains NaN")
+    qs = np.quantile(samples, list(levels))
+    return {_quantile_key(lvl): float(qs[i]) for i, lvl in enumerate(levels)}
+
+
+def chain_credible_interval(
+    samples: np.ndarray,
+    level: float = 0.99,
+) -> dict[str, float]:
+    """
+    Compute the central credible interval at the given level.
+
+    Parameters
+    ----------
+    samples : np.ndarray, shape (n,)
+        1-D array of posterior samples.  Empty arrays raise
+        :class:`ValueError`.
+    level : float in (0, 1), default 0.99
+        Credible-interval fraction (e.g., 0.68, 0.95, 0.99).
+        Computes quantiles at ``(1-level)/2`` and ``(1+level)/2``.
+
+    Returns
+    -------
+    dict[str, float]
+        ``{"low": float, "high": float, "level": float}``.
+    """
+    samples = np.asarray(samples, dtype=np.float64)
+    if samples.size == 0:
+        raise ValueError("samples is empty")
+    if np.any(np.isnan(samples)):
+        raise ValueError("samples contains NaN")
+    lo_q = (1.0 - level) / 2.0
+    hi_q = (1.0 + level) / 2.0
+    lo, hi = np.quantile(samples, [lo_q, hi_q])
+    return {"low": float(lo), "high": float(hi), "level": float(level)}
+
+
+def chain_circular_summary(samples: np.ndarray) -> dict[str, float]:
+    """
+    Compute circular-mean and circular-std of an angular chain.
+
+    Parameters
+    ----------
+    samples : np.ndarray, shape (n,)
+        1-D array of angle samples in radians.  Convention may be
+        ``[0, 2π)`` or ``[-π, π)``; the function pre-wraps to
+        ``[0, 2π)`` before computing.
+
+    Returns
+    -------
+    dict[str, float]
+        ``{"circmean_rad": float, "circstd_rad": float}``.  Boundary
+        convention: ``circmean_rad`` is in ``[0, 2π)``.
+
+    Notes
+    -----
+    Pre-wrap to ``[0, 2π)`` is required because
+    :func:`scipy.stats.circstd` does NOT internally re-wrap; values
+    outside ``[low, high]`` are not folded.
+
+    For posteriors known to be multi-modal (e.g., ω with the
+    ``(ω + π, Ω + π)`` 180° degeneracy in retrograde-inclination
+    astrometric fits), the circular mean collapses bimodal
+    posteriors to a meaningless midpoint.  Report per-mode summaries
+    instead in those cases.
+    """
+    samples = np.asarray(samples, dtype=np.float64)
+    if samples.size == 0:
+        raise ValueError("samples is empty")
+    if np.any(np.isnan(samples)):
+        raise ValueError("samples contains NaN")
+    # Pre-wrap to [0, 2π) — scipy.stats.circstd does not re-wrap.
+    samples = samples % (2.0 * np.pi)
+    # Circular mean via complex-exponential trick: arg(mean(exp(i·θ)))
+    # is the canonical wrap-aware estimator and ignores additive 2π
+    # shifts in any individual sample.
+    circmean = float(np.angle(np.mean(np.exp(1j * samples))) % (2.0 * np.pi))
+    circstd_val = float(circstd(samples, low=0.0, high=2.0 * np.pi))
+    return {"circmean_rad": circmean, "circstd_rad": circstd_val}
+
+
+def chain_summary_table(
+    chain_dict: dict[str, np.ndarray],
+    levels: tuple[float, ...] = (0.025, 0.5, 0.975),
+) -> pd.DataFrame:
+    """
+    Tabular summary across multiple chain keys.
+
+    Parameters
+    ----------
+    chain_dict : dict[str, np.ndarray]
+        Mapping of parameter name → 1-D sample array.
+    levels : tuple of floats, default ``(0.025, 0.5, 0.975)``
+        Same as :func:`chain_quantiles`.
+
+    Returns
+    -------
+    pd.DataFrame
+        Columns: ``parameter``, plus one column per level (integer-
+        encoded keys: ``q025``, ``q500``, ``q975``).
+    """
+    # pandas is imported HERE, not at module scope, so that the three helpers
+    # above — which need only numpy and scipy — stay usable when pandas is not
+    # installed.  A top-of-file import would take all four down over this one
+    # call.  See the module docstring for the install story.
+    try:
+        import pandas as pd
+    except ImportError as exc:  # pragma: no cover - depends on the install
+        raise ImportError(
+            "chain_summary_table returns a pandas DataFrame, and orblet does "
+            "not install pandas by default. Either `pip install orblet[tables]`, "
+            "or call chain_quantiles per parameter — it returns the same "
+            "numbers as a plain dict and needs only numpy."
+        ) from exc
+
+    rows = []
+    for name, samples in chain_dict.items():
+        qs = chain_quantiles(samples, levels=levels)
+        row = {"parameter": name}
+        row.update(qs)
+        rows.append(row)
+    columns = ["parameter"] + [_quantile_key(lvl) for lvl in levels]
+    return pd.DataFrame(rows, columns=columns)