diff --git a/docs/release-notes/4204.perf.md b/docs/release-notes/4204.perf.md new file mode 100644 index 0000000000..e4c0e98790 --- /dev/null +++ b/docs/release-notes/4204.perf.md @@ -0,0 +1,3 @@ +Speed up {func}`scanpy.tl.rank_genes_groups` `vs_rest` statistics via a single {func}`scanpy.get.aggregate` pass with [Chan's][] cancellation-free leave-one-out variance {smaller}`Z Boldyga` + +[Chan's]: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm diff --git a/src/scanpy/tools/_rank_genes_groups.py b/src/scanpy/tools/_rank_genes_groups.py index 6b487f39ea..80f60f2b8a 100644 --- a/src/scanpy/tools/_rank_genes_groups.py +++ b/src/scanpy/tools/_rank_genes_groups.py @@ -9,12 +9,11 @@ import pandas as pd from anndata import AnnData from fast_array_utils.numba import njit -from fast_array_utils.stats import mean_var from scipy import sparse from .. import _utils from .. import logging as logg -from .._compat import CSBase, warn +from .._compat import CSBase, DaskArray, warn from .._settings import Default, Preset from .._settings.presets import DETest from .._utils import ( @@ -23,7 +22,8 @@ get_literal_vals, raise_not_implemented_error_if_backed_type, ) -from ..get import _check_mask, _get_obs_rep +from ..get import _check_mask, _get_obs_rep, aggregate +from ..get._aggregated import _chan_combine if TYPE_CHECKING: from collections.abc import Generator, Iterable @@ -33,6 +33,7 @@ type _CorrMethod = Literal["benjamini-hochberg", "bonferroni"] +type _TestResult = tuple[int, NDArray[np.floating], NDArray[np.floating] | None] _CONST_MAX_SIZE: int = 10_000_000 @@ -53,7 +54,7 @@ def _illico_results_to_iter( ireference: int | None, *, copy_pvalues: bool, -) -> Generator[tuple[int, NDArray[np.floating], NDArray[np.floating]]]: +) -> Generator[_TestResult]: """Yield per-group ``(index, z, p)`` tuples from illico's long-form output. illico returns a DataFrame with a 2-level MultiIndex ``(pert, feature)`` @@ -152,6 +153,96 @@ def _ranks( yield ranks, left, right +def _apply_expm1_preserving_sparsity(x, expm1_func): + """Apply ``expm1`` to ``x`` while keeping sparse data sparse. + + Applied lazily and chunk-wise for dask; ``expm1(0) == 0`` preserves sparsity. + """ + if isinstance(x, DaskArray): + return x.map_blocks( + _apply_expm1_preserving_sparsity, + expm1_func, + dtype=x.dtype, + meta=x._meta, + ) + if isinstance(x, CSBase): + xp = x.copy() + xp.data = expm1_func(xp.data) + return xp + return expm1_func(x) + + +@numba.njit # noqa: TID251 (inner kernel called from _vars_rest's nopython loop) +def _chan_accumulate( + group_counts: NDArray[np.float64], + mean: NDArray[np.float64], + m2: NDArray[np.float64], + j: int, + direction: int, +) -> NDArray[np.float64]: + """Accumulate a running Chan combine of the groups for gene ``j``. + + ``acc[i]`` holds the combined ``(count, mean, M2)`` of group ``i`` together with + every group toward ``direction``: ``+1`` gives groups ``0..i`` (forward), + ``-1`` gives groups ``i..end`` (backward). + """ + n_groups = group_counts.shape[0] + acc = np.empty((n_groups, 3)) + # accumulated (count, mean, M2) of the groups visited so far + acc_count = acc_mean = acc_m2 = 0.0 + # visit groups forward (direction +1) or backward (-1) + group_order = range(n_groups) if direction == 1 else range(n_groups - 1, -1, -1) + for i in group_order: + acc_count, acc_mean, acc_m2 = _chan_combine( + acc_count, acc_mean, acc_m2, group_counts[i], mean[i, j], m2[i, j] + ) + acc[i, 0], acc[i, 1], acc[i, 2] = acc_count, acc_mean, acc_m2 + return acc + + +@njit +def _vars_rest( + group_counts: NDArray[np.float64], + mean: NDArray[np.float64], + m2: NDArray[np.float64], + k: int, +) -> NDArray[np.float64]: + """Leave-one-out variance for each selected group, parallel over genes. + + Group ``g``'s "rest" is every other group combined — the groups up to ``g - 1`` + pooled with the groups from ``g + 1`` — so variances are never subtracted + (Chan's cancellation-free combine). + """ + n_genes = mean.shape[1] + vars_rest = np.zeros((k, n_genes)) + for j in numba.prange(n_genes): + # combined (count, mean, M2) of groups 0..i (forward) and i..end (backward) + combined_stats_upto_group = _chan_accumulate(group_counts, mean, m2, j, 1) + combined_stats_from_group = _chan_accumulate(group_counts, mean, m2, j, -1) + + # each group g's "rest" stats = the groups before g pooled with the groups after g + for g in range(k): + stats_after_g = combined_stats_from_group[g + 1] + if g >= 1: + stats_before_g = combined_stats_upto_group[g - 1] + # each stats row is (count, mean, M2) + n_r, _, m2_r = _chan_combine( + n_a=stats_before_g[0], + mean_a=stats_before_g[1], + m2_a=stats_before_g[2], + n_b=stats_after_g[0], + mean_b=stats_after_g[1], + m2_b=stats_after_g[2], + ) + else: + # g == 0 has no groups before it, so its rest is just the groups after + n_r, m2_r = stats_after_g[0], stats_after_g[2] + denom = n_r - 1.0 if n_r >= 2.0 else 1.0 + v = m2_r / denom + vars_rest[g, j] = v if v > 0.0 else 0.0 + return vars_rest + + class _RankGenes: def __init__( self, @@ -218,7 +309,6 @@ def __init__( self.means = None self.vars = None - self.means_rest = None self.vars_rest = None @@ -232,68 +322,144 @@ def __init__( self.grouping_mask = adata.obs[groupby].isin(self.groups_order) self.grouping = adata.obs.loc[self.grouping_mask, groupby] - def _basic_stats(self, *, exponentiate_values: bool = False) -> None: - """Set self.{means,vars,pts}{,_rest} depending on X.""" - n_genes = self.X.shape[1] - n_groups = self.groups_masks_obs.shape[0] - - self.means = np.zeros((n_groups, n_genes)) - self.vars = np.zeros((n_groups, n_genes)) - self.pts = np.zeros((n_groups, n_genes)) if self.comp_pts else None - + def _basic_stats( + self, *, exponentiate_values: bool = False, need_var: bool = False + ) -> None: + """Populate per-group stats, and (in vs_rest mode) rest-group stats. + + ``need_var`` controls whether variance (per-group and per-rest) is + computed; only the t-test family reads it. In vs_rest mode every cell + is assigned to its selected group or a single "remainder" group (cells + in no selected group), and each group's "rest" is the forward + Chan-combine of all other groups — a sum of non-negative terms, hence + free of catastrophic cancellation for any group sizes. + """ + x = ( + _apply_expm1_preserving_sparsity(self.X, self.expm1_func) + if exponentiate_values + else self.X + ) if self.ireference is None: - self.means_rest = np.zeros((n_groups, n_genes)) - self.vars_rest = np.zeros((n_groups, n_genes)) - self.pts_rest = np.zeros((n_groups, n_genes)) if self.comp_pts else None + self._stats_vs_rest(x, need_var=need_var) else: - mask_rest = self.groups_masks_obs[self.ireference] - x_rest = self.X[mask_rest] - if exponentiate_values: - x_rest = self.expm1_func(x_rest) - self.means[self.ireference], self.vars[self.ireference] = mean_var( - x_rest, axis=0, correction=1 - ) - # deleting the next line causes a memory leak for some reason - del x_rest - - if isinstance(self.X, CSBase): - get_nonzeros = lambda x: x.getnnz(axis=0) + self._stats_vs_reference(x, need_var=need_var) + + def _aggregate_group_stats( + self, x_used, codes: NDArray[np.int64], n_groups: int, *, need_var: bool + ): + """Aggregate ``x_used`` in one batched :func:`scanpy.get.aggregate`. + + Grouped by ``codes`` (values ``0 .. n_groups-1``). Returns ``(mean, var, + nnz)`` of shape ``(n_groups, n_genes)``, zero-filled for groups with + no cells. ``var`` is ``None`` unless ``need_var``; ``nnz`` is ``None`` + unless ``self.comp_pts``. + """ + n_genes = x_used.shape[1] + mean = np.zeros((n_groups, n_genes)) + var = np.zeros((n_groups, n_genes)) if need_var else None + nnz = np.zeros((n_groups, n_genes)) if self.comp_pts else None + + funcs = ["mean"] + if need_var: + funcs.append("var") + if self.comp_pts: + funcs.append("count_nonzero") + agg_adata = AnnData( + X=x_used, + obs=pd.DataFrame( + {"_g": pd.Categorical(codes, categories=range(n_groups))}, + index=pd.RangeIndex(len(codes)).astype(str), + ), + ) + out = aggregate(agg_adata, by="_g", func=funcs, dof=1) + idx = out.obs_names.astype(int).to_numpy() + mean[idx] = np.asarray(out.layers["mean"]) + if need_var: + var[idx] = np.asarray(out.layers["var"]) + if self.comp_pts: + nnz[idx] = np.asarray(out.layers["count_nonzero"]) + return mean, var, nnz + + def _stats_vs_reference(self, x, *, need_var: bool) -> None: + """Aggregate the selected-group cells only (vs-reference; no rest derivation). + + The reference is itself one of the selected groups. + """ + mask = self.grouping_mask.to_numpy() + x_used = x if mask.all() else x[mask] + codes = pd.Index(self.groups_order).get_indexer(self.grouping) + k = self.groups_masks_obs.shape[0] + + self.means, self.vars, nnz = self._aggregate_group_stats( + x_used, codes, k, need_var=need_var + ) + if self.comp_pts: + n_per_group = self.groups_masks_obs.sum(axis=1) + self.pts = nnz / n_per_group[:, None] else: - get_nonzeros = lambda x: np.count_nonzero(x, axis=0) - - for group_index, mask_obs in enumerate(self.groups_masks_obs): - x_mask = self.X[mask_obs] - if exponentiate_values: - x_mask = self.expm1_func(x_mask) + self.pts = None + + def _stats_vs_rest(self, x, *, need_var: bool) -> None: + """Assign every cell to one of the ``k`` selected groups or a remainder group. + + The remainder group holds cells in no selected group (non-selected + groups and unassigned/NaN). Each group's "rest" is the forward + Chan-combine of every other group. + """ + k = self.groups_masks_obs.shape[0] + + # each cell's selected-group index, or `k` (the remainder group) for + # cells in no selected group (non-selected / NaN) + sel = pd.Index(self.groups_order).get_indexer(self.group_col) + codes = np.where(sel >= 0, sel, k).astype(np.int64) + group_counts = np.bincount(codes, minlength=k + 1) # group k == remainder + n_sel = group_counts[:k] + + mean, var, nnz = self._aggregate_group_stats(x, codes, k + 1, need_var=need_var) + + # selected-group arm of the test (the remainder group is excluded) + self.means = mean[:k] + self.vars = var[:k] if need_var else None + self.pts = nnz[:k] / n_sel[:, None] if self.comp_pts else None + + # m2 = var * (n - 1); forced to 0 for groups with <= 1 cell so a + # singleton remainder (aggregate var undefined there) is harmless + if need_var: + with np.errstate(invalid="ignore"): + m2 = var * (group_counts[:, None] - 1) + m2[group_counts <= 1] = 0.0 + else: + m2 = None - if self.comp_pts: - self.pts[group_index] = get_nonzeros(x_mask) / x_mask.shape[0] + self._derive_rest_stats(group_counts, mean, m2, nnz, k, need_var=need_var) - if self.ireference is not None and group_index == self.ireference: - continue - - self.means[group_index], self.vars[group_index] = mean_var( - x_mask, axis=0, correction=1 + def _derive_rest_stats( + self, group_counts, mean, m2, nnz, k: int, *, need_var: bool + ) -> None: + """Set ``means_rest``/``vars_rest``/``pts_rest`` for each selected group ``g``. + + Statistics are over every cell *not* in ``g``. ``means_rest`` and + ``pts_rest`` are linear in the groups, so they are the exact + total-minus-group difference. Variance would lose precision under such + subtraction, so ``vars_rest`` uses the cancellation-free Chan + leave-one-out scan (:func:`_vars_rest`) — only the variance-based tests + request it. + """ + n_rest = (self.X.shape[0] - group_counts[:k])[:, None] + total = (group_counts[:, None] * mean).sum(axis=0) + self.means_rest = (total - group_counts[:k, None] * mean[:k]) / n_rest + self.vars_rest = ( + _vars_rest( + np.ascontiguousarray(group_counts, dtype=np.float64), mean, m2, k ) - - if self.ireference is None: - mask_rest = ~mask_obs - x_rest = self.X[mask_rest] - if exponentiate_values: - x_rest = self.expm1_func(x_rest) - ( - self.means_rest[group_index], - self.vars_rest[group_index], - ) = mean_var(x_rest, axis=0, correction=1) - # this can be costly for sparse data - if self.comp_pts: - self.pts_rest[group_index] = get_nonzeros(x_rest) / x_rest.shape[0] - # deleting the next line causes a memory leak for some reason - del x_rest + if need_var + else None + ) + self.pts_rest = (nnz.sum(axis=0) - nnz[:k]) / n_rest if self.comp_pts else None def t_test( self, method: Literal["t-test", "t-test_overestim_var"] - ) -> Generator[tuple[int, NDArray[np.floating], NDArray[np.floating]], None, None]: + ) -> Generator[_TestResult, None, None]: from scipy import stats for group_index, (mask_obs, mean_group, var_group) in enumerate( @@ -342,9 +508,7 @@ def t_test( yield group_index, scores, pvals - def wilcoxon( - self, *, tie_correct: bool - ) -> Generator[tuple[int, NDArray[np.floating], NDArray[np.floating]], None, None]: + def wilcoxon(self, *, tie_correct: bool) -> Generator[_TestResult, None, None]: from scipy import stats n_genes = self.X.shape[1] @@ -421,9 +585,7 @@ def wilcoxon( yield group_index, scores[group_index], pvals - def logreg( - self, **kwds - ) -> Generator[tuple[int, NDArray[np.floating], None], None, None]: + def logreg(self, **kwds) -> Generator[_TestResult, None, None]: # if reference is not set, then the groups listed will be compared to the rest # if reference is set, then the groups listed will be compared only to the other groups listed from sklearn.linear_model import LogisticRegression @@ -456,7 +618,7 @@ def logreg( def illico( self, *, tie_correct: bool, corr_method: _CorrMethod - ) -> Generator[tuple[int, NDArray[np.floating], NDArray[np.floating]], None, None]: + ) -> Generator[_TestResult, None, None]: from illico import asymptotic_wilcoxon illico_df = asymptotic_wilcoxon( @@ -488,7 +650,7 @@ def illico( copy_pvalues=corr_method == "benjamini-hochberg", ) - def compute_statistics( # noqa: PLR0912 + def compute_statistics( self, method: DETest, *, @@ -500,11 +662,8 @@ def compute_statistics( # noqa: PLR0912 **kwds, ) -> None: if method in {"t-test", "t-test_overestim_var"}: - self._basic_stats(exponentiate_values=False) + self._basic_stats(exponentiate_values=not mean_in_log_space, need_var=True) generate_test_results = self.t_test(method) - if not mean_in_log_space: - # If we are not exponentiating after the mean aggregation, we need to recalculate the stats. - self._basic_stats(exponentiate_values=True) elif "wilcoxon" in method: generate_test_results = ( self.illico(tie_correct=tie_correct, corr_method=corr_method) @@ -516,60 +675,81 @@ def compute_statistics( # noqa: PLR0912 elif method == "logreg": generate_test_results = self.logreg(**kwds) - self.stats = None - - n_genes = self.X.shape[1] - for group_index, scores, pvals in generate_test_results: - group_name = str(self.groups_order[group_index]) - - if n_genes_user is not None: - scores_sort = np.abs(scores) if rankby_abs else scores - global_indices = _select_top_n(scores_sort, n_genes_user) - first_col = "names" - else: - global_indices = slice(None) - first_col = "scores" - - if self.stats is None: - idx = pd.MultiIndex.from_tuples([(group_name, first_col)]) - self.stats = pd.DataFrame(columns=idx) + self.stats = _build_stats_dataframe( + self, + generate_test_results, + corr_method=corr_method, + n_genes_user=n_genes_user, + rankby_abs=rankby_abs, + mean_in_log_space=mean_in_log_space, + ) - if n_genes_user is not None: - self.stats[group_name, "names"] = self.var_names[global_indices] - self.stats[group_name, "scores"] = scores[global_indices] +def _build_stats_dataframe( + rg: _RankGenes, + results: Iterable[_TestResult], + *, + corr_method: _CorrMethod, + n_genes_user: int | None, + rankby_abs: bool, + mean_in_log_space: bool, +) -> pd.DataFrame | None: + """Drain the per-group ``(group_index, scores, pvals)`` iterator into a DataFrame. + + Builds a wide-form ``(group, statistic)`` MultiIndex frame: top-N selection, + multiple-testing correction, and (when ``rg.means`` is set) log2 + fold-change. Read-only on ``rg``. + """ + from statsmodels.stats.multitest import multipletests - if pvals is not None: - self.stats[group_name, "pvals"] = pvals[global_indices] - if corr_method == "benjamini-hochberg": - from statsmodels.stats.multitest import multipletests + n_genes_total = rg.X.shape[1] + df: pd.DataFrame | None = None - pvals[np.isnan(pvals)] = 1 - _, pvals_adj, _, _ = multipletests( - pvals, alpha=0.05, method="fdr_bh" - ) - elif corr_method == "bonferroni": - pvals_adj = np.minimum(pvals * n_genes, 1.0) - self.stats[group_name, "pvals_adj"] = pvals_adj[global_indices] - - if self.means is not None: - mean_group = self.means[group_index] - if self.ireference is None: - mean_rest = self.means_rest[group_index] - else: - mean_rest = self.means[self.ireference] - foldchanges = ( - (self.expm1_func(mean_group) + 1e-9) - / (self.expm1_func(mean_rest) + 1e-9) - if mean_in_log_space - else (mean_group + 1e-9) / (mean_rest + 1e-9) - ) # add small value to avoid zeros - self.stats[group_name, "logfoldchanges"] = np.log2( - foldchanges[global_indices] - ) + for group_index, scores, pvals in results: + group_name = str(rg.groups_order[group_index]) - if n_genes_user is None: - self.stats.index = self.var_names + if n_genes_user is not None: + scores_sort = np.abs(scores) if rankby_abs else scores + global_indices = _select_top_n(scores_sort, n_genes_user) + first_col = "names" + else: + global_indices = slice(None) + first_col = "scores" + + if df is None: + idx = pd.MultiIndex.from_tuples([(group_name, first_col)]) + df = pd.DataFrame(columns=idx) + + if n_genes_user is not None: + df[group_name, "names"] = rg.var_names[global_indices] + df[group_name, "scores"] = scores[global_indices] + + if pvals is not None: + df[group_name, "pvals"] = pvals[global_indices] + if corr_method == "benjamini-hochberg": + pvals[np.isnan(pvals)] = 1 + _, pvals_adj, _, _ = multipletests(pvals, alpha=0.05, method="fdr_bh") + elif corr_method == "bonferroni": + pvals_adj = np.minimum(pvals * n_genes_total, 1.0) + df[group_name, "pvals_adj"] = pvals_adj[global_indices] + + if rg.means is not None: + mean_group = rg.means[group_index] + mean_rest = ( + rg.means_rest[group_index] + if rg.ireference is None + else rg.means[rg.ireference] + ) + foldchanges = ( + (rg.expm1_func(mean_group) + 1e-9) / (rg.expm1_func(mean_rest) + 1e-9) + if mean_in_log_space + else (mean_group + 1e-9) / (mean_rest + 1e-9) + ) # add small value to avoid zeros + df[group_name, "logfoldchanges"] = np.log2(foldchanges[global_indices]) + + if df is not None and n_genes_user is None: + df.index = rg.var_names + return df def rank_genes_groups( # noqa: PLR0912, PLR0913, PLR0915 diff --git a/tests/test_rank_genes_groups.py b/tests/test_rank_genes_groups.py index f3d58f5b01..119bacd307 100644 --- a/tests/test_rank_genes_groups.py +++ b/tests/test_rank_genes_groups.py @@ -116,8 +116,14 @@ def test_results( for g in range(expected["names"].shape[0]): with subtests.test(group=g): + # atol guards against ULP-level golden vs new-code bit-pattern + # differences at near-zero t-scores (e.g., genes with equal + # group/rest means produce 0.0 in one path vs ~1e-15 in another). np.testing.assert_allclose( - expected["scores"][g, :n], results["scores"][str(g)][:n], rtol=1e-5 + expected["scores"][g, :n], + results["scores"][str(g)][:n], + rtol=1e-5, + atol=1e-10, ) np.testing.assert_array_equal( expected["names"][g, :n], results["names"][str(g)][:n]