Skip to content

Orbital transforms

Orbital data may need a different coordinate representation or scale for prediction, comparison, or subsequent physical calculations. These transforms operate on the Torch data structures used for batching, while retaining the basis ordering and physical conventions of the underlying quantities.

Clebsch–Gordan coupling resolves each pair of orbital angular momenta into spherical components. Affine transforms apply fixed scales and symmetry-compatible shifts. A Bloch transform combines cell-shift-indexed blocks into matrices at selected k-points. Each operation has its own role and can be composed where needed.

Coupled orbital coordinates

The CG transform preserves the atom-pair block mapping and changes the coordinates within each block. It provides both decomposition into coupled components and composition back into orbital matrix elements.

BlockSparseCGTransform 

BlockSparseCGTransform(basis_set: BasisSet)

Bases: Module

Change packed block-sparse values to and from coupled coordinates.

The transform is defined by one BasisSet. For every ordered pair of shells with angular momenta \((\ell_1, \ell_2)\), it replaces the flattened product basis \(m_1 \otimes m_2\) by the canonical sequence

\[ L = |\ell_1-\ell_2|,\ldots,\ell_1+\ell_2, \qquad M=-L,\ldots,L. \]

Both orbital axes and the coupled M components use ELFES's Wikipedia real spherical-harmonic basis, ordered from -l to +l (l=1: y, z, x). No Cartesian permutation or conversion to another component basis is required. Coefficients come from the vendored NumPy Wigner-3j generator also used by S2NN tensor products, with its real coupling-path phases. Multiplying each unit-Frobenius-norm block by \(\sqrt{2L+1}\) makes the complete matrix orthogonal, so the inverse is its transpose.

Input and output use the same BlockSparseOrbData mapping. Block order, block lengths, and all metadata are preserved; only values_real and an optional values_imag are replaced. The output of decompose is therefore a transient coupled-coordinate view of the same ragged storage, not another physical-data type. Any trailing value axes are carried independently through the spatial transform.

Parameters:

  • basis_set (BasisSet) –

    Basis set defining species, shell, and orbital order.

decompose 

decompose(matrix: BlockSparseOrbData, species_indices: Tensor) -> BlockSparseOrbData

Decompose physical matrix blocks into coupled coordinates.

compose 

compose(
    coupled_matrix: BlockSparseOrbData, species_indices: Tensor
) -> BlockSparseOrbData

Compose coupled coordinates into physical matrix blocks.

Affine transforms

Orbital vectors are transformed shell by shell. Block-sparse matrices are transformed in coupled coordinates, with shared scales within each spherical multiplet and shifts restricted to scalar components. Both directions are available for normalizing data or restoring physical scales.

OrbVectorAffineTransform 

OrbVectorAffineTransform(
    basis_set: BasisSet, shell_shifts: Sequence[Tensor], shell_scales: Sequence[Tensor]
)

Bases: Module

Apply a fixed invertible affine transform to real orbital vectors.

Parameters are specified once per atomic basis and shell in the canonical species order of one BasisSet, and every scale is nonzero. Let atomic site \(i\) belong to species \(s\). A shell is labeled by \((n,\ell)\), where \(n\) enumerates radial shells with the same \(\ell\) in their AtomicBasis order; \(c\) is any trailing channel multi-index that is not transformed by spatial rotations.

For \(\ell=0\), \(x^i_{n0c}\) and \(y^i_{n0c}\) are scalars. For \(\ell>0\), \(\mathbf{x}^i_{n\ell c}\) and \(\mathbf{y}^i_{n\ell c}\) are complete real spherical vectors in the \((2\ell+1)\)-dimensional orbital representation \(V_{\ell,(-1)^\ell}\). The scale \(A^s_{n\ell c}\) and scalar-shell shift \(B^s_{n0c}\) are scalars. The forward transform is

\[ y^i_{n0c} = A^s_{n0c}x^i_{n0c}+B^s_{n0c}, \qquad \ell=0, \]

and

\[ \mathbf{y}^i_{n\ell c} = A^s_{n\ell c}\mathbf{x}^i_{n\ell c}, \qquad \ell>0. \]

Its inverse is

