TGPA.py

# TGPA.py
# ============================================================
# TGPA: Takens Gaussian Process Autoregressor.
#
# Purpose:
#   A nonparametric baseline that removes NLSA entirely.
#   Uses raw Takens delay vectors z_t in R^{L*D}.
#
# Two modes:
#   (A) "direct_ic": ICM-only. Fit GP on prefix pairs and forecast.
#   (B) "global+ic_resid": optional global GPLM head trained on F_train,
#       plus in-context GP residual correction (LISA-style).
#
# Dependencies:
#   numpy, scipy
#   (optional) GPLM.py if you want the global head
# ============================================================

from __future__ import annotations
from typing import Optional, Tuple, Literal

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

from scipy.linalg import cho_factor, cho_solve, solve_triangular

try:
    from GPLM import GPLM  # type: ignore
except Exception:
    try:
        from gplm import GPLM  # type: ignore
    except Exception:
        GPLM = None  # allow "direct_ic" mode without GPLM


# -----------------------------
# utils
# -----------------------------
def _as_2d(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=float)
    if X.ndim == 1:
        X = X[:, None]
    return X


def _sliding_windows(F_tD: np.ndarray, L: int) -> np.ndarray:
    """
    Return windows W (K,L,D) from F (N,D) with K=N-L+1.
    """
    F = _as_2d(F_tD)
    W = sliding_window_view(F, window_shape=int(L), axis=0)

    a, b = W.shape[1], W.shape[2]
    if (a, b) == (L, F.shape[1]):
        return np.ascontiguousarray(W)
    if (a, b) == (F.shape[1], L):
        return np.ascontiguousarray(np.transpose(W, (0, 2, 1)))

    raise ValueError(f"Unexpected window shape {W.shape} for L={L}, D={F.shape[1]}.")


def _flatten_windows(W_BLD: np.ndarray) -> np.ndarray:
    """
    (B,L,D) -> (B, L*D)
    """
    W = np.asarray(W_BLD, dtype=np.float64)
    if W.ndim == 2:
        W = W[None, :, :]
    B, L, D = W.shape
    return np.ascontiguousarray(W.reshape(B, L * D))


def _pairwise_sq_dists(X: np.ndarray) -> np.ndarray:
    """
    Dense pairwise squared Euclidean distances (n,n).
    """
    X = np.asarray(X, dtype=np.float64)
    x2 = np.sum(X * X, axis=1, keepdims=True)
    d2 = x2 + x2.T - 2.0 * (X @ X.T)
    np.maximum(d2, 0.0, out=d2)
    return d2


def _estimate_rbf_ell_from_d2(d2_mat: np.ndarray, q: float = 0.5, eps: float = 1e-12) -> float:
    """
    Heuristic for RBF k = exp(-||x-y||^2/(2 ell^2)):
      ell^2 ~ quantile(d^2)/2
    """
    iu = np.triu_indices_from(d2_mat, k=1)
    vals = d2_mat[iu]
    if vals.size == 0:
        return 1.0
    qv = float(np.quantile(vals, q))
    return float(np.sqrt(max(qv, 0.0) / 2.0 + eps))


# ============================================================
# TDGP
# ============================================================

ICMode = Literal["direct_ic", "global+ic_resid"]


