Skip to content

Samples and batches

One PyG Data describes one atomic sample. A Batch combines several samples by concatenating their atoms and physical values, keeping the boundaries needed to recover each sample. The dataset holds shared basis definitions and quantity descriptions; each sample carries the tensors and references needed for computation.

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.

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.

The layouts on this page describe the CPU tensors returned by a dataset or SampleToData. Floating values use float32, atom indices and counts use int64, and periodic-axis flags use bool. Moving or casting a batch later is a separate runtime operation. The storage precision recorded in a physical-data description refers to the HDF5 data.

Forming a batch

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.

In this section, n_atoms, n_blocks, and n_values refer to the totals across the batch. Sample boundaries and atom assignments are explicit:

batch.batch                           int64   [n_atoms]
batch.ptr                             int64   [n_samples + 1]
batch.pos                             float32 [n_atoms, 3]
batch.cell                            float32 [n_samples, 3, 3]
batch.pbc                             bool    [n_samples, 3]
batch.basis_map[role]["atomic_basis_id"] int64 [n_atoms]
batch.basis_map[role]["orb_counts"]      int64 [n_atoms]

For an electronic-data entry, num_values becomes [n_samples], while values_real and optional values_imag are concatenated along their leading axis. For block-sparse entries, num_blocks becomes [n_samples], and atom_pair_index has shape [2, n_blocks] with batch-global atom indices. Cell shifts and per-block lengths remain aligned with the concatenated blocks.

PyG collates nested dictionaries recursively. Under its default rules, fields are concatenated along dimension zero without incrementing values, except keys containing index: these concatenate along the last dimension and are incremented by the number of preceding atoms. This is what turns sample-local atom_pair_index into batch-global indices.

The layouts for each quantity are documented in Physical data structures. Connectivity adds model edges to the resulting batch, and orbital data preparation independently adds output requests or calculated matrices.