xnn.gnn.models.allegro.Allegro#

class xnn.gnn.models.allegro.Allegro(species, cutoff=6.0, l_max=1, parity='o3_full', n_rbf=8, num_layers=2, num_tensor_features=32, two_body_latent=None, latent=None, env_embed=None, edge_eng=None, initial_scalar_embedding_dim=None, avg_num_neighbors=None, latent_resnet=True, num_polynomial_cutoff=6, trainable_rbf=True, atomic_energies=None, atomic_scales=None)[source]#

Bases: EquivariantGNN

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

The energy is decomposed into pair energies, E = sum_i [ sigma_Z_i (sum_j E_ij / sqrt(lambda)) + mu_Z_i ] (paper eqs 5-6). Each ordered pair ij carries an invariant scalar latent and an equivariant tensor latent that interact at every layer: the scalar latent generates the weights embedding atom i’s environment, a per-channel tensor product couples the pair tensors with that embedded environment (eqs 11-14), its scalar outputs feed back into the scalar latent (eq 15), and a linear layer mixes the tensor channels (eq 16). There is no message passing – every pair energy is a function of the fixed local environment, which is what makes Allegro scalable.

All architecture options are read from ModelConfig.extra (see from_config()); upstream allegro yaml spellings are translated to the xnn names at config-load time by 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, by default 6.0.

  • l_max (int, optional) – Maximum rotation order of the tensor track, by default 1.

  • parity (str, optional) – "o3_full" (default; all irreps up to l_max), "o3_restricted" (only SH irreps) or "so3" (no parity).

  • n_rbf (int, optional) – Bessel basis size, by default 8.

  • num_layers (int, optional) – Number of tensor-product layers (>= 1), by default 2.

  • num_tensor_features (int, optional) – Channel multiplicity of the tensor track, by default 32.

  • two_body_latent (list of int, optional) – Hidden+output widths of the two-body scalar MLP (upstream two_body_latent_mlp_latent_dimensions), default [32, 64, 128].

  • latent (list of int, optional) – Hidden+output widths of the later latent MLPs (upstream latent_mlp_latent_dimensions), default [128]. The final width must match two_body_latent[-1] when latent_resnet is on.

  • env_embed (list of int, optional) – Hidden widths of the environment-embedding MLPs (upstream env_embed_mlp_latent_dimensions), default [] (linear).

  • edge_eng (list of int, optional) – Hidden widths of the final edge-energy MLP (upstream edge_eng_mlp_latent_dimensions), default [32].

  • initial_scalar_embedding_dim (int or None, optional) – Width of the two-body radial-chemical product embedding; None (default) uses two_body_latent[0] (upstream default). Must be even.

  • avg_num_neighbors (float or None, optional) – Environment/energy-sum normalization; None (default) disables it.

  • latent_resnet (bool, optional) – Cumulative-softmax residual updates of the scalar latent, default True.

  • num_polynomial_cutoff (int, optional) – Polynomial cutoff degree p, by default 6.

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

  • atomic_energies (array-like or None, optional) – Per-species energy shifts / scales (as in the xnn NequIP).

  • atomic_scales (array-like or None, optional) – Per-species energy shifts / scales (as in the xnn NequIP).

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

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

Parameters:

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

  • edge_index (torch.Tensor) – Edge index (2, E); row 0 is the neighbour, row 1 the centre.

  • edge_vec (torch.Tensor) – Edge vectors pos[dst] - pos[src], shape (E, 3).

Returns:

The invariant node features (N, node_feature_dim) – the final per-edge scalar latents summed onto their centre atoms (what LatentEwald consumes) – 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.

Parameters:

data (xnn.common.data.AtomicGraph) – The input atomic graph.

Returns:

"node_energy" (per-atom energies) and "energy" (per-structure totals).

Return type:

dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct an Allegro from a core model config.

Reads the Allegro hyper-parameters from cfg.extra; upstream yaml spellings are translated by xnn.common.config.translate and value forms coerced by xnn.common.config.coerce.

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The core model config.

Returns:

The instantiated model.

Return type:

Allegro