Create LISA.py

LISA.py

@@ -0,0 +1,535 @@
+# LISA.py
+# ============================================================
+# LISA = NLSA encoder + GPLM baseline forecaster
+#      + in-context full Gaussian Process Regression (GPR)
+#        on residuals in diffusion (psi) space.
+#
+# Auto behavior:
+#   - if prefix length ell == L: baseline-only AR rollout
+#   - if ell > L: baseline + GP residual IC correction (LISA)
+#
+# Repo expectations:
+#   - NLSAEncoder available as:
+#        from nlsa_encoder import NLSAEncoder
+#     (or from NLSA import NLSAEncoder depending on your naming)
+#   - GPLM available as:
+#        from gplm import GPLM
+#
+# Dependencies:
+#   numpy, scipy
+# ============================================================
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import scipy.linalg as la
+from numpy.lib.stride_tricks import sliding_window_view
+
+
+# ------------------------------------------------------------
+# Imports from your repo (robust fallbacks)
+# ------------------------------------------------------------
+
+# NLSA encoder
+try:
+    from .nlsa_encoder import NLSAEncoder  # type: ignore
+except Exception:
+    try:
+        from nlsa_encoder import NLSAEncoder  # type: ignore
+    except Exception:
+        # if you named the file NLSA.py
+        from NLSA import NLSAEncoder  # type: ignore
+
+# GPLM decoder (GP/KRR mean map)
+try:
+    from .gplm import GPLM  # type: ignore
+except Exception:
+    try:
+        from gplm import GPLM  # type: ignore
+    except Exception:
+        # if you named the file GPLM.py
+        from GPLM import GPLM  # type: ignore
+
+
+Array = np.ndarray
+
+
+# ============================================================
+# Utilities
+# ============================================================
+
+def _as_2d(X: Array) -> Array:
+    X = np.asarray(X, dtype=float)
+    if X.ndim == 1:
+        X = X[:, None]
+    return X
+
+
+def _sliding_windows(F_tD: Array, L: int) -> Array:
+    """
+    Return windows W (K,L,D) from F (N,D) with K=N-L+1.
+    Handles numpy stride ordering differences.
+    """
+    F = _as_2d(F_tD)
+    W = sliding_window_view(F, window_shape=int(L), axis=0)
+
+    # W can be (K,L,D) or (K,D,L)
+    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 _pairwise_d2(X: Array) -> Array:
+    """
+    Dense pairwise squared distances (n,n) for moderate 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 _median_ell_from_d2(d2_mat: Array, q: float = 0.5, eps: float = 1e-12, min_ell: float = 1e-6) -> float:
+    """
+    If k(x,y) = exp(-||x-y||^2/(2 ell^2)), heuristic:
+        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
+    v = float(np.quantile(vals, q))
+    ell = np.sqrt(max(v / 2.0, eps))
+    return float(max(ell, min_ell))
+
+
+# ============================================================
+# Kernel helpers (for the in-context GP on residuals)
+# ============================================================
+
+def _kernel_matrix(Psi: Array, *, kind: str, ell: Optional[float]) -> Tuple[Array, Optional[float]]:
+    """
+    Build dense kernel matrix K(Psi,Psi).
+    kind: "linear" or "rbf"
+    ell: lengthscale for RBF. If None -> caller estimates separately.
+    """
+    kind = kind.lower().strip()
+    Psi = np.asarray(Psi, dtype=np.float64)
+
+    if kind == "linear":
+        return Psi @ Psi.T, None
+
+    if kind != "rbf":
+        raise ValueError("kernel kind must be 'linear' or 'rbf'.")
+
+    if ell is None:
+        raise ValueError("RBF kernel requires ell != None (estimate ell before calling).")
+
+    d2 = _pairwise_d2(Psi)
+    K = np.exp(-0.5 * d2 / (float(ell) ** 2 + 1e-12))
+    return K, float(ell)
+
+
+def _kernel_eval(Psi_ctx: Array, psi: Array, *, kind: str, ell: Optional[float]) -> Array:
+    """
+    Vector k(Psi_ctx, psi) shape (K_ctx,).
+    """
+    kind = kind.lower().strip()
+    Psi_ctx = np.asarray(Psi_ctx, dtype=np.float64)
+    psi = np.asarray(psi, dtype=np.float64).reshape(-1)
+
+    if kind == "linear":
+        return Psi_ctx @ psi
+
+    if kind != "rbf":
+        raise ValueError("kernel kind must be 'linear' or 'rbf'.")
+
+    if ell is None:
+        raise ValueError("RBF kernel requires ell != None.")
+
+    diff = Psi_ctx - psi[None, :]
+    d2 = np.einsum("kr,kr->k", diff, diff, optimize=True)
+    return np.exp(-0.5 * d2 / (float(ell) ** 2 + 1e-12))
+
+
+def _k_qq(psi: Array, *, kind: str) -> float:
+    """k(psi,psi)."""
+    kind = kind.lower().strip()
+    if kind == "linear":
+        return float(np.dot(psi, psi))
+    return 1.0  # RBF
+
+
+def _gate_from_var(var: float, *, tau2: float, mode: str) -> float:
+    """
+    Convert predictive variance into a [0,1] trust weight.
+    - "rational": tau2/(tau2+var)
+    - "exp": exp(-var/tau2)
+    """
+    v = max(float(var), 0.0)
+    tau2 = max(float(tau2), 1e-18)
+    mode = mode.lower().strip()
+
+    if mode == "exp":
+        return float(np.exp(-v / tau2))
+    if mode == "rational":
+        return float(tau2 / (tau2 + v))
+    raise ValueError("gate_mode must be 'rational' or 'exp'.")
+
+
+# ============================================================
+# LISA
+# ============================================================
+
+@dataclass
+class LISAConfig:
+    # NLSA encoder
+    L: int = 128
+    rank: int = 32
+    beta: Optional[float] = None
+    alpha: float = 1.0
+    center: bool = True
+    drop_first: bool = True
+    max_K_dense: int = 6000
+    seed: int = 0
+
+    # GPLM baseline (psi -> next sample)
+    gplm_kwargs: Optional[Dict[str, Any]] = None
+
+    # In-context GP residual model
+    ctx_min_windows: Optional[int] = None  # default r+1
+    ctx_k0: float = 10.0                  # base blending gate: K_ctx/(K_ctx+k0)
+
+    gp_kernel: str = "rbf"                # "rbf" or "linear"
+    gp_noise2: float = 1e-3               # σ_n^2
+    gp_rbf_ell: Optional[float] = None    # if None -> estimate from context
+    gp_rbf_q: float = 0.5                 # quantile for ell estimation
+    gp_jitter: float = 1e-10              # add to diag for numeric stability
+
+    # Variance gating
+    use_var_gate: bool = True
+    gate_tau2: float = 1.0
+    gate_mode: str = "rational"           # "rational" or "exp"
+
+
+class LISA:
+    """
+    LISA = NLSAEncoder + GPLM baseline + in-context full GP residual correction.
+
+    Training (`__init__`):
+      1) Train NLSAEncoder on F_train
+      2) Fit GPLM baseline: psi(window) -> next sample (centered space)
+
+    Inference (`__call__`):
+      Given prefix F_prefix (ell,D):
+        - If ell == L: baseline-only AR rollout (NLSA+GPLM)
+        - If ell > L: use prefix context to fit residual GP, then rollout:
+            y = y_glob + w_eff * e_gp
+          where w_eff = w_ctx_base * trust(var_gp)
+
+    Notes:
+      - GP residual is multi-output via shared kernel and independent outputs:
+          E ~ GP(0, k(ψ,ψ')) per output dimension.
+      - Predictive variance is scalar (same for all dims) because kernel is shared.
+    """
+
+    def __init__(self, F_train: Array, *, config: Optional[LISAConfig] = None, **kwargs: Any):
+        if config is None:
+            config = LISAConfig(**kwargs)
+        self.cfg = config
+
+        F_train = _as_2d(F_train)
+        self.D = int(F_train.shape[1])
+        self.L = int(self.cfg.L)
+
+        # ---- 1) NLSA encoder ----
+        self.enc = NLSAEncoder(
+            F_train,
+            L=int(self.cfg.L),
+            rank=int(self.cfg.rank),
+            beta=self.cfg.beta,
+            alpha=float(self.cfg.alpha),
+            center=bool(self.cfg.center),
+            drop_first=bool(self.cfg.drop_first),
+            max_K_dense=int(self.cfg.max_K_dense),
+            seed=int(self.cfg.seed),
+        )
+
+        self.r = 0 if self.enc.psi_ is None else int(self.enc.psi_.shape[1])
+        if self.r <= 0:
+            raise ValueError("NLSAEncoder produced r=0 diffusion dims. Increase rank or set drop_first=False.")
+
+        # mean (same mean used inside encoder)
+        self.center = bool(self.cfg.center)
+        self.mu_X = self.enc.mu_.reshape(-1).astype(np.float64) if self.center else np.zeros((self.D,), dtype=np.float64)
+
+        # ---- 2) GPLM baseline: psi -> next sample ----
+        # Training pairs: for each training window T=0..K-2, target is F[T+L]
+        # We can use the training diffusion coords directly: Psi_tr = enc.psi_[:K-1]
+        K = int(self.enc.K_)
+        n_pairs = K - 1
+        Psi_tr = np.asarray(self.enc.psi_[:n_pairs, :], dtype=np.float64)  # (K-1, r)
+
+        # Targets in centered space consistent with encoder centering
+        F_centered = _as_2d(F_train) - self.mu_X[None, :] if self.center else _as_2d(F_train)
+        Y_tr = np.asarray(F_centered[self.L : self.L + n_pairs, :], dtype=np.float64)  # (K-1, D)
+
+        gkw = {} if self.cfg.gplm_kwargs is None else dict(self.cfg.gplm_kwargs)
+        # Critical: Y_tr is already centered (if center=True), so do NOT let GPLM recenter outputs.
+        gkw.setdefault("center_X", False)
+        gkw.setdefault("seed", int(self.cfg.seed))
+
+        # A sane default regularization if user didn't provide one:
+        gkw.setdefault("sigma2", 1e-4)
+        gkw.setdefault("jitter", 1e-8)
+
+        # inducing size default
+        gkw.setdefault("m", min(1024, n_pairs))
+        gkw.setdefault("inducing", "kmeans_medoids")
+
+        self.baseline = GPLM(Psi_tr, Y_tr, **gkw)
+
+        # IC defaults
+        if self.cfg.ctx_min_windows is None:
+            self.ctx_min_windows = max(1, self.r + 1)
+        else:
+            self.ctx_min_windows = int(self.cfg.ctx_min_windows)
+
+        self._rng = np.random.default_rng(int(self.cfg.seed))
+
+        # for diagnostics
+        self.last_gp_ell_ = None
+
+    # -------------------------
+    # baseline step
+    # -------------------------
+
+    def predict_one_step_baseline(self, W_LD: Array) -> Array:
+        """
+        One-step baseline prediction from a window (L,D) -> (D,).
+        Uses: psi = enc.encode_window(window), then GPLM(psi).
+        """
+        W = np.asarray(W_LD, dtype=float)
+        if W.ndim == 1:
+            W = W[:, None]
+        if W.shape != (self.L, self.D):
+            raise ValueError(f"Expected window shape {(self.L, self.D)}, got {W.shape}")
+
+        psi = self.enc.encode_window(W)  # (r,)
+        y_c = np.asarray(self.baseline(psi), dtype=np.float64).reshape(-1)  # centered
+        y = y_c + self.mu_X if self.center else y_c
+        return y
+
+    # -------------------------
+    # main rollout
+    # -------------------------
+
+    def __call__(
+        self,
+        prefix: Array,
+        steps: int,
+        *,
+        return_var: bool = False,
+        sample: bool = False,
+        rng: Optional[np.random.Generator] = None,
+        include_obs_noise: bool = True,
+    ) -> Union[Array, Tuple[Array, Array]]:
+        """
+        Forecast from prefix context.
+
+        prefix: (ell,D), ell >= L
+        steps: horizon H
+
+        If ell == L: baseline-only AR rollout.
+        If ell > L : full LISA with dense GP residual IC correction.
+
+        return_var:
+          returns scalar residual GP variance per step (H,)
+
+        sample:
+          samples residual from GP posterior at each step (generative mode).
+        """
+        prefix = _as_2d(prefix)
+        ell, D = prefix.shape
+        if D != self.D:
+            raise ValueError(f"LISA trained with D={self.D}, got prefix D={D}.")
+        if ell < self.L:
+            raise ValueError(f"Need prefix length ell >= L={self.L}.")
+
+        H = int(steps)
+        if H <= 0:
+            preds = np.zeros((0, self.D), dtype=np.float64)
+            return (preds, np.zeros((0,), dtype=np.float64)) if return_var else preds
+
+        if rng is None:
+            rng = self._rng
+
+        # seed window always last L
+        cur = prefix[-self.L :, :].copy()  # (L,D)
+
+        # ----------------------------------------------------
+        # If no extra context, run baseline-only AR
+        # ----------------------------------------------------
+        if ell == self.L or self.r <= 0:
+            preds = self._rollout_baseline(cur, H)
+            if return_var:
+                return preds, np.zeros((H,), dtype=np.float64)
+            return preds
+
+        # ----------------------------------------------------
+        # Build context windows and targets from prefix
+        # ----------------------------------------------------
+        W_all = _sliding_windows(prefix, self.L)  # (K_n, L, D)
+        K_n = int(W_all.shape[0])
+        K_ctx = K_n - 1  # number of (window -> next) pairs inside prefix
+
+        if K_ctx < self.ctx_min_windows:
+            preds = self._rollout_baseline(cur, H)
+            if return_var:
+                return preds, np.zeros((H,), dtype=np.float64)
+            return preds
+
+        W_ctx = np.ascontiguousarray(W_all[:K_ctx, :, :])             # (K_ctx,L,D)
+        Y_true = np.asarray(prefix[self.L : self.L + K_ctx, :], float)  # (K_ctx,D)
+
+        # centered targets
+        Y_true_c = (Y_true - self.mu_X[None, :]) if self.center else Y_true
+
+        # encode context
+        Psi_ctx = np.zeros((K_ctx, self.r), dtype=np.float64)
+        for i in range(K_ctx):
+            Psi_ctx[i] = self.enc.encode_window(W_ctx[i])
+
+        # baseline on context
+        Y_glob_ctx_c = np.asarray(self.baseline(Psi_ctx), dtype=np.float64)  # (K_ctx,D) centered
+
+        # residual table
+        E_ctx = Y_true_c - Y_glob_ctx_c  # (K_ctx,D)
+
+        # ----------------------------------------------------
+        # Fit dense GP on residuals in psi-space
+        # ----------------------------------------------------
+        kernel = self.cfg.gp_kernel.lower().strip()
+
+        if kernel == "rbf":
+            ell_used = self.cfg.gp_rbf_ell
+            if ell_used is None:
+                d2 = _pairwise_d2(Psi_ctx)
+                ell_used = _median_ell_from_d2(d2, q=float(self.cfg.gp_rbf_q))
+            ell_used = float(ell_used)
+            self.last_gp_ell_ = ell_used
+        else:
+            ell_used = None
+            self.last_gp_ell_ = None
+
+        K_mat, _ = _kernel_matrix(Psi_ctx, kind=kernel, ell=ell_used)
+
+        # (K + σ_n^2 I + jitter I)
+        noise2 = float(self.cfg.gp_noise2)
+        jitter = float(self.cfg.gp_jitter)
+        K_reg = K_mat + (noise2 + jitter) * np.eye(K_ctx, dtype=np.float64)
+
+        # cholesky
+        try:
+            cF = la.cho_factor(K_reg, lower=True, check_finite=False)
+            alpha = la.cho_solve(cF, E_ctx, check_finite=False)  # (K_ctx, D)
+            Lfac, lower = cF
+            chol_ok = True
+        except la.LinAlgError:
+            # fall back to solve (no variance)
+            alpha = np.linalg.solve(K_reg, E_ctx)
+            Lfac, lower = None, True
+            chol_ok = False
+
+        # base mixing (more context => stronger)
+        k0 = float(self.cfg.ctx_k0)
+        w_ctx_base = float(K_ctx) / float(K_ctx + k0) if k0 > 0 else 1.0
+
+        # ----------------------------------------------------
+        # Rollout with IC correction
+        # ----------------------------------------------------
+        preds = np.zeros((H, self.D), dtype=np.float64)
+        vars_out = np.zeros((H,), dtype=np.float64) if return_var else None
+
+        for h in range(H):
+            psi = self.enc.encode_window(cur)  # (r,)
+
+            # baseline
+            y_glob_c = np.asarray(self.baseline(psi), dtype=np.float64).reshape(-1)  # (D,)
+
+            # GP residual mean
+            k_eval = _kernel_eval(Psi_ctx, psi, kind=kernel, ell=ell_used)  # (K_ctx,)
+            e_mean = k_eval @ alpha  # (D,)
+
+            # GP residual variance (scalar)
+            var_f = 0.0
+            if chol_ok and Lfac is not None:
+                # u = L^{-1} k ; quad = ||u||^2
+                u = la.solve_triangular(Lfac, k_eval, lower=True, check_finite=False)
+                quad = float(np.dot(u, u))
+                var_f = max(0.0, _k_qq(psi, kind=kernel) - quad)
+
+            if return_var:
+                vars_out[h] = float(var_f)
+
+            # variance trust gate
+            if self.cfg.use_var_gate:
+                w_gate = _gate_from_var(var_f, tau2=self.cfg.gate_tau2, mode=self.cfg.gate_mode)
+            else:
+                w_gate = 1.0
+
+            w_eff = w_ctx_base * float(w_gate)
+
+            # optional sampling (generative)
+            e_use = e_mean
+            if sample:
+                var_y = var_f + (noise2 if include_obs_noise else 0.0)
+                var_y = max(0.0, float(var_y))
+                if var_y > 0:
+                    e_use = e_mean + np.sqrt(var_y) * rng.standard_normal(size=(self.D,))
+
+            # combine
+            y_c = y_glob_c + w_eff * e_use
+            y = y_c + self.mu_X if self.center else y_c
+
+            preds[h] = y
+
+            # advance seed window
+            if self.L > 1:
+                cur[:-1] = cur[1:]
+            cur[-1] = y
+
+        if return_var:
+            return preds, vars_out
+        return preds
+
+    def _rollout_baseline(self, seed_LD: Array, H: int) -> Array:
+        """
+        Baseline-only autoregressive rollout using (NLSA encode) + GPLM decode.
+        """
+        cur = np.asarray(seed_LD, dtype=float).copy()
+        out = np.zeros((int(H), self.D), dtype=np.float64)
+        for h in range(int(H)):
+            y = self.predict_one_step_baseline(cur)  # (D,)
+            out[h] = y
+            if self.L > 1:
+                cur[:-1] = cur[1:]
+            cur[-1] = y
+        return out
+
+    def __repr__(self) -> str:
+        return (
+            f"LISA(L={self.L}, D={self.D}, r={self.r}, "
+            f"center={self.center}, gp_kernel={self.cfg.gp_kernel}, "
+            f"gp_noise2={self.cfg.gp_noise2})"
+        )
+
+
+__all__ = ["LISA", "LISAConfig"]