PyAETNA 4.0.0__py3-none-any.whl

PyAETNA/reporting/residual_plotting.py

@@ -0,0 +1,419 @@
+"""
+residual_plotting.py
+~~~~~~~~~~~~~~~~~~~
+Residual panel plotting and KDE visualization.
+
+Responsibilities:
+- Plot residuals with error bars
+- Build and render KDE (kernel density estimation)
+- Format residual panels with proper labeling and references
+"""
+
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.stats import gaussian_kde
+
+
+def plot_residuals_with_kde(
+    ax_residual: Any,
+    ax_kde: Any,
+    residuals: np.ndarray,
+    time: np.ndarray,
+    dy: np.ndarray,
+    visit_sorted: np.ndarray | None = None,
+    visit_colors: dict[Any, str] | None = None,
+    data_type: str = "LC",
+    clusters: list[tuple[int, int]] | None = None,
+    cluster_idx: int = 0,
+    show_reference_bars: bool = False,
+    med_raw: float | None = None,
+    med_tot: float | None = None,
+) -> dict[str, Any]:
+    """
+    Plot residuals (time-series) and optional KDE on dual-axis panel.
+
+    Args:
+        ax_residual: Matplotlib axis for time-series
+        ax_kde: Matplotlib axis for KDE (right inset)
+        residuals: Residual values (data - model)
+        time: Time array
+        dy: Error bars
+        visit_sorted: Visits (optional, for per-visit coloring)
+        visit_colors: Dict {visit_id: color} for per-visit plots
+        data_type: 'RV' or 'LC'
+        clusters: List of (start, end) tuples for per-cluster plotting
+        cluster_idx: Current cluster index (0 for first)
+        show_reference_bars: Whether to add reference uncertainty bars
+        med_raw: Median raw uncertainty (for reference)
+        med_tot: Median total uncertainty (for reference)
+
+    Returns:
+        dict: {'rms': float, 'kde_visible': bool}
+    """
+    is_per_visit = visit_colors is not None
+
+    # Plot time-series residuals
+    if is_per_visit and visit_sorted is not None:
+        unique_visits = np.unique(visit_sorted)
+        for visit_id in unique_visits:
+            sel = visit_sorted == visit_id
+            color = (
+                visit_colors.get(visit_id, "gray")
+                if visit_colors is not None
+                else "gray"
+            )
+            label = f"Visit {visit_id}" if cluster_idx == 0 else "_nolegend_"
+            ax_residual.errorbar(
+                time[sel],
+                residuals[sel],
+                yerr=dy[sel],
+                fmt=".",
+                markersize=3,
+                elinewidth=0.5,
+                alpha=0.6,
+                color=color,
+                label=label,
+                zorder=1,
+            )
+    else:
+        label = "Data" if cluster_idx == 0 else "_nolegend_"
+        ax_residual.errorbar(
+            time,
+            residuals,
+            yerr=dy,
+            fmt=".",
+            markersize=3,
+            elinewidth=0.5,
+            alpha=0.6,
+            color="gray",
+            label=label,
+            zorder=1,
+        )
+
+    ax_residual.axhline(0, color="k", linestyle="--", linewidth=0.5, zorder=2)
+
+    # RMS text (first cluster only)
+    rms = None
+    if cluster_idx == 0:
+        rms = np.std(residuals)
+        if data_type == "RV":
+            rms_str = f"RMS = {rms:.3f} m/s"
+        else:
+            rms_str = f"RMS = {rms * 1e6:.2f} ppm"
+        ax_residual.text(
+            0.02,
+            0.95,
+            rms_str,
+            transform=ax_residual.transAxes,
+            fontsize=8,
+            verticalalignment="top",
+            bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
+        )
+
+    # Reference uncertainty bars (first cluster only)
+    if show_reference_bars and cluster_idx == 0 and med_raw is not None:
+        from .plot_utils import _add_unc_ref_bars
+
+        x_ref = time[len(time) // 2] if len(time) > 0 else 0.0
+        fmt_unc = "m/s" if data_type == "RV" else "ppm"
+        _add_unc_ref_bars(ax_residual, x_ref, 0.0, med_raw, med_tot, fmt_unc)
+
+    # KDE visualization
+    kde_visible = _plot_kde(ax_kde, residuals, visit_sorted, visit_colors, is_per_visit)
+
+    return {"rms": rms, "kde_visible": kde_visible}
+
+
+def _plot_kde(
+    ax_kde: Any,
+    residuals: np.ndarray,
+    visit_sorted: np.ndarray | None = None,
+    visit_colors: dict[Any, str] | None = None,
+    is_per_visit: bool = False,
+) -> bool:
+    """
+    Render KDE panel (right inset).
+
+    Args:
+        ax_kde: Matplotlib axis
+        residuals: Residual values
+        visit_sorted: Visits (optional)
+        visit_colors: Visit color dict (optional)
+        is_per_visit: Bool
+
+    Returns:
+        bool: True if KDE was plotted, False otherwise
+    """
+    if len(residuals) < 4:
+        return False
+
+    try:
+        # Determine y-range for KDE sampling
+        res_span = float(residuals.max() - residuals.min())
+        margin = 0.15 * res_span if res_span > 0 else 1e-6
+        y_k = np.linspace(
+            float(residuals.min()) - margin, float(residuals.max()) + margin, 300
+        )
+
+        if is_per_visit and visit_sorted is not None:
+            # Overall combined KDE
+            kde_all = gaussian_kde(residuals)
+            ax_kde.plot(
+                kde_all(y_k), y_k, color="black", linewidth=1.2, label="All data"
+            )
+
+            # Reference Gaussian
+            mu = float(np.mean(residuals))
+            sig = float(np.std(residuals))
+            if sig > 0:
+                gauss = np.exp(-0.5 * ((y_k - mu) / sig) ** 2) / (
+                    sig * np.sqrt(2.0 * np.pi)
+                )
+                ax_kde.plot(
+                    gauss,
+                    y_k,
+                    color="C3",
+                    linewidth=1.0,
+                    linestyle="--",
+                    alpha=0.8,
+                    label="Gaussian ref.",
+                )
+
+            # Per-visit KDEs
+            unique_visits = np.unique(visit_sorted)
+            if len(unique_visits) > 1:
+                for visit_id in unique_visits:
+                    sel = visit_sorted == visit_id
+                    res_visit = residuals[sel]
+                    if len(res_visit) > 3:
+                        kde_visit = gaussian_kde(res_visit)
+                        color = (
+                            visit_colors.get(visit_id, "gray")
+                            if visit_colors is not None
+                            else "gray"
+                        )
+                        ax_kde.plot(
+                            kde_visit(y_k),
+                            y_k,
+                            color=color,
+                            linewidth=0.5,
+                            alpha=0.8,
+                            label=f"Visit {visit_id}",
+                        )
+
+            ax_kde.legend(fontsize=6, loc="upper right")
+        else:
+            # Single KDE
+            kde_all = gaussian_kde(residuals)
+            ax_kde.plot(kde_all(y_k), y_k, color="black", linewidth=1.2)
+
+            # Reference Gaussian
+            mu = float(np.mean(residuals))
+            sig = float(np.std(residuals))
+            if sig > 0:
+                gauss = np.exp(-0.5 * ((y_k - mu) / sig) ** 2) / (
+                    sig * np.sqrt(2.0 * np.pi)
+                )
+                ax_kde.plot(
+                    gauss,
+                    y_k,
+                    color="C3",
+                    linewidth=1.0,
+                    linestyle="--",
+                    alpha=0.8,
+                    label="Gaussian ref.",
+                )
+
+        ax_kde.axhline(0, color="k", linestyle="--", linewidth=0.5)
+        ax_kde.set_xlim(left=0)
+
+        return True
+
+    except (ImportError, Exception):
+        return False
+
+
+def style_residual_axes(
+    axes: list[Any], clusters: list[tuple[int, int]] | None = None
+) -> None:
+    """
+    Apply consistent styling to residual axes.
+
+    Args:
+        axes: List of residual axes
+        clusters: List of (start, end) tuples for multi-panel layout
+    """
+    for idx, ax in enumerate(axes):
+        ax.set_xlabel("Time [d]")
+        ax.set_ylabel("Residuals (σ_orig error bars)")
+
+        # Hide y-labels on non-leftmost columns
+        if idx > 0:
+            plt.setp(ax.get_yticklabels(), visible=False)
+            plt.setp(ax.get_yaxis().get_label(), visible=False)
+
+
+def style_kde_axis(ax_kde: Any, kde_visible: bool) -> None:
+    """
+    Apply consistent styling to KDE axis.
+
+    Args:
+        ax_kde: KDE axis
+        kde_visible: Bool
+    """
+    if kde_visible:
+        ax_kde.tick_params(labelleft=False, labelsize=7)
+        ax_kde.set_xlabel("KDE", fontsize=7)
+        ax_kde.xaxis.set_major_locator(plt.MaxNLocator(2))
+    else:
+        ax_kde.set_visible(False)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Self-contained KDE inset renderer used by plot_models
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def render_kde_inset(
+    ax_kde: Any,
+    residuals: np.ndarray,
+    vis_sorted: np.ndarray,
+    unique_vis: np.ndarray,
+    visit_colors: dict[Any, str] | None,
+    is_per_inst_plot: bool,
+    ref_sigma: float | None = None,
+) -> bool:
+    """Render the KDE residual distribution inset.
+
+    The KDE panel shares its y-axis with the residuals time-series so the
+    residual-value scale is identical on both sides.  The x-axis shows the
+    probability density (arbitrary units, increasing to the right), making
+    it easy to see at a glance whether the residuals are Gaussian, skewed,
+    or have heavy tails, without leaving the time-series context.
+
+    Parameters
+    ----------
+    ax_kde : matplotlib.axes.Axes
+        The right-inset axis (y-axis shared with residuals panel).
+    residuals : array
+        Residual values (data − total model).
+    vis_sorted, unique_vis : arrays
+        Per-point visit labels and unique visit IDs (used for per-visit KDEs).
+    visit_colors : dict or None
+        ``{visit_id: color}`` mapping; ``None`` disables per-visit KDEs.
+    is_per_inst_plot : bool
+        When True and ``visit_colors`` is set, draws one KDE per visit.
+    ref_sigma : float, optional
+        Width of the dashed "Gaussian ref." curve, in residual units.  This
+        should be the fitted **total** uncertainty (raw ⊕ jitter) — e.g.
+        ``ctx.med_tot`` or a channel's ``med_tot`` — so the curve answers
+        "does the residual scatter match what the noise model (including the
+        fitted jitter) predicts?", not merely "what Gaussian best fits the
+        residuals themselves".  When several rows in view have different
+        total uncertainties (heterogeneous raw errors, or a per-visit jitter
+        on an ``is_per_inst_plot`` page), pass a representative value (the
+        median, as the noise-budget panels and the reference bars do) — the
+        legend label says so.  If omitted or non-positive, falls back to the
+        empirical standard deviation of the residuals, matching the curve's
+        pre-existing behaviour.
+
+    Returns
+    -------
+    kde_visible : bool
+        ``True`` if KDE was rendered, ``False`` if suppressed (too few points
+        or scipy not available).
+    """
+    kde_visible = False
+    try:
+        from scipy.stats import gaussian_kde as _gkde
+
+        if len(residuals) >= 4:
+            kde_visible = True
+            _res_span = float(residuals.max() - residuals.min())
+            _margin = 0.15 * _res_span if _res_span > 0 else 1e-6
+            _y_k = np.linspace(
+                float(residuals.min()) - _margin, float(residuals.max()) + _margin, 300
+            )
+
+            def _draw_gaussian_ref(label: str) -> None:
+                _mu = float(np.mean(residuals))
+                _sig = (
+                    float(ref_sigma)
+                    if ref_sigma is not None and ref_sigma > 0
+                    else float(np.std(residuals))
+                )
+                if _sig <= 0:
+                    return
+                _gauss = np.exp(-0.5 * ((_y_k - _mu) / _sig) ** 2) / (
+                    _sig * np.sqrt(2.0 * np.pi)
+                )
+                ax_kde.plot(
+                    _gauss,
+                    _y_k,
+                    color="C3",
+                    linewidth=1.0,
+                    linestyle="--",
+                    alpha=0.8,
+                    label=label,
+                )
+
+            _gauss_label = (
+                "Gaussian ref. (median σ_tot)"
+                if ref_sigma is not None and ref_sigma > 0
+                else "Gaussian ref."
+            )
+
+            if is_per_inst_plot and visit_colors is not None:
+                try:
+                    _kde_v = _gkde(residuals)(_y_k)
+                    ax_kde.fill_betweenx(_y_k, 0, _kde_v, color="gray", alpha=0.30)
+                    ax_kde.plot(
+                        _kde_v, _y_k, color="gray", linewidth=1.2, label="All visits"
+                    )
+                except Exception:
+                    pass
+
+                _draw_gaussian_ref(_gauss_label)
+
+                for uv in unique_vis:
+                    sel = vis_sorted == uv
+                    res_uv = residuals[sel]
+                    if len(res_uv) < 4:
+                        continue
+                    try:
+                        kde = _gkde(res_uv)
+                        ax_kde.plot(
+                            kde(_y_k),
+                            _y_k,
+                            color=visit_colors[uv],
+                            linewidth=1.2,
+                            alpha=0.8,
+                            label=f"Visit {uv}",
+                        )
+                    except Exception:
+                        pass
+            else:
+                _kde_v = _gkde(residuals)(_y_k)
+                ax_kde.fill_betweenx(_y_k, 0, _kde_v, color="gray", alpha=0.30)
+                ax_kde.plot(_kde_v, _y_k, color="gray", linewidth=1.2, label="Data")
+
+                _draw_gaussian_ref(_gauss_label)
+
+            ax_kde.axhline(0, color="k", linestyle="--", linewidth=0.5)
+            ax_kde.set_xlim(left=0)
+            ax_kde.legend(fontsize=6, loc="upper right")
+    except ImportError:
+        pass
+    except Exception:
+        kde_visible = False
+
+    if kde_visible:
+        ax_kde.tick_params(labelleft=False, labelsize=7)
+        ax_kde.set_xlabel("KDE", fontsize=7)
+        ax_kde.xaxis.set_major_locator(plt.MaxNLocator(2))
+    else:
+        ax_kde.set_visible(False)
+
+    return kde_visible

PyAETNA/reporting/trend_detection.py

@@ -0,0 +1,128 @@
+"""
+trend_detection.py
+~~~~~~~~~~~~~~~~~
+Detect and compute non-time detrending components.
+
+Responsibilities:
+- Identify whether polynomial/Fourier/GP use non-time variables
+- Build non-time trend arrays
+- Handle per-instrument GP contributions
+"""
+
+from typing import Any
+
+import numpy as np
+
+from .model_math import _compute_nontime_trend
+
+
+def detect_nontime_xvariables(
+    poly_cfgs: list[dict[str, Any]],
+    four_cfgs: list[dict[str, Any]],
+    gp_cfgs: dict[Any, dict[str, Any]],
+    eiv_cfgs: list[dict[str, Any]] | None = None,
+) -> tuple[bool, bool]:
+    """
+    Check if any detrending component uses a non-time variable.
+
+    Args:
+        poly_cfgs: List of polynomial configs
+        four_cfgs: List of Fourier configs
+        gp_cfgs: Dict of {inst_id: gp_config}
+        eiv_cfgs: List of EIV configs (optional)
+
+    Returns:
+        tuple: (has_poly_fourier, has_gp_nontime) booleans
+    """
+    if eiv_cfgs is None:
+        eiv_cfgs = []
+    has_poly_fourier = (
+        any(cfg.get("x_variable", "TIME") != "TIME" for cfg in poly_cfgs)
+        or any(cfg.get("x_variable", "TIME") != "TIME" for cfg in four_cfgs)
+        or any(cfg.get("x_variable", "TIME") != "TIME" for cfg in eiv_cfgs)
+    )
+
+    has_gp_nontime = any(_gp_uses_nontime_xvar(cfg) for cfg in gp_cfgs.values())
+
+    return has_poly_fourier, has_gp_nontime
+
+
+def _gp_uses_nontime_xvar(gp_cfg: dict[str, Any]) -> bool:
+    """Check if a single GP config uses non-time variable."""
+    xvars = gp_cfg.get("x_variables", gp_cfg.get("x_variable", ["TIME"]))
+    if isinstance(xvars, str):
+        xvars = [xvars]
+    return any(xv != "TIME" for xv in xvars)
+
+
+def build_nontime_trend_at_data_points(
+    model_comp: dict[str, Any],
+    plData: Any,
+    param_config: dict[str, Any],
+    poly_cfgs: list[dict[str, Any]],
+    four_cfgs: list[dict[str, Any]],
+    gp_cfgs: dict[Any, dict[str, Any]],
+) -> np.ndarray | None:
+    """
+    Pre-compute non-time trend component at every data point.
+
+    This combines polynomial/Fourier non-time terms and any GP predictions
+    on non-time variables. Used for Panel 0 (raw data + model) overlay.
+
+    Args:
+        model_comp: Best-fit model components
+        plData: Data object
+        param_config: Configuration
+        poly_cfgs, four_cfgs, gp_cfgs: Configuration dicts
+
+    Returns:
+        ndarray: Non-time trend at each data point, or None if no non-time trends
+    """
+    has_poly_fourier, has_gp = detect_nontime_xvariables(poly_cfgs, four_cfgs, gp_cfgs)
+
+    if not (has_poly_fourier or has_gp):
+        return None
+
+    # Start with polynomial/Fourier/EIV non-time contributions
+    trend = _compute_nontime_trend(
+        model_comp["poly_terms"],
+        model_comp["fourier_terms"],
+        model_comp.get("eiv_terms", {}),
+        poly_cfgs,
+        four_cfgs,
+        model_comp.get("eiv_configs", []),
+        len(plData.time),
+    )
+
+    # Add non-time GP contributions
+    if has_gp:
+        inst_arr = np.array(plData.instrument)
+        for gp_inst_id, gp_cfg in gp_cfgs.items():
+            if _gp_uses_nontime_xvar(gp_cfg):
+                mask = inst_arr == gp_inst_id
+                # model_comp['gp_mean'] contains the full GP prediction
+                # for this instrument; add it where it applies
+                trend[mask] += model_comp["gp_mean"][mask]
+
+    return trend
+
+
+def apply_nontime_detrending(
+    y: np.ndarray, nontime_trend: np.ndarray | None, mask: np.ndarray | None = None
+) -> np.ndarray:
+    """
+    Remove non-time trend from data.
+
+    Args:
+        y: Data array
+        nontime_trend: Non-time trend to subtract
+        mask: Optional subset mask
+
+    Returns:
+        ndarray: Detrended data
+    """
+    if nontime_trend is None:
+        return y
+
+    trend_to_use = nontime_trend if mask is None else nontime_trend[mask]
+    return y - trend_to_use