Speed up categorical regressor with numba

docs/release-notes/3353.performance.md

@@ -0,0 +1 @@
+* Speed up for a categorical regressor in {func}`~scanpy.pp.regress_out` {smaller}`S Dicks`

src/scanpy/preprocessing/_simple.py

@@ -628,6 +628,21 @@
 DT = TypeVar("DT")
 
 
+@njit
+def _create_regressor_categorical(
+    X: np.ndarray, cats: int, filters: np.ndarray
+) -> np.ndarray:
+    # create regressor matrix faster for categorical variables
+    regressors = np.zeros(X.shape, dtype=X.dtype)
+    XT = X.T
+    for category in range(cats):
+        mask = category == filters
+        for ix in numba.prange(XT.shape[0]):
+            x = XT[ix]
+            regressors[mask, ix] = x[mask].mean()
+    return regressors
+
+
 @njit
 def get_resid(
     data: np.ndarray,
@@ -722,13 +737,13 @@
                 "we regress on the mean for each category."
             )
         logg.debug("... regressing on per-gene means within categories")
-        regressors = np.zeros(X.shape, dtype="float32")
+        # Create numpy array's from categorical variable
+        cats = np.int64(len(adata.obs[keys[0]].cat.categories))
+        filters = adata.obs[keys[0]].cat.codes.to_numpy()
+        cats = cats.astype(filters.dtype)
+
         X = _to_dense(X, order="F") if issparse(X) else X
-        # TODO figure out if we should use a numba kernel for this
-        for category in adata.obs[keys[0]].cat.categories:
-            mask = (category == adata.obs[keys[0]]).values
-            for ix, x in enumerate(X.T):
-                regressors[mask, ix] = x[mask].mean()
+        regressors = _create_regressor_categorical(X, cats, filters)
         variable_is_categorical = True
     # regress on one or several ordinal variables
     else:

tests/test_preprocessing.py

@@ -13,6 +13,8 @@
 from scipy.sparse import coo_matrix, csc_matrix, csr_matrix, issparse
 
 import scanpy as sc
+from scanpy.preprocessing._simple import _create_regressor_categorical
+from scanpy.preprocessing._utils import _to_dense
 from testing.scanpy._helpers import (
     anndata_v0_8_constructor_compat,
     check_rep_mutation,
@@ -327,14 +329,51 @@ def test_regress_out_constants():
     assert_equal(adata, adata_copy)
 
 
-def test_regress_out_reproducible():
-    adata = pbmc68k_reduced()
+@pytest.mark.parametrize(
+    ("keys", "tester_file"),
+    [
+        (["n_counts", "percent_mito"], "regress_test_small.npy"),
+        (["bulk_labels"], "regress_test_small_cat.npy"),
+    ],
+)
+def test_regress_out_reproducible(keys, tester_file):
+    adata = sc.datasets.pbmc68k_reduced()
     adata = adata.raw.to_adata()[:200, :200].copy()
-    sc.pp.regress_out(adata, keys=["n_counts", "percent_mito"])
+    sc.pp.regress_out(adata, keys=keys)
     # This file was generated from the original implementation in version 1.10.3
-    # Now we compare new implementation with the old one
-    tester = np.load(DATA_PATH / "regress_test_small.npy")
-    np.testing.assert_allclose(adata.X, tester)
+    # Now we compare the new implementation with the old one
+    tester = np.load(DATA_PATH / tester_file)
+    np.testing.assert_array_almost_equal(adata.X, tester)
+
+
+def _gen_org_regressors(adata, keys, X_org):
+    # helper function to generate the original regressors
+    regressors = np.zeros(X_org.shape, dtype=X_org.dtype)
+    X = _to_dense(X_org, order="F")
+    for category in adata.obs[keys[0]].cat.categories:
+        mask = (category == adata.obs[keys[0]]).values
+        for ix, x in enumerate(X.T):
+            regressors[mask, ix] = x[mask].mean()
+    return regressors
+
+
+def test_regressor_categorical():
+    adata = sc.datasets.pbmc68k_reduced()
+    adata = adata.raw.to_adata()[:200, :200]
+    X_org = adata.X.copy().astype(np.float64)
+    keys = ["bulk_labels"]
+    # Create org regressors
+    regressors = _gen_org_regressors(adata, keys, X_org)
+
+    # Create new regressors
+    cats = np.int64(len(adata.obs[keys[0]].cat.categories))
+    filters = adata.obs[keys[0]].cat.codes.to_numpy()
+    cats = cats.astype(filters.dtype)
+    X = _to_dense(X_org, order="F")
+    new_reg = _create_regressor_categorical(X, cats, filters)
+
+    # Compare the two implementations
+    np.testing.assert_allclose(new_reg, regressors)
 
 
 def test_regress_out_constants_equivalent():