xnn.gnn.models.nequip#

NequIP (Batzner et al. 2022): E(3)-equivariant message-passing potential.

A faithful, self-contained NequIP built on the xnn equivariant-GNN abstractions: it subclasses EquivariantGNN (species bookkeeping, per-element reference energy atom_ref, and the SphericalHarmonicEdgeEmbedding edge featurizer) and adds the genuinely NequIP-specific pieces – the InteractionBlock convolution, the gated ConvNetLayer, and the per-species energy scale/shift of the deployed upstream model.

Architecture (the upstream EnergyModel): one-hot species -> linear chemical embedding -> num_layers convnet layers (each an equivariant convolution of the node features with the edge spherical harmonics, radially weighted, with an element-dependent self-connection and a gated equivariant nonlinearity) -> two linear atom-wise readouts to a per-atom energy, finally per-species scale/shifted (E_i = sigma_Z * eps_i + E0_Z).

Given the same weights this reproduces the original nequip package (mir-group/nequip) to machine precision – the interaction blocks even share the upstream parameter names (linear_1/fc/tp/linear_2/sc) so load_state_dict transplants work directly (see tests/test_nequip.py and examples/fidelity_checks/nequip_verification.ipynb). Only e3nn is required – no nequip / torch_runstats.

Upstream conventions preserved here:

  • the Bessel radial basis is trainable and normalized by 2/r_max (MACE uses fixed weights and sqrt(2/r_max));

  • messages are divided by sqrt(avg_num_neighbors) (MACE divides by the full count);

  • the spherical harmonics are evaluated on r_j - r_i (neighbour minus centre) – the opposite orientation to the xnn/MACE edge vector, so the model flips the edge vectors internally;

  • hidden feature irreps carry both parities per l when parity=True and are pruned per layer to irreps reachable by some tensor-product path.

The model is TorchScript-deployable like MACE: the tensor-only NequIP.node_energy() core compiles under torch.jit.script (used by the LAMMPS/TorchScript exporters in xnn.common.deploy). The e3nn Gate itself does not script on torch 2.x, so _Gate re-implements it exactly (same sorted input layout, same normalize2mom activations).

Functions

nequip_hidden_irreps(num_features, l_max[, ...])

The NequIP hidden feature irreps (upstream feature_irreps_hidden).

Classes

ConvNetLayer(irreps_in, ...[, resnet, ...])

One NequIP layer: InteractionBlock + gated nonlinearity (+ resnet).

InteractionBlock(irreps_in, irreps_out, ...)

The NequIP equivariant convolution (upstream nequip.nn.InteractionBlock).

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

Faithful NequIP (Batzner et al. 2022) with a flexible number of layers.
Show Source