scale function(_get_mean_var) updated for dense array, speedup upto ~4.65x

docs/release-notes/1.10.3.md

@@ -1,16 +1,14 @@
-(v1.10.3)=
-### 1.10.3 {small}`2024-09-17`
+### 1.10.3 {small}`the future`
 
-#### Bug fixes
+```{rubric} Development features
+```
 
-- Prevent empty control gene set in {func}`~scanpy.tl.score_genes` {smaller}`M Müller` ({pr}`2875`)
-- Fix `subset=True` of {func}`~scanpy.pp.highly_variable_genes` when `flavor` is `seurat` or `cell_ranger`, and `batch_key!=None` {smaller}`E Roellin` ({pr}`3042`)
-- Add compatibility with {mod}`numpy` 2.0 {smaller}`P Angerer` {pr}`3065` and ({pr}`3115`)
-- Fix `legend_loc` argument in {func}`scanpy.pl.embedding` not accepting matplotlib parameters {smaller}`P Angerer` ({pr}`3163`)
-- Fix dispersion cutoff in {func}`~scanpy.pp.highly_variable_genes` in presence of `NaN`s {smaller}`P Angerer` ({pr}`3176`)
-- Fix axis labeling for swapped axes in {func}`~scanpy.pl.rank_genes_groups_stacked_violin` {smaller}`Ilan Gold` ({pr}`3196`)
-- Upper bound dask on account of {issue}`scverse/anndata#1579` {smaller}`Ilan Gold` ({pr}`3217`)
-- The [fa2-modified][] package replaces [forceatlas2][] for the latter’s lack of maintenance {smaller}`A Alam` ({pr}`3220`)
+```{rubric} Docs
+```
 
-  [fa2-modified]: https://github.com/AminAlam/fa2_modified
-  [forceatlas2]: https://github.com/bhargavchippada/forceatlas2
+```{rubric} Bug fixes
+```
+
+```{rubric} Performance
+```
+* Speed up _get_mean_var used in {func}`~scanpy.pp.scale` {pr}`3099` {smaller}`P Ashish & S Dicks`

src/scanpy/preprocessing/_utils.py

@@ -33,17 +33,53 @@
 def _get_mean_var(
     X: _SupportedArray, *, axis: Literal[0, 1] = 0
 ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
-    if isinstance(X, sparse.spmatrix):
-        mean, var = sparse_mean_variance_axis(X, axis=axis)
+    if isinstance(X, np.ndarray):
+        n_threads = numba.get_num_threads()
+        mean, var = _compute_mean_var(X, axis=axis, n_threads=n_threads)
     else:
         mean = axis_mean(X, axis=axis, dtype=np.float64)
         mean_sq = axis_mean(elem_mul(X, X), axis=axis, dtype=np.float64)
         var = mean_sq - mean**2
     # enforce R convention (unbiased estimator) for variance
     if X.shape[axis] != 1:
         var *= X.shape[axis] / (X.shape[axis] - 1)
+
     return mean, var
 
+@numba.njit(cache=True, parallel=True)
+def _compute_mean_var(
+    X: _SupportedArray, axis: Literal[0, 1] = 0, n_threads=1
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    if axis == 0:
+        axis_i = 1
+        sums = np.zeros((n_threads, X.shape[axis_i]), dtype=np.float64)
+        sums_squared = np.zeros((n_threads, X.shape[axis_i]), dtype=np.float64)
+        mean = np.zeros(X.shape[axis_i], dtype=np.float64)
+        var = np.zeros(X.shape[axis_i], dtype=np.float64)
+        n = X.shape[axis]
+        for i in numba.prange(n_threads):
+            for r in range(i, n, n_threads):
+                for c in range(X.shape[axis_i]):
+                    value = X[r, c]
+                    sums[i, c] += value
+                    sums_squared[i, c] += value * value
+        for c in numba.prange(X.shape[axis_i]):
+            sum_ = sums[:, c].sum()
+            mean[c] = sum_ / n
+            var[c] = (sums_squared[:, c].sum() - sum_ * sum_ / n) / (n - 1)
+    else:
+        axis_i = 0
+        mean = np.zeros(X.shape[axis_i], dtype=np.float64)
+        var = np.zeros(X.shape[axis_i], dtype=np.float64)
+        for r in numba.prange(X.shape[0]):
+            for c in range(X.shape[1]):
+                value = X[r, c]
+                mean[r] += value
+                var[r] += value * value
+        for c in numba.prange(X.shape[0]):
+            mean[c] = mean[c] / X.shape[1]
+            var[c] = (var[c] - mean[c] ** 2) / (X.shape[1] - 1)
+
 
 def sparse_mean_variance_axis(mtx: sparse.spmatrix, axis: int):
     """
@@ -158,4 +194,4 @@
     idx = sample_without_replacement(
         np.prod(dims), nsamp, random_state=random_state, method=method
     )
-    return np.vstack(np.unravel_index(idx, dims)).T
+    return np.vstack(np.unravel_index(idx, dims)).T