Create cat regressor

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,20 @@
 DT = TypeVar("DT")
 
 
+@njit
+def _create_regressor_categorical(
+    X: np.ndarray, number_categories: 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(number_categories):
+        mask = category == filters
+        for ix in numba.prange(XT.shape[0]):
+            regressors[mask, ix] = XT[ix][mask].mean()
+    return regressors
+
+
 @njit
 def get_resid(
     data: np.ndarray,
@@ -722,13 +736,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
+        number_categories = np.int64(len(adata.obs[keys[0]].cat.categories))
+        filters = adata.obs[keys[0]].cat.codes.to_numpy()
+        number_categories = number_categories.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, number_categories, filters)
         variable_is_categorical = True
     # regress on one or several ordinal variables
     else:

tests/test_preprocessing.py

@@ -337,6 +337,16 @@ def test_regress_out_reproducible():
     np.testing.assert_allclose(adata.X, tester)
 
 
+def test_regress_out_reproducible_category():
+    adata = sc.datasets.pbmc68k_reduced()
+    adata = adata.raw.to_adata()[:200, :200].copy()
+    sc.pp.regress_out(adata, keys=["bulk_labels"])
+    # 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_cat.npy")
+    np.testing.assert_array_almost_equal(adata.X, tester)
+
+
 def test_regress_out_constants_equivalent():
     # Tests that constant values don't change results
     # (since support for constant values is implemented by us)