Update src/dima/gplm.py

src/dima/gplm.py

@@ -1,15 +1,15 @@
 # src/dima/gplm.py
 from __future__ import annotations
 
-from typing import Any, Dict, Literal, Optional, Tuple, Union, overload
+from typing import Any, Dict, Literal, Optional, Tuple, Union
 
 import numpy as np
 import scipy.linalg as la
 
-
 from .ann import ANNBackend, make_ann
 from .utils import fps_indices, median_eps_from_knn_d2, sqdist_ab
 
+
 InducingMode = Literal["random_subset", "fps", "kmeans_medoids", "given"]
 
 
@@ -22,6 +22,7 @@ def _kmeans2_safe(Z: np.ndarray, m: int, seed: int = 0) -> np.ndarray:
     m = int(min(max(1, m), Z.shape[0]))
     try:
         from scipy.cluster.vq import kmeans2  # type: ignore
+        # minit="points" picks initial centers from data -> stable for medoids snapping
         C, _ = kmeans2(Z.astype(np.float64, copy=False), m, minit="points", seed=seed)
         return C.astype(Z.dtype, copy=False)
     except Exception:
@@ -29,105 +30,254 @@ def _kmeans2_safe(Z: np.ndarray, m: int, seed: int = 0) -> np.ndarray:
         idx = rng.choice(Z.shape[0], size=m, replace=False)
         return Z[idx]
 
-from dataclasses import dataclass
 
-def sqdist_ab(A: np.ndarray, B: np.ndarray) -> np.ndarray:
-    """
-    Pairwise squared Euclidean distances between rows:
-      A: (a,d), B: (b,d)  ->  D2: (a,b)
+class GPLM:
     """
-    A = np.asarray(A, dtype=np.float64)
-    B = np.asarray(B, dtype=np.float64)
-    AA = np.sum(A * A, axis=1, keepdims=True)          # (a,1)
-    BB = np.sum(B * B, axis=1, keepdims=True).T        # (1,b)
-    D2 = AA + BB - 2.0 * (A @ B.T)
-    return np.maximum(D2, 0.0)
+    Inducing-point / Nyström GP (kernel ridge) decoder on latents.
 
+    Training:
+      R_ix: (N,d) latents
+      R_iX: (N,D) ambients
 
-class GPLM:
-    """
-    GPLM decoder with optional ANN sparse inducing prediction, plus a geodesic "flow" integrator.
-
-    Required stored arrays/params:
-      - Z_mx_w : (m,d) inducing points in *whitened latent coords*
-      - M_mX   : (m,D) inducing outputs
-      - mean_X : (D,) output mean added to decode
-      - lat_mean_x : (d,) latent mean (for whitening)
-      - lat_std_x  : (d,) latent std  (for whitening)
-      - beta, eps  : RBF hyperparams used exactly as in your _decode
-      - pred_k : if not None and < m, use ANN sparse neighbors
-      - ann_Z : object with method: search(query: (a,d), k:int) -> (idx: (a,k), D2: (a,k))
+      Choose inducing Z_mx (m << N), typically subset (medoids) of R_ix.
+
+      Latent kernel (unnormalized Gaussian affinity):
+        C_im = exp(-β * ||R_ix - Z_mx||^2 / ε)      (N,m)
+        W_mn = exp(-β * ||Z_mx - Z_nx||^2 / ε)      (m,m)
+
+      Nyström KRR/GP mean reduced solve:
+        M_mX = (C^T C + σ2 W + jitter I)^-1 (C^T (R_iX - mean_X))
+
+      Predict:
+        For novel R_ax:
+          find κ inducing neighbors (or all m if pred_κ=None)
+          C_am = exp(-β * ||R_ax - Z_mx||^2 / ε)
+          R_aX = C_am M_mX + mean_X
 
     Notes:
-      - __call__ and _decode are kept behavior-identical to your pasted version.
-      - flow() integrates geodesics of the pullback metric g = J^T J induced by this decoder.
+      - This implementation supports BOTH ascii kwargs and unicode kwargs
+        (β, ε, κ_eps, σ2, pred_κ, ε_use_kth, ε_mul).
+      - If whiten_latent=True, distances are computed in whitened latent space.
     """
 
     def __init__(
         self,
+        R_ix: np.ndarray,
+        R_iX: np.ndarray,
         *,
-        Z_mx_w: np.ndarray,
-        M_mX: np.ndarray,
-        mean_X: np.ndarray,
-        lat_mean_x: Optional[np.ndarray] = None,
-        lat_std_x: Optional[np.ndarray] = None,
+        # ASCII names (preferred for library APIs)
         beta: float = 1.0,
-        eps: float = 1.0,
-        pred_k: Optional[int] = None,
-        ann_Z=None,
-        dtype=np.float32 ):
-        self.dtype = dtype
 
