xnn.gnn.models.gemnet.GemNet#

class xnn.gnn.models.gemnet.GemNet(n_features=128, n_edge_features=128, n_interactions=4, n_rbf=6, n_spherical=7, cutoff=5.0, int_cutoff=10.0, p=5, n_triplet_features=64, n_quadruplet_features=32, n_rbf_features=16, n_cbf_features=16, n_sbf_features=32, n_bilinear_triplet=64, n_bilinear_quadruplet=32, n_before_skip=1, n_after_skip=1, n_concat=1, n_atom_residual=2, quadruplets=True, direct_forces=False, forces_coupled=False, output_init='he_orthogonal', trainable_rbf=True, species=None, atomic_energies=None)[source]#

Bases: InteratomicPotential

GemNet: geometric message passing (Gasteiger, Becker, Guennemann, NeurIPS 2021).

Directed edge embeddings m_ca within the embedding cutoff are updated by one-hop messages from m_ba (distance, angle), two-hop messages from m_db via an interaction edge a - b (distances, two angles, dihedral; quadruplets=True) and atom self-interactions; variance scaling factors keep activations normalized. The defaults are the paper’s GemNet-Q (App. F, Table 6).

Parameters:

  • n_features (int, optional) – Atom embedding size, by default 128.

  • n_edge_features (int, optional) – Edge embedding size, by default 128.

  • n_interactions (int, optional) – Interaction blocks, by default 4.

  • n_rbf (int, optional) – Radial basis size N_RBF (and zeros per degree N_SRBF of the circular and spherical bases), by default 6.

  • n_spherical (int, optional) – Degrees N_SHBF of the circular and spherical bases, by default 7.

  • cutoff (float, optional) – Embedding cutoff c_emb in Angstrom, by default 5.0.

  • int_cutoff (float, optional) – Interaction cutoff c_int of the quadruplets’ middle edge, by default 10.0 (unused without quadruplets).

  • p (int, optional) – Exponent of the polynomial envelope, by default 5.

  • n_triplet_features (int, optional) – Down-projected sizes in the triplet / quadruplet interaction, by default 64 and 32.

  • n_quadruplet_features (int, optional) – Down-projected sizes in the triplet / quadruplet interaction, by default 64 and 32.

  • n_rbf_features (int, optional) – Sizes of the shared down-projections of the radial, circular and spherical bases, by default 16, 16 and 32.

  • n_cbf_features (int, optional) – Sizes of the shared down-projections of the radial, circular and spherical bases, by default 16, 16 and 32.

  • n_sbf_features (int, optional) – Sizes of the shared down-projections of the radial, circular and spherical bases, by default 16, 16 and 32.

  • n_bilinear_triplet (int, optional) – Output sizes of the bilinear layers, by default 64 and 32.

  • n_bilinear_quadruplet (int, optional) – Output sizes of the bilinear layers, by default 64 and 32.

  • n_before_skip (int, optional) – Residual blocks before / after the first skip connection and after the atom-edge concatenation, by default 1, 1, 1.

  • n_after_skip (int, optional) – Residual blocks before / after the first skip connection and after the atom-edge concatenation, by default 1, 1, 1.

  • n_concat (int, optional) – Residual blocks before / after the first skip connection and after the atom-edge concatenation, by default 1, 1, 1.

  • n_atom_residual (int, optional) – Residual blocks of the atom embedding and output blocks, by default 2.

  • quadruplets (bool, optional) – Two-hop message passing (GemNet-Q, default) or triplets only (GemNet-T).

  • direct_forces (bool, optional) – Predict forces directly from the edge embeddings (GemNet-dQ / -dT) instead of as the energy gradient; by default False.

  • forces_coupled (bool, optional) – With direct_forces, average the force magnitudes of c -> a and a -> c; by default False.

  • output_init (str, optional) – Initialization of the final layers, "he_orthogonal" (default) or "zeros".

  • trainable_rbf (bool, optional) – Learn the Bessel frequencies, by default True.

  • species (list of int or None, optional) – Only used to interpret atomic_energies.

  • atomic_energies (array-like or None, optional) – Per-species reference energies loaded into atom_ref.

Variables:

  • cutoff (float) – Neighbor-list radius, max(c_emb, c_int) with quadruplets.

  • int_cutoff (emb_cutoff,) – The embedding and interaction cutoffs.

  • scaling_fitted (Tensor) – Boolean buffer, whether fit_scaling_factors() has run.

  • node_feature_dim (int) – Width of the "node_features" output (the summed output-block atom features before their final layer).

set_atomic_energies(species, values)[source]#

Initialize the per-element reference energies atom_ref.

Parameters:

  • species (list of int) – Atomic numbers the values refer to, in order.

  • values (array-like) – One reference energy per entry of species.

Raises:

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

Return type:

None

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

TorchScript-compatible core: tensors in, features + per-atom energy out.

Parameters:

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

  • edge_index (Tensor) – Edge index of shape (2, E) (row 0 source c, row 1 destination a), within cutoff.

  • edge_vec (Tensor) – Edge vectors x_a - x_c, shape (E, 3) (cell shifts included).

Returns:

Per-atom features (N, node_feature_dim) and energies (N,).

Return type:

tuple of Tensor

Raises:

RuntimeError – For the direct-force variants, whose forces are not the energy gradient.

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

Per-atom energy, shape (N,) (the deploy entry point).

Parameters:
Return type:

Tensor

forward(data)[source]#

Compute per-atom and total energies (and direct forces) for a batch.

Parameters:

data (AtomicGraph) – Batched atomic graph built with cutoff.

Returns:

"node_energy" (N,), "energy" (B,), "node_features" (N, node_feature_dim) and, with direct_forces, "forces" (N, 3).

Return type:

dict[str, Tensor]

scale_factors()[source]#

The scaling factors in forward order, keyed by the reference code’s names.

Returns:

QuadInteraction_i_had_rbf, TripInteraction_i_sum_cbf, AtomUpdate_i_sum, OutBlock_i_sum, … (blocks counted from 1, output blocks from 0).

Return type:

dict[str, ScaleFactor]

fit_scaling_factors(batches)[source]#

Fit the scaling factors on data, one after the other (paper section 6).

Each factor is set so that the mean feature variance of its output matches that of its input over all batches, with the factors before it already fitted. Labels are not used.

Parameters:

batches (sequence of AtomicGraph) – Batches of structures (iterated once per factor), on the model’s device.

Returns:

The fitted factors.

Return type:

dict[str, float]

load_scaling_factors(factors)[source]#

Set the scaling factors from a dict or a JSON file of the reference code.

Parameters:

factors (dict or str or pathlib.Path) – {name: value} with the names of scale_factors(), or the path of a scaling_factors.json; entries of absent blocks (e.g. the quadruplet factors for GemNet-T) are ignored.

Raises:

KeyError – If a factor of the model has no value.

Return type:

None

classmethod from_config(cfg)[source]#

Build a GemNet from a configuration object.

Core fields: cfg.n_features (atom embedding size), cfg.n_interactions, cfg.n_rbf and cfg.cutoff (the embedding cutoff). cfg.extra may hold every other constructor option, scaling_factors (a dict or the path of a reference-code JSON file) and the reference code’s triplets_only and envelope_exponent; its other key spellings (emb_size_edge, num_atom, …) are translated by xnn.common.config.translate.

Parameters:

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

Returns:

Instantiated model.

Return type:

GemNet