Skip to content

elfes.modules

Torch operations for geometry, physical data, model outputs, and learning.

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.

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.

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.

AtomCartesianVectorDecoder 

AtomCartesianVectorDecoder(node_channels: int)

Bases: Module

Decode l=1 node features into atom-level Cartesian vectors.

The same channel-mixing weights act on all three l=1 components, which makes the readout SO(3)-equivariant. ELFES spherical tensors use (y, z, x) order at l=1; returned Cartesian vectors use ordinary (x, y, z) order.

Parameters:

  • node_channels

    (int) –

    Channel multiplicity of the spherical node features.

forward 

forward(node_features: Tensor) -> Tensor

Parameters:

  • node_features

    (Tensor) –

    [N_atom, (l_max+1)**2, C_node].

Returns:

  • vectors ( Tensor ) –

    [N_atom, 3].

BlockSparseOrbMatrixDecoder 

BlockSparseOrbMatrixDecoder(
    node_channels: int,
    edge_channels: int,
    basis_set: BasisSet,
    l_max: int,
    m_max: int | None = None,
    *,
    hermitian: bool = False,
    output_transform: BlockSparseAffineTransform | None = None,
)

Bases: Module

Decode real spinless atom-pair blocks, optionally enforcing Hermiticity.

Nodes use global l-major features; edges use local m-major features and global-to-local Wigner matrices. The basis set determines ordered species, shell pairs, coupled irreps and packed output sizes. Outputs contain only requested atom pairs, in request order, independent of connectivity order. With hermitian=True, both directed edges must exist even when only one member of a Hermitian pair is requested. Zero-shift onsite blocks read nodes. Nonzero-shift pairs with equal atom indices read edges.

l_max must cover all output coupled degrees. Truncated local modes are zero-filled before rotation; this restricts the offsite readout function. U(1) real/imaginary components are converted to real spherical components, not physical values_imag. No spin or SOC coupling is performed.

output_transform optionally restores physical scales in global coupled coordinates before CG composition and Hermitian averaging. It must use the same basis set and scalar values (extra_shape=()). Its canonical path policy applies to Hermitian targets only.

forward 

forward(
    node_features: Tensor,
    edge_features: Tensor,
    species_indices: Tensor,
    edge_index: Tensor,
    edge_shifts: Tensor,
    wigner_matrices: Tensor,
    atom_pair_index: Tensor,
    pair_shifts: Tensor,
    batch: Tensor,
    ptr: Tensor,
) -> BlockSparseOrbData

Parameters:

  • node_features

    (Tensor) –

    [N, (l_max+1)**2, C_node].

  • edge_features

    (Tensor) –

    [E, K, C_edge] in m-major layout.

  • species_indices

    (Tensor) –

    [N], in basis-set species order.

  • edge_index

    (Tensor) –

    [2, E].

  • edge_shifts

    (Tensor) –

    [E, 3], integer cell shifts.

  • wigner_matrices

    (Tensor) –

    [E, sum_l (2*l+1)**2].

  • atom_pair_index

    (Tensor) –

    [2, P], batch-global atom indices.

  • pair_shifts

    (Tensor) –

    [P, 3], integer cell shifts.

  • batch

    (Tensor) –

    [N].

  • ptr

    (Tensor) –

    [B+1].

Returns:

  • matrix ( BlockSparseOrbData ) –

    Requested packed blocks, values_real: [V]. Requests are grouped by sample, with arbitrary order within it.

OrbVectorDecoder 

OrbVectorDecoder(node_channels: int, basis_set: BasisSet)

Bases: Module

Decode dense spherical node features into packed orbital vectors.

Each output shell uses a species- and shell-specific linear combination of the input channels. The same weights act on every component within the shell's l block, so the readout is SO(3)-equivariant. Output components follow the atomic-basis shell order and the ELFES real spherical-harmonic convention within each shell.

Attributes:

  • weight –

    Channel-mixing weights for every species shell, with shape [sum_s n_shells(s), node_channels].

  • orb_counts (Tensor) –

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

  • species_offsets (Tensor) –

    Leading-zero prefix sum of orb_counts, with shape [num_species + 1].

  • value_component_indices (Tensor) –

    Spherical component for every species orbital component, with shape [sum_s n_orb(s)].

  • value_weight_indices (Tensor) –

    Shell-weight row for every species orbital component, with shape [sum_s n_orb(s)].

forward 

forward(
    node_features: Tensor, species_indices: Tensor, batch: Tensor, ptr: Tensor
) -> OrbData

Parameters:

  • node_features

    (Tensor) –

    [N_atom, (l_max+1)**2, C_node].

  • species_indices

    (Tensor) –

    [N_atom].

  • batch

    (Tensor) –

    [N_atom].

  • ptr

    (Tensor) –

    [B+1].

Returns:

  • orb_data ( OrbData ) –

    values_real: [N_orb]. num_values: [B].

ScalarDecoder 

ScalarDecoder(node_channels: int, hidden_channels: int)

Bases: Module

Decode spherical node features into one additive scalar per sample.

The decoder maps each node's l=0 features to a scalar contribution and sums the contributions over the nodes belonging to each sample.

Parameters:

  • node_channels

    (int) –

    Channel multiplicity of the spherical node features.

  • hidden_channels

    (int) –

    Width of the scalar readout MLP.

forward 

forward(node_features: Tensor, batch: Tensor, ptr: Tensor) -> Tensor

