Repository navigation
perf: numbafy spatial_autocorr permutations - #1303
Open
selmanozleyen wants to merge 54 commits into
Open
selmanozleyen wants to merge 54 commits into
selmanozleyen wants to merge 54 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.95% 78.78% -0.18%
==========================================
Files 63 63
Lines 9439 9574 +135
Branches 1574 1596 +22
==========================================
+ Hits 7453 7543 +90
- Misses 1469 1521 +52
+ Partials 517 510 -7
🚀 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.
selmanozleyen
marked this pull request as draft
October 12, 2026 09:51
The observed score now comes out of the same kernel as the permuted ones, so ties are bit-exact. A sparse gene whose non-zero cells touch no other non-zero cell sits at the floor of its null and ties most permutations; folding sims >= obs alone then gives every such gene the minimum p-value (1/(n_perms+1)). Before, float noise between scanpy's observed and permuted scores split those ties at random. Take the smaller of the two tails with ties counted in both, which equals the old fold when nothing ties.
# Conflicts: # src/squidpy/gr/_ppatterns.py
Instead of a fixed 2**28-entry block, the buffer is working_memory worth of int32, as in scikit-learn: a number of MiB or a string like '16GB'. The default is an eighth of the memory the process may use (cgroup-aware, so SLURM and container limits count), never below the old block size. Every feature is revisited once per block, so at 100M cells the fixed block meant a pass per two permutations.
A fixed default, the old block size, instead of an eighth of free memory: predictable, and no cgroup probing. The buffer is still clipped to n_perms x n_cells, so small inputs allocate only what they need.
Observed score from the same kernel pass as the permutations; NaN pval_sim for constant features; canonical CSR features; one byte constant for the default permutation buffer; drop the dead numba thread cap, the explicit full-gene-list fast path and duplicated tests.
…es, which is being deprecated
flying-sheep
reviewed
Nov 3, 2026
flying-sheep
left a comment
flying-sheep
left a comment
Member
There was a problem hiding this comment.
Can’t really review the algorithm as I don’t know it, but the memory budged is awesome.
Only question os why the p_norm tests are gone, but there’s probably a good reason.
| # depend on how the features were split across workers | ||
| df_parallel = df_parallel.loc[df.index] # align in case the stat-based sort ties differently | ||
| for col in ("pval_sim", "pval_z_sim", "var_sim"): | ||
| np.testing.assert_allclose(df[col].values, df_parallel[col].values, atol=1e-12) |
Member
There was a problem hiding this comment.
use the built-in subtests fixture here to see what actually failed (if something does)
| import numba | ||
| import numpy as np | ||
| import xarray as xr | ||
| from humanfriendly import InvalidSize, parse_size |
| def _p_value_calc( | ||
| score: NDArrayA, | ||
| sims: NDArrayA | None, | ||
| sims: dict[str, Any] | None, |
Member
There was a problem hiding this comment.
just use a private dataclass, it’s nicer to use for autocomplete
…lmanozleyen/squidpy into perf/autocorr-reduce
for more information, see https://pre-commit.ci
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
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.
my notes:
disclaimer AI report:
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.