Update src/dima/dima.py

src/dima/dima.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 import json
 import time
 from dataclasses import asdict, dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import numpy as np
 
@@ -14,22 +14,25 @@ from jax import random
 
 from flax import serialization as flax_ser
 
+from .ann import ANNBackend, make_ann
+from .ddpm import DDPM
 from .dmap import DMAP
 from .gplm import GPLM
-from .ddpm import DDPM
-from .ann import ANNBackend, make_ann
 
 
 # ----------------------------
 # Optional: Hugging Face Hub
 # ----------------------------
 try:
-    from huggingface_hub import HfApi, HfFolder, upload_file, hf_hub_download
+    from huggingface_hub import HfApi, HfFolder, upload_file, hf_hub_download  # type: ignore
     _HAS_HF = True
 except Exception:
     _HAS_HF = False
 
 
+_UNSET = object()
+
+
 def _select_device(prefer: str = "auto"):
     """
     Safe JAX device selection.
@@ -44,7 +47,6 @@ def _select_device(prefer: str = "auto"):
         return gpu[0] if gpu else (cpu[0] if cpu else devs[0])
     if prefer == "cpu":
         return cpu[0] if cpu else devs[0]
-    # unknown -> auto
     return gpu[0] if gpu else (cpu[0] if cpu else devs[0])
 
 
@@ -74,13 +76,18 @@ class FrozenDMAP:
     ):
         self.k = int(state["k"])
         self.beta = float(state["beta"])
+        self.β = self.beta
         self.alpha = float(state["alpha"])
+        self.α = self.alpha
         self.eps = float(state["eps"])
+        self.ε = self.eps
         self.dtype = np.dtype(state.get("dtype", "float32"))
 
         self.R_iX = np.ascontiguousarray(np.asarray(state["R_iX"]).astype(self.dtype, copy=False))
         self.qalpha_i = np.asarray(state["qalpha_i"], dtype=np.float64)
+        self.qα_i = self.qalpha_i  # alias
         self.R_over_lam_ix = np.asarray(state["R_over_lam_ix"], dtype=np.float64)  # (Nref,d)
+        self.R_over_λ_ix = self.R_over_lam_ix  # alias
 
         self.ann, self.ann_backend = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
         self.ann.build(self.R_iX)
@@ -122,81 +129,41 @@ class FrozenDMAP:
         return Z_ax
 
 
-class FrozenGPLM:
+def _restore_gplm_as_object(
+    state: Dict[str, Any],
+    *,
+    ann_backend: ANNBackend = "auto",
+    ann_params: Optional[Dict[str, Any]] = None,
+    n_jobs: int = -1,
+) -> GPLM:
     """
-    Inference-only GPLM decoder built from saved GPLM state.
-    Supports optional pred_k via ANN over inducing points.
+    Rehydrate a GPLM instance from saved state WITHOUT retraining.
+    This is deliberately done as a true GPLM instance so you also get GPLM.flow(...)
+    (assuming your GPLM class implements .flow()).
     """
+    obj = GPLM.__new__(GPLM)  # type: ignore
 
-    def __init__(
-        self,
-        state: Dict[str, Any],
-        *,
-        ann_backend: ANNBackend = "auto",
-        ann_params: Optional[Dict[str, Any]] = None,
-        n_jobs: int = -1,
-    ):
-        self.beta = float(state["beta"])
-        self.eps = float(state["eps"])
-        self.pred_k = None if state.get("pred_k", None) is None else int(state["pred_k"])
-        self.dtype = np.dtype(state.get("dtype", "float32"))
-
-        self.mean_X = np.asarray(state["mean_X"], dtype=np.float64)
-        self.M_mX = np.asarray(state["M_mX"], dtype=np.float64)
+    obj.beta = float(state["beta"])
+    obj.β = obj.beta
+    obj.eps = float(state["eps"])
+    obj.ε = obj.eps
+    obj.pred_k = None if state.get("pred_k", None) is None else int(state["pred_k"])
+    obj.pred_κ = obj.pred_k
+    obj.dtype = np.dtype(state.get("dtype", "float32"))
 
-        # latent whitening used at training time
-        self.lat_mean_x = np.asarray(state["lat_mean_x"], dtype=np.float64)
-        self.lat_std_x = np.asarray(state["lat_std_x"], dtype=np.float64)
+    obj.mean_X = np.asarray(state["mean_X"], dtype=np.float64)
+    obj.M_mX = np.asarray(state["M_mX"], dtype=np.float64)
 
-        # inducing points in whitened latent space
-        self.Z_mx_w = np.ascontiguousarray(np.asarray(state["Z_mx_w"]).astype(np.float64, copy=False))
-        self.m = int(self.Z_mx_w.shape[0])
+    obj.lat_mean_x = np.asarray(state["lat_mean_x"], dtype=np.float64)
+    obj.lat_std_x = np.asarray(state["lat_std_x"], dtype=np.float64)
 
-        self.ann_Z, self.ann_backend = 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))
+    obj.Z_mx_w = np.ascontiguousarray(np.asarray(state["Z_mx_w"]).astype(np.float64, copy=False))
+    obj.m = int(obj.Z_mx_w.shape[0])
 
-    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))
+    obj.ann_Z, _ = make_ann(ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+    obj.ann_Z.build(obj.Z_mx_w.astype(obj.dtype, copy=False))
 
-        if batch_size is None:
-            Y = self._decode(R_ax)
-        else:
-            out = []
-            bs = int(batch_size)
-            for s in range(0, R_ax.shape[0], bs):
-                out.append(self._decode(R_ax[s : s + bs]))
-            Y = np.vstack(out)
-
-        return Y[0] if single else Y
-
-    @staticmethod
-    def _sqdist_ab(A: np.ndarray, B: np.ndarray) -> np.ndarray:
-        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)
-
-    def _decode(self, R_ax: np.ndarray) -> np.ndarray:
-        Za = R_ax.astype(np.float64)
-        Za_w = (Za - self.lat_mean_x[None, :]) / self.lat_std_x[None, :]
-
-        if self.pred_k is None or self.pred_k >= self.m:
-            D2_am = self._sqdist_ab(Za_w, self.Z_mx_w)
-            C_am = np.exp(-self.beta * (D2_am / 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)
-
-        return Y + self.mean_X[None, :]
+    return obj
 
 
 # ----------------------------
@@ -205,9 +172,9 @@ class FrozenGPLM:
 @dataclass
 class DIMAConfig:
     d: int = 32
+    beta: float = 1.0
     ddpm_device: str = "auto"   # "auto" | "cpu" | "gpu"
-    # purely informational defaults (actual settings live in kwargs dicts)
-    version: str = "0.1.0"
+    version: str = "0.2.0"
 
 
 # ----------------------------
@@ -215,99 +182,196 @@ class DIMAConfig:
 # ----------------------------
 class DIMA:
     """
