Skip to content

Datasets and loading

A dataset presents a logical collection of physical samples through a common PyG interface. The reading strategy determines how much numerical data stays in memory; the returned sample fields and their physical meaning are shared by both strategies.

Reading ELFES HDF5

HDF5Dataset and InMemoryDataset use different HDF5 reading paths but return the same edge-free, single-sample torch_geometric.data.Data structure. The Data contains batchable sample tensors and compact basis references; complete basis sets and physical-data descriptions remain on the owning Dataset. Standard PyTorch DataLoaders invoke each Dataset's __getitems__() for one batch of indices: the on-disk Dataset reads all selected physical data directly into batch-contiguous storage without constructing physical Sample objects, while the in-memory Dataset separates views from its existing collated storage before ordinary PyG collation.

The input paths are ordered shards of one logical dataset. physical_data_selection=None reads all physical quantities, while an empty selection reads geometry and basis references without quantity values. The dataset exposes definitions for the selected quantities and complete basis sets. Field shapes and batch behavior are described in Samples and batches.

HDF5Dataset opens file handles within the process that reads them and provides close() for releasing them. InMemoryDataset.from_hdf5() loads the selected arrays into CPU tensor storage and does not retain HDF5 handles.

HDF5Dataset 

HDF5Dataset(
    paths: StrPath | Sequence[StrPath],
    *,
    physical_data_selection: Collection[str] | None = None,
)

Bases: Dataset[Data]

On-disk PyG view of one logical ELFES HDF5 dataset.

The input paths are ordered shards of the same logical dataset. Construction reads their sample IDs and shared dataset definition, but numerical sample data remains on disk. Integer and batched indexing read the requested quantity-major arrays directly into continuous storage and return lightweight edge-free Data views in the requested order without constructing physical Sample objects. The common returned Data structure is documented in Samples and batches.

HDF5 handles are opened on first access within each process, so DataLoader workers do not share live handles. Multi-worker DataLoaders should use the forkserver multiprocessing context and persistent workers. Call close() when the dataset is no longer needed.

Parameters:

  • paths (StrPath | Sequence[StrPath]) –

    One HDF5 path or an ordered sequence of shard paths.

  • physical_data_selection (Collection[str] | None, default: None ) –

    Physical data to read, selected by globally unique name. None selects all data; an empty collection selects none.

Attributes:

  • sample_ids –

    Sample IDs in logical dataset order.

  • basis_sets –

    Dataset-level basis sets keyed by role.

  • physical_data_descriptions –

    Definitions of physical data returned by the Dataset.

  • physical_data_names –

    Physical data returned by the Dataset.

  • atomic_basis_tables –

    Atomic bases indexed by the atomic-basis IDs stored in Data.

close 

close() -> None

Close HDF5 handles opened by the current process.

InMemoryDataset 

InMemoryDataset(
    collated_data: Data,
    slice_dict: dict[str, Any],
    *,
    sample_ids: tuple[str, ...],
    basis_sets: dict[str, BasisSet],
    physical_data_descriptions: dict[str, PhysicalDataDescription],
    physical_data_names: Collection[str],
)

Bases: Dataset[Data]

In-memory ELFES dataset backed by one collated PyG tensor store.

from_hdf5() bulk-reads the selected quantity-major arrays from all ordered shards and converts them directly into a combined PyG Data plus nested sample boundaries. It does not retain HDF5 handles or construct a physical Sample for every row.

Integer indexing uses PyG separate() to return an edge-free sample whose tensors view the combined storage. Atom indices remain sample-local until a standard PyG DataLoader constructs a real batch; batched indexing applies the same operation to the complete requested index sequence. The common returned Data structure is documented in Samples and batches.

Parameters:

  • collated_data (Data) –

    Combined PyG data for every sample.

  • slice_dict (dict[str, Any]) –

    Nested sample boundaries for fields in collated_data.

  • sample_ids (tuple[str, ...]) –

    Sample IDs in logical dataset order.

  • basis_sets (dict[str, BasisSet]) –

    Dataset-level basis sets keyed by role.

  • physical_data_descriptions (dict[str, PhysicalDataDescription]) –

    Definitions of physical data in collated_data.

  • physical_data_names (Collection[str]) –

    Physical data present in collated_data.

Attributes:

  • sample_ids –

    Sample IDs in logical dataset order.

  • basis_sets –

    Dataset-level basis sets keyed by role.

  • physical_data_descriptions –

    Definitions of physical data returned by the Dataset.

  • physical_data_names –

    Physical data returned by the Dataset.

  • atomic_basis_tables –

    Atomic bases indexed by the atomic-basis IDs stored in Data.

from_hdf5 classmethod 

from_hdf5(
    paths: StrPath | Sequence[StrPath],
    *,
    physical_data_selection: Collection[str] | None = None,
) -> Self

Read and convert complete HDF5 shards into memory.

Parameters:

  • paths (StrPath | Sequence[StrPath]) –

    One HDF5 path or an ordered sequence of shard paths.

  • physical_data_selection (Collection[str] | None, default: None ) –

    Physical data to read, selected by globally unique name. None selects all data; an empty collection selects none.

Converting physical samples

SampleToData provides the physical conversion for Sample objects that are already in memory. Because a physical Sample has no intrinsic dataset membership, this direct conversion does not add data.sample_id. Construct one converter from the basis sets shared by a logical dataset, then reuse it for every compatible sample:

from elfes.data import SampleToData

converter = SampleToData(basis_sets)
data = converter(sample)

SampleToData dataclass 

SampleToData(basis_sets: Mapping[str, BasisSet])

Convert one physical Sample into edge-free PyG data.

Geometry and electronic quantities become Torch tensors, while every atom receives its atomic-basis-table index and orbital count for each available basis role. Atom-pair indices remain local to the sample; a later PyG DataLoader performs the index increments required for a real batch.

Parameters:

  • basis_sets (Mapping[str, BasisSet]) –

    Dataset-level basis sets keyed by role. Every converted sample must use the same definitions.

Forming batches

create_data_loader() uses PyG collation. With with_connectivity=False, it returns edge-free batches. With with_connectivity=True, a cutoff is required and directed connectivity is built during CPU collation. Worker-based loading uses forkserver and persistent workers; the connectivity reference explains how to add edges after loading instead.

For an existing ELFES HDF5 file, a minimal edge-free loading path is:

from elfes.data import HDF5Dataset, create_data_loader

dataset = HDF5Dataset("samples.h5", physical_data_selection=["energy", "forces"])
loader = create_data_loader(dataset, batch_size=8, with_connectivity=False)
try:
    for batch in loader:
        positions = batch.pos
        forces = batch.physical_data["forces"]
finally:
    dataset.close()

Here energy and forces are names defined by the example dataset. Conversion retains ELFES physical units and conventions. Target normalization and other preparation steps are composed explicitly with the loader.

create_data_loader 

create_data_loader(
    dataset: Dataset[Data] | Sequence[Data],
    *,
    batch_size: int,
    with_connectivity: bool,
    cutoff: float | None = None,
    cpu_threads: int | None = None,
    shuffle: bool = False,
    num_workers: int = 0,
    pin_memory: bool = False,
    prefetch_factor: int | None = None,
    drop_last: bool = False,
) -> DataLoader[Batch]

Return a DataLoader that forms PyG batches with explicit connectivity.

Worker-based loading uses forkserver so on-disk datasets never inherit open HDF5 handles, and keeps workers alive across iterations. When with_connectivity is true, cutoff is required and connectivity is built independently inside each worker using cpu_threads threads. Otherwise, batches remain edge-free and cutoff and cpu_threads must not be set.