torchcell.datamodules.CellDataModule

class CellDataModule(dataset: Any, cache_dir: str = 'cache', batch_size: int = 32, random_seed: int = 42, num_workers: int = 0, pin_memory: bool = False, prefetch: bool = False, prefetch_factor: int = 2, persistent_workers: bool = True, split_indices: str | list[str] | None = None, follow_batch: list[str] | None = None, train_shuffle: bool = True, collate_fn: Callable[[list[Any]], Any] | None = None, val_batch_size: int | None = None, pinned_test_indices: Iterable[int] | None = None, pinned_split_indices: Mapping[str, Iterable[int]] | None = None, index_subset: Iterable[int] | None = None)

Bases: LightningDataModule

Lightning data module that splits a cell dataset and builds dataloaders.

property index: DataModuleIndex

Return the train/val/test split index, computing or loading it if needed.

property index_details: DataModuleIndexDetails

Return the detailed split-index breakdown, computing or loading if needed.

setup(stage: str | None = None) → None

Build train/val/test Subset datasets from the computed split indices.

IDEMPOTENT BY DESIGN. Lightning calls setup() again inside trainer.fit() even when the caller already invoked it, so an unconditional rebuild silently discarded any post-setup narrowing of Subset.indices (e.g. the require_modalities filter in experiments/019, which logged “4074 -> 1161 rows” while every epoch still ran 128 batches = 4074/32). Rebuilding only when the subsets do not yet exist makes the caller’s filter survive the second call.

train_dataloader() → DataLoader[Any] | PrefetchLoader

Return a dataloader over the training subset.

val_dataloader() → DataLoader[Any] | PrefetchLoader

Return a dataloader over the validation subset.

test_dataloader() → DataLoader[Any] | PrefetchLoader

Return a dataloader over the test subset.

all_dataloader() → DataLoader[Any] | PrefetchLoader

Return a dataloader over the entire dataset.