-    DIMA: Diffusion–Intrinsic Manifold Autoencoder
-
-    Pipeline:
-      - DMAP encoder (NumPy/SciPy, CPU): ambient -> raw latent
-      - GPLM decoder (NumPy/SciPy, CPU): raw latent -> ambient
-      - DDPM latent model (JAX/Flax, CPU or GPU): trained on *normalized* latents
-
-    API:
-      model = DIMA(R_iX, dmap_kwargs=..., gplm_kwargs=..., ddpm_kwargs=..., ddpm_device="auto")
-      Z = model.encode(X)           # normalized latents (jnp on ddpm device)
-      Xhat = model.decode(Z)        # ambient (np on CPU), optionally with DDPM refinement
-      Xhat2 = model(X)              # polymorphic: ambient -> latent
-      Xhat3 = model(Z)              # polymorphic: latent -> ambient (refine=True default)
-      Xsamp = model.sample(1000)    # unconditional samples (ambient np)
+    DIMA: DMAP encoder + (latent DDPM) + GPLM decoder.
+
+    Public “user-facing” convention in this wrapper:
+
+      - Raw DMAP coordinates are the *public latent* (np.ndarray): R_ax (a,d)
+      - Normalized latents are the DDPM coordinates (jax/np):      Z_ax (a,d)
+
+    Minimal user API (what you asked for):
+
+        dima = DIMA(R_iX, d=20, beta=2.0)
+        R_ax = dima(R_aX)        # encode ambient -> raw latents
+        Q_aX = dima(R_ax)        # decode raw latents -> ambient
+
+    You can still pass full dict overrides for any submodule:
+        dima = DIMA(..., dmap_kwargs={...}, gplm_kwargs={...}, ddpm_kwargs={...})
+    and you can also tweak the “headline” DDPM knobs directly in __init__ (below).
     """
 
     def __init__(
         self,
         R_iX: np.ndarray,
         *,
+        # main knobs
         d: int = 32,
+        beta: float = 1.0,
+
+        # allow per-module override (if None -> uses global beta)
+        dmap_beta: Optional[float] = None,
+        gplm_beta: Optional[float] = None,
+
+        # DMAP headline knobs (everything else via dmap_kwargs)
+        dmap_alpha: float = 0.0,
+        dmap_t: float = 1.0,
+        dmap_k: Optional[int] = None,
+
+        # GPLM headline knobs (everything else via gplm_kwargs)
+        gplm_m: int = 1024,
+        gplm_pred_k: Optional[int] = None,
+
+        # DDPM headline knobs (the ones worth surfacing)
+        ddpm_T: int = 200,
+        ddpm_hidden_dim: int = 128,
+        ddpm_t_embed_dim: int = 64,
+        ddpm_learning_rate: float = 3e-4,
+        ddpm_n_iter: int = 200_000,
+        ddpm_ema_decay: float = 0.999,
+        ddpm_beta_max: float = 0.02,
+        ddpm_batch_size: int = 256,
+        ddpm_verbose_every: int = 0,
+        ddpm_eps: float = 1e-5,
+
+        # runtime
+        ddpm_device: str = "auto",
+        key: Optional[jax.Array] = None,
+
+        # ann
+        ann_backend: ANNBackend = "auto",
+        ann_params: Optional[Dict[str, Any]] = None,
+        n_jobs: int = -1,
+
+        # “escape hatches”
         dmap_kwargs: Optional[Dict[str, Any]] = None,
         gplm_kwargs: Optional[Dict[str, Any]] = None,
         ddpm_kwargs: Optional[Dict[str, Any]] = None,
-        ddpm_device: str = "auto",
-        key: Optional[jax.Array] = None ):
-        self.config = DIMAConfig(d=int(d), ddpm_device=str(ddpm_device))
-        self.training_time = 0.0
 
-        t0 = time.time()
+        # unicode aliases (β)
+        **kwargs: Any,
+    ):
+        # ---- unicode aliases ----
+        if "β" in kwargs:
+            beta = float(kwargs.pop("β"))
+        if "β_dmap" in kwargs:
+            dmap_beta = float(kwargs.pop("β_dmap"))
+        if "β_gplm" in kwargs:
+            gplm_beta = float(kwargs.pop("β_gplm"))
+        if kwargs:
+            raise TypeError(f"Unexpected kwargs: {sorted(kwargs.keys())}")
+
+        self.config = DIMAConfig(d=int(d), beta=float(beta), ddpm_device=str(ddpm_device))
+
+        # devices + rng
+        self.ddpm_device = _select_device(ddpm_device)
+        self.cpu_device = _select_device("cpu")
+        self.rng = random.PRNGKey(0) if key is None else key
 
-        # data
+        # training data
         self.R_iX = np.asarray(R_iX)
         if self.R_iX.ndim != 2:
             raise ValueError("R_iX must be 2D (N,D).")
         self.N, self.D = self.R_iX.shape
         self.d = int(d)
 
-        dmap_kwargs = {} if dmap_kwargs is None else dict(dmap_kwargs)
-        gplm_kwargs = {} if gplm_kwargs is None else dict(gplm_kwargs)
-        ddpm_kwargs = {} if ddpm_kwargs is None else dict(ddpm_kwargs)
+        # betas
+        self.beta = float(beta)
+        self.β = self.beta
 
-        # choose device for DDPM
-        self.ddpm_device = _select_device(self.config.ddpm_device)
-        self.cpu_device = _select_device("cpu")
+        self.dmap_beta = float(self.beta if dmap_beta is None else dmap_beta)
+        self.gplm_beta = float(self.beta if gplm_beta is None else gplm_beta)
 
-        if key is None:
-            key = random.PRNGKey(0)
-        self.rng = key
+        t0 = time.time()
 
         # -------------------------
         # 1) Train DMAP (CPU)
         # -------------------------
-        self.enc = DMAP(self.R_iX, d=self.d, **dmap_kwargs)
-        R_ix = np.asarray(self.enc(self.R_iX))  # raw latents (N,d) float64
+        dmap_init = dict(
+            d=self.d,
+            beta=self.dmap_beta,
+            alpha=float(dmap_alpha),
+            t=float(dmap_t),
+            k=dmap_k,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+        )
+        if dmap_kwargs:
+            dmap_init.update(dict(dmap_kwargs))
+        # enforce headline knobs
+        dmap_init["d"] = self.d
+        dmap_init["beta"] = self.dmap_beta
+        dmap_init["alpha"] = float(dmap_alpha)
+        dmap_init["t"] = float(dmap_t)
+        dmap_init["k"] = dmap_k
+
+        self.enc = DMAP(self.R_iX, **dmap_init)
+
+        # raw DMAP coordinates for *all* training points (Nyström OOS on training set)
+        R_ix = np.asarray(self.enc(self.R_iX), dtype=np.float64)  # (N,d)
 
         # -------------------------
-        # 2) Fit latent normalization (ONCE)
+        # 2) Latent normalization for DDPM
         # -------------------------
-        self.lat_mean_np = R_ix.mean(axis=0, keepdims=True).astype(np.float64)
-        self.lat_std_np = (R_ix.std(axis=0, keepdims=True) + 1e-6).astype(np.float64)
-
-        Z_ix = (R_ix - self.lat_mean_np) / self.lat_std_np  # normalized latents (np)
+        self.lat_mean_np = R_ix.mean(axis=0)
+        self.lat_std_np = np.maximum(R_ix.std(axis=0), 1e-12)
 
-        # keep norms on ddpm device too
         self.lat_mean_j = jax.device_put(jnp.asarray(self.lat_mean_np, dtype=jnp.float32), self.ddpm_device)
         self.lat_std_j = jax.device_put(jnp.asarray(self.lat_std_np, dtype=jnp.float32), self.ddpm_device)
 
+        Z_ix = (R_ix - self.lat_mean_np) / self.lat_std_np  # (N,d)
+
         # -------------------------
         # 3) Train GPLM (CPU)
         # -------------------------
-        self.dec = GPLM(R_ix, self.R_iX, **gplm_kwargs)
+        gplm_init = dict(
+            beta=self.gplm_beta,
+            m=int(gplm_m),
+            pred_k=gplm_pred_k,
+            ann_backend=ann_backend,
+            ann_params=ann_params,
+            n_jobs=n_jobs,
+        )
+        if gplm_kwargs:
+            gplm_init.update(dict(gplm_kwargs))
+        # enforce headline knobs
+        gplm_init["beta"] = self.gplm_beta
+        gplm_init["m"] = int(gplm_m)
+        gplm_init["pred_k"] = gplm_pred_k
+
+        self.dec = GPLM(R_ix.astype(np.float32, copy=False), self.R_iX, **gplm_init)
 
         # -------------------------
         # 4) Train DDPM on normalized latents (DDPM device)
         # -------------------------
         Z_ix_j = jax.device_put(jnp.asarray(Z_ix, dtype=jnp.float32), self.ddpm_device)
 
-        # defaults (can be overridden by ddpm_kwargs)
         ddpm_init = dict(
-            T=200,
-            hidden_dim=128,
-            t_embed_dim=64,
-            learning_rate=3e-4,
-            n_iter=200_000,
-            ema_decay=0.999,
-            beta_max=0.02,
-            batch_size=256,
+            T=int(ddpm_T),
+            hidden_dim=int(ddpm_hidden_dim),
+            t_embed_dim=int(ddpm_t_embed_dim),
+            learning_rate=float(ddpm_learning_rate),
+            n_iter=int(ddpm_n_iter),
+            ema_decay=float(ddpm_ema_decay),
+            beta_max=float(ddpm_beta_max),
+            batch_size=int(ddpm_batch_size),
             key=self.rng,
-            verbose_every=0,
-            eps=1e-5,
+            verbose_every=int(ddpm_verbose_every),
+            eps=float(ddpm_eps),
         )
-        ddpm_init.update(ddpm_kwargs)
+        if ddpm_kwargs:
+            ddpm_init.update(dict(ddpm_kwargs))
+        # enforce headline knobs
+        ddpm_init["T"] = int(ddpm_T)
+        ddpm_init["hidden_dim"] = int(ddpm_hidden_dim)
+        ddpm_init["t_embed_dim"] = int(ddpm_t_embed_dim)
+        ddpm_init["learning_rate"] = float(ddpm_learning_rate)
+        ddpm_init["n_iter"] = int(ddpm_n_iter)
+        ddpm_init["ema_decay"] = float(ddpm_ema_decay)
+        ddpm_init["beta_max"] = float(ddpm_beta_max)
+        ddpm_init["batch_size"] = int(ddpm_batch_size)
+        ddpm_init["verbose_every"] = int(ddpm_verbose_every)
+        ddpm_init["eps"] = float(ddpm_eps)
 
         with jax.default_device(self.ddpm_device):
             self.dm = DDPM(Z_ix_j, **ddpm_init)
@@ -315,47 +379,69 @@ class DIMA:
         self.training_time = time.time() - t0
 
     # -------------------------
-    # Encode / Decode
+    # Latent conversions
+    # -------------------------
+    def normalize(self, R_ax: Union[np.ndarray, jnp.ndarray]) -> jnp.ndarray:
+        """raw latents (R) -> normalized latents (Z) on ddpm_device."""
+        R = np.asarray(R_ax, dtype=np.float64)
+        if R.ndim == 1:
+            R = R[None, :]
+        Z = (R - self.lat_mean_np) / self.lat_std_np
+        Zj = jnp.asarray(Z, dtype=jnp.float32)
+        return jax.device_put(Zj, self.ddpm_device)
+
+    def unnormalize(self, Z_ax: Union[np.ndarray, jnp.ndarray]) -> np.ndarray:
+        """normalized latents (Z) -> raw latents (R) on CPU (np)."""
+        if isinstance(Z_ax, jax.Array):
+            Z_np = np.asarray(jax.device_get(Z_ax))
+        else:
+            Z_np = np.asarray(Z_ax)
+        if Z_np.ndim == 1:
+            Z_np = Z_np[None, :]
+        R = Z_np * self.lat_std_np + self.lat_mean_np
+        return np.asarray(R)
+
+    # -------------------------
+    # Encode / Decode (public: raw latents)
     # -------------------------
-    def encode(self, R_aX: Union[np.ndarray, jnp.ndarray], *, to_device: bool = True) -> jnp.ndarray:
+    def encode(self, R_aX: Union[np.ndarray, jnp.ndarray], *, normalize: bool = False) -> Union[np.ndarray, jnp.ndarray]:
         """