-        self.Z_mx_w = np.asarray(Z_mx_w, dtype=np.float64)  # (m,d) whitened inducing
-        self.M_mX = np.asarray(M_mX, dtype=np.float64)      # (m,D)
-        self.mean_X = np.asarray(mean_X, dtype=np.float64)  # (D,)
+        # ε estimation
+        eps: Optional[float] = None,
+        k_eps: int = 256,
+        eps_use_kth: bool = True,
+        eps_mul: float = 1.0,
+
+        # regularization
+        sigma2: float = 1e-5,
+        jitter: float = 1e-8,
 
-        self.m = int(self.M_mX.shape[0])
-        self.d = int(self.Z_mx_w.shape[1])
-        self.D = int(self.M_mX.shape[1])
+        # inducing
+        m: int = 1024,
+        inducing: InducingMode = "kmeans_medoids",
+        Z_mx: Optional[np.ndarray] = None,
+        seed: int = 0,
 
-        if lat_mean_x is None:
-            lat_mean_x = np.zeros(self.d, dtype=np.float64)
-        if lat_std_x is None:
-            lat_std_x = np.ones(self.d, dtype=np.float64)
+        # preprocess
+        center_X: bool = True,
+        whiten_latent: bool = False,
+        dtype: Any = np.float32,
 
-        self.lat_mean_x = np.asarray(lat_mean_x, dtype=np.float64)  # (d,)
-        self.lat_std_x = np.asarray(lat_std_x, dtype=np.float64)    # (d,)
+        # compute/memory
+        fit_block: int = 8192,
 
+        # inference
+        pred_k: Optional[int] = None,
+        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_κ")
+
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        # ---- store params (provide both spellings) ----
         self.beta = float(beta)
-        self.eps = float(eps)
+        self.β = self.beta
+
+        self.sigma2 = float(sigma2)
+        self.σ2 = self.sigma2
+
+        self.jitter = float(jitter)
+        self.seed = int(seed)
+        self.dtype = dtype
+        self.fit_block = int(fit_block)
+
+        # ---- 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 ----
+        if 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) ----
+        Ztrain = R_ix.astype(np.float64)
+        if 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 = Ztrain_w  # (N,d) in float64
+
+        # ---- ANN on training latents (for eps + medoids snapping) ----
+        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 on latents ----
+        if eps is None:
+            k_eps = int(min(max(8, int(k_eps)), self.N - 1))
+            # ask for k_eps+1 to try to include self
+            j_iK1, D2_iK1 = self.ann_train.search(self.R_ix_w.astype(self.dtype, copy=False), k_eps + 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), dtype=np.float64)
+                for ii in range(self.N):
+                    keep = (j_iK1[ii] != ii)
+                    D2_iK[ii] = D2_iK1[ii][keep][:k_eps]
+            else:
+                D2_iK = D2_iK1[:, :k_eps].astype(np.float64, copy=False)
 
-        self.pred_k = pred_k
-        self.ann_Z = ann_Z
+            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
+
+        # ---- choose inducing points (in whitened latent space) ----
+        rng = np.random.default_rng(self.seed)
+        m = int(min(max(1, int(m)), self.N))
+
+        if Z_mx is not None:
+            Zm = np.asarray(Z_mx, dtype=np.float64)
+            if Zm.ndim != 2 or Zm.shape[1] != self.d_lat:
+                raise ValueError("Z_mx must be (m, d_lat).")
+            Zm_w = (Zm - self.lat_mean_x) / self.lat_std_x
+        else:
+            if inducing == "random_subset":
+                idx = rng.choice(self.N, size=m, replace=False)
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "fps":
+                idx = fps_indices(self.R_ix_w, m=m, seed=self.seed)
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "kmeans_medoids":
+                C = _kmeans2_safe(self.R_ix_w, m, seed=self.seed).astype(np.float64, copy=False)
+                j_cm, _ = self.ann_train.search(C.astype(self.dtype, copy=False), 1)
+                idx = j_cm.reshape(-1).astype(np.int64)
+
+                # de-duplicate and refill if needed
+                idx_u = np.unique(idx)
+                if idx_u.size < m:
+                    needed = m - idx_u.size
+                    pool = np.setdiff1d(np.arange(self.N), idx_u, assume_unique=False)
+                    if pool.size >= needed:
+                        extra = rng.choice(pool, size=needed, replace=False)
+                    else:
+                        extra = rng.choice(self.N, size=needed, replace=True)
+                    idx = np.concatenate([idx_u, extra])
+                else:
+                    idx = idx_u[:m]
 
-        # matches your code: scale used in dW = scale * W * diff
-        self._rbf_grad_scale = float(-2.0 * self.beta / self.eps)
+                Zm_w = self.R_ix_w[idx]
+            elif inducing == "given":
+                raise ValueError("Provide Z_mx when inducing='given'.")
+            else:
+                raise ValueError(f"Unknown inducing mode: {inducing!r}")
 
