xnn.gnn.models.base.GNNPotential#

class xnn.gnn.models.base.GNNPotential(species, cutoff)[source]#

Bases: InteratomicPotential

Species bookkeeping shared by every GNN potential in this package.

Handles what all graph potentials need regardless of angular basis: an atomic-number -> element-index lookup table, one-hot species node attributes, and the per-element reference energy atom_ref. The spherical-harmonic models subclass it via EquivariantGNN; CACE (Cartesian angular basis) subclasses it directly.

Parameters:

  • species (list[int]) – Atomic numbers of the elements the model supports, in order. The list index defines the element-index (channel) of each species.

  • cutoff (float) – Radial cutoff in Angstrom (used to build neighbor lists).

Variables:

  • species (list[int]) – The supported atomic numbers.

  • z_to_index (Tensor) – Registered long buffer of shape (200,) mapping atomic number to element index (-1 for unsupported elements).

  • atom_ref (torch.nn.Embedding) – Per-element reference-energy embedding (200 entries, one scalar each), initialized to zero.

set_atomic_energies(values)[source]#

Initialise the per-element reference energies atom_ref.

Parameters:

values (array-like) – One reference energy per entry of self.species, in order (MACE E0s, NequIP per_species_rescale_shifts).

Raises:

ValueError – If the number of values does not match the number of species.

Return type:

None

node_attr(atomic_numbers)[source]#

Build the one-hot species node attributes.

Parameters:

atomic_numbers (Tensor) – Long tensor of shape (N,) giving the atomic number of each of the N nodes.

Returns:

One-hot node attributes of shape (N, n_species), cast to the dtype of atom_ref. These are the invariant node_attr inputs consumed by the interaction blocks’ self-connections.

Return type:

Tensor