-        ambient -> normalized latent (DDPM space)
-        Returns jnp on ddpm_device by default.
+        ambient -> raw DMAP latents (np) by default.
+        If normalize=True, returns normalized latents (jnp) on ddpm_device.
         """
         X = np.asarray(R_aX)
-        Z_raw = np.asarray(self.enc(X))  # CPU
-        Z_norm = (Z_raw - self.lat_mean_np) / self.lat_std_np
-        Zj = jnp.asarray(Z_norm, dtype=jnp.float32)
-        return jax.device_put(Zj, self.ddpm_device) if to_device else Zj
+        R_raw = np.asarray(self.enc(X))  # CPU, (a,d)
+        if not normalize:
+            return R_raw
+        return self.normalize(R_raw)
 
     def decode(
         self,
-        Z_ax: Union[np.ndarray, jnp.ndarray],
+        R_ax: Union[np.ndarray, jnp.ndarray],
         *,
-        refine: bool = True,
+        refine: bool = False,
         t_start: int = 10,
         add_noise: bool = True,
         key: Optional[jax.Array] = None,
         batch_size: Optional[int] = None,
     ) -> np.ndarray:
         """
-        normalized latent -> (optional DDPM refine) -> raw latent -> ambient
-        Returns numpy array (CPU).
+        raw latent -> (optional DDPM refine in normalized coords) -> raw latent -> ambient.
+        Returns ambient np.ndarray on CPU.
         """
-        Z = jnp.asarray(Z_ax, dtype=jnp.float32)
-        if Z.ndim == 1:
-            Z = Z[None, :]
-        Z = jax.device_put(Z, self.ddpm_device)
+        R_raw = np.asarray(R_ax, dtype=np.float64)
+        single = (R_raw.ndim == 1)
+        if single:
+            R_raw = R_raw[None, :]
 
         if refine:
+            Z = self.normalize(R_raw)  # on device
             Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
+            R_raw = self.unnormalize(Z)  # back to CPU raw
 
-        # unnormalize to raw latent (still on device)
-        R_ax = (Z * self.lat_std_j) + self.lat_mean_j
-        R_ax_np = np.asarray(jax.device_get(R_ax))  # to CPU
-
-        X_hat = self.dec(R_ax_np, batch_size=batch_size)  # CPU decode
-        return np.asarray(X_hat)
+        X_hat = self.dec(R_raw.astype(np.float32, copy=False), batch_size=batch_size)
+        X_hat = np.asarray(X_hat)
+        return X_hat[0] if single else X_hat
 
     def reconstruct(
         self,
@@ -367,248 +453,141 @@ class DIMA:
         key: Optional[jax.Array] = None,
         batch_size: Optional[int] = None,
     ) -> np.ndarray:
-        """Convenience: decode(encode(X))."""
-        Z = self.encode(R_aX, to_device=True)
-        return self.decode(Z, refine=refine, t_start=t_start, add_noise=add_noise, key=key, batch_size=batch_size)
+        """decode(encode(X))."""
+        R_raw = self.encode(R_aX, normalize=False)
+        return self.decode(R_raw, refine=refine, t_start=t_start, add_noise=add_noise, key=key, batch_size=batch_size)
 
-    def sample(self, n: int, *, decode: bool = True, batch_size: Optional[int] = None) -> Union[jnp.ndarray, np.ndarray]:
+    def sample(
+        self,
+        n: int,
+        *,
+        decode: bool = True,
+        batch_size: Optional[int] = None,
+    ) -> Union[np.ndarray, np.ndarray]:
         """
         Unconditional samples from latent DDPM.