-        if self.pred_k is not None:
-            if int(self.pred_k) <= 0:
-                raise ValueError("pred_k must be positive or None.")
-            if self.ann_Z is None:
-                raise ValueError("ann_Z must be provided when pred_k is not None.")
+        self.Z_mx_w = np.ascontiguousarray(Zm_w.astype(np.float64, copy=False))
+        self.m = int(self.Z_mx_w.shape[0])
 
-    # -------------------------
-    # Your original public API
-    # -------------------------
-    def __call__(
-        self,
-        R_ax: Union[np.ndarray, list],
-        *,
-        batch_size: Optional[int] = None,
-        jacobian: bool = False,
-        metric: bool = False ):
-        """
-        Decode latents to ambient.
+        # also store raw inducing points (unwhitened) for serialization convenience
+        self.Z_mx = (self.Z_mx_w * self.lat_std_x[None, :]) + self.lat_mean_x[None, :]
 
-        Returns:
-          - if jacobian=False and metric=False: Y (a,D)
-          - if jacobian=True, metric=False: (Y, J) where J is (a,D,d)
-          - if jacobian=False, metric=True: (Y, g) where g is (a,d,d)
-          - if jacobian=True, metric=True: (Y, J, g)
-        """
+        # ---- ANN on inducing points for fast prediction ----
+        self.ann_Z, _ = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        self.ann_Z.build(self.Z_mx_w.astype(self.dtype, copy=False))
+
+        # pred_k
+        if pred_k is None:
+            self.pred_k = None
+        else:
+            self.pred_k = int(min(max(1, int(pred_k)), self.m))
+        self.pred_κ = self.pred_k  # unicode alias
+
+        # ---- W_mm ----
+        D2_mm = sqdist_ab(self.Z_mx_w, self.Z_mx_w)
+        W_mm = np.exp(-self.beta * (D2_mm.astype(np.float64) / self.eps))
+        W_mm.flat[:: self.m + 1] += self.jitter
+        self.W_mm = W_mm  # (m,m)
+
+        # ---- accumulate G=C^T C and B=C^T Y ----
+        G_mm = np.zeros((self.m, self.m), dtype=np.float64)
+        B_mX = np.zeros((self.m, self.D), dtype=np.float64)
+
+        bs = int(self.fit_block)
+        for i0 in range(0, self.N, bs):
+            i1 = min(self.N, i0 + bs)
+            Zi = self.R_ix_w[i0:i1]  # (b,d) float64
+            D2_im = sqdist_ab(Zi, self.Z_mx_w)
+            C_im = np.exp(-self.beta * (D2_im.astype(np.float64) / self.eps))
+            G_mm += C_im.T @ C_im
+            B_mX += C_im.T @ Y[i0:i1]
+
+        A_mm = G_mm + self.sigma2 * W_mm
+        A_mm.flat[:: self.m + 1] += self.jitter
+
+        cF = la.cho_factor(A_mm, lower=True, check_finite=False)
+        self.M_mX = la.cho_solve(cF, B_mX, check_finite=False)  # (m,D)
+
+    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:
@@ -135,319 +285,218 @@ class GPLM:
         R_ax = np.ascontiguousarray(R_ax.astype(self.dtype, copy=False))
 
         if batch_size is None:
-            out = self._decode(R_ax, jacobian=jacobian, metric=metric)
+            Y = self._decode(R_ax)
         else:
             bs = int(batch_size)
-            Ys = []
-            Js = [] if jacobian else None
-            Gs = [] if metric else None
+            out = []
             for s in range(0, R_ax.shape[0], bs):
-                chunk = self._decode(R_ax[s:s + bs], jacobian=jacobian, metric=metric)
-                if (not jacobian) and (not metric):
-                    Ys.append(chunk)
-                elif jacobian and (not metric):
-                    Yc, Jc = chunk
-                    Ys.append(Yc); Js.append(Jc)  # type: ignore[arg-type]
-                elif (not jacobian) and metric:
-                    Yc, gc = chunk
-                    Ys.append(Yc); Gs.append(gc)  # type: ignore[arg-type]
-                else:
-                    Yc, Jc, gc = chunk
-                    Ys.append(Yc); Js.append(Jc); Gs.append(gc)  # type: ignore[arg-type]
-
-            Y = np.vstack(Ys)
-            if (not jacobian) and (not metric):
-                out = Y
-            elif jacobian and (not metric):
-                out = (Y, np.vstack(Js))  # type: ignore[arg-type]
-            elif (not jacobian) and metric:
-                out = (Y, np.vstack(Gs))  # type: ignore[arg-type]
-            else:
-                out = (Y, np.vstack(Js), np.vstack(Gs))  # type: ignore[arg-type]
+                out.append(self._decode(R_ax[s:s + bs]))
+            Y = np.vstack(out)
 
