xnn.gnn.models.aimnet2.AIMNet2#

class xnn.gnn.models.aimnet2.AIMNet2(species, cutoff=5.0, n_features=16, n_rbf=16, rbf_start=0.8, gaussian_width=None, hidden=None, n_passes=3, n_vector_combinations=12, aim_size=256, readout_hidden=(128, 128), charge_channels=1, coulomb='simple', coulomb_sr_cutoff=4.6, coulomb_sr_envelope='exp', lr_cutoff=15.0, dsf_alpha=0.2, ewald_accuracy=1e-06, pme_spline_order=4, atomic_energies=None)[source]#

Bases: GNNPotential

The AIMNet2 potential (Anstine et al. 2025); see the module docstring.

Subclasses GNNPotential for the species bookkeeping and the per-element reference energy atom_ref (the paper’s atomic energy shifts, kept in float64 whatever the model dtype).

Parameters:

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

  • cutoff (float, optional) – Radius of the local environment (the Gaussian basis and its cosine cutoff), by default 5.0 Angstrom. The neighbor-list radius self.cutoff is this, or lr_cutoff with a long-range Coulomb method ("dsf", "ewald", "pme").

  • n_features (int, optional) – Embedding channels d per radial shell, by default 16.

  • n_rbf (int, optional) – Gaussian basis functions (radial shells s), by default 16.

  • rbf_start (float, optional) – Center of the first Gaussian, by default 0.8 Angstrom; the centers are rbf_start + s (cutoff - rbf_start) / n_rbf.

  • gaussian_width (float or None, optional) – Exponent eta of the Gaussians; None (default) uses the reference choice (n_rbf / (cutoff - rbf_start))^2 (the inverse squared spacing).

  • hidden (sequence of sequences of int or None, optional) – Hidden widths of the MLP of each message pass; the number of passes is their count (at least two: the first pass predicts the charges, the last the AIM vector). None gives the published architecture [[512, 380], [512, 380], [512, 380, 380]] for n_passes.

  • n_passes (int, optional) – Number of message passes when hidden is None, by default 3.

  • n_vector_combinations (int, optional) – Learned shell combinations h of the vector features, by default 12.

  • aim_size (int, optional) – Width of the AIM vector (the node_features), by default 256.

  • readout_hidden (sequence of int, optional) – Hidden widths of the energy MLP, by default (128, 128).

  • charge_channels (int, optional) – 1 (closed shell, default) or 2 (AIMNet2-NSE: one channel per spin).

  • coulomb (str or None, optional) – "simple" (default): the all-pairs Coulomb energy of molecular structures; "dsf": the damped shifted-force truncation at lr_cutoff over the neighbor list (periodic structures, large molecules); "ewald" / "pme": the Ewald sum of periodic structures, its real-space part over the neighbor list within lr_cutoff and its reciprocal part over the lattice vectors or on a particle mesh (non-periodic structures of the batch get the all-pairs sum); None: no electrostatics (AIMNet2-Pd).

  • coulomb_sr_cutoff (float, optional) – Radius within which the implicitly learned Coulomb energy is subtracted, by default 4.6 Angstrom.

  • coulomb_sr_envelope (str, optional) – Its envelope, "exp" (the mollifier, default) or "cosine".

  • lr_cutoff (float, optional) – Real-space cutoff of the long-range methods, by default 15.0 Angstrom.

  • dsf_alpha (float, optional) – Damping parameter of coulomb="dsf", by default 0.2 per Angstrom.

  • ewald_accuracy (float, optional) – Target relative accuracy of the Ewald and PME sums, by default 1e-6; sets the splitting parameter and the reciprocal cutoff (ewald_parameters()) and the PME mesh.

  • pme_spline_order (int, optional) – B-spline order of the PME charge assignment, by default 4 (cubic).

  • atomic_energies (array-like or None, optional) – Per-element energy shifts, one per entry of species.

Raises:

ValueError – For fewer than two passes, an unknown coulomb method or envelope, charge_channels outside {1, 2}, an accuracy outside (0, 1) or a spline order below 3.

equilibrate(charges, weights, total, batch, num_graphs, eps=1e-06)[source]#

Neural charge equilibration (paper eq 6).

Parameters:

  • charges (Tensor) – Charges before equilibration (N, C).

  • weights (Tensor) – Non-negative weights f (N, C).

  • total (Tensor) – Target total per structure and channel (B, C).

  • batch (Tensor) – Structure index per atom (N,).

  • num_graphs (int) – Number of structures B.

  • eps (float, optional) – Regularization of the weight sum, by default NQE_EPSILON.

Returns:

Charges (N, C) whose structure sums equal total (up to the regularization of the weight sum).

Return type:

Tensor

