Update LISAx.py

LISAx.py

@@ -1,13 +1,18 @@
 # LISA.py
 # ============================================================
-# LISA: NLSA encoder + GPLM nonlinear baseline decoder
-#       + dense in-context Gaussian Process Regression (GPR)
+# LISA: NLSA encoder + GPLM baseline decoder
+#       + (optional) dense in-context Gaussian Process Regression (GPR)
 #       on residuals in diffusion coordinate space.
 #
-# This version adds:
-#   - Automatic inducing/anchor selection for GPLM (m auto-scales with N)
-#     based on a compute budget (default 1000**3).
-#   - Optional cap on IC context windows to avoid O(K_ctx^2) blowups at call-time.
+# This version is compatible with the "sparse-kernel KRR" GPLM.py:
+#   - GPLM trains by building a sparse kNN kernel on training latents ψ
+#   - solves (K + sigma2 I) S = Y (mean GP / KRR)
+#   - predicts using only pred_k neighbors per query
+#
+# Key advantages vs Nyström GPLM:
+#   - No anchors/inducing points required
+#   - Strong nonlinearity => very local sparse kernel => scalable
+#
 # ============================================================
 
 from __future__ import annotations
@@ -29,12 +34,9 @@ except Exception:
     from NLSA import NLSA  # type: ignore
 
 try:
-    from .GPLM import GPLM  # type: ignore
+    from .GPLMx import GPLM  # type: ignore
 except Exception:
-    try:
-        from GPLM import GPLM  # type: ignore
-    except Exception:
-        from gplm import GPLM  # type: ignore
+    from GPLMx import GPLM  # type: ignore
 
 
 # ============================================================
