Data

class cellarium.ml.data.AnnDataSchema(adata: AnnData, obs_columns_to_validate: Sequence[str] | None = None)[source]

Bases: object

Store reference AnnData attributes for a collection of distributed AnnData objects.

Validate AnnData objects against reference attributes.

Example:

>>> ref_adata = AnnData(X=np.zeros(2, 3))
>>> ref_adata.obs["batch"] = [0, 1]
>>> ref_adata.var["mu"] = ["a", "b", "c"]
>>> adata = AnnData(X=np.zeros(4, 3))
>>> adata.obs["batch"] = [2, 3, 4, 5]
>>> adata.var["mu"] = ["a", "b", "c"]
>>> schema = AnnDataSchema(ref_adata)
>>> schema.validate_anndata(adata)

Parameters:

adata (AnnData) – Reference AnnData object.
obs_columns_to_validate (Sequence[str] | None) – Subset of columns to validate in the .obs attribute. If None, all columns are validated.

validate_anndata(adata: AnnData) → None[source]

Validate anndata has proper attributes.

Parameters:: adata (AnnData)
Return type:: None

class cellarium.ml.data.DeltaCellsBatch(collection: DistributedDeltaCellsCollection, indices: ndarray)[source]

Bases: object

A batch of cells that behaves like an AnnData for the attributes AnnDataField reads: X (float32 CSR), obs (pandas, with categorical columns carrying the dataset’s global categories), obs_names, var, var_names.

Everything is read lazily, only when the attribute is accessed, and X once.

Parameters:

collection (DistributedDeltaCellsCollection)
indices (ndarray)

class cellarium.ml.data.DistributedAnnDataCollection(filenames: Sequence[str] | str, limits: Iterable[int] | None = None, shard_size: int | None = None, last_shard_size: int | None = None, max_cache_size: int = 1, cache_size_strictly_enforced: bool = True, label: str | None = None, keys: Sequence[str] | None = None, index_unique: str | None = None, convert: Callable | dict[str, Callable | dict[str, Callable]] | None = None, indices_strict: bool = True, obs_columns_to_validate: Sequence[str] | None = None)[source]

Bases: AnnCollection, DistributedCollection

Distributed AnnData Collection.

This class is a wrapper around AnnCollection where adatas is a list of LazyAnnData objects.

Underlying anndata files must conform to the same schema (see validate_anndata). The schema is inferred from the first AnnData file in the collection. Individual AnnData files may otherwise vary in the number of cells, and the actual content stored in X, layers, obs and obsm.

Example 1:

>>> dadc = DistributedAnnDataCollection(
...     "gs://bucket-name/folder/adata{000..005}.h5ad",
...     shard_size=10000,  # use if shards are sized evenly
...     max_cache_size=2)

Example 2:

>>> dadc = DistributedAnnDataCollection(
...     "gs://bucket-name/folder/adata{000..005}.h5ad",
...     shard_size=10000,
...     last_shard_size=6000,  # use if the size of the last shard is different
...     max_cache_size=2)

Example 3:

>>> dadc = DistributedAnnDataCollection(
...     "gs://bucket-name/folder/adata{000..005}.h5ad",
...     limits=[500, 1000, 2000, 2500, 3000, 4000],  # use if shards are sized unevenly
...     max_cache_size=2)

Parameters:

filenames (Sequence[str] | str) – Names of anndata files.
limits (list[int]) – List of global cell indices (limits) for the last cells in each shard. If None, the limits are inferred from shard_size and last_shard_size.
shard_size (int | None) – The number of cells in each anndata file (shard). Must be specified if the limits is not provided.
last_shard_size (int | None) – Last shard size. If not None, the last shard will have this size possibly different from shard_size.
max_cache_size (int) – Max size of the cache.
cache_size_strictly_enforced (bool) – Assert that the number of retrieved anndatas is not more than maxsize.
label (str | None) – Column in obs to place batch information in. If it’s None, no column is added.
keys (Sequence[str] | None) – Names for each object being added. These values are used for column values for label or appended to the index if index_unique is not None. If None, keys are set to filenames.
index_unique (str | None) – Whether to make the index unique by using the keys. If provided, this is the delimeter between {orig_idx}{index_unique}{key}. When None, the original indices are kept.
convert (Callable | dict[str, Callable | dict[str, Callable]] | None) – You can pass a function or a Mapping of functions which will be applied to the values of attributes (obs, obsm, layers, X) or to specific keys of these attributes in the subset object. Specify an attribute and a key (if needed) as keys of the passed Mapping and a function to be applied as a value.
indices_strict (bool) – If True, arrays from the subset objects will always have the same order of indices as in selection used to subset. This parameter can be set to False if the order in the returned arrays is not important, for example, when using them for stochastic gradient descent. In this case the performance of subsetting can be a bit better.
obs_columns_to_validate (Sequence[str] | None) – Subset of columns to validate in the obs attribute. If None, all columns are validated.