-        # unwrap single
-        if single:
-            if (not jacobian) and (not metric):
-                return out[0]  # type: ignore[index]
-            if jacobian and (not metric):
-                Y, J = out  # type: ignore[misc]
-                return Y[0], J[0]
-            if (not jacobian) and metric:
-                Y, g = out  # type: ignore[misc]
-                return Y[0], g[0]
-            Y, J, g = out  # type: ignore[misc]
-            return Y[0], J[0], g[0]
-
-        return out
-
-    def _decode(
-        self,
-        R_ax: np.ndarray,
-        *,
-        jacobian: bool,
-        metric: bool ):
-        """
-        Internal decode for a batch (a,d).
-
-        Jacobian is w.r.t. *unwhitened* input R_ax (chain rule applied if whiten_latent=True).
-        """
-        Za = R_ax.astype(np.float64)                            # (a,d)
-        Za_w = (Za - self.lat_mean_x) / self.lat_std_x          # (a,d)
+        return Y[0] if single else Y
 
-        # We'll need this for chain rule in gradients:
-        inv_std = (1.0 / self.lat_std_x).astype(np.float64)     # (d,)
+    def _decode(self, R_ax: np.ndarray) -> np.ndarray:
+        Za = R_ax.astype(np.float64)
+        Za_w = (Za - self.lat_mean_x) / self.lat_std_x
 
         if self.pred_k is None or self.pred_k == self.m:
-            # full inducing
-            D2_am = sqdist_ab(Za_w, self.Z_mx_w)                                # (a,m)
-            W_am = np.exp(-self.beta * (D2_am.astype(np.float64) / self.eps))   # (a,m)
-            Y = W_am @ self.M_mX                                                # (a,D)
-
-            if not (jacobian or metric):
-                return Y + self.mean_X[None, :]
-
-            # diff in whitened latent space: (a,m,d)
-            diff_amd = Za_w[:, None, :] - self.Z_mx_w[None, :, :]
-
-            # dW/dZa_w: (a,m,d)
-            dW_amd = (self._rbf_grad_scale * W_am[:, :, None]) * diff_amd
-
-            # chain to unwhitened: dW/dZa = dW/dZa_w * dZa_w/dZa = * (1/std)
-            dW_amd *= inv_std[None, None, :]
-
-            if jacobian:
-                # J_{a,D,d} = Σ_m dW_{a,m,d} * M_{m,D}
-                J_aDd = np.einsum("amd,mD->aDd", dW_amd, self.M_mX, optimize=True)
-                if metric:
-                    g_add = np.einsum("aDd,aDe->ade", J_aDd, J_aDd, optimize=True)
-                    return (Y + self.mean_X[None, :], J_aDd, g_add)
-                return (Y + self.mean_X[None, :], J_aDd)
-
-            # metric-only
-            # Build J implicitly then g; for simplicity we materialize J here.
-            J_aDd = np.einsum("amd,mD->aDd", dW_amd, self.M_mX, optimize=True)
-            g_add = np.einsum("aDd,aDe->ade", J_aDd, J_aDd, optimize=True)
-            return (Y + self.mean_X[None, :], g_add)
-
-        # sparse inducing neighbors for speed
-        j_aK, D2_aK = self.ann_Z.search(Za_w.astype(self.dtype, copy=False), self.pred_k)  # (a,k), (a,k)
-        W_aK = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))                  # (a,k)
-        M_aKD = self.M_mX[j_aK]                                                            # (a,k,D)
-        Y = np.sum(W_aK[:, :, None] * M_aKD, axis=1)                                       # (a,D)
-
-        if not (jacobian or metric):
-            return Y + self.mean_X[None, :]
-
-        Zsel_aKd = self.Z_mx_w[j_aK]               # (a,k,d) in whitened latent
-        diff_aKd = Za_w[:, None, :] - Zsel_aKd     # (a,k,d)
-
-        dW_aKd = (self._rbf_grad_scale * W_aK[:, :, None]) * diff_aKd
-        dW_aKd *= inv_std[None, None, :]
-
-        if jacobian:
-            # J_{a,D,d} = Σ_k dW_{a,k,d} * M_{a,k,D}
-            J_aDd = np.einsum("akd,akD->aDd", dW_aKd, M_aKD, optimize=True)
-            if metric:
-                g_add = np.einsum("aDd,aDe->ade", J_aDd, J_aDd, optimize=True)
-                return (Y + self.mean_X[None, :], J_aDd, g_add)
-            return (Y + self.mean_X[None, :], J_aDd)
-
-        # metric-only
-        J_aDd = np.einsum("akd,akD->aDd", dW_aKd, M_aKD, optimize=True)
-        g_add = np.einsum("aDd,aDe->ade", J_aDd, J_aDd, optimize=True)
-        return (Y + self.mean_X[None, :], g_add)
-
-    # -------------------------
-    # NEW: Geodesic flow (analytic Christoffels) with generalized velocity-Verlet
-    # -------------------------
-    def flow(
-        self,
-        R_ax: Union[np.ndarray, list],
-        v0_ax: Union[np.ndarray, list],
-        *,
-        dt: float = 0.05,
-        reg: float = 1e-8,
-        keep_speed: bool = True,
-        ) -> Tuple[np.ndarray, np.ndarray]:
-        """
-        One generalized velocity-Verlet / leapfrog step for the geodesic ODE:
-          r_dot = v
-          v_dot^x = - Γ^x_{yz}(r) v^y v^z
-
-        Returns:
-          Q_ax : next latent positions (a,d)
-          v_ax : next latent velocities (a,d)
-
-        Use in a loop:
-          Q, v = R0, v0
-          for _ in range(T):
-              Q, v = gplm.flow(Q, v, dt=...)
-        """
-        Z = np.asarray(R_ax, dtype=np.float64)
-        v = np.asarray(v0_ax, dtype=np.float64)
-
-        single = (Z.ndim == 1)
-        if single:
-            Z = Z[None, :]
-        if v.ndim == 1:
-            v = v[None, :]
+            D2_am = sqdist_ab(Za_w, self.Z_mx_w)
+            C_am = np.exp(-self.beta * (D2_am.astype(np.float64) / self.eps))
+            Y = C_am @ self.M_mX
+        else:
+            j_aK, D2_aK = self.ann_Z.search(Za_w.astype(self.dtype, copy=False), self.pred_k)
+            W = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))  # (a,k)
+            M = self.M_mX[j_aK]                                              # (a,k,D)
+            Y = np.sum(W[:, :, None] * M, axis=1)                             # (a,D)
 