-        If decode=True, returns ambient samples as np.ndarray.
-        Otherwise returns normalized latents as jnp.ndarray.
+        If decode=True: returns ambient samples (np) on CPU.
+        If decode=False: returns raw latents (np) on CPU.
         """
         with jax.default_device(self.ddpm_device):
             Z = self.dm.sample(int(n))
+        R = self.unnormalize(Z)  # raw (np)
         if not decode:
-            return Z
-        return self.decode(Z, refine=False, batch_size=batch_size)
-
-    # -------------------------
-    # helpers: normalize <-> raw
-    # -------------------------
-    def _to_np(self, A: Union[np.ndarray, jnp.ndarray]) -> np.ndarray:
-        if isinstance(A, jax.Array):
-            return np.asarray(jax.device_get(A))
-        return np.asarray(A)
-
-    def _norm_to_raw_np(self, Z_ax: np.ndarray) -> np.ndarray:
-        # Z_ax: normalized (a,d) -> raw r_ax (a,d)
-        return Z_ax * self.lat_std_np + self.lat_mean_np
-
-    def _raw_to_norm_np(self, R_ax: np.ndarray) -> np.ndarray:
-        # raw r_ax (a,d) -> normalized (a,d)
-        return (R_ax - self.lat_mean_np) / self.lat_std_np
-
-    def _metric_raw_to_norm(self, g_raw_add: np.ndarray) -> np.ndarray:
-        # g_norm_xy = sigma_x sigma_y g_raw_xy
-        sig = self.lat_std_np.reshape(-1)  # (d,)
-        return g_raw_add * (sig[None, :, None] * sig[None, None, :])
-
-    def _jac_raw_to_norm(self, J_raw_aDd: np.ndarray) -> np.ndarray:
-        # J_norm_{D,x} = dR/dz_x = dR/dr_y * dr_y/dz_x = J_raw_{D,x} * sigma_x
-        sig = self.lat_std_np.reshape(-1)  # (d,)
-        return J_raw_aDd * sig[None, None, :]
+            return R
+        return self.dec(R.astype(np.float32, copy=False), batch_size=batch_size)
 
     # -------------------------
-    # new: flow wrapper
+    # Geodesic-ish flow wrapper (delegates to GPLM.flow)
     # -------------------------
     def flow(
         self,
-        Z_ax: Union[np.ndarray, jnp.ndarray],
-        v0_ax: Union[np.ndarray, jnp.ndarray],
+        R_ax: Union[np.ndarray, jnp.ndarray],
+        v_ax: Union[np.ndarray, jnp.ndarray],
         *,
         dt: float = 0.05,
         reg: float = 1e-8,
         keep_speed: bool = True,
-        # optional: decode + metric/J outputs at the new point
-        decode: bool = True,
-        jacobian: bool = False,
-        metric: bool = False,
-        metric_in: str = "norm",  # "norm" or "raw"
-        # optional DDPM projection (stochastic): OFF by default for deterministic geodesics
+        # optional DDPM projection step (in normalized coords)
         refine: bool = False,
-        ddpm_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        # decode return
+        decode: bool = False,
+        batch_size: Optional[int] = None,
+    ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray]]:
         """
