updates sparse scale

docs/release-notes/1.10.1.md

@@ -0,0 +1,5 @@
+### 1.10.1 {small}`the future`
+
+```{rubric} Features
+```
+* {func}`~scanpy.pp.scale` now uses numba kernels for `sparse.csr_matrix` and `sparse.csc_matrix` when `zero_center==False` and `mask_obs` is provided. This greatly speed up execution {pr}`2942` {smaller}`S Dicks`

docs/release-notes/index.md

@@ -9,6 +9,9 @@
 
 ## Version 1.10
 
+```{include} /release-notes/1.10.1.md
+```
+
 ```{include} /release-notes/1.10.0.md
 ```
 

scanpy/preprocessing/_simple.py

@@ -14,7 +14,7 @@
 import scipy as sp
 from anndata import AnnData
 from pandas.api.types import CategoricalDtype
-from scipy.sparse import csr_matrix, issparse, isspmatrix_csr, spmatrix
+from scipy.sparse import csr_matrix, issparse, isspmatrix_csc, isspmatrix_csr, spmatrix
 from sklearn.utils import check_array, sparsefuncs
 
 from .. import logging as logg
@@ -29,7 +29,7 @@
 )
 from ..get import _check_mask, _get_obs_rep, _set_obs_rep
 from ._distributed import materialize_as_ndarray
-from ._utils import _get_mean_var
+from ._utils import _get_mean_var, update_spmatrix_inplace
 
 # install dask if available
 try:
@@ -788,6 +788,7 @@ def scale(
         Restrict both the derivation of scaling parameters and the scaling itself
         to a certain set of observations. The mask is specified as a boolean array
         or a string referring to an array in :attr:`~anndata.AnnData.obs`.
+        This will transform data from csc to csr format if `issparse(data)`.
 
     Returns
     -------
@@ -830,6 +831,8 @@ def scale_array(
         X = X.copy()
     if mask_obs is not None:
         mask_obs = _check_mask(X, mask_obs, "obs")
+        if isspmatrix_csc(X):
+            X = X.tocsr()
         scale_rv = scale_array(
             X[mask_obs, :],
             zero_center=zero_center,
@@ -838,12 +841,21 @@ def scale_array(
             return_mean_std=return_mean_std,
             mask_obs=None,
         )
-        if return_mean_std:
-            X[mask_obs, :], mean, std = scale_rv
-            return X, mean, std
-        else:
-            X[mask_obs, :] = scale_rv
-            return X
+        if isinstance(X, np.ndarray):
+            if return_mean_std:
+                X[mask_obs, :], mean, std = scale_rv
+                return X, mean, std
+            else:
+                X[mask_obs, :] = scale_rv
+                return X
+        elif issparse(X):
+            if return_mean_std:
+                scaled, mean, std = scale_rv
+                update_spmatrix_inplace(X, scaled, mask_obs)
+                return X, mean, std
+            else:
+                update_spmatrix_inplace(X, scale_rv, mask_obs)
+                return X
 
     if not zero_center and max_value is not None:
         logg.info(  # Be careful of what? This should be more specific

scanpy/preprocessing/_utils.py

@@ -138,3 +138,42 @@
         np.prod(dims), nsamp, random_state=random_state, method=method
     )
     return np.vstack(np.unravel_index(idx, dims)).T
+
+
+def update_spmatrix_inplace(X, update, mask):
+    """
+    Update the values in a sparse matrix inplace.
+
+    Parameters
+    ----------
+    X
+        The sparse matrix to update.
+    update
+        The values to update with.
+    mask
+        The mask of the values to update.
+    """
+    subset_mask = np.where(mask)[0]
+
+    _update_csr_inplace(
+        X.indptr,
+        X.data,
+        update_indptr=update.indptr,
+        update_data=update.data,
+        mask=subset_mask,
+    )
+
+
+@numba.njit()
+def _update_csr_inplace(indptr, data, *, update_indptr, update_data, mask):
+    for i in range(len(update_indptr) - 1):
+        sub_start_idx = update_indptr[i]
+        sub_stop_idx = update_indptr[i + 1]
+        subidx = mask[i]
+
+        start_idx = indptr[subidx]
+        stop_idx = indptr[subidx + 1]
+
+        if sub_stop_idx - sub_start_idx == stop_idx - start_idx:
+            for j in range(sub_stop_idx - sub_start_idx):
+                data[start_idx + j] = update_data[sub_start_idx + j]

scanpy/tests/test_scaling.py

@@ -3,9 +3,10 @@
 import numpy as np
 import pytest
 from anndata import AnnData
-from scipy.sparse import csr_matrix
+from scipy.sparse import csc_matrix, csr_matrix
 
 import scanpy as sc
+from scanpy.preprocessing._utils import _update_csr_inplace, update_spmatrix_inplace
 
 # test "data" for 3 cells * 4 genes
 X_original = [
@@ -53,7 +54,9 @@
 ]
 
 
-@pytest.mark.parametrize("typ", [np.array, csr_matrix], ids=lambda x: x.__name__)
+@pytest.mark.parametrize(
+    "typ", [np.array, csr_matrix, csc_matrix], ids=lambda x: x.__name__
+)
 @pytest.mark.parametrize("dtype", ["float32", "int64"])
 @pytest.mark.parametrize(
     ("mask_obs", "X", "X_centered", "X_scaled"),
@@ -113,3 +116,23 @@ def test_clip(zero_center):
     if zero_center:
         assert adata.X.min() >= -1
     assert adata.X.max() <= 1
+
+
+def test_inplace_sparse():
+    org = csr_matrix(np.array([[2, 0, 2], [2, 0, 2], [2, 0, 2], [2, 0, 2]]))
+    sub = csr_matrix(np.array([[1, 0, 1], [1, 0, 1]]))
+    res = csr_matrix(np.array([[2, 0, 2], [1, 0, 1], [1, 0, 1], [2, 0, 2]]))
+    mask = np.array([False, True, True, False])
+    update_spmatrix_inplace(org, sub, mask)
+    np.testing.assert_equal(org.toarray(), res.toarray())
+
+
+def test_inplace_sparse_kernel():
+    org = csr_matrix(np.array([[2, 0, 2], [2, 0, 2], [2, 0, 2], [2, 0, 2]]))
+    sub = csr_matrix(np.array([[1, 0, 1], [1, 0, 1]]))
+    res = csr_matrix(np.array([[2, 0, 2], [1, 0, 1], [1, 0, 1], [2, 0, 2]]))
+    mask = np.array([1, 2])
+    _update_csr_inplace(
+        org.indptr, org.data, update_indptr=sub.indptr, update_data=sub.data, mask=mask
+    )
+    np.testing.assert_equal(org.toarray(), res.toarray())