Parameters:

  • node_features

    (Tensor) –

    [N_atom, (l_max+1)**2, C_node].

  • batch

    (Tensor) –

    [N_atom].

  • ptr

    (Tensor) –

    [B+1].

Returns:

  • scalars ( Tensor ) –

    [B].

SymmetricTensorDecoder 

SymmetricTensorDecoder(node_channels: int, hidden_channels: int)

Bases: Module

Decode intensive symmetric Cartesian tensors [B,3,3] by atom mean.

A scalar MLP predicts the isotropic coefficient; a shared channel linear maps l=2 to the five traceless coefficients in ELFES real harmonics. The six Cartesian basis tensors are orthonormal under the Frobenius product.

neighbor_list 

neighbor_list(
    quantities: str,
    positions: Tensor,
    cell: Tensor,
    pbc: Tensor,
    cutoff: float,
    batch_ptr: Tensor | None = None,
    *,
    algorithm: Literal["auto", "brute_force", "cell_list"] = "auto",
    cpu_threads: int | None = None,
    sorted: bool = False,
    half_list: bool = False,
    include_self: bool = False,
) -> tuple[Tensor, ...]

Build a Torch neighbor list within a strict distance cutoff.

Parameters:

  • quantities

    (str) –

    Returned quantities in order: "i" and "j" select source and target indices, "P" paired indices, "S" integer cell shifts, "d" distances, and "D" displacement vectors.

  • positions

    (Tensor) –

    Cartesian positions with shape (n_atoms, 3) for one structure or (n_total_atoms, 3) for a batch. Tensors may reside on CPU or CUDA and must have dtype float32 or float64.

  • cell

    (Tensor) –

    Cell vectors stored as rows, with shape (3, 3) for one structure or (n_structures, 3, 3) for a batch. Its dtype and device must match positions.

  • pbc

    (Tensor) –

    Periodic-axis flags with shape (3,) for one structure or (n_structures, 3) for a batch, dtype bool, and the same device as positions.

  • cutoff

    (float) –

    Strict, finite, positive distance cutoff in the same length unit as positions and cell.

  • batch_ptr

    (Tensor | None, default: None ) –

    Optional int64 structure boundaries with shape (n_structures + 1,). None denotes one structure.

  • algorithm

    (Literal['auto', 'brute_force', 'cell_list'], default: 'auto' ) –

    Native search algorithm.

  • cpu_threads

    (int | None, default: None ) –

    CPU thread count. None uses one CPU thread and leaves the CUDA option unspecified.

  • sorted

    (bool, default: False ) –

    Whether to sort pairs by source index.

  • half_list

    (bool, default: False ) –

    Whether to retain one member of each reverse-pair class.

  • include_self

    (bool, default: False ) –

    Whether to include zero-shift self pairs.

Returns:

  • Tensor –

    One tensor per character in quantities. For n_edges pairs, i and

  • ... –

    j have shape (n_edges,), P has shape (n_edges, 2), S has

  • tuple[Tensor, ...] –

    shape (n_edges, 3), d has shape (n_edges,), and D has shape

  • tuple[Tensor, ...] –

    (n_edges, 3). Results remain on the input device. Distances and

  • tuple[Tensor, ...] –

    displacements preserve the floating dtype and remain differentiable

  • tuple[Tensor, ...] –

    with respect to positions and cell while the discrete neighbor

  • tuple[Tensor, ...] –

    identity is fixed.

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.

spherical_harmonics 

spherical_harmonics(l_max: int, directions: Tensor) -> Tensor

Evaluate normalized real spherical harmonics through l_max.

The harmonics follow the Wikipedia real spherical harmonic convention and are ordered in consecutive \(l\) blocks, with m = -l, ..., l within each block.

Parameters:

  • l_max

    (int) –

    Highest \(l\) to calculate. Results include every \(l\) from zero through l_max.

  • directions

    (Tensor) –

    Nonzero Cartesian directions with shape (..., 3) on CPU or CUDA and dtype float32 or float64.

Returns:

  • Tensor –

    Harmonic values with shape (..., (l_max + 1) ** 2), preserving the

  • Tensor –

    dtype and device of directions and supporting reverse-mode

  • Tensor –

    differentiation through second derivatives.

wigner_D 

wigner_D(l_max: int, rotation: Tensor) -> Tensor

Return flattened real Wigner \(D\) matrices through l_max.

The matrices use the Wikipedia real spherical-harmonic basis, with each block ordered by m = -l, ..., l and the l = 1 block ordered as \((y,z,x)\). The input acts actively on Cartesian column vectors and the result satisfies \(Y_l(R^\mathsf{T}x)=Y_l(x)D^{(l)}(R)\) for row harmonic values.

Parameters:

  • l_max

    (int) –

    Highest non-negative l. This is a static model configuration under torch.compile.

  • rotation

    (Tensor) –

    Proper rotation matrices in SO(3), with shape (..., 3, 3), dtype float32 or float64, and device CPU or CUDA. Matrix entries are treated as independent variables when differentiating.

Returns:

  • Tensor –

    Consecutive row-major flattened blocks

  • Tensor –

    \(D^{(0)},\ldots,D^{(l_{\max})}\),

  • Tensor –

    with shape

  • Tensor –

    (..., sum((2 * l + 1) ** 2 for l in range(l_max + 1))). Values

  • Tensor –

    preserve the input dtype and device and support reverse-mode

  • Tensor –

    differentiation through second derivatives.