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 electronic-data descriptions remain on the owning Dataset.

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 either HDF5-backed Dataset 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.electronic_data                  ElectronicDataDict

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. electronic_data is always present and may be empty when no electronic data was selected.

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 electronic data

data.electronic_data maps each selected electronic-data name to an ordinary Python dictionary. Its four possible dictionary structures are published as OrbData, BlockSparseOrbMatrixData, UniformVolumetricData, and QuadratureVolumetricData; ElectronicDataDict is the named outer mapping. The dense OrbMatrix hierarchy and OrbVector share OrbData because their PyG fields are identical. For example, select the Hamiltonian data with:

hamiltonian_data = data.electronic_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.electronic_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.data_descriptions[name] rather than repeated in every sample.

A typical access path therefore looks like:

data = dataset[0]
positions = data.pos

hamiltonian = data.electronic_data["hamiltonian"]
values_real = hamiltonian["values_real"]
atom_pair_index = hamiltonian["atom_pair_index"]
description = dataset.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 electronic-data names are dataset-defined, callers use the corresponding 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.

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

ConnectivityCollator or add_connectivity() may then attach directed model connectivity to a Batch 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_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 BlockSparseOrbMatrixData. 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 retain ELFES physical units and conventions. Dataset conversion does not normalize targets, generate model connectivity, or change the real spherical-harmonic basis.

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

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

close() -> None

InMemoryDataset(
    collated_data: Data,
    slice_dict: dict[str, Any],
    *,
    sample_ids: tuple[str, ...],
    basis_sets: dict[str, BasisSet],
    data_descriptions: dict[str, ElectronicDataDescription],
    metadata: dict[str, str],
    electronic_data_names: Collection[str],
)

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

SampleToData(basis_sets: Mapping[str, BasisSet])

add_connectivity(
    batch: Batch, cutoff: float, *, cpu_threads: int | None = None
) -> 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