\[ x^i_{n0c} = \frac{y^i_{n0c}-B^s_{n0c}}{A^s_{n0c}}, \qquad \ell=0, \]

and

\[ \mathbf{x}^i_{n\ell c} = \frac{\mathbf{y}^i_{n\ell c}}{A^s_{n\ell c}}, \qquad \ell>0. \]

For \(g\in\mathrm{O}(3)\) and \(\ell>0\), the spherical vector transforms by the real orbital irrep matrix \(D^{(\ell,(-1)^\ell)}(g)\). Because \(A^s_{n\ell c}\) is a scalar,

\[ \begin{aligned} T\!\left(D^{(\ell,(-1)^\ell)}(g)\mathbf{x}^i_{n\ell c}\right) &= A^s_{n\ell c}D^{(\ell,(-1)^\ell)}(g) \mathbf{x}^i_{n\ell c} \\ &= D^{(\ell,(-1)^\ell)}(g) \left(A^s_{n\ell c}\mathbf{x}^i_{n\ell c}\right) \\ &= D^{(\ell,(-1)^\ell)}(g)T(\mathbf{x}^i_{n\ell c}). \end{aligned} \]

For \(\ell=0\), \(D^{(0,+)}(g)=1\), so

\[ T\!\left(D^{(0,+)}(g)x^i_{n0c}\right) = A^s_{n0c}x^i_{n0c}+B^s_{n0c} = D^{(0,+)}(g)T(x^i_{n0c}). \]

Thus the shift is allowed only for scalar shells. The inverse replaces \(A\) by \(A^{-1}\) and \(B\) by \(-A^{-1}B\), so it obeys the same argument and is O(3)-equivariant as well.

Parameters:

  • basis_set (BasisSet) –

    Basis set defining the canonical model species order and orbital layouts.

  • shell_shifts (Sequence[Tensor]) –

    One tensor per entry of basis_set.atomic_bases, with shape [n_shells, *extra_shape].

  • shell_scales (Sequence[Tensor]) –

    Nonzero scales with the same shapes as shell_shifts.

Attributes:

  • orb_shifts (Tensor) –

    Shifts expanded from shells to orbital components, with shape [sum_s n_orb(s), *extra_shape]. The first axis concatenates species in basis_set.atomic_bases order, then each species in shell order and the 2 * ell + 1 spherical components within each shell.

  • orb_scales (Tensor) –

    Scales expanded in the same layout and with the same shape as orb_shifts.

  • orb_counts (Tensor) –

    Number of orbital components per atom for each model species, with shape [num_species] and dtype torch.long.

  • species_offsets (Tensor) –

    Leading-zero prefix sum of orb_counts, with shape [num_species + 1] and dtype torch.long. Species s occupies orb_shifts[species_offsets[s]:species_offsets[s + 1]], and likewise for orb_scales.

from_stats classmethod 

from_stats(basis_set: BasisSet, stats: OrbVectorStats) -> Self

Construct the transform represented by dataset statistics.

identity classmethod 

identity(basis_set: BasisSet, *, extra_shape: tuple[int, ...] = ()) -> Self

Construct an identity transform with the requested value shape.

forward 

forward(vector: OrbData, species_indices: Tensor) -> OrbData

Return values * scale + shift in the model's packed order.

Parameters:

  • vector (OrbData) –

    Real orbital vectors in packed atom order.

  • species_indices (Tensor) –

    Model species index of every atom, with shape [n_atoms] and dtype torch.long.

inverse 

inverse(vector: OrbData, species_indices: Tensor) -> OrbData

Return (values - shift) / scale in the model's packed order.

Parameters:

  • vector (OrbData) –

    Real orbital vectors in packed atom order.

  • species_indices (Tensor) –

    Model species index of every atom, with shape [n_atoms] and dtype torch.long.

BlockSparseAffineTransform 

BlockSparseAffineTransform(
    basis_set: BasisSet,
    onsite_shifts: Tensor,
    onsite_scales: Tensor,
    offsite_shifts: Tensor,
    offsite_scales: Tensor,
)

Bases: Module

Apply a fixed invertible affine transform in coupled coordinates.

