diff --git a/src/cellpin/models/cellpin_model.py b/src/cellpin/models/cellpin_model.py index a716a7a..4181833 100644 --- a/src/cellpin/models/cellpin_model.py +++ b/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: diff --git a/src/cellpin/models/utils/__init__.py b/src/cellpin/models/utils/__init__.py index 0c37588..a493a61 100644 --- a/src/cellpin/models/utils/__init__.py +++ b/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", +] diff --git a/src/cellpin/models/utils/nb_sampling.py b/src/cellpin/models/utils/nb_sampling.py new file mode 100644 index 0000000..86534ca --- /dev/null +++ b/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 diff --git a/tests/models/test_cellpin.py b/tests/models/test_cellpin.py index 69c1972..4f902c5 100644 --- a/tests/models/test_cellpin.py +++ b/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") diff --git a/tests/models/utils/test_nb_sampling.py b/tests/models/utils/test_nb_sampling.py new file mode 100644 index 0000000..4a0d540 --- /dev/null +++ b/tests/models/utils/test_nb_sampling.py @@ -0,0 +1,134 @@ +"""Tests for the Monte-Carlo NB normalisation helper.""" + +import numpy as np +import pytest +import torch + +from cellpin.models.utils import mc_log1p_norm +from cellpin.models.utils.nb_sampling import _cell_chunks, _chunk_torch + +TARGET = 1e3 + + +@pytest.fixture +def mu_theta(): + rng = np.random.default_rng(0) + mu = rng.gamma(0.4, 2.5, size=(120, 60)).astype(np.float32) + theta = np.exp(rng.normal(size=60)) + return mu, theta + + +def _reference(mu, theta, scale, n_samples, seed): + """Original impute() implementation, kept as an independent reference.""" + rng = np.random.default_rng(seed) + mu64 = mu.astype(np.float64) + p = theta / (theta + np.clip(mu64, 1e-8, None)) + acc = np.zeros_like(mu64) + for _ in range(n_samples): + draw = rng.negative_binomial(theta, p).astype(np.float64) + if scale is not None: + normed = draw * scale[:, np.newaxis] + else: + normed = draw * (TARGET / draw.sum(axis=1, keepdims=True).clip(1e-12)) + acc += np.log1p(normed) + return (acc / n_samples).astype(np.float32) + + +def test_shape_and_dtype(mu_theta): + mu, theta = mu_theta + out = mc_log1p_norm(mu, theta, 5, TARGET, seed=0) + assert out.shape == mu.shape + assert out.dtype == np.float32 + assert np.isfinite(out).all() + assert (out >= 0).all() + + +def test_seed_is_reproducible_across_thread_counts(mu_theta): + mu, theta = mu_theta + a = mc_log1p_norm(mu, theta, 5, TARGET, seed=7, n_threads=1) + b = mc_log1p_norm(mu, theta, 5, TARGET, seed=7, n_threads=8) + c = mc_log1p_norm(mu, theta, 5, TARGET, seed=8, n_threads=8) + np.testing.assert_array_equal(a, b) + assert not np.array_equal(a, c) + + +def test_unseeded_runs_differ(mu_theta): + mu, theta = mu_theta + a = mc_log1p_norm(mu, theta, 5, TARGET) + b = mc_log1p_norm(mu, theta, 5, TARGET) + assert not np.array_equal(a, b) + + +@pytest.mark.parametrize("use_scale", [True, False]) +def test_matches_reference_within_mc_error(mu_theta, use_scale): + """New estimator must sit within the reference's own run-to-run spread.""" + mu, theta = mu_theta + scale = np.full(mu.shape[0], 4.0) if use_scale else None + k = 400 + ref_a = _reference(mu, theta, scale, k, seed=1) + ref_b = _reference(mu, theta, scale, k, seed=2) + new = mc_log1p_norm(mu, theta, k, TARGET, scale=scale, seed=3) + + baseline = np.sqrt(((ref_a - ref_b) ** 2).mean()) + cross = np.sqrt(((new - ref_a) ** 2).mean()) + assert cross < baseline * 1.2 + assert abs(float(np.mean(new - ref_a))) < 0.1 * baseline + + +@pytest.mark.parametrize("use_scale", [True, False]) +def test_torch_backend_matches_reference(mu_theta, use_scale): + """Gamma-Poisson backend targets the same expectation as the NB sampler.""" + mu, theta = mu_theta + scale = np.full(mu.shape[0], 4.0, dtype=np.float32) if use_scale else None + k = 400 + ref_a = _reference(mu, theta, None if scale is None else scale.astype(np.float64), k, seed=1) + ref_b = _reference(mu, theta, None if scale is None else scale.astype(np.float64), k, seed=2) + + generator = torch.Generator(device="cpu") + generator.manual_seed(3) + new = _chunk_torch( + np.ascontiguousarray(mu, np.float32), + torch.from_numpy(theta).float(), + scale, + k, + TARGET, + generator, + torch.device("cpu"), + ) + + baseline = np.sqrt(((ref_a - ref_b) ** 2).mean()) + assert np.sqrt(((new - ref_a) ** 2).mean()) < baseline * 1.2 + + +def test_jensen_correction_is_downward(mu_theta): + """Sampling inside log1p must sit below the plug-in log1p(norm(E[X])).""" + mu, theta = mu_theta + scale = np.full(mu.shape[0], 4.0) + corrected = mc_log1p_norm(mu, theta, 400, TARGET, scale=scale, seed=0) + plugin = np.log1p(mu.astype(np.float64) * scale[:, None]) + assert corrected.mean() < plugin.mean() + + +def test_chunks_tile_the_cell_axis(): + for n_cells, n_genes in ((1, 10), (7, 10), (5000, 2000), (200_000, 5000)): + chunks = _cell_chunks(n_cells, n_genes) + covered = np.concatenate([np.arange(n_cells)[s] for s in chunks]) + np.testing.assert_array_equal(covered, np.arange(n_cells)) + assert all(s.stop > s.start for s in chunks) + + +def test_zero_mu_gives_zero(mu_theta): + _, theta = mu_theta + mu = np.zeros((10, 60), dtype=np.float32) + out = mc_log1p_norm(mu, theta, 5, TARGET, scale=np.ones(10), seed=0) + np.testing.assert_array_equal(out, np.zeros_like(out)) + + +def test_invalid_arguments(mu_theta): + mu, theta = mu_theta + with pytest.raises(ValueError, match="n_samples"): + mc_log1p_norm(mu, theta, 0, TARGET) + with pytest.raises(ValueError, match="theta"): + mc_log1p_norm(mu, theta[:5], 2, TARGET) + with pytest.raises(ValueError, match="scale"): + mc_log1p_norm(mu, theta, 2, TARGET, scale=np.ones(3))