Create GPLMx.py

GPLMx.py

@@ -0,0 +1,580 @@
+# GPLM.py
+# ============================================================
+# GPLM (drop-in replacement):
+#   Sparse-kernel KRR / mean-GP decoder on latent coordinates.
+#
+# This version:
+#   - Builds a sparse kernel matrix K using kNN (ANN) on training latents.
+#   - Solves (K + sigma2 * I) S = Y  for S (N,D) via iterative sparse solvers.
+#   - Predicts for new points using only k_pred nearest training points:
+#       y(x) ≈ sum_{j in kNN(x)} k(x, x_j) * S_j
+#
+# Key properties:
+#   - No Nyström anchors / no inducing points.
+#   - "Nonlinearity" corresponds to strong locality (small eps, small k_graph):
+#       sparse + high-rank-ish operator, but scalable because it's sparse.
+#
+# API compatibility:
+#   - class GPLM
+#   - __init__(...), fit(...)
+#   - __call__(R_ax), predict(..., return_var=True)
+#   - kernel_mass(...)
+#   - flow(...) present but NotImplemented (optional advanced geometry)
+#
+# Dependencies:
+#   numpy, scipy
+#   ann.py + utils.py (same repo assumptions as your previous GPLM)
+# ============================================================
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+import scipy.sparse as sp
+import scipy.sparse.linalg as spla
+
+# --- same repo dependencies as before ---
+from ann import ANNBackend, make_ann
+from utils import median_eps_from_knn_d2
+
+InducingMode = Literal["random_subset", "fps", "kmeans_medoids", "given"]  # kept for API compat
+SolveMethod = Literal["cg", "minres", "auto"]
+
+
+def _as_2d(x: np.ndarray) -> np.ndarray:
+    x = np.asarray(x)
+    if x.ndim == 1:
+        return x[None, :]
+    return x
+
+
+def _row_norm2(X: np.ndarray) -> np.ndarray:
+    return np.sum(X * X, axis=1)
+
+
+def _rbf_weights_from_d2(D2: np.ndarray, beta: float, eps: float) -> np.ndarray:
+    # weight = exp(-beta * d^2 / eps)
+    return np.exp(-float(beta) * (D2.astype(np.float64) / float(eps)))
+
+
+def _symmetrize_coo(i: np.ndarray, j: np.ndarray, v: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Add transpose entries and return concatenated arrays."""
+    i2 = np.concatenate([i, j], axis=0)
+    j2 = np.concatenate([j, i], axis=0)
+    v2 = np.concatenate([v, v], axis=0)
+    return i2, j2, v2
+
+
+def _unique_coo_sum(N: int, i: np.ndarray, j: np.ndarray, v: np.ndarray) -> sp.csr_matrix:
+    """Build CSR matrix with duplicates summed."""
+    K = sp.coo_matrix((v, (i, j)), shape=(N, N), dtype=np.float64).tocsr()
+    K.sum_duplicates()
+    return K
+
+
+def _solve_multi_rhs(
+    A: sp.csr_matrix,
+    Y: np.ndarray,
+    *,
+    method: SolveMethod = "auto",
+    tol: float = 1e-6,
+    maxiter: int = 500,
+    verbose: bool = False,
+) -> np.ndarray:
+    """
+    Solve A X = Y for X with multiple RHS columns using iterative solvers.
+    A is expected sparse and (typically) symmetric.
+    """
+    Y = np.asarray(Y, dtype=np.float64, order="C")
+    N, D = Y.shape
+    X = np.zeros((N, D), dtype=np.float64)
+
+    # Choose solver
+    if method == "auto":
+        # CG is fastest if SPD; minres is safer if indefinite
+        method_use: SolveMethod = "cg"
+    else:
+        method_use = method
+
+    # Wrapper per RHS
+    for d in range(D):
+        b = Y[:, d]
+        x0 = None  # could warm-start if you want
+
+        if method_use == "cg":
+            x, info = spla.cg(A, b, x0=x0, tol=tol, maxiter=maxiter)
+            if info != 0:
+                # fallback to MINRES
+                x, info2 = spla.minres(A, b, x0=x0, tol=tol, maxiter=maxiter)
+                if verbose:
+                    print(f"[GPLM sparse] cg failed info={info}, minres info={info2} on dim {d}")
+            X[:, d] = x
+        elif method_use == "minres":
+            x, info = spla.minres(A, b, x0=x0, tol=tol, maxiter=maxiter)
+            if verbose and info != 0:
+                print(f"[GPLM sparse] minres info={info} on dim {d}")
+            X[:, d] = x
+        else:
+            raise ValueError(f"Unknown solve method: {method_use}")
+
+    return X
+
+
+@dataclass
+class _SparseKernelConfig:
+    k_graph: int = 64          # neighbors per training point for building sparse K
+    mutual: bool = True        # mutual-kNN symmetrization (recommended)
+    include_self: bool = True  # ensure diagonal has 1.0
+    normalize_rows: bool = False  # optional row-normalization for stability
+
+
+class GPLM:
+    """
+    GPLM (Sparse-kernel KRR decoder)
+
+    Training:
+      Inputs:  R_ix (N,d) latents, R_iX (N,D) outputs
+      Build sparse kernel K via kNN on R_ix:
+        K_ij = exp(-beta ||R_i-R_j||^2 / eps)  for neighbors only
+
+      Solve for weights S (N,D):
+        (K + sigma2 * I + jitter*I) S = Y_centered
+
+    Inference:
+      For query R_ax:
+        Find k_pred nearest training points j_aK
+        Compute weights w_aK = exp(-beta d2/eps)
+        Predict:
+          Yc = sum_k w[a,k] * S[j_aK[a,k], :]
+        Return Y = Yc + mean_X
+
+    Variance proxy:
+      Not full GP variance; return a support-based scalar:
+        mass = sum_k w[a,k]
+        var ≈ sigma2 / (mass + 1e-12)
+    """
+
+    def __init__(
+        self,
+        R_ix: np.ndarray,
+        R_iX: np.ndarray,
+        *,
+        # Kernel params
+        beta: float = 1.0,
+        eps: Optional[float] = None,
+        k_eps: int = 256,
+        eps_use_kth: bool = True,
+        eps_mul: float = 1.0,
+        # Regularization (acts like ridge lambda)
+        sigma2: float = 1e-5,
+        jitter: float = 1e-8,
+        # "Inducing" params kept for API compat (ignored)
+        m: int = 1024,
+        inducing: InducingMode = "kmeans_medoids",
+        Z_mx: Optional[np.ndarray] = None,
+        seed: int = 0,
+        # Preprocess
+        center_X: bool = True,
+        whiten_latent: bool = False,
+        dtype: Any = np.float32,
+        # Sparse kernel build
+        k_graph: int = 64,
+        mutual_knn: bool = True,
+        include_self: bool = True,
+        normalize_rows: bool = False,
+        # Solve
+        solve_method: SolveMethod = "auto",
+        solve_tol: float = 1e-6,
+        solve_maxiter: int = 800,
+        solve_verbose: bool = False,
+        # Inference neighbor truncation
+        pred_k: Optional[int] = 128,
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+        # accept unicode kwargs (β, ε, κ_eps, σ2, pred_κ, ...)
+        **kwargs: Any,
+    ):
+        # ---- map unicode kwargs -> ascii ----
+        if "β" in kwargs:
+            beta = kwargs.pop("β")
+        if "ε" in kwargs:
+            eps = kwargs.pop("ε")
+        if "κ_eps" in kwargs:
+            k_eps = kwargs.pop("κ_eps")
+        if "ε_use_kth" in kwargs:
+            eps_use_kth = kwargs.pop("ε_use_kth")
+        if "ε_mul" in kwargs:
+            eps_mul = kwargs.pop("ε_mul")
+        if "σ2" in kwargs:
+            sigma2 = kwargs.pop("σ2")
+        if "pred_κ" in kwargs:
+            pred_k = kwargs.pop("pred_κ")
+
+        # ignore anchor arguments quietly (compat)
+        _ = (m, inducing, Z_mx)
+
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        self.beta = float(beta)
+        self.β = self.beta
+
+        self.sigma2 = float(sigma2)   # ridge lambda in (K + sigma2 I)
+        self.σ2 = self.sigma2
+
+        self.jitter = float(jitter)
+        self.seed = int(seed)
+        self.dtype = dtype
+
+        # ---- validate / cast ----
+        R_ix = np.ascontiguousarray(np.asarray(R_ix).astype(self.dtype, copy=False))
+        R_iX = np.ascontiguousarray(np.asarray(R_iX).astype(self.dtype, copy=False))
+        if R_ix.ndim != 2 or R_iX.ndim != 2 or R_ix.shape[0] != R_iX.shape[0]:
+            raise ValueError("R_ix must be (N,d) and R_iX must be (N,D) with same N.")
+
+        self.R_ix = R_ix
+        self.R_iX = R_iX
+        self.N, self.d_lat = R_ix.shape
+        _, self.D = R_iX.shape
+
+        # ---- center output ----
+        self.center_X = bool(center_X)
+        if self.center_X:
+            self.mean_X = R_iX.mean(axis=0).astype(np.float64)
+            Y = (R_iX.astype(np.float64) - self.mean_X[None, :])
+        else:
+            self.mean_X = np.zeros((self.D,), dtype=np.float64)
+            Y = R_iX.astype(np.float64)
+
+        # ---- latent whitening (optional) ----
+        self.whiten_latent = bool(whiten_latent)
+        Ztrain = R_ix.astype(np.float64)
+        if self.whiten_latent:
+            self.lat_mean_x = Ztrain.mean(axis=0)
+            self.lat_std_x = np.maximum(Ztrain.std(axis=0), 1e-12)
+            Ztrain_w = (Ztrain - self.lat_mean_x) / self.lat_std_x
+        else:
+            self.lat_mean_x = np.zeros((self.d_lat,), dtype=np.float64)
+            self.lat_std_x = np.ones((self.d_lat,), dtype=np.float64)
+            Ztrain_w = Ztrain
+
+        self.R_ix_w = np.ascontiguousarray(Ztrain_w.astype(np.float64, copy=False))  # (N,d) float64
+
+        # ---- ANN on training latents ----
+        self.ann_train, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        self.ann_train.build(self.R_ix_w.astype(self.dtype, copy=False))
+
+        # ---- eps via kNN distances ----
+        if eps is None:
+            k_eps_eff = int(min(max(8, int(k_eps)), self.N - 1))
+            # ask for k_eps+1 to try include self
+            j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), k_eps_eff + 1)
+
+            i = np.arange(self.N)[:, None]
+            is_self = (j_iK1 == i)
+
+            if np.any(is_self):
+                D2_iK = np.empty((self.N, k_eps_eff), dtype=np.float64)
+                for ii in range(self.N):
+                    keep = (j_iK1[ii] != ii)
+                    D2_iK[ii] = D2_iK1[ii][keep][:k_eps_eff]
+            else:
+                D2_iK = D2_iK1[:, :k_eps_eff].astype(np.float64, copy=False)
+
+            eps_hat = median_eps_from_knn_d2(D2_iK, use_kth=bool(eps_use_kth))
+        else:
+            eps_hat = float(eps)
+
+        eps_hat *= float(eps_mul)
+        if eps_hat <= 0:
+            raise ValueError("eps must be > 0.")
+        self.eps = float(eps_hat)
+        self.ε = self.eps
+
+        # ---- build sparse kernel K ----
+        cfg = _SparseKernelConfig(
+            k_graph=int(min(max(4, int(k_graph)), self.N - 1)),
+            mutual=bool(mutual_knn),
+            include_self=bool(include_self),
+            normalize_rows=bool(normalize_rows),
+        )
+        self._cfg = cfg
+
+        # query kNN on training set for graph edges
+        j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), cfg.k_graph + 1)
+
+        # drop self if present
+        rows = []
+        cols = []
+        vals = []
+
+        for i in range(self.N):
+            nbrs = j_iK1[i]
+            d2 = D2_iK1[i].astype(np.float64, copy=False)
+            # remove self
+            mask = (nbrs != i)
+            nbrs = nbrs[mask][: cfg.k_graph]
+            d2 = d2[mask][: cfg.k_graph]
+
+            w = _rbf_weights_from_d2(d2, beta=self.beta, eps=self.eps)
+
+            rows.append(np.full(nbrs.shape[0], i, dtype=np.int64))
+            cols.append(nbrs.astype(np.int64, copy=False))
+            vals.append(w.astype(np.float64, copy=False))
+
+        i_idx = np.concatenate(rows, axis=0)
+        j_idx = np.concatenate(cols, axis=0)
+        v_idx = np.concatenate(vals, axis=0)
+
+        # symmetric adjacency
+        if cfg.mutual:
+            i_idx, j_idx, v_idx = _symmetrize_coo(i_idx, j_idx, v_idx)
+
+        # build sparse K (CSR)
+        K = _unique_coo_sum(self.N, i_idx, j_idx, v_idx)
+
+        # set diagonal to 1 (kernel self-sim), improves conditioning
+        if cfg.include_self:
+            K = K.tolil(copy=False)
+            diag = K.diagonal()
+            # if diagonal already has values from sym edges, top it up to 1
+            diag_new = np.maximum(np.asarray(diag).reshape(-1), 1.0)
+            for ii in range(self.N):
+                K[ii, ii] = float(diag_new[ii])
+            K = K.tocsr(copy=False)
+
+        # optional row normalization (turns kernel into a diffusion-like operator)
+        if cfg.normalize_rows:
+            rs = np.asarray(K.sum(axis=1)).reshape(-1)
+            rs = np.maximum(rs, 1e-12)
+            inv = 1.0 / rs
+            K = sp.diags(inv, format="csr") @ K
+
+        K.sum_duplicates()
+        self.K = K  # (N,N) sparse
+
+        # ---- form A = K + (sigma2 + jitter) I ----
+        lam = float(self.sigma2)
+        jit = float(self.jitter)
+        A = self.K.tocsr(copy=True)
+        A = A + sp.diags((lam + jit) * np.ones(self.N), format="csr")
+        self._A = A.tocsr(copy=False)
+
+        # ---- solve for S: A S = Y ----
+        self.solve_method = str(solve_method)
+        self.solve_tol = float(solve_tol)
+        self.solve_maxiter = int(solve_maxiter)
+        self.solve_verbose = bool(solve_verbose)
+
+        self.S_iX = _solve_multi_rhs(
+            self._A,
+            Y,
+            method=self.solve_method,  # type: ignore
+            tol=self.solve_tol,
+            maxiter=self.solve_maxiter,
+            verbose=self.solve_verbose,
+        ).astype(np.float64)
+
+        # store float32 copy for fast inference
+        self.S_iX_f32 = self.S_iX.astype(np.float32, copy=False)
+
+        # ---- inference neighbor count ----
+        if pred_k is None:
+            self.pred_k = int(min(128, self.N - 1))
+        else:
+            self.pred_k = int(min(max(1, int(pred_k)), self.N - 1))
+        self.pred_κ = self.pred_k  # unicode alias
+
+    # ------------------------------------------------------------
+    # Convenience alternate constructor
+    # ------------------------------------------------------------
+    @classmethod
+    def fit(cls, R_ix: np.ndarray, R_iX: np.ndarray, **kwargs: Any) -> "GPLM":
+        return cls(R_ix, R_iX, **kwargs)
+
+    # ------------------------------------------------------------
+    # Internal prediction helpers
+    # ------------------------------------------------------------
+    def _whiten_query(self, R_ax: np.ndarray) -> np.ndarray:
+        R_ax = np.asarray(R_ax, dtype=np.float64)
+        if self.whiten_latent:
+            return (R_ax - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
+        return R_ax
+
+    def _predict_mean(self, R_ax: np.ndarray) -> np.ndarray:
+        """
+        Mean prediction using only pred_k nearest training points.
+        """
+        R_ax = np.asarray(R_ax)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+        R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))
+
+        # whiten query for ANN
+        Rw = self._whiten_query(R_ax.astype(np.float64, copy=False)).astype(self.dtype, copy=False)
+
+        # neighbors in training set
+        j_aK, D2_aK = self.ann_train.search(Rw, self.pred_k)
+
+        # weights (A,k)
+        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)
+
+        # gather S for neighbors -> (A,k,D)
+        S = self.S_iX_f32  # (N,D)
+        Sj = S[j_aK]       # (A,k,D)
+
+        # weighted sum -> (A,D)
+        Yc = np.einsum("ak,akd->ad", W.astype(np.float32, copy=False), Sj, optimize=True).astype(np.float64)
+
+        # add mean
+        Y = Yc + self.mean_X[None, :]
+        return Y[0] if single else Y
+
+    def _predict_mean_var(self, R_ax: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Mean + cheap scalar uncertainty proxy based on kernel mass on neighbors.
+        This is *not* exact GP posterior variance.
+        """
+        R_ax = np.asarray(R_ax)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+        R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))
+
+        Rw = self._whiten_query(R_ax.astype(np.float64, copy=False)).astype(self.dtype, copy=False)
+        j_aK, D2_aK = self.ann_train.search(Rw, self.pred_k)
+
+        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)  # (A,k)
+        mass = np.sum(W, axis=1)  # (A,)
+
+        S = self.S_iX_f32
+        Sj = S[j_aK]  # (A,k,D)
+        Yc = np.einsum("ak,akd->ad", W.astype(np.float32, copy=False), Sj, optimize=True).astype(np.float64)
+        mean = Yc + self.mean_X[None, :]
+
+        # heuristic: low mass => off-support => higher uncertainty
+        var = (self.sigma2 / (mass + 1e-12)).astype(np.float64)
+
+        if single:
+            return mean[0], var[0]
+        return mean, var
+
+    # ------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------
+    def __call__(self, R_ax: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
+        """
+        Mean prediction only. For uncertainty, use predict(..., return_var=True).
+        """
+        R_ax = np.asarray(R_ax)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+
+        if batch_size is None:
+            Y = self._predict_mean(R_ax)
+        else:
+            bs = int(batch_size)
+            out = []
+            for s in range(0, R_ax.shape[0], bs):
+                out.append(self._predict_mean(R_ax[s : s + bs]))
+            Y = np.vstack(out)
+
+        return Y[0] if single else Y
+
+    def predict(
+        self,
+        R_ax: Union[np.ndarray, list],
+        *,
+        return_var: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+        """
+        Predict mean and (optional) scalar variance proxy per query.
+        """
+        R_ax = np.asarray(R_ax)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+
+        if not return_var:
+            mean = self.__call__(R_ax, batch_size=batch_size)
+            if single and mean.ndim == 1:
+                return mean[None, :]
+            return mean
+
+        if batch_size is None:
+            mean, var = self._predict_mean_var(R_ax)
+        else:
+            bs = int(batch_size)
+            ms = []
+            vs = []
+            for s in range(0, R_ax.shape[0], bs):
+                m, v = self._predict_mean_var(R_ax[s : s + bs])
+                ms.append(m)
+                vs.append(np.atleast_1d(v))
+            mean = np.vstack(ms)
+            var = np.concatenate(vs, axis=0)
+
+        if single:
+            return mean, float(var[0]) if np.ndim(var) > 0 else float(var)
+        return mean, var
+
+    def kernel_mass(self, R_ax: Union[np.ndarray, list]) -> np.ndarray:
+        """
+        Support diagnostic:
+          mass(x) = sum_{j in kNN(x)} exp(-beta ||x-x_j||^2 / eps)
+        """
+        R_ax = np.asarray(R_ax, dtype=np.float64)
+        single = (R_ax.ndim == 1)
+        if single:
+            R_ax = R_ax[None, :]
+
+        Rw = self._whiten_query(R_ax).astype(self.dtype, copy=False)
+        _, D2_aK = self.ann_train.search(Rw, self.pred_k)
+        W = _rbf_weights_from_d2(D2_aK.astype(np.float64, copy=False), beta=self.beta, eps=self.eps)
+        mass = np.sum(W, axis=1)
+        return float(mass[0]) if single else mass
+
+    # ------------------------------------------------------------
+    # Geometry / flow (kept for API compatibility)
+    # ------------------------------------------------------------
+    def flow(
+        self,
+        R_ax: Union[np.ndarray, list],
+        v_ax: Union[np.ndarray, list],
+        *,
+        dt: float = 1e-2,
+        K_p: int = 5,
+        K_q: int = 5,
+        D_block: int = 8192,
+        lam: float = 1e-8,
+        metric_solver: str = "auto",
+        eig_clip: float = 1e-12,
+        k_tangent: Optional[int] = None,
+        force_fn: Optional[Any] = None,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Placeholder for compatibility with the original GPLM API.
+
+        Sparse-kernel KRR decoder has a well-defined Jacobian via kernel gradients,
+        but a robust geodesic integrator is non-trivial and model-specific.
+
+        If you truly need flow() (geodesic-ish latent motion), you can:
+          - implement it using local neighbor gradients, OR
+          - keep the original GPLM for geometry tasks.
+
+        For now: not implemented.
+        """
+        raise NotImplementedError(
+            "flow() is not implemented in sparse-kernel GPLM. "
+            "Use the original Nyström GPLM for geodesic flow, or implement "
+            "a local-neighborhood gradient-based flow."
+        )
+
+
+__all__ = ["GPLM"]