AIMNet2-rxn

AIMNet2-rxn is a neural network interatomic potential specialized for generalized organic reaction modeling — energies, forces, charges, and electrostatic moments along reaction coordinates at near-ωB97M-V/def2-TZVPP accuracy. Trained on ~4.7M reaction-relevant geometries to handle bond breaking/forming, transition states, and minimum-energy pathways out of the box, without system-specific fine-tuning.

Highlights

• Reaction-aware: ~1 kcal mol⁻¹ across IRC profiles, ~1.6 kcal mol⁻¹ RMSD on reaction energies, ~2 kcal mol⁻¹ RMSD on diverse-reaction barrier heights, ~3 kcal mol⁻¹ RMSD on the Grambow benchmark — all without retraining or fine-tuning. (Numbers from the paper; see below.)
• High-throughput pathway search: Enables minimum-energy-pathway searches at ~10⁶ reactions/day via batched NEB on a single GPU — roughly 10⁶× the throughput of serial NEB-DFT at the reference level. Per-call speedup is closer to 10³–10⁴×; the 10⁶× figure is end-to-end batched NEB throughput, not per-energy-evaluation.
• TS- and NEB-ready: Designed for transition-state optimization and (batched) nudged-elastic-band minimum-energy-pathway searches.
• Ensemble: 4 ensemble members for uncertainty estimation.
• Outputs: energy, forces (autograd), atomic charges, molecular dipole, molecular quadrupole.

Element coverage

H, C, N, O only. The model was trained on closed-shell organics. Unsupported elements (F, P, S, …) are not validated by the calculator and will silently produce undefined predictions — read implemented_species from this repo's config.json ([1, 6, 7, 8]) and check input atomic numbers yourself before running on novel systems. For broader element coverage on equilibrium structures, use isayevlab/aimnet2-wb97m-d3 or related families.

Energy convention

AIMNet2-rxn outputs are on a learned per-element shifted electronic-energy scale (eV). They are designed to be accurate for energy differences between geometries that share atomic composition — barrier heights, reaction energies, IRC profiles, conformer energies. They are not absolute ωB97M-V/def2-TZVPP electronic energies and not formal atomization energies.

Practical implications:

• Methane sits near a few eV on this scale, not the −1100 eV scale used by aimnet2-wb97m-d3 and the other AIMNet2 families. Do not mix or compare energies across families.
• Single-point energies in isolation are not physically meaningful. Always interpret as a difference vs. another geometry from this same model.

Installation

pip install "aimnet[hf]"

Quick Start

from aimnet.calculators import AIMNet2Calculator

# Load from Hugging Face (downloads and caches automatically)
calc = AIMNet2Calculator("isayevlab/aimnet2-rxn")

# Single-point calculation
results = calc(
    {"coord": coords, "numbers": atomic_numbers, "charge": 0.0},
    forces=True,
)
print(results["energy"])      # Energy in eV (atomization-style scale, see "Energy convention")
print(results["forces"])      # Forces in eV/A
print(results["charges"])     # Partial charges in e

With ASE — minimum energy path via NEB

from aimnet.calculators.aimnet2ase import AIMNet2ASE
from ase.mep import NEB
from ase.optimize import LBFGS

# images: list[ase.Atoms] interpolated between reactant and product
calc_factory = lambda: AIMNet2ASE("isayevlab/aimnet2-rxn")
for img in images:
    img.calc = calc_factory()

neb = NEB(images, climb=True)
LBFGS(neb).run(fmax=0.05)

For batched-NEB (BNEB) on millions of pathways, see the AIMNet2-rxn paper for the recommended workflow with pysisyphus.

Transition-state and Hessian workflows

For TS optimization or any workflow that requires a Hessian (e.g. P-RFO, dimer, IRC, vibrational analysis), set compile_model=False when constructing the calculator:

calc = AIMNet2Calculator("isayevlab/aimnet2-rxn", compile_model=False)

torch.compile (Dynamo) is known to hang on the double-backward path through GELU activations that Hessian autograd requires. The default avoids the issue, but explicit compile_model=True for speed should be reserved for pure energy/force evaluations.

Ensemble Members

This repo contains 4 ensemble members (ensemble_0.safetensors through ensemble_3.safetensors). The default loads member 0. For uncertainty estimation, load all 4 and average predictions.

