nvalchemi.data.datapipes.MultiDataset

class nvalchemi.data.datapipes.MultiDataset(*datasets, output_strict=True, num_workers=2)

Compose multiple batch-loadable datasets behind one global index space.

MultiDataset concatenates several Dataset objects so one run can draw from more than one source (for example several Zarr stores, or datasets with different chemistries or labels). Child order defines a global index space: a global index maps to a (child, local_index) pair, and to_global_index() performs the reverse mapping the samplers rely on. It exposes the same length and batch-read surface as a single Dataset, so a DataLoader consumes it interchangeably.

On its own a MultiDataset behaves like uniform concatenation; the mixing policy – how often each child is drawn – comes from a multi-dataset sampler passed to the loader. MultiDatasetSampler mixes at per-sample rates (as sampler=), while MultiDatasetBatchSampler fixes each batch’s composition (as batch_sampler=); both read child offsets from this wrapper through to_global_index().

With output_strict=True (default) every non-empty child must expose identical field names, so collated batches are homogeneous; empty children are skipped. Use output_strict=False for heterogeneous sources, where only the first child’s field names are reported and a custom training loop or collator must handle source-specific fields. num_workers sizes the thread pool for the mixed-dataset fused prefetch. Compared with torch.utils.data.ConcatDataset, MultiDataset adds this field-name contract, the fused batch API, and AtomicData semantics.

Parameters:

• *datasets (BatchDatasetProtocol) – One or more nvalchemi datasets. Order defines the global index mapping.

• output_strict (bool, default=True) – If True, require all datasets to expose identical field names.

• num_workers (int, default=2) – Thread pool size for mixed-dataset fused prefetches.

Examples

Concatenate two datasets and draw each batch as three samples from the first source and one from the second:

>>> from nvalchemi.data.datapipes import DataLoader
>>> from nvalchemi.data.datapipes.multidataset import MultiDataset
>>> from nvalchemi.data.datapipes.samplers import (
...     MultiDatasetBatchSampler,
... )
>>> multi = MultiDataset(dataset_a, dataset_b)
>>> sampler = MultiDatasetBatchSampler(
...     multi, batch_size=4, samples_per_dataset=(3, 1)
... )
>>> loader = DataLoader(multi, batch_sampler=sampler)

For stochastic per-sample mixing instead, use a MultiDatasetSampler as the loader’s sampler= argument.

cancel_prefetch(index=None)

Cancel prefetch for one global index or all child datasets.

Parameters:

index (int | None)

Return type:

None

close()

Close all child datasets and release wrapper resources.

Return type:

None

property datasets: tuple[BatchDatasetProtocol, ...]

Child datasets in global index order.

property field_names: list[str]

Return field names exposed by child datasets.

get_fused_batches()

Consume one pending fused prefetch chunk.

Return type:

Iterator[Batch]

get_metadata(index)

Return lightweight metadata for a sample by global index.

Parameters:

index (int)

Return type:

tuple[int, int]

has_pending_fused_batches()

Return whether a fused prefetch chunk is waiting to be consumed.

Return type:

bool

load_batches(batch_index_lists, stream=None)

Load several global batches immediately.

This is the synchronous counterpart to prefetch_fused_batches()/get_fused_batches(). Same-child chunks are delegated directly to the owning child dataset, while mixed chunks are routed per child and recombined in the requested batch order.

Parameters:

• batch_index_lists (Sequence[Sequence[int]]) – Per-batch global sample indices.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for child dataset transfers when supported.

Returns:

One Batch per input batch-index list.

Return type:

list[Batch]

property offsets: tuple[int, ...]

Cumulative global index offsets for child datasets.

prefetch(index, stream=None)

Start prefetching one sample by global index.

Parameters:

• index (int)

• stream (Stream | None)

Return type:

None

prefetch_batch(indices, streams=None)

Start prefetching multiple samples by global index.

Parameters:

• indices (Sequence[int])

• streams (Sequence[Stream] | None)

Return type:

None

property prefetch_count: int

Return queued prefetch count across this wrapper and children.

prefetch_fused_batches(batch_index_lists, stream=None)

Submit multiple global batches as one fused async read.

Parameters:

• batch_index_lists (Sequence[Sequence[int]])

• stream (Stream | None)

Return type:

None

to_global_index(dataset_index, local_index)

Map a child dataset index and local index to one global index.

Parameters:

• dataset_index (int)

• local_index (int)

Return type:

int

to_local_index(index)

Map one global index to (dataset_index, local_index).

Parameters:

index (int)

Return type:

tuple[int, int]

validate_field_names(output_strict=None)

Validate and return the field names exposed by this wrapper.

Parameters:

output_strict (bool | None, default=None) – Strictness mode to use for validation. None uses the mode passed to MultiDataset at construction time.

Returns:

Field names this multidataset exposes.

Return type:

list[str]

Raises:

ValueError – If output_strict=True and non-empty child datasets expose different field names.

Notes

With output_strict=True, all non-empty child datasets must expose identical field names. Empty children are skipped, matching the standalone MultiDataset strict-output behavior.

With output_strict=False, no cross-dataset validation is performed and the first child dataset’s field names are returned. Use this mode for heterogeneous datasets where a custom training loop or collator handles source-specific fields.