channel_totals(total_charge, spin_multiplicity)[source]#

The net charge of every structure split over the charge channels, (B, C).

One channel carries Q; two channels carry (Q +- (M - 1)) / 2.

Parameters:

  • total_charge (Tensor) – Net charges (B,).

  • spin_multiplicity (Tensor) – Multiplicities (B,) (read by the two-channel models only).

Return type:

Tensor

node_features_energy_charges(atomic_numbers, edge_index, edge_vec, pos, cell, pbc, total_charge, spin_multiplicity)[source]#

TorchScript core of one structure: AIM features, atomic energies and charges.

The entry point TorchScriptPotential exports; the net charge and spin multiplicity are fixed at export.

Parameters:

  • atomic_numbers (Tensor) – (N,).

  • edge_index (Tensor) – (2, E) within self.cutoff.

  • edge_vec (Tensor) – (E, 3).

  • pos (Tensor) – (N, 3).

  • cell (Tensor) – (3, 3) (ignored unless pbc has a true entry).

  • pbc (Tensor) – (3,) bool.

  • total_charge (float) – Net charge of the structure.

  • spin_multiplicity (float) – Its spin multiplicity (two-channel models).

Returns:

node_features (N, aim_size), node_energy (N,) (float64, Coulomb energy included) and charges (N,).

Return type:

tuple of Tensor

forward(data)[source]#

Predict energies, partial charges and the dipole of a (batched) graph.

Parameters:

data (xnn.common.data.AtomicGraph) – The input graph, built with self.cutoff. total_charge (None = neutral) sets the net charge the partial charges sum to; spin_multiplicity (None = 1) is read by the two-channel (NSE) models.

Returns:

"node_energy" (N,) and "energy" (B,) in float64 (as every xnn model), "node_features" (N, aim_size) (the AIM vector), "charges" (N,), "dipole" (B, 3) in e Angstrom, "energy_coulomb" (B,) (the electrostatic part of the energy; zeros without a Coulomb term) and, for the two-channel models, "spin_charges" (N,).

Return type:

dict of str to torch.Tensor

Raises:

ValueError – For an element outside species, a periodic structure with coulomb="simple" (use "dsf", "ewald" or "pme"), or a partially periodic structure with "ewald" / "pme".

classmethod from_config(cfg)[source]#

Construct an AIMNet2 from a core model config.

Core fields: cfg.cutoff (the local cutoff), cfg.n_features (embedding channels), cfg.n_rbf (radial shells) and cfg.n_interactions (message passes, when extra["hidden"] does not list the MLP widths of every pass). The remaining options are the constructor arguments, read from cfg.extra; the reference code’s spellings (nfeature, nshifts_s, rc_s, ncomb_v, num_charge_channels) are translated at config-load time by xnn.common.config.translate.

Alternatively extra["foundation"] names a published AIMNet2 checkpoint (see from_foundation()); the bare network is returned and every architecture key comes from the checkpoint, while the Coulomb options (coulomb, lr_cutoff, dsf_alpha, ewald_accuracy, pme_spline_order) may be overridden.

Parameters:

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

Returns:

The instantiated model.

Return type:

AIMNet2

Raises:

ValueError – If cfg.cutoff differs from a foundation checkpoint’s neighbor-list cutoff.

classmethod from_foundation(source, dtype=None, cache_dir=None, dispersion=None, **model_options)[source]#

Load a published AIMNet2 model.

Downloads (and caches) the requested checkpoint if needed and converts it weight for weight into this implementation; see xnn.gnn.models.aimnet2_foundation for the registry and the conversion. The published models are served with the D3(BJ) dispersion they were trained without, so by default the returned potential is the network wrapped in the shared D3Dispersion term.

Parameters:

  • source (str, Path or Mapping) – A registered name or alias ("aimnet2", "aimnet2-nse", "aimnet2-wb97m-d3-2", …; see FOUNDATION_MODELS), a checkpoint URL or local path (the reference .pt artifact, or a directory of config.json + ensemble_<k>.safetensors), or an already-loaded artifact mapping.

  • dtype (torch.dtype or str, optional) – Final dtype; None keeps the checkpoint’s float32 (the per-element shifts stay float64 either way).

  • cache_dir (str or Path, optional) – Model hub cache directory.

  • dispersion (dict, str, bool or None, optional) – None (default) adds the D3(BJ) term recorded for the model, False returns the bare AIMNet2; see load_pretrained().

  • **model_options – Deployment knobs applied before the model is built, e.g. coulomb="dsf" with lr_cutoff=15.0, or coulomb="ewald" / "pme" with ewald_accuracy=1e-6, for periodic structures.

Returns:

The AIMNet2 model, inside its dispersion term unless dispersion=False.

Return type:

InteratomicPotential