For each canonical onsite or offsite path \(a=(s_1,s_2,n_1\ell_1, n_2\ell_2,L)\), all \(2L+1\) components share one nonzero scale \(A_a\). Only \(L=0\) paths may have a shift \(B_a\). For coupled input \(x\) and output \(y\), the forward transform is

\[ y_{aM}=A_a x_{aM}+B_a\delta_{L0}, \]

and its inverse is

\[ x_{aM}=\frac{y_{aM}-B_a\delta_{L0}}{A_a}. \]

Onsite and offsite parameters are distinct, while Hermitian reverse species and shell pairs use the same path parameters.

Both directions accept and return coupled-coordinate BlockSparseOrbData. This module does not couple or decouple block-sparse values and does not assign normalization semantics to either direction; a model or experiment composes it with BlockSparseCGTransform according to its own data flow.

Parameters:

  • basis_set (BasisSet) –

    Basis set defining species, shells, and orbital order.

  • onsite_shifts (Tensor) –

    Scalar-path shifts with shape [n_onsite_paths, *extra_shape]; higher-\(L\) entries must be zero.

  • onsite_scales (Tensor) –

    Nonzero onsite scales with the same shape.

  • offsite_shifts (Tensor) –

    Offsite shifts with shape [n_offsite_paths, *extra_shape]; higher-\(L\) entries must be zero.

  • offsite_scales (Tensor) –

    Nonzero offsite scales with the same shape.

from_stats classmethod 

from_stats(basis_set: BasisSet, stats: BlockSparseStats) -> Self

Construct the transform represented by dataset statistics.

identity classmethod 

identity(basis_set: BasisSet, *, extra_shape: tuple[int, ...] = ()) -> Self

Construct an identity transform with the requested value shape.

forward 

forward(
    coupled_matrix: BlockSparseOrbData, species_indices: Tensor
) -> BlockSparseOrbData

Return values * scale + shift in coupled coordinates.

inverse 

inverse(
    coupled_matrix: BlockSparseOrbData, species_indices: Tensor
) -> BlockSparseOrbData

Return (values - shift) / scale in coupled coordinates.

Bloch transform

The Torch transform is differentiable with respect to matrix values and supports samples with different orbital dimensions. NumPy transforms and generalized eigensolvers are documented under Physical calculations.

bloch_transform 

bloch_transform(
    matrix_r: BlockSparseOrbData,
    orb_counts: Tensor,
    batch_ptr: Tensor,
    kpoints: Tensor,
    *,
    hermitian: bool,
) -> OrbData

Transform cell-shift blocks into compact dense Bloch matrices.

The transform uses fractional reciprocal coordinates and the ELFES forward phase

\[ X(\mathbf k)=\sum_{\mathbf R} e^{+2\pi i\mathbf k\cdot\mathbf R}X(\mathbf R). \]

matrix_r follows the batched BlockSparseOrbData layout, with batch-global atom indices. When hermitian=True, its stored blocks are the independent rows of a Hermitian matrix and the omitted conjugate-transpose partners are included in the result.

Each sample's full dense matrix is flattened in C order along the leading values axis. The sample matrices remain consecutive but are not padded to a common orbital dimension. The operation is expressed entirely in ordinary PyTorch and is differentiable with respect to values_real and optional values_imag.

Parameters:

  • matrix_r (BlockSparseOrbData) –

    Real or complex block values. values_real and optional values_imag have shape [n_values, *extra_shape].

  • orb_counts (Tensor) –

    Number of orbitals on every atom in the batch, with shape [n_atoms].

  • batch_ptr (Tensor) –

    Sample boundaries on the atom axis, with shape [n_samples + 1].

  • kpoints (Tensor) –

    One fractional k-point with shape [3] or multiple k-points with shape [n_kpoints, 3].

  • hermitian (bool) –

    Whether matrix_r stores only independent Hermitian blocks.

Returns:

  • OrbData –

    Full dense matrices as OrbData. num_values contains \(N_n^2\) for

  • OrbData –

    each sample. For multiple k-points, values_real and values_imag

  • OrbData –

    have shape [sum(num_values), n_kpoints, *extra_shape]; the k-point

  • OrbData –

    axis is absent when one [3] k-point is given.