AIMNet2

Fast neural-network interatomic potentials for organic, elemental-organic, open-shell, reactive, and palladium chemistry.

• Documentation: https://isayevlab.github.io/aimnetcentral/
• Repository: https://github.com/isayevlab/aimnetcentral
• PyPI: https://pypi.org/project/aimnet/
• Demo: https://huggingface.co/spaces/isayevlab/aimnet2-demo

What It Does

AIMNet2 predicts energies, forces, atomic charges, stress tensors, and Hessians for molecular and periodic atomistic simulations. It combines AIMNet2 neural potentials with adaptive dense/sparse neighbor handling, optional long-range electrostatics, DFT-D3 dispersion, and integrations for simulation workflows.

Highlights

• Pretrained model families for general organic chemistry, B97-3c, open-shell chemistry, palladium catalysis, and reactive paths
• Core inference through AIMNet2Calculator
• ASE, PySisyphus, Sella, TorchSim, Hugging Face Hub, and CLI workflows
• Periodic systems with DSF, Ewald, and PME Coulomb methods
• DFT-D3(BJ) dispersion through GPU-accelerated nvalchemiops kernels
• Python 3.11, 3.12, 3.13, and 3.14 support; TorchSim extra is available on Python 3.12+

Install

AIMNet2 requires Python 3.11+ and PyTorch 2.10+.

# CPU or PyTorch-default install
pip install aimnet

# CUDA example: install PyTorch first, then AIMNet2
pip install torch --index-url https://download.pytorch.org/whl/cu126
pip install aimnet

With uv:

uv pip install torch --index-url https://download.pytorch.org/whl/cu126
uv pip install aimnet

Optional integrations:

pip install "aimnet[ase]"              # ASE calculator
pip install "aimnet[pysis]"            # PySisyphus reaction paths
pip install "aimnet[sella]"            # Sella TS optimization, includes ASE
pip install "aimnet[hf]"               # Hugging Face Hub model loading
pip install "aimnet[torchsim,ase]"     # TorchSim, Python 3.12+
pip install "aimnet[train]"            # training and export commands

For a full development environment:

git clone https://github.com/isayevlab/aimnetcentral.git
cd aimnetcentral
make install
source .venv/bin/activate

Quick Start

from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("aimnet2")

result = calc(
    {"coord": coordinates, "numbers": atomic_numbers, "charge": 0.0},
    forces=True,
)

print(result["energy"])   # eV
print(result["forces"])   # eV/Angstrom
print(result["charges"])  # e

compile_model=False and cuda_graph=False are the defaults. Opt into CUDA-graph replay for supported repeated evaluations with cuda_graph=True; use compile_model=True, cuda_graph=False for a compile-only comparison. See CUDA-Graph Inference for supported models, fallback behavior, capture costs, and cache memory.

ASE

from ase.io import read
from aimnet.calculators import AIMNet2ASE

atoms = read("molecule.xyz")
atoms.calc = AIMNet2ASE("aimnet2", charge=0)

energy = atoms.get_potential_energy()
forces = atoms.get_forces()

TorchSim

import ase.io
import torch_sim as ts

from aimnet.calculators import AIMNet2Calculator, AIMNet2TorchSim

atoms = ase.io.read("molecule.xyz")
model = AIMNet2TorchSim(AIMNet2Calculator("aimnet2"))

results = ts.static(system=atoms, model=model)
print(results[0]["potential_energy"], results[0]["forces"])

Periodic Systems

calc = AIMNet2Calculator("aimnet2")
calc.set_lrcoulomb_method("ewald", ewald_accuracy=1e-6)

result = calc(
    {
        "coord": coordinates,
        "numbers": atomic_numbers,
        "charge": 0.0,
        "cell": cell_vectors,
        "pbc": True,
    },
    forces=True,
    stress=True,
)

Model Families

Model	Elements	Best for
aimnet2	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	General organic and elemental-organic chemistry, wB97M-D3
aimnet2-2025	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	B97-3c with improved intermolecular interactions
aimnet2-b973c	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	Legacy B97-3c reproducibility
aimnet2-nse	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	Open-shell systems and radicals
aimnet2-pd	H, B, C, N, O, F, Si, P, S, Cl, Se, Br, Pd, I	Pd catalysis with B97-3c/CPCM(THF) reference data
aimnet2-rxn	H, C, N, O	Reactive chemistry, transition states, NEB, IRC; net-neutral systems only

Each family has four ensemble members indexed 0 through 3. Short aliases load member 0 by default. Previously published aliases such as aimnet2_2025, aimnet2nse, and aimnet2pd continue to resolve.

See the model selection guide for limitations, aliases, citations, and download links.

Hugging Face Models

Install the HF extra and pass a repo ID directly:

