PyAETNA 4.0.0__py3-none-any.whl

PyAETNA/gp.py

@@ -0,0 +1,791 @@
+"""
+PyAETNA GP Module - Shared Gaussian Process utilities
+
+This module provides GP kernel building and prediction functions used by both
+the sampler (PyAETNA/sampling/engine/runner.py) and reporting modules.
+
+Centralizing GP code here ensures consistency between sampling and visualization.
+
+Per-visit block-diagonal GP
+----------------------------
+When ``per_visit: true`` is set in a GP section, the covariance matrix is
+factorised into independent blocks—one per visit—while **sharing** the same
+kernel hyperparameters across all visits.  This is a *computational strategy*,
+not a parameter-scope expansion:
+
+.. math::
+
+    \\log p(\\mathbf{y} \\mid \\theta)
+    = \\sum_v \\log p(\\mathbf{y}_v \\mid \\theta)
+
+This factorisation is valid **only when the off-diagonal inter-visit blocks of
+the covariance matrix are negligible**, i.e. when the kernel decays to zero over
+the inter-visit gap.  Kernels that violate this condition make the factorisation
+mathematically incorrect:
+
+* **``ExpSineSquared`` (strictly periodic kernel)** — covariance never decays;
+  $k(nP) = \\sigma^2$ at any integer multiple of the period.  Per-visit is
+  **contraindicated** for this kernel.
+* **``SHO_Quality`` / ``SHO`` with very high quality factor** — damping is slow;
+  cross-visit covariances remain significant when the inter-visit gap is much
+  shorter than the damping time $Q / \\omega_0$.  Use with caution and verify
+  that the fitted Q implies fast damping relative to the gap.
+* **``QuasiPeriodic``** — the SE envelope makes it conditionally valid when the
+  inter-visit gap $\\gg \\ell_\\text{SE}$.  Safe for typical multi-visit
+  photometry (gaps of days, SE scales of hours).
+
+For all monotonically decaying kernels (Matérn32, Matérn52, ExpSquared, SHO
+with moderate Q) the benefits are:
+
+* **O(N³) dense kernels** — wall-clock speedup proportional to V² (number of
+  visits squared), since each block is ``(N/V)³`` instead of ``N³``.
+* **SHO quasisep kernels** — prevents the ``-inf`` numerical catastrophe that
+  occurs when Kalman transition-matrix elements ``exp(-α Δt)`` underflow at
+  large inter-visit gaps.
+* **O(N) Matérn quasisep** — no benefit (already numerically robust and
+  linear); a small overhead from the Python loop is incurred.
+
+See ``docs/gp_per_visit_analysis/GP_PER_VISIT_ANALYSIS.md`` for benchmarks.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from tinygp import GaussianProcess
+
+# Import kernel functions from the new dedicated kernels/factory.py module
+from PyAETNA.kernels.factory import (
+    build_gp_kernel_for_plotting,
+    build_gp_kernel_for_sampling,
+    build_gp_kernel_from_specs,
+)
+from PyAETNA.logger import _log
+
+__all__ = [
+    "build_gp_kernel_for_plotting",
+    "build_gp_kernel_for_sampling",
+    "build_gp_kernel_from_specs",
+    "compute_kernel_theoretical_psd",
+    "compute_windowed_psd",
+    "get_gp_x_variable",
+    "gp_log_probability_per_visit",
+    "gp_predict",
+    "gp_predict_per_visit",
+]
+
+# =============================================================================
+# GP Prediction (JIT-compiled for speed)
+# =============================================================================
+
+
+def gp_predict(
+    kernel: Any,
+    x: Any,
+    y: Any,
+    dy_sq: Any,
+    x_test: Any,
+    log_func: Callable[..., None] | None = None,
+    log_target: Any = None,
+) -> tuple[jax.Array, jax.Array]:
+    """
+    GP prediction wrapper. Handles shape validation and logging before
+    calling the JIT-compiled core.
+
+    Args:
+        kernel: tinygp kernel object
+        x: Training inputs (1D or 2D array)
+        y: Training outputs (residuals)
+        dy_sq: Squared uncertainties (diagonal noise)
+        x_test: Test inputs (usually same as x for interpolation)
+        log_func: Logging callable (expected signature: log_func(target, msg))
+        log_target: Target for logging (e.g. log file path or PlModel object)
+
+    Returns:
+        mean, std: Predicted mean and standard deviation at x_test
+    """
+    # Convert to arrays for shape checking
+    x = jnp.asarray(x)
+    y = jnp.asarray(y)
+    dy_sq = jnp.asarray(dy_sq)
+    x_test = jnp.asarray(x_test)
+
+    # Shape repair and validation
+    if dy_sq.ndim > 1:
+        flat = dy_sq.reshape(-1)
+        if flat.shape[0] == y.shape[0]:
+            dy_sq = flat
+        else:
+            msg = f"ERROR: dy_sq shape {dy_sq.shape} cannot be flattened to match y length {y.shape[0]}"
+            if log_func:
+                log_func(log_target, msg)
+            raise ValueError(msg)
+
+    if dy_sq.shape[0] != y.shape[0]:
+        msg = f"ERROR: dy_sq length {dy_sq.shape[0]} must match y length {y.shape[0]}"
+        if log_func:
+            log_func(log_target, msg)
+        raise ValueError(msg)
+
+    return _gp_predict_jit(kernel, x, y, dy_sq, x_test)
+
+
+@jax.jit
+def _gp_predict_jit(
+    kernel: Any, x: jax.Array, y: jax.Array, dy_sq: jax.Array, x_test: jax.Array
+) -> tuple[jax.Array, jax.Array]:
+    """Core JIT-compiled GP prediction logic."""
+    gp = GaussianProcess(kernel, x, diag=dy_sq)
+    gp_cond = gp.condition(y, x_test)
+    return gp_cond.gp.loc, jnp.sqrt(gp_cond.gp.variance)
+
+
+# =============================================================================
+# GP x-variable extraction helper
+# =============================================================================
+
+
+def get_gp_x_variable(
+    plData: Any, gp_cfg: dict[str, Any], inst_mask: Any, use_jax: bool = True
+) -> jax.Array | np.ndarray:
+    """
+    Extract GP input variable(s) for an instrument.
+
+    Args:
+        plData: PlData object with time and dBasis
+        gp_cfg: GP configuration dict with 'x_variable' or 'x_variables'
+        inst_mask: Boolean mask or index array for the instrument
+        use_jax: If True, return jax arrays. If False, return numpy arrays.
+
+    Returns:
+        x_inst: Input array for GP (1D or 2D depending on number of variables)
+    """
+    xp = jnp if use_jax else np
+
+    # Support both x_variable (string) and x_variables (list)
+    x_variables = gp_cfg.get("x_variables", gp_cfg.get("x_variable", ["TIME"]))
+    if isinstance(x_variables, str):
+        x_variables = [x_variables]
+
+    if len(x_variables) == 1:
+        # Single variable - 1D input
+        xv = x_variables[0]
+        if xv == "TIME":
+            x_var = xp.array(plData.time) if use_jax else np.array(plData.time)
+        else:
+            x_var = (
+                xp.array(plData.dBasis[xv]) if use_jax else np.array(plData.dBasis[xv])
+            )
+        x_inst = x_var[inst_mask]
+    else:
+        # Multiple variables - stack into 2D input
+        x_arrays = []
+        for xv in x_variables:
+            if xv == "TIME":
+                x_arrays.append(np.array(plData.time))
+            else:
+                x_arrays.append(np.array(plData.dBasis[xv]))
+        x_var = np.column_stack(x_arrays)
+        x_inst = x_var[inst_mask]
+        if use_jax:
+            x_inst = jnp.array(x_inst)
+
+    return x_inst
+
+
+# =============================================================================
+# Per-visit block-diagonal GP helpers
+# =============================================================================
+
+
+# Variance of the fake points that pad per-visit GP blocks to a common length,
+# and the log-probability each of them contributes (see
+# gp_log_probability_per_visit).
+_PAD_VARIANCE = 1e30
+_PAD_LOGPROB_CORRECTION = 0.5 * float(np.log(2.0 * np.pi * _PAD_VARIANCE))
+
+
+def make_gp(
+    kernel: Any, x: Any, diag: Any, *, assume_sorted: bool = False
+) -> GaussianProcess:
+    """Build a tinygp ``GaussianProcess``, optionally skipping the sortedness check.
+
+    For quasiseparable kernels tinygp (0.3.1) checks, on every evaluation,
+    that the inputs are sorted, through a ``jax.debug.callback``: each likelihood call then
+    waits for a device-to-host transfer, which is costly on a GPU (and fails
+    when no CPU device is visible).  PyAETNA's inputs are fixed data, so the
+    caller verifies the ordering once, before tracing, and passes
+    ``assume_sorted=True``.  Dense kernels have no such check.
+    """
+    from tinygp.kernels.quasisep import Quasisep
+
+    if assume_sorted and isinstance(kernel, Quasisep):
+        return GaussianProcess(kernel, x, diag=diag, assume_sorted=True)
+    return GaussianProcess(kernel, x, diag=diag)
+
+
+def is_quasiseparable(kernel: Any) -> bool:
+    from tinygp.kernels.quasisep import Quasisep
+
+    return isinstance(kernel, Quasisep)
+
+
+def gp_log_probability(
+    kernel: Any,
+    x: Any,
+    diag: Any,
+    y: Any,
+    *,
+    assume_sorted: bool = False,
+    parallel: bool = False,
+) -> Any:
+    """GP marginal log-likelihood of residuals ``y``.
+
+    With ``parallel=True``, a quasiseparable kernel on sorted 1-D inputs is
+    evaluated with the parallel-in-time Kalman filter
+    (:mod:`PyAETNA.gp_parallel`), which gives the same value as tinygp but in
+    ~2 log2 N sequential steps instead of N: much faster on GPUs, slower on
+    CPUs. Any other case uses tinygp.
+    """
+    if parallel and assume_sorted and is_quasiseparable(kernel) and jnp.ndim(x) == 1:
+        from PyAETNA.gp_parallel import log_probability_parallel
+
+        return log_probability_parallel(kernel, x, diag, y)
+    return make_gp(kernel, x, diag, assume_sorted=assume_sorted).log_probability(y)
+
+
+def resolve_gp_solver(solver: str) -> bool:
+    """True if the parallel quasiseparable solver should be used.
+
+    ``auto`` selects it on GPUs, where it is 40-150x faster than tinygp's
+    sequential scan, and not on CPUs, where it is 2-4x slower
+    (``paper_benchmarks/gp_solver``).
+    """
+    solver = str(solver).lower()
+    if solver not in ("auto", "sequential", "parallel"):
+        raise ValueError(
+            f"GP solver must be 'auto', 'sequential' or 'parallel', got {solver!r}"
+        )
+    if solver == "auto":
+        return jax.default_backend() == "gpu"
+    return solver == "parallel"
+
+
+def inputs_sorted(x: Any) -> bool:
+    """True if a 1-D input array is non-decreasing (checked on concrete data)."""
+    arr = np.asarray(x)
+    return arr.ndim == 1 and bool(np.all(np.diff(arr) >= 0))
+
+
+def gp_log_probability_per_visit(
+    kernel: Any,
+    x_inst: jax.Array,
+    res_inst: jax.Array,
+    diag_inst: jax.Array,
+    visit_indices: list[Any],
+    assume_sorted: bool = False,
+    parallel: bool = False,
+) -> jax.Array:
+    """
+    Compute the GP marginal log-probability using a per-visit block-diagonal
+    factorisation.  Hyperparameters are **shared** across all visits.
+
+    This is mathematically equivalent to the full-matrix computation when
+    inter-visit covariance is negligible (gap >> kernel length scale), which
+    is the expected regime for space-based photometry (e.g. CHEOPS visits
+    separated by days, kernel length scales of minutes).
+
+    .. math::
+
+        \\log p(\\mathbf{y} \\mid \\theta)
+        = \\sum_v \\log p(\\mathbf{y}_v \\mid \\theta)
+
+    Each visit contributes an independent block of size ``n_v``, so the
+    total cost for O(N³) dense kernels drops from ``N³`` to
+    ``\\sum_v n_v^3 \\approx V \\cdot (N/V)^3 = N^3 / V^2``.
+
+    For SHO quasisep kernels this factorisation is a **numerical necessity**:
+    the Kalman state-transition matrix contains ``exp(-\\alpha \\Delta t)``
+    terms that underflow for large inter-visit gaps, producing ``-inf``.
+
+    **Compilation optimization** — All visits are padded to the same length
+    and evaluated via a single ``jax.vmap``-ed GP call.  This replaces the
+    previous Python for-loop (which JAX unrolled into *n_visits* separate GP
+    subgraphs) with one compiled GP kernel batched across visits,
+    dramatically reducing XLA graph size and compilation time.  Padded
+    positions use ``diag = 1e30`` (effectively infinite observation noise):
+    the Kalman gain is driven to zero, so padding cannot corrupt the
+    real-data likelihood contribution—it only adds a constant independent
+    of the kernel hyperparameters.
+
+    .. warning::
+
+        This factorisation is only valid when the kernel decays sufficiently
+        over the inter-visit gap, i.e. the off-diagonal cross-visit covariance
+        blocks must be negligible.  Three specific cases break this assumption:
+
+        * **``ExpSineSquared``** (strictly periodic): covariance never decays;
+          ``k(n*P) == sigma^2`` for any integer ``n``.  Using per-visit with
+          this kernel is **mathematically incorrect** and will produce biased
+          hyperparameter estimates.  Do not combine ``ExpSineSquared`` with
+          ``per_visit: true``.
+        * **``SHO_Quality`` / ``SHO`` with high Q factor**: decay is slow;
+          cross-visit covariances are significant when the inter-visit gap is
+          much less than the damping time ``Q / omega_0``.  Verify that the
+          posterior Q implies fast damping relative to the gap.
+        * **``QuasiPeriodic``**: conditionally valid when the inter-visit gap
+          is much larger than ``exp(log_gp_se_scale)`` (the SE length scale
+          that provides the decaying envelope).
+
+    Args:
+        kernel: tinygp kernel object (built from shared hyperparameters).
+        x_inst: GP input array for the instrument (1-D, sorted within visits).
+        res_inst: Residuals for the instrument.
+        diag_inst: Diagonal variance (data error² + jitter²) for the instrument.
+        visit_indices: List of JAX index arrays, one per visit, giving the
+            positions within the *instrument* arrays (not the global data array).
+
+    Returns:
+        Scalar GP marginal log-probability (sum over visits).
+    """
+    n_visits = len(visit_indices)
+
+    if n_visits == 0:
+        return jnp.zeros(())
+
+    # Single visit: no padding or vmap overhead needed.
+    if n_visits == 1:
+        v_idx = visit_indices[0]
+        return gp_log_probability(
+            kernel,
+            x_inst[v_idx],
+            diag_inst[v_idx],
+            res_inst[v_idx],
+            assume_sorted=assume_sorted,
+            parallel=parallel,
+        )
+
+    # --- Pad all visits to uniform length for a single vmap-ed GP call ------
+    # Instead of a Python for-loop (which JAX unrolls into n_visits separate
+    # GP subgraphs), we build (n_visits, max_len) arrays and vmap a single
+    # GP evaluation.  This compiles ONE GP solve and batches it, reducing
+    # XLA graph size by ~n_visits×.
+    #
+    # Padding conventions:
+    #   x:    sorted values placed far after the last real point (gap = 1e4)
+    #         so that exp(-α*Δt) → 0 and the Kalman state resets.
+    #   r:    zero residuals (no signal contribution).
+    #   diag: 1e30 (effectively infinite noise → Kalman gain ≈ 0 → no
+    #         state update; log-probability contribution is a constant
+    #         that does not depend on kernel hyperparameters).
+    max_len = max(len(v_idx) for v_idx in visit_indices)
+
+    x_padded = jnp.zeros((n_visits, max_len), dtype=x_inst.dtype)
+    r_padded = jnp.zeros((n_visits, max_len), dtype=res_inst.dtype)
+    d_padded = jnp.full((n_visits, max_len), _PAD_VARIANCE, dtype=diag_inst.dtype)
+
+    for i, v_idx in enumerate(visit_indices):
+        n_v = len(v_idx)  # static (concrete) at trace time
+        x_padded = x_padded.at[i, :n_v].set(x_inst[v_idx])
+        r_padded = r_padded.at[i, :n_v].set(res_inst[v_idx])
+        d_padded = d_padded.at[i, :n_v].set(diag_inst[v_idx])
+        # Padding x-values: sorted, placed far after the last real point.
+        # The large gap (1e4) ensures exp(-α*Δt) → 0, resetting the
+        # Kalman state so padding cannot corrupt the real-data block.
+        if n_v < max_len:
+            n_pad = max_len - n_v
+            last_x = x_inst[v_idx[-1]]
+            pad_start = last_x + jnp.asarray(1e4, dtype=x_inst.dtype)
+            pad_x = pad_start + jnp.arange(n_pad, dtype=x_inst.dtype) * 0.001
+            x_padded = x_padded.at[i, n_v:].set(pad_x)
+
+    def _visit_logprob(
+        x_v: jnp.ndarray, r_v: jnp.ndarray, d_v: jnp.ndarray
+    ) -> jnp.ndarray:
+        # Padding appends points after the last real one, so a sorted visit
+        # stays sorted.
+        return gp_log_probability(
+            kernel, x_v, d_v, r_v, assume_sorted=assume_sorted, parallel=parallel
+        )
+
+    # Single compiled GP kernel, vmapped over the visit dimension.
+    # The kernel object (captured from outer scope) is shared—hyperparameters
+    # are NOT batched, only x/r/d are.
+    logprobs = jax.vmap(_visit_logprob)(x_padded, r_padded, d_padded)
+    # Each padded point (zero residual, variance _PAD_VARIANCE, uncorrelated
+    # with the data) adds -0.5 ln(2 pi _PAD_VARIANCE) to the log-probability.
+    # The term is constant, so posteriors are unaffected, but it would shift
+    # the Bayesian evidence (by ~35.5 nats per padded point): remove it.
+    n_padded = sum(max_len - len(v_idx) for v_idx in visit_indices)
+    return jnp.sum(logprobs) + n_padded * _PAD_LOGPROB_CORRECTION
+
+
+def gp_predict_per_visit(
+    kernel: Any,
+    x_inst: Any,
+    res_inst: Any,
+    diag_inst: Any,
+    visit_indices: list[Any],
+    x_pred_inst: Any = None,
+) -> tuple[jax.Array, jax.Array]:
+    """
+    Compute GP posterior predictions using per-visit block-diagonal
+    conditioning, with **shared** kernel hyperparameters.
+
+    Each visit is conditioned independently (no cross-visit covariance is used
+    for the conditioning), then results are concatenated back into the
+    instrument-level arrays.  The same validity conditions as
+    :func:`gp_log_probability_per_visit` apply: the kernel must decay over
+    the inter-visit gap.  In particular, ``ExpSineSquared`` (strictly
+    periodic, non-decaying) must **not** be combined with this function, as
+    ignoring cross-visit covariance will corrupt the posterior mean and
+    variance estimates.
+
+    Args:
+        kernel: tinygp kernel object.
+        x_inst: Training inputs for the instrument.
+        res_inst: Training residuals.
+        diag_inst: Diagonal variance (data error² + jitter²).
+        visit_indices: List of JAX index arrays (per-visit positions within
+            the instrument arrays).
+        x_pred_inst: (optional) Prediction inputs.  If ``None``, predicts at
+            the training locations.
+
+    Returns:
+        (mean, std) — predicted GP posterior mean and standard deviation,
+        same length as ``x_inst`` (or ``x_pred_inst`` if provided).
+    """
+    n_out = len(x_inst) if x_pred_inst is None else len(x_pred_inst)
+    mean_out = jnp.zeros(n_out)
+    std_out = jnp.zeros(n_out)
+
+    for v_idx in visit_indices:
+        # Convert to JAX arrays — dtype follows the global x64 setting,
+        # keeping data and kernel params consistently float32 or float64.
+        x_v = jnp.asarray(x_inst[v_idx])
+        r_v = jnp.asarray(res_inst[v_idx])
+        d_v = jnp.asarray(diag_inst[v_idx])
+        gp_v = GaussianProcess(kernel, x_v, diag=d_v)
+
+        if x_pred_inst is None:
+            cond = gp_v.condition(r_v, x_v)
+        else:
+            x_p_v = jnp.asarray(x_pred_inst[v_idx])
+            cond = gp_v.condition(r_v, x_p_v)
+
+        mean_out = mean_out.at[v_idx].set(cond.gp.loc)
+        std_out = std_out.at[v_idx].set(jnp.sqrt(cond.gp.variance))
+
+    return mean_out, std_out
+
+
+# =============================================================================
+# Kernel theoretical PSD computation (for periodogram comparison)
+# =============================================================================
+
+
+def compute_kernel_theoretical_psd(
+    kernel: Any,
+    freqs: np.ndarray,
+    n_acf_lags: int = 8192,
+    oversample: float = 20.0,
+    max_refine: int = 6,
+    rtol: float = 1e-3,
+    max_n_lags: int = 2**20,
+    alias_rtol: float = 1e-4,
+) -> np.ndarray:
+    """
+    Compute the one-sided power spectral density (PSD) of a tinygp kernel
+    at a set of requested cyclic frequencies via numerical FFT of the
+    autocovariance function.
+
+    The approach is kernel-agnostic: it builds a regularly-spaced lag grid,
+    evaluates the autocovariance k(0, tau) at each lag via ``jax.vmap``, then
+    computes the PSD via a real FFT. Works for dense kernels (ExpSquared,
+    ExpSineSquared, Matern, ...), quasisep kernels, and arbitrary sums/
+    products thereof.
+
+    **Aliasing safeguard**: the ACF lag spacing ``dt_lag`` must be fine
+    enough to resolve the *kernel's own* frequency content, which for
+    periodic/quasi-periodic kernels (ExpSineSquared and friends) can extend
+    to harmonics far above the highest frequency you actually care about in
+    ``freqs``. Coupling ``dt_lag`` directly to ``max(freqs)`` (as a naive
+    Nyquist argument would suggest) undersamples that structure and folds
+    high harmonics back into the requested band, producing spurious/
+    intermittent extra peaks. Instead, this function starts from a lag
+    spacing oversampled by ``oversample`` beyond ``max(freqs)`` and
+    adaptively halves it (doubling the number of lags to keep the total lag
+    span, and hence frequency resolution, fixed) until **both** of the
+    following hold:
+
+    1. **There is no power left to fold**: the PSD in the top octave of the
+       lag grid's own Nyquist band,
+       :math:`[f_{\\text{nyq}}/2,\\, f_{\\text{nyq}}]` with
+       :math:`f_{\\text{nyq}} = 1/(2\\,\\Delta\\tau)`, has dropped below
+       ``alias_rtol`` times the peak PSD *in the requested band*. Whatever
+       lies above :math:`f_{\\text{nyq}}` and folds back down is bounded by
+       that level, so it cannot contaminate the result. Being measured on
+       the spectrum actually computed, this detects kernel-specific
+       high-frequency content empirically rather than assuming it is
+       bounded by ``max(freqs)``.
+    2. **Refining no longer changes the answer**: the in-band PSD moved by
+       less than ``rtol`` of its own peak between successive refinements.
+       This catches the *other* dt-dependent error, the quadrature error of
+       the discrete transform, which shows up on the steep flanks of sharp
+       kernel peaks and is not aliasing at all — criterion 1 is blind to it.
+
+    Both criteria are scaled by the in-band **peak** PSD rather than applied
+    pointwise-relative. A pointwise-relative test is meaningless in the far
+    wings, where the PSD of a smooth kernel sits 10-12 orders of magnitude
+    below the peak, on the FFT's own round-off floor: that floor reshuffles
+    at O(1) relative level on every refinement and never settles, so a
+    pointwise test reports "did not converge" for even the most benign
+    kernel while the physically meaningful part of the spectrum is already
+    accurate to far better than ``rtol``.
+
+    A pure ``Constant`` kernel is the one genuinely non-convergent case: all
+    its power sits at DC, so it has no meaningful in-band PSD and the
+    warning below is the honest answer. In any sum with a real kernel its
+    contribution is stable and the warning does not fire.
+
+    See `localDocs/2026_08_20_ExpSineSquared_PSD_Aliasing_Analysis.pdf` for
+    a detailed analysis of the underlying aliasing issue.
+
+    **Normalisation convention** (matches Lomb-Scargle PSD in reports.py):
+
+    .. math::
+
+        S_{\\text{1-sided}}(f) = 2 \\,|\\tilde{k}(f)|\\,\\Delta\\tau
+        \\qquad \\text{for } f > 0
+
+    so that :math:`\\int_0^\\infty S(f)\\,df \\approx k(0)` (the kernel
+    variance).
+
+    Parameters
+    ----------
+    kernel : tinygp kernel object
+        Any tinygp kernel (dense or quasisep) exposing ``.evaluate(x1, x2)``.
+    freqs : 1-D array of float
+        Cyclic frequencies [d⁻¹] at which to evaluate the PSD. Must be
+        positive.
+    n_acf_lags : int, optional
+        Minimum number of lags used for the FFT at the *starting*
+        resolution (before adaptive refinement). Default 8192.
+    oversample : float, optional
+        Factor by which the initial lag-spacing Nyquist frequency exceeds
+        ``max(freqs)``. Larger values are safer for sharply peaked/high-Γ
+        periodic kernels but cost more compute on the first pass. Default
+        20.
+    max_refine : int, optional
+        Maximum number of dt_lag-halving refinement iterations. Default 6
+        (i.e. up to 64x finer than the initial spacing).
+    rtol : float, optional
+        Stability tolerance: the PSD at the requested ``freqs`` counts as
+        converged once successive refinements move it by less than ``rtol``
+        times its own in-band peak. Default 1e-3.
+    max_n_lags : int, optional
+        Hard cap on the number of lags, to bound memory/compute if
+        convergence is never reached. A warning is emitted if this cap is
+        hit without convergence. Default 2**20.
+    alias_rtol : float, optional
+        Aliasing tolerance: the PSD near the lag grid's Nyquist frequency
+        must fall below this fraction of the in-band peak for the lag
+        spacing to be declared fine enough. Default 1e-4, i.e. an order of
+        magnitude stricter than ``rtol``, so residual folded power stays
+        well below the level at which the result is declared converged —
+        while remaining orders of magnitude above the FFT's own round-off
+        floor (~1e-12 of the peak), so the test responds to real spectral
+        content rather than to numerical noise.
+
+    Returns
+    -------
+    psd : 1-D array of float
+        One-sided PSD evaluated at ``freqs``, in units of
+        ``[data_units]² · [x_units]`` (e.g. flux² · day when x is TIME).
+    """
+    _prev_x64 = jax.config.jax_enable_x64
+    jax.config.update("jax_enable_x64", True)
+    try:
+        freqs = np.asarray(freqs, dtype=np.float64)
+        f_max = float(np.max(freqs))
+        f_min = float(np.min(freqs[freqs > 0]))
+
+        def _eval_psd(dt_lag: float, n_lags: int) -> tuple[np.ndarray, np.ndarray]:
+            """FFT-based one-sided PSD at a given lag resolution."""
+            x_grid = jnp.arange(n_lags, dtype=jnp.float64) * dt_lag
+            origin = x_grid[0]  # = 0.0
+            k_tau = np.asarray(
+                jax.vmap(lambda tau: kernel.evaluate(origin, tau))(x_grid)
+            )
+
+            # Cosine taper to suppress Gibbs ringing from truncating the ACF.
+            window = 0.5 * (1.0 + np.cos(np.pi * np.arange(n_lags) / n_lags))
+            k_tau_w = k_tau * window
+
+            # Two-sided ACF: k(-tau) = k(tau)
+            k_sym = np.concatenate([k_tau_w[::-1][:-1], k_tau_w])
+            fft_vals = np.fft.fft(k_sym)
+            psd_twosided = np.abs(fft_vals) * dt_lag
+            fft_freqs = np.fft.fftfreq(len(k_sym), d=dt_lag)
+
+            pos_mask = fft_freqs > 0
+            f_pos = fft_freqs[pos_mask]
+            psd_onesided = 2.0 * psd_twosided[pos_mask]
+
+            sort_idx = np.argsort(f_pos)
+            return f_pos[sort_idx], psd_onesided[sort_idx]
+
+        def _alias_level(
+            f_grid: np.ndarray, psd_grid: np.ndarray, peak: float
+        ) -> float:
+            """Power near the lag grid's Nyquist, as a fraction of the
+            in-band peak.
+
+            Aliasing folds the spectrum above ``f_nyq = 1/(2 dt_lag)`` back
+            down into the computed band. That folded-in power is bounded by
+            the level the PSD has already reached as it approaches
+            ``f_nyq``, so probing the top octave of the computed grid gives
+            a direct, single-pass measure of how much contamination is still
+            possible — no iteration-to-iteration differencing needed.
+            """
+            if peak <= 0.0:
+                return 0.0
+            f_nyq = float(f_grid[-1])
+            # Keep the probe band clear of the requested band itself, which
+            # it would otherwise overlap for ``oversample`` near 1.
+            tail = psd_grid[f_grid >= max(0.5 * f_nyq, f_max)]
+            if tail.size == 0:
+                return 0.0
+            return float(np.max(tail)) / peak
+
+        # ── Initial resolution: oversampled well beyond max(freqs) ──────────
+        f_nyq_resolve = oversample * f_max
+        dt_lag = 1.0 / (2.0 * f_nyq_resolve)
+        # Ensure total lag span covers at least one cycle of the lowest
+        # requested frequency (needed for low-frequency resolution).
+        min_lags_needed = int(np.ceil(1.0 / (f_min * dt_lag))) + 1
+        n_lags = max(n_acf_lags, min_lags_needed)
+
+        # Initial evaluation
+        f_grid, psd_grid = _eval_psd(dt_lag, n_lags)
+        psd_prev = np.interp(freqs, f_grid, psd_grid, left=0.0, right=0.0)
+        peak = float(np.max(psd_prev))
+
+        alias_level = _alias_level(f_grid, psd_grid, peak)
+        stab_level = np.inf
+
+        # ``peak <= 0`` means the kernel has no power anywhere in the
+        # requested band (e.g. a multi-GP channel whose projections are both
+        # zero): trivially converged, since there is no scale to test
+        # anything against.
+        converged = peak <= 0.0
+
+        # ── Adaptive refinement: halve dt_lag, double n_lags (span fixed) ───
+        # At least one refinement has to run before the result can be called
+        # converged: ``stab_level`` is a difference between successive
+        # resolutions, so the first pass alone cannot establish it.
+        n_refined = 0
+        if not converged:
+            for _ in range(max_refine):
+                if n_lags * 2 > max_n_lags:
+                    break
+                dt_lag /= 2.0
+                n_lags *= 2
+                n_refined += 1
+
+                f_grid, psd_grid = _eval_psd(dt_lag, n_lags)
+                psd_cur = np.interp(freqs, f_grid, psd_grid, left=0.0, right=0.0)
+                peak = float(np.max(psd_cur))
+
+                if peak > 0.0:
+                    alias_level = _alias_level(f_grid, psd_grid, peak)
+                    stab_level = float(np.max(np.abs(psd_cur - psd_prev))) / peak
+                psd_prev = psd_cur
+
+                if peak <= 0.0 or (alias_level <= alias_rtol and stab_level <= rtol):
+                    converged = True
+                    break
+
+            if n_refined == 0:
+                # Refinement was disabled (``max_refine <= 0``) or blocked
+                # outright by ``max_n_lags``, so stability could never be
+                # measured; the aliasing test is then all there is to go on.
+                converged = alias_level <= alias_rtol
+
+        if not converged:
+            _log(
+                f"compute_kernel_theoretical_psd: PSD not converged after "
+                f"{n_refined} of max_refine={max_refine} refinements (final "
+                f"dt_lag={dt_lag:.3g}, n_lags={n_lags}): power near the "
+                f"lag-grid Nyquist is still {alias_level:.2e} of the in-band "
+                f"peak (target {alias_rtol:.0e}), and the last refinement "
+                f"still moved the in-band PSD by {stab_level:.2e} of that "
+                f"peak (target {rtol:.0e}). Result may still contain aliasing "
+                f"from unresolved kernel structure (e.g. very high Γ in a "
+                f"quasi-periodic term). Consider raising oversample, "
+                f"max_refine, or max_n_lags.",
+                level="WARNING",
+            )
+
+        return psd_prev
+    finally:
+        jax.config.update("jax_enable_x64", _prev_x64)
+
+
+# =============================================================================
+# Kernel realistic PSD computation taking into account the observational time window (for periodogram comparison)
+# =============================================================================
+
+
+def compute_windowed_psd(kernel: Any, freqs: Any, t: Any) -> jax.Array:
+    """
+    Compute the discrete one-sided PSD approximation of a tinygp kernel,
+    fully JAX-compatible and differentiable.
+    """
+    _prev_x64 = jax.config.jax_enable_x64
+    jax.config.update("jax_enable_x64", True)
+
+    try:
+        t = jnp.asarray(t, dtype=jnp.float64).reshape(-1)
+        freqs = jnp.asarray(freqs, dtype=jnp.float64)
+
+        # 1. Build the Covariance Matrix
+        K = kernel(t, t)
+
+        # # Safe JAX-compatible dimensionality check
+        # K = jnp.atleast_2d(K)
+        # if K.shape == (1, 1) and t.shape[0] > 1:
+        #     K = jnp.full((len(t), len(t)), K[0, 0])
+
+        # 2. Get Intrinsic Kernel Variance
+        kernel_variance = jnp.mean(jnp.diagonal(K))
+
+        # 3. Project into frequency space
+        # Using a more memory-efficient/readable outer product for the exponent
+        exponent = -2j * jnp.pi * jnp.outer(t, freqs)
+        projection_matrix = jnp.exp(exponent)
+
+        # Matrix quadratic form: raw_power = diag(P^H @ K @ P)
+        raw_power = jnp.real(
+            jnp.sum(jnp.conj(projection_matrix) * (K @ projection_matrix), axis=0)
+        )
+
+        # 4. Pure JAX Trapezoidal Integration
+        # ∫ P(f) df ≈ ∑ 0.5 * (P_i + P_{i+1}) * Δf
+        df = jnp.diff(freqs)
+        zero_padded_power = 0.5 * (raw_power[:-1] + raw_power[1:])
+        raw_integral = jnp.sum(zero_padded_power * df)
+
+        # 5. Normalize
+        psd_normalized = jnp.where(
+            raw_integral > 0, raw_power * (kernel_variance / raw_integral), raw_power
+        )
+
+        return psd_normalized
+
+    finally:
+        jax.config.update("jax_enable_x64", _prev_x64)