diff --git a/docs/release-notes/1.10.2.md b/docs/release-notes/1.10.2.md index d1d9212e8..c7dc9720d 100644 --- a/docs/release-notes/1.10.2.md +++ b/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` diff --git a/src/scanpy/preprocessing/_simple.py b/src/scanpy/preprocessing/_simple.py index 1e61b60ac..314382351 100644 --- a/src/scanpy/preprocessing/_simple.py +++ b/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