spatialrna 0.2.0__tar.gz → 0.3.0.dev1__tar.gz

spatialrna-0.3.0.dev1/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: spatialrna
+Version: 0.3.0.dev1
+Summary: Spatial RNA graph data processing and analysis for iST (Xenium / Atera): graphs, GNN training / embedding, clustering, spatial maps.
+Author-email: Anne Pan <anne.pan@svi.edu.au>
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: scipy
+Requires-Dist: zarr>=3
+Requires-Dist: numcodecs
+Requires-Dist: tqdm
+Requires-Dist: scikit-learn
+Requires-Dist: matplotlib
+Provides-Extra: gnn
+Requires-Dist: torch; extra == "gnn"
+Requires-Dist: torch_geometric; extra == "gnn"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+
+description = "A package for spatial RNA graph data processing and analysis."
+readme = "README.md"
+version = "0.3.0.dev1"
+authors = [{ name = "Anne Pan", email = "anne.pan@svi.edu.au" },{ name = "Ruqian Lyu" }]
+license = { text = "MIT" }
+requires-python = ">=3.12"
+
+[project.urls]
+Repository = "https://gitlab.svi.edu.au/biocellgen-public/spatialrna_pypi_public"
+Documentation = "https://ruqianl.github.io/spatialrna_docs/"

spatialrna-0.3.0.dev1/README.md

@@ -0,0 +1,10 @@
+description = "A package for spatial RNA graph data processing and analysis."
+readme = "README.md"
+version = "0.3.0.dev1"
+authors = [{ name = "Anne Pan", email = "anne.pan@svi.edu.au" },{ name = "Ruqian Lyu" }]
+license = { text = "MIT" }
+requires-python = ">=3.12"
+
+[project.urls]
+Repository = "https://gitlab.svi.edu.au/biocellgen-public/spatialrna_pypi_public"
+Documentation = "https://ruqianl.github.io/spatialrna_docs/"

