Skip to content

elfes.nn

Reusable neural-network representations and building blocks.

GraphIrreps dataclass 

GraphIrreps(
    node_features: Irreps,
    node_attributes: Irreps,
    edge_spherical_harmonics: Irreps,
    edge_radial_features: Irreps,
    edge_cutoff: Irreps,
)

Static representation contract for the fields in GraphTensors.

with_node_features 

with_node_features(irreps: Irreps) -> GraphIrreps

Return the same graph contract with a new learned node state.

GraphTensors dataclass 

GraphTensors(
    node_features: Tensor,
    node_attributes: Tensor,
    edge_index: Tensor,
    edge_spherical_harmonics: Tensor,
    edge_radial_features: Tensor,
    edge_cutoff: Tensor,
)

Graph connectivity and generic e3nn-style node and edge tensors.

The tensors contain no target-specific electronic quantity. Representation metadata belongs to the producing and consuming modules rather than this runtime bundle. Those modules interpret the trailing dimensions as:

  • node_features are the learned node state;
  • node_attributes are fixed even-scalar categorical features;
  • edge_spherical_harmonics carry angular edge features;
  • edge_radial_features and edge_cutoff are invariant edge features.

edge_index[0] is the center and edge_index[1] its neighbor. Messages therefore read neighbor features and aggregate to the center.

BesselBasis 

BesselBasis(cutoff: float, num_basis: int)

Bases: Module

Fixed orthonormal radial sine basis.

For cutoff radius \(r_\mathrm{c}\) and \(n = 1, \ldots, N\), the basis functions are

\[ B_n(r) = \sqrt{\frac{2}{r_\mathrm{c}}} \frac{\sin(n\pi r / r_\mathrm{c})}{r}. \]

They are orthonormal on \([0, r_\mathrm{c}]\) with radial measure \(r^2\,\mathrm{d}r\). At the origin, the implementation uses the analytic limit

\[ B_n(0) = \sqrt{\frac{2}{r_\mathrm{c}}}\frac{n\pi}{r_\mathrm{c}}. \]

The frequencies \(n\pi\) are fixed.

Parameters:

  • cutoff

    (float) –

    Cutoff radius \(r_\mathrm{c}\).

  • num_basis

    (int) –

    Number of basis functions \(N\).

GaussianBasis 

GaussianBasis(cutoff: float, num_basis: int, *, width_scale: float = 2.0)

Bases: Module

Fixed Gaussian basis with uniformly spaced centers.

For \(n = 0, \ldots, N - 1\), define

\[ \begin{aligned} \mu_n &= \frac{n r_\mathrm{c}}{N - 1}, \\ \Delta &= \frac{r_\mathrm{c}}{N - 1}, \\ \sigma &= s\Delta, \end{aligned} \]

The basis functions are

\[ G_n(r) = \exp\left[-\frac{1}{2}\left(\frac{r - \mu_n}{\sigma}\right)^2\right]. \]

The centers and widths are fixed rather than learned. Gaussian functions do not vanish at \(r_\mathrm{c}\), so a separate cutoff remains necessary when edge messages must go smoothly to zero.

Parameters:

  • cutoff

    (float) –

    Cutoff radius \(r_\mathrm{c}\) and location of the final center.

  • num_basis

    (int) –

    Number of Gaussian functions \(N\); must be at least two.

  • width_scale

    (float, default: 2.0 ) –

    Width scale \(s\) relative to the center spacing.

PolynomialCutoff 

PolynomialCutoff(cutoff: float, power: int = 6)

Bases: Module

DimeNet polynomial cutoff with two vanishing derivatives.

With \(x = r / r_\mathrm{c}\) and positive integer power \(p\), the cutoff is

\[ C(r) = \begin{cases} 1 - \dfrac{(p+1)(p+2)}{2}x^p + p(p+2)x^{p+1} - \dfrac{p(p+1)}{2}x^{p+2}, & r < r_\mathrm{c}, \\ 0, & r \ge r_\mathrm{c}. \end{cases} \]

Its value, slope, and curvature all reach zero at \(r_\mathrm{c}\), giving a smooth transition to zero. The default \(p = 6\) is the common NequIP-style choice.

Parameters:

  • cutoff

    (float) –

    Cutoff radius \(r_\mathrm{c}\).

  • power

    (int, default: 6 ) –

    Polynomial power \(p\).

WeightedTensorProduct 

WeightedTensorProduct(
    irreps_in1: Irreps,
    irreps_in2: Irreps,
    irreps_out: Irreps,
    instructions: Sequence[TensorProductInstruction],
)

Bases: Module

External-weight tensor product with automatic provider selection.

Automatic selection prefers OpenEquivariance on supported CUDA inputs and otherwise uses e3nn.

require_matching_graph_irreps 

require_matching_graph_irreps(
    produced: GraphIrreps, consumed: GraphIrreps, *, connection: str
) -> None

Raise when two connected graph modules disagree on representations.

coefficient_squared_l2_loss 

coefficient_squared_l2_loss(
    prediction: OrbData,
    target: OrbData,
    *,
    reduction: Literal["none", "mean", "sum"] = "mean",
) -> Tensor

Return the squared Euclidean error between density coefficients.

For sample \(b\), this function evaluates

\[ \ell_b = \lVert \mathsf d_b - \mathsf d'_b \rVert_2^2. \]

All trailing value dimensions, such as a Pauli component axis, are part of the norm. "mean" averages the complete sample losses rather than individual coefficients. No factor of \(1/2\) is applied.

Parameters:

  • prediction

    (OrbData) –

    Real predicted coefficients in ragged OrbData layout.

  • target

    (OrbData) –

    Real target coefficients with the same layout.

  • reduction

    (Literal['none', 'mean', 'sum'], default: 'mean' ) –

    "none" returns one value per sample; "mean" or "sum" reduces those sample losses.

density_squared_l2_loss 

density_squared_l2_loss(
    prediction: OrbData,
    target: OrbData,
    cholesky_factor: OrbData,
    *,
    reduction: Literal["none", "mean", "sum"] = "mean",
) -> Tensor

Return the squared \(L^2\) error between represented densities.

For sample \(b\), let \(\mathsf U_b^\chi\) be the upper Cholesky factor of the auxiliary-basis overlap,

\[ \mathsf S_b^\chi = \mathsf U_b^{\chi\mathsf T}\mathsf U_b^\chi. \]

The loss is

\[ \ell_b = \left\lVert \mathsf U_b^\chi(\mathsf d_b - \mathsf d'_b) \right\rVert_2^2. \]

cholesky_factor uses the existing batched triu_orb_matrix layout: every sample contributes \(n_b(n_b+1)/2\) upper-triangle values in torch.triu_indices/NumPy triu_indices row-major order. The factor acts independently on every trailing value component. "mean" averages complete sample density errors, and no factor of \(1/2\) is applied.

Parameters:

  • prediction

    (OrbData) –

    Real predicted coefficients in ragged OrbData layout.

  • target

    (OrbData) –

    Real target coefficients with the same layout.

  • cholesky_factor

    (OrbData) –

    Packed real upper Cholesky factors for the same samples.

  • reduction

    (Literal['none', 'mean', 'sum'], default: 'mean' ) –

    "none" returns one value per sample; "mean" or "sum" reduces those sample losses.