anndata-oom

Out-of-memory AnnData powered by Rust — a drop-in replacement for anndata.AnnData that keeps the expression matrix on disk and runs entire preprocessing pipelines (normalize, log1p, scale, PCA, neighbors, UMAP, Leiden) as lazy transforms or chunked operations. The full matrix is never loaded into memory.

Built on top of scverse/anndata-rs, the Rust implementation of AnnData.

Why?

Standard anndata.AnnData loads the entire expression matrix into RAM. For a million-cell atlas this can mean 100+ GB of memory — beyond what most workstations have.

anndataoom keeps X on disk (HDF5) and streams it through the preprocessing pipeline in chunks. Peak RAM is independent of dataset size.

Memory comparison

How?

Each preprocessing step adds a small "transform descriptor" (a vector or flag) to a lazy computation chain. Data is computed on-the-fly during chunked reads from the HDF5 file:

X (HDF5 on disk, Rust I/O via anndata-rs)
  → TransformedBackedArray      (normalize: ÷ per-cell size factors)
    → TransformedBackedArray    (log1p: on-the-fly)
      → _SubsetBackedArray      (HVG: select 2,000 gene columns)
        → ScaledBackedArray     (z-score: stores only mean/std vectors)
          → Randomized SVD      (chunked matrix products)
            → X_pca             (n_obs × 50, in memory)
              → Neighbors / UMAP / Leiden (operate on X_pca only)

For a deeper look at the lazy operator model behind this — how nodes compose, how data flows through the chain, and why subsetting stays O(1) — see docs/transform-chain.md.

Documentation

• The Transform Chain — design walkthrough with ASCII diagrams of the lazy operator model (BackedArray, _SubsetBackedArray, TransformedBackedArray, ScaledBackedArray) and how chunked reads flow through it.

Installation

Prebuilt wheels (recommended)

pip install anndataoom

Platform support (verified end-to-end by the CI matrix — builds the Rust extension from source and runs the full test suite on each cell):

Windows support landed in 0.1.8. anndata-rs is now vendored in-tree (vendor/anndata-rs/) and its anndata-hdf5 gates the HDF5 threadsafe feature off on Windows — HDF5's CMake refuses thread-safety with a static library there, which previously broke the build. The vendored static HDF5 now compiles and the full test suite passes on Windows in CI. Linux aarch64 wheels are still pending (the ring crate fails to cross-compile).

Wheels bundle a statically-linked HDF5 — on the supported platforms no system dependencies are needed and no Rust toolchain is required.

Build from source

If no prebuilt wheel matches your system, pip falls back to source. You'll need a Rust toolchain:

# Install Rust (if needed)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y

# Build + install
pip install anndataoom

Or for development:

git clone https://github.com/Starlitnightly/anndata-oom
cd anndata-oom
pip install maturin
maturin develop --release

Quick start

import anndataoom as oom

# Read an h5ad file — matrix stays on disk
adata = oom.read("large_dataset.h5ad")
print(adata)

AnnDataOOM                                 [Rust · out-of-memory · backed]
Dimensions:  n_obs: 100,000    n_vars: 30,000

┌───────────┬──────────────────────────────────────┐
│ File      │ large_dataset.h5ad  (1.2 GB on disk) │
│ X         │ csr_matrix · float32 · 5.3% density   │
│ Chunk I/O │ ~20 MB per 1,000-row chunk            │
└───────────┴──────────────────────────────────────┘

▸ obs     (8)    batch · cell_type · n_counts · ...
▸ var     (3)    gene_name · highly_variable · ...
▸ obsm    (–)
▸ layers  (–)
▸ raw     (–)

Chunked operations

# Sum over all cells — streams the matrix in 1000-row chunks
row_sums = adata.X.sum(axis=1)       # ndarray of shape (n_obs,)

# Per-gene means — one-pass chunked Welford's
mean, var = oom.chunked_mean_var(adata)

# Iterate chunks manually
for start, end, chunk in adata.X.chunked(5000):
    # chunk is a csr_matrix (or ndarray) of shape (≤5000, n_vars)
    ...

Subsetting