class TGPA:
    """
    TGPA: Takens GP Autoregressor.

    - direct_ic:
        Fit GP on prefix pairs (z_t -> x_{t+1}) and rollout.
        This is "ICM alone" (no training beyond kernel hypers).
    - global+ic_resid:
        Train a global GPLM head on training data once,
        then do in-context residual GP correction like LISA, but in raw delay space.
    """

    def __init__(
        self,
        F_train: Optional[np.ndarray] = None,
        *,
        L: int,
        ic_mode: ICMode = "direct_ic",

        # ---------- global head (optional) ----------
        gplm_kwargs: Optional[dict] = None,
        center_outputs: bool = True,

        # ---------- IC GP settings ----------
        gp_noise2: float = 1e-3,
        gp_rbf_ell: Optional[float] = None,
        gp_rbf_q: float = 0.5,

        # cap context points for dense IC solve
        ctx_max_points: int = 1000,
        ctx_k0: float = 10.0,
        ctx_min_windows: Optional[int] = None,

        # variance trust gating
        use_var_gate: bool = True,
        gate_tau2: float = 1.0,
        gate_mode: Literal["rational", "exp"] = "rational",

        seed: int = 0,
    ):
        self.L = int(L)
        self.ic_mode = str(ic_mode)

        self._rng = np.random.default_rng(int(seed))

        self.gp_noise2 = float(gp_noise2)
        self.gp_rbf_ell = None if gp_rbf_ell is None else float(gp_rbf_ell)
        self.gp_rbf_q = float(gp_rbf_q)

        self.ctx_max_points = int(ctx_max_points)
        self.ctx_k0 = float(ctx_k0)
        self.ctx_min_windows = int(ctx_min_windows) if ctx_min_windows is not None else max(2, 8)

        self.use_var_gate = bool(use_var_gate)
        self.gate_tau2 = float(gate_tau2)
        self.gate_mode = str(gate_mode).lower().strip()

        # global head (optional)
        self.has_global = False
        self.center_outputs = bool(center_outputs)
        self.mu_X: Optional[np.ndarray] = None
        self.gplm: Optional[GPLM] = None

        if self.ic_mode == "global+ic_resid":
            if F_train is None:
                raise ValueError("global+ic_resid requires F_train.")
            if GPLM is None:
                raise ImportError("GPLM not importable, but global+ic_resid requires GPLM.")

            F_train = _as_2d(F_train)
            self.D = int(F_train.shape[1])

            # build training pairs (z_t -> x_{t+1})
            W_all = _sliding_windows(F_train, self.L)  # (K,L,D), K=N-L+1
            K = W_all.shape[0]
            N_pairs = K - 1
            X_train = _flatten_windows(W_all[:N_pairs])  # (K-1, L*D)
            Y_train = np.asarray(F_train[self.L:self.L + N_pairs, :], dtype=np.float64)  # (K-1, D)

            if self.center_outputs:
                self.mu_X = Y_train.mean(axis=0)
                Yc = Y_train - self.mu_X[None, :]
            else:
                self.mu_X = np.zeros((self.D,), dtype=np.float64)
                Yc = Y_train

            if gplm_kwargs is None:
                gplm_kwargs = {}
            gkw = dict(gplm_kwargs)
            gkw.setdefault("center_X", False)
            gkw.setdefault("sigma2", 1e-5)
            gkw.setdefault("jitter", 1e-8)
            gkw.setdefault("m", min(2048, X_train.shape[0]))
            gkw.setdefault("inducing", "fps")
            gkw.setdefault("seed", int(seed))

            self.gplm = GPLM(X_train, Yc, **gkw)
            self.has_global = True

        else:
            # direct_ic mode: we only need D at call-time
            self.D = -1

    # -----------------------------
    # kernel bits
    # -----------------------------
    def _rbf_kernel_matrix(self, X: np.ndarray) -> Tuple[np.ndarray, float]:
        d2 = _pairwise_sq_dists(X)
        ell = _estimate_rbf_ell_from_d2(d2, q=self.gp_rbf_q) if self.gp_rbf_ell is None else float(self.gp_rbf_ell)
        K = np.exp(-0.5 * d2 / (ell**2 + 1e-12))
        return K, ell

    def _rbf_kernel_eval(self, X: np.ndarray, xq: np.ndarray, ell: float) -> np.ndarray:
        diff = X - xq[None, :]
        d2 = np.einsum("nd,nd->n", diff, diff, optimize=True)
        return np.exp(-0.5 * d2 / (ell**2 + 1e-12))

    def _gate_from_var(self, var_f: float) -> float:
        if not self.use_var_gate:
            return 1.0
        v = max(float(var_f), 0.0)
        tau2 = max(float(self.gate_tau2), 1e-18)
        if self.gate_mode == "exp":
            return float(np.exp(-v / tau2))
        return float(tau2 / (tau2 + v))

    # -----------------------------
    # global prediction
    # -----------------------------
    def _global_pred(self, Z_Bp: np.ndarray) -> np.ndarray:
        assert self.gplm is not None
        Yc = self.gplm(Z_Bp)  # centered
        if self.center_outputs:
            return Yc + self.mu_X[None, :]
        return Yc

    # ============================================================
    # main forecast
    # ============================================================
    def __call__(self, prefix: np.ndarray, *, steps: int = 1, return_var: bool = False):
        """
        prefix: (ell,D), ell>=L
        returns preds: (steps,D)
        """
        prefix = _as_2d(prefix)
        ell, D = prefix.shape
        if ell < self.L:
            raise ValueError(f"Need prefix length ell >= L={self.L}.")

        if self.D < 0:
            self.D = int(D)
        if D != self.D:
            raise ValueError(f"TDGP expects D={self.D}, got D={D}.")

        H = int(steps)
        if H <= 0:
            out = np.zeros((0, D), dtype=np.float64)
            return (out, np.zeros((0,), dtype=np.float64)) if return_var else out

        # seed rolling window
        cur = prefix[-self.L:, :].copy()

        # build context pairs from prefix
        K_ctx = ell - self.L
        if K_ctx < self.ctx_min_windows:
            # not enough context: fall back
            if self.ic_mode == "global+ic_resid":
                return self._rollout_global(cur, H, return_var=return_var)
            raise ValueError("direct_ic needs ell > L with enough context pairs.")

        W_all = _sliding_windows(prefix, self.L)          # (ell-L+1, L, D)
        W_ctx = np.ascontiguousarray(W_all[:K_ctx, :, :])  # windows with observed next step
        X_ctx = _flatten_windows(W_ctx)                  # (K_ctx, L*D)
        Y_ctx = np.asarray(prefix[self.L:self.L + K_ctx, :], dtype=np.float64)  # (K_ctx, D)

        # possibly subsample context to keep dense GP affordable
        M = min(int(K_ctx), int(self.ctx_max_points))
        if M < K_ctx:
            # simple but effective: pick evenly spaced indices
            idx = np.linspace(0, K_ctx - 1, M).round().astype(np.int64)
            X_fit = X_ctx[idx]
            Y_fit = Y_ctx[idx]
        else:
            X_fit = X_ctx
            Y_fit = Y_ctx

        # choose targets for IC solve
        if self.ic_mode == "global+ic_resid":
            # residual targets
            Y_glob_fit = self._global_pred(X_fit)  # (M,D)
            T_fit = Y_fit - Y_glob_fit
        else:
            # direct map
            T_fit = Y_fit

        # fit dense GP on (X_fit -> T_fit)
        K_mat, ell_used = self._rbf_kernel_matrix(X_fit)
        K_reg = K_mat + self.gp_noise2 * np.eye(M, dtype=np.float64)

        cf = cho_factor(K_reg, lower=True, check_finite=False)
        alpha = cho_solve(cf, T_fit, check_finite=False)  # (M,D)
        Lfac, lower = cf

        # base context mixing weight
        w_ctx_base = float(M) / float(M + self.ctx_k0) if self.ctx_k0 > 0 else 1.0

        preds = np.zeros((H, D), dtype=np.float64)
        vars_out = np.zeros((H,), dtype=np.float64) if return_var else None

        # rollout
        for h in range(H):
            zq = _flatten_windows(cur)[0]  # (L*D,)

            if self.ic_mode == "global+ic_resid":
                y_base = self._global_pred(zq[None, :])[0]
            else:
                y_base = np.zeros((D,), dtype=np.float64)

            k_eval = self._rbf_kernel_eval(X_fit, zq, ell_used)     # (M,)
            t_mean = k_eval @ alpha                                  # (D,)

            # function variance proxy
            u = solve_triangular(Lfac, k_eval, lower=lower, check_finite=False)
            quad = float(np.dot(u, u))
            var_f = max(0.0, 1.0 - quad)

            if return_var:
                vars_out[h] = var_f

            w_gate = self._gate_from_var(var_f)
            w_eff = w_ctx_base * w_gate

            if self.ic_mode == "global+ic_resid":
                y = y_base + w_eff * t_mean
            else:
                # direct_ic: optionally apply mixing too (stabilizes rollout)
                y = w_eff * t_mean + (1.0 - w_eff) * y_base  # y_base is zeros here

            preds[h] = y

            # update rolling window
            if self.L > 1:
                cur[:-1] = cur[1:]
            cur[-1] = y

        if return_var:
            return preds, vars_out
        return preds

    def _rollout_global(self, seed_LD: np.ndarray, H: int, return_var: bool = False):
        """
        Pure global rollout (only for global+ic_resid mode).
        """
        assert self.gplm is not None
        cur = np.asarray(seed_LD, dtype=np.float64).copy()
        out = np.zeros((H, self.D), dtype=np.float64)

        for h in range(int(H)):
            zq = _flatten_windows(cur)[0]
            y = self._global_pred(zq[None, :])[0]
            out[h] = y
            if self.L > 1:
                cur[:-1] = cur[1:]
            cur[-1] = y

        if return_var:
            return out, np.zeros((H,), dtype=np.float64)
        return out


__all__ = ["TGPA"]