xnn.gnn.models.allegro#

Allegro (Musaelian et al. 2023): strictly local equivariant potential.

A faithful, self-contained re-implementation of the original mir-group/allegro (v0.3.0, the e3nn-era reference of the paper) on the xnn equivariant-GNN abstractions: it subclasses EquivariantGNN (species bookkeeping, per-species energy shift atom_ref, shared edge featurizer) and adds the genuinely Allegro-specific pieces – the two-body scalar embedding, the per-edge scalar / tensor latent tracks, and the iterated weighted-environment tensor products (paper eqs 7-16). Given the same weights it reproduces the original package to machine precision (tests/test_allegro.py and the block-by-block notebook), needing only e3nn – no nequip/allegro/opt_einsum_fx.

Upstream conventions preserved (defaults of the reference uuulin mode):

  • Bessel basis is trainable with the Allegro prefactor r_max / pi (the “normalized sinc” AllegroBesselBasis);

  • spherical harmonics on r_j - r_i (flipped from the xnn edge vector);

  • the environment sum is normalized by 1/sqrt(avg_num_neighbors - 1) (the current edge is subtracted out of its own environment) and the edgewise energy sum by 1/sqrt(avg_num_neighbors);

  • per-channel (“uuu”) weightless tensor products whose Wigner-3j blocks are scaled by sqrt(2 l_out + 1), followed by a strided linear mix with 1/sqrt(mul * n_paths) normalization;

  • variance-preserving scalar MLPs (e3nn.nn.FullyConnectedNet is transplant-identical to upstream’s ScalarMLPFunction);

  • the cumulative-softmax latent resnet coefficients;

  • per-species scale/shift E_i = sigma_Z eps_i + mu_Z of the deployed upstream model (atom_ref shift + atom_scale buffer, as in the xnn NequIP).

The tensor-only Allegro.node_energy() core compiles under torch.jit.script for the LAMMPS/TorchScript exporters in xnn.common.deploy (the pair_allegro deployment path).

Classes

Allegro(species[, cutoff, l_max, parity, ...])

Faithful Allegro (Musaelian et al. 2023): strictly local pair energies.
Show Source