elfes.data

PyG sample data, datasets, and derived connectivity.

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.

Shape names

• n_atoms: number of atoms in the sample.
• n_blocks: number of stored blocks in one block-sparse orbital matrix.
• n_values: number of flattened value rows in one electronic-data entry.
• extra_shape: the entry's fixed trailing shape, possibly empty.
• grid_shape: the three-dimensional shape of a uniform grid.
• n_samples: number of samples in a PyG batch.
• n_edges: number of directed connectivity edges in a batch.

Reading one sample

Every item returned by the on-disk and in-memory Datasets is a Data. The Dataset adds its stable member ID for tracing after shuffle or batching; the remaining fields come from the physical Sample:

data.sample_id                         str
data.num_nodes                        int = n_atoms
data.atomic_numbers                   int64   [n_atoms]
data.pos                              float32 [n_atoms, 3]
data.cell                             float32 [1, 3, 3]
data.pbc                              bool    [1, 3]
data.magmoms                          float32 [n_atoms] or [n_atoms, 3]  # optional
data.basis_map                        dict[str, dict[str, Tensor]]
data.physical_data                    PhysicalDataDict

PyG also permits mapping-style access such as data["pos"], but attribute access is the usual form for these top-level fields.

data.basis_map maps every basis role available in the Dataset to its per-atom Torch data. Role names are open strings; common conventions include ao, aux, and paw_coupled:

role_basis_map = data.basis_map[role]
role_basis_map["atomic_basis_id"]    int64 [n_atoms]
role_basis_map["orb_counts"]         int64 [n_atoms]

For atom atom_idx, role_basis_map["atomic_basis_id"][atom_idx] indexes dataset.atomic_basis_tables[role], while role_basis_map["orb_counts"][atom_idx] is the number of basis functions contributed by that atom. dataset.basis_sets[role] retains the complete BasisSet; full definitions are not copied into each Data.

The leading length-one axes of cell and pbc are sample axes: PyG batching concatenates them into [n_samples, 3, 3] and [n_samples, 3]. When the logical dataset carries input atomic magnetic moments, magmoms is a node-level tensor in μB and PyG concatenates it along the atom axis. It is absent for datasets without this input; one dataset does not mix absent, collinear, and noncollinear forms. physical_data is always present and may be empty when no physical data was selected. Its globally unique names combine the atomistic and electronic data stored separately in the physical Sample and ELFES HDF5.

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:

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

Reading atomistic data

For selected atomistic data, data.physical_data[name] is its values tensor. The owning Dataset exposes physical_data_descriptions[name] with its data_type, precision, and extra_shape.

data_type="scalar"                   values float32 [1, *extra_shape]
data_type="atom_scalar"              values float32 [n_atoms, *extra_shape]
data_type="cartesian_vector"         values float32 [1, *extra_shape, 3]
data_type="atom_cartesian_vector"    values float32 [n_atoms, *extra_shape, 3]
data_type="cartesian_tensor_2"       values float32 [1, *extra_shape, 3, 3]
data_type="atom_cartesian_tensor_2"  values float32 [n_atoms, *extra_shape, 3, 3]

The length-one leading axes mark structure-level values so PyG concatenates them along the sample axis. AtomScalar, AtomCartesianVector, and AtomCartesianTensor2 instead use their atom axis as the leading runtime axis and are concatenated in the same order as pos. Their physical arrays place n_atoms immediately before any Cartesian axes; conversion moves that atom axis forward for batching.

Reading electronic data

For selected electronic data, data.physical_data[name] is an ordinary Python dictionary. Its four possible dictionary structures are published as OrbData, BlockSparseOrbData, UniformVolumetricData, and QuadratureVolumetricData; PhysicalDataDict is the named outer mapping shared with atomistic tensors. The dense OrbMatrix hierarchy and OrbVector share OrbData because their PyG fields are identical. For example, select the Hamiltonian data with:

hamiltonian_data = data.physical_data["hamiltonian"]

The data dictionary always contains:

hamiltonian_data["num_values"]          int64   [1]
hamiltonian_data["values_real"]         float32 [n_values, *extra_shape]
hamiltonian_data["values_imag"]         float32 [n_values, *extra_shape]  # complex data only

Thus, for example, the Hamiltonian's real values are read as data.physical_data["hamiltonian"]["values_real"]; fields inside an entry use dictionary keys rather than attribute access.

The entry's data_type, basis_role, pauli, complexity, and fixed extra_shape describe the whole Dataset and are available from dataset.physical_data_descriptions[name] rather than repeated in every sample.

A typical access path therefore looks like:

data = dataset[0]
positions = data.pos

hamiltonian = data.physical_data["hamiltonian"]
values_real = hamiltonian["values_real"]
atom_pair_index = hamiltonian["atom_pair_index"]
description = dataset.physical_data_descriptions["hamiltonian"]

Consumers such as losses can accept one typed electronic-data entry without depending on the complete Data:

def orbital_loss(prediction: OrbData, target: OrbData) -> Tensor:
    ...

Because physical-data names are dataset-defined, callers use the corresponding physical_data_descriptions[name].data_type when choosing or statically narrowing one entry to a specific dictionary type.

The leading values axis is the variable-size axis concatenated by PyG. A complex entry uses parallel real tensors; ELFES does not place complex tensors in Data. When Pauli components are present, their component axis is the final axis of extra_shape.

BlockSparseOrbData similarly names a runtime structure rather than a coordinate system. Dataset entries contain physical orbital-matrix components, while model-side operations may reuse the same pair blocks and boundaries for coupled components.