property var: DataFrame: Per-variable annotations, as in the first anndata file.

property reference_adata: LazyAnnData: The first anndata file (read lazily), for attributes that are the same in every file such as var.

obs_categories(key: str) → ndarray[source]

The categories of the categorical obs column key, as in the first anndata file.

Parameters:: key (str)
Return type:: ndarray

obs_key_nunique(key: str) → int[source]

Number of distinct values of the obs column key in the first anndata file.

Parameters:: key (str)
Return type:: int

reset_cache() → None[source]

Forget the cached anndata files.

Return type:: None

property cache_miss_count: int: Number of anndata files read so far by this process.

Materialize and gather anndata files at given indices from the list of lazy anndatas.

LazyAnnData instances corresponding to cells in the index are materialized.

Parameters:: index (int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, EllipsisType] | tuple[EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array])
Return type:: AnnData

Like indexing, but cells that lie in more files than the cache holds (for example a batch of scattered cells) are read in groups of at most max_cache_size files, copying the cells of each group out before the next group is read, so no more than max_cache_size files are ever in memory (and cache_size_strictly_enforced is never violated).

Parameters:: index (int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array | EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, EllipsisType] | tuple[EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array] | tuple[int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array, EllipsisType, int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array])
Return type:: AnnData

Buffer and return anndata files at given indices from the list of lazy anndatas.

This efficiently first retrieves cached files and only then caches new files.

Parameters:

adatas_oidx (list[ndarray | None])
vidx (int | str | int64 | slice | ndarray[Any, dtype[bool_]] | ndarray[Any, dtype[integer]] | Sequence[int] | Sequence[str] | Sequence[bool] | Series | Index | ExtensionArray | ndarray[Any, dtype[str_]] | matrix | csr_matrix | csc_matrix | csr_array | csc_array)

Return type:

list[AnnData]

class cellarium.ml.data.DistributedAnnDataCollectionView(reference, convert, resolved_idx)[source]

Bases: AnnCollectionView

Distributed AnnData Collection View.

This class is a wrapper around AnnCollectionView where adatas is a list of LazyAnnData objects.

property obs_names: Index: Gather and return the obs_names from all AnnData objects in the collection.

class cellarium.ml.data.DistributedCollection[source]

Bases: ABC

Interface of the data sources that IterableDistributedAnnDataCollectionDataset and CellariumAnnDataDataModule read from.

A collection is a sequence of n_obs cells split into consecutive shards (for example h5ad files or deltacells tiles). Indexing it with global cell indices returns an object that behaves like an AnnData for the attributes that AnnDataField reads (X, obs, var_names …).

Subclasses must provide n_obs, n_vars, var and __getitem__(), and should set the following attributes:

limits

Cumulative number of cells at the end of each shard, e.g. [10000, 20000, 25000]. Shards are the units that the dataset shuffles and splits between workers, so cells of one shard should be cheap to read together.

Type:: list[int]

var_names

Names of the variables (genes), in column order.

Type:: Any

The remaining members have defaults that do nothing, so a source only overrides what it can use: prefetch() (a hint about which cells will be requested soon), prepare() (one-time preparation before training, e.g. making metadata local) and the cache hooks used by the dataset’s test_mode.

abstract property n_obs: int: Number of cells.

abstract property n_vars: int: Number of variables (genes).

abstract property var: DataFrame: Per-variable annotations (one row per variable, in column order).

abstractmethod __getitem__(index: Any) → Any[source]

Cells with the given global indices (an int, a slice or a sequence of ints), as an AnnData-like object.

Parameters:: index (Any)
Return type:: Any

read(index: Any) → Any[source]