-        One geodesic step in latent space using GPLM.flow (analytic Christoffels + velocity-Verlet),
-        where Z_ax and v0_ax are in *normalized* latent coordinates.
+        One step of latent flow in *raw* coordinates.
 
-        Returns (if decode=True):
-            Xhat_aX, Q_ax, v_ax                  (and optionally J/g depending on flags)
-        Returns (if decode=False):
-            Q_ax, v_ax
-        """
-        ddpm_kwargs = {} if ddpm_kwargs is None else dict(ddpm_kwargs)
+        Requires: your GPLM class implements:
+            R_next, v_next = gplm.flow(R, v, dt=..., reg=..., keep_speed=...)
 
-        Z_np = self._to_np(Z_ax)
-        v_np = self._to_np(v0_ax)
+        If refine=True, we project the *position* through DDPM in normalized coords after the step.
+        (Velocity after projection is left unchanged—projection isn’t a deterministic diffeo.)
 
-        single = (Z_np.ndim == 1)
+        Returns:
+          if decode=False:
+            (R_next, v_next)        both np arrays
+          if decode=True:
+            (X_next, R_next, v_next)
+        """
+        R = np.asarray(R_ax, dtype=np.float64)
+        v = np.asarray(v_ax, dtype=np.float64)
+        single = (R.ndim == 1)
         if single:
-            Z_np = Z_np[None, :]
-        if v_np.ndim == 1:
-            v_np = v_np[None, :]
-
-        if Z_np.shape != v_np.shape:
-            raise ValueError(f"Z_ax and v0_ax must have same shape; got {Z_np.shape} vs {v_np.shape}")
+            R = R[None, :]
+            v = v[None, :]
 
-        # normalized -> raw for GPLM.flow
-        R_raw = self._norm_to_raw_np(Z_np)                         # (a,d)
-        v_raw = v_np * self.lat_std_np                             # (a,d)  (dr/dt = sigma * dz/dt)
-
-        # one step geodesic update in raw coordinates
-        Q_raw, v_raw_new = self.dec.flow(R_raw, v_raw, dt=dt, reg=reg, keep_speed=keep_speed)
+        if not hasattr(self.dec, "flow"):
+            raise AttributeError(
+                "Decoder does not have .flow(). Make sure you updated GPLM to include flow()."
+            )
 
-        # raw -> normalized
-        Q_norm = self._raw_to_norm_np(Q_raw)
-        v_norm = v_raw_new / self.lat_std_np
+        Rn, vn = self.dec.flow(R, v, dt=float(dt), reg=float(reg), keep_speed=bool(keep_speed))
 
-        # optional: DDPM projection in normalized coords
         if refine:
-            # keep this minimal; adapt to your DDPM API if needed
-            Zj = jax.device_put(jnp.asarray(Q_norm, dtype=jnp.float32), self.ddpm_device)
-            if hasattr(self.dm, "refine_latents"):
-                Zj2 = self.dm.refine_latents(Zj, **ddpm_kwargs)
-            else:
-                Zj2 = self.dm(Zj, **ddpm_kwargs)
-            Q_norm = np.asarray(jax.device_get(Zj2))
-            # velocity after projection is ambiguous; keep v as-is in normalized chart
-
-        # return latents on ddpm device (consistent with encode())
-        Q_out = jax.device_put(jnp.asarray(Q_norm, dtype=jnp.float32), self.ddpm_device)
-        v_out = jax.device_put(jnp.asarray(v_norm, dtype=jnp.float32), self.ddpm_device)
+            Z = self.normalize(Rn)  # device
+            Z = self.dm.refine_latents(Z, t_start=int(t_start), key=key, add_noise=bool(add_noise))
+            Rn = self.unnormalize(Z)
 
         if not decode:
             if single:
-                return Q_out[0], v_out[0]
-            return Q_out, v_out
-
-        # decode at Q_norm: need raw coords for GPLM.__call__
-        Q_raw2 = self._norm_to_raw_np(np.asarray(Q_norm))
+                return np.asarray(Rn[0]), np.asarray(vn[0])
+            return np.asarray(Rn), np.asarray(vn)
 
-        if (not jacobian) and (not metric):
-            Xhat = self.dec(Q_raw2)
-            if single:
-                return Xhat[0], Q_out[0], v_out[0]
-            return Xhat, Q_out, v_out
-
-        # request J and/or g from GPLM in raw coords
-        out = self.dec(Q_raw2, jacobian=jacobian, metric=metric)
-
-        # normalize J/g if requested
-        if jacobian and metric:
-            Xhat, J_raw, g_raw = out
-            if metric_in == "norm":
-                J_use = self._jac_raw_to_norm(J_raw)
-                g_use = self._metric_raw_to_norm(g_raw)
-            else:
-                J_use, g_use = J_raw, g_raw
-            if single:
-                return Xhat[0], J_use[0], g_use[0], Q_out[0], v_out[0]
-            return Xhat, J_use, g_use, Q_out, v_out
-
-        if jacobian and (not metric):
-            Xhat, J_raw = out
-            J_use = self._jac_raw_to_norm(J_raw) if metric_in == "norm" else J_raw
-            if single:
-                return Xhat[0], J_use[0], Q_out[0], v_out[0]
-            return Xhat, J_use, Q_out, v_out
-
-        # metric only
-        Xhat, g_raw = out
-        g_use = self._metric_raw_to_norm(g_raw) if metric_in == "norm" else g_raw
+        Xn = self.dec(Rn.astype(np.float32, copy=False), batch_size=batch_size)
+        Xn = np.asarray(Xn)
         if single:
-            return Xhat[0], g_use[0], Q_out[0], v_out[0]
-        return Xhat, g_use, Q_out, v_out
+            return Xn[0], np.asarray(Rn[0]), np.asarray(vn[0])
+        return Xn, np.asarray(Rn), np.asarray(vn)
 
     # -------------------------