Block-sparse orbital matrices

data_type="block_sparse_orb_matrix" and data_type="herm_block_sparse_orb_matrix" add:

block_data["num_blocks"]               int64 [1]
block_data["atom_pair_index"]          int64 [2, n_blocks]
block_data["pair_shifts"]              int64 [n_blocks, 3]
block_data["block_lengths"]            int64 [n_blocks]

block_data["atom_pair_index"][:, b] contains sample-local row and column atom indices, while block_data["pair_shifts"][b] identifies the cell image of the column atom. General matrices contain actual directed blocks; Hermitian matrices contain one independent block from each partner pair. The C-order-flattened spatial values of consecutive blocks are concatenated in block_data["values_real"] and, when present, block_data["values_imag"]; block_data["block_lengths"] retains their boundaries. Both use the same PyG dictionary fields, while dataset.physical_data_descriptions[name].data_type retains the mathematical distinction.

Dense orbital matrices

All dense orbital-matrix data types use only the common value fields. For an orbital matrix with n_orbitals orbitals:

• data_type="orb_matrix" uses n_orbitals**2 values in full-matrix C order.
• data_type="herm_orb_matrix" uses n_orbitals * (n_orbitals + 1) // 2 values from the upper triangle, including the diagonal, in NumPy triu_indices row-major order; the omitted lower triangle is its conjugate transpose.
• data_type="triu_orb_matrix" uses the same packed upper order and value count; the omitted strict lower triangle is zero.

Orbital vectors

data_type="orb_vector" also uses only the common value fields. Its orbital axis becomes the leading values axis, so n_values is the total orbital count for the entry's basis role.

Uniform volumetric data

data_type="uniform_volumetric" is real and adds:

uniform_data["origin"]                 float32 [1, 3]
uniform_data["step_vectors"]           float32 [1, 3, 3]
uniform_data["shape"]                  int64   [1, 3]

The three grid axes are flattened in C order into n_values = grid_shape[0] * grid_shape[1] * grid_shape[2] value rows. The stored shape reconstructs those axes; grid periodicity is the sample's pbc.

Quadrature volumetric data

data_type="quadrature_volumetric" is real and adds:

quadrature_data["coordinates"]         float32 [n_values, 3]
quadrature_data["weights"]             float32 [n_values]

Each value row is aligned with one explicit Cartesian quadrature point and weight.

Batching and connectivity

The Dataset output is edge-free: it has no edge_index. A PyG DataLoader concatenates atom, block, and value fields; converts the length-one count fields into per-sample arrays; and increments atom_pair_index from sample-local to batch-global atom indices. Standard PyG batch and ptr fields identify the atom partition.

create_data_loader() forms edge-free batches or explicitly uses ConnectivityCollator to attach directed model connectivity during CPU collation, while add_connectivity() attaches it to an existing CPU or CUDA Batch. Connectivity is stored as edge_index with shape [2, n_edges] and integer edge_shifts with shape [n_edges, 3]. Models reconstruct differentiable edge displacements from pos, cell, and edge_shifts; displacement vectors are not stored by the Dataset.

add_basis_atom_pairs(batch, basis_set) independently generates output requests from geometry and complete numerical-basis cutoffs. It attaches top-level atom_pair_index and pair_shifts, retaining onsite and Hermitian-half pairs within the sum of the two atomic radii. It never reads labels or connectivity. reindex_herm_block_sparse() aligns Hermitian matrix labels to those keys, conjugate-transposes reverse partners, fills absent partner pairs with zeros, and recomputes packed counts. The basis-radius rule is an output truncation choice rather than an exact sparsity theorem for every operator.

add_nao_overlap() may similarly calculate numerical-basis overlap for an already collated CPU Batch. It passes the batch's packed atom arrays through a Spline or Uniform calculator's internal batch execution and attaches the result as BlockSparseOrbData. When a consumer only needs Γ, add_nao_gamma_overlap() attaches its OrbData(values_real, num_values) representation; add_nao_cholesky() directly attaches the packed upper factor. Temporary NumPy batch arrays remain an internal bridge rather than a parallel public matrix hierarchy, and the physics and native modules do not depend on PyG.

All tensors returned by a Dataset retain ELFES physical units and conventions. Dataset conversion does not normalize targets, generate model connectivity, or change the real spherical-harmonic basis. Training code can explicitly compose the DataLoader and connectivity APIs published here with the accumulators in elfes.data.stats without changing the stored samples.

ConnectivityCollator(cutoff: float, cpu_threads: int | None = None)

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

close() -> None

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],
)

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

SampleToData(basis_sets: Mapping[str, BasisSet])

add_basis_atom_pairs(batch: Batch, basis_set: BasisSet) -> Batch

reindex_herm_block_sparse(
    data: BlockSparseOrbData,
    atom_pair_index: Tensor,
    pair_shifts: Tensor,
    orb_counts: Tensor,
    ptr: Tensor,
) -> BlockSparseOrbData

add_connectivity(
    batch: Batch, cutoff: float, *, cpu_threads: int | None = None
) -> Batch

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]

add_nao_cholesky(
    batch: Batch,
    calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
    *,
    name: str = "cholesky",
) -> Batch

add_nao_gamma_overlap(
    batch: Batch,
    calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
    *,
    name: str = "overlap",
) -> Batch

add_nao_overlap(
    batch: Batch,
    calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
    *,
    name: str = "overlap",
) -> Batch