Add clr pca

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

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},
+}

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`

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
 ```
 

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

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,13 +17,18 @@
     _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
 
     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,
@@ -124,6 +129,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 +251,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 +303,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 +359,160 @@ 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 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.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 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
+        materialized, not the matrix. Passing `target_sum` explicitly keeps it lazy.
 
-        def __sum(X_part):
-            return X_part.sum(axis=1)
+    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 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
+            Whether to return a copy or update `adata`. Not compatible with
+            `inplace=False`.
+
+    Returns
+    -------
+        Depending on `inplace`:
+
+        - `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.
+    """
+    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)
+    # 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:
-        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
+        X = X.copy()
+    X, cell_depths, residuals = _normalize_clr(X, target_sum=target_sum, alpha=alpha)
+    if inplace:
+        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:
+        return X, cell_depths, residuals
+    return None
+
+
+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:
+        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: