Repository navigation
perf: numbafy spatial_autocorr permutations - #1303
Draft
selmanozleyen wants to merge 33 commits into
Draft
selmanozleyen wants to merge 33 commits into
selmanozleyen wants to merge 33 commits into
Conversation
selmanozleyen
force-pushed
the
perf/numbafy-spatial-autocorr
branch
3 times, most recently
from
October 1, 2026 12:15
4667f60 to
bca5e96
Compare
Ruff's PLR0917 (preview, selected explicitly) now flags any function with more than five positional arguments. - Public functions take only their data object positionally; everything after it is keyword-only. Released functions keep accepting the old positional calls through legacy-api-wrap, with a FutureWarning naming the arguments. Experimental functions don't: that API promises no stability. - Private functions change only where the rule flags them: arguments with a default become keyword-only, and helpers whose arguments are all required keep only their first argument positional. Numba kernels are exempted, since numba needs positional arguments. - Left as on main: the deprecated calculate_niche and its validator (removed in v1.9.0), and ImageContainer until it is deprecated. - Internal and test callers pass the moved arguments by name. Tests are exempted from the rule, since pytest passes fixtures by name. Preview mode also caught an undefined apply_stain_normalization in squidpy.experimental.im.__all__, removed here. Also documents these rules and the rng conventions in the contribution guide, and drops the unused docs/release/notes-dev.md: release notes come from the PR changelog.
… name v1.9.0 in the warning
selmanozleyen
force-pushed
the
chore/kwonly-args
branch
from
October 1, 2026 12:22
e523269 to
75b4beb
Compare
selmanozleyen
force-pushed
the
perf/numbafy-spatial-autocorr
branch
from
October 1, 2026 12:22
bca5e96 to
02ab58a
Compare
One njit(parallel=True, nogil=True) kernel over the permutations replaces the loky pool. Both statistics sum, over the rows of g[perm, :], a per-row quantity the permutation only reindexes, so accumulating it once leaves O(n_cells) per permutation instead of O(nnz): 21-34x end to end. pval_sim, var_sim and the FDR columns are bit-identical; the continuous columns agree to <=5e-10 relative, which is the float reordering scanpy's own kernels already have across thread counts. Reductions are explicit serial loops so the output no longer depends on n_jobs. Sparse vals arrive as CSC, whose row slicing is O(nnz); one tocsr saves 7.9s of a 19s call at 40k cells x 2000 genes. backend is deprecated as co_occurrence and nhood_enrichment deprecate it, and n_jobs now caps the numba thread pool.
…ermutations, keep backend in the shim
selmanozleyen
force-pushed
the
perf/numbafy-spatial-autocorr
branch
from
October 5, 2026 12:26
02ab58a to
b97a1f3
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1303 +/- ##
==========================================
- Coverage 78.88% 78.58% -0.31%
==========================================
Files 65 65
Lines 9611 9722 +111
Branches 1591 1612 +21
==========================================
+ Hits 7582 7640 +58
- Misses 1497 1550 +53
Partials 532 532
🚀 New features to boost your workflow:
|
The permutation kernel indexes indptr/indices/data as CSR, and sklearn's row-normalize silently no-ops on a CSC graph. tocsr(copy=True) keeps the copy normalize(copy=False) needs.
Each k owns its output slot and sums its own row in index order, so prange does not reorder any accumulation. Bit-identical; worth 1.4-1.7x once the permutations are blocked and the serial precompute repeats.
selmanozleyen
marked this pull request as draft
October 9, 2026 16:42
Both statistics split into a term that only touches a feature's non-zero cells and a term the permutation leaves invariant, so a permutation costs O(nnz_feature) rather than O(n_cells) and Geary's wx2 term leaves the loop entirely. The precompute scatters g's columns for the same reason, which also means a feature is never densified. `w_sum` and `col_sum` do not depend on the feature, so they move out of the kernel; the O(n_cells) running totals become O(nnz) through `col_sum`. Features are independent and the kernel is nogil, so parallelism moves up to a thread pool over features -- measured faster than an in-kernel prange even when there are fewer features than workers. Permutation draws fan out the same way. The permutation scores are only ever reduced along the permutation axis, so `_score_perms` keeps a tally and a running mean/M2 per feature instead of the (n_perms, n_features) matrix. pval_sim and the FDR columns stay bit-identical; var_sim and pval_z_sim move by ~1e-11 relative on Moran's I, whose expectation sits near zero.
The kernel is compiled serial now, the graph transpose is unconditional, and the block cap is 2**28, so the notes describing the previous shapes were contradicting the code beside them. Inline the one-use source function into its decorator and drop the kernel alias.
The permuted scores cluster around the observed one, so shifting by it turns the mean and variance into plain sums: they merge across blocks with `+=` instead of a Chan combine, and `s2/n - (s1/n)**2` has nothing to cancel because `s1/n` is the small residual. Also more accurate than the combine it replaces -- 2e-15 against 2e-09 relative on tightly clustered scores, where the naive `E[x^2] - E[x]^2` is wrong outright.
`extract_X` skips the subset when the selection is every gene in var_names order. Without the order check a reordered full-length list takes that path too, returning X in var order against the caller's labels. Nothing caught that; this does.
The constant said what it was, not why that value. It exists to keep memory off `n_perms` -- 1000 permutations over 100M cells would want 372 GB unbatched -- and 2**28 comes from measuring the cost of splitting, not from any cache size: ~1.03x at 1M cells, ~1.10x at 10M, and 0.86x at 10M with two blocks, where splitting is the faster arm. Also corrects the reduction note: the kernel is serial, so there is no parallel reduction to avoid, and independence from `n_jobs` comes from the RNG binding and per-feature ownership instead.
The ratios are what transfer; the byte counts were specific to the machine and dataset they were measured on.
Seventeen lines of justification for one constant was more than it needs.
…nto perf/autocorr-reduce
`_build` and `_sepal` already map over a thread pool through this helper; `_score_perms` was managing its own. Measured equal or slightly faster interleaved, so the pool it builds per call costs nothing here. Also notes that the v1.8.3 positional test goes when the `backend` shim does, since the signature it pins only exists for that shim.
The `_score_perms` docstring still described a running mean and M2, which the shifted sums replaced. The `extract_X` note claimed `genes is None` means every gene, which is false when `highly_variable` is set. `nogil` was credited with avoiding the `parallel=True` abort, which is what `parallel=False` does. `arr.sum()` materializes nothing, so that argument only applied to the centred square. The fold in `_p_value_calc` was called a swap, and a test docstring still referred to numba threads.
selmanozleyen
marked this pull request as ready for review
October 10, 2026 09:23
Shifting the variance by the observed score makes it exactly zero when every permutation scores alike, where the old two-pass form left a denormal residual. `z_sim` is then NaN, and neither mask in `_p_value_calc` covers NaN, so `np.empty` left those entries at whatever memory held -- often 0.0, the most significant p-value there is. A gene expressed in one cell hits this, which is routine in raw counts. Both p-value arrays start as NaN now. `multipletests` spreads a NaN over its whole input, so the correction runs over the defined p-values alone and leaves the rest NaN. Also drops the n_jobs coverage duplicated across `seq_par`, `perm_blocks` and a separate invariance test, and the `pval_norm`/`var_norm` assertions that `assert_frame_equal` already makes.
`s1/n` is the observed score's distance from the permutation mean, so it is largest for the features with the most autocorrelation, not a small residual. The cancellation it causes is real, about eps*z**2.
The note claimed splitting only gets expensive past ~8 blocks, which held at 10M and not at 100M: the block there is two permutations, so the scatter is re-paid every two and the measured cost is 1.3-1.6x at three to five blocks. Block size drives it, not block count.
selmanozleyen
marked this pull request as draft
October 12, 2026 09:51
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
spatial_autocorr's permutations run in one numba kernel (njit(nogil=True), with features across a thread pool) instead of a loky pool. Both statistics sum, over the rows ofg[perm, :], a per-row quantity the permutation only reindexes, so it is accumulated once and each permutation costs O(nnz in the feature) instead of O(nnz).pval_simand the FDR columns are bit-identical;var_simandpval_z_simagree to ~1e-11 relative on Moran's I, whose expectation sits near zero (float reordering)n_jobs, which now sizes the thread pool;backendis deprecated like inco_occurrenceandnhood_enrichmentvalsare converted to CSR once, since row slicing CSC is O(nnz)wx2term leaves the loop entirely(n_perms, n_features)matrixFollow up:
ImageContainerremoves the last use ofparallelizeBenchmarks
Real datasets, full
sq.gr.spatial_autocorr, exclusive HMGU nodes,mainvs this branch.pval_simis bit-identical in every row.Xenium on 32 cores, MERFISH on 90. At 1000 permutations this branch still finishes ~13x faster than
maindoes at 100, so finer FDR resolution is affordable.Scaling, on MERFISH tiled by offsetting copies (so the graph stays a real neighbourhood graph), against the first two commits of this PR rather than
main:Permutations are drawn in blocks so memory does not track
n_perms; splitting costs ~3% at 1M cells and ~10% at 10M, and at two blocks it measured faster than a single buffer. 100M cells x 161 genes runs in ~96 GB.