# All forms of indexing work
sub = adata[0:1000]                          # first 1000 cells
sub = adata[adata.obs["cell_type"] == "B"]   # boolean mask
sub = adata[:, ["GENE1", "GENE2"]]           # by gene name
sub = adata[:, adata.var["highly_variable"]] # after HVG selection

# Returns a new AnnDataOOM — still lazy
print(sub.shape)   # e.g. (17003, 2000)

Single-gene access

# obs_vector reads exactly one column from disk
expr = adata.obs_vector("CD3D")   # ndarray of shape (n_obs,)

Integration with omicverse

omicverse automatically detects anndataoom and uses it as the backend for ov.read(..., backend="rust"):

import omicverse as ov

# Read — returns AnnDataOOM if anndataoom is installed
adata = ov.read("data.h5ad", backend="rust")

# Full preprocessing pipeline — all chunked/lazy
adata = ov.pp.qc(adata,
                 tresh={"mito_perc": 0.2, "nUMIs": 500, "detected_genes": 250},
                 doublets=False)
adata = ov.pp.preprocess(adata, mode="shiftlog|pearson",
                         n_HVGs=2000, target_sum=50 * 1e4)

# HVG subset — returns a new AnnDataOOM
adata.raw = adata
adata = adata[:, adata.var.highly_variable_features]

# Scale + PCA — lazy z-score + chunked randomized SVD
ov.pp.scale(adata)
ov.pp.pca(adata, layer="scaled", n_pcs=50)

# Neighbors / UMAP / Leiden — operate on obsm['X_pca'], no matrix touch
ov.pp.neighbors(adata, n_neighbors=15, n_pcs=50,
                use_rep="scaled|original|X_pca")
ov.pp.umap(adata)
ov.pp.leiden(adata, resolution=1)

# Plotting — all ov.pl.* functions work directly, incl. use_raw=True
ov.pl.embedding(adata, basis="X_umap", color="leiden")
ov.pl.dotplot(adata, marker_genes, groupby="leiden")
ov.pl.violin(adata, keys="CD3D", groupby="leiden", use_raw=True)

CPU–GPU mixed mode

anndataoom (storage backend) is orthogonal to omicverse's execution mode (compute), so you can flip on GPU acceleration with an OOM-backed AnnData and pay no penalty:

import omicverse as ov
ov.settings.cpu_gpu_mixed_init()       # route PCA/neighbors/UMAP to torch-GPU
adata = oom.read("data.h5ad")          # still out-of-memory
# … identical qc → preprocess → scale → pca → neighbors → umap pipeline …

What changes (measured on the Tabula Sapiens benchmark, H100):

• Memory-bound preprocessing is mode-invariant. qc → preprocess → scale → pca run at the same wall-clock and RSS in cpu and cpu-gpu-mixed, with zero GPU memory used. This is by design: anndataoom's chunked operators are pure-CPU, and omicverse routes the OOM PCA through anndataoom.chunked_pca (never the torch-GPU solver). Peak RSS stays flat and bounded by chunk size regardless of mode.
• Downstream embedding gets the GPU. Steps that operate on the small (n_obs × 50) PCA embedding do offload: ov.pp.neighbors runs on a CUDA PyG-kNN backend (device memory allocated) and ov.pp.umap, ov.pp.mde run on the GPU. The kNN itself is sub-second on a warm GPU; the exact wall-clock gain at small scale is measurement-sensitive, so we just note the offload is real. PCA results are bit-identical between modes (|cos| = 1.0).
• In-memory backend, by contrast, gets a big PCA win. With a plain anndata.AnnData (not OOM), ov.pp.pca dispatches to the GPU torch_pca solver: PCA dropped 13× (13.9 s → 1.1 s on TS-5k), ~1.6× on the whole pipeline. anndataoom forgoes this by design — it routes through CPU chunked_pca to keep peak RSS flat.

omicverse function compatibility

Compatibility of ov.pp.* against an AnnDataOOM backend — 14 of 24 probed functions run on the OOM path (probed in both cpu and cpu-gpu-mixed; ᴳ = offloads to GPU in mixed mode):

Failures raise a clear exception at the call site — they never silently mis-compute. Reproduce with benchmark/scripts/compat_matrix.py.

Full API reference

Top-level

