nvalchemi.data.Dataset

class nvalchemi.data.Dataset(reader, *, device=None, num_workers=2, skip_validation=False, transforms=None)

AtomicData-native, map-style dataset that bypasses TensorDict conversion.

Dataset is the entry point of nvalchemi’s data pipeline. It wraps a single Reader (such as AtomicDataZarrReader) and turns stored records into AtomicData graphs. Indexing returns a (AtomicData, metadata) pair (ds[i]) and len(ds) reports the sample count, so it satisfies the map-style dataset protocol that PyTorch samplers expect. The pipeline is Reader -> Dataset (-> MultiDataset) -> DataLoader: DataLoader collates samples into batched Batch graphs, and MultiDataset concatenates several datasets behind one index space.

Unlike torch.utils.data.Dataset, this class owns collation and device movement rather than deferring them to forked workers:

• it returns validated AtomicData directly, not raw tensors or a TensorDict;

• it resolves and transfers each sample to device itself;

• it prefetches on background threads and, when CUDA is available, overlapping CUDA streams – num_workers sizes a thread pool, not a multiprocessing fork – and DataLoader reads batches in fused windows through a private batch API rather than one __getitem__ per sample.

Two nuances are worth noting. skip_validation=True bypasses AtomicData Pydantic validation on the fused batch path and is only safe for trusted stores (for example those written by AtomicDataZarrWriter). And transforms are applied per sample on the prefetch CUDA stream, so they must be stream-safe (avoid .item(), .cpu(), and synchronization); for per-batch work use DataLoader’s batch_transforms instead.

Dataset implements BatchDatasetProtocol, the batch-loading contract that DataLoader and MultiDataset consume. For a fully-resident alternative that materializes the whole dataset once and trades memory for read speed, see InMemoryDataset.

Parameters:

• reader (Reader | ReaderProtocol) – Reader providing raw tensor dicts from a data source.

• device (str | torch.device | None, default=None) – Target device. "auto" picks CUDA if available, otherwise CPU.

• num_workers (int, default=2) – Thread pool size for async prefetch.

• transforms (Sequence[SampleTransform] | None, default=None) – Optional per-sample transforms applied after device transfer. See __init__() for details.

• skip_validation (bool)

reader

The underlying data reader.

Type:

Reader | ReaderProtocol

target_device

Resolved target device for data transfer.

Type:

torch.device | None

num_workers

Number of worker threads for prefetching.

Type:

int

Examples

>>> from nvalchemi.data.datapipes.dataset import Dataset
>>> from nvalchemi.data.datapipes.backends.base import Reader
>>> # Assuming a concrete Reader implementation exists:
>>> # reader = MyReader("dataset.zarr")
>>> # ds = Dataset(reader, device="cpu")
>>> # atomic_data, meta = ds[0]

With a user-supplied per-sample transform:

>>> def shift(data, metadata):
...     return data.replace(positions=data.positions + 1.0), metadata
>>> ds = Dataset(reader, device="cpu", transforms=[shift])
>>> atomic_data, meta = ds[0]

cancel_prefetch(index=None)

Cancel pending prefetch operations.

Parameters:

index (int | None, default=None) – Specific index to cancel, or None to cancel all.

Return type:

None

close()

Release resources held by the dataset.

Drains pending prefetch futures, shuts down the thread pool executor, and closes the underlying reader.

Return type:

None

property field_names: list[str]

Return field names available in reader samples.

Returns:

Field names exposed by the backing reader.

Return type:

list[str]

get_batch(indices)

Read sample indices and return a validated Batch.

Parameters:

indices (Sequence[int]) – Sample indices to batch in order.

Returns:

Batched AtomicData as a disjoint graph.

Return type:

Batch

get_fused_batches()

Consume the pending fused prefetch and yield per-batch results.

Blocks until the fused read completes, then splits the flat result list according to the original batch sizes and yields one Batch per sub-batch.

Yields:

Batch – One batch per sub-batch from the fused read.

Raises:

• RuntimeError – If no fused prefetch is pending.

• Exception – If the background read failed, re-raises the original error.

Return type:

Iterator[Batch]

get_metadata(index)

Return lightweight metadata for a sample without full construction.

Delegates to the reader when it provides lightweight metadata; otherwise loads the raw tensor dictionary and extracts shape information for atom and edge counts, avoiding the overhead of full AtomicData construction and validation.

Parameters:

index (int) – Sample index.

Returns:

(num_atoms, num_edges) for the sample.

Return type:

tuple[int, int]

Raises:

• IndexError – If index is out of range.

• KeyError – If the sample dict does not contain "atomic_numbers".

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 batches immediately.

This is the synchronous counterpart to prefetch_fused_batches()/get_fused_batches(). The provided batch index lists are read through one fused reader request so backends can coalesce I/O while returning one Batch per input list.

Parameters:

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

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for device transfer when supported.

Returns:

One Batch per input batch-index list.

Return type:

list[Batch]

property pin_memory: bool

Whether the underlying reader should return pinned CPU tensors.

prefetch(index, stream=None)

Submit a sample for async prefetching.

If the sample is already being prefetched, this is a no-op.

Parameters:

• index (int) – Sample index.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for GPU operations.

Return type:

None

prefetch_batch(indices, streams=None)

Prefetch multiple samples asynchronously.

Parameters:

• indices (Sequence[int]) – Sample indices to prefetch.

• streams (Sequence[torch.cuda.Stream] | None, default=None) – CUDA streams to distribute across. Streams are assigned round-robin to the indices.

Return type:

None

property prefetch_count: int

Return the number of pending prefetch requests.

Returns:

Count of queued single-sample and fused-batch prefetches.

Return type:

int

prefetch_fused_batches(batch_index_lists, stream=None)

Submit multiple batches as one fused async read.

All indices across the provided batch lists are concatenated into a single read_many call, amortizing Zarr I/O overhead. Use get_fused_batches() to consume the results.

Parameters:

• batch_index_lists (Sequence[Sequence[int]]) – Per-batch index lists.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for GPU operations.

Return type:

None

prefetch_many(indices, stream=None)

Submit one batch of sample indices as a fused async prefetch.

Parameters:

• indices (Sequence[int]) – Sample indices to prefetch as one batch.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for GPU operations.

Return type:

None

read_many(indices)

Read and validate multiple samples in one dataset request.

Parameters:

indices (Sequence[int]) – Sample indices to load in order.

Returns:

Ordered (AtomicData, metadata) pairs.

Return type:

list[tuple[AtomicData, dict[str, Any]]]