Speedup (~20x) of scanpy.pp.regress_out function using Linear Least Square method.

docs/release-notes/1.10.2.md

@@ -31,3 +31,4 @@
 * `pp.highly_variable_genes` with `flavor=seurat_v3` now uses a numba kernel {pr}`3017` {smaller}`S Dicks`
 * Speed up {func}`~scanpy.pp.scrublet` {pr}`3044` {smaller}`S Dicks` and {pr}`3056` {smaller}`P Angerer`
 * Speed up clipping of array in {func}`~scanpy.pp.scale` {pr}`3100` {smaller}`P Ashish & S Dicks`
+* Speed up {func}`~scanpy.pp.regress_out` {pr}`3110` {smaller}`P Ashish & P Angerer`

src/scanpy/preprocessing/_simple.py

@@ -7,7 +7,7 @@
 
 import warnings
 from functools import singledispatch
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 import numba
 import numpy as np
@@ -612,6 +612,53 @@ def normalize_per_cell(  # noqa: PLR0917
     return X if copy else None
 
 
+DT = TypeVar("DT")
+
+
+@numba.njit(cache=True, parallel=True)
+def to_dense(
+    shape: tuple[int, int],
+    indptr: NDArray[np.integer],
+    indices: NDArray[np.integer],
+    data: NDArray[DT],
+) -> NDArray[DT]:
+    """\
+    Numba kernel for np.toarray() function
+    """
+    X = np.empty(shape, dtype=data.dtype)
+
+    for r in numba.prange(shape[0]):
+        X[r] = 0
+        for i in range(indptr[r], indptr[r + 1]):
+            X[r, indices[i]] = data[i]
+    return X
+
+
+def numpy_regress_out(
+    data: np.ndarray,
+    regressor: np.ndarray,
+) -> np.ndarray:
+    """\
+    Numba kernel for regress out unwanted sorces of variantion.
+    Finding coefficient using Linear regression (Linear Least Squares).
+    """
+
+    @numba.njit(cache=True, parallel=True)
+    def get_resid(
+        data: np.ndarray,
+        regressor: np.ndarray,
+        coeff: np.ndarray,
+    ) -> np.ndarray:
+        for i in numba.prange(data.shape[0]):
+            data[i] -= regressor[i] @ coeff
+        return data
+
+    inv_gram_matrix = np.linalg.inv(regressor.T @ regressor)
+    coeff = inv_gram_matrix @ (regressor.T @ data)
+    data = get_resid(data, regressor, coeff)
+    return data
+
+
 @old_positionals("layer", "n_jobs", "copy")
 def regress_out(
     adata: AnnData,
@@ -664,7 +711,7 @@ def regress_out(
 
     if issparse(X):
         logg.info("    sparse input is densified and may " "lead to high memory use")
-        X = X.toarray()
+        X = to_dense(X.shape, X.indptr, X.indices, X.data)
 
     n_jobs = sett.n_jobs if n_jobs is None else n_jobs
 
@@ -715,14 +762,28 @@ def regress_out(
             regres = regressors
         tasks.append(tuple((data_chunk, regres, variable_is_categorical)))
 
-    from joblib import Parallel, delayed
+    res = None
+    if not variable_is_categorical:
+        A = regres.to_numpy()
+        # if det(A.T@A) != 0 we can take the inverse and regress using a fast method.
+        if np.linalg.det(A.T @ A) != 0:
+            res = numpy_regress_out(X, A)
+
+    # for a categorical variable or if the above checks failed,
+    # we fall back to the GLM implemetation of regression.
+    if variable_is_categorical or res is None:
+        from joblib import Parallel, delayed
+
+        # TODO: figure out how to test that this doesn't oversubscribe resources
+        res = Parallel(n_jobs=n_jobs)(
+            delayed(_regress_out_chunk)(task) for task in tasks
+        )
 
-    # TODO: figure out how to test that this doesn't oversubscribe resources
-    res = Parallel(n_jobs=n_jobs)(delayed(_regress_out_chunk)(task) for task in tasks)
+        # res is a list of vectors (each corresponding to a regressed gene column).
+        # The transpose is needed to get the matrix in the shape needed
+        res = np.vstack(res).T
 
-    # res is a list of vectors (each corresponding to a regressed gene column).
-    # The transpose is needed to get the matrix in the shape needed
-    _set_obs_rep(adata, np.vstack(res).T, layer=layer)
+    _set_obs_rep(adata, res, layer=layer)
     logg.info("    finished", time=start)
     return adata if copy else None