-        if Z.shape != v.shape:
-            raise ValueError(f"R_ax and v0_ax must have same shape; got {Z.shape} vs {v.shape}")
+        return Y + self.mean_X[None, :]
 
-        # a_n = a(z_n, v_n)
-        a0, g0 = self._geodesic_accel_analytic(Z, v, reg=reg)
+    # -----------------------------
+    # Geometry helpers (RBF kernel)
+    # -----------------------------
+    def _whiten(self, R_ax: np.ndarray) -> np.ndarray:
+        R = np.asarray(R_ax, dtype=np.float64)
+        return (R - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
 
-        # v_{n+1/2}
-        v_half = v + 0.5 * float(dt) * a0
+    def _inducing_idx(self, Rw_ax: np.ndarray) -> Optional[np.ndarray]:
+        """
+        Return inducing indices per query if pred_k is set; else None means "use all m".
+        Shape: (A, k) int64 if pred_k; else None.
+        """
+        if self.pred_k is None or self.pred_k == self.m:
+            return None
+        j_aK, _ = self.ann_Z.search(Rw_ax.astype(self.dtype, copy=False), self.pred_k)
+        return j_aK.astype(np.int64, copy=False)
 
-        # z_{n+1}
-        Q = Z + float(dt) * v_half
+    def _rbf_terms_single(
+        self,
+        r_x: np.ndarray,
+        v_x: np.ndarray,
+        *,
+        idx_m: Optional[np.ndarray],
+        ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        """
+        For a single point (no batch), compute:
+          g_dd: pullback metric in latent coords (d,d)
+          cF:   Cholesky factorization handle for solves with g
+          f_d:  geodesic momentum force term (d,)
+          p_d:  momentum corresponding to input velocity (p = g v)
 
-        # a_{n+1} = a(z_{n+1}, v_{n+1/2})
-        a1, g1 = self._geodesic_accel_analytic(Q, v_half, reg=reg)
+        Uses efficient contractions with C_mm (no dependence on ambient dimension D).
+        """
+        d = self.d_lat
+        c = self.beta / self.eps  # = beta/eps
+        inv_std = 1.0 / self.lat_std_x            # (d,)
+        inv_std2 = inv_std * inv_std              # (d,)
+
+        # whitened coordinates (float64)
+        rw = (r_x.astype(np.float64) - self.lat_mean_x) / self.lat_std_x  # (d,)
+
+        # select inducing points
+        if idx_m is None:
+            Zw = self.Z_mx_w                      # (m,d)
+            C = self.C_mm                         # (m,m)
+        else:
+            Zw = self.Z_mx_w[idx_m]               # (k,d)
+            C = self.C_mm[np.ix_(idx_m, idx_m)]   # (k,k)
 
-        # v_{n+1}
-        v_new = v_half + 0.5 * float(dt) * a1
+        Dw = rw[None, :] - Zw                     # (k,d)  (k=m if idx_m is None)
+        D2 = np.sum(Dw * Dw, axis=1)              # (k,)
+        k_m = np.exp(-c * D2)                     # (k,)
 
-        # Optional: keep Riemannian speed sqrt(v^T g v) constant per sample
-        if keep_speed:
-            s0 = np.sqrt(np.maximum(np.einsum("ai,aij,aj->a", v, g0, v, optimize=True), 1e-16))
-            s1 = np.sqrt(np.maximum(np.einsum("ai,aij,aj->a", v_new, g1, v_new, optimize=True), 1e-16))
-            v_new = v_new * (s0 / s1)[:, None]
+        # Dw_over = Dw / std  (this is the extra chain-rule factor for derivatives wrt unwhitened R)
+        Dw_over = Dw * inv_std[None, :]           # (k,d)
 
-        if single:
-            return Q[0], v_new[0]
-        return Q, v_new
+        # s_m = sum_y v_y * (Dw_y / std_y)   (k,)
+        v = v_x.astype(np.float64, copy=False)
+        s_m = Dw_over @ v                         # (k,)
 
-    def _geodesic_accel_analytic(
-        self,
-        R_ax: np.ndarray,
-        v_ax: np.ndarray,
-        *,
-        reg: float,
-    ) -> Tuple[np.ndarray, np.ndarray]:
-        """
-        Compute acceleration a_ax = -Γ^x_{yz}(r) v^y v^z using analytic Christoffels,
-        and return (a_ax, g_aij).
-
-        Shapes:
-          R_ax: (a,d)
-          v_ax: (a,d)
-          a_ax: (a,d)
-          g:    (a,d,d)
-        """
-        Z = np.asarray(R_ax, dtype=np.float64)
-        v = np.asarray(v_ax, dtype=np.float64)
-        a, d = Z.shape
-        if v.shape != (a, d):
-            raise ValueError(f"v_ax must have shape {(a, d)}, got {v.shape}")
+        # grad_k[m,x] = d/dR_x k_m
+        # = -(2c) * k_m * (Dw_x / std_x)
+        grad_k = -(2.0 * c) * (k_m[:, None] * Dw_over)   # (k,d)
 
-        # Precompute whitening + chain rule
-        Za_w = (Z - self.lat_mean_x) / self.lat_std_x         # (a,d)
-        inv_std = (1.0 / self.lat_std_x).astype(np.float64)   # (d,)
-        alpha = float(self._rbf_grad_scale)
+        # pullback metric: g = grad^T C grad   (d,d)
+        # (this equals J^T J with J_Xx = sum_m (dk_m/dR_x) M_mX)
+        Cg = C @ grad_k                             # (k,d)
+        g = grad_k.T @ Cg                           # (d,d)
+        g = 0.5 * (g + g.T)                         # symmetrize for numerical stability
+        # regularize to ensure PD for Cholesky
+        lam = 1e-10
+        g.flat[:: d + 1] += lam
 
-        I = np.eye(d, dtype=np.float64)
+        cF = la.cho_factor(g, lower=True, check_finite=False)
 
-        if self.pred_k is None or self.pred_k == self.m:
-            # ----- full inducing -----
-            diff_amd = Za_w[:, None, :] - self.Z_mx_w[None, :, :]         # (a,m,d)
-            D2_am = np.sum(diff_amd * diff_amd, axis=2)                   # (a,m)
-            k_am = np.exp(-self.beta * (D2_am / self.eps))                # (a,m)
+        # momentum from velocity: p = g v
+        p = g @ v
 
-            # ∂_y k_am = alpha * k_am * D_amy * (1/σ_y)
-            dk_amd = (alpha * k_am[:, :, None]) * diff_amd                # (a,m,d) (still whitened diffs)
-            dk_amd *= inv_std[None, None, :]                              # chain to unwhitened
+        # S_m = v^T grad_k[m,:] = -(2c) k_m s_m
+        S_m = -(2.0 * c) * (k_m * s_m)              # (k,)
 
-            # J_{aXy} = Σ_m (∂_y k_am) M_{mX}
-            J_aXy = np.einsum("amd,mX->aXd", dk_amd, self.M_mX, optimize=True)  # (a,D,d)
+        # T_xm = sum_y v_y * d^2/dR_x dR_y k_m
+        # Hessian_{xy} = k_m * [ 4c^2 * (Dw_x/std_x)(Dw_y/std_y) - 2c * delta_xy / std_x^2 ]
+        # contraction over y gives:
+        # T_xm = k_m * [ 4c^2*(Dw_x/std_x)*s_m - 2c*(v_x/std_x^2) ]
+        term1 = (4.0 * c * c) * (k_m * s_m)[:, None] * Dw_over          # (k,d)
+        term2 = (2.0 * c) * k_m[:, None] * (v[None, :] * inv_std2[None, :])  # (k,d)
+        T = (term1 - term2).T                                           # (d,k)
 
-            # metric g_{ayz} = J_{aXy} J_{aXz}
-            g_ayz = np.einsum("aXd,aXe->ade", J_aXy, J_aXy, optimize=True)      # (a,d,d)
-            g_ayz = 0.5 * (g_ayz + np.swapaxes(g_ayz, 1, 2))                    # symmetrize
+        # force: f = T @ (C @ S)
+        w = C @ S_m                                                     # (k,)
+        f = T @ w                                                       # (d,)
 
-            g_inv = np.linalg.inv(g_ayz + reg * I[None, :, :])                  # (a,d,d)
+        return g, cF, f, p
 
-            # dg[a, w, y, z] = ∂_w g_{yz}
-            dg_awyz = np.zeros((a, d, d, d), dtype=np.float64)
+    def _solve(self, cF: Tuple[np.ndarray, bool], b: np.ndarray) -> np.ndarray:
+        return la.cho_solve(cF, b, check_finite=False)
 
-            # For each derivative index w, build Hessian slice H_{X w y} (as (a,D,y))
-            for w in range(d):
-                diff_w_am = diff_amd[:, :, w]  # (a,m)
+    # -----------------------------
+    # Geodesic generalized leapfrog
+    # -----------------------------
+    def flow(
+        self,
+        R_ax: Union[np.ndarray, list],
+        v_ax: Union[np.ndarray, list],
+        *,
+        eps: float = 1e-2,
+        K_p: int = 5,
+        K_q: int = 5,
+        ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        One generalized-leapfrog step for geodesic flow on the pullback manifold.
 
-                # ∂_{w y} k = k * [ alpha^2 * D_w * D_y + alpha * δ_{wy} ] * (1/σ_w)(1/σ_y)
-                d2k_amy = k_am[:, :, None] * (
-                    (alpha * alpha) * diff_w_am[:, :, None] * diff_amd
-                    + alpha * I[w][None, None, :]
-                )  # (a,m,d) in whitened diffs
-                d2k_amy *= (inv_std[w] * inv_std[None])[None, None, :]  # chain factors (1/σ_w)(1/σ_y)
+        Input:
+          R_ax: (A,d) or (d,)
+          v_ax: (A,d) or (d,)  (velocity in *unwhitened* latent coordinates)
 
-                # H_{aXy} for this w: H_{aXwy} = Σ_m (∂_{w y} k_am) M_{mX}
-                H_w_aXy = np.einsum("amy,mX->aXd", d2k_amy, self.M_mX, optimize=True)  # (a,D,d)
+        Output:
+          (R_next, v_next) with same shapes as input.
 
-                # ∂_w g_{yz} = H_{X w y} J_{X z} + J_{X y} H_{X w z}
-                dg_w = (
-                    np.einsum("aXd,aXe->ade", H_w_aXy, J_aXy, optimize=True)
-                    + np.einsum("aXd,aXe->ade", J_aXy, H_w_aXy, optimize=True)
-                )  # (a,d,d)
-                dg_awyz[:, w, :, :] = dg_w
+        Notes:
+          - Uses fixed-point iterations with fixed K_p, K_q for determinism/time-reversibility.
+          - If pred_k is set, uses only the k nearest inducing points per query for geometry.
+        """
+        R = np.asarray(R_ax, dtype=np.float64)
+        v = np.asarray(v_ax, dtype=np.float64)
 
-        else:
-            # ----- sparse ANN inducing neighbors -----
-            j_aK, D2_aK = self.ann_Z.search(Za_w.astype(self.dtype, copy=False), int(self.pred_k))  # (a,k),(a,k)
-            k_aK = np.exp(-self.beta * (D2_aK.astype(np.float64) / self.eps))                        # (a,k)
+        single = (R.ndim == 1)
+        if single:
+            R = R[None, :]
+            v = v[None, :]
+        if R.shape != v.shape or R.ndim != 2 or R.shape[1] != self.d_lat:
+            raise ValueError(f"Expected R_ax and v_ax to have shape (A,{self.d_lat}) (or ({self.d_lat},)).")
 
-            Zsel_aKd = self.Z_mx_w[j_aK]                    # (a,k,d)
-            diff_aKd = Za_w[:, None, :] - Zsel_aKd          # (a,k,d)
-            M_aKX = self.M_mX[j_aK]                         # (a,k,D)
+        A, d = R.shape
 
-            dk_aKd = (alpha * k_aK[:, :, None]) * diff_aKd
-            dk_aKd *= inv_std[None, None, :]
+        # inducing neighbor indices per point (optional)
+        Rw = self._whiten(R)
+        idx_aK = self._inducing_idx(Rw)  # None or (A,k)
 
-            # J_{aXy} = Σ_k (∂_y k_aK) M_{aKX}
-            J_aXy = np.einsum("akd,akX->aXd", dk_aKd, M_aKX, optimize=True)  # (a,D,d)
+        R_next = np.empty_like(R)
+        v_next = np.empty_like(v)
 
-            g_ayz = np.einsum("aXd,aXe->ade", J_aXy, J_aXy, optimize=True)   # (a,d,d)
-            g_ayz = 0.5 * (g_ayz + np.swapaxes(g_ayz, 1, 2))
+        for a in range(A):
+            idx = None if idx_aK is None else idx_aK[a]
 
-            g_inv = np.linalg.inv(g_ayz + reg * I[None, :, :])               # (a,d,d)
+            r_n = R[a]
+            v_n = v[a]
 
-            dg_awyz = np.zeros((a, d, d, d), dtype=np.float64)
+            # metric/force at step start (for implicit p-half)
+            g_n, cF_n, f_n, p_n = self._rbf_terms_single(r_n, v_n, idx_m=idx)
 
-            for w in range(d):
-                diff_w_aK = diff_aKd[:, :, w]  # (a,k)
+            # --- (1) implicit half-step in momentum via fixed-point on p ---
+            p = p_n.copy()
+            for _ in range(int(K_p)):
+                v_k = self._solve(cF_n, p)                     # v = g_n^{-1} p
+                _, _, f_k, _ = self._rbf_terms_single(r_n, v_k, idx_m=idx)
+                p = p_n + 0.5 * eps * f_k
+            p_half = p
 
-                d2k_aKy = k_aK[:, :, None] * (
-                    (alpha * alpha) * diff_w_aK[:, :, None] * diff_aKd
-                    + alpha * I[w][None, None, :]
-                )  # (a,k,d)
-                d2k_aKy *= (inv_std[w] * inv_std[None])[None, None, :]
+            # --- (2) implicit position update via fixed-point on r ---
+            v_half_n = self._solve(cF_n, p_half)               # M(r_n) p_half
+            r = r_n + eps * v_half_n                           # init
 
-                H_w_aXy = np.einsum("akd,akX->aXd", d2k_aKy, M_aKX, optimize=True)  # (a,D,d)
+            for _ in range(int(K_q)):
+                g_r, cF_r, _, _ = self._rbf_terms_single(r, v_half_n, idx_m=idx)  # metric at current r
+                v_half_r = self._solve(cF_r, p_half)            # M(r) p_half
+                r = r_n + 0.5 * eps * (v_half_n + v_half_r)
 
-                dg_w = (
-                    np.einsum("aXd,aXe->ade", H_w_aXy, J_aXy, optimize=True)
-                    + np.einsum("aXd,aXe->ade", J_aXy, H_w_aXy, optimize=True)
-                )
-                dg_awyz[:, w, :, :] = dg_w
+            r_np1 = r
 
-        # Christoffels:
-        # A[y,z,w] = ∂_y g_{zw} + ∂_z g_{yw} - ∂_w g_{yz}
-        # with dg stored as dg[w,y,z] = ∂_w g_{yz}
-        A_ayzw = (
-            dg_awyz
-            + dg_awyz.transpose(0, 2, 1, 3)
-            - dg_awyz.transpose(0, 2, 3, 1)
-        )  # (a,y,z,w)
+            # --- (3) explicit half-step in momentum at r_{n+1} ---
+            # recompute metric at r_{n+1} and compute force using v_mid = M(r_{n+1}) p_half
+            g_np1, cF_np1, _, _ = self._rbf_terms_single(r_np1, v_half_n, idx_m=idx)
+            v_mid = self._solve(cF_np1, p_half)
+            _, _, f_np1, _ = self._rbf_terms_single(r_np1, v_mid, idx_m=idx)
 
-        # Gamma[a,x,y,z] = 1/2 * g_inv[a,x,w] * A[a,y,z,w]
-        Gamma_axyz = 0.5 * np.einsum("axw,ayzw->axyz", g_inv, A_ayzw, optimize=True)
+            p_np1 = p_half + 0.5 * eps * f_np1
+            v_np1 = self._solve(cF_np1, p_np1)
 
-        # acceleration: a^x = - Γ^x_{yz} v^y v^z
-        a_ax = -np.einsum("axyz,ay,az->ax", Gamma_axyz, v, v, optimize=True)
+            R_next[a] = r_np1
+            v_next[a] = v_np1
 
-        return a_ax, g_ayz
+        if single:
+            return R_next[0], v_next[0]
+        return R_next, v_next
 
 
-__all__ = ["GPLM", "InducingMode"]
+__all__ = ["GPLM", "InducingMode"]