Chunked preprocessing

AnnDataOOM methods

All anndata.AnnData methods and properties are supported. Key ones:

Benchmark: PBMC 8k (7,750 cells × 20,939 genes)

Full preprocessing pipeline (QC → normalize → HVG → scale → PCA → neighbors → UMAP → Leiden):

→ 27.8× memory savings on this small dataset; ratio grows with scale.

Supported h5ad formats

anndataoom automatically preserves sparsity through normalize and log1p (sparse → sparse), and materializes to dense only where algorithmically necessary (z-score, PCA).

Architecture

anndataoom is a thin Python wrapper over scverse/anndata-rs:

┌──────────────────────────────────────────────┐
│  anndataoom (Python package)                 │
│  ┌────────────────────────────────────────┐  │
│  │  AnnDataOOM                            │  │
│  │  ├─ obs, var (pandas.DataFrame)        │  │
│  │  ├─ obsm, varm (dict of ndarray)       │  │
│  │  ├─ layers (BackedLayers — sidecar H5) │  │
│  │  └─ X (BackedArray — wraps ↓)          │  │
│  └────────────────────────────────────────┘  │
│            │                                  │
│            ▼                                  │
│  ┌────────────────────────────────────────┐  │
│  │  anndataoom._backend  (Rust extension)│  │
│  │  ├─ AnnData (pyanndata)                │  │
│  │  ├─ PyArrayElem (chunked() iterator)   │  │
│  │  └─ Statically linked:                 │  │
│  │     ├─ anndata (Rust crate)            │  │
│  │     ├─ anndata-hdf5                    │  │
│  │     └─ HDF5 C library                  │  │
│  └────────────────────────────────────────┘  │
└──────────────────────────────────────────────┘

The Rust extension (anndataoom._backend) is pinned to a specific commit of scverse/anndata-rs for reproducible builds (the same commit used by SnapATAC2).

Limitations and caveats

• Writing back to X is lazy — modifications via adata[mask] = value materialize X in memory. Use adata.obs, adata.obsm, or adata.write(path) to persist changes.
• PCA accuracy: chunked_pca uses randomized SVD with 4 power iterations — top ~20 PCs are nearly identical to standard PCA; PCs 30+ start to deviate. For publication-quality analyses, consider adata.to_adata() + standard PCA.
• Some ops require materialization: score_genes_cell_cycle, find_markers, non-Harmony batch correction, etc. These auto-materialize with a warning.
• File mode: Default backed='r' (read-only) protects the source file. Use backed='r+' if you need to write back (advanced).
• Concurrent access: HDF5 files default to exclusive locking. Set HDF5_USE_FILE_LOCKING=FALSE in the environment if multiple processes need to read the same file.

Comparison with alternatives

Development

git clone https://github.com/Starlitnightly/anndata-oom
cd anndata-oom

# Install Rust (first time)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
source $HOME/.cargo/env

# Build in editable mode
pip install maturin
maturin develop --release

# Run tests
pip install pytest
pytest tests/

Contributing

Contributions welcome! Areas of interest:

• More lazy transforms: regress-out, harmony, scVI integration
• Zarr backend: currently only HDF5 supported
• Dask interop: expose BackedArray as a dask.array
• Query engine: SQL-like filtering over chunks

Release process

1. Bump version in pyproject.toml and Cargo.toml
2. Update CHANGELOG.md
3. Commit, tag, push:

   git commit -am "Release v0.x.0"
   git tag v0.x.0
   git push && git push --tags
   

4. GitHub Actions builds wheels for all platforms and publishes to PyPI (via trusted publishing)

License

MIT License — see LICENSE.

Built on scverse/anndata-rs (MIT, © Kai Zhang).

Citation

If you use anndataoom in published research, please cite:

@software{omicverse,
  title  = {OmicVerse: A framework for multi-omic data analysis},
  author = {Zeng, Z. et al.},
  url    = {https://github.com/Starlitnightly/omicverse},
  year   = {2024},
}

@software{anndata_rs,
  title  = {anndata-rs: Rust implementation of AnnData},
  author = {Zhang, Kai},
  url    = {https://github.com/scverse/anndata-rs},
  year   = {2022},
}