xnn.gnn.models.aimnet2#

AIMNet2 (Anstine, Zubatyuk and Isayev, 2025).

Chem. Sci. 16, 10228 (2025).

The atoms-in-molecules neural network potential for neutral and charged organic and elemental-organic molecules, written from the paper on the xnn GNN abstractions. The total energy is (paper eq 1)

U = U_local + U_disp + U_coul

with U_local the output of a message-passing network over a radial Gaussian basis, U_coul the Coulomb energy of the partial charges the network predicts (made exact in the net charge by neural charge equilibration, eq 6) and U_disp a DFT-D3(BJ) correction. Block by block:

  • radial basis (eq 2): g_ijs = exp(-eta (r_ij - r_s)^2) f_c(r_ij), n_rbf Gaussians on [rbf_start, cutoff) under a cosine cutoff (GaussianRBF, CosineCutoff);

  • atomic embedding (eq 3): one learned (n_features, n_rbf) matrix a_ds per element, a feature vector per radial shell;

  • convolution (eqs 4-5): the scalar features v_isd = sum_j g_ijs a_jds and the vector features v_ihd = | sum_js g_ijs u_ij a_jds w_dsh |^2 (the squared norm of a learned combination of the shells along the unit bond vectors u_ij), see ShellConvolution;

  • message passing: each pass feeds the atom’s own embedding and its convolved environment to an MLP that returns an embedding update, a partial-charge update and a non-negative weight. From the second pass on the charges enter the input too, through the same convolution. Every charge update is followed by the equilibration (eq 6)

    q_i = q~_i + f_i / sum_j f_j (Q - sum_j q~_j)

    that redistributes the surplus charge according to the predicted weights f; the final pass yields the AIM vector;

  • readout: an MLP from the AIM vector to the atomic energy, plus a per-element shift (the standard xnn atom_ref, kept in float64 as the reference code does since the shifts are O(1000 eV));

  • electrostatics: the published models learn the Coulomb interaction within coulomb_sr_cutoff implicitly, so the short-range part is subtracted under a smooth envelope and the Coulomb energy of the predicted charges is added back: the full 1/r sum over every pair (coulomb="simple", molecules), a damped shifted-force truncation at lr_cutoff (coulomb="dsf", what the paper uses for its periodic CO2 simulation), or the Ewald sum of a periodic cell with the real-space part over the neighbor list within lr_cutoff and the reciprocal part either over the lattice vectors (coulomb="ewald") or on a B-spline charge mesh (coulomb="pme"), both to the relative accuracy ewald_accuracy (electrostatics).

The open-shell variant (AIMNet2-NSE, charge_channels=2) carries two charge channels equilibrated separately to (Q +- (M - 1)) / 2 for a structure of net charge Q and spin multiplicity M (AtomicGraph.spin_multiplicity); their sum is the partial charge and their difference the spin density.

The D3(BJ) correction the published models were fitted without and are served with is not part of this class: the foundation loader (aimnet2_foundation) records it as the config’s subtracted_dispersion and the model hub wraps the model in the shared D3Dispersion term, as for every xnn model.

The network and its Coulomb sums are tensor code that compiles under TorchScript: AIMNet2.node_features_energy_charges() is the core the deploy wrapper (TorchScriptPotential) exports, with the net charge and spin multiplicity fixed at export time.

Fidelity: an independent implementation that follows the conventions of the reference code isayevlab/aimnetcentral where the paper leaves them open (Gaussian centers r_s = rbf_start + s (cutoff - rbf_start) / n_rbf with eta = (n_rbf / (cutoff - rbf_start))^2, the squared vector norm, the 1e-6 regularization of the equilibration weights, the shifted-force Coulomb with its self term, the Ewald/PME sums with their self and background terms and the Kolafa-Perram parameter balance, the float64 accumulation of the pair sums), so that the published checkpoints transplant weight for weight and reproduce the reference energies, forces and charges to float32 round-off, and the Ewald and PME energies to the requested accuracy (tests/test_aimnet2.py, examples/fidelity_checks/aimnet2_verification.ipynb). One documented difference: the reference chooses the Ewald real-space cutoff per structure from the accuracy, xnn keeps it at lr_cutoff (the radius of the graph the data pipeline builds) and balances the splitting parameter to it. Units: eV, Angstrom, elementary charges.

Module Attributes

COULOMB_METHODS

the all-pairs sum of molecules, the damped shifted-force truncation, the Ewald and particle-mesh Ewald sums of periodic cells, or none.

LONG_RANGE_METHODS

The methods whose neighbor list reaches lr_cutoff.

SR_ENVELOPES

Envelopes that switch off the implicitly learned short-range Coulomb energy.

NQE_EPSILON

Regularization of the equilibration weights, sum_j f_j + eps (eq 6).

COULOMB_OPTIONS

The constructor options a deployment may override on a published model.

Classes

AIMNet2(species[, cutoff, n_features, ...])

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

ShellConvolution(n_channels, n_shells, ...)

The AIMNet2 convolution of per-atom features over the radial shells (paper eqs 4-5).
Show Source