spatialrna-0.3.0.dev1/pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spatialrna"
+version = "0.3.0.dev1"
+description = "Spatial RNA graph data processing and analysis for iST (Xenium / Atera): graphs, GNN training / embedding, clustering, spatial maps."
+readme = "README.md"
+requires-python = ">=3.12"
+authors = [{ name = "Anne Pan", email = "anne.pan@svi.edu.au" }]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+# Core: readers, graph builders, clustering, maps. No torch needed.
+dependencies = [
+    "numpy",
+    "pandas",
+    "scipy",
+    "zarr>=3",
+    "numcodecs",
+    "tqdm",
+    "scikit-learn",
+    "matplotlib",
+]
+
+[project.optional-dependencies]
+# SpatialRNA dataset, datasets, train_val. Install torch and the PyG wheels matching
+# your CUDA first -- see README.
+gnn = ["torch", "torch_geometric"]
+dev = ["pytest"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["spatialrna*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

spatialrna-0.3.0.dev1/src/spatialrna/__init__.py

@@ -0,0 +1,60 @@
+"""spatialrna: spatial RNA graph data processing and analysis for single-molecule
+spatial transcriptomics (Xenium / Atera): transcripts -> tile graphs / subgraphs ->
+any GNN (train / validate / embed) -> clusters -> spatial maps.
+
+    import spatialrna as sr
+    df = sr.readers.load_transcripts("/path/to/xenium_or_atera_output")
+    from spatialrna import SpatialRNA                 # tile graphs / subgraphs; needs torch + PyG
+
+The ISTIST pipeline builds on this with its own model (CrossSampleGAT), transcript
+annotation, statistics and report.
+
+Submodules (imported on first use, so the light ones need neither torch nor a GPU):
+
+    readers         Xenium / Atera transcripts.zarr.zip -> transcripts
+    graph_build     radius graphs (KD-tree) and subgraph sampling
+    graph_zarr_io   SpatialRNA tile dataset (strips or grid tiles) + zarr I/O [torch]
+    datasets        ZarrGraphDataset over many tile / subgraph stores        [torch]
+    train_val       model-agnostic link-prediction training, validation, tile embedding [torch]
+    cluster         Gaussian mixture (GPU or numpy) and k-means on embeddings
+    viz             pixel maps of labels / niches / cell types
+
+"""
+from __future__ import annotations
+
+import importlib
+
+__version__ = "0.3.0.dev0"
+
+_SUBMODULES = {"readers", "graph_build", "graph_zarr_io", "datasets", "train_val", "cluster", "viz"}
+
+# name -> submodule, for `from spatialrna import SpatialRNA` etc. (resolved lazily, so
+# `import spatialrna` works without torch; the graph_zarr_io names need it).
+_EXPORTS = {
+    **{n: "graph_zarr_io" for n in (
+        "SpatialRNA", "save_data_zarr", "load_data_zarr", "zarr_group_exists", "save_raw_transcripts_zarr",
+        "load_raw_transcripts_zarr", "store_exists", "tile_store_path", "sample_store_path",
+        "ensure_sample_store", "subgraph_store_path", "read_tile_graph", "write_tile_graph",
+        "read_subgraph", "write_subgraph")},
+    "load_transcripts": "readers",
+    "ZarrGraphDataset": "datasets",
+    "plan_tiles": "graph_build",
+}
+
+__all__ = sorted(_SUBMODULES | set(_EXPORTS)) + ["__version__"]
+
+
+def __getattr__(name):
+    if name in _SUBMODULES:
+        mod = importlib.import_module(f".{name}", __name__)
+        globals()[name] = mod
+        return mod
+    if name in _EXPORTS:
+        value = getattr(importlib.import_module(f".{_EXPORTS[name]}", __name__), name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module 'spatialrna' has no attribute {name!r}")
+
+
+def __dir__():
+    return __all__

spatialrna-0.3.0.dev1/src/spatialrna/cluster.py

@@ -0,0 +1,473 @@
+"""Clustering transcript embeddings into niches (spatialrna 0.2's run_kmeans, extended).
+
+    from spatialrna import cluster
+    labels, model = cluster.gmm(X, 12)          # Gaussian mixture, GPU if available
+    labels, model = cluster.kmeans(X, 12)       # scikit-learn k-means
+
+GaussianMixtureGPU -- a Gaussian mixture model (EM) that runs on the GPU with PyTorch.
+
+Replaces PyCave (torch < 2 only, unmaintained) for niche clustering. Same method:
+greedy k-means++ seeding and a few k-means iterations, then expectation-maximisation, with diagonal ("diag") or full
+covariance matrices. Data stay on the GPU in float32 and are processed in chunks;
+sufficient statistics (sums) are accumulated in float64 for numerical stability.
+
+The algorithm is written once against a few array operations (`_TorchOps` for the
+pipeline, `_NumpyOps` for testing it against scikit-learn without a GPU), so both
+run the identical code path.
+
+    gmm = GaussianMixtureGPU(n_components=12, covariance_type="diag", device="cuda")
+    gmm.fit(X)                 # X: (n, d) numpy float32
+    labels = gmm.predict(X)    # numpy int64
+    gmm.save("gmm.pt")
+"""
+from __future__ import annotations
+
+import math
+import time
+from typing import Optional
+
+import numpy as np
+
+
+# ---------------------------------------------------------------------------
+# Array backends
+# ---------------------------------------------------------------------------
+class _TorchOps:
+    def __init__(self, device: str):
+        import torch
+        self.t = torch
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device.startswith("cuda") and not torch.cuda.is_available():
+            raise RuntimeError("GMM device 'cuda' requested but no GPU is visible "
+                               "(torch.cuda.is_available() is False).")
+        self.device = torch.device(device)
+        self.f32, self.f64, self.i64 = torch.float32, torch.float64, torch.int64
+
+    def from_numpy(self, a, dtype=None):
+        return self.t.as_tensor(a, dtype=dtype or self.f32, device=self.device)
+
+    def to_numpy(self, a):
+        return a.detach().cpu().numpy()
+
+    def zeros(self, shape, dtype):
+        return self.t.zeros(shape, dtype=dtype, device=self.device)
+
+    def eye(self, n, dtype):
+        return self.t.eye(n, dtype=dtype, device=self.device)
+
+    def astype(self, a, dtype):
+        return a.to(dtype)
+
+    def sum(self, a, axis):
+        return a.sum(dim=axis)
+
+    def logsumexp(self, a, axis):
+        return self.t.logsumexp(a, dim=axis)
+
+    def argmax(self, a, axis):
+        return a.argmax(dim=axis)
+
+    def min(self, a, axis):
+        return a.min(dim=axis).values
+
+    def log(self, a):
+        return self.t.log(a)
+
+    def exp(self, a):
+        return self.t.exp(a)
+
+    def clip_min(self, a, v):
+        return a.clamp_min(v)
+
+    def cholesky(self, a):
+        return self.t.linalg.cholesky(a)
+
+    def inv(self, a):
+        return self.t.linalg.inv(a)
+
+    def diagonal(self, a):  # batched (k, d, d) -> (k, d)
+        return self.t.diagonal(a, dim1=-2, dim2=-1)
+
+    def transpose_last(self, a):
+        return a.transpose(-1, -2)
+
+    def bincount(self, a, n):
+        return self.t.bincount(a, minlength=n)
+
+    def item(self, a):
+        return float(a.item())
+
+    def sync(self):
+        if self.device.type == "cuda":
+            self.t.cuda.synchronize()
+
+
+class _NumpyOps:
+    """Same interface on NumPy -- used to test the algorithm without a GPU."""
+
+    def __init__(self, device: str = "cpu"):
+        self.device = "cpu"
+        self.f32, self.f64, self.i64 = np.float32, np.float64, np.int64
+
+    def from_numpy(self, a, dtype=None):
+        return np.asarray(a, dtype=dtype or self.f32)
+
+    def to_numpy(self, a):
+        return np.asarray(a)
+
+    def zeros(self, shape, dtype):
+        return np.zeros(shape, dtype=dtype)
+
+    def eye(self, n, dtype):
+        return np.eye(n, dtype=dtype)
+
+    def astype(self, a, dtype):
+        return a.astype(dtype)
+
+    def sum(self, a, axis):
+        return a.sum(axis=axis)
+
+    def logsumexp(self, a, axis):
+        m = a.max(axis=axis, keepdims=True)
+        return (m + np.log(np.exp(a - m).sum(axis=axis, keepdims=True))).squeeze(axis)
+
+    def argmax(self, a, axis):
+        return a.argmax(axis=axis)
+
+    def min(self, a, axis):
+        return a.min(axis=axis)
+
+    def log(self, a):
+        return np.log(a)
+
+    def exp(self, a):
+        return np.exp(a)
+
+    def clip_min(self, a, v):
+        return np.maximum(a, v)
+
+    def cholesky(self, a):
+        return np.linalg.cholesky(a)
+
+    def inv(self, a):
+        return np.linalg.inv(a)
+
+    def diagonal(self, a):
+        return np.diagonal(a, axis1=-2, axis2=-1)
+
+    def transpose_last(self, a):
+        return np.swapaxes(a, -1, -2)
+
+    def bincount(self, a, n):
+        return np.bincount(a, minlength=n)
+
+    def item(self, a):
+        return float(a)
+
+    def sync(self):
+        pass
+
+
+# ---------------------------------------------------------------------------
+# GMM
+# ---------------------------------------------------------------------------
+class GaussianMixtureGPU:
+    def __init__(self, n_components: int, covariance_type: str = "diag", max_iter: int = 500,
+                 tol: float = 1e-3, reg_covar: float = 1e-6, seed: int = 0, device: str = "auto",
+                 chunk_size: int = 1_000_000, init_sample: int = 200_000, verbose: bool = True,
+                 backend: str = "torch", n_init: int = 1, kmeans_iter: int = 10):
+        if covariance_type not in ("diag", "full"):
+            raise ValueError("covariance_type must be 'diag' or 'full'")
+        self.k = int(n_components)
+        self.covariance_type = covariance_type
+        self.max_iter, self.tol, self.reg_covar = max_iter, tol, reg_covar
+        self.seed, self.chunk_size, self.init_sample = seed, chunk_size, init_sample
+        self.verbose = verbose
+        self.n_init, self.kmeans_iter = max(1, int(n_init)), int(kmeans_iter)
+        self.ops = _TorchOps(device) if backend == "torch" else _NumpyOps()
+        self.converged_, self.n_iter_, self.lower_bound_ = False, 0, -math.inf
+
+    # -- internals ----------------------------------------------------------
+    def _log(self, msg):
+        if self.verbose:
+            print(msg, flush=True)
+
+    def _chunks(self, n):
+        for s in range(0, n, self.chunk_size):
+            yield s, min(n, s + self.chunk_size)
+
+    def _weighted_log_prob(self, Xc):
+        """(rows, k): log pi_k + log N(x | mu_k, Sigma_k), from the cached precisions."""
+        o, d = self.ops, Xc.shape[1]
+        const = -0.5 * d * math.log(2 * math.pi)
+        if self.covariance_type == "diag":
+            prec = self._prec                                  # (k, d) = 1 / var
+            mp = self.means_ * prec                            # (k, d)
+            maha = (Xc * Xc) @ o.transpose_last(prec) - 2.0 * (Xc @ o.transpose_last(mp)) \
+                + o.sum(self.means_ * mp, 1)[None, :]
+            maha = o.clip_min(maha, 0.0)
+        else:
+            cols = []
+            for j in range(self.k):                            # whitened residuals
+                y = (Xc - self.means_[j][None, :]) @ self._prec_chol[j]
+                cols.append(o.sum(y * y, 1)[:, None])
+            maha = cols[0] if self.k == 1 else _concat(o, cols)
+        return const - 0.5 * maha + self._half_logdet[None, :] + self._log_w[None, :]
+
+    def _refresh_precisions(self):
+        o = self.ops
+        if self.covariance_type == "diag":
+            self._prec = 1.0 / self.covariances_
+            self._half_logdet = 0.5 * o.sum(o.log(self._prec), 1)
+        else:
+            L = o.cholesky(self.covariances_)                  # cov = L L^T
+            self._prec_chol = o.transpose_last(o.inv(L))       # (x-mu) @ prec_chol = whitened
+            self._half_logdet = -o.sum(o.log(o.diagonal(L)), 1)
+        self._log_w = o.log(self.weights_)
+
+    def _m_step(self, nk, sx, sxx, n):
+        """Parameters from float64 sufficient statistics."""
+        o, d = self.ops, sx.shape[1]
+        nk = o.clip_min(nk, 10 * np.finfo(np.float64).eps)
+        means = sx / nk[:, None]
+        if self.covariance_type == "diag":
+            var = sxx / nk[:, None] - means * means
+            var = o.clip_min(var, 0.0) + self.reg_covar
+        else:
+            outer = means[:, :, None] * means[:, None, :]
+            var = sxx / nk[:, None, None] - outer
+            var = 0.5 * (var + o.transpose_last(var)) + self.reg_covar * o.eye(d, o.f64)[None]
+        self.weights_ = o.astype(nk / n, o.f32)
+        self.means_ = o.astype(means, o.f32)
+        self.covariances_ = o.astype(var, o.f32)
+        self._refresh_precisions()
+
+    def _stats(self, X, resp_fn):
+        """One pass over X: per-component counts, sums, (co)variance sums (float64)
+        under responsibilities resp_fn(chunk) -> (rows, k); plus total log-likelihood."""
+        o, d = self.ops, X.shape[1]
+        nk, sx = o.zeros((self.k,), o.f64), o.zeros((self.k, d), o.f64)
+        sxx = o.zeros((self.k, d), o.f64) if self.covariance_type == "diag" else o.zeros((self.k, d, d), o.f64)
+        ll = 0.0
+        for s, e in self._chunks(X.shape[0]):
+            Xc = X[s:e]
+            resp, chunk_ll = resp_fn(Xc)
+            ll += chunk_ll
+            rT = o.transpose_last(resp)                        # (k, rows)
+            nk = nk + o.astype(o.sum(resp, 0), o.f64)
+            sx = sx + o.astype(rT @ Xc, o.f64)
+            if self.covariance_type == "diag":
+                sxx = sxx + o.astype(rT @ (Xc * Xc), o.f64)
+            else:
+                for j in range(self.k):
+                    Xw = Xc * resp[:, j][:, None]
+                    sxx[j] = sxx[j] + o.astype(o.transpose_last(Xw) @ Xc, o.f64)
+        return nk, sx, sxx, ll
+
+    def _e_resp(self, Xc):
+        o = self.ops
+        wlp = self._weighted_log_prob(Xc)
+        lse = o.logsumexp(wlp, 1)
+        return o.exp(wlp - lse[:, None]), o.item(o.sum(o.astype(lse, o.f64), 0))
+
+    def _kmeanspp_init(self, X, rng):
+        """Greedy k-means++ seeds (as scikit-learn: 2 + log(k) candidates per step, keep
+        the one that most reduces inertia) on a random subsample, refined by a few
+        k-means (Lloyd) iterations; then hard-assign ALL points to the nearest centre
+        to get the starting weights / means / covariances."""
+        o = self.ops
+        n = X.shape[0]
+        idx = np.sort(rng.choice(n, size=min(n, self.init_sample), replace=False))
+        S = X[o.from_numpy(idx, o.i64)] if isinstance(o, _TorchOps) else X[idx]
+        m = S.shape[0]
+        n_trials = 2 + int(math.log(self.k))
+        S_sq = o.sum(S * S, 1)
+
+        def sqdist_to(P):                                      # (m, p) squared distances
+            return o.clip_min(S_sq[:, None] - 2.0 * (S @ o.transpose_last(P))
+                              + o.sum(P * P, 1)[None, :], 0.0)
+
+        first = int(rng.integers(m))
+        centres = [first]
+        d2 = sqdist_to(S[first][None, :])[:, 0]
+        for _ in range(1, self.k):
+            p = o.to_numpy(o.astype(d2, o.f64))
+            total = p.sum()
+            p = p / total if total > 0 else np.full(m, 1.0 / m)
+            cand = rng.choice(m, size=n_trials, p=p)
+            cand_idx = o.from_numpy(cand, o.i64) if isinstance(o, _TorchOps) else cand
+            dc = _minimum(o, d2[:, None], sqdist_to(S[cand_idx]))   # (m, trials)
+            best = int(np.argmin(o.to_numpy(o.astype(o.sum(dc, 0), o.f64))))
+            centres.append(int(cand[best]))
+            d2 = dc[:, best]
+        C = _stack_rows(o, [S[c] for c in centres])            # (k, d)
+
+        for _ in range(self.kmeans_iter):                      # Lloyd refinement on the subsample
+            lab = o.argmax(-sqdist_to(C), 1)
+            oh = _one_hot(o, lab, self.k)
+            cnt = o.sum(oh, 0)
+            newC = (o.transpose_last(oh) @ S) / o.clip_min(cnt, 1.0)[:, None]
+            empty = o.to_numpy(cnt) == 0                       # keep old centre if a cluster empties
+            if empty.any():
+                newC = _where_rows(o, empty, C, newC)
+            C = newC
+
+        def hard_resp(Xc):
+            dist = o.sum(Xc * Xc, 1)[:, None] - 2.0 * (Xc @ o.transpose_last(C)) + o.sum(C * C, 1)[None, :]
+            return _one_hot(o, o.argmax(-dist, 1), self.k), 0.0
+
+        nk, sx, sxx, _ = self._stats(X, hard_resp)
+        self._m_step(nk, sx, sxx, X.shape[0])
+
+    def _fit_once(self, X, rng):
+        o, n = self.ops, X.shape[0]
+        t0 = time.time()
+        self._kmeanspp_init(X, rng)
+        self._log(f"  init (greedy k-means++ + {self.kmeans_iter} k-means iterations) "
+                  f"{time.time() - t0:.1f}s on {o.device}")
+        prev, converged, it = -math.inf, False, 0
+        for it in range(1, self.max_iter + 1):
+            nk, sx, sxx, ll = self._stats(X, self._e_resp)
+            mean_ll = ll / n                                   # under the parameters BEFORE this M-step
+            self._m_step(nk, sx, sxx, n)
+            change = mean_ll - prev
+            if self.verbose and (it == 1 or it % 10 == 0):
+                self._log(f"  iter {it:4d}  mean log-likelihood {mean_ll:.6f}  change {change:.2e}")
+            prev = mean_ll
+            if abs(change) < self.tol:
+                converged = True
+                break
+        final_ll = self._mean_ll(X)                            # under the final parameters
+        o.sync()
+        self._log(f"  EM {'converged' if converged else 'stopped at max_iter'} after {it} iterations "
+                  f"({time.time() - t0:.1f}s); mean log-likelihood {final_ll:.6f}")
+        return final_ll, converged, it
+
+    def _mean_ll(self, X):
+        o, total = self.ops, 0.0
+        for s, e in self._chunks(X.shape[0]):
+            total += o.item(o.sum(o.astype(o.logsumexp(self._weighted_log_prob(X[s:e]), 1), o.f64), 0))
+        return total / X.shape[0]
+
+    # -- public API ---------------------------------------------------------
+    def fit(self, X_np: np.ndarray):
+        """Fit with n_init independent starts; keep the one with the highest likelihood."""
+        o = self.ops
+        X = o.from_numpy(np.ascontiguousarray(X_np, dtype=np.float32))
+        rng = np.random.default_rng(self.seed)
+        best = None
+        for run in range(self.n_init):
+            if self.n_init > 1:
+                self._log(f" start {run + 1}/{self.n_init}")
+            ll, conv, n_it = self._fit_once(X, rng)
+            if best is None or ll > best[0]:
+                best = (ll, conv, n_it, self.weights_, self.means_, self.covariances_)
+        self.lower_bound_, self.converged_, self.n_iter_ = best[0], best[1], best[2]
+        self.weights_, self.means_, self.covariances_ = best[3], best[4], best[5]
+        self._refresh_precisions()
+        return self
+
+    def predict(self, X_np: np.ndarray) -> np.ndarray:
+        o = self.ops
+        X = o.from_numpy(np.ascontiguousarray(X_np, dtype=np.float32))
+        out = []
+        for s, e in self._chunks(X.shape[0]):
+            out.append(o.to_numpy(o.argmax(self._weighted_log_prob(X[s:e]), 1)))
+        return np.concatenate(out).astype(np.int64)
+
+    def score(self, X_np: np.ndarray) -> float:
+        """Mean log-likelihood per sample."""
+        return self._mean_ll(self.ops.from_numpy(np.ascontiguousarray(X_np, dtype=np.float32)))
+
+    def n_parameters(self, d: int) -> int:
+        cov = self.k * d if self.covariance_type == "diag" else self.k * d * (d + 1) // 2
+        return int(cov + self.k * d + self.k - 1)
+
+    def bic(self, X_np) -> float:
+        n, d = X_np.shape
+        return -2 * self.score(X_np) * n + self.n_parameters(d) * math.log(n)
+
+    def aic(self, X_np) -> float:
+        n, d = X_np.shape
+        return -2 * self.score(X_np) * n + 2 * self.n_parameters(d)
+
+    def state(self) -> dict:
+        o = self.ops
+        return {
+            "n_components": self.k, "covariance_type": self.covariance_type,
+            "weights": o.to_numpy(self.weights_), "means": o.to_numpy(self.means_),
+            "covariances": o.to_numpy(self.covariances_),
+            "converged": self.converged_, "n_iter": self.n_iter_,
+            "reg_covar": self.reg_covar, "seed": self.seed,
+        }
+
+    def save(self, path: str) -> None:
+        np.savez(path, **{k: v for k, v in self.state().items()})
+
+
+# ---------------------------------------------------------------------------
+# small backend-neutral helpers
+# ---------------------------------------------------------------------------
+def _concat(o, cols):
+    if isinstance(o, _TorchOps):
+        return o.t.cat(cols, dim=1)
+    return np.concatenate(cols, axis=1)
+
+
+def _stack_rows(o, rows):
+    if isinstance(o, _TorchOps):
+        return o.t.stack(rows, dim=0)
+    return np.stack(rows, axis=0)
+
+
+def _minimum(o, a, b):
+    if isinstance(o, _TorchOps):
+        return o.t.minimum(a, b)
+    return np.minimum(a, b)
+
+
+def _one_hot(o, labels, k):
+    if isinstance(o, _TorchOps):
+        return o.t.nn.functional.one_hot(labels, num_classes=k).to(o.f32)
+    out = np.zeros((labels.shape[0], k), dtype=np.float32)
+    out[np.arange(labels.shape[0]), labels] = 1.0
+    return out
+
+
+def _where_rows(o, mask_np, a, b):
+    """Rows of `a` where mask (numpy bool, length k) is True, else rows of `b`."""
+    if isinstance(o, _TorchOps):
+        m = o.t.as_tensor(mask_np, device=o.device)[:, None]
+        return o.t.where(m, a, b)
+    return np.where(mask_np[:, None], a, b)
+
+
+# ---------------------------------------------------------------------------
+# convenience functions on an (n, d) embedding array
+# ---------------------------------------------------------------------------
+def _has_torch_cuda() -> bool:
+    try:
+        import torch
+        return bool(torch.cuda.is_available())
+    except ImportError:
+        return False
+
+
+def gmm(X, n_components: int, covariance_type: str = "diag", backend: str = "auto", seed: int = 0, **kwargs):
+    """Fit a Gaussian mixture to X (n x d) -> (labels, model). backend "auto" = torch on
+    a GPU if there is one, else the numpy implementation (same algorithm)."""
+    if backend == "auto":
+        backend = "torch" if _has_torch_cuda() else "numpy"
+    model = GaussianMixtureGPU(n_components, covariance_type=covariance_type, seed=seed,
+                               backend=backend, **kwargs)
+    X = np.ascontiguousarray(X, dtype=np.float32)
+    model.fit(X)
+    return np.asarray(model.predict(X)), model
+
+
+def kmeans(X, n_clusters: int = 10, seed: int = 0, **kwargs):
+    """scikit-learn KMeans on X (n x d) -> (labels, model)."""
+    from sklearn.cluster import KMeans
+    model = KMeans(n_clusters=n_clusters, random_state=seed, **kwargs).fit(X)
+    return model.labels_.astype(np.int64), model

spatialrna-0.3.0.dev1/src/spatialrna/datasets.py

@@ -0,0 +1,51 @@
+"""Datasets over many stored tile / subgraph graphs (spatialrna 0.2's on-disk datasets,
+for zarr stores). Graphs are read lazily, one at a time.
+
+    from spatialrna.datasets import ZarrGraphDataset
+    ds = ZarrGraphDataset.from_sample_store("data/S01.zarr", kind="subgraphs")
+    ds = ZarrGraphDataset.from_manifest("data/manifest.json")
+    for data in ds: ...            # torch_geometric Data per tile / subgraph
+"""
+from __future__ import annotations
+
+import json
+import os
+import os.path as osp
+import re
+from typing import List, Optional, Sequence
+
+from torch.utils.data import Dataset
+
+from .graph_zarr_io import load_data_zarr, zarr_group_exists
+
+
+class ZarrGraphDataset(Dataset):
+    """One torch_geometric Data per zarr graph group (tile or subgraph store).
+    keys: only these arrays (e.g. ["x", "edge_index"]); None = all."""
+
+    def __init__(self, paths: Sequence[str], keys: Optional[Sequence[str]] = None):
+        self.paths: List[str] = [str(p) for p in paths]
+        self.keys = list(keys) if keys else None
+
+    def __len__(self) -> int:
+        return len(self.paths)
+
+    def __getitem__(self, i: int):
+        return load_data_zarr(self.paths[i], keys=self.keys) if self.keys else load_data_zarr(self.paths[i])
+
+    @classmethod
+    def from_manifest(cls, manifest_path: str, keys=None) -> "ZarrGraphDataset":
+        """A manifest written by the pipeline: a JSON list of {"path": ...} entries."""
+        with open(manifest_path) as f:
+            entries = json.load(f)
+        return cls([e["path"] if isinstance(e, dict) else e for e in entries], keys)
+
+    @classmethod
+    def from_sample_store(cls, sample_store: str, kind: str = "tiles", keys=None) -> "ZarrGraphDataset":
+        """Every tile<t> group under <sample>.zarr/tiles (kind="tiles") or /subgraphs,
+        in tile order; empty / unfinished groups are skipped."""
+        root = osp.join(sample_store, kind)
+        names = [n for n in os.listdir(root) if re.match(r"^tile\d+$", n)] if osp.isdir(root) else []
+        names.sort(key=lambda n: int(n[4:]))
+        paths = [osp.join(root, n) for n in names if zarr_group_exists(osp.join(root, n))]
+        return cls(paths, keys)