Orbital data preparation
Orbital models may need an explicit set of matrix blocks to predict, or an overlap matrix for a physical loss. These operations prepare such data for an existing batch. Output atom pairs are selected from geometry and basis cutoffs independently of the model's connectivity.
Output atom pairs and label alignment
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.
The top-level batch.atom_pair_index and batch.pair_shifts describe requested outputs. A reference matrix keeps its own block keys inside batch.physical_data[name] until it is explicitly aligned to those requests.
add_basis_atom_pairs(batch: Batch, basis_set: BasisSet) -> Batch
Attach Hermitian-half atom pairs within the sum of numerical-basis radii.
Each atomic radius is the largest shell cutoff in Å. Include all onsite
blocks and pairs satisfying distance < radius_i + radius_j, retaining
the lexicographically first (i, j, S) of each Hermitian partner pair.
This is an explicit output truncation rule, not a claim that every
Hamiltonian vanishes outside these distances.
Only geometry (pos, cell, pbc, atomic_numbers, ptr) is read.
Attach atom_pair_index: [2, P] and pair_shifts: [P, 3], grouped by
sample, on the input device. Labels and model connectivity are untouched.
A layout-only or Gaussian basis cannot supply finite numerical cutoffs.
reindex_herm_block_sparse(
data: BlockSparseOrbData,
atom_pair_index: Tensor,
pair_shifts: Tensor,
orb_counts: Tensor,
ptr: Tensor,
) -> BlockSparseOrbData
Gather requested blocks, conjugate-transposing partners and filling zeros.
data stores unique lexicographically first Hermitian partner keys, as
returned for an ELFES Hermitian block-sparse matrix. Missing partner pairs
denote zeros. Requested keys may use either direction, in arbitrary order
within each sample; samples must remain grouped in batch order. All atom
indices are batch-global. orb_counts: [N] gives the spatial block shapes
and ptr: [B+1] partitions the atoms. Real/imaginary extra axes are retained.
Returns values in exactly the requested order with recomputed block/value counts. Supports CPU/CUDA and gradients with respect to stored values.
Numerical-basis overlaps
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.
These helpers attach entries to batch.physical_data under the requested name and require a CPU batch. The reusable calculators and their basis requirements are documented under Numerical basis calculations. The Γ overlap and Cholesky forms connect to the quadratic orbital-vector losses.
add_nao_overlap(
batch: Batch,
calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
*,
name: str = "overlap",
) -> Batch
Calculate and attach numerical-basis overlap to a CPU PyG batch.
add_nao_gamma_overlap(
batch: Batch,
calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
*,
name: str = "overlap",
) -> Batch
Calculate and attach packed real Γ overlap to a CPU PyG batch.
add_nao_cholesky(
batch: Batch,
calculator: SplineNumericalOverlapCalculator | UniformNumericalOverlapCalculator,
*,
name: str = "cholesky",
) -> Batch
Calculate and attach packed upper Cholesky factors to a CPU PyG batch.