From 0a848254fe99e6e734cfd7fcc37c8d2b87e346bc Mon Sep 17 00:00:00 2001 From: SID Date: Sun, 14 Jun 2026 03:43:24 -0400 Subject: [PATCH] perf(score_genes): avoid sparse copies in _sparse_nanmean --- src/scanpy/tools/_score_genes.py | 39 +++++++++++++++++++------------- tests/test_score_genes.py | 6 ++++- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/scanpy/tools/_score_genes.py b/src/scanpy/tools/_score_genes.py index 6ba7f51a4c..9e459d442f 100644 --- a/src/scanpy/tools/_score_genes.py +++ b/src/scanpy/tools/_score_genes.py @@ -34,23 +34,30 @@ def _sparse_nanmean(x: CSBase, /, axis: Literal[0, 1]) -> NDArray[np.float64]: if not isinstance(x, CSBase): msg = "X must be a compressed sparse matrix" raise TypeError(msg) + if axis not in (0, 1): + msg = "axis must be 0 or 1" + raise ValueError(msg) - # count the number of nan elements per row/column (dep. on axis) - z = x.copy() - z.data = np.isnan(z.data) - z.eliminate_zeros() - n_elements = z.shape[axis] - z.sum(axis) - - # set the nans to 0, so that a normal .sum() works - y = x.copy() - y.data[np.isnan(y.data)] = 0 - y.eliminate_zeros() - - # the average - s = y.sum(axis, dtype="float64") # float64 for score_genes function compatibility) - m = s / n_elements - - return m + # Work in the compressed format aligned with the reduction axis and aggregate + # directly from index pointers to avoid matrix copies and eliminate_zeros(). + mat = x.tocsc(copy=False) if axis == 0 else x.tocsr(copy=False) + segment_lengths = np.diff(mat.indptr) + out_size = mat.shape[1] if axis == 0 else mat.shape[0] + full_length = mat.shape[0] if axis == 0 else mat.shape[1] + + segment_ids = np.repeat(np.arange(out_size), segment_lengths) + isnan = np.isnan(mat.data) + + sums = np.bincount( + segment_ids[~isnan], + weights=mat.data[~isnan], + minlength=out_size, + ).astype(np.float64, copy=False) + nan_counts = np.bincount(segment_ids[isnan], minlength=out_size) + counts = full_length - nan_counts + + with np.errstate(invalid="ignore", divide="ignore"): + return sums / counts @_doc_params(rng=doc_rng) diff --git a/tests/test_score_genes.py b/tests/test_score_genes.py index addd384f18..a564138a29 100644 --- a/tests/test_score_genes.py +++ b/tests/test_score_genes.py @@ -106,6 +106,7 @@ def test_add_score(): @pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("matrix_format", ["csr", "csc"]) @pytest.mark.parametrize( "mk_arr", [ @@ -130,7 +131,9 @@ def test_add_score(): ], ) def test_sparse_nanmean( - mk_arr: Callable[[], CSBase | np.ndarray], axis: Literal[0, 1] + mk_arr: Callable[[], CSBase | np.ndarray], + axis: Literal[0, 1], + matrix_format: Literal["csr", "csc"], ) -> None: """Check that _sparse_nanmean() is equivalent to np.nanmean().""" from scanpy.tools._score_genes import _sparse_nanmean @@ -138,6 +141,7 @@ def test_sparse_nanmean( arr_or_mat = mk_arr() arr = conv.to_dense(arr_or_mat) mat = sparse.csr_matrix(arr) if not isinstance(arr, CSBase) else arr # noqa: TID251 + mat = mat.asformat(matrix_format) np.testing.assert_allclose( np.nanmean(arr, axis), np.array(_sparse_nanmean(mat, axis)).flatten() )