annbatch.Loader

class annbatch.Loader(*, batch_sampler=None, chunk_size=None, preload_nchunks=None, shuffle=None, return_index=False, batch_size=None, preload_to_gpu=False, drop_last=None, to=None, rng=None)

A loader for on-disk data anndata stores.

This loader by default batches together slice requests (chunk_size parameter) to the underlying stores to achieve higher performance. You can also use chunk_size==1 for perfect random sampling (for relevant samplers), although this comes at a performance penalty for on-disk (and likely in-memory) data as well. Custom samplers are supported via the batch_sampler argument. We thus recommend using DatasetCollection to preshuffle your data (or pre-shuffling in-memory). The loader is agnostic to the on-disk chunking/sharding, but it may be advisable to align with the in-memory chunk size for dense.

When providing a custom sampler, chunk_size, preload_nchunks, batch_size, shuffle, drop_last, and rng must not be set (they are controlled by the batch_sampler instead). When providing these arguments and no batch_sampler, they are used to construct a RandomSampler (if shuffle=True) or SequentialSampler.

Parameters:

batch_sampler Sampler | None (default: None)

If not provided, a default SequentialSampler or RandomSampler will be used with the same defaults below.

chunk_size int | None (default: None)

The obs size (i.e., axis 0) of contiguous array data to fetch. Mutually exclusive with batch_sampler. Defaults to 512.

preload_nchunks int | None (default: None)

The number of chunks of contiguous array data to fetch. Mutually exclusive with batch_sampler. Defaults to 32.

shuffle bool | None (default: None)

Whether or not to shuffle the data. Mutually exclusive with batch_sampler. Defaults to False.

batch_size int | None (default: None)

Batch size to yield from the dataset. Mutually exclusive with batch_sampler. Defaults to 1.

drop_last bool | None (default: None)

Set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. Leave as False when using in conjunction with a torch.utils.data.DataLoader. Mutually exclusive with batch_sampler. Defaults to False.

rng Generator | None (default: None)

Random number generator for shuffling. Mutually exclusive with batch_sampler. Defaults to np.random.default_rng() if not provided.

return_index bool (default: False)

Whether or not to yield the index on each iteration.

preload_to_gpu bool (default: False)

Whether or not to use cupy for non-io array operations like vstack and indexing once the data is in memory internally (and batches yielded if to is not set). This option entails greater GPU memory usage, but is faster at least for sparse operations. torch.vstack() does not support CSR sparse matrices, hence the current use of cupy internally (which also means torch is an optional dep). Furthermore, there is no way to allocate pinned memory for jax arrays. Setting this to False is advisable when using the torch.utils.data.DataLoader wrapper or potentially with dense data due to memory pressure. For top performance, this should be used in conjunction with to="torch" and then torch.Tensor.to_dense() if you wish to densify (or similar for jax, although jax may be very slow for sparse->dense transform). cupy.cuda.MemoryPool.free_all_blocks() (i.e., the method of the pool of cupy.get_default_memory_pool()) is called aggressively to keep memory usage low. If you are using your own memory pool or allocator, you may have to free blocks on your own.

to Literal['torch', 'jax'] | None (default: None)

The output library for which you would like your array output.

Examples

>>> from annbatch import Loader
>>> ds = Loader(
        batch_size=4096,
        chunk_size=32,
        preload_nchunks=512,
    ).add_adata(my_anndata)
>>> for batch in ds:
        # optionally convert to dense
        # batch = batch.to_dense()
        do_fit(batch)

Attributes table

batch_sampler	The sampler used to generate batches.
dataset_type	The type of on-disk data used in this loader.
n_obs	The total number of observations in this instance i.e., the sum of the first axis of all added datasets.
n_var	The total number of variables in this instance i.e., the second axis (which is the same) across all datasets.
var	The var annotations for the variables in this loader.

Methods table

add_adata(adata)	Append an adata to this dataset.
add_adatas(adatas)	Append adatas to this dataset.
add_dataset(dataset[, obs, var])	Append a dataset to this dataset.
add_datasets(datasets[, obs, var])	Append datasets to this dataset.
use_collection(collection, *[, load_adata])	Load from an existing annbatch.DatasetCollection.

Attributes

Loader.batch_sampler

The sampler used to generate batches.

Returns:

The sampler.

Loader.dataset_type

The type of on-disk data used in this loader.

Returns:

The type used.

Loader.n_obs

The total number of observations in this instance i.e., the sum of the first axis of all added datasets.

Returns:

The number of observations.

Loader.n_var

The total number of variables in this instance i.e., the second axis (which is the same) across all datasets.

Returns:

The number of variables.

Loader.var

The var annotations for the variables in this loader.

Returns:

The var DataFrame or None if no var annotations were provided.

Methods

Loader.add_adata(adata)

Append an adata to this dataset.

Parameters:

adata AnnData

A anndata.AnnData object, with zarr.Array, scipy.sparse.csr_matrix, scipy.sparse.csr_array, numpy.ndarray, or anndata.abc.CSRDataset as the data matrix in X, and obs containing annotations to yield in a pandas.DataFrame. var must match the var of any previously added datasets. Only X, obs, and var are kept for now: any obsm and layers elements are ignored and a FutureWarning is emitted (a future release will additionally load and yield them if present in the passed in adata).

Return type:

Self

Loader.add_adatas(adatas)

Append adatas to this dataset.

Parameters:

adatas list[AnnData]

List of anndata.AnnData objects, with zarr.Array, scipy.sparse.csr_matrix, scipy.sparse.csr_array, numpy.ndarray, or anndata.abc.CSRDataset as the data matrix in X, and obs containing annotations to yield in a pandas.DataFrame. Only X, obs, and var are kept for now: any obsm and layers elements are ignored and a FutureWarning is emitted (a future release will additionally load and yield them).

Return type:

Self

Loader.add_dataset(dataset, obs=None, var=None)

Append a dataset to this dataset.

Parameters:

dataset CSRDataset | Array | csr_array | csr_matrix | ndarray

A zarr.Array or anndata.abc.CSRDataset object, generally from anndata.AnnData.X.

obs DataFrame | None (default: None)

DataFrame obs, generally from anndata.AnnData.obs.

var DataFrame | None (default: None)

DataFrame var, generally from anndata.AnnData.var. var must match the var of any previously added datasets.

Return type:

Self

Loader.add_datasets(datasets, obs=None, var=None)

Append datasets to this dataset.

Parameters:

datasets list[CSRDataset | Array | csr_array | csr_matrix | ndarray]

List of zarr.Array or anndata.abc.CSRDataset objects, generally from anndata.AnnData.X. They must all be of the same type and match that of any already added datasets.

obs list[DataFrame] | None (default: None)

List of DataFrame for annotating observations (i.e., samples), generally from anndata.AnnData.obs.

var list[DataFrame] | None (default: None)

List of DataFrame for annotating features, generally from anndata.AnnData.var. All var DataFrames must be identical.

Return type:

Self

Loader.use_collection(collection, *, load_adata=<function load_all_aligned>)

Load from an existing annbatch.DatasetCollection.

This function can only be called once. If you want to manually add more data, use Loader.add_adatas() or open an issue.

Parameters:

collection DatasetCollection

The collection whose on-disk datasets should be used in this loader.

load_adata Callable[[Group], AnnData] (default: <function load_all_aligned at 0x743989a28720>)

A custom load function - recall that whatever is found in X and obs will be yielded in batches. Only X, obs, and var are yielded for now; a future release will additionally yield observation-aligned obsm and layers elements.

Return type:

Self