xnn.gnn.models.nequip.NequIP#

class xnn.gnn.models.nequip.NequIP(species, cutoff=4.0, l_max=2, parity=True, n_rbf=8, n_layers=3, num_features=32, invariant_layers=2, invariant_neurons=64, avg_num_neighbors=None, use_sc=True, resnet=False, nonlinearity_scalars=None, nonlinearity_gates=None, num_polynomial_cutoff=6, trainable_rbf=True, conv_to_output_hidden=None, atomic_energies=None, atomic_scales=None)[source]#

Bases: EquivariantGNN

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

Subclasses EquivariantGNN, inheriting species bookkeeping, the per-element reference energy atom_ref (the NequIP per-species shift) and the shared SphericalHarmonicEdgeEmbedding edge featurizer (configured with the NequIP radial conventions), and adds the NequIP layers and readout. All architecture options are read from ModelConfig.extra (see NequIP.from_config()); upstream NequIP yaml spellings (r_max, num_layers, num_basis, …) are translated to the xnn names at config-load time by the key-translation registry in xnn.common.config.translate.

Parameters:

  • species (list of int) – Atomic numbers of the supported elements, in channel order.

  • cutoff (float, optional) – Radial cutoff r_max in angstrom, by default 4.0.

  • l_max (int, optional) – Maximum rotation order of the hidden features and edge spherical harmonics, by default 2.

  • parity (bool, optional) – Use both parities per l (the full O(3) model), by default True.

  • n_rbf (int, optional) – Number of Bessel basis functions (num_basis), by default 8.

  • n_layers (int, optional) – Number of convnet layers (num_layers), by default 3. 0 gives a pure per-species baseline.

  • num_features (int, optional) – Channel multiplicity of the hidden irreps, by default 32.

  • invariant_layers (int, optional) – Hidden layers of the radial MLP, by default 2.

  • invariant_neurons (int, optional) – Hidden width of the radial MLP, by default 64.

  • avg_num_neighbors (float or None, optional) – Message normalization sqrt(avg_num_neighbors); None (default) disables it. Pass the training-set average (upstream auto).

  • use_sc (bool, optional) – Use the element-dependent self-connection, by default True.

  • resnet (bool, optional) – Residual updates between layers of equal irreps, by default False.

  • nonlinearity_scalars (dict or None, optional) – Per-parity activation names, upstream defaults {"e": "silu", "o": "tanh"}.

  • nonlinearity_gates (dict or None, optional) – Per-parity activation names, upstream defaults {"e": "silu", "o": "tanh"}.

  • num_polynomial_cutoff (int, optional) – Degree p of the polynomial cutoff envelope, by default 6.

  • trainable_rbf (bool, optional) – Learnable Bessel frequencies (upstream default), by default True.

  • conv_to_output_hidden (int or None, optional) – Width of the scalar readout hidden layer; None (default) uses the upstream max(1, num_features // 2).

  • atomic_energies (torch.Tensor or None, optional) – Per-species energy shifts (upstream per_species_rescale_shifts, in the same order as species), used to initialise atom_ref.

  • atomic_scales (torch.Tensor or None, optional) – Per-species energy scales (upstream per_species_rescale_scales with any global rescale folded in). Default is 1 for every species.

Raises:

ValueError – If n_layers is negative.

node_features_energy(atomic_numbers, edge_index, edge_vec)[source]#

TorchScript-compatible core: tensors in, features + energies out.

The single implementation reused by node_energy() (the deploy entry point) and forward(); it avoids the AtomicGraph dataclass. Embeds the species, runs the convnet layers, and reads out both the invariant conv-to-output features (what LatentEwald consumes) and the per-atom energy with the per-species scale/shift applied.

Parameters:

  • atomic_numbers (torch.Tensor) – Per-atom atomic numbers, shape (N,).

  • edge_index (torch.Tensor) – Edge index of shape (2, E); row 0 is the source (neighbour) and row 1 the destination (centre) node of each edge.

  • edge_vec (torch.Tensor) – Edge displacement vectors pos[dst] - pos[src] of shape (E, 3) (already accounting for periodic cell shifts).

Returns:

The invariant conv-to-output features (N, node_feature_dim) and the per-atom energy (N,).

Return type:

tuple of torch.Tensor

node_energy(atomic_numbers, edge_index, edge_vec)[source]#

Per-atom energy, shape (N,) (thin wrapper over node_features_energy(); the deploy wrappers call this).

Parameters:
Return type:

Tensor

forward(data)[source]#

Predict per-node and total energy for an atomic graph.

Thin wrapper over node_energy() (the scriptable tensor core): builds the edge displacement vectors from the graph and aggregates the per-node energies to a per-structure total.

Parameters:

data (xnn.common.data.AtomicGraph) – The input atomic graph (atomic numbers, edge index, positions, …).

Returns:

"node_energy" (per-node energies) and "energy" (per-structure total energy).

Return type:

dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct a NequIP from a core model config.

Reads the NequIP-specific hyper-parameters from cfg.extra (falling back to the upstream defaults) and the shared fields from cfg directly. Only the xnn canonical key names are read here; upstream NequIP yaml spellings (r_max, num_layers, num_basis, chemical_symbols, per_species_rescale_shifts, …) are translated to these names at config-load time by xnn.common.config.translate. Values copied from an upstream yaml are coerced: species accepts atomic numbers or chemical symbols (also as a string form), and atomic_energies / atomic_scales accept a list aligned with species, a {Z: value} dict, a single number (broadcast), or the string form of any of these.

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The core model config, whose extra dict carries the NequIP architecture options.

Returns:

The instantiated model.

Return type:

NequIP