aimnet

• Documentation: https://isayevlab.github.io/aimnetcentral/
• Repository: https://github.com/isayevlab/aimnetcentral/

AIMNet2 : ML Interatomic Potential for Fast and Accurate Atomistic Simulations

AIMNet2 is a neural network interatomic potential that predicts energies, forces, atomic charges, stress tensors, and Hessians for organic and elemental-organic molecules. It supports 14 elements (H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I) with specialized models for open-shell chemistry and palladium catalysis.

Key Features

• Accurate for neutral, charged, organic, and elemental-organic systems
• GPU-accelerated with NVIDIA Warp CUDA kernels and torch.compile support
• ASE, PySisyphus, and TorchSim calculator interfaces
• Periodic boundary conditions with DSF and Ewald Coulomb methods
• DFT-D3 dispersion corrections (BJ damping, GPU-accelerated)
• Adaptive neighbor lists with automatic dense/sparse mode selection

Installation

Requirements

• Python 3.11+
• PyTorch 2.8+ (pytorch.org)

Using pip

# Install PyTorch first (with CUDA if you have a GPU)
pip install torch --index-url https://download.pytorch.org/whl/cu126

# Install AIMNet2
pip install aimnet

Using uv (recommended for fast installs)

# Install PyTorch + AIMNet2
uv pip install torch --index-url https://download.pytorch.org/whl/cu126
uv pip install aimnet

Using conda/mamba

# Create environment with PyTorch from conda-forge
mamba create -n aimnet python=3.12 pytorch pytorch-cuda=12.6 -c pytorch -c nvidia -c conda-forge
mamba activate aimnet

# Install AIMNet2 via pip (not yet on conda-forge)
pip install aimnet

Optional Extras

pip install "aimnet[ase]"             # ASE calculator interface
pip install "aimnet[pysis]"           # PySisyphus reaction path calculator
pip install "aimnet[sella]"           # Sella TS optimizer (includes ASE)
pip install "aimnet[hf]"              # Hugging Face Hub model loading
pip install "aimnet[torchsim,ase]"    # TorchSim integration (Python 3.12+)
pip install "aimnet[train]"           # Training pipeline (W&B, ignite)
pip install "aimnet[ase,pysis,sella,hf,torchsim,train]" # All extras available on this Python

Development Setup

git clone https://github.com/isayevlab/aimnetcentral.git
cd aimnetcentral
make install        # Creates venv, installs all extras + dev tools
source .venv/bin/activate

Available Models

Each model has 4 ensemble members (0-3). Models are auto-downloaded on first use. Previously published aliases (aimnet2_2025, aimnet2nse, aimnet2pd, etc.) continue to resolve.

Quick Start

Core Calculator

from aimnet.calculators import AIMNet2Calculator

calc = AIMNet2Calculator("aimnet2")

results = calc(
    {"coord": coordinates, "numbers": atomic_numbers, "charge": 0.0},
    forces=True,
)
print(results["energy"], results["forces"])

ASE Integration

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()

Periodic Systems

data = {
    "coord": coordinates,
    "numbers": atomic_numbers,
    "charge": 0.0,
    "cell": cell_vectors,  # 3x3 array in Angstrom
}
results = calc(data, forces=True, stress=True)

# Configure Coulomb method for periodic systems
calc.set_lrcoulomb_method("dsf", cutoff=15.0, dsf_alpha=0.2)
# or Ewald summation with the default nvalchemiops accuracy
calc.set_lrcoulomb_method("ewald", ewald_accuracy=1e-6)

Performance: torch.compile

For molecular dynamics, compile_model=True gives ~5x speedup (requires CUDA):

calc = AIMNet2Calculator("aimnet2", compile_model=True)

Output Reference

Loading from Hugging Face

AIMNet2 models are available on Hugging Face. Install the optional HF extras:

pip install "aimnet[hf]"

from aimnet.calculators import AIMNet2Calculator

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

# All available HF repos:
# isayevlab/aimnet2-wb97m-d3   general purpose (wB97M-D3)
# isayevlab/aimnet2-2025       improved intermolecular (B97-3c)
# isayevlab/aimnet2-nse        open-shell / radicals
# isayevlab/aimnet2-pd         palladium chemistry
# isayevlab/aimnet2-rxn        reactive chemistry / TS / IRC

# Load a specific ensemble member (0–3) or a pinned revision:
calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", ensemble_member=2)
calc = AIMNet2Calculator("isayevlab/aimnet2-wb97m-d3", revision="v1.0")

# Private repos — pass a HF token:
calc = AIMNet2Calculator("myorg/private-model", token="hf_...")

# Local directory in HF repo layout (config.json + ensemble_N.safetensors):
calc = AIMNet2Calculator("/path/to/local/repo")

# Existing registry aliases still work without any HF deps:
calc = AIMNet2Calculator("aimnet2")

Try the interactive demo!

How It Works

Architecture

AIMNet2 uses a message-passing neural network with iterative charge equilibration:

1. AEVSV - Gaussian basis expansion of pairwise distances and displacement vectors
2. ConvSV - Sparse indexed convolution combining atomic features with local geometry (GPU-accelerated via NVIDIA Warp kernels)
3. MLP passes - Iterative refinement with charge prediction and Coulomb-aware features
4. Output modules - Energy, forces (via autograd), charges, stress, Hessian

Dense vs Sparse Mode

The calculator automatically selects the optimal strategy:

• Dense mode (O(N^2)) - Small molecules on GPU. Fully connected graph, maximum parallelism.
• Sparse mode (O(N)) - Large systems or CPU. Adaptive neighbor lists with ~75% buffer utilization, 16-byte aligned allocations, automatic overflow handling.

The threshold is configurable via nb_threshold (default: 120 atoms).

Long-Range Corrections

• DFT-D3 dispersion with BJ damping (GPU-accelerated via nvalchemiops)
• Coulomb: Simple (all-pairs), DSF (damped-shifted force), or Ewald summation
• Long-range modules support inference forces and stress where documented; DSF and DFT-D3 use specialized backend paths, with Hessian support limited to the documented pure-torch DFT-D3 path

Training

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

See the training documentation for dataset preparation, configuration, and W&B integration.

Development

make check       # Linters and code quality (ruff, markdownlint, prettier)
make test        # Tests with coverage (pytest, parallel)
make docs        # Build and serve documentation (mkdocs)
make docs-test   # Validate docs build

Citation

If you use AIMNet2 in your research, please cite:

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-NSE: ChemRxiv preprint

AIMNet2-Pd: ChemRxiv preprint

License

See LICENSE file for details.