Out-of-core processing (qc/normalize/log1p/hvg/pca/pseudobulk) with deterministic parallelism

[iandriver]

Summary

Implements native out-of-core (disk-backed) processing for SingleRust — the flagship item from
the performance roadmap. The whole Tier-1 pipeline (QC → normalize_total → log1p → HVG → PCA) plus
decoupler-style pseudobulk aggregation now run streamed against an .h5ad in bounded
memory, the larger-than-RAM niche scanpy fills with Dask (where, per scanpy's own issues, the
sparse path is weak). All compute-bound passes are parallelized with a determinism guarantee:
results are bit-identical regardless of thread count.

Stacked on demo/scverse-integration, so this PR's diff is just the OOC engine + benchmarks.

What's new

OOC engine (src/backed/processing/) — streams X in row-chunks (x().iter) and writes back
via set_x_from_iter (extendable HDF5 dataset); peak memory ≈ one chunk + small accumulators:

• transformation: log1p_backed, normalize_total_backed
• qc: qc_metrics_backed (one pass; obs/var written in place)
• pipeline: preprocess_backed — fused QC+normalize+log1p (the per-cell total is reused as the
  normalization row-sum, so it's computed once)
• hvg: highly_variable_genes_backed (Seurat) — shares seurat_select with the in-memory path
• pca: pca_backed — exact covariance/eigendecomposition PCA over the HVGs (pass-1 Gram matrix,
  symmetric eigendecomp, pass-2 projection → obsm["X_pca"])
• pseudobulk: pseudobulk_backed — sample × group aggregation (sum/mean) in a single
  streaming sparse scatter-add, matching dc.pp.pseudobulk's outputs (obs["psbulk_cells"],
  obs["psbulk_counts"], layers["psbulk_props"], group-major row order)

Deterministic parallelism — naive parallel FP reduction is not reproducible (non-associative
addition + work-stealing), so two schemes are used depending on the access pattern:

• Reductions (backed::processing::det): fixed-block, ordered-merge — rows are cut into
  fixed-size blocks at fixed indices, each folded sequentially, partials merged in block order. The
  summation order depends only on data size + a constant block size, never on threads/scheduling.
  Applied to PCA Gram, HVG sum/sumsq, and QC accumulation.
• Scatter-add (pseudobulk): groups are partitioned across a fixed 16 buckets (group % 16),
  each owned by one thread. Partitions are disjoint slices of the output (no data race, no unsafe,
  no per-thread copy of the full accumulator), and every group is summed by a single thread in cell
  order — so no cross-thread merge ever happens. The partition count is a constant, so results
  reproduce across machines, not just across thread counts.

Drop-in scanpy interface — examples/sr_ooc.rs CLI
(qc|normalize_total|log1p|preprocess|hvg|pca|pseudobulk, in-place via temp+rename) and
demo/singlerust.py mirroring sc.pp.*:

scanpy / decoupler	SingleRust (out-of-core)
sc.pp.calculate_qc_metrics(adata)	sr.pp.calculate_qc_metrics(adata)
sc.pp.normalize_total(adata, target_sum=1e4)	sr.pp.normalize_total(adata, target_sum=1e4)
sc.pp.log1p(adata)	sr.pp.log1p(adata)
sc.pp.highly_variable_genes(adata, …)	sr.pp.highly_variable_genes(adata, …)
sc.pp.pca(adata, …)	sr.pp.pca(adata, …)
dc.pp.pseudobulk(adata, sample_col, groups_col)	sr.pp.pseudobulk(adata, sample_col, groups_col)

adata is a path or a backed AnnData; the shim releases the HDF5 handle before invoking Rust.

Benchmarks — demo/bench_ooc.py (3 lanes: scanpy in-memory vs scanpy+Dask vs SingleRust OOC)
and demo/bench_pseudobulk.py (decoupler vs SingleRust), each lane run under /usr/bin/time -l
for peak RSS.

Results

All on 500k cells × 48,788 genes (48 GB / 18-core).

Preprocessing (QC + normalize_total + log1p):

lane	time	peak RSS
scanpy in-memory	13.5 s	13.6 GB
scanpy + Dask OOC	35.4 s	14.5 GB
SingleRust OOC (fused)	45.0 s	2.5 GB

SingleRust OOC uses ~5.4× less memory than scanpy in-memory and ~5.8× less than scanpy+Dask,
with a chunk-bounded footprint that stays flat as cells grow (scanpy OOMs ~46 GB at 2M). Dask's
sparse OOC path gave no memory benefit (RSS ≥ in-memory) at 2.6× the time — the immaturity
scanpy's issues describe.

