Quickstart

annbatch turns a collection of AnnData files into a shuffled, on-disk zarr collection and streams shuffled mini-batches from it. This lets you train on data that does not fit in memory while keeping the GPU busy.

This quickstart uses a small synthetic dataset to show the three steps every workflow shares: convert → load → iterate. For real, end-to-end walkthroughs on actual data, see the guides by data type:

• Single-cell RNA-seq — single-cell RNA-seq count matrices

• Genetics (VCF / whole-genome sequencing) — genotypes from VCF / whole-genome sequencing

• Single-cell microscopy images — single-cell microscopy images

• Distributed / multi-GPU training — multi-GPU / distributed training

Configure zarrs

Pointing zarr at the zarrs-python codec pipeline is what makes local reads fast. Do this once, before reading or writing any store.

import warnings
import zarr

for msg in ("Consolidated metadata is currently not part in the Zarr format 3 specification.*",):
    warnings.filterwarnings("ignore", message=msg)

zarr.config.set({"codec_pipeline.path": "zarrs.ZarrsCodecPipeline"})

A small example collection

Real workflows start from your own .zarr files. Here we generate two tiny ones with sparse counts in X and a cell_type label in obs.

import anndata as ad
import numpy as np
import pandas as pd
import scipy.sparse as sp

n_genes = 200
var = pd.DataFrame(index=[f"gene_{i}" for i in range(n_genes)])


def _make_adata(n_cells: int, seed: int) -> ad.AnnData:
    rng = np.random.default_rng(seed)
    counts = rng.poisson(0.3, size=(n_cells, n_genes)).astype("float32")
    obs = pd.DataFrame(
        {"cell_type": pd.Categorical(rng.choice(["B", "T", "NK"], n_cells))},
        index=[f"cell{seed}_{i}" for i in range(n_cells)],
    )
    return ad.AnnData(X=sp.csr_matrix(counts), obs=obs, var=var.copy())


# AnnData 0.13.0 defaults to the following settings
ad.settings.zarr_write_format = 3
ad.settings.auto_shard_zarr_v3 = True

adata_paths = []
for i, n_cells in enumerate([3000, 2000]):
    path = f"example_{i}.zarr"
    _make_adata(n_cells, seed=i).write_zarr(path)
    adata_paths.append(path)

Convert into a shuffled on-disk collection

add_adatas() aligns the var spaces, shuffles cells across the inputs, and writes a sharded zarr collection. Shuffling at write time is what lets the loader read large contiguous chunks at train time without sacrificing randomness.

from annbatch import DatasetCollection

collection = DatasetCollection(zarr.open("quickstart_collection.zarr", mode="w"))
collection.add_adatas(adata_paths=adata_paths, shuffle=True)

Show code cell output

Hide code cell output

<annbatch.io.DatasetCollection at 0x761dcd37ecf0>

Stream shuffled mini-batches

The Loader is a drop-in replacement for a PyTorch DataLoader. Passing a load_adata that keeps only the columns you need (here cell_type) keeps loading fast.

from annbatch import Loader


def _load_adata(group: zarr.Group) -> ad.AnnData:
    return ad.AnnData(
        X=ad.io.sparse_dataset(group["X"]),
        obs=ad.experimental.read_lazy(group).obs[["cell_type"]].to_memory(),
    )


loader = Loader(batch_size=512, chunk_size=64, preload_nchunks=32, preload_to_gpu=False).use_collection(
    collection, load_adata=_load_adata
)

Each batch yields a sparse X and the requested obs columns. Convert the sparse batch to dense on the GPU (.cuda() then .to_dense()). This is much faster than densifying on the CPU.

for batch in loader:
    x = batch["X"].cuda().to_dense()
    cell_type = batch["obs"]["cell_type"]
    # ... your training step here ...
    break

print(f"batch X: {tuple(x.shape)} {x.dtype} on {x.device}")
print(f"labels:  {cell_type[:5].tolist()}")

batch X: (512, 200) torch.float32 on cuda:0
labels:  ['T', 'NK', 'NK', 'NK', 'NK']