Like indexing, for requests of any size or spread. Collections whose indexing limits how many shards one request may touch (for example DistributedAnnDataCollection, with a strictly enforced cache size) override this to read in groups. Used for batches of scattered cells.

Parameters:: index (Any)
Return type:: Any

property reference_adata: Any: An AnnData-like object to read the attributes that are the same for every cell (var, var_names) from, as cheaply as the source allows. Defaults to the first cell.

obs_categories(key: str) → ndarray[source]

The categories of the categorical obs column key (code i of a batch means categories[i]).

Parameters:: key (str)
Return type:: ndarray

obs_key_nunique(key: str) → int[source]

Number of distinct values of the obs column key.

Parameters:: key (str)
Return type:: int

prefetch(indices: ndarray) → None[source]

Hint that cells indices will be requested soon (in about that order), so that their shards can be fetched in the background. Must be cheap, non-blocking and idempotent. Does nothing by default.

Parameters:: indices (ndarray)
Return type:: None

prepare(batch_keys: Any) → None[source]

One-time preparation for reading the fields in batch_keys (a pytree of AnnDataField), called by the datamodule’s prepare_data. Does nothing by default.

Parameters:: batch_keys (Any)
Return type:: None

reset_cache() → None[source]

Forget cached shards (used by the dataset’s test_mode). Does nothing by default.

Return type:: None

property cache_miss_count: int: Number of shards fetched so far by this process (used by the dataset’s test_mode).

class cellarium.ml.data.DistributedDeltaCellsCollection(uri: str, max_cached_tiles: int = 1, max_prefetch_tiles: int = 8, io_threads: int = 4, decode_threads: int = 1, cache_dir: str | None = None, verify: bool = False, obs_columns: Sequence[str] | None = None)[source]

Bases: DistributedCollection

A collection of cells stored as a deltacells dataset: one object per tile of consecutive cells (the shards of this collection), compressed so that a tile is read and decoded in about a tenth of a second, with the cell metadata (obs) kept beside the tiles and copied to a node-local cache, one column at a time, only when a field needs it.

Indexing returns a DeltaCellsBatch, an AnnData-like batch, so AnnDataField and all of its convert_fn s work as they do with h5ad files. Differences from h5ad data to be aware of:

X holds integer counts as float32 (exact up to 2**24) and the genes are in the dataset’s stored order, which differs from the original order; models identify genes by var_names_g.
The categories of categorical obs columns are global (the same in every batch and every shard) and are strings.
Only the attributes X, obs, obs_names, var and var_names exist.

Example:

>>> from cellarium.ml import CellariumAnnDataDataModule
>>> from cellarium.ml.data import DistributedDeltaCellsCollection
>>> from cellarium.ml.utilities.data import AnnDataField, categories_to_codes, to_torch_sparse_csr

>>> dadc = DistributedDeltaCellsCollection("gs://bucket-name/dataset")
>>> dm = CellariumAnnDataDataModule(
...     dadc,
...     batch_keys={
...         "x_ng": AnnDataField(attr="X", convert_fn=to_torch_sparse_csr),
...         "var_names_g": AnnDataField(attr="var_names"),
...         "y_n": AnnDataField(attr="obs", key="cell_type", convert_fn=categories_to_codes),
...     },
...     batch_size=5000,
...     shuffle=True,
...     num_workers=4,
... )

The obs columns used by batch_keys are copied to the local cache by the datamodule’s prepare_data.

Parameters:

uri (str) – Dataset location: a local path or gs://bucket/prefix.
max_cached_tiles (int) – Decoded tiles kept in memory per process (about 8 bytes per nonzero each).
max_prefetch_tiles (int) – Compressed tiles kept (or being fetched) per process.
io_threads (int) – Concurrent tile fetches per process.
decode_threads (int) – Threads used to decode one tile (1 is right when each worker has about one core).
cache_dir (str | None) – Where obs columns are cached (default: $DELTACELLS_CACHE or ~/.cache/deltacells).
verify (bool) – Check the CRC32 of every tile that is read.
obs_columns (Sequence[str] | None) – Additional obs columns to make local in prepare() (for example for plotting later); the columns that batch_keys use are always made local.

property obs_columns_available: list[str]: Names of the obs columns stored with the dataset.

obs_categories(key: str) → ndarray[source]

The global categories of the categorical obs column key.