-    # modify __call__ to support metric/jacobian and velocity
+    # Convenience __call__
     # -------------------------
     def __call__(
         self,
-        R_any: Union[np.ndarray, jnp.ndarray],
+        A: Union[np.ndarray, jnp.ndarray],
         *,
-        refine: bool = True,
-        jacobian: bool = False,
-        metric: bool = False,
-        metric_in: str = "norm",     # for latent->ambient: return g in "norm" or "raw"
-        velocity: Optional[Union[np.ndarray, jnp.ndarray]] = None,
-        dt: float = 0.05,
-        reg: float = 1e-8,
-        keep_speed: bool = True,
-        ddpm_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+        refine: bool = False,
+        t_start: int = 10,
+        add_noise: bool = True,
+        key: Optional[jax.Array] = None,
+        batch_size: Optional[int] = None,
+        normalize_latent: bool = False,
+    ) -> Union[np.ndarray, jnp.ndarray]:
         """
-        Polymorphic call:
-          - if trailing dim == D: ambient -> normalized latent (jnp on ddpm device)
-          - if trailing dim == d: normalized latent -> ambient (np), optionally returning J/g
-          - if trailing dim == d and velocity is provided: one flow step then decode
+        Dispatch by last dimension:
 
-        Returns:
-          ambient->latent: Z (jnp)
-          latent->ambient:
-            - if no velocity: Xhat (or (Xhat,J)/(Xhat,g)/(Xhat,J,g))
-            - if velocity:    Xhat, Q, v  (and optionally J/g depending on flags)
+          - if A is (a,D): encode -> raw latents (np) by default
+          - if A is (a,d): decode -> ambient (np)
+
+        Options:
+          - normalize_latent=True only affects encoding (returns Z on device)
+          - refine/t_start/add_noise/key only affect decoding
         """
-        ddpm_kwargs = {} if ddpm_kwargs is None else dict(ddpm_kwargs)
-
-        X = self._to_np(R_any)
-        last = X.shape[-1]
-
-        if last == self.D:
-            # keep whatever your encode() does; here’s the minimal consistent behavior:
-            Z = np.asarray(self.enc(X))                      # raw latents (np)
-            Z = (Z - self.lat_mean_np) / self.lat_std_np     # normalized
-            return jax.device_put(jnp.asarray(Z, dtype=jnp.float32), self.ddpm_device)
-
-        if last == self.d:
-            # If velocity: do one geodesic step (normalized chart), then decode
-            if velocity is not None:
-                return self.flow(
-                    X, velocity,
-                    dt=dt, reg=reg, keep_speed=keep_speed,
-                    decode=True, jacobian=jacobian, metric=metric, metric_in=metric_in,
-                    refine=refine, ddpm_kwargs=ddpm_kwargs,
-                )
-
-            # Otherwise: decode directly (optionally refine in normalized coords)
-            Z_norm = X
-            if refine:
-                Zj = jax.device_put(jnp.asarray(Z_norm, dtype=jnp.float32), self.ddpm_device)
-                if hasattr(self.dm, "refine_latents"):
-                    Zj2 = self.dm.refine_latents(Zj, **ddpm_kwargs)
-                else:
-                    Zj2 = self.dm(Zj, **ddpm_kwargs)
-                Z_norm = np.asarray(jax.device_get(Zj2))
-
-            R_raw = self._norm_to_raw_np(np.asarray(Z_norm))
-            out = self.dec(R_raw, jacobian=jacobian, metric=metric)
-
-            if (not jacobian) and (not metric):
-                return out
-
-            if jacobian and metric:
-                Xhat, J_raw, g_raw = out
-                if metric_in == "norm":
-                    return Xhat, self._jac_raw_to_norm(J_raw), self._metric_raw_to_norm(g_raw)
-                return out
-
-            if jacobian and (not metric):
-                Xhat, J_raw = out
-                return (Xhat, self._jac_raw_to_norm(J_raw)) if metric_in == "norm" else out
-
-            # metric only
-            Xhat, g_raw = out
-            return (Xhat, self._metric_raw_to_norm(g_raw)) if metric_in == "norm" else out
-
-        raise ValueError(f"Trailing dim must be D={self.D} (ambient) or d={self.d} (latent), got {last}.")
+        A_np = np.asarray(A)
+        if A_np.ndim == 1:
+            A_np = A_np[None, :]
+
+        if A_np.shape[1] == self.D:
+            return self.encode(A_np, normalize=bool(normalize_latent))
+
+        if A_np.shape[1] == self.d:
+            return self.decode(
+                A_np,
+                refine=bool(refine),
+                t_start=int(t_start),
+                add_noise=bool(add_noise),
+                key=key,
+                batch_size=batch_size,
+            )
 
+        raise ValueError(f"Input has last-dim {A_np.shape[1]}, expected D={self.D} or d={self.d}.")
 
     # -------------------------
     # Save / Load
     # -------------------------
     def _pack_encoder(self) -> Dict[str, Any]:
-        # minimal state for Nyström OOS
         return dict(
             R_iX=np.asarray(self.enc.R_iX),
             qalpha_i=np.asarray(getattr(self.enc, "qalpha_i", getattr(self.enc, "qα_i"))),
@@ -622,7 +601,7 @@ class DIMA:
 
     def _pack_decoder(self) -> Dict[str, Any]:
         return dict(
-            Z_mx_w=np.asarray(getattr(self.dec, "Z_mx_w", getattr(self.dec, "Z_mx_w", None))),
+            Z_mx_w=np.asarray(getattr(self.dec, "Z_mx_w", None)),
             M_mX=np.asarray(self.dec.M_mX),
             mean_X=np.asarray(getattr(self.dec, "mean_X", np.zeros((self.D,), dtype=np.float64))),
             lat_mean_x=np.asarray(getattr(self.dec, "lat_mean_x", np.zeros((self.d,), dtype=np.float64))),
@@ -634,9 +613,6 @@ class DIMA:
         )
 
     def state_dict(self) -> Dict[str, Any]:
-        """
-        Serializable state (msgpack via flax.serialization).
-        """
         dd = dict(
             T=int(self.dm.T),
             D=int(self.dm.D),
@@ -654,7 +630,7 @@ class DIMA:
                 N=int(self.N),
                 D=int(self.D),
                 d=int(self.d),
-                training_time=float(self.training_time),
+                training_time=float(getattr(self, "training_time", 0.0)),
             ),
             config=asdict(self.config),
             latent_norm=dict(
@@ -702,9 +678,11 @@ class DIMA:
         obj.d = int(state["meta"]["d"])
 
         # RNG
-        if key is None:
-            key = random.PRNGKey(0)
-        obj.rng = key
+        obj.rng = random.PRNGKey(0) if key is None else key
+
+        # beta (for display / convenience)
+        obj.beta = float(obj.config.beta)
+        obj.β = obj.beta
 
         # latent norm
         obj.lat_mean_np = np.asarray(state["latent_norm"]["mean"], dtype=np.float64)
@@ -713,9 +691,9 @@ class DIMA:
         obj.lat_mean_j = jax.device_put(jnp.asarray(obj.lat_mean_np, dtype=jnp.float32), obj.ddpm_device)
         obj.lat_std_j = jax.device_put(jnp.asarray(obj.lat_std_np, dtype=jnp.float32), obj.ddpm_device)
 
-        # encoder/decoder frozen
+        # frozen encoder + rehydrated decoder-as-GPLM (so flow works if GPLM.flow exists)
         obj.enc = FrozenDMAP(state["encoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
-        obj.dec = FrozenGPLM(state["decoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
+        obj.dec = _restore_gplm_as_object(state["decoder"], ann_backend=ann_backend, ann_params=ann_params, n_jobs=n_jobs)
 
         # rebuild DDPM skeleton with dummy data, then load params
         dd = state["ddpm"]
@@ -749,36 +727,26 @@ class DIMA:
         return obj
 
     # -------------------------
-    # Hugging Face helpers
+    # (Optional) HF helpers
     # -------------------------
-    def upload_to_huggingface(
-        self,
-        repo_id: str,
-        *,
-        hf_token: Optional[str] = None,
-        weights_file: str = "dima.msgpack",
-        config_file: str = "config.json",
-    ) -> None:
+    def save_hf(self, repo_id: str, weights_file: str = "dima.msgpack", config_file: str = "config.json") -> None:
         if not _HAS_HF:
-            raise ImportError("huggingface_hub not installed. pip install dima[hf] (or huggingface_hub).")
-        if hf_token is not None:
-            HfFolder.save_token(hf_token)
-        api = HfApi()
-        _ = api.whoami(token=hf_token)  # sanity check token (no-op if public)
-
+            raise RuntimeError("huggingface_hub not installed.")
         self.save_local(weights_file=weights_file, config_file=config_file)
-        for fname in (weights_file, config_file):
-            upload_file(path_or_fileobj=fname, path_in_repo=fname, repo_id=repo_id, token=hf_token)
-        return None
+        token = HfFolder.get_token()
+        if token is None:
+            raise RuntimeError("No HF token found. Run `huggingface-cli login`.")
+        api = HfApi()
+        api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, token=token)
+        upload_file(path_or_fileobj=weights_file, path_in_repo=weights_file, repo_id=repo_id, token=token)
+        upload_file(path_or_fileobj=config_file, path_in_repo=config_file, repo_id=repo_id, token=token)
 
     @classmethod
-    def load_from_huggingface(
+    def load_hf(
         cls,
         repo_id: str,
         *,
-        hf_token: Optional[str] = None,
         weights_file: str = "dima.msgpack",
-        config_file: str = "config.json",
         ddpm_device: str = "auto",
         ann_backend: ANNBackend = "auto",
         ann_params: Optional[Dict[str, Any]] = None,
@@ -786,13 +754,10 @@ class DIMA:
         key: Optional[jax.Array] = None,
     ) -> "DIMA":
         if not _HAS_HF:
-            raise ImportError("huggingface_hub not installed. pip install dima[hf] (or huggingface_hub).")
-
-        wpath = hf_hub_download(repo_id, weights_file, token=hf_token)
-        _ = hf_hub_download(repo_id, config_file, token=hf_token)  # for humans / future-proofing
-
+            raise RuntimeError("huggingface_hub not installed.")
+        path = hf_hub_download(repo_id=repo_id, filename=weights_file)
         return cls.load_local(
-            wpath,
+            path,
             ddpm_device=ddpm_device,
             ann_backend=ann_backend,
             ann_params=ann_params,