perf: optimization of rank_genes_groups

[zboldyga]

• Closes #
• Tests included or not required because:

• Release notes not necessary because:

@ilan-gold replayed changes of #4118 onto main

Edit on 07/15/2026: @ilan-gold I widened the scope of this PR to include two things:

• the stats computation improvements (original focus)
• and added illico's return_as_scanpy

I decided to combine these into one PR because they have some overlap, and I think it would be more cumbersome to have to come back again and do a separate PR review. Especially since rank_genes_groups overall is kinda obfuscated / not pythonic and I'm always losing track of what's going on in this file lol

[zboldyga]

@ilan-gold regarding these points, perhaps I open separate PR(s):

• addressing the scanpy 2.0 change in how cells with group NaN/unassigned are handled (uniform treatment across aggregate, rank_genes_groups, anywhere else relevant?, + docs update for this).

• adding benchmark code for rank_genes_groups paths that are fast-enough to run quickly (e.g. < 1 min).

[codecov]

Codecov Report

❌ Patch coverage is 96.52174% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.94%. Comparing base (8673b4d) to head (fa070ae).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/scanpy/tools/_rank_genes_groups.py	96.52%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4204      +/-   ##
==========================================
+ Coverage   79.92%   79.94%   +0.02%     
==========================================
  Files         121      121              
  Lines       12958    12987      +29     
==========================================
+ Hits        10357    10383      +26     
- Misses       2601     2604       +3     

Flag	Coverage Δ
hatch-test.low-vers	79.07% <96.52%> (+0.03%)	⬆️
hatch-test.pre	79.90% <96.52%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/scanpy/tools/_rank_genes_groups.py	93.59% <96.52%> (-0.57%)	⬇️

... and 1 file with indirect coverage changes

[scverse-benchmark]

Benchmark changes

Change	Before [c52da39]	After [e1a8863]	Ratio	Benchmark (Parameter)
-	7.79±0.01s	4.75±0.01s	0.61	preprocessing_counts.RankGenesGroups.time_agg('ovo', True)
-	10.5±0.02s	7.25±0.01s	0.69	preprocessing_counts.RankGenesGroups.time_agg('ovr', False)
-	14.3±0.04s	5.85±0s	0.41	preprocessing_counts.RankGenesGroups.time_agg('ovr', True)
-	69.5±2ms	47.5±1ms	0.68	tools.ToolsSuite.time_rank_genes_groups

Comparison: https://github.com/scverse/scanpy/compare/c52da398956c8de9a5bc1945a7966ad86b8431be..e1a88630cd6cef277acedcda56b9e5a6ab22e500
Last changed: Fri, 17 Jul 2026 11:05:15 +0000

More details: https://github.com/scverse/scanpy/pull/4204/checks?check_run_id=87850791862

[ilan-gold]

I don't fully understand this - could you explain more? Why do we need to use numba.njit here?

[zboldyga]

@ilan-gold in general njit is giving a significant speedup on the means/vars work for stats, and we use @numba.njit specifically because fau's njit is used on _vars_rest, and it's a python-based wrapper (https://github.com/scverse/fast-array-utils/blob/main/src/fast_array_utils/numba/__init__.py#L113), so when _chan_aggregate is called inside of that it can't reference the python wrapper code -- it has to call njit directly. @numba.njit is the same pattern used here elsewhere in scanpy, with the noqa exclusion

this is another thing I was curious to look closer at w/ fau / njit, as it means fairly common patterns of an njit function calling another njit function have to have this workaround w/ #noqa

[ilan-gold]

That makes total sense ok! I really think that fau njit function is mostly for dask but since so many of our kernels fall into that category....

[ilan-gold]

P.S Any thought why the csr-ovo case didn't improve?

[zboldyga]

@ilan-gold FYI there is also a patch I'll need to make so we use illico's return_as_scanpy. We're currently doing without it and it adds quite a bit of time (it's like 80% of the runtime, with illico being ~10%). I am debating about whether that work belongs in this PR, there might be some overlap with the stats code. Will report back on that later today

(this is an important change because it's still way faster to use illico directly. e.g. on reploglek562gwps it was something like 7 mins vs < 1 min).

[zboldyga]

P.S Any thought why the csr-ovo case didn't improve?

yes, so ovo is already fine. the issue in the stats code was vs_rest calculates things like means_rest, vars_rest, pts_rest, which are different for each group. That was done very inefficiently. And the performance impact was much bigger than it shows in the benchmark because those are small amounts of data and the runtime was ~quadratic. e.g. on real datasets see the tables here: #4118 . I'm rerunning a benchmark now and will share those results, as some things are unclear about the old benchmark I had run, and maybe incorrect. Will report back on that later today + share the script(s) with you

[zboldyga]

Edited on 07/16/2026 to reflect latest PR state

@ilan-gold so here are some more comprehensive benchmarks on small perturb seq datasets, which are still quite a bit larger than the benchmark data. This also runs more paths than we do in the built in benchmark, and I used cProfile to show a breakdown of where timing improvements come from.

Note that most datasets I typically use take hours on the old code in the vs_rest cases, and that's now down to minutes on the new code. But I didn't time those out here b/c too slow

dataset	cells	genes	groups
wessels23	30,636	16,775	184
adamson16	65,286	20,616	90
sunshine23	54,576	23,570	195
norman19	111,545	22,608	238

wessels23 (groups=184)

method	mode	old	new	speedup
t-test	vs_rest	18.76	1.69	11.1×
wilcoxon_illico	vs_rest	39.46	1.93	20.4×
t-test	vs_rest, pts	54.86	1.73	31.7×
wilcoxon_illico	vs_rest, pts	72.99	2.09	34.9×
t-test	vs_reference	1.75	1.49	1.2×
wilcoxon_illico	vs_reference	7.25	1.66	4.4×
wilcoxon_illico	vs_rest, rankby_abs	39.55	1.96	20.2×
wilcoxon_illico	vs_rest, non-log(base2)	40.82	1.92	21.3×

adamson16 (groups=90)

method	mode	old	new	speedup
t-test	vs_rest	20.18	2.23	9.0×
wilcoxon_illico	vs_rest	61.8	3.85	16.0×
t-test	vs_rest, pts	65.23	2.51	26.0×
wilcoxon_illico	vs_rest, pts	102.0	4.03	25.3×
t-test	vs_reference	2.4	2.02	1.2×
wilcoxon_illico	vs_reference	21.8	3.26	6.7×
wilcoxon_illico	vs_rest, rankby_abs	62.4	3.73	16.7×
wilcoxon_illico	vs_rest, non-log(base2)	62.3	3.92	15.9×

sunshine23 (groups=195)

method	mode	old	new	speedup
t-test	vs_rest	34.48	2.75	12.5×
wilcoxon_illico	vs_rest	80.8	3.05	26.5×
t-test	vs_rest, pts	110.11	3.06	36.0×
wilcoxon_illico	vs_rest, pts	146.7	3.48	42.2×
t-test	vs_reference	2.7	2.49	1.1×
wilcoxon_illico	vs_reference	18.1	2.91	6.2×
wilcoxon_illico	vs_rest, rankby_abs	80.7	3.45	23.4×
wilcoxon_illico	vs_rest, non-log(base2)	81.5	3.45	23.6×

norman19 (groups=238)

method	mode	old	new	speedup
t-test	vs_rest	91.13	4.71	19.3×
wilcoxon_illico	vs_rest	203.9	6.17	33.0×
t-test	vs_rest, pts	252.13	4.78	52.7×
wilcoxon_illico	vs_rest, pts	371.0	6.90	53.8×
t-test	vs_reference	4.6	3.98	1.2×
wilcoxon_illico	vs_reference	47.9	6.90	6.9×
wilcoxon_illico	vs_rest, rankby_abs	204.4	6.64	30.8×
wilcoxon_illico	vs_rest, non-log(base2)	204.5	6.64	30.8×

Component breakdown — wessels23, vs_rest (cProfile, profile_rgg.py, current code)

Cumulative seconds (nested, so not additive; each row is the largest matching frame).

t-test (rest_t) — profiled total 1.58s

component	cumtime	frame
STATS _basic_stats	0.05	_rank_genes_groups.py:325
↳ aggregate / _stats_vs_rest	0.04 / 0.05	_aggregate_group_stats / _stats_vs_rest
TEST t_test	0.37	_rank_genes_groups.py:460 (185 calls)
XFORM p-value correction (multipletests)	0.19	multitest.py:63
XFORM _build_stats_dataframe	0.84	_rank_genes_groups.py:731
ASSEMBLE to_records	0.10	frame.py:2557

wilcoxon_illico (rest_i) — profiled total 2.11s

component	cumtime	frame
TEST illico asymptotic_wilcoxon	0.69	asymptotic_wilcoxon.py:231
STATS _basic_stats (+aggregate/vs_rest)	0.03	_rank_genes_groups.py:325
XFORM _illico_results_to_iter	0.09	_rank_genes_groups.py:53
XFORM p-value correction (multipletests)	0.18	multitest.py:63
XFORM _build_stats_dataframe	0.46	_rank_genes_groups.py:731
ASSEMBLE to_records	0.10	frame.py:2557

_vars_rest numba kernel — njit vs pure-Python

Isolated microbenchmark of the leave-one-out variance kernel; identical inputs, JIT toggled via NUMBA_DISABLE_JIT; result checksums identical.

kernel size (k × genes)	njit	pure-Python	speedup
184 × 1000	0.71 ms	307 ms	435×
500 × 1000	1.86 ms	836 ms	449×
1819 × 1000	2.84 ms	3009 ms	1060×
1819 × 9623	34.6 ms	~29,000 ms (extrap.)	~840×
--> and this just clarifies the impact of njit here. It's negligible on small datasets but does move the needle on typical public perturb seq datasets, especially the ones from the past few years

Contribution of each mechanism

Also, there are 3 distinct mechanisms in this PR that give performance improvements in different circumstances.

1. passing the number of threads to illico so it can multithread. previously it was just single threading.

2. the means/vars stats computation code, which impacts vs_rest scenarios.

3. build_stats_dataframe improvement, which significantly impacted any rank_genes_groups work on datasets with a few thousand+ groups.

Here's a vs_rest breakdown:

wilcoxon_illico

phase	old	new	saved	mechanism
_basic_stats (means/vars)	30.13	0.03	30.1	① means/vars rewrite
illico asymptotic_wilcoxon	8.14	0.67	7.5	② illico threading
output build + correction + records	~1.4	~1.4	~0 @184g	③ dataframe construction (scale)
total	39.7	2.1	37.6

t-test

phase	old	new	saved	mechanism
_basic_stats (means/vars)	17.75	0.05	17.7	① means/vars rewrite
t-test + output build + correction	~1.8	~1.5	~0 @184g	③ dataframe construction (scale)
total	19.5	1.6	17.9

_build_stats_dataframe time vs #groups

Because running bigger datasets on the old code takes a long time, and the build stats df issue mostly appears there, I've isolated that piece with a separate benchmark here. So you can see once we start getting into bigger group sizes this grows. It will be substantial on gwps, multiple minutes for something that can take a few seconds.

Synthetic (many groups, 4 cells/group so illico is ~instant), wilcoxon_illico vs_rest,
8,248 genes, n_genes=None, 16 threads. Timer wraps _build_stats_dataframe only
(includes the inherent per-group p-value correction).

groups	OLD (per-column)	NEW (one-shot)	speedup	OLD growth / 2× groups
250	384 ms	299 ms	1.3×	—
500	961 ms	585 ms	1.6×	×2.5
1,000	2,295 ms	1,239 ms	1.9×	×2.4
2,000	6,730 ms	2,666 ms	2.5×	×2.9
4,000	22,228 ms	5,501 ms	4.0×	×3.3

[ilan-gold]

--> to clarify here: stats can be skipped here because we use illico's return_as_scanpy which already does that work. And illico before this PR was always running single threaded, so that's why the illico call itself is faster.

@zboldyga Do you have a sense where the time difference is coming from? The return_as_scanpy only gives what we are calculating here + p values (I think). So is our var calculation really slow? Something else? Like, what is return_as_scanpy doing that we're not? It seems that if the problem is aggregation speed, that's a deeper issue (and also why I prefer to dogfood our own stuff internally)

[ilan-gold]

This explanation is exactly why I preferred to rely on scanpy's machinery outside illico given the complexity - for me, the illico - scanpy gap is quite small i.e., calculating the stats isn't even that slow, it appears to be largely coming from our nonnegativity check, which I would argue we should just abandon in favor of "is this a unsigned integer" and throw an error otherwise. Would love to hear your feedback, but for my conversation with claude, this appeared to be the case.

[zboldyga]

--> to clarify here: stats can be skipped here because we use illico's return_as_scanpy which already does that work. And illico before this PR was always running single threaded, so that's why the illico call itself is faster.

@zboldyga Do you have a sense where the time difference is coming from? The return_as_scanpy only gives what we are calculating here + p values (I think). So is our var calculation really slow? Something else? Like, what is return_as_scanpy doing that we're not? It seems that if the problem is aggregation speed, that's a deeper issue (and also why I prefer to dogfood our own stuff internally)

@ilan-gold thb I should have approached this differently. I was initially thinking, illico is kind of getting absorbed into scanpy, may as well use this supported feature b/c it was already implemented. But that doesn't really make sense

I'm glad you pushed on this. build_stats_dataframe was the issue that was giving longer runtimes on real datasets which using return_as_scanpy was inadvertently solving. But the build_stats_dataframe fix is actually relevant to any rank_genes_groups call with thousands+ groups: dataframe uses BlockManager internally, and setting properties like df[group, "scores"] I think rebuilds the whole BlockManager memory thus far, so it's super inefficient (basically O(groups^2)). So I've reverted to not using return_as_scanpy, tweaked the dataframe build (which also impacts the runtime of other de methods -- especially with bigger datasets), and re-run benchmarks to clarify.

I updated above benchmarks to reflect latest code. will push momentarily

[zboldyga]

@ilan-gold see above message with my benchmark timing and breakdowns of the 3 optimization areas in this PR.

all of this said, I do think expanding the benchmarks to capture some of these pieces would make sense. And we can do it on some isolated levels, e.g. build stats dataframe, other helpers. But perhaps we do that after this PR for two reasons:

1. Some things we should test can't reasonably have a before/after test right now because the before part is too slow on realistic data sizes.

2. I really, really want to improve the code quality of rank_genes_groups. I find it disorienting. May be a little easier to do that first, then once more confident about the internal function signatures and return types, some of the internal functions could be profiled directly? (like building the stats dataframe, and likely other pieces)... I can coordinate with Philip here as I think you mentioned he might already be working on this

[ilan-gold]

I'm glad you pushed on this. build_stats_dataframe was the issue that was giving longer runtimes on real datasets which using return_as_scanpy was inadvertently solving. But the build_stats_dataframe fix is actually relevant to any rank_genes_groups call with thousands+ groups: dataframe uses BlockManager internally, and setting properties like df[group, "scores"] I think rebuilds the whole BlockManager memory thus far, so it's super inefficient (basically O(groups^2)). So I've reverted to not using return_as_scanpy, tweaked the dataframe build (which also impacts the runtime of other de methods -- especially with bigger datasets), and re-run benchmarks to clarify.

Yesssssssss, great deep dive!

I really, really want to improve the code quality of rank_genes_groups. I find it disorienting. May be a little easier to do that first, then once more confident about the internal function signatures and return types, some of the internal functions could be profiled directly? (like building the stats dataframe, and likely other pieces)... I can coordinate with Philip here as I think you mentioned he might already be working on this

all of this said, I do think expanding the benchmarks to capture some of these pieces would make sense. And we can do it on some isolated levels, e.g. build stats dataframe, other helpers. But perhaps we do that after this PR for two reasons:

In general, like with testing, not a huge fan of benchmarking private APIs. But possibly the constituent steps become public (would not be unreasonable to e.g., request from a class-based RGG to rebuild the p-values with a different correction method etc.).

For this, you will definitely want to talk to Phil as I think scanpy will expose a class-based 2.0 API but am not sure.

[ilan-gold]

So in its current form, I think this is ready for the light-of-day. Thanks a million. Feel free to go ham in a follow-up PR updating RGG - if you want to run API stuff by someone, independent of whether I am around, Phil is probably best anyway.

[lumberbot-app]

Owee, I'm MrMeeseeks, Look at me.

There seem to be a conflict, please backport manually. Here are approximate instructions:

1. Checkout backport branch and update it.

git checkout 1.12.x
git pull

2. Cherry pick the first parent branch of the this PR on top of the older branch:

git cherry-pick -x -m1 6c3204c0c4e1938a7bfe10c9631d43461b22a351

3. You will likely have some merge/cherry-pick conflict here, fix them and commit:

git commit -am 'Backport PR #4204: perf: optimization of rank_genes_groups'

4. Push to a named branch:

git push YOURFORK 1.12.x:auto-backport-of-pr-4204-on-1.12.x

5. Create a PR against branch 1.12.x, I would have named this PR:

"Backport PR #4204 on branch 1.12.x (perf: optimization of rank_genes_groups)"

And apply the correct labels and milestones.

Congratulations — you did some good work! Hopefully your backport PR will be tested by the continuous integration and merged soon!

Remember to remove the Still Needs Manual Backport label once the PR gets merged.

If these instructions are inaccurate, feel free to suggest an improvement.