Parameters:: key (str)
Return type:: ndarray

obs_key_nunique(key: str) → int[source]

Number of distinct values of the obs column key: the number of categories for a categorical column.

Parameters:: key (str)
Return type:: int

take_obs(indices: ndarray, columns: Sequence[str]) → DataFrame[source]

Obs columns for the given cells (fetched on first use if they were not made local in prepare()).

Parameters:

indices (ndarray)
columns (Sequence[str])

Return type:

DataFrame

class cellarium.ml.data.IterableDistributedAnnDataCollectionDataset(dadc: DistributedCollection | AnnData, batch_keys: dict[str, dict[str, AnnDataField] | AnnDataField], batch_size: int = 1, iteration_strategy: Literal['same_order', 'cache_efficient'] = 'cache_efficient', shuffle: bool = False, shuffle_seed: int = 0, drop_last_indices: bool = False, drop_incomplete_batch: bool = False, start_idx: int | None = None, end_idx: int | None = None, worker_seed: int | None = None, test_mode: bool = False, prefetch_lookahead: int = 2)[source]

Bases: IterableDataset

Iterable DistributedAnnDataCollection Dataset.

When shuffle is set to True then the iterator yields datapoints that are uniformly sampled from the entire dataset. Typical use cases include training variational models using the stochastic gradient descent algorithm.

In order to maximize buffer usage, we only shuffle shards and datapoints within individual shards (and not across shards). Therefore, to achieve unbiased pseudo-random uniform sampling, it is imperative that the shards themselves contain datapoints that are uniformly sampled from the entire dataset. If correlations exist between datapoints in a given shard (e.g. all cells coming from the same tissue or experiment), then this assumption is violated. It is the user’s responsibility to prepare appropriately shuffled data shards.

Example:

>>> from cellarium.ml.data import (
...     DistributedAnnDataCollection,
...     IterableDistributedAnnDataCollectionDataset,
... )
>>> from cellarium.ml.utilities.data import AnnDataField, densify

>>> dadc = DistributedAnnDataCollection(
...     "gs://bucket-name/folder/adata{000..005}.h5ad",
...     shard_size=10_000,
...     max_cache_size=2)

>>> dataset = IterableDistributedAnnDataCollectionDataset(
...     dadc,
...     batch_keys={
...         "x_ng": AnnDataField(attr="X", convert_fn=densify),
...         "var_names_g": AnnDataField(attr="var_names"),
...     },
...     batch_size=5000,
...     iteration_strategy="cache_efficient",
...     shuffle=True,
...     shuffle_seed=0,
...     drop_last_indices=True,
... )

Parameters:

dadc (DistributedCollection | AnnData) – DistributedCollection (for example a DistributedAnnDataCollection) or AnnData from which to load the data.
batch_keys (dict[str, dict[str, AnnDataField] | AnnDataField]) – Dictionary that specifies which attributes and keys of the dadc to return in the batch data and how to convert them. Keys must correspond to the input keys of the transforms or the model. Values must be instances of cellarium.ml.utilities.data.AnnDataField.
batch_size (int) – How many samples per batch to load.
iteration_strategy (Literal['same_order', 'cache_efficient']) – Strategy to use for iterating through the dataset. Options are same_order and cache_efficient. same_order will iterate through the dataset in the same order independent of the number of replicas and workers. cache_efficient will try to minimize the amount of anndata files fetched by each worker.
shuffle (bool) – If True, the data is reshuffled at every epoch.
shuffle_seed (int) – Random seed used to shuffle the sampler if shuffle=True.
drop_last_indices (bool) – If True, then the sampler will drop the tail of the data to make it evenly divisible across the number of replicas. If False, the sampler will add extra indices to make the data evenly divisible across the replicas.
drop_incomplete_batch (bool) – If True, the dataloader will drop the incomplete batch if the dataset size is not divisible by the batch size.
start_idx (int | None) – The starting index of the dataset. If None, then the dataset will start from the first index.
end_idx (int | None) – The ending index (exclusive) of the dataset. If None, then the dataset will end at the last index (inclusive).
worker_seed (int | None) – Random seed used to seed the workers. If None, then the workers will not be seeded. The seed of the individual worker is computed based on the worker_seed, global worker id, and the epoch. Note that the this seed affects cpu_transforms when they are used. When resuming training, the seed should be set to a different value to ensure that the workers are not seeded with the same seed as the previous run.
test_mode (bool) – If True, then tracking of cache and worker informations will be enabled.
prefetch_lookahead (int) – For collections that can prefetch (dadc.supports_prefetch), the number of upcoming shards, beyond those of the current batch, that the dataset announces to dadc.prefetch while it iterates. It has no effect on the order of the data or on collections that cannot prefetch.

