perf: welford's algorithm for mean-var aggregation

benchmarks/benchmarks/preprocessing_counts.py

@@ -151,16 +151,21 @@ def peakmem_log1p(self, *_) -> None:
 
 
 class Agg:  # noqa: D101
-    params: tuple[AggType] = tuple(get_literal_vals(AggType))
-    param_names = ("agg_name",)
+    params: tuple[list[AggType], list[bool]] = (
+        list(get_literal_vals(AggType)),
+        [True, False],
+    )
+    param_names = ("agg_name", "use_csc")
 
     def setup_cache(self) -> None:
         """Without this caching, asv was running several processes which meant the data was repeatedly downloaded."""
         adata, _ = get_dataset("lung93k")
         adata.write_h5ad("lung93k.h5ad")
 
-    def setup(self, agg_name: AggType) -> None:
+    def setup(self, agg_name: AggType, use_csc: bool) -> None:  # noqa: FBT001
         self.adata = ad.read_h5ad("lung93k.h5ad")
+        if use_csc:
+            self.adata.layers["counts"] = self.adata.layers["counts"].tocsc()
         self.agg_name = agg_name
 
     def time_agg(self, *_) -> None:

docs/release-notes/4147.perf.md

@@ -0,0 +1,3 @@
+Use [Welford's algorithm][] for mean-var calculation in {func}`scanpy.get.aggregate` for in-memory (i.e., non-dask) arrays {smaller}`I Gold`
+
+[Welford's algorithm]: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm

src/scanpy/get/_aggregated.py

@@ -13,7 +13,13 @@
 from scanpy._compat import CSBase, CSRBase, DaskArray
 
 from .._utils import _resolve_axis, get_literal_vals
-from ._kernels import agg_sum_csc, agg_sum_csr, mean_var_csc, mean_var_csr
+from ._kernels import (
+    agg_sum_csc,
+    agg_sum_csr,
+    mean_var_csc,
+    mean_var_csr,
+    mean_var_dense,
+)
 from .get import _check_mask
 
 if TYPE_CHECKING:
@@ -117,11 +123,9 @@ def mean(self) -> Array:
     def mean_var(self, dof: int = 1) -> tuple[np.ndarray, np.ndarray]:
         """Compute the count, as well as mean and variance per feature, per group of observations.
 
-        The formula `Var(X) = E(X^2) - E(X)^2` suffers loss of precision when the variance is a
-        very small fraction of the squared mean. In particular, when X is constant, the formula may
-        nonetheless be non-zero. By default, our implementation resets the variance to exactly zero
-        when the computed variance, relative to the squared mean, nears limit of precision of the
-        floating-point significand.
+        Mean and variance are computed with Welford's online algorithm, which is
+        numerically stable for constant or near-constant inputs
+        compared to subtracting E[X^2] - E[X]^2 since both values will be so close.
 
         Params
         ------
@@ -137,21 +141,11 @@ def mean_var(self, dof: int = 1) -> tuple[np.ndarray, np.ndarray]:
 
         group_counts = np.bincount(self.groupby.codes)
         if isinstance(self.data, np.ndarray):
-            mean_ = self.mean()
-            # sparse matrices do not support ** for elementwise power.
-            mean_sq = self._sum(_power(self.data, 2)) / group_counts[:, None]
-            sq_mean = mean_**2
-            var_ = mean_sq - sq_mean
+            mean_, var_ = mean_var_dense(self.indicator_matrix.tocsr(), self.data)
         else:
             mean_, var_ = (
                 mean_var_csr if isinstance(self.data, CSRBase) else mean_var_csc
             )(self.indicator_matrix, self.data)
-            sq_mean = mean_**2
-        # TODO: Why these values exactly? Because they are high relative to the datatype?
-        # (unchanged from original code: https://github.com/scverse/anndata/pull/564)
-        precision = 2 << (42 if self.data.dtype == np.float64 else 20)
-        # detects loss of precision in mean_sq - sq_mean, which suggests variance is 0
-        var_[precision * var_ < sq_mean] = 0
         if dof != 0:
             var_ *= (group_counts / (group_counts - dof))[:, np.newaxis]
         return mean_, var_

src/scanpy/get/_kernels.py

@@ -48,69 +48,136 @@ def agg_sum_csc(indicator: CSRBase, data: CSCBase, out: np.ndarray) -> None:
                 out[cat, col] += data.data[j]
 
 
+@njit
+def mean_var_dense(
+    indicator: CSRBase, data: NDArray
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    # Welford's online algorithm, parallelized over categories. The indicator
+    # CSR lists which observations belong to each category, allowing mask
+    # handling to be folded in naturally.
+    n_cats = indicator.shape[0]
+    n_features = data.shape[1]
+    mean = np.zeros((n_cats, n_features), dtype="float64")
+    var = np.zeros((n_cats, n_features), dtype="float64")
+
+    for cat in numba.prange(n_cats):
+        start = indicator.indptr[cat]
+        stop = indicator.indptr[cat + 1]
+        n = 0
+        for row_num in range(start, stop):
+            obs = indicator.indices[row_num]
+            n += 1
+            for col in range(n_features):
+                value = np.float64(data[obs, col])
+                delta = value - mean[cat, col]
+                mean[cat, col] += delta / n
+                delta2 = value - mean[cat, col]
+                var[cat, col] += delta * delta2
+        if n > 0:
+            for col in range(n_features):
+                var[cat, col] /= n
+    return mean, var
+
+
 @njit
 def mean_var_csr(
     indicator: CSRBase,
     data: CSCBase,
 ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
-    mean = np.zeros((indicator.shape[0], data.shape[1]), dtype="float64")
-    var = np.zeros((indicator.shape[0], data.shape[1]), dtype="float64")
-
-    for cat_num in numba.prange(indicator.shape[0]):
+    # Welford's online algorithm over nonzeros, then merge with the block of
+    # implicit zeros per (category, feature). Merging a Welford accumulator
+    # (n_A, mean_A, M2_A) with k zeros gives:
+    #   mean   = mean_A * n_A / (n_A + k)
+    #   M2_new = M2_A + mean_A^2 * n_A * k / (n_A + k)
+    n_cats = indicator.shape[0]
+    n_features = data.shape[1]
+    mean = np.zeros((n_cats, n_features), dtype="float64")
+    var = np.zeros((n_cats, n_features), dtype="float64")
+
+    for cat_num in numba.prange(n_cats):
         start_cat_idx = indicator.indptr[cat_num]
         stop_cat_idx = indicator.indptr[cat_num + 1]
+        n_obs = stop_cat_idx - start_cat_idx
+        if n_obs == 0:
+            continue
+
+        n_nonzero = np.zeros(n_features, dtype=np.int64)
+
         for row_num in range(start_cat_idx, stop_cat_idx):
             obs_per_cat = indicator.indices[row_num]
-
             start_obs = data.indptr[obs_per_cat]
             end_obs = data.indptr[obs_per_cat + 1]
 
             for j in range(start_obs, end_obs):
                 col = data.indices[j]
                 value = np.float64(data.data[j])
-                value = data.data[j]
-                mean[cat_num, col] += value
-                var[cat_num, col] += value * value
-
-        n_obs = stop_cat_idx - start_cat_idx
-        mean_cat = mean[cat_num, :] / n_obs
-        mean[cat_num, :] = mean_cat
-        var[cat_num, :] = (var[cat_num, :] / n_obs) - (mean_cat * mean_cat)
+                n = n_nonzero[col] + 1
+                n_nonzero[col] = n
+                m = mean[cat_num, col]
+                delta = value - m
+                m += delta / n
+                mean[cat_num, col] = m
+                var[cat_num, col] += delta * (value - m)
+
+        for col in range(n_features):
+            n_nz = n_nonzero[col]
+            k = n_obs - n_nz
+            if k > 0 and n_nz > 0:
+                mean_a = mean[cat_num, col]
+                mean[cat_num, col] = mean_a * n_nz / n_obs
+                var[cat_num, col] += mean_a * mean_a * n_nz * k / n_obs
+            var[cat_num, col] /= n_obs
     return mean, var
 
 
 @njit
 def mean_var_csc(
     indicator: CSRBase, data: CSCBase
 ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    # Welford's online algorithm, parallelized over columns. For each column
+    # we accumulate per-category over the explicit nonzeros, then merge each
+    # category's accumulator with its block of implicit zeros (see merge
+    # formula in `mean_var_csr`).
+    n_cats = indicator.shape[0]
+    n_features = data.shape[1]
     obs_to_cat = np.full(data.shape[0], -1, dtype=np.int64)
-
-    mean = np.zeros((indicator.shape[0], data.shape[1]), dtype="float64")
-    var = np.zeros((indicator.shape[0], data.shape[1]), dtype="float64")
-
-    for cat in range(indicator.shape[0]):
+    n_obs_per_cat = np.zeros(n_cats, dtype=np.int64)
+    for cat in range(n_cats):
+        n_obs_per_cat[cat] = indicator.indptr[cat + 1] - indicator.indptr[cat]
         for k in range(indicator.indptr[cat], indicator.indptr[cat + 1]):
             obs_to_cat[indicator.indices[k]] = cat
 
-    for col in numba.prange(data.shape[1]):
+    mean = np.zeros((n_cats, n_features), dtype="float64")
+    var = np.zeros((n_cats, n_features), dtype="float64")
+
+    for col in numba.prange(n_features):
+        n_nonzero = np.zeros(n_cats, dtype=np.int64)
         start = data.indptr[col]
         end = data.indptr[col + 1]
 
         for j in range(start, end):
             obs = data.indices[j]
             cat = obs_to_cat[obs]
-
-            if cat != -1:
-                value = np.float64(data.data[j])
-                value = data.data[j]
-                mean[cat, col] += value
-                var[cat, col] += value * value
-
-    for cat_num in numba.prange(indicator.shape[0]):
-        start_cat_idx = indicator.indptr[cat_num]
-        stop_cat_idx = indicator.indptr[cat_num + 1]
-        n_obs = stop_cat_idx - start_cat_idx
-        mean_cat = mean[cat_num, :] / n_obs
-        mean[cat_num, :] = mean_cat
-        var[cat_num, :] = (var[cat_num, :] / n_obs) - (mean_cat * mean_cat)
+            if cat == -1:
+                continue
+            value = np.float64(data.data[j])
+            n = n_nonzero[cat] + 1
+            n_nonzero[cat] = n
+            m = mean[cat, col]
+            delta = value - m
+            m += delta / n
+            mean[cat, col] = m
+            var[cat, col] += delta * (value - m)
+
+        for cat in range(n_cats):
+            n_obs = n_obs_per_cat[cat]
+            if n_obs == 0:
+                continue
+            n_nz = n_nonzero[cat]
+            k = n_obs - n_nz
+            if k > 0 and n_nz > 0:
+                mean_a = mean[cat, col]
+                mean[cat, col] = mean_a * n_nz / n_obs
+                var[cat, col] += mean_a * mean_a * n_nz * k / n_obs
+            var[cat, col] /= n_obs
     return mean, var

tests/test_aggregated.py

@@ -16,6 +16,7 @@
 from testing.scanpy._helpers.data import pbmc3k_processed
 from testing.scanpy._pytest.marks import needs
 from testing.scanpy._pytest.params import ARRAY_TYPES as ARRAY_TYPES_ALL
+from testing.scanpy._pytest.params import ARRAY_TYPES_MEM
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -544,3 +545,49 @@ def test_nan() -> None:
         "s2_control_C",
     ]
     assert adata_agg.obs["n_obs_aggregated"].tolist() == [1, 2, 1]
+
+
+@pytest.mark.parametrize("array_type", ARRAY_TYPES_MEM)
+def test_var_no_catastrophic_cancellation(array_type) -> None:
+    # Values of the form `offset + tiny_noise` make the textbook two-pass
+    # formula sum(x**2)/n - (sum(x)/n)**2 lose ~all precision: both terms are
+    # ~n*offset**2 ≈ 1e19 in float64 (precision ~1e3) but their difference is
+    # the variance ~1e-3, far below the rounding noise. Welford's online
+    # algorithm avoids the subtraction entirely.
+    n_per_group, n_features = 1000, 4
+    offset, std = 1e8, 1e-3
+    groups = ["a", "b"]
+    x = np.vstack([
+        offset
+        + std * np.random.default_rng().standard_normal((n_per_group, n_features))
+        for _ in groups
+    ])
+    obs = pd.DataFrame(
+        {"group": pd.Categorical(np.repeat(groups, n_per_group))},
+        index=[f"cell_{i}" for i in range(x.shape[0])],
+    )
+    adata = ad.AnnData(X=array_type(x), obs=obs)
+
+    expected = np.vstack([
+        np.var(x[i * n_per_group : (i + 1) * n_per_group], axis=0, ddof=0)
+        for i in range(len(groups))
+    ])
+    # Sanity: textbook formula on this data is either catastrophically wrong by a large magnitude relative to the expected
+    # or the sum-sq and sq-sum in naive are literally identical due to precision errors at the upper bound of the range.
+    naive = np.vstack([
+        (xg**2).mean(axis=0) - xg.mean(axis=0) ** 2
+        for xg in (
+            x[i * n_per_group : (i + 1) * n_per_group] for i in range(len(groups))
+        )
+    ])
+    diff_magnitude = np.abs(naive - expected) / expected
+    all_large = (diff_magnitude > 1e5).all()
+    if not all_large:
+        does_naive_fully_cancel = naive == 0
+        assert does_naive_fully_cancel.any()
+        assert (diff_magnitude[does_naive_fully_cancel] == 1).all()
+
+    result = sc.get.aggregate(adata, by="group", func="var", dof=0).layers["var"]
+    if isinstance(result, DaskArray):
+        result = result.compute()
+    np.testing.assert_allclose(result, expected, rtol=1e-4)