significant speed up of imputation

src/cellpin/models/cellpin_model.py

@@ -39,6 +39,7 @@
 from cellpin.models.utils import (
     build_data_loaders,
     load_config_and_checkpoint,
+    mc_log1p_norm,
     save_checkpoint,
 )
 from cellpin.models.vae import CellPinVAE
@@ -1093,6 +1094,7 @@ def impute(
         norm_target_sum: float = 1e3,
         area_key: str | None = None,
         nb_count_samples: int = 100,
+        nb_seed: int | None = None,
         return_int: bool = False,
         return_sparse: bool = True,
         table_key: str = "table",
@@ -1122,6 +1124,12 @@ def impute(
                 Because log1p is concave, Jensen's inequality means
                 ``log1p(norm(E[X])) > E[log1p(norm(X))]``; sampling inside the
                 transform corrects this bias.  More samples → lower variance.
+            nb_seed: Seed for the NB draws (only used when ``return_norm=True``).
+                Pass an integer for a reproducible ``layers['imputed_norm']``;
+                ``None`` (default) draws fresh entropy on every call.  For a
+                fixed seed the result is independent of thread count, though
+                CPU and GPU use different samplers and so agree only up to
+                Monte-Carlo error.
             return_int: If ``True``, round ``X`` to integer counts (``int32``).
             return_sparse: If ``True`` (default), store ``X``, ``layers['imputed']``,
                 and ``layers['imputed_norm']`` as :class:`scipy.sparse.csr_matrix`.
@@ -1175,9 +1183,6 @@ def impute(
             # drawing K samples from NB, normalising and log1p-ing each draw,
             # then averaging — so norm+log1p go *inside* the MC loop.
             px_r_np = np.exp(self.vae.px_r.detach().cpu().numpy()).astype(np.float64)  # (n_genes,)
-            mu = counts.astype(np.float64)
-            # NB success-probability: p = theta / (theta + mu)
-            p = px_r_np / (px_r_np + np.clip(mu, 1e-8, None))
 
             # Resolve area column: explicit > auto-detect "cell_area" > None
             resolved_area_key = area_key
@@ -1194,18 +1199,15 @@ def impute(
             else:
                 scale = None  # use per-draw library size
 
-            K = max(nb_count_samples, 1)
-            log1p_acc = np.zeros_like(mu)
-            for _ in range(K):
-                draw = np.random.negative_binomial(px_r_np, p).astype(np.float64)
-                if scale is not None:
-                    normed = draw * scale[:, np.newaxis]
-                else:
-                    lib = draw.sum(axis=1, keepdims=True).clip(1e-12)
-                    normed = draw * (norm_target_sum / lib)
-                log1p_acc += np.log1p(normed)
-
-            norm_layer = (log1p_acc / K).astype(np.float32)
+            norm_layer = mc_log1p_norm(
+                counts,
+                px_r_np,
+                n_samples=max(nb_count_samples, 1),
+                norm_target_sum=norm_target_sum,
+                scale=scale,
+                seed=nb_seed,
+                device=self.device,
+            )
             adata_out.layers["imputed_norm"] = sp.csr_matrix(norm_layer) if return_sparse else norm_layer
 
         if sdata is not None:

src/cellpin/models/utils/__init__.py

@@ -3,5 +3,11 @@
     load_config_and_checkpoint,
     save_checkpoint,
 )
+from cellpin.models.utils.nb_sampling import mc_log1p_norm
 
-__all__ = ["build_data_loaders", "load_config_and_checkpoint", "save_checkpoint"]
+__all__ = [
+    "build_data_loaders",
+    "load_config_and_checkpoint",
+    "mc_log1p_norm",
+    "save_checkpoint",
+]

src/cellpin/models/utils/nb_sampling.py

@@ -0,0 +1,202 @@
+"""Monte-Carlo estimation of ``E[log1p(norm(X))]`` for ``X ~ NB(mu, theta)``.
+
+Used by :meth:`cellpin.CellPin.impute` to build ``layers['imputed_norm']``.
+Because ``log1p`` is concave, Jensen's inequality gives
+``log1p(norm(E[X])) > E[log1p(norm(X))]``; drawing counts *inside* the
+transform removes that bias.
+
+The estimator is parallel over cells, so the work is split into
+cell chunks that are sampled independently.  Two backends are available:
+
+* ``numpy`` — chunks run on a thread pool.  :class:`numpy.random.Generator`
+  releases the GIL while sampling, so this scales close to linearly with cores.
+* ``torch`` — chunks run sequentially on a CUDA device, using the
+  Gamma-Poisson representation of the NegativeBinomial.
+
+Both backends draw one independent RNG stream per chunk, and the chunking is a
+function of the input shape alone, so for a fixed seed the result is
+reproducible and independent of thread count.  The two backends use different
+samplers and therefore agree only up to Monte-Carlo error.
+"""
+
+from __future__ import annotations
+
+import os
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import torch
+
+# Elements (cells x genes) per chunk.  Each chunk holds a handful of arrays of
+# this size, so ~2M elements keeps a worker's working set around 50 MB.
+_CHUNK_ELEMENTS = 2_000_000
+
+# Lower bound on the number of chunks, so a thread pool this wide stays busy.
+_MIN_CHUNKS = 64
+
+# Floor on mu before forming the NB success probability / Gamma rate.
+_MU_FLOOR = 1e-8
+
+
+def _cell_chunks(n_cells: int, n_genes: int) -> list[slice]:
+    """Split the cell axis into chunks bounded by :data:`_CHUNK_ELEMENTS`.
+
+    Chunking on the *cell* axis (never the gene axis) is required: the
+    library-size branch normalises by a per-cell sum over all genes.
+
+    Depends on the input shape only, never on the worker count, so that a
+    seeded run reproduces regardless of how many threads it gets.
+    """
+    rows_per_chunk = max(1, _CHUNK_ELEMENTS // max(n_genes, 1))
+    rows_per_chunk = min(rows_per_chunk, max(1, -(-n_cells // _MIN_CHUNKS)))
+    n_chunks = max(1, -(-n_cells // rows_per_chunk))
+    bounds = np.linspace(0, n_cells, n_chunks + 1).astype(int)
+    return [slice(lo, hi) for lo, hi in zip(bounds[:-1], bounds[1:], strict=True) if hi > lo]
+
+
+def _chunk_numpy(
+    mu: np.ndarray,
+    theta: np.ndarray,
+    scale: np.ndarray | None,
+    n_samples: int,
+    norm_target_sum: float,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Accumulate ``mean_k log1p(norm(draw_k))`` for one chunk of cells."""
+    # p = theta / (theta + mu); float64 because it parameterises the sampler.
+    p = theta / (theta + np.clip(mu, _MU_FLOOR, None).astype(np.float64))
+    acc = np.zeros(mu.shape, dtype=np.float32)
+    for _ in range(n_samples):
+        draw = rng.negative_binomial(theta, p).astype(np.float32)
+        if scale is not None:
+            normed = draw * scale[:, np.newaxis]
+        else:
+            lib = draw.sum(axis=1, keepdims=True).clip(1e-12)
+            normed = draw * (norm_target_sum / lib)
+        acc += np.log1p(normed)
+    return acc / n_samples
+
+
+def _chunk_torch(
+    mu: np.ndarray,
+    theta: torch.Tensor,
+    scale: np.ndarray | None,
+    n_samples: int,
+    norm_target_sum: float,
+    generator: torch.Generator,
+    device: torch.device,
+) -> np.ndarray:
+    """GPU counterpart of :func:`_chunk_numpy` via the Gamma-Poisson mixture.
+
+    ``X ~ NB(mu, theta)`` iff ``X | lam ~ Poisson(lam)`` with
+    ``lam ~ Gamma(shape=theta, rate=theta/mu)``.
+    """
+    mu_t = torch.from_numpy(mu).to(device=device, dtype=torch.float32)
+    concentration = theta.expand_as(mu_t).contiguous()
+    rate = theta / mu_t.clamp_min(_MU_FLOOR)
+    scale_t = None if scale is None else torch.from_numpy(scale).to(device=device, dtype=torch.float32)[:, None]
+
+    acc = torch.zeros_like(mu_t)
+    for _ in range(n_samples):
+        lam = torch._standard_gamma(concentration, generator) / rate
+        draw = torch.poisson(lam, generator=generator)
+        if scale_t is not None:
+            normed = draw * scale_t
+        else:
+            lib = draw.sum(dim=1, keepdim=True).clamp_min(1e-12)
+            normed = draw * (norm_target_sum / lib)
+        acc += torch.log1p(normed)
+    return (acc / n_samples).cpu().numpy()
+
+
+def mc_log1p_norm(
+    mu: np.ndarray,
+    theta: np.ndarray,
+    n_samples: int,
+    norm_target_sum: float,
+    scale: np.ndarray | None = None,
+    seed: int | None = None,
+    device: torch.device | str | None = None,
+    n_threads: int | None = None,
+) -> np.ndarray:
+    """Estimate ``E[log1p(norm(X))]`` with ``X ~ NB(mu, theta)`` per element.
+
+    Args:
+        mu: Expected counts ``(n_cells, n_genes)``.
+        theta: Per-gene inverse dispersion ``(n_genes,)``, strictly positive.
+        n_samples: Number of NB draws averaged per element.
+        norm_target_sum: Target total counts after normalisation.
+        scale: Per-cell normalisation factor ``(n_cells,)`` for area-based
+            normalisation.  When ``None``, each draw is normalised by its own
+            library size.
+        seed: Seed for the per-chunk RNG streams.  ``None`` draws entropy from
+            the OS, matching unseeded behaviour.
+        device: Torch device for the GPU backend.  ``None`` or a CPU device
+            selects the threaded numpy backend, which is faster on CPU because
+            torch's CPU RNG is single-threaded.
+        n_threads: Worker threads for the numpy backend.  Defaults to the
+            number of CPUs, capped at 32.  Ignored by the torch backend.
+
+    Returns:
+    -------
+        Float32 array ``(n_cells, n_genes)``.
+    """
+    if n_samples < 1:
+        raise ValueError(f"n_samples must be >= 1, got {n_samples}")
+
+    n_cells, n_genes = mu.shape
+    theta = np.asarray(theta, dtype=np.float64)
+    if theta.shape != (n_genes,):
+        raise ValueError(f"theta must have shape ({n_genes},), got {theta.shape}")
+    if scale is not None:
+        scale = np.asarray(scale, dtype=np.float32)
+        if scale.shape != (n_cells,):
+            raise ValueError(f"scale must have shape ({n_cells},), got {scale.shape}")
+
+    mu = np.ascontiguousarray(mu, dtype=np.float32)
+
+    device = torch.device(device) if device is not None else None
+    use_torch = device is not None and device.type != "cpu"
+
+    n_workers = 1 if use_torch else (n_threads or min(32, os.cpu_count() or 1))
+    chunks = _cell_chunks(n_cells, n_genes)
+    seeds = np.random.SeedSequence(seed).spawn(len(chunks))
+
+    out = np.empty((n_cells, n_genes), dtype=np.float32)
+
+    if use_torch:
+        theta_t = torch.from_numpy(theta).to(device=device, dtype=torch.float32)
+        for sl, seed_seq in zip(chunks, seeds, strict=True):
+            generator = torch.Generator(device=device)
+            generator.manual_seed(int(seed_seq.generate_state(2, dtype=np.uint64)[0] >> 1))
+            out[sl] = _chunk_torch(
+                mu[sl],
+                theta_t,
+                None if scale is None else scale[sl],
+                n_samples,
+                norm_target_sum,
+                generator,
+                device,
+            )
+        return out
+
+    def _run(job: tuple[slice, np.random.SeedSequence]) -> None:
+        sl, seed_seq = job
+        out[sl] = _chunk_numpy(
+            mu[sl],
+            theta,
+            None if scale is None else scale[sl],
+            n_samples,
+            norm_target_sum,
+            np.random.default_rng(seed_seq),
+        )
+
+    jobs = list(zip(chunks, seeds, strict=True))
+    if n_workers == 1 or len(jobs) == 1:
+        for job in jobs:
+            _run(job)
+    else:
+        with ThreadPoolExecutor(max_workers=n_workers) as pool:
+            list(pool.map(_run, jobs))
+
+    return out

tests/models/test_cellpin.py

@@ -3,6 +3,7 @@
 import anndata as ad
 import numpy as np
 import pytest
+import torch
 from torch.utils.data import DataLoader
 
 from cellpin.models import CellPin
@@ -181,3 +182,79 @@ def test_impute_return_int_sparse(small_datasets):
     assert sp.issparse(adata_out.layers["imputed"])
     assert adata_out.layers["imputed"].dtype == np.int32
     assert adata_out.X.min() >= 0
+
+
+def test_impute_return_norm(small_datasets):
+    """return_norm=True adds a finite, non-negative imputed_norm layer."""
+    import scipy.sparse as sp
+
+    sc_ds, st_ds = small_datasets
+    model = CellPin(sc_dataset=sc_ds, config=MINIMAL_CONFIG)
+    model.eval()
+
+    loader = DataLoader(st_ds, batch_size=4, shuffle=False)
+    adata_out = model.impute(loader, mc_samples=2, return_norm=True, nb_count_samples=10, return_sparse=False)
+
+    norm = adata_out.layers["imputed_norm"]
+    assert not sp.issparse(norm)
+    assert norm.shape == (len(st_ds), sc_ds.X.shape[1])
+    assert norm.dtype == np.float32
+    assert np.isfinite(norm).all()
+    assert (norm >= 0).all()
+
+    adata_sparse = model.impute(loader, mc_samples=2, return_norm=True, nb_count_samples=10, return_sparse=True)
+    assert sp.issparse(adata_sparse.layers["imputed_norm"])
+
+
+def test_impute_nb_seed_is_reproducible(small_datasets):
+    """nb_seed pins the imputed_norm layer; omitting it does not."""
+    sc_ds, st_ds = small_datasets
+    model = CellPin(sc_dataset=sc_ds, config=MINIMAL_CONFIG)
+    model.eval()
+
+    loader = DataLoader(st_ds, batch_size=4, shuffle=False)
+    kwargs = {
+        "mc_samples": 2,
+        "return_norm": True,
+        "nb_count_samples": 10,
+        "return_sparse": False,
+    }
+    # mc_impute makes the counts themselves stochastic, so pin the whole run.
+    torch.manual_seed(0)
+    a = model.impute(loader, nb_seed=123, **kwargs).layers["imputed_norm"]
+    torch.manual_seed(0)
+    b = model.impute(loader, nb_seed=123, **kwargs).layers["imputed_norm"]
+    torch.manual_seed(0)
+    c = model.impute(loader, nb_seed=456, **kwargs).layers["imputed_norm"]
+
+    np.testing.assert_array_equal(a, b)
+    assert not np.array_equal(a, c)
+
+
+def test_impute_return_norm_area_key(small_datasets):
+    """area_key switches to area normalisation and rejects non-positive areas."""
+    sc_ds, st_ds = small_datasets
+    model = CellPin(sc_dataset=sc_ds, config=MINIMAL_CONFIG)
+    model.eval()
+
+    loader = DataLoader(st_ds, batch_size=4, shuffle=False)
+    obs_adata = ad.AnnData(X=np.zeros((len(st_ds), 0), dtype=np.float32))
+    obs_adata.obs["cell_area"] = np.linspace(20.0, 100.0, len(st_ds))
+
+    adata_out = model.impute(
+        loader,
+        obs_adata=obs_adata,
+        mc_samples=2,
+        return_norm=True,
+        nb_count_samples=10,
+        area_key="cell_area",
+        return_sparse=False,
+    )
+    assert np.isfinite(adata_out.layers["imputed_norm"]).all()
+
+    with pytest.raises(ValueError, match="not found in adata.obs"):
+        model.impute(loader, obs_adata=obs_adata, mc_samples=2, return_norm=True, area_key="missing")
+
+    obs_adata.obs["cell_area"] = 0.0
+    with pytest.raises(ValueError, match="areas must be positive"):
+        model.impute(loader, obs_adata=obs_adata, mc_samples=2, return_norm=True, area_key="cell_area")