@@ -95,65 +97,25 @@ def _estimate_rbf_ell_from_d2(
     return float(max(ell, min_ell))
 
 
-def _auto_gplm_m(
-    n_train: int,
-    *,
-    budget: int = 1000**3,
-    m_min: int = 32,
-    m_max: Optional[int] = None,
-) -> int:
-    """
-    Choose inducing count m for GPLM given training size n_train,
-    constrained by an approximate compute budget.
-
-    GPLM training costs (dominant terms):
-      - Accumulating G = C^T C:     O(n_train * m^2)
-      - Cholesky on A (m x m):      O(m^3)
-
-    We choose:
-      m <= budget^(1/3)    (controls m^3)
-      m <= sqrt(budget/n)  (controls n*m^2)
-
-    and of course:
-      m <= n_train
-      m >= m_min
-
-    This yields:
-      - n=1000  -> m=1000 (exact within budget)
-      - n=5000  -> m≈447
-      - n=50000 -> m≈141
-    """
-    n = int(max(1, n_train))
-    b = int(max(1, budget))
-
-    m_cubic = int(np.floor(b ** (1.0 / 3.0)))              # budget for m^3
-    m_nm2 = int(np.floor(np.sqrt(b / float(n))))           # budget for n*m^2
-
-    m = min(n, m_cubic, m_nm2)
-    m = max(int(m_min), int(m))
-
-    if m_max is not None:
-        m = min(m, int(m_max))
-
-    # final sanity
-    m = int(max(1, min(m, n)))
-    return m
-
-
 # ============================================================
 # LISA main class
 # ============================================================
 
 class LISA:
     """
-    LISA: NLSA encoder + GPLM baseline + dense GPR residual correction.
+    LISA: NLSA encoder + GPLM baseline + optional dense GPR residual correction.
+
+    - Baseline:
+        window W (L,D) -> ψ(W) in R^r (via NLSA Nyström OOS)
+        y_glob = GPLM(ψ) in R^D
 
-    - If prefix length ℓ == L: no in-context pairs -> baseline AR only.
-    - If ℓ > L: residual GP is fit on the prefix windows and applied during rollout.
+    - In-context (optional):
+        Given prefix length ℓ > L, fit dense GP on residuals e(ψ)
+        using context windows inside prefix, then apply during rollout.
 
-    Notes:
-    - Residual GP is multi-output with independent output dims sharing a kernel.
-      That means GP mean is vector in R^D, but variance is a scalar (same for all dims).
+    Note:
+      The IC GP here is dense and scales O(K_ctx^2) memory / O(K_ctx^3) time.
+      A default cap gp_max_ctx prevents accidental blowups with huge prefixes.
     """
 
     def __init__(
@@ -171,11 +133,6 @@ class LISA:
         nlsa_kwargs: Optional[dict] = None,
         # ------------------ GPLM baseline decoder hyperparams ----------
         gplm_kwargs: Optional[dict] = None,
-        # --- NEW: auto anchor/inducing scaling for large N -------------
-        gplm_auto_m: bool = True,
-        gplm_budget: int = 1000**3,   # your requested budget baseline
-        gplm_m_min: int = 32,
-        gplm_m_max: Optional[int] = None,
         # ------------------ IC / GPR controls -------------------------
         ctx_min_windows: Optional[int] = None,
         ctx_k0: float = 10.0,                 # base mixing K_ctx/(K_ctx+ctx_k0)
@@ -183,8 +140,8 @@ class LISA:
         gp_kernel: str = "rbf",               # "rbf" or "linear"
         gp_rbf_ell: float | None = None,      # fixed, or auto from context
         gp_rbf_q: float = 0.5,
-        # --- NEW: prevent call-time blowups when prefix is huge --------
-        gp_max_ctx: Optional[int] = 4096,     # cap context windows used by dense GP
+        # ---- NEW safety: cap dense GP context used in __call__ -------
+        gp_max_ctx: Optional[int] = 4096,     # None disables cap
         gp_ctx_mode: str = "recent",          # "recent" or "uniform"
         # ------------------ stability / trust gating -----------------
         use_var_gate: bool = True,
@@ -219,7 +176,7 @@ class LISA:
         self.D = int(base.D_)
         self.r = int(base.psi_.shape[1])
 
-        # mean handling: baseline GPLM will be trained in *centered output space*
+        # mean handling (consistent w/ NLSA centering)
         self.center_outputs = bool(base.center)
         self.mu_X = base.mu_.reshape(-1).astype(np.float64) if self.center_outputs else np.zeros((self.D,), dtype=np.float64)
 
@@ -235,32 +192,46 @@ class LISA:
             dtype=np.float64
         )  # (K-1,D) centered if base.center
 
+        # ---- Defaults tuned for sparse-kernel GPLM ----
         gkw = dict(gplm_kwargs)
         gkw.setdefault("seed", int(seed))
-        gkw.setdefault("center_X", False)         # IMPORTANT: Y_train is already centered
-        gkw.setdefault("sigma2", 1e-5)
+
+        # IMPORTANT: LISA already provides centered Y, so disable GPLM centering
+        gkw.setdefault("center_X", False)
+
+        # Regularization for sparse KRR (lambda in K + lambda I)
+        # (You may tune this up if solves are noisy/unstable.)
+        gkw.setdefault("sigma2", 1e-4)
+
+        # Numerical diagonal stabilizer
         gkw.setdefault("jitter", 1e-8)
-        gkw.setdefault("inducing", "kmeans_medoids")
-        gkw.setdefault("dtype", np.float32)
-        gkw.setdefault("fit_block", 8192)
-
-        # ------------------------------------------------------------
-        # NEW: automatic inducing m scaling (for N up to 50k+)
-        # ------------------------------------------------------------
-        if "m" not in gkw:
-            if bool(gplm_auto_m):
-                m_auto = _auto_gplm_m(
-                    X_train.shape[0],
-                    budget=int(gplm_budget),
-                    m_min=int(gplm_m_min),
-                    m_max=gplm_m_max,
-                )
-                gkw["m"] = int(m_auto)
-            else:
-                gkw["m"] = min(1024, X_train.shape[0])
 
-        self.gplm_m_ = int(gkw["m"])
-        self.gplm_budget_ = int(gplm_budget)
+        # Sparse kernel build on ψ_train
+        gkw.setdefault("k_graph", 64)          # training graph sparsity
+        gkw.setdefault("mutual_knn", True)
+        gkw.setdefault("include_self", True)
+        gkw.setdefault("normalize_rows", False)
+
+        # eps estimation for RBF weights (median kNN distance)
+        gkw.setdefault("k_eps", 256)
+        gkw.setdefault("eps_use_kth", True)
+        gkw.setdefault("eps_mul", 1.0)
+
+        # Solver controls for (K + lambda I) S = Y
+        gkw.setdefault("solve_method", "auto")
+        gkw.setdefault("solve_tol", 1e-6)
+        gkw.setdefault("solve_maxiter", 800)
+        gkw.setdefault("solve_verbose", False)
+
+        # Inference neighbor truncation
+        gkw.setdefault("pred_k", 128)
+
+        # Latent preconditioning (optional)
+        # On ψ it sometimes helps, sometimes hurts. Leave default False unless needed.
+        gkw.setdefault("whiten_latent", False)
+
+        # dtype for ANN storage
+        gkw.setdefault("dtype", np.float32)
 
         self.gplm = GPLM(X_train, Y_train_c, **gkw)
 
@@ -278,7 +249,7 @@ class LISA:
         self.gp_rbf_q = float(gp_rbf_q)
         self.last_gp_rbf_ell_: Optional[float] = None
 
-        # NEW: cap call-time dense GP context size
+        # Dense GP context cap
         self.gp_max_ctx = None if gp_max_ctx is None else int(gp_max_ctx)
         self.gp_ctx_mode = str(gp_ctx_mode).lower().strip()
         if self.gp_ctx_mode not in ("recent", "uniform"):
@@ -315,7 +286,7 @@ class LISA:
         Psi = np.asarray(Psi_Br, dtype=np.float64)
         if Psi.ndim == 1:
             Psi = Psi[None, :]
-        Yc = self.gplm(Psi)  # (B,D)
+        Yc = self.gplm(Psi)  # (B,D) (already centered because center_X=False + mean_X=0)
         return np.asarray(Yc, dtype=np.float64)
 
     def _baseline_centered_one(self, psi_r: np.ndarray) -> np.ndarray:
@@ -335,7 +306,6 @@ class LISA:
         if self.gp_kernel == "linear":
             return Psi_ctx @ Psi_ctx.T, None
 
-        # RBF
         d2 = _pairwise_sq_dists(Psi_ctx)
         if self.gp_rbf_ell is None:
             ell_used = _estimate_rbf_ell_from_d2(d2, q=self.gp_rbf_q)
@@ -382,7 +352,6 @@ class LISA:
 
         if self.gate_mode == "exp":
             return float(np.exp(-v / tau2))
-        # rational
         return float(tau2 / (tau2 + v))
 
     # ============================================================
@@ -417,22 +386,8 @@ class LISA:
         """
         Autoregressive forecast from prefix (ℓ,D), ℓ >= L.
 
-        If ℓ == L (IC nullset): baseline AR only (NLSA+GPLM).
-        If ℓ >  L: uses in-context dense GPR on residuals.
-
-        Parameters
-        ----------
-        prefix : (ℓ,D)
-        steps : horizon H
-        return_var : return scalar GP function variance per step (H,)
-        sample : sample GP residual instead of mean (generative mode)
-        rng : RNG for sampling
-        include_obs_noise : if sample=True, draw using var_y = var_f + gp_noise2
-
-        Returns
-        -------
-        preds : (H,D)  (or (D,) if H==1)
-        vars  : (H,) if return_var=True
+        If ℓ == L: baseline AR only (NLSA + sparse GPLM).
+        If ℓ >  L: uses dense in-context GPR on residuals (optional).
         """
         prefix = _as_2d(prefix)
         ell, D = prefix.shape
@@ -469,7 +424,7 @@ class LISA:
             return preds[0] if H == 1 else preds
 
         # ---------------------------
-        # Dense GP context cap (NEW)
+        # Dense GP context cap (safety)
         # ---------------------------
         if self.gp_max_ctx is None:
             K_ctx = int(K_ctx_full)
@@ -481,17 +436,16 @@ class LISA:
                 start = int(K_ctx_full - K_ctx)
                 ctx_idx = None
             else:
-                # uniform subsample of context indices
                 ctx_idx = np.linspace(0, K_ctx_full - 1, num=K_ctx, dtype=np.int64)
-                start = 0  # unused
+                start = 0
 
         # ---------------------------
         # Build context windows + targets
         # ---------------------------
-        W_all = _sliding_windows(prefix, self.L)    # (K_n,L,D), K_n = ell-L+1
+        W_all = _sliding_windows(prefix, self.L)  # (ell-L+1,L,D)
 
         if ctx_idx is None:
-            W_ctx = np.ascontiguousarray(W_all[start : start + K_ctx, :, :])  # (K_ctx,L,D)
+            W_ctx = np.ascontiguousarray(W_all[start : start + K_ctx, :, :])
             Y_ctx = np.asarray(prefix[self.L + start : self.L + start + K_ctx, :], dtype=np.float64)
         else:
             W_ctx = np.ascontiguousarray(W_all[ctx_idx, :, :])
@@ -501,26 +455,24 @@ class LISA:
         Y_ctx_c = (Y_ctx - self.mu_X[None, :]) if self.center_outputs else Y_ctx
 
         # encode ψ for context
-        Psi_ctx = self._encode_batch(W_ctx)                        # (K_ctx,r)
+        Psi_ctx = self._encode_batch(W_ctx)  # (K_ctx,r)
 
         # baseline predictions on context
-        Y_glob_ctx_c = self._baseline_centered_batch(Psi_ctx)      # (K_ctx,D)
+        Y_glob_ctx_c = self._baseline_centered_batch(Psi_ctx)  # (K_ctx,D)
 
         # residual table
-        E_ctx = Y_ctx_c - Y_glob_ctx_c                              # (K_ctx,D)
+        E_ctx = Y_ctx_c - Y_glob_ctx_c  # (K_ctx,D)
 
         # ---------------------------
         # Fit dense GP on residuals
         # ---------------------------
-        K_mat, ell_used = self._gp_kernel_matrix(Psi_ctx)          # (K_ctx,K_ctx)
+        K_mat, ell_used = self._gp_kernel_matrix(Psi_ctx)  # (K_ctx,K_ctx)
         K_reg = K_mat + self.gp_noise2 * np.eye(K_ctx, dtype=np.float64)
 
-        # Cholesky factorization
         cf = cho_factor(K_reg, lower=True, check_finite=False)
-        alpha = cho_solve(cf, E_ctx, check_finite=False)           # (K_ctx,D)
+        alpha = cho_solve(cf, E_ctx, check_finite=False)  # (K_ctx,D)
         Lfac, lower = cf
 
-        # base context mixing weight
         w_ctx_base = float(K_ctx) / float(K_ctx + self.ctx_k0) if self.ctx_k0 > 0 else 1.0
 
         preds = np.zeros((H, self.D), dtype=np.float64)
@@ -530,14 +482,14 @@ class LISA:
         # AR rollout with GP residual
         # ---------------------------
         for h in range(H):
-            psi_q = self.base.encode_window(cur)                   # (r,)
+            psi_q = self.base.encode_window(cur)
 
             # baseline
-            y_glob_c = self._baseline_centered_one(psi_q)          # (D,)
+            y_glob_c = self._baseline_centered_one(psi_q)
 
             # GP residual mean
-            k_eval = self._gp_kernel_eval(Psi_ctx, psi_q, ell_used)  # (K_ctx,)
-            e_mean = k_eval @ alpha                                 # (D,)
+            k_eval = self._gp_kernel_eval(Psi_ctx, psi_q, ell_used)
+            e_mean = k_eval @ alpha
 
             # GP residual variance (function variance)
             u = solve_triangular(Lfac, k_eval, lower=lower, check_finite=False)
@@ -547,11 +499,9 @@ class LISA:
             if return_var:
                 vars_out[h] = float(var_f)
 
-            # trust gating from variance
             w_gate = self._gate_from_var(var_f)
             w_eff = w_ctx_base * w_gate
 
-            # optionally sample residual
             e_use = e_mean
             if sample:
                 var_y = var_f + (self.gp_noise2 if include_obs_noise else 0.0)
@@ -578,7 +528,7 @@ class LISA:
 
     def _rollout_baseline(self, seed_LD: np.ndarray, H: int) -> np.ndarray:
         """
-        Baseline AR rollout only: NLSA encode + GPLM decode.
+        Baseline AR rollout only: NLSA encode + sparse GPLM decode.
         """
         cur = np.asarray(seed_LD, dtype=np.float64).copy()
         out = np.zeros((H, self.D), dtype=np.float64)