elfes.modules.decoders

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

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(node_features: Tensor) -> Tensor

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(node_features: Tensor, batch: Tensor, ptr: Tensor) -> Tensor

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.

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(
    node_features: Tensor, species_indices: Tensor, batch: Tensor, ptr: Tensor
) -> OrbData