Edge encoder: `radial_angular_encoder`

Role

The edge encoder produces the initial pair state \(\bm e_{ij}^{(0)}\). Unlike nodes, edges have a direction, so their initial state can already contain non-scalar irreps built from the spherical harmonics of \(\hat{\bm r}_{ij}\).

Built-in members of this kind: pair_radial_angular_embedding (MANDALA’s rich encoder, a fully connected tensor product of a pair-type embedding with \(\bm\varphi\oplus\mathrm{TS}(Y)\)), and gaussian_basis used directly as a scalar edge state (DeepH-E3 preset).

The new layer

The radial–angular encoder gives every output copy \((\ell,k)\) its own learned radial profile:

\[\bm e_{ij}^{(\ell,k)}=R_{\ell k}\big(\bm\varphi(r_{ij})\big)\;Y^{\ell}(\hat{\bm r}_{ij}), \qquad R=\mathrm{MLP}(\bm\varphi):\mathbb R^{n}\to\mathbb R^{\sum_\ell m_\ell}.\]

This is the most general equivariant function of a single displacement vector that is linear in the angular basis: by Schur’s lemma any equivariant map from a vector to a copy of degree \(\ell\) is \(f(r)\,Y^\ell(\hat{\bm r})\). The encoder therefore spans exactly the space of “two-body” edge features, with \(m_\ell\) independent radial functions per degree.

Symmetry. \(R_{\ell k}\) is a function of invariants; multiplying the irrep \(Y^\ell\) by an invariant gives a copy of the same irrep (checklist items 2 and 3c). Requested irreps must appear in the angular basis; the factory rejects, e.g., 1e from harmonics that only contain 1o.

Implementation

class RadialAngularEncoder(nn.Module):
    r"""``e_ij^{(l,k)} = R_{l,k}(phi(r_ij)) Y_l(r_hat_ij)``.

    A radial MLP produces one invariant amplitude per output copy ``(l, k)``;
    each copy carries the spherical harmonic of its degree.
    """

    def __init__(self, radial_dim, angular_irreps, irreps_out, hidden=32):
        super().__init__()
        self.angular_irreps = o3.Irreps(angular_irreps)
        self.irreps_out = o3.Irreps(irreps_out)
        angular = {
            ir: sl
            for (_, ir), sl in zip(self.angular_irreps, self.angular_irreps.slices())
        }
        missing = [str(ir) for _, ir in self.irreps_out if ir not in angular]
        if missing:
            raise ValueError(f"Angular basis lacks output irreps {missing}")
        self.slices = [angular[ir] for _, ir in self.irreps_out]
        self.radial = nn.Sequential(
            nn.Linear(radial_dim, hidden),
            nn.SiLU(),
            nn.Linear(hidden, self.irreps_out.num_irreps),
        )

    def forward(self, radial, angular):
        amplitudes = self.radial(radial)
        blocks, cursor = [], 0
        for (mul, ir), sl in zip(self.irreps_out, self.slices):
            weight = amplitudes[:, cursor : cursor + mul, None]
            blocks.append(
                (weight * angular[:, None, sl]).reshape(len(radial), mul * ir.dim)
            )
            cursor += mul
        return torch.cat(blocks, dim=-1)

@register_transform("radial_angular_encoder")
def _radial_angular_encoder(options, inputs, mapper):
    radial, angular = inputs
    radial_irreps = global_irreps(radial, "edge")
    if not invariant_scalars(radial_irreps):
        raise ValueError("Radial features must be invariant scalars")
    support = require_same_support(radial, angular)
    module = RadialAngularEncoder(
        radial_irreps.dim, global_irreps(angular, "edge"), **options
    )
    return module, [
        FeatureSpec("edge", "irreps", str(module.irreps_out), support=support)
    ]

Usage

- id: angular
  transform: spherical_harmonics
  inputs: [edge_vectors]
  outputs: [angular]
  options: {irreps: 1x0e+1x1o+1x2e+1x3o+1x4e}
- id: edge_encoder
  role: encoder
  transform: radial_angular_encoder
  inputs: [radial, angular]
  outputs: [edge_0]
  options: {irreps_out: 16x0e+8x1o+8x2e+4x3o+4x4e}

Variations

Condition the radial MLP on the species pair (concatenate a pair-type embedding to \(\bm\varphi\)), giving each bond type its own radial profiles.
Append the tensor square of the harmonics to reach 1e, 2o… directly in the encoder.

Edge encoder: radial_angular_encoder