__len__() → int[source]

Returns the number of batches per replica.

Return type:: int

set_epoch(epoch: int) → None[source]

Sets the epoch for the iterator. When shuffle=True, this ensures all replicas use a different random ordering for each epoch.

Parameters:: epoch (int)
Return type:: None

set_resume_step(resume_step: int | None) → None[source]

Sets the resume step for the iterator. When resuming from a checkpoint, this ensures that the iterator skips the batches that have already been processed.

Parameters:: resume_step (int | None)
Return type:: None

__getitem__(idx: int | list[int] | slice) → dict[str, dict[str, ndarray] | ndarray][source]

Returns a dictionary containing the data from the dadc with keys specified by the batch_keys at the given index idx.

Parameters:: idx (int | list[int] | slice)
Return type:: dict[str, dict[str, ndarray] | ndarray]

__iter__()[source]

Iterate through the dataset by trying to minimize the amount of anndata files fetched by each worker. Iterated indices are evenly divided between replicas (see drop_last_indices).

Note

For both strategies the amount of anndata files fetched is reduced by shuffling the shards first and then the datapoints within the shards.
same_order strategy will iterate through the dataset in the same order independent of the number of replicas and workers.
For cache_efficient strategy the amount of anndata files fetched is further reduced by assigning to each worker a contiguous chunk of the dataset. The returned iterator is determined by the torch.utils.data.get_worker_info() and torch.distributed contexts.

Example 1:

indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
num_replicas=1
batch_size=2
num_workers=3

Same order:

batch idx	0	1	2	3	4	5
indices	(0,1)	(2,3)	(4,5)	(6,7)	(8,9)	(10,11)
worker id	0	1	2	0	1	2

Cache efficient:

batch idx	0	1	2	3	4	5
indices	(0,1)	(4,5)	(8,9)	(2,3)	(6,7)	(10,11)
worker id	0	1	2	0	1	2

Example 2:

indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
num_replicas=1
batch_size=2
num_workers=2

Same order:

batch idx	0	1	2	3	4	5
indices	(0,1)	(2,3)	(4,5)	(6,7)	(8,9)	(10,)
worker id	0	1	0	1	0	1

Cache efficient:

batch idx	0	1	2	3	4	5
indices	(0,1)	(6,7)	(2,3)	(8,9)	(4,5)	(10,)
worker id	0	1	0	1	0	1

Example 3:

indices=[0, 1, 2, 3, 4, 5, 6, 7]
num_replicas=1
batch_size=3
num_workers=2

Same order:

batch idx	0	1	2
indices	(0,1,2)	(3,4,5)	(6,7)
worker id	0	1	0

Cache efficient:

batch idx	0	1	2
indices	(0,1,2)	(6,7)	(3,4,5)
worker id	0	1	0

Example 4:

indices=[0, 1, 2, 3, 4, 5, 6, 7]
num_replicas=1
batch_size=3
drop_incomplete_batch=True
num_workers=2

Same order:

batch idx	0	1
indices	(0,1,2)	(3,4,5)
worker id	0	1

Cache efficient:

batch idx	0	1
indices	(0,1,2)	(3,4,5)
worker id	0	1

Example 5:

indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
num_replicas=2
drop_last_indices=True
batch_size=2
num_workers=1

Same order:

Replica 1

batch idx	0	1	2
indices	(0,2)	(4,6)	(8,)
worker id	0	0	0

Replica 2

batch idx	0	1	2
indices	(1,3)	(5,7)	(9,)
worker id	0	0	0

Cache efficient:

Replica 1

batch idx	0	1	2
indices	(0,1)	(2,3)	(4,)
worker id	0	0	0

Replica 2

batch idx	0	1	2
indices	(5,6)	(7,8)	(9,)
worker id	0	0	0

Example 6:

indices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
num_replicas=2
drop_last_indices=False
batch_size=2
num_workers=1

Same order:

Replica 1

