src/dima/dmap.py

# src/dima/dmap.py
from __future__ import annotations

from typing import Any, Dict, Optional, Tuple, Union

import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import eigsh, LinearOperator, lobpcg

from .ann import ANNBackend, make_ann
from .utils import median_eps_from_knn_d2


def k_ideal(d: int, N: int) -> int:
    """
    Heuristic for kNN graph size in diffusion maps.
    Stable default: grows slowly with N and linearly with d.
    """
    d = int(max(1, d))
    N = int(max(2, N))
    k = int(np.ceil(2.0 * d * np.log2(N)))
    return int(min(max(8, k), N - 1))


def _sqdist_ab(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """
    Squared Euclidean distances between rows:
      A: (a,d), B: (b,d) -> D2: (a,b)
    """
    A = np.asarray(A, dtype=np.float64)
    B = np.asarray(B, dtype=np.float64)
    A2 = np.sum(A * A, axis=1, keepdims=True)
    B2 = np.sum(B * B, axis=1, keepdims=True).T
    G = A @ B.T
    return np.maximum(A2 + B2 - 2.0 * G, 0.0)


class DMAP:
    """
    Diffusion Maps encoder with Nyström out-of-sample extension.

    Notation (arrays named by indices):

      R_iX: reference ambient data
      K_ij: kernel on graph edges (sparse CSR)
      q_i = Σ_j K_ij
      qα_i = (q_i)^α
      Kα_ij = K_ij / (qα_i qα_j)
      d_i = Σ_j Kα_ij
      A_ij = Kα_ij / sqrt(d_i d_j)  (symmetric)

      eigsh(A) -> λ_x, u_ix
      ψ_ix = u_ix / sqrt(d_i)
      R_ix = (λ_x)^t ψ_ix

    Nyström OOS for novel ambient R_aX:
      K_ai = exp(-β * D2_ai / ε)
      q_a = Σ_i K_ai,  qα_a = (q_a)^α
      Kα_ai = K_ai / (qα_a qα_i)
      d_a = Σ_i Kα_ai
      P_ai = Kα_ai / d_a
      R_ax = Σ_i P_ai * (R_ix / λ_x)

    Extras (from your latest version):
      - refine_dense: warm-start with sparse kNN eigsh, then refine with streaming-matvec LOBPCG
      - stream_block: tile size for dense kernel streaming
      - lobpcg_maxiter / lobpcg_tol: refinement controls
      - use_symmetry: exploit K symmetry in dense matvec (roughly ~2x fewer tiles)
    """

    def __init__(
        self,
        R_iX: np.ndarray,
        *,
        # preferred ASCII names
        d: int = 6,
        k: Optional[int] = None,
        beta: float = 1.0,
        alpha: float = 0.0,
        t: float = 1.0,
        drop_trivial: bool = True,
        seed: int = 0,
        ann_backend: ANNBackend = "auto",
        ann_params: Optional[Dict[str, Any]] = None,
        n_jobs: int = -1,
        eps: Optional[float] = None,
        eps_use_kth: bool = True,
        eps_mul: float = 1.0,
        sym: str = "max",
        nL: Optional[int] = None,
        L_iX: Optional[np.ndarray] = None,
        dtype: Any = np.float32,
        # dense refinement (O(N^2) compute, streaming memory)
        refine_dense: bool = False,
        stream_block: int = 4096,
        lobpcg_maxiter: int = 3,
        lobpcg_tol: float = 1e-6,
        use_symmetry: bool = True,
        # allow unicode kwargs (β, α, ε, ε_mul, ε_use_kth, ...)
        **kwargs: Any,
    ):
        # ---- map unicode kwargs -> ascii ----
        if "β" in kwargs:
            beta = kwargs.pop("β")
        if "α" in kwargs:
            alpha = kwargs.pop("α")
        if "ε" in kwargs:
            eps = kwargs.pop("ε")
        if "ε_use_kth" in kwargs:
            eps_use_kth = kwargs.pop("ε_use_kth")
        if "ε_mul" in kwargs:
            eps_mul = kwargs.pop("ε_mul")
        if "sym" in kwargs:
            sym = kwargs.pop("sym")
        if kwargs:
            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")

        self.d = int(d)
        self.k = int(k_ideal(self.d, int(np.asarray(R_iX).shape[0])) if k is None else int(k))

        self.beta = float(beta)
        self.alpha = float(alpha)
        self.t = float(t)
        self.drop_trivial = bool(drop_trivial)
        self.seed = int(seed)
        self.sym = str(sym)
        self.dtype = dtype

        # unicode aliases (so older code + your packer can find them)
        self.β = self.beta
        self.α = self.alpha

        self.refine_dense = bool(refine_dense)
        self.stream_block = int(stream_block)
        self.lobpcg_maxiter = int(lobpcg_maxiter)
        self.lobpcg_tol = float(lobpcg_tol)
        self.use_symmetry = bool(use_symmetry)

        rng = np.random.default_rng(self.seed)

        # ---- choose reference set (landmarks optional) ----
        R_iX = np.asarray(R_iX)
        if L_iX is not None:
            R_iXref = np.asarray(L_iX)
        elif nL is not None:
            nL = int(nL)
            if nL <= 0 or nL > R_iX.shape[0]:
                raise ValueError("nL must be in [1, N].")
            sel = rng.choice(R_iX.shape[0], size=nL, replace=False)
            R_iXref = R_iX[sel]
        else:
            R_iXref = R_iX

        R_iXref = np.ascontiguousarray(R_iXref.astype(self.dtype, copy=False))
        Nref = int(R_iXref.shape[0])
        if self.k >= Nref:
            raise ValueError(f"k={self.k} must be < Nref={Nref}.")

        # ---- ANN on reference ----
        self.ann, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
        self.ann.build(R_iXref)
        self.R_iX = R_iXref  # store reference set

        # ---- kNN on reference (ask k+1 to try to include self) ----
        j_iK1, D2_iK1 = self.ann.search(R_iXref, self.k + 1)  # (Nref,k+1)

        # drop self neighbor if present
        i = np.arange(Nref)[:, None]
        is_self = (j_iK1 == i)
        if np.any(is_self):
            j_iK = np.empty((Nref, self.k), dtype=np.int64)
            D2_iK = np.empty((Nref, self.k), dtype=np.float64)
            for ii in range(Nref):
                keep = (j_iK1[ii] != ii)
                jj = j_iK1[ii][keep][: self.k]
                dd = D2_iK1[ii][keep][: self.k]
                if jj.shape[0] < self.k:
                    pad = self.k - jj.shape[0]
                    jj = np.pad(jj, (0, pad), mode="edge")
                    dd = np.pad(dd, (0, pad), mode="edge")
                j_iK[ii] = jj
                D2_iK[ii] = dd
        else:
            j_iK = j_iK1[:, : self.k].astype(np.int64, copy=False)
            D2_iK = D2_iK1[:, : self.k].astype(np.float64, copy=False)

        # ---- eps from kNN statistics ----
        if eps is None:
            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(f"eps must be > 0, got {eps_hat}")
        self.eps = float(eps_hat)
        self.ε = self.eps  # unicode alias

        # -------------------------
        # 1) Warm start: sparse kNN DMAP operator + eigsh
        # -------------------------
        K_iK = np.exp(-self.beta * (D2_iK / self.eps)).astype(np.float64, copy=False)

        indptr = (np.arange(Nref + 1, dtype=np.int64) * self.k)
        indices = j_iK.reshape(-1).astype(np.int64, copy=False)
        data = K_iK.reshape(-1)

        K_ij = sp.csr_matrix((data, indices, indptr), shape=(Nref, Nref), dtype=np.float64)

        # symmetrize
        if self.sym == "max":
            K_ij = K_ij.maximum(K_ij.T)
        elif self.sym == "mean":
            K_ij = (K_ij + K_ij.T) * 0.5
        else:
            raise ValueError(f"Unknown sym={self.sym!r}")

        # degrees q_i and qalpha_i (warm)
        q_i_warm = np.asarray(K_ij.sum(axis=1)).ravel()
        q_i_warm = np.maximum(q_i_warm, 1e-30)
        qalpha_i_warm = np.maximum(np.power(q_i_warm, self.alpha), 1e-30)

        Qinv = sp.diags(1.0 / qalpha_i_warm, format="csr")
        Kalpha_ij = Qinv @ K_ij @ Qinv

        d_i_warm = np.asarray(Kalpha_ij.sum(axis=1)).ravel()
        d_i_warm = np.maximum(d_i_warm, 1e-30)

        Dinv_sqrt = sp.diags(1.0 / np.sqrt(d_i_warm), format="csr")
        A_ij = Dinv_sqrt @ Kalpha_ij @ Dinv_sqrt

        nev = self.d + (1 if self.drop_trivial else 0)
        v0 = rng.normal(size=Nref).astype(np.float64)
        lam0, u0 = eigsh(A_ij, k=nev, which="LA", v0=v0)

        ord0 = np.argsort(lam0)[::-1]
        lam0 = lam0[ord0]
        u0 = u0[:, ord0]

        # LOBPCG warm-start block (orthonormalize)
        X0, _ = np.linalg.qr(u0.astype(np.float64, copy=False))

        # -------------------------
        # 2) Optional refinement: streaming dense LOBPCG on dense PSD operator
        # -------------------------
        if self.refine_dense:
            # full dense K uses all pairs; streaming avoids materializing K
            self._R2_i = np.sum(self.R_iX.astype(np.float64) ** 2, axis=1)  # (Nref,)

            ones = np.ones((Nref, 1), dtype=np.float64)
            q_i = self._K_matmat_dense(ones).ravel()
            q_i = np.maximum(q_i, 1e-30)
            qalpha_i = np.maximum(np.power(q_i, self.alpha), 1e-30)
            u_i = 1.0 / qalpha_i  # q^{-alpha}

            Ku = self._K_matmat_dense(u_i[:, None]).ravel()
            d_i = np.maximum(u_i * Ku, 1e-30)
            s_i = 1.0 / np.sqrt(d_i)

            def A_matmat(V: np.ndarray) -> np.ndarray:
                V = V.astype(np.float64, copy=False)
                V1 = s_i[:, None] * V
                V2 = u_i[:, None] * V1
                V3 = self._K_matmat_dense(V2)
                V4 = u_i[:, None] * V3
                V5 = s_i[:, None] * V4
                return V5

            Aop = LinearOperator(
                (Nref, Nref),
                matvec=lambda v: A_matmat(v[:, None])[:, 0],
                matmat=A_matmat,
                dtype=np.float64,
            )

            try:
                lam, u = lobpcg(
                    Aop,
                    X0,
                    largest=True,
                    maxiter=self.lobpcg_maxiter,
                    tol=self.lobpcg_tol,
                )
                ord1 = np.argsort(lam)[::-1]
                lam = lam[ord1]
                u = u[:, ord1]

                self.q_i = q_i.astype(np.float64, copy=False)
                self.qalpha_i = qalpha_i.astype(np.float64, copy=False)
                self.d_i = d_i.astype(np.float64, copy=False)
            except Exception:
                # fallback to warm start if refinement fails
                lam, u = lam0, u0
                self.q_i = q_i_warm.astype(np.float64, copy=False)
                self.qalpha_i = qalpha_i_warm.astype(np.float64, copy=False)
                self.d_i = d_i_warm.astype(np.float64, copy=False)
        else:
            lam, u = lam0, u0
            self.q_i = q_i_warm.astype(np.float64, copy=False)
            self.qalpha_i = qalpha_i_warm.astype(np.float64, copy=False)
            self.d_i = d_i_warm.astype(np.float64, copy=False)

        # provide unicode aliases for packers / older code
        self.qα_i = self.qalpha_i
        self.λ_x = lam.astype(np.float64, copy=False)

        # psi and drop trivial
        psi = u / np.sqrt(self.d_i)[:, None]
        if self.drop_trivial:
            lam = lam[1:]
            psi = psi[:, 1:]
            u = u[:, 1:]

        # diffusion coords
        R_ix = psi * (lam ** self.t)[None, :]

        # store
        self.λ_x = lam.astype(np.float64, copy=False)   # (d,)
        self.u_ix = u.astype(np.float64, copy=False)    # (Nref,d)
        self.ψ_ix = psi.astype(np.float64, copy=False)  # (Nref,d)
        self.R_ix = R_ix.astype(np.float64, copy=False) # (Nref,d)
        self.π_i = (self.d_i / self.d_i.sum()).astype(np.float64, copy=False)

        # for Nyström: R_ix / λ_x
        self.R_over_λ_ix = (self.R_ix / self.λ_x[None, :]).astype(np.float64, copy=False)

    # --------- streaming dense kernel primitives (only used if refine_dense=True) ----------

    def _rbf_block(self, Xb: np.ndarray, Xc: np.ndarray, X2b: np.ndarray, X2c: np.ndarray) -> np.ndarray:
        # squared distances: ||b||^2 + ||c||^2 - 2 b c^T
        G = Xb @ Xc.T
        D2 = np.maximum(X2b[:, None] + X2c[None, :] - 2.0 * G, 0.0)
        return np.exp(-self.beta * (D2 / self.eps))

    def _K_matmat_dense(self, V: np.ndarray) -> np.ndarray:
        """
        Streaming matmat for dense K:
          out = K @ V
        Does NOT materialize K. Optionally exploits symmetry by only computing upper-tri tiles.
        """
        X = self.R_iX.astype(np.float64, copy=False)
        X2 = self._R2_i
        N = X.shape[0]
        bs = self.stream_block
        V = V.astype(np.float64, copy=False)
        out = np.zeros((N, V.shape[1]), dtype=np.float64)

        if not self.use_symmetry:
            for i0 in range(0, N, bs):
                i1 = min(N, i0 + bs)
                Xi = X[i0:i1]
                X2i = X2[i0:i1]
                acc = np.zeros((i1 - i0, V.shape[1]), dtype=np.float64)
                for j0 in range(0, N, bs):
                    j1 = min(N, j0 + bs)
                    Xj = X[j0:j1]
                    X2j = X2[j0:j1]
                    Kij = self._rbf_block(Xi, Xj, X2i, X2j)
                    acc += Kij @ V[j0:j1]
                out[i0:i1] = acc
            return out

        # symmetric tiling
        for i0 in range(0, N, bs):
            i1 = min(N, i0 + bs)
            Xi = X[i0:i1]
            X2i = X2[i0:i1]
            Vi = V[i0:i1]

            # diagonal tile
            Kii = self._rbf_block(Xi, Xi, X2i, X2i)
            out[i0:i1] += Kii @ Vi

            for j0 in range(i1, N, bs):
                j1 = min(N, j0 + bs)
                Xj = X[j0:j1]
                X2j = X2[j0:j1]
                Vj = V[j0:j1]

                Kij = self._rbf_block(Xi, Xj, X2i, X2j)
                out[i0:i1] += Kij @ Vj
                out[j0:j1] += Kij.T @ Vi

        return out

    # --------- Nyström embedding ----------

    def __call__(self, R_aX: Union[np.ndarray, list], *, batch_size: Optional[int] = None) -> np.ndarray:
        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))

        if batch_size is None:
            R_ax = self._embed(R_aX)
        else:
            bs = int(batch_size)
            out = []
            for s in range(0, R_aX.shape[0], bs):
                out.append(self._embed(R_aX[s:s + bs]))
            R_ax = np.vstack(out)

        return R_ax[0] if single else R_ax

    def _embed(self, R_aX: np.ndarray) -> np.ndarray:
        # kNN for novel points
        j_aK, D2_aK = self.ann.search(R_aX, self.k)  # (a,k)

        K_ai = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)

        q_a = np.maximum(K_ai.sum(axis=1), 1e-30)
        qalpha_a = np.maximum(np.power(q_a, self.alpha), 1e-30)

        qalpha_i = np.maximum(self.qalpha_i[j_aK], 1e-30)
        Kalpha_ai = K_ai / (qalpha_a[:, None] * qalpha_i)

        d_a = np.maximum(Kalpha_ai.sum(axis=1), 1e-30)
        P_ai = Kalpha_ai / d_a[:, None]

        R_over = self.R_over_λ_ix[j_aK, :]  # (a,k,d)
        R_ax = (P_ai[:, :, None] * R_over).sum(axis=1)
        return R_ax


__all__ = ["DMAP", "k_ideal"]