pip install "aimnet[hf]"

from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3")
calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", ensemble_member=2)
calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", revision="v1.0")

Published repositories:

HF repo	Alias	Use
isayevlab/aimnet2-wb97m-d3	aimnet2	General organic chemistry
isayevlab/aimnet2-2025	aimnet2-2025	Improved intermolecular interactions
isayevlab/aimnet2-nse	aimnet2-nse	Open-shell chemistry
isayevlab/aimnet2-pd	aimnet2-pd	Palladium chemistry
isayevlab/aimnet2-rxn	aimnet2-rxn	Reactive paths and transition states

Private repos can be loaded with token= or the HF_TOKEN environment variable. Local HF-style directories with config.json and ensemble_N.safetensors are also supported.

Outputs

Key	Shape	Units
energy	scalar or (B,)	eV
forces	(N, 3) or (B, N, 3)	eV/Angstrom
charges	(N,) or (B, N)	electron charge
stress	(3, 3)	eV/Angstrom^3
hessian	(N, 3, N, 3)	eV/Angstrom^2

Hessians are single-molecule only. With compile_model=True, Hessian and HVP requests use the original eager model rather than the compiled inference forward. Long-range backends differ in derivative support; see the calculator and long-range docs for the exact contracts.

Training and CLI

pip install "aimnet[train]"
aimnet train --config my_config.yaml --model aimnet2.yaml

The aimnet entry point is installed with the core package. Training, export, and self-atomic-energy commands require the train extra.

To compile the AIMNet2 energy, force, and stress computation during training, set trainer.compile: true. Compiled training requires CUDA and supports DDP; each process or rank owns one compiled derivative graph. The first batch fixes the requested properties, neighbor mode, and ordered input keys, ranks, dtypes, and device type. Later batches may change their batch, atom, and neighbor dimensions. Forces are supported in modes 0, 1, and 2. Stress requires an explicit mode-1 or mode-2 neighbor matrix, a cell, and an aligned shift tensor for every neighbor matrix.

Stress labels must be full (B, 3, 3) tensors in eV/Å^3 with the ASE sign convention (σ = (1/V) ∂E/∂ε, negative under compression), as returned by atoms.get_stress(voigt=False). VASP reports stress with the opposite sign in kBar (1 eV/Å^3 = 1602.1766 kBar). The built-in SizeGroupedDataset loader produces dense batches without neighbor lists, so stress training requires a custom dataset; aimnet train checks the first batch and stops with an explanation otherwise. See Training.

The original AIMNet2 module remains the sole parameter and checkpoint owner. Loss evaluation, gradient clipping, and the optimizer step remain eager. The trainer invokes loss.backward() from eager Python, which enters the compiled AOTAutograd backward for the model and derivative graph. Hessian and HVP requests made through a calculator with compiled inference enabled continue to use the original eager model.

Development

make check       # formatting, linting, dependency checks
make test        # pytest suite
make docs        # serve docs locally
make docs-test   # strict docs build

CI runs the core test suite on Python 3.11, 3.13, and 3.14 and separate optional-extra lanes for ASE, PySisyphus, Sella, Hugging Face, TorchSim, docs, and security checks.

Citation

If AIMNet2 is useful in your work, please cite the relevant model papers:

AIMNet2

@article{aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet Your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan M and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  volume={16},
  pages={10228--10244},
  year={2025},
  doi={10.1039/D4SC08572H}
}

AIMNet2-2025

If you use AIMNet2-2025 (aimnet2-2025), please also cite the AIMNet2-2025 paper:

@article{nayal2026critical,
  title={Critical benchmarking of machine-learned interatomic potentials for intermolecular and noncovalent interactions},
  author={Nayal, Kamal and Cho, IlKwon and Isayev, Olexandr},
  journal={Machine Learning: Science and Technology},
  year={2026},
  doi={10.1088/2632-2153/aea39f}
}

AIMNet2-NSE: Kalita, B.; Zubatyuk, R.; Anstine, D. M.; Bergeler, M.; Settels, V.; Stork, C.; Spicher, S.; Isayev, O. AIMNet2-NSE: A Transferable Reactive Neural Network Potential for Open-Shell Chemistry. Angew. Chem. Int. Ed. 2026. DOI: 10.1002/anie.202516763

AIMNet2-Pd: Anstine, D. M.; Zubatyuk, R.; Gallegos, L.; Paton, R.; Wiest, O.; Nebgen, B.; Jones, T.; Gomes, G.; Tretiak, S.; Isayev, O. Transferable Machine Learning Interatomic Potential for Pd-Catalyzed Cross-Coupling Reactions. ChemRxiv 2025. DOI: 10.26434/chemrxiv-2025-n36r6

License

See LICENSE.