batch idx	0	1	2
indices	(0,2)	(4,6)	(8,10)
worker id	0	0	0

Replica 2

batch idx	0	1	2
indices	(1,3)	(5,7)	(9,0)
worker id	0	0	0

Cache efficient:

Replica 1

batch idx	0	1	2
indices	(0,1)	(2,3)	(4,5)
worker id	0	0	0

Replica 2

batch idx	0	1	2
indices	(6,7)	(8,9)	(10,0)
worker id	0	0	0

Resuming from a checkpoint:

For persistent workers the state (epoch and resume_step) is initially set by the load_state_dict() method. At the end of the iteration, the epoch is incremented and the resume_step is set to None.
For non-persistent workers the state is initially set by the load_state_dict() method. The epoch is updated by the on_train_epoch_start hook and the resume_step is set to None by the on_train_epoch_end hook.
If the resume_step is not None, then the worker will skip the batches that have already been processed. The workers are shifted based on the global step.

load_state_dict(state_dict: dict[str, Any]) → None[source]

Loads the state of the dataset from the given state dictionary.

Parameters:: state_dict (dict[str, Any]) – State dictionary containing the state of the dataset.
Return type:: None

class cellarium.ml.data.LazyAnnData(filename: str, limits: tuple[int, int], schema: AnnDataSchema, cache: LRU | None = None)[source]

Bases: object

Lazy AnnData backed by a file.

Accessing attributes under lazy_getattr() context returns schema attributes.

Parameters:

filename (str) – Name of anndata file.
limits (tuple[int, int]) – Limits of cell indices (inclusive, exclusive).
schema (AnnDataSchema) – Schema used as a reference for lazy attributes.
cache (LRU | None) – Shared LRU cache storing buffered anndatas.

property n_obs: int: Number of observations.

property n_vars: int: Number of variables/features.

property shape: tuple[int, int]: Shape of the data matrix.

property obs_names: Index: Return the observation names.

property cached: bool: Return whether the anndata is cached.

property adata: AnnData: Return backed anndata from the filename

class cellarium.ml.data.PyTreeDataset(pytree: Any)[source]

Bases: Dataset

A dataset that wraps a PyTree of tensors and ndarrays.

Example:

import torch
from cellarium.ml.data import PyTreeDataset
from cellarium.ml.utilities.data import collate_fn

data = {
    "gene_token_nc_dict": {
        "gene_id": torch.randint(0, 10, (10, 3)),
        "gene_value": torch.randint(0, 10, (10, 3)),
    },
    "gene_token_mask_nc": torch.randint(0, 10, (10, 3)),
    "metadata_token_nc_dict": {
        "cell_type": torch.randint(0, 10, (10, 3)),
    },
    "metadata_token_mask_nc_dict": {
        "cell_type": torch.randint(0, 10, (10, 3)),
    },
    "prompt_mask_nc": torch.randint(0, 10, (10, 3)),
}
dataset = PyTreeDataset(data)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=batch_size,
    collate_fn=collate_fn,
)
for batch in dataloader:
    ...

Parameters:: pytree (Any) – A PyTree of tensors and ndarrays.

cellarium.ml.data.read_h5ad_file(filename: str, **kwargs) → AnnData[source]

Read .h5ad-formatted hdf5 file from a filename.

Parameters:: filename (str) – Path to the data file.
Return type:: AnnData

cellarium.ml.data.read_h5ad_gcs(filename: str, storage_client: Client | None = None) → AnnData[source]

Read .h5ad-formatted hdf5 file from the Google Cloud Storage.

Example:

>>> adata = read_h5ad_gcs("gs://dsp-cellarium-cas-public/test-data/test_0.h5ad")

Parameters:

filename (str) – Path to the data file in Cloud Storage.
storage_client (Client | None)

Return type:

AnnData

cellarium.ml.data.read_h5ad_local(filename: str) → AnnData[source]

Read .h5ad-formatted hdf5 file from the local disk.

Parameters:: filename (str) – Path to the local data file.
Return type:: AnnData

cellarium.ml.data.read_h5ad_url(filename: str) → AnnData[source]

Read .h5ad-formatted hdf5 file from the URL.

Example:

>>> adata = read_h5ad_url(
...     "https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad"
... )

Parameters:: filename (str) – URL of the data file.
Return type:: AnnData