File	Description
ensemble_0.safetensors	Ensemble member 0 (default)
ensemble_1.safetensors	Ensemble member 1
ensemble_2.safetensors	Ensemble member 2
ensemble_3.safetensors	Ensemble member 3
config.json	Shared model configuration and metadata

Model Details

• Architecture: AIMNet2 (Atomic Environment Vectors + message-passing MLPs) with dipole and quadrupole heads
• Cutoff radius: 5.0 Å
• DFT reference: ωB97M-V / def2-TZVPP (range-separated hybrid meta-GGA with VV10 nonlocal correlation; Mardirossian & Head-Gordon, J. Chem. Phys. 144, 214110, 2016)
• Dispersion: Supplied by ωB97M-V's VV10 nonlocal correlation kernel, included in every training-set energy and force. No separate D3/D4 correction is added at inference (needs_dispersion: false in config.json); doing so would double-count the dispersion already baked into the reference data. AIMNet2's external D3 path remains available for other family members whose reference functional does not carry its own dispersion (e.g. aimnet2-wb97m-d3).
• Coulomb: Short-range embedded (≤4.6 Å) + external long-range (≥4.6 Å, applied by the calculator via LRCoulomb). The 4.6 Å crossover is physically frozen — it is the SR/LR cancellation point used during training. Do not alter coulomb_sr_rc at inference.
• Training set: 4,685,165 geometries combining the AIMNet2 base set with the RGD-1 reaction database, all recomputed at ωB97M-V/def2-TZVPP
• Format: safetensors (converted from JIT TorchScript via reconstructed YAML)

Intended use

Mechanism enumeration, transition-state searches, NEB / BNEB, reaction-coordinate energy profiles, high-throughput reactivity screening on closed-shell organic systems.

Calculator-enforced safeguards (with aimnet>=NEXT_VERSION)

The calculator validates inputs and refuses to silently produce undefined output:

• Element scope — passing atomic numbers outside [1, 6, 7, 8] raises ValueError (bypass with validate_species=False).
• Net charge — charge != 0 raises ValueError (zwitterions are OK as long as the net charge is zero; for net-charged species use isayevlab/aimnet2-wb97m-d3).
• Hessian + torch.compile — combining both raises RuntimeError (Dynamo and GELU double-backward hang); reconstruct the calculator with compile_model=False for TS / IRC / vibrational analysis.
• Coulomb cutoff — set_lrcoulomb_method(...) warns if you change the SR/LR crossover from 4.6 Å (physically frozen during training).
• Cross-family mixing — constructing calculators from two AIMNet2 families in one process emits a one-time UserWarning about energy-scale incompatibility.

afv.weight rows for elements outside [1, 6, 7, 8] are NaN-padded at the safetensors level, so even with validate_species=False the model produces NaN (not plausible-looking nonsense) for unsupported elements. Two layers of defense.

Limitations

• Element scope: H, C, N, O only — passing other elements is not validated by the calculator and yields undefined output.
• Closed-shell only: Open-shell species (radicals, doublets) are not handled by AIMNet2-rxn. For open-shell reactions and radical chemistry, use isayevlab/aimnet2-nse — a fully reactive model with parametrization focused on open-shell systems and radicals.
• Net-charged species (anions, cations): Out of scope for this release; future work per the paper. Net-neutral zwitterions are within the training distribution and supported — local charge separation in a globally neutral molecule (e.g. amino acid neutral forms, proton-transfer intermediates) is a deliberate part of training.
• Transition metals & heteroatoms beyond H/C/N/O: Not supported.
• Equilibrium MD on large organics: While accurate near equilibrium, the wb97m-d3 family is generally preferred when reaction information is not needed.
• Hessian/TS workflows under torch.compile: Set compile_model=False (see Quick Start). torch.compile and Hessian autograd through GELU activations are known to interact poorly.

Citation

@misc{anstine2025aimnet2rxn,
  title  = {AIMNet2-rxn: A Machine Learned Potential for Generalized Reaction Modeling on a Millions-of-Pathways Scale},
  author = {Anstine, Dylan M. and Zhao, Qiyuan and Zubatiuk, Roman and Zhang, Shuhao and Singla, Veerupaksh and Nikitin, Filipp and Savoie, Brett M. and Isayev, Olexandr},
  year   = {2025},
  doi    = {10.26434/chemrxiv-2025-hpdmg},
  note   = {ChemRxiv preprint}
}

License

MIT License