PERF: Speed up permutation cluster tests via Numba union-find + compact graph

[sharifhsn]

Reference issue

Related: #5439, #7784, #8095, #12609

What does this implement/fix?

Speeds up spatio_temporal_cluster_1samp_test (and the other permutation_cluster_* functions) by ~5-10x on realistic data. The PR is split into 7 incremental commits. Maintainers can accept or reject each layer independently.

Commit 1 — Precompute sum-of-squares for sign-flip t-test (+29/−9 lines, 3.2x)
For the default ttest_1samp_no_p, s²=1 means sum(X²) is constant across permutations. Each permutation becomes a single signs @ X dot product instead of calling stat_fun. Also skips buffer_size verification for built-in stat functions.

Commit 2 — Numba union-find for spatio-temporal CCL (+226/−11 lines, 10.3x cumulative)
JIT-compiled union-find kernel (_st_fused_ccl) with path compression and union-by-rank, replacing the Python BFS in _get_clusters_st. Bundles tightly-coupled pieces: pre-computed CSR adjacency arrays, _sums_only flag to skip cluster list construction (uses np.bincount instead), and _csr_data parameter threading. These are bundled because _sums_only only fires inside if has_numba: and CSR data is only consumed by the Numba kernel.

Commit 3 — Extract _union helper + simplify (+36/−72 lines)
Extract duplicated find+union logic into _union() with inline="always", simplify _sum_cluster_data, trim docstrings/comments to match codebase style.

Commit 4 — Fix step-down reshape (+1/−1 lines)
Pre-existing bug: adjacency is None and adjacency is not False was equivalent to just adjacency is None, missing the adjacency is False case where step_down_include still needs reshaping.

Commit 5 — Changelog entries

Commit 6 — Test fixture (+1 line)
Patch has_numba in numba_conditional fixture so the "NumPy" test variant actually exercises the Python BFS fallback path for spatio-temporal clustering.

Commit 7 — Docstring (+18/−1 lines)
Expand _st_fused_ccl docstring with algorithm description, complexity analysis, and Wikipedia reference, per reviewer request.

Commits 3-7 are cleanup, bugfix, docs, and tests — they don't affect performance. All optimizations fall back to the original code paths when Numba is not installed. No public API changes.

Benchmarks

Per-commit cumulative speedup (local, Apple M-series, spatio_temporal_cluster_1samp_test, ico-5, 15 subjects x 15 timepoints x 20,484 vertices, threshold=3.0, 512 permutations, median of 3 runs):

Cumulative through	ms/perm	Speedup	Net lines
main (baseline)	16.94	1.0x	—
commit 1 (precomputed sum_sq)	5.37	3.2x	+20
commit 2 (Numba union-find)	1.64	10.3x	+235

AWS HPC end-to-end (AMD EPYC 7R13, same data dimensions):

Permutations	Before	After	Speedup
256	4.12 s	0.86 s	4.8x
1024	16.35 s	3.25 s	5.0x
4096	65.00 s	12.64 s	5.1x

Per-permutation cost: 15.8 ms → 3.1 ms (5.2x). Projected 10,000 permutations: 31 s vs 159 s.

Reproduce benchmarks locally

"""Quick benchmark: perf-opt vs baseline on realistic source-space data."""
import time
import numpy as np
import mne
from mne.stats import spatio_temporal_cluster_1samp_test
from mne.stats import cluster_level as cl

# Load fsaverage ico-5 adjacency
subjects_dir = mne.datasets.sample.data_path() / "subjects"
src = mne.setup_source_space(
    "fsaverage", spacing="ico5", subjects_dir=subjects_dir, add_dist=False
)
adjacency = mne.spatial_src_adjacency(src)

# Synthetic data: 15 subjects x 15 timepoints x 20,484 vertices
rng = np.random.default_rng(42)
X = rng.standard_normal((15, 15, adjacency.shape[0]))
X[:, 5:10, 1000:1100] += 1.0  # inject focal activation

# Warmup JIT
spatio_temporal_cluster_1samp_test(
    X, adjacency=adjacency, n_permutations=64,
    threshold=3.0, tail=1, verbose=False, seed=42
)

# Optimized
t0 = time.perf_counter()
spatio_temporal_cluster_1samp_test(
    X, adjacency=adjacency, n_permutations=512,
    threshold=3.0, tail=1, verbose=False, seed=42
)
t_opt = time.perf_counter() - t0

# Baseline (disable Numba path)
saved = cl.has_numba
cl.has_numba = False
t0 = time.perf_counter()
spatio_temporal_cluster_1samp_test(
    X, adjacency=adjacency, n_permutations=512,
    threshold=3.0, tail=1, verbose=False, seed=42
)
t_base = time.perf_counter() - t0
cl.has_numba = saved

print(f"Optimized: {t_opt:.2f}s  Baseline: {t_base:.2f}s  Speedup: {t_base/t_opt:.1f}x")

Additional information

• Numba JIT warmup happens once on first call; subsequent calls pay no warmup cost
• TFCE (threshold=dict(...)) correctly falls back to the original code path
• Custom stat functions still benefit from the CCL and overhead optimizations but not the precomputed sum-of-squares
• AI (Claude) was used to generate the code, which was checked over manually

[drammock]

this is a pretty large diff. Before we invest time in a review, a few questions:

1. is this ready for review? if not, please mark as draft until it is ready
2. have you run relevant tests and doc build locally? are they all passing?
3. please disclose the way(s) in which you used AI to assist in this contribution (if any)

tip: next time, if you name the changelog newfeature.rst instead of XXXXX.newfeature.rst then one of our CIs will automatically rename it to include the PR number for you.

[sharifhsn]

Note: The PR has since been restructured and is now ready for review. See the latest comments and PR body for the current state.

My bad, I meant to publish it as draft. Yes it's still in progress.

Tests are all passing.

AI was used to generate the code which was checked over manually to catch bugs.

I'm going to try to reduce the amount of code as much as possible while still maintaining the main speedups.

[sharifhsn]

Note: This comment describes an earlier version of the PR. The current 7-commit structure is in the updated PR body — commits 1-4 from this comment have been dropped per reviewer feedback.

Restructured this into 6 incremental commits ordered roughly by bang-for-buck, so you can review (and accept/reject) each optimization independently.

Quick summary of what each one does and what it costs:

1. Vectorize p-value computation: np.searchsorted replacing O(n·k) list comp
2. Compact graph reindexing : subgraph of supra-threshold vertices only
3. Vectorize cluster sums: np.add.reduceat replacing per-cluster loop
4. Vectorize sign-order generation: bit-shifting instead of np.binary_repr
5. Precompute sum-of-squares for sign-flip t-test: 3.2× end-to-end
6. Numba union-find for spatio-temporal CCL (+215 lines): 10.3× end-to-end

Commits 1-4 are small, pure NumPy/SciPy, and basically free in terms of complexity. They mainly help the global-adjacency path rather than the spatio-temporal one, so they don't show up much in the benchmark above.

Commit 5 is where the first big measurable win kicks in (3.2×), and it's still straightforward NumPy.

Commit 6 is the heavy one. It adds a Numba JIT union-find kernel with CSR precomputation and a bincount shortcut for cluster sums. It's responsible for most of the remaining speedup (3.2× → 10.3×), but it's also +215 lines and adds a Numba dependency to this code path. I have a stripped-down scipy-only version that gets ~3.8× in about +84 lines if you'd prefer to keep it simpler. Happy to swap that in or simplify further if the Numba approach feels like too much for this module. Would appreciate your guidance on where you'd want to draw that line.

[sharifhsn]

Note: This comment describes an earlier version. Commits 1-4 have been dropped per reviewer feedback. The current 7-commit structure is in the updated PR body.

Update: Restructured from the initial 6 commits:

• Commit 7 (new): Extracted duplicated union-find code into _union() helper with inline="always", simplified _sum_cluster_data, and trimmed docstrings/comments to match codebase conventions. Net −36 lines from commit 6.
• Commit 8 (new): Split out a pre-existing bugfix — adjacency is None and adjacency is not False → adjacency is None or adjacency is False — the old condition always evaluated to just adjacency is None, missing the adjacency is False case in step-down-in-jumps. This is a 1-line fix to existing code, separated so it can be merged independently.
• Commit 9: Changelog entries (newfeature + bugfix).
• Commit 10 (new): Patch has_numba in the numba_conditional test fixture so the "NumPy" variant actually exercises the Python BFS fallback for spatio-temporal clustering.

PR body updated accordingly.

[sharifhsn]

The PR has been restructured extensively since that initial comment and is now ready for review. All 55 relevant cluster-level tests pass locally (both Numba and NumPy variants).

[wmvanvliet]

Looking at it quickly, commits 1-4 don't really make much of an impact in speed while slightly increasing the code complexity. However, commits 5 and especially 6 look interesting. Speeding up the cluster permutation test by 10x is a huge improvement! Since this is AI generated, could the AI generate some more information in the docstring about the role of union_find and how it goes about doing it exactly? It's hard to tell from the code alone what is actually going on and why it is so fast. If this algorithm has been published somewhere, it would be good to link to an article. Or if it's some well-known algorithm with a wikipedia page we could link there. Something to aid human readers of the code.

[sharifhsn]

Thanks for taking a look!

Commits 1-4: Agreed — they don't measurably speed up the spatio-temporal path and add unnecessary complexity. Dropped all four; the PR is now 7 commits focused on the two performance wins (precomputed sum-of-squares + Numba union-find) plus cleanup, bugfix, docs, and test.

Union-find documentation: Added an expanded docstring to _st_fused_ccl explaining what the algorithm does, why it's faster than the BFS it replaces, and a link to the Wikipedia article on disjoint-set data structures (commit 7).

[larsoner]

Can you speak to the cost to speed when numba is not available? I suspect all of these nested loops will perhaps have a noticeable additional performance cost in those cases rather than a benefit.

If that ends up being true, then:

I have a stripped-down scipy-only version that gets ~3.8× in about +84 lines if you'd prefer to keep it simpler

We should probably have a numba path that uses the fastest code, and a non-numba path that uses the simpler SciPy code that should work for everyone. We do this sort of thing elsewhere already, so you should be able to pick up on how to test this properly for equivalence across the two implementations (see for example stats where we triage based on Numba presence a bit). In particular numba_conditional does some magical swapping of Numba functions to their non-Numba fallbacks so that the same test is run twice on systems with and without Numba -- once with the Numba code, and then without Numba (what we call in the param "Numpy", though your code using SciPy will make this a bit of a misnomer it's okay).

[sharifhsn]

The scipy fallback is implemented. I did another benchmark and honestly it might not be worth it to have a Numba path.

Per-call in the clustering function, for 200 iterations, numba is the clear winner.

Method	Per-call (ms)	vs BFS	vs SciPy
BFS (original)	17.91	1.00x	3.2x slower
SciPy (new fallback)	5.61	3.2x faster	1.00x
Numba (union-find)	1.39	12.9x faster	4.0x faster

But using an actual workflow, spatio_temporal_cluster_1samp_test with 200 permutations, Amdahl's law makes it seem not really worth it. The simple t-test optimization impacts both paths and ends up being a bigger deciding factor.

Command	Mean [s]	Min [s]	Max [s]	Relative
BFS (original)	5.650 ± 0.063	5.579	5.698	5.39 ± 0.08
SciPy (new fallback)	1.196 ± 0.050	1.154	1.251	1.14 ± 0.05
Numba (union-find)	1.048 ± 0.012	1.039	1.061	1.00

Keep in mind that the numba path is about 100 lines of extra code, compared to the 50 line scipy path.

[larsoner]

200 extra lines for a 4x speedup in real-world use cases seems worth it to me

@wmvanvliet ?

[larsoner]

Ohhh wait the actual use case is the second where where it's ~10% speedup. Agreed that's probably not worth it.

My vote would be to remove the numba code for now but link in the code to this GitHub URL so we can find it again easily in case we want to try it in the future

[sharifhsn]

I agree, I've removed the code. Here's a permalink to the Numba path commit.

[wmvanvliet]

We can look at the speed up, but could a human please look at this and not just an AI? Currently the new code does not integrate nicely with the existing code, and the naming and abstraction boundaries are not good. Hence my earlier question: what is the core idea of the actual speedup? If we know that, we can implement it much cleaner I think.

[sharifhsn]

Apologies for the issues. I understand it must be frustrating to deal with AI-generated PRs. I have cleaned up the code to reduce complexity.

As for correctness: the two optimizations are pretty easy to understand.

1. Precomputing sum of squares: this relies on the fact that squared signs are always 1, so we can rearrange the computation of the t-statistic like so:

The sample variance formula can be written as $\text{Var}(y) = \frac{\sum y_i^2 - n\bar{y}^2}{n-1}$. Since $y_i = s_i \cdot x_i$ and $s_i^2 = 1$, we get $\sum y_i^2 = \sum x_i^2$ regardless of which signs are picked. So we precompute $sum = \sum x_i^2$ once, and each permutation only needs to compute $\bar{y} = (\mathbf{s} \cdot \mathbf{x}) / n$ (a single BLAS call) and derive $t$ from $sum$ and $\bar{y}$ algebraically. This replaces three full passes over the data matrix that have element-wise computation with one dot product + cheap arithmetic.

2. BFS: the algorithm is the exact same, BFS through an adjacency graph, it's just done through scipy instead of Python natively. The numba implementation used union-find, which is more complex for only a slight gain in performance with respect to the actual workflows.

[larsoner]

It would indeed be great to move to using connected_components! I also like the speedup for ttest_1samp, but have some ideas for a cleaner implementation

FYI at the MNE level we have decided to allow AI-facilitated PRs, and now that this is simple/small enough to be understandable, my sense is that it's worth seeing if we can get this over the finish line given the performance improvements. Okay with you @wmvanvliet ?

[larsoner]

We shouldn't need this, as order is something that we construct

[larsoner]

Suggested change

	t_obs_surr = stat_fun(X)
	# Set X back to previous state (trade memory eff. for CPU use)
	X *= signs
	try:
	t_obs_surr = stat_fun(X)
	finally:
	# Set X back to previous state (trade memory eff. for CPU use)
	X *= signs

[larsoner]

Instead of doing this, I think it would be nicer to have some pattern that uses a class. So something like
a class TTestReorder that we instantiate, and stores inv_n, neg_n, etc. We won't be able to use signs @ X but instead later still X *= signs followed by stat_fun(X) but it should be almost as fast, and require only one conditional at the beginning like:

class TTestReordered:
    def __init__(self, X):
        self.sum_sq = ...
        self.sqrt_n_nm1 = ...
        ...

    def __call__(self, X):
        dot = np.sum(X, axis=...)
        ...

and here

if stat_fun is TTestReordered:
    stat_fun = TTestReordered(X)

and then the code below gets simplified (no extra conditional on use_fast_ttest

[larsoner]

The only place _get_clusters_st_1step is used is in this function, so now that its usage is gone, it should be deleted entirely. See on main:

$ git grep _get_clusters_st_1step
mne/stats/cluster_level.py:def _get_clusters_st_1step(keepers, neighbors):
mne/stats/cluster_level.py:            return _get_clusters_st_1step(keepers, neighbors)

Same is true for _get_clusters_st_multistep

[larsoner]

Our neighbor lists are usually themselves derived from adjacency which itself is typically a sparse matrix or array. So maybe there is another optimization to be done where we don't pass neighbors here but instead just the adjacency itself

[larsoner]

Okay I've gone through and made my suggested edits. I also wrote some unit tests (drafted with the help of Claude Opus 4.8) that caught a real bug in main as well 🤦 And I had Opus benchmark our real-world tutorials and the speedup there is ~2x 🚀

Marking for merge-when-green, thanks for iterating on this @sharifhsn !

[welcome]

🎉 Congrats on merging your first pull request! 🥳 Looking forward to seeing more from you in the future! 💪