From dfb8b1c97144827d6602751430cbb26991a187f7 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 29 Jun 2026 17:32:51 +0200 Subject: [PATCH 1/3] add CLR Signed-off-by: Intron7 --- docs/api/scanpy_gpu.md | 1 + docs/references.bib | 10 + docs/release-notes/0.16.1.md | 5 + docs/release-notes/index.md | 2 + .../preprocessing/__init__.py | 8 +- .../preprocessing/_normalize.py | 242 +++++++++++++++--- src/rapids_singlecell/preprocessing/_utils.py | 42 +-- tests/dask/test_dask_mean_var.py | 7 +- tests/test_mean_var.py | 7 +- tests/test_normalization.py | 219 ++++++++++++++++ 10 files changed, 476 insertions(+), 67 deletions(-) create mode 100644 docs/release-notes/0.16.1.md diff --git a/docs/api/scanpy_gpu.md b/docs/api/scanpy_gpu.md index a4b624051..d3d030b08 100644 --- a/docs/api/scanpy_gpu.md +++ b/docs/api/scanpy_gpu.md @@ -19,6 +19,7 @@ Other than `tools`, preprocessing steps usually don’t return an easily interpr pp.filter_cells pp.filter_genes pp.normalize_total + pp.normalize_clr pp.log1p pp.highly_variable_genes pp.regress_out diff --git a/docs/references.bib b/docs/references.bib index 564ac8540..602760b92 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -392,3 +392,13 @@ @article{Lambiotte2014 month = {jul}, pages = {76--90}, } + +@article{Booeshaghi2026, + author = {Booeshaghi, A. Sina and Hallgrímsdóttir, Ingileif B. and Gálvez-Merchán, Ángel and Pachter, Lior}, + title = {Depth normalization for single-cell genomics count data}, + year = {2026}, + url = {https://doi.org/10.1101/2022.05.06.490859}, + doi = {10.1101/2022.05.06.490859}, + publisher = {Cold Spring Harbor Laboratory}, + journal = {bioRxiv}, +} diff --git a/docs/release-notes/0.16.1.md b/docs/release-notes/0.16.1.md new file mode 100644 index 000000000..7e994fb13 --- /dev/null +++ b/docs/release-notes/0.16.1.md @@ -0,0 +1,5 @@ +### 0.16.0 {small}`the-future` + +```{rubric} Features +``` +* Add {func}`~rapids_singlecell.pp.normalize_clr` for shifted centered log-ratio (PFlog1pPF) normalization, a variance-stabilizing, depth-invariant and rank-preserving count transform {cite:p}`Booeshaghi2026` {pr}`702` {smaller}`S Dicks` diff --git a/docs/release-notes/index.md b/docs/release-notes/index.md index 1f01cc8aa..95a605a66 100644 --- a/docs/release-notes/index.md +++ b/docs/release-notes/index.md @@ -4,6 +4,8 @@ ## Version 0.16.0 +```{include} /release-notes/0.16.1.md +``` ```{include} /release-notes/0.16.0.md ``` diff --git a/src/rapids_singlecell/preprocessing/__init__.py b/src/rapids_singlecell/preprocessing/__init__.py index 0e25d8c5f..407fd203b 100644 --- a/src/rapids_singlecell/preprocessing/__init__.py +++ b/src/rapids_singlecell/preprocessing/__init__.py @@ -3,7 +3,13 @@ from ._harmony_integrate import harmony_integrate from ._hvg import highly_variable_genes from ._neighbors import bbknn, neighbors -from ._normalize import log1p, normalize_pearson_residuals, normalize_total, sqrt +from ._normalize import ( + log1p, + normalize_clr, + normalize_pearson_residuals, + normalize_total, + sqrt, +) from ._pca import pca from ._qc import calculate_qc_metrics from ._regress_out import regress_out diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index e89a564e5..b41f61e22 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -3,7 +3,7 @@ import math import warnings from functools import partial -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Literal, Union import cupy as cp from anndata import AnnData @@ -17,7 +17,7 @@ _meta_sparse, ) -from ._utils import _check_gpu_X, _check_nonnegative_integers +from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var if TYPE_CHECKING: from cupyx.scipy.sparse import spmatrix @@ -124,6 +124,61 @@ def normalize_total( return X +def _normalize_total( + X: ArrayTypesDask, + target_sum: float | None, + *, + exclude_highly_expressed: bool = False, + max_fraction: float = 0.05, +) -> ArrayTypesDask: + if isinstance(X, DaskArray): + return _normalize_total_dask(X, target_sum) + elif isinstance(X, sparse.csr_matrix): + X = _normalize_total_csr( + X, + target_sum, + exclude_highly_expressed=exclude_highly_expressed, + max_fraction=max_fraction, + ) + elif isinstance(X, cp.ndarray): + X = _normalize_total_dense( + X, + target_sum, + exclude_highly_expressed=exclude_highly_expressed, + max_fraction=max_fraction, + ) + else: + raise ValueError(f"Cannot normalize {type(X)}") + return X + + +def _sum_axis1(X: ArrayTypesDask) -> cp.ndarray | DaskArray: + """Per-cell counts (sum over axis=1) for a CSR matrix, dense cupy array, or Dask array.""" + if isinstance(X, DaskArray): + return X.map_blocks( + _sum_axis1, + meta=cp.array((1.0,), dtype=X.dtype), + dtype=X.dtype, + chunks=(X.chunksize[0],), + drop_axis=1, + ) + if isinstance(X, sparse.csr_matrix): + from rapids_singlecell._cuda import _norm_cuda as _nc + + counts = cp.zeros(X.shape[0], dtype=X.dtype) + _nc.sum_major( + X.indptr, + X.data, + sums=counts, + major=X.shape[0], + stream=cp.cuda.get_current_stream().ptr, + ) + return counts + elif isinstance(X, cp.ndarray): + return X.sum(axis=1) + raise ValueError(f"Cannot compute row sums for {type(X)}") + + def _counts_to_scales( counts_per_cell: cp.ndarray, target_sum: float | None = None ) -> cp.ndarray: @@ -191,14 +246,7 @@ def _normalize_total_csr( from rapids_singlecell._cuda import _norm_cuda as _nc if gene_is_hi is None: - counts = cp.zeros(n_cells, dtype=X.dtype) - _nc.sum_major( - X.indptr, - X.data, - sums=counts, - major=n_cells, - stream=cp.cuda.get_current_stream().ptr, - ) + counts = _sum_axis1(X) else: counts = cp.zeros(n_cells, dtype=X.dtype) _nc.masked_sum_major( @@ -250,11 +298,11 @@ def _normalize_total_dense( # Compute per-cell counts, then prescaled multiply from rapids_singlecell._cuda import _norm_cuda as _nc - counts_per_cell = X.sum(axis=1) + counts_per_cell = _sum_axis1(X) if exclude_highly_expressed: hi_exp = X > max_fraction * counts_per_cell.reshape(-1, 1) gene_subset = ~hi_exp.any(axis=0) - counts_per_cell = X[:, gene_subset].sum(axis=1) + counts_per_cell = _sum_axis1(X[:, gene_subset]) scales = _counts_to_scales(counts_per_cell, target_sum) _nc.prescaled_mul_dense( @@ -306,37 +354,153 @@ def __mul(X_part): def _get_target_sum_dask(X: DaskArray) -> int: - if isinstance(X._meta, sparse.csr_matrix): - from rapids_singlecell._cuda import _norm_cuda as _nc + counts_per_cell = _sum_axis1(X).compute() + counts_per_cell = counts_per_cell[counts_per_cell > 0] + target_sum = cp.median(counts_per_cell) + return target_sum - def __sum(X_part): - counts_per_cell = cp.zeros(X_part.shape[0], dtype=X_part.dtype) - _nc.sum_major( - X_part.indptr, - X_part.data, - sums=counts_per_cell, - major=X_part.shape[0], - stream=cp.cuda.get_current_stream().ptr, - ) - return counts_per_cell - elif isinstance(X._meta, cp.ndarray): +def normalize_clr( + adata: AnnData, + *, + target_sum: float | None = None, + alpha: float | Literal["auto"] | None = None, + layer: str | None = None, + inplace: bool = True, + copy: bool = False, +) -> Union[AnnData, spmatrix, cp.ndarray, None]: # noqa: UP007 + r"""\ + Normalize counts with the shifted centered log-ratio (PFlog1pPF) transform. + + Computes the shifted centered log-ratio (CLR) transform + + .. math:: + T(x)_i = \log(u_i + 1) - \frac{1}{D} \sum_{j=1}^D \log(u_j + 1), + + where :math:`u_i = K \, x_i / \sum_j x_j` are the depth-normalized counts + (proportional fitting to a target depth :math:`K`) and :math:`D` is the number + of genes. Equivalently this is proportional fitting, then ``log1p``, then + per-cell mean-centering in log space (the centered-log-ratio step). The + transform is simultaneously variance-stabilizing, depth-invariant, and + rank-preserving :cite:p:`Booeshaghi2026`. + + To avoid densifying the matrix, the centering term is *not* subtracted in + place: ``adata.X`` (or `layer`) holds the sparse :math:`\log(u + 1)`, while the + per-cell centering offset :math:`\frac{1}{D}\sum_j \log(u_j + 1)` is written to + ``adata.obsm["clr_residuals"]`` and the raw per-cell depths to + ``adata.obsm["clr_cell_depths"]``. The full centered CLR is recovered as + ``adata.X - adata.obsm["clr_residuals"][:, None]``. + + .. note:: + When `adata.X` is a Dask array, deriving the proportional-fitting target + :math:`K` from the data requires a global reduction and therefore triggers + a blocking ``.compute()`` (for the default mean-depth target, and for + `alpha`, including ``alpha="auto"``). Only the scalar reduction is + materialized, not the matrix. Passing `target_sum` explicitly keeps it lazy. + + Parameters + ---------- + adata + The annotated data matrix of shape `n_obs` × `n_vars`. Rows correspond + to cells and columns to genes. + target_sum + Target depth :math:`K` for the proportional-fitting step. If `None` + (and `alpha` is not given), the empirical mean cell depth is used. This + is only an *intermediate* target: the subsequent log and centering steps + put each cell on the zero-sum hyperplane regardless of `K`. + alpha + Negative-binomial overdispersion of the dataset (``var = μ + α·μ²``). + When given, it overrides `target_sum` and sets :math:`K = 4 \cdot α \cdot s` + by the delta method, where :math:`s` is the mean cell depth, calibrating + the count-scale pseudocount to the variance-stabilizing value + ``y0 = 1/(4·α)`` :cite:p:`Booeshaghi2026`. Pass ``"auto"`` to estimate + :math:`α` from the data (closed-form least squares of ``var = μ + α·μ²`` + across genes). Raises a :class:`ValueError` if the estimated or supplied + :math:`α` is not positive (e.g. underdispersed data); pass `target_sum` + instead. + layer + Layer to normalize instead of `X`. If `None`, `X` is normalized. + inplace + Whether to update `adata` or return the result. + copy + Whether to return a copy or update `adata`. Not compatible with + `inplace=False`. + + Returns + ------- + Depending on `inplace`: + + - `inplace=True` (default): updates `adata.X` (or `layer`) with the sparse + :math:`\log(u + 1)`, writes ``adata.obsm["clr_cell_depths"]`` and + ``adata.obsm["clr_residuals"]``, and returns `None`. + - `copy=True`: performs the in-place update on a copy and returns it. + - `inplace=False`: returns the tuple ``(X, cell_depths, residuals)`` and + leaves `adata` untouched. + """ + if copy: + if not inplace: + msg = "`copy=True` cannot be used with `inplace=False`." + raise ValueError(msg) + adata = adata.copy() + X = _get_obs_rep(adata, layer=layer) + _check_gpu_X(X, allow_dask=True) + if not inplace: + X = X.copy() + if sparse.isspmatrix_csc(X): + X = X.tocsr() + X, cell_depths, residuals = _normalize_clr(X, target_sum=target_sum, alpha=alpha) + if inplace: + _set_obs_rep(adata, X, layer=layer) + adata.obsm["clr_cell_depths"] = cell_depths + adata.obsm["clr_residuals"] = residuals + if copy: + return adata + if not inplace: + return X, cell_depths, residuals + return None + - def __sum(X_part): - return X_part.sum(axis=1) +def _estimate_overdispersion(X: ArrayTypesDask) -> tuple[float, cp.ndarray]: + mean, var = _get_mean_var(X, axis=0, correction=0) + cell_depths = _sum_axis1(X) + if isinstance(X, DaskArray): + import dask + + mean, var, cell_depths = dask.compute(mean, var, cell_depths) + mean_sq = mean**2 + numerator = cp.sum((var - mean) * mean_sq) + denominator = cp.sum(mean_sq**2) + if float(denominator) == 0: + msg = "Cannot estimate overdispersion: all gene means are zero." + raise ValueError(msg) + return numerator / denominator, cell_depths + + +def _normalize_clr( + X: ArrayTypesDask, target_sum: float | None, alpha: float | Literal["auto"] | None +) -> tuple[ArrayTypesDask, cp.ndarray, cp.ndarray]: + if alpha == "auto": + alpha, cell_depths = _estimate_overdispersion(X) else: - raise ValueError(f"Cannot compute target sum for {type(X)}") - target_sum_chunk_matrices = X.map_blocks( - __sum, - meta=cp.array((1.0,), dtype=X.dtype), - dtype=X.dtype, - chunks=(X.chunksize[0],), - drop_axis=1, - ) - counts_per_cell = target_sum_chunk_matrices.compute() - counts_per_cell = counts_per_cell[counts_per_cell > 0] - target_sum = cp.median(counts_per_cell) - return target_sum + cell_depths = _sum_axis1(X) + if isinstance(cell_depths, DaskArray): + cell_depths = cell_depths.compute() + if bool((cell_depths == 0).any()): + warnings.warn("Some cells have zero counts", UserWarning) + if alpha is not None: + if alpha <= 0: + raise ValueError("alpha must be positive") + target_sum = 4.0 * alpha * float(cell_depths.mean()) + elif target_sum is None: + target_sum = float(cell_depths.mean()) + + X = _normalize_total(X, target_sum) + X = _calc_log1p(X) + # Centering offset = per-cell mean of log1p(PF) = row sum / n_genes. + # `_sum_axis1` already covers CSR/dense/Dask; avoids the unused variance pass. + residuals = _sum_axis1(X) / X.shape[1] + + return X, cell_depths, residuals def _calc_log1p(X: ArrayTypesDask, base: float | None = None) -> ArrayTypesDask: diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index e9bb49f9b..562d7fe79 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -63,7 +63,7 @@ def _sanitize_column(adata: AnnData, column: str): adata.obs[column] = c -def _mean_var_major(X, major, minor): +def _mean_var_major(X, major, minor, *, correction=1): from rapids_singlecell._cuda import _mean_var_cuda as _mv mean = cp.zeros(major, dtype=cp.float64) @@ -81,11 +81,11 @@ def _mean_var_major(X, major, minor): mean = mean / minor var = var / minor var -= cp.power(mean, 2) - var *= minor / (minor - 1) + var *= minor / (minor - correction) return mean, var -def _mean_var_minor(X, major, minor): +def _mean_var_minor(X, major, minor, *, correction=1): from rapids_singlecell._cuda import _mean_var_cuda as _mv mean = cp.zeros(minor, dtype=cp.float64) @@ -101,11 +101,11 @@ def _mean_var_minor(X, major, minor): mean /= major var /= major var -= mean**2 - var *= major / (major - 1) + var *= major / (major - correction) return mean, var -def _mean_var_minor_dask(X, major, minor): +def _mean_var_minor_dask(X, major, minor, *, correction=1): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ @@ -135,12 +135,12 @@ def __mean_var(X_part): ).sum(axis=0) mean /= major var /= major - var = (var - mean**2) * (major / (major - 1)) + var = (var - mean**2) * (major / (major - correction)) return mean, var # todo: Implement this dynamically for csc matrix as well -def _mean_var_major_dask(X, major, minor): +def _mean_var_major_dask(X, major, minor, *, correction=1): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ @@ -173,11 +173,11 @@ def __mean_var(X_part): mean = mean / minor var = var / minor var -= mean**2 - var *= minor / (minor - 1) + var *= minor / (minor - correction) return mean, var -def _mean_var_dense_dask(X, axis): +def _mean_var_dense_dask(X, axis, *, correction=1): """ Implements sum operation for dask array when the backend is cupy dense matrix """ @@ -209,11 +209,11 @@ def __mean_var(X_part): mean = mean / X.shape[axis] var = var / X.shape[axis] var -= mean**2 - var *= X.shape[axis] / (X.shape[axis] - 1) + var *= X.shape[axis] / (X.shape[axis] - correction) return mean, var -def _mean_var_dense(X, axis): +def _mean_var_dense(X, axis, *, correction=1): from ._kernels._mean_var_kernel import mean_sum, sq_sum var = sq_sum(X, axis=axis) @@ -221,30 +221,30 @@ def _mean_var_dense(X, axis): mean = mean / X.shape[axis] var = var / X.shape[axis] var -= cp.power(mean, 2) - var *= X.shape[axis] / (X.shape[axis] - 1) + var *= X.shape[axis] / (X.shape[axis] - correction) return mean, var -def _get_mean_var(X, axis=0): +def _get_mean_var(X, axis=0, *, correction=1): if issparse(X): if axis == 0: if isspmatrix_csr(X): major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_minor(X, major, minor) + mean, var = _mean_var_minor(X, major, minor, correction=correction) elif isspmatrix_csc(X): major = X.shape[1] minor = X.shape[0] - mean, var = _mean_var_major(X, major, minor) + mean, var = _mean_var_major(X, major, minor, correction=correction) elif axis == 1: if isspmatrix_csr(X): major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_major(X, major, minor) + mean, var = _mean_var_major(X, major, minor, correction=correction) elif isspmatrix_csc(X): major = X.shape[1] minor = X.shape[0] - mean, var = _mean_var_minor(X, major, minor) + mean, var = _mean_var_minor(X, major, minor, correction=correction) else: raise ValueError("axis must be either 0 or 1") elif isinstance(X, DaskArray): @@ -252,19 +252,19 @@ def _get_mean_var(X, axis=0): if axis == 0: major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_minor_dask(X, major, minor) + mean, var = _mean_var_minor_dask(X, major, minor, correction=correction) if axis == 1: major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_major_dask(X, major, minor) + mean, var = _mean_var_major_dask(X, major, minor, correction=correction) elif isinstance(X._meta, cp.ndarray): - mean, var = _mean_var_dense_dask(X, axis) + mean, var = _mean_var_dense_dask(X, axis, correction=correction) else: raise ValueError( "Type not supported. Please provide a CuPy ndarray or a CuPy sparse matrix. Or a Dask array with a CuPy ndarray or a CuPy sparse matrix as meta." ) else: - mean, var = _mean_var_dense(X, axis) + mean, var = _mean_var_dense(X, axis, correction=correction) return mean, var diff --git a/tests/dask/test_dask_mean_var.py b/tests/dask/test_dask_mean_var.py index b24f90b11..119da1082 100644 --- a/tests/dask/test_dask_mean_var.py +++ b/tests/dask/test_dask_mean_var.py @@ -17,7 +17,8 @@ @pytest.mark.parametrize("data_kind", ["sparse", "dense"]) @pytest.mark.parametrize("axis", [0, 1]) @pytest.mark.parametrize("dtype", [cp.float32, cp.float64]) -def test_mean_var(client, data_kind, axis, dtype): +@pytest.mark.parametrize("correction", [0, 1]) +def test_mean_var(client, data_kind, axis, dtype, correction): if data_kind == "dense": adata = pbmc68k_reduced() adata.X = adata.X.astype(dtype) @@ -31,8 +32,8 @@ def test_mean_var(client, data_kind, axis, dtype): dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() rsc.get.anndata_to_GPU(adata) - mean, var = _get_mean_var(adata.X, axis=axis) - dask_mean, dask_var = _get_mean_var(dask_data.X, axis=axis) + mean, var = _get_mean_var(adata.X, axis=axis, correction=correction) + dask_mean, dask_var = _get_mean_var(dask_data.X, axis=axis, correction=correction) dask_mean, dask_var = dask_mean.compute(), dask_var.compute() cp.testing.assert_allclose(mean, dask_mean) diff --git a/tests/test_mean_var.py b/tests/test_mean_var.py index 96c6f0d39..f6aeb18ae 100644 --- a/tests/test_mean_var.py +++ b/tests/test_mean_var.py @@ -13,7 +13,8 @@ @pytest.mark.parametrize("data_kind", ["csc", "csr", "dense"]) @pytest.mark.parametrize("axis", [0, 1]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) -def test_mean_var(data_kind, axis, dtype): +@pytest.mark.parametrize("correction", [0, 1]) +def test_mean_var(data_kind, axis, dtype, correction): if data_kind == "dense": adata = pbmc68k_reduced() else: @@ -24,8 +25,8 @@ def test_mean_var(data_kind, axis, dtype): adata.X = adata.X.astype(dtype) cudata = rsc.get.anndata_to_GPU(adata, copy=True) - mean, var = sc_get_mean_var(adata.X, axis=axis, correction=1) - rsc_mean, rsc_var = rsc_get_mean_var(cudata.X, axis=axis) + mean, var = sc_get_mean_var(adata.X, axis=axis, correction=correction) + rsc_mean, rsc_var = rsc_get_mean_var(cudata.X, axis=axis, correction=correction) cp.testing.assert_allclose(mean, rsc_mean) cp.testing.assert_allclose(var, rsc_var) diff --git a/tests/test_normalization.py b/tests/test_normalization.py index 75932358c..edd5ac297 100644 --- a/tests/test_normalization.py +++ b/tests/test_normalization.py @@ -261,3 +261,222 @@ def test_normalize_total_max_fraction_validation(): with pytest.raises(ValueError, match="`max_fraction` must be between 0 and 1"): rsc.pp.normalize_total(cudata, exclude_highly_expressed=True, max_fraction=-0.1) + + +# ------------------------------------------------------------------------------ +# normalize_clr (shifted CLR / PFlog1pPF) +# ------------------------------------------------------------------------------ + +# A small count matrix with no empty cells, used for the value/equivalence tests. +X_clr = np.array( + [[5, 0, 3, 2], [1, 1, 0, 4], [0, 7, 2, 1], [3, 3, 3, 3]], dtype="float32" +) + +CLR_ARRAY_TYPES = [cp.array, csr_matrix, csc_matrix] + + +def _to_np(x): + """cupy dense / cupy sparse / numpy -> dense numpy.""" + if hasattr(x, "toarray"): + x = x.toarray() + return cp.asnumpy(x) + + +def _estimate_alpha_reference(x) -> float: + """Closed-form OLS overdispersion, independent of the implementation.""" + x = np.asarray(x, dtype=np.float64) + mu = x.mean(axis=0) + var = (x**2).mean(axis=0) - mu**2 + mu2 = mu**2 + return float(np.sum((var - mu) * mu2) / np.sum(mu2 * mu2)) + + +def _clr_reference(x, *, target_sum=None, alpha=None) -> np.ndarray: + """Self-contained dense shifted-CLR, independent of the implementation. + + PF to a target depth, log1p, then subtract the per-cell mean. Empty cells + (zero depth) are left as all-zero rows, matching `normalize_clr`. + """ + x = np.asarray(x, dtype=np.float64) + depths = x.sum(axis=1) + if alpha is not None: + if alpha == "auto": + alpha = _estimate_alpha_reference(x) + target_sum = 4.0 * alpha * depths.mean() + elif target_sum is None: + target_sum = depths.mean() + safe_depths = np.where(depths == 0, 1.0, depths) + u = x * (target_sum / safe_depths)[:, None] + log_u = np.log1p(u) + return log_u - log_u.mean(axis=1, keepdims=True) + + +def _reconstruct_clr(X, residuals) -> np.ndarray: + """rsc keeps X = log1p(PF) sparse and the per-cell centering offset apart. + + The centered CLR is `X - offset[:, None]`; only the test materializes it. + """ + return _to_np(X) - cp.asnumpy(residuals).reshape(-1, 1) + + +@pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_normalize_clr_values(array_type, dtype): + """Reconstructed CLR matches the reference and every cell sums to zero. + + Asserting that dense / csr / csc inputs all equal the same dense reference + also proves the sparse "offset trick" matches the dense path. + """ + adata = AnnData(array_type(cp.asarray(X_clr).astype(dtype))) + rsc.pp.normalize_clr(adata) + + result = _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]) + np.testing.assert_allclose(result, _clr_reference(X_clr), rtol=1e-5, atol=1e-5) + # zero-sum (Aitchison) hyperplane + np.testing.assert_allclose(result.sum(axis=1), 0.0, atol=1e-5) + + +@pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) +@pytest.mark.parametrize( + "kwargs", + [{}, {"target_sum": 1e4}, {"alpha": 0.5}, {"alpha": "auto"}], + ids=["default", "target_sum", "alpha", "alpha_auto"], +) +def test_normalize_clr_params(array_type, kwargs): + adata = AnnData(array_type(cp.asarray(X_clr))) + rsc.pp.normalize_clr(adata, **kwargs) + np.testing.assert_allclose( + _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]), + _clr_reference(X_clr, **kwargs), + rtol=1e-5, + atol=1e-5, + ) + + +def test_normalize_clr_alpha_overrides_target_sum(): + """`alpha` sets target_sum = 4*alpha*scale and overrides any given `target_sum`.""" + alpha = 0.5 + scale = X_clr.sum(axis=1).mean() + + via_alpha = AnnData(csr_matrix(cp.asarray(X_clr))) + rsc.pp.normalize_clr(via_alpha, alpha=alpha) + + via_target = AnnData(csr_matrix(cp.asarray(X_clr))) + rsc.pp.normalize_clr(via_target, target_sum=4.0 * alpha * scale) + np.testing.assert_allclose( + _reconstruct_clr(via_alpha.X, via_alpha.obsm["clr_residuals"]), + _reconstruct_clr(via_target.X, via_target.obsm["clr_residuals"]), + rtol=1e-5, + atol=1e-5, + ) + + # passing both -> alpha wins, target_sum ignored + both = AnnData(csr_matrix(cp.asarray(X_clr))) + rsc.pp.normalize_clr(both, alpha=alpha, target_sum=999.0) + np.testing.assert_allclose( + _reconstruct_clr(both.X, both.obsm["clr_residuals"]), + _reconstruct_clr(via_alpha.X, via_alpha.obsm["clr_residuals"]), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) +def test_normalize_clr_alpha_auto(array_type): + """`alpha="auto"` estimates the overdispersion and matches an explicit alpha.""" + estimated = _estimate_alpha_reference(X_clr) + assert estimated > 0 + + auto = AnnData(array_type(cp.asarray(X_clr))) + rsc.pp.normalize_clr(auto, alpha="auto") + + explicit = AnnData(array_type(cp.asarray(X_clr))) + rsc.pp.normalize_clr(explicit, alpha=estimated) + np.testing.assert_allclose( + _reconstruct_clr(auto.X, auto.obsm["clr_residuals"]), + _reconstruct_clr(explicit.X, explicit.obsm["clr_residuals"]), + rtol=1e-5, + atol=1e-5, + ) + + +@pytest.mark.parametrize("alpha", [0.0, -0.5], ids=["zero", "negative"]) +def test_normalize_clr_nonpositive_alpha_raises(alpha): + """A non-positive `alpha` cannot derive K = 4*alpha*s and raises.""" + adata = AnnData(csr_matrix(cp.asarray(X_clr))) + with pytest.raises(ValueError, match=r"alpha.*positive"): + rsc.pp.normalize_clr(adata, alpha=alpha) + + +def test_normalize_clr_alpha_auto_zero_mean_raises(): + """`alpha="auto"` cannot estimate overdispersion when every gene mean is zero.""" + adata = AnnData(cp.zeros((3, 4), dtype=cp.float32)) + with pytest.raises(ValueError, match="Cannot estimate overdispersion"): + rsc.pp.normalize_clr(adata, alpha="auto") + + +@pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) +def test_normalize_clr_zero_cell(array_type): + """An empty cell stays all-zero, stays finite, and triggers a warning.""" + x = X_clr.copy() + x[1] = 0 # make the second cell empty + adata = AnnData(array_type(cp.asarray(x))) + with pytest.warns(UserWarning, match="zero counts"): + rsc.pp.normalize_clr(adata) + result = _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]) + assert np.isfinite(result).all() + np.testing.assert_allclose(result[1], 0.0, atol=1e-6) + + +def test_normalize_clr_inplace_false(): + adata = AnnData(csr_matrix(cp.asarray(X_clr))) + x_before = _to_np(adata.X).copy() + out = rsc.pp.normalize_clr(adata, inplace=False) + + # factored design: inplace=False returns (X, cell_depths, residuals) + X, _cell_depths, residuals = out + np.testing.assert_allclose( + _reconstruct_clr(X, residuals), _clr_reference(X_clr), rtol=1e-5, atol=1e-5 + ) + # input is left untouched + np.testing.assert_array_equal(_to_np(adata.X), x_before) + + +def test_normalize_clr_copy(): + adata = AnnData(csr_matrix(cp.asarray(X_clr))) + returned = rsc.pp.normalize_clr(adata, copy=True) + + assert isinstance(returned, AnnData) + assert returned is not adata + np.testing.assert_allclose( + _reconstruct_clr(returned.X, returned.obsm["clr_residuals"]), + _clr_reference(X_clr), + rtol=1e-5, + atol=1e-5, + ) + + +def test_normalize_clr_copy_inplace_error(): + adata = AnnData(csr_matrix(cp.asarray(X_clr))) + with pytest.raises( + ValueError, match="`copy=True` cannot be used with `inplace=False`" + ): + rsc.pp.normalize_clr(adata, copy=True, inplace=False) + + +def test_normalize_clr_layer(): + """`layer` targets that layer and leaves `X` untouched.""" + adata = AnnData( + csr_matrix(cp.asarray(X_clr)), + layers={"counts": csr_matrix(cp.asarray(X_clr))}, + ) + x_before = _to_np(adata.X).copy() + rsc.pp.normalize_clr(adata, layer="counts") + + np.testing.assert_array_equal(_to_np(adata.X), x_before) + np.testing.assert_allclose( + _reconstruct_clr(adata.layers["counts"], adata.obsm["clr_residuals"]), + _clr_reference(X_clr), + rtol=1e-5, + atol=1e-5, + ) From ecec334a1823ef2e67f0c12f1dfae666810d889e Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 29 Jun 2026 18:26:06 +0200 Subject: [PATCH 2/3] fix output Signed-off-by: Intron7 --- .../preprocessing/_normalize.py | 42 +++++++++++------- tests/test_normalization.py | 43 +++++++++++++------ 2 files changed, 57 insertions(+), 28 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index b41f61e22..2a0e13c51 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -24,6 +24,11 @@ from rapids_singlecell._utils import ArrayTypesDask +# normalize_clr writes its factored output to these fixed keys. +CLR_LAYER = "clr" +CLR_CELL_DEPTHS_KEY = "clr_cell_depths" +CLR_RESIDUALS_KEY = "clr_residuals" + def normalize_total( adata: AnnData, @@ -385,14 +390,17 @@ def normalize_clr( rank-preserving :cite:p:`Booeshaghi2026`. To avoid densifying the matrix, the centering term is *not* subtracted in - place: ``adata.X`` (or `layer`) holds the sparse :math:`\log(u + 1)`, while the - per-cell centering offset :math:`\frac{1}{D}\sum_j \log(u_j + 1)` is written to - ``adata.obsm["clr_residuals"]`` and the raw per-cell depths to + place and the source matrix is left untouched. ``adata.layers["clr"]`` holds + the sparse :math:`\log(u + 1)`, the per-cell centering offset + :math:`\frac{1}{D}\sum_j \log(u_j + 1)` is written to + ``adata.obsm["clr_residuals"]``, and the raw per-cell depths to ``adata.obsm["clr_cell_depths"]``. The full centered CLR is recovered as - ``adata.X - adata.obsm["clr_residuals"][:, None]``. + ``adata.layers["clr"] - adata.obsm["clr_residuals"][:, None]``; + :func:`~rapids_singlecell.pp.pca` consumes this factored form directly via + ``layer="clr"`` without ever materializing it. .. note:: - When `adata.X` is a Dask array, deriving the proportional-fitting target + When the input is a Dask array, deriving the proportional-fitting target :math:`K` from the data requires a global reduction and therefore triggers a blocking ``.compute()`` (for the default mean-depth target, and for `alpha`, including ``alpha="auto"``). Only the scalar reduction is @@ -419,7 +427,8 @@ def normalize_clr( :math:`α` is not positive (e.g. underdispersed data); pass `target_sum` instead. layer - Layer to normalize instead of `X`. If `None`, `X` is normalized. + Layer to read the counts from. If `None`, `X` is used. The result is + always written to ``adata.layers["clr"]``; the source is not modified. inplace Whether to update `adata` or return the result. copy @@ -430,10 +439,11 @@ def normalize_clr( ------- Depending on `inplace`: - - `inplace=True` (default): updates `adata.X` (or `layer`) with the sparse - :math:`\log(u + 1)`, writes ``adata.obsm["clr_cell_depths"]`` and - ``adata.obsm["clr_residuals"]``, and returns `None`. - - `copy=True`: performs the in-place update on a copy and returns it. + - `inplace=True` (default): writes the sparse :math:`\log(u + 1)` to + ``adata.layers["clr"]`` and the centering offsets / depths to + ``adata.obsm["clr_residuals"]`` / ``adata.obsm["clr_cell_depths"]``, + leaving the source matrix untouched, and returns `None`. + - `copy=True`: performs the update on a copy and returns it. - `inplace=False`: returns the tuple ``(X, cell_depths, residuals)`` and leaves `adata` untouched. """ @@ -444,15 +454,17 @@ def normalize_clr( adata = adata.copy() X = _get_obs_rep(adata, layer=layer) _check_gpu_X(X, allow_dask=True) - if not inplace: - X = X.copy() + # The PF step mutates the matrix in place; copy so the source (X / `layer`) + # is preserved (CSC `.tocsr()` already yields a fresh matrix). if sparse.isspmatrix_csc(X): X = X.tocsr() + else: + X = X.copy() X, cell_depths, residuals = _normalize_clr(X, target_sum=target_sum, alpha=alpha) if inplace: - _set_obs_rep(adata, X, layer=layer) - adata.obsm["clr_cell_depths"] = cell_depths - adata.obsm["clr_residuals"] = residuals + adata.layers[CLR_LAYER] = X + adata.obsm[CLR_CELL_DEPTHS_KEY] = cell_depths + adata.obsm[CLR_RESIDUALS_KEY] = residuals if copy: return adata if not inplace: diff --git a/tests/test_normalization.py b/tests/test_normalization.py index edd5ac297..0c61b4436 100644 --- a/tests/test_normalization.py +++ b/tests/test_normalization.py @@ -319,6 +319,11 @@ def _reconstruct_clr(X, residuals) -> np.ndarray: return _to_np(X) - cp.asnumpy(residuals).reshape(-1, 1) +def _clr_result(adata) -> np.ndarray: + """Centered CLR from the factored output written to `layers["clr"]`.""" + return _reconstruct_clr(adata.layers["clr"], adata.obsm["clr_residuals"]) + + @pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) @pytest.mark.parametrize("dtype", ["float32", "float64"]) def test_normalize_clr_values(array_type, dtype): @@ -328,12 +333,15 @@ def test_normalize_clr_values(array_type, dtype): also proves the sparse "offset trick" matches the dense path. """ adata = AnnData(array_type(cp.asarray(X_clr).astype(dtype))) + x_before = _to_np(adata.X).copy() rsc.pp.normalize_clr(adata) - result = _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]) + result = _clr_result(adata) np.testing.assert_allclose(result, _clr_reference(X_clr), rtol=1e-5, atol=1e-5) # zero-sum (Aitchison) hyperplane np.testing.assert_allclose(result.sum(axis=1), 0.0, atol=1e-5) + # source matrix is left untouched (result goes to layers["clr"]) + np.testing.assert_array_equal(_to_np(adata.X), x_before) @pytest.mark.parametrize("array_type", CLR_ARRAY_TYPES, ids=lambda f: f.__name__) @@ -346,7 +354,7 @@ def test_normalize_clr_params(array_type, kwargs): adata = AnnData(array_type(cp.asarray(X_clr))) rsc.pp.normalize_clr(adata, **kwargs) np.testing.assert_allclose( - _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]), + _clr_result(adata), _clr_reference(X_clr, **kwargs), rtol=1e-5, atol=1e-5, @@ -364,8 +372,8 @@ def test_normalize_clr_alpha_overrides_target_sum(): via_target = AnnData(csr_matrix(cp.asarray(X_clr))) rsc.pp.normalize_clr(via_target, target_sum=4.0 * alpha * scale) np.testing.assert_allclose( - _reconstruct_clr(via_alpha.X, via_alpha.obsm["clr_residuals"]), - _reconstruct_clr(via_target.X, via_target.obsm["clr_residuals"]), + _clr_result(via_alpha), + _clr_result(via_target), rtol=1e-5, atol=1e-5, ) @@ -374,8 +382,8 @@ def test_normalize_clr_alpha_overrides_target_sum(): both = AnnData(csr_matrix(cp.asarray(X_clr))) rsc.pp.normalize_clr(both, alpha=alpha, target_sum=999.0) np.testing.assert_allclose( - _reconstruct_clr(both.X, both.obsm["clr_residuals"]), - _reconstruct_clr(via_alpha.X, via_alpha.obsm["clr_residuals"]), + _clr_result(both), + _clr_result(via_alpha), rtol=1e-5, atol=1e-5, ) @@ -393,8 +401,8 @@ def test_normalize_clr_alpha_auto(array_type): explicit = AnnData(array_type(cp.asarray(X_clr))) rsc.pp.normalize_clr(explicit, alpha=estimated) np.testing.assert_allclose( - _reconstruct_clr(auto.X, auto.obsm["clr_residuals"]), - _reconstruct_clr(explicit.X, explicit.obsm["clr_residuals"]), + _clr_result(auto), + _clr_result(explicit), rtol=1e-5, atol=1e-5, ) @@ -423,7 +431,7 @@ def test_normalize_clr_zero_cell(array_type): adata = AnnData(array_type(cp.asarray(x))) with pytest.warns(UserWarning, match="zero counts"): rsc.pp.normalize_clr(adata) - result = _reconstruct_clr(adata.X, adata.obsm["clr_residuals"]) + result = _clr_result(adata) assert np.isfinite(result).all() np.testing.assert_allclose(result[1], 0.0, atol=1e-6) @@ -438,22 +446,28 @@ def test_normalize_clr_inplace_false(): np.testing.assert_allclose( _reconstruct_clr(X, residuals), _clr_reference(X_clr), rtol=1e-5, atol=1e-5 ) - # input is left untouched + # input is left untouched and nothing is written to the object np.testing.assert_array_equal(_to_np(adata.X), x_before) + assert "clr" not in adata.layers + assert "clr_residuals" not in adata.obsm def test_normalize_clr_copy(): adata = AnnData(csr_matrix(cp.asarray(X_clr))) + x_before = _to_np(adata.X).copy() returned = rsc.pp.normalize_clr(adata, copy=True) assert isinstance(returned, AnnData) assert returned is not adata np.testing.assert_allclose( - _reconstruct_clr(returned.X, returned.obsm["clr_residuals"]), + _clr_result(returned), _clr_reference(X_clr), rtol=1e-5, atol=1e-5, ) + # source matrix on the copy is preserved; original object untouched + np.testing.assert_array_equal(_to_np(returned.X), x_before) + assert "clr" not in adata.layers def test_normalize_clr_copy_inplace_error(): @@ -465,17 +479,20 @@ def test_normalize_clr_copy_inplace_error(): def test_normalize_clr_layer(): - """`layer` targets that layer and leaves `X` untouched.""" + """`layer` selects the input; output always goes to layers["clr"], sources kept.""" adata = AnnData( csr_matrix(cp.asarray(X_clr)), layers={"counts": csr_matrix(cp.asarray(X_clr))}, ) x_before = _to_np(adata.X).copy() + counts_before = _to_np(adata.layers["counts"]).copy() rsc.pp.normalize_clr(adata, layer="counts") + # both X and the source layer are untouched; result lands in layers["clr"] np.testing.assert_array_equal(_to_np(adata.X), x_before) + np.testing.assert_array_equal(_to_np(adata.layers["counts"]), counts_before) np.testing.assert_allclose( - _reconstruct_clr(adata.layers["counts"], adata.obsm["clr_residuals"]), + _clr_result(adata), _clr_reference(X_clr), rtol=1e-5, atol=1e-5, From 821d72760c0ecf6309c585e4beda32ac92df38fd Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 29 Jun 2026 18:45:48 +0200 Subject: [PATCH 3/3] add CLR PCA Signed-off-by: Intron7 --- src/rapids_singlecell/preprocessing/_pca.py | 35 +++++++++++- .../preprocessing/_sparse_pca/_operators.py | 56 +++++++++++++++---- .../preprocessing/_sparse_pca/_sparse_pca.py | 40 +++++++++++-- .../_sparse_pca/_sparse_svd_pca.py | 32 ++++++++++- tests/test_pca.py | 44 +++++++++++++++ 5 files changed, 185 insertions(+), 22 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 08f2a3233..13a6f3c68 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -15,6 +15,7 @@ from rapids_singlecell._compat import DaskArray from rapids_singlecell.get import _check_mask, _get_obs_rep +from ._normalize import CLR_LAYER, CLR_RESIDUALS_KEY from ._utils import _check_gpu_X if TYPE_CHECKING: @@ -265,6 +266,27 @@ def pca( del use_highly_variable X = X[:, mask_var] if mask_var is not None else X + # CLR (`normalize_clr`) stores a factored log1p(PF) in `layers["clr"]` plus a + # per-cell centering offset in `obsm`; fold it in so PCA sees the centered CLR. + obs_offset = None + if layer == CLR_LAYER and CLR_RESIDUALS_KEY in adata.obsm: + if not zero_center: + raise ValueError( + "The `clr` layer requires `zero_center=True`; truncated SVD " + "(`zero_center=False`) is not meaningful for centered log-ratio data." + ) + if isinstance(X, DaskArray): + raise NotImplementedError( + "PCA on a Dask `clr` layer is not implemented yet." + ) + offset = cp.asarray(adata.obsm[CLR_RESIDUALS_KEY]).reshape(-1) + if isinstance(X, cp.ndarray): + # Dense: collapse the per-cell offset in, then run standard PCA. + X = X - offset[:, None] + else: + # Sparse: pass it to the matrix-free composite path (no densify). + obs_offset = offset + pca_func, X_pca, n_comps = _pca_compute( X, n_comps, @@ -275,6 +297,7 @@ def pca( chunk_size=chunk_size, dtype=dtype, kwargs=kwargs, + obs_offset=obs_offset, ) key_obsm, key_varm, key_uns = ( @@ -312,6 +335,7 @@ def _pca_compute( chunk_size: int | None, dtype: str, kwargs: dict, + obs_offset: cp.ndarray | None = None, ): if n_comps is None: min_dim = min(X.shape[0], X.shape[1]) @@ -365,9 +389,12 @@ def _pca_compute( random_state=random_state, n_oversamples=kwargs.get("n_oversamples"), n_iter=kwargs.get("n_iter"), + offset=obs_offset, ) else: - pca_func, X_pca = _run_covariance_pca(X, n_comps, zero_center) + pca_func, X_pca = _run_covariance_pca( + X, n_comps, zero_center, offset=obs_offset + ) else: pca_func, X_pca = _run_cuml_pca(X, n_comps, svd_solver=svd_solver) @@ -401,7 +428,7 @@ def _as_numpy(X): return X -def _run_covariance_pca(X, n_comps, zero_center): +def _run_covariance_pca(X, n_comps, zero_center, *, offset=None): """Run PCA using covariance matrix eigendecomposition.""" if issparse(X): X = sparse_scipy_to_cp(X, dtype=X.dtype) @@ -412,7 +439,7 @@ def _run_covariance_pca(X, n_comps, zero_center): if issparse_cupy(X): X.sort_indices() - pca_func = PCA_sparse(n_components=n_comps, zero_center=zero_center) + pca_func = PCA_sparse(n_components=n_comps, zero_center=zero_center, offset=offset) X_pca = pca_func.fit_transform(X) return pca_func, X_pca @@ -426,6 +453,7 @@ def _run_sparse_svd_pca( random_state: int = 0, n_oversamples: int | None = None, n_iter: int | None = None, + offset=None, ): """Run PCA using SVD solvers (lanczos, randomized).""" if issparse(X): @@ -443,6 +471,7 @@ def _run_sparse_svd_pca( "svd_solver": svd_solver, "zero_center": zero_center, "random_state": random_state, + "offset": offset, } if n_oversamples is not None: kwargs["n_oversamples"] = n_oversamples diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_operators.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_operators.py index 9ef9d9bc6..cbc9a3748 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_operators.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_operators.py @@ -11,11 +11,16 @@ from cupyx.scipy.sparse.linalg import LinearOperator -def mean_centered_operator(X, mean: cp.ndarray) -> LinearOperator: +def mean_centered_operator( + X, mean: cp.ndarray, row_offset: cp.ndarray | None = None +) -> LinearOperator: """ - Create a linear operator for mean-centered sparse matrix. + Create a linear operator for a centered sparse matrix. - Computes products with (X - 1*mean.T) without forming the dense matrix. + Computes products with ``A = X - 1*mean.T - row_offset*1.T`` without forming + the dense matrix. The optional per-row offset (constant across features, e.g. + the CLR per-cell centering term) is a rank-1 correction along the all-ones + feature direction, handled exactly like the column-mean term. Parameters ---------- @@ -23,26 +28,53 @@ def mean_centered_operator(X, mean: cp.ndarray) -> LinearOperator: Sparse matrix in CSR format. mean Column means of shape (n_features,). + row_offset + Optional per-row offset of shape (n_samples,). If `None`, the operator is + plain mean-centering. Returns ------- LinearOperator - Operator that computes mean-centered matrix-vector products. + Operator that computes the centered matrix-vector products. """ n_samples, n_features = X.shape XT = X.T # CSC view - no copy - def matvec(v): - return X.dot(v) - cp.dot(mean, v) + if row_offset is None: - def rmatvec(v): - return XT.dot(v) - mean * cp.sum(v) + def matvec(v): + return X.dot(v) - cp.dot(mean, v) - def matmat(V): - return X.dot(V) - cp.dot(mean, V)[cp.newaxis, :] + def rmatvec(v): + return XT.dot(v) - mean * cp.sum(v) - def rmatmat(V): - return XT.dot(V) - cp.outer(mean, cp.sum(V, axis=0)) + def matmat(V): + return X.dot(V) - cp.dot(mean, V)[cp.newaxis, :] + + def rmatmat(V): + return XT.dot(V) - cp.outer(mean, cp.sum(V, axis=0)) + else: + r = row_offset + + def matvec(v): + return X.dot(v) - r * cp.sum(v) - cp.dot(mean, v) + + def rmatvec(v): + return XT.dot(v) - cp.dot(r, v) - mean * cp.sum(v) + + def matmat(V): + return ( + X.dot(V) + - cp.outer(r, cp.sum(V, axis=0)) + - cp.dot(mean, V)[cp.newaxis, :] + ) + + def rmatmat(V): + return ( + XT.dot(V) + - cp.dot(r, V)[cp.newaxis, :] + - cp.outer(mean, cp.sum(V, axis=0)) + ) return LinearOperator( shape=(n_samples, n_features), diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index 66527485e..58c948b93 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -19,9 +19,18 @@ class PCA_sparse: - def __init__(self, n_components: int | None, *, zero_center: bool = True) -> None: + def __init__( + self, + n_components: int | None, + *, + zero_center: bool = True, + offset: cp.ndarray | None = None, + ) -> None: self.n_components = n_components self.zero_center = zero_center + # Optional per-cell offset (CLR centering term); folded into the gram, + # mean and transform as a rank-1 composite, never densified. + self.offset_ = offset def fit(self, x: spmatrix | DaskArray) -> Self: if self.n_components is None: @@ -37,7 +46,7 @@ def fit(self, x: spmatrix | DaskArray) -> Self: self.dtype = x.dtype if self.zero_center: - covariance, self.mean_ = _cov_sparse(x) + covariance, self.mean_ = _cov_sparse(x, offset=self.offset_) else: # For truncated SVD (uncentered), operate on the Gram matrix (1/n * X^T X) # We don't subtract the mean in this path @@ -88,6 +97,10 @@ def _transform_cupy(self, X: spmatrix) -> spmatrix: (X.shape[0], 1), dtype=cp.float32 ) @ precomputed_mean_impact.reshape(1, -1) X_transformed = X.dot(self.components_.T) - mean_impact + if self.offset_ is not None: + # Per-cell offset impact: r ⊗ (V·1), the row-sum of each component. + v_sum = self.components_.sum(axis=1) + X_transformed -= self.offset_[:, None] * v_sum[None, :] else: # Uncentered projection for truncated SVD X_transformed = X.dot(self.components_.T) @@ -127,16 +140,22 @@ def fit_transform(self, X, y=None): @overload def _cov_sparse( - x: spmatrix | DaskArray, *, return_gram: Literal[False] = False + x: spmatrix | DaskArray, + *, + return_gram: Literal[False] = False, + offset: cp.ndarray | None = None, ) -> tuple[cp.ndarray, cp.ndarray]: ... @overload def _cov_sparse( - x: spmatrix | DaskArray, *, return_gram: Literal[True] + x: spmatrix | DaskArray, *, return_gram: Literal[True], offset: None = None ) -> cp.ndarray: ... def _cov_sparse( - x: spmatrix | DaskArray, *, return_gram: bool = False + x: spmatrix | DaskArray, + *, + return_gram: bool = False, + offset: cp.ndarray | None = None, ) -> cp.ndarray | tuple[cp.ndarray, cp.ndarray]: """ Computes the mean and the covariance of matrix X of @@ -192,6 +211,17 @@ def _cov_sparse( gram_matrix = _copy_gram(gram_matrix, x.shape[1]) mean_x = mean_x.astype(x.dtype) + if offset is not None: + # Correct the gram of M = X - r·1ᵀ on the symmetrized G: + # MᵀM = G - g·1ᵀ - 1·gᵀ + s·11ᵀ, g = Xᵀr, s = Σrᵢ², + # and shift the column mean by mean(r). No densification. + offset = offset.astype(x.dtype) + g = x.T.dot(offset) + s = float(offset @ offset) + gram_matrix -= g[:, None] + gram_matrix -= g[None, :] + gram_matrix += s + mean_x = mean_x - offset.mean().astype(x.dtype) gram_matrix *= 1 / x.shape[0] cov_result = gram_matrix diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_svd_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_svd_pca.py index 780f99c21..5e66e3aaf 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_svd_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_svd_pca.py @@ -65,6 +65,7 @@ def __init__( n_oversamples: int = 10, n_iter: int | None = None, random_state: int | None = 0, + offset: cp.ndarray | None = None, ) -> None: self.n_components = n_components self.svd_solver = svd_solver @@ -72,6 +73,9 @@ def __init__( self.n_oversamples = n_oversamples self.n_iter = n_iter self.random_state = random_state + # Optional per-cell offset (CLR centering term); applied via the operator + # and the transform as a rank-1 composite, never densified. + self.offset_ = offset def fit(self, X: spmatrix) -> Self: """ @@ -104,12 +108,16 @@ def fit(self, X: spmatrix) -> Self: if self.zero_center: self.mean_, _ = _get_mean_var(X, axis=0) self.mean_ = self.mean_.astype(X.dtype) + if self.offset_ is not None: + # Centering M = X - r·1ᵀ shifts every gene mean by mean(r). + self.offset_ = self.offset_.astype(X.dtype) + self.mean_ = self.mean_ - self.offset_.mean().astype(X.dtype) else: self.mean_ = None # Create operator (centered or raw) if self.zero_center: - X_op = mean_centered_operator(X, self.mean_) + X_op = mean_centered_operator(X, self.mean_, row_offset=self.offset_) else: X_op = X @@ -121,7 +129,9 @@ def fit(self, X: spmatrix) -> Self: self.explained_variance_ = (S**2) / (self.n_samples_ - 1) # Compute total variance for variance ratio - if self.zero_center: + if self.zero_center and self.offset_ is not None: + total_variance = self._total_variance_with_offset(X) + elif self.zero_center: _, var_x = _get_mean_var(X, axis=0) total_variance = cp.sum(var_x) else: @@ -134,6 +144,20 @@ def fit(self, X: spmatrix) -> Self: return self + def _total_variance_with_offset(self, X) -> cp.ndarray: + """Total variance of the centered M = X - r·1ᵀ, without densifying. + + ``‖M‖_F² = ‖X‖_F² - 2·Σ rᵢ·rowsumᵢ(X) + n_features·Σ rᵢ²`` and + ``total_var = (‖M‖_F² - n·‖mean_M‖²) / (n - 1)`` (``self.mean_`` already + carries the ``-mean(r)`` shift). + """ + r = self.offset_ + x_norm2 = cp.sum(X.data**2) if hasattr(X, "data") else cp.sum(X**2) + row_sums = cp.asarray(X.sum(axis=1)).ravel() + m_norm2 = x_norm2 - 2.0 * cp.dot(r, row_sums) + X.shape[1] * cp.dot(r, r) + n = self.n_samples_ + return (m_norm2 - n * cp.dot(self.mean_, self.mean_)) / (n - 1) + def _run_svd(self, X_op): """Run the selected SVD solver.""" if self.svd_solver == "lanczos": @@ -176,6 +200,10 @@ def transform(self, X: spmatrix) -> cp.ndarray: X_transformed = X.dot(self.components_.T) mean_projection = cp.dot(self.mean_, self.components_.T) X_transformed -= mean_projection + if self.offset_ is not None: + # Per-cell offset impact: r ⊗ (V·1), the row-sum of each component. + v_sum = self.components_.sum(axis=1) + X_transformed -= self.offset_[:, None] * v_sum[None, :] return X_transformed else: return X.dot(self.components_.T) diff --git a/tests/test_pca.py b/tests/test_pca.py index 40923ceef..b8573ded7 100644 --- a/tests/test_pca.py +++ b/tests/test_pca.py @@ -1021,3 +1021,47 @@ def test_svd_small_k(self, svd_func): Vt_np = Vt.get() np.testing.assert_allclose(U_np.T @ U_np, np.eye(k), atol=1e-6) np.testing.assert_allclose(Vt_np @ Vt_np.T, np.eye(k), atol=1e-6) + + +@pytest.mark.parametrize("svd_solver", ["covariance_eigh", "lanczos"]) +def test_pca_clr_sparse_matches_dense(svd_solver): + """PCA on the factored `clr` layer matches PCA on the densely-collapsed CLR. + + `normalize_clr` keeps a sparse `log1p(PF)` in `layers["clr"]` plus a per-cell + offset in `obsm`; PCA must fold the offset in via the gram/operator composites + (no densify) and reproduce the result of PCA on the explicit centered matrix. + """ + rng = np.random.default_rng(0) + counts = rng.poisson(2.0, size=(200, 60)).astype("float32") + n_comps = 10 + + # Dense reference: collapse the offset, then standard PCA. + ref = AnnData(cp.asarray(counts.copy())) + rsc.pp.normalize_clr(ref) + rsc.pp.pca(ref, layer="clr", n_comps=n_comps) + ref_vr = np.asarray(ref.uns["pca"]["variance_ratio"]) + ref_emb = cp.asnumpy(ref.obsm["X_pca"]) + + # Sparse factored path: offset applied as a composite, never densified. + adata = AnnData(cusparse.csr_matrix(cp.asarray(counts.copy()))) + rsc.pp.normalize_clr(adata) + rsc.pp.pca(adata, layer="clr", n_comps=n_comps, svd_solver=svd_solver) + + np.testing.assert_allclose(adata.uns["pca"]["variance_ratio"], ref_vr, atol=1e-4) + # components carry an arbitrary sign; compare magnitudes + emb = cp.asnumpy(adata.obsm["X_pca"]) + np.testing.assert_allclose(np.abs(emb), np.abs(ref_emb), atol=1e-2) + # the sparse `clr` layer is untouched (no densification in place) + assert cusparse.issparse(adata.layers["clr"]) + + +@pytest.mark.parametrize( + "array_type", [cp.asarray, lambda c: cusparse.csr_matrix(cp.asarray(c))] +) +def test_pca_clr_tsvd_raises(array_type): + """CLR is centered log-ratio data; truncated SVD (`zero_center=False`) is blocked.""" + counts = np.random.default_rng(0).poisson(2.0, size=(80, 20)).astype("float32") + adata = AnnData(array_type(counts.copy())) + rsc.pp.normalize_clr(adata) + with pytest.raises(ValueError, match="zero_center=True"): + rsc.pp.pca(adata, layer="clr", zero_center=False)