Pseudobulk, scaled 500k → 2M cells (12 donors × 143 cell types = 1,716 groups held constant, so
this isolates cell-count scaling; 1M/2M inputs are the 500k dataset tiled along obs by
demo/_replicate_cells.py):

cells	decoupler time	decoupler peak RSS	SingleRust time	SingleRust peak RSS
500k	29.9 s (compute)	21.6 GB	12.2 s (wall, incl. I/O)	3.4 GB
1M	167.8 s ⚠	19.2 GB ⚠	19.3 s	2.4 GB
2M	not run (see caveats)	—	26.1 s	3.6 GB

SingleRust's memory stays flat (~2.4–3.6 GB) from 500k to 2M and its time is sub-linear — 4× the
cells costs 2.1× the time. Both follow from the design: peak RSS is the fixed groups × genes
accumulator plus one chunk (neither depends on cell count), and the fixed output allocate/write cost
amortizes as cells grow. decoupler loops the full sample×group cartesian product and, for each,
boolean-masks all cells and densifies the submatrix — O(n_obs · n_groups) masking plus
per-group densification with the whole input resident — so its cost grows with cells even at
constant group count.

Results are exact at every size: aggregate-sum max abs diff = 0 vs decoupler at 1M over all
1,716 groups. Tiling also gives a free self-check — the 1M/2M sums must equal exactly 2×/4× the 500k
sums, and they do (diff 0, psbulk_cells likewise).

⚠ Please read the caveats below before quoting the decoupler comparison.

Correctness & determinism

Each step is unit-tested against its same-algorithm counterpart: log1p/normalize match scanpy's X
exactly; QC matches; HVG selects the same genes as in-memory HVG; PCA matches a dense
covariance-PCA reference (variance ratios exact, embedding norms match up to sign); pseudobulk
matches decoupler's aggregate sums exactly.

Determinism verified two ways:

• unit tests run QC/HVG/preprocess/PCA/pseudobulk under 1 vs 8 rayon threads → bit-identical;
• at scale, the full pipeline on 50k real cells and pseudobulk on 500k real cells, each with
  RAYON_NUM_THREADS=1 vs 18 → bit-identical X, obs/var, obsm["X_pca"], variance ratios,
  and pseudobulk aggregates/counts/props.

Benchmark caveats

Recording these rather than smoothing them over — the SingleRust columns are clean, but two of the
decoupler figures are not, and one previously-reported number was wrong.

• ⚠ The 1M decoupler run is contaminated — don't quote a ratio from it. An unrelated 8–16 GB
  process was on the machine for part of it and decoupler was paging: 185 s wall against only 122 s
  CPU (71 user + 51 sys). Its true quiet-machine time is lower than 167.8 s. Read the 1M row as
  "decoupler degrades sharply once it stops fitting," not as a precise 8.7×.
• ⚠ Peak RSS understates decoupler's demand at 1M. Its resident set fell from 13.5 GB to ~5 GB
  mid-run as the OS evicted pages, which is why 1M reads lower than 500k despite twice the data.
  Once a lane swaps, RSS is a ceiling on residency, not on demand.
• The 500k SingleRust figure is corrected. An earlier revision of this PR reported 5.1 s /
  4.6 GB. That run followed decoupler over the same file and so read from a warm page cache; cold,
  it is 12.2 s. The honest gap at 500k is ~2.4×, not ~5.9×.
• 2M decoupler was not run. It needs ~24 GB for the CSR matrix before any aggregation; at the
  time a neighbor process held 16.2 GB and swap was 98% used, so attempting it risked destabilizing
  the machine rather than producing a usable number.
• Time remains not apples-to-apples, in decoupler's favor on I/O: its number is compute-only
  (excludes its data load), SingleRust's is wall time including read and write.

Notes for reviewers

• Adds rayon as a direct dependency. 23 lib tests pass; clippy clean (one pre-existing dead-code
  warning unrelated to this PR).
• nnz overflows int32 past ~1.4M cells at this density. The 2M input (nnz = 3.01e9) requires an
  int64 CSR indptr; SingleRust reads it fine, but it is a real cliff for any 32-bit index path.
• The pseudobulk partition buffers cost ~1.4 GB over the sequential version in exchange for a ~1.6×
  speedup; this scales with group count, so it is negligible for the coarser groupings typical of
  pseudobulk DE.
• End-to-end PCA axes differ from scanpy only because SingleRust's Seurat HVG selects a different
  gene set than scanpy's — a pre-existing in-memory difference, not introduced here.
• Surfaced (and flagged separately) a pre-existing convert_to_array_f64 CSR densify bug; the OOC
  tests densify manually to avoid it.
• The Dask benchmark lane needs an isolated .venv-dask (anndata ≥ 0.11 / scanpy ≥ 1.11); the main
  3.9 demo env is untouched.

🤖 Generated with Claude Code