From fdcad6cbf715c63b45d5a0d93c7841ba178dffd8 Mon Sep 17 00:00:00 2001 From: ashish615 Date: Tue, 18 Jun 2024 06:16:29 +0000 Subject: [PATCH 1/4] Finding coefficient using Linear regression (Linear Least Squares) --- src/scanpy/preprocessing/_simple.py | 68 ++++++++++++++++++++++++++--- 1 file changed, 61 insertions(+), 7 deletions(-) diff --git a/src/scanpy/preprocessing/_simple.py b/src/scanpy/preprocessing/_simple.py index 1e61b60ace..911747b8e5 100644 --- a/src/scanpy/preprocessing/_simple.py +++ b/src/scanpy/preprocessing/_simple.py @@ -612,6 +612,47 @@ def normalize_per_cell( # noqa: PLR0917 return X if copy else None +@numba.njit(cache=True, parallel=True) +def to_dense(shape, indptr, indices, data): + """\ + 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( + X: np.ndarray, + A: 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( + X: np.ndarray, + A: np.ndarray, + coeff, + ) -> np.ndarray: + for i in numba.prange(X.shape[0]): + X[i] -= A[i] @ coeff + return X + + tmp0 = A.T @ A + tmp = np.linalg.inv(tmp0) + tmp0 = A.T @ X + coeff = tmp @ tmp0 + X = get_resid(X, A, coeff) + return X + + @old_positionals("layer", "n_jobs", "copy") def regress_out( adata: AnnData, @@ -664,7 +705,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 +756,27 @@ def regress_out( regres = regressors tasks.append(tuple((data_chunk, regres, variable_is_categorical))) - from joblib import Parallel, delayed + flag = None + if not variable_is_categorical: + A = regres.to_numpy() + # checking if inverse of A.T@A matix is zero or not. + if np.linalg.det(A.T @ A) != 0: + res = numpy_regress_out(X, A) + flag = 1 + # for categorical variable and failed the above code then run original code. + if variable_is_categorical or flag 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 From 3e3ca9dcb2e77e72b75095ff895cb55aeb7f98bc Mon Sep 17 00:00:00 2001 From: ashish615 Date: Wed, 19 Jun 2024 06:31:41 +0000 Subject: [PATCH 2/4] changes done as per suggestions --- docs/release-notes/1.10.2.md | 1 + src/scanpy/preprocessing/_simple.py | 49 +++++++++++++++++------------ 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/docs/release-notes/1.10.2.md b/docs/release-notes/1.10.2.md index d1d9212e8a..c7dc9720d5 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 911747b8e5..e21a10fe0a 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,8 +612,16 @@ 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, indptr, indices, data): +def to_dense( + shape: tuple[int, int], + indptr: NDArray[np.integer], + indices: NDArray[np.integer], + data: NDArray[np.number], +) -> NDArray[DT]: """\ Numba kernel for np.toarray() function """ @@ -627,8 +635,8 @@ def to_dense(shape, indptr, indices, data): def numpy_regress_out( - X: np.ndarray, - A: np.ndarray, + data: np.ndarray, + regressor: np.ndarray, ) -> np.ndarray: """\ Numba kernel for regress out unwanted sorces of variantion. @@ -637,20 +645,18 @@ def numpy_regress_out( @numba.njit(cache=True, parallel=True) def get_resid( - X: np.ndarray, - A: np.ndarray, - coeff, + data: np.ndarray, + regressor: np.ndarray, + coeff: np.ndarray, ) -> np.ndarray: - for i in numba.prange(X.shape[0]): - X[i] -= A[i] @ coeff - return X + for i in numba.prange(data.shape[0]): + data[i] -= regressor[i] @ coeff + return data - tmp0 = A.T @ A - tmp = np.linalg.inv(tmp0) - tmp0 = A.T @ X - coeff = tmp @ tmp0 - X = get_resid(X, A, coeff) - return X + 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") @@ -756,15 +762,18 @@ def regress_out( regres = regressors tasks.append(tuple((data_chunk, regres, variable_is_categorical))) - flag = None + res = None if not variable_is_categorical: A = regres.to_numpy() - # checking if inverse of A.T@A matix is zero or not. + # if det(A.T@A) zero, then can not take inverse and regression fails. + # if fails then fall back to GLM implemetation of regression. + # Numba kernel is used to speedup regression using Linear Least Square method + # on regressor of type numpy array. if np.linalg.det(A.T @ A) != 0: res = numpy_regress_out(X, A) - flag = 1 + # for categorical variable and failed the above code then run original code. - if variable_is_categorical or flag is None: + 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 From ebdf441333da6fe589c3f47c7e87c03f4d3890bd Mon Sep 17 00:00:00 2001 From: Ashish Patel Date: Mon, 24 Jun 2024 10:20:33 +0530 Subject: [PATCH 3/4] Apply suggestions from code review Co-authored-by: Philipp A. --- src/scanpy/preprocessing/_simple.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/scanpy/preprocessing/_simple.py b/src/scanpy/preprocessing/_simple.py index e21a10fe0a..131e169f1d 100644 --- a/src/scanpy/preprocessing/_simple.py +++ b/src/scanpy/preprocessing/_simple.py @@ -620,7 +620,7 @@ def to_dense( shape: tuple[int, int], indptr: NDArray[np.integer], indices: NDArray[np.integer], - data: NDArray[np.number], + data: NDArray[DT], ) -> NDArray[DT]: """\ Numba kernel for np.toarray() function From 5e4df8c8f24f3066e5d11974793f7dba33859f3b Mon Sep 17 00:00:00 2001 From: Ashish Patel Date: Mon, 24 Jun 2024 10:21:41 +0530 Subject: [PATCH 4/4] Update src/scanpy/preprocessing/_simple.py Co-authored-by: Philipp A. --- src/scanpy/preprocessing/_simple.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/scanpy/preprocessing/_simple.py b/src/scanpy/preprocessing/_simple.py index 131e169f1d..314382351c 100644 --- a/src/scanpy/preprocessing/_simple.py +++ b/src/scanpy/preprocessing/_simple.py @@ -765,14 +765,12 @@ def regress_out( res = None if not variable_is_categorical: A = regres.to_numpy() - # if det(A.T@A) zero, then can not take inverse and regression fails. - # if fails then fall back to GLM implemetation of regression. - # Numba kernel is used to speedup regression using Linear Least Square method - # on regressor of type numpy array. + # 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 categorical variable and failed the above code then run original code. + # 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