nvalchemi.data.Reader

class nvalchemi.data.Reader(*, pin_memory=False, include_index_in_metadata=True, coordinated_subsampling=None)

Abstract base class for data readers.

Readers are intentionally simple and transactional:

• Load data from a source (file, database, etc.)

• Return (dict[str, torch.Tensor], metadata_dict) tuples with CPU tensors

• No threading, no prefetching, no device transfers

Subclasses must implement __len__() and at least one loading hook: _load_sample() for simple single-sample readers, or _load_many_samples() for readers that can amortize I/O across a group of samples.

Parameters:

• pin_memory (bool, default=False) – If True, pin loaded tensors to page-locked memory for faster async CPU→GPU transfers.

• include_index_in_metadata (bool, default=True) – If True, automatically add "index" to each sample’s metadata dict.

• coordinated_subsampling (dict[str, Any] | None)

Examples

>>> class MyReader(Reader):
...     def _load_sample(self, index: int) -> dict[str, torch.Tensor]:
...         return {"x": torch.randn(3)}
...     def __len__(self) -> int:
...         return 10
>>> reader = MyReader()
>>> data, meta = reader[0]

close()

Release resources held by the reader.

Override in subclasses to close file handles, connections, etc.

Return type:

None

property field_levels: dict[str, str]

"atom", "edge", or "system".

Override in subclasses that store explicit level metadata (e.g. Zarr stores). The default returns an empty dict, which causes downstream consumers to fall back to AtomicData._default_*_keys for classification.

Returns:

Mapping of field name to level string.

Return type:

dict[str, str]

Type:

Per-field level classification

property field_names: list[str]

Field names available in each sample.

Returns:

Field names.

Return type:

list[str]

read(index)

Load a sample and its metadata by index.

Handles optional pin-memory and automatic index injection into metadata. Index validity is determined by the concrete reader.

Parameters:

index (int) – Sample index. Concrete readers determine supported values.

Returns:

(data_dict, metadata) pair with CPU tensors.

Return type:

tuple[dict[str, torch.Tensor], dict[str, Any]]

Raises:

IndexError – If the concrete reader considers index out of range.

read_many(indices)

Load multiple samples and their metadata.

The default implementation delegates raw tensor loading to _load_many_samples(), and then attaches metadata and optional pinned memory. Backend implementations should override _load_many_samples() instead of this method. Index validity is determined by the concrete reader.

Parameters:

indices (Sequence[int]) – Sample indices to load. Concrete readers determine supported values.

Returns:

Ordered (data_dict, metadata) pairs with CPU tensors.

Return type:

list[tuple[dict[str, torch.Tensor], dict[str, Any]]]

Raises:

IndexError – If the concrete reader considers any requested index out of range.