Skip to content

isayevlab/aimnetcentral

Repository files navigation

Release Python Build status codecov License

AIMNet2 : ML interatomic potential for fast and accurate atomistic simulations

Key Features

  • Accurate and Versatile: AIMNet2 excels at modeling neutral, charged, organic, and elemental-organic systems.
  • Flexible Interfaces: Use AIMNet2 through convenient calculators for popular simulation packages like ASE and PySisyphus.
  • Flexible Long-Range Interactions: Optionally employ the Damped-Shifted Force (DSF) or Ewald summation Coulomb models for accurate calculations in large or periodic systems.

Requirements

Python Version

AIMNet2 requires Python 3.11 or 3.12.

GPU Support (Optional)

AIMNet2 works on CPU out of the box. For GPU acceleration:

  • CUDA GPU: Install PyTorch with CUDA support from pytorch.org
  • compile_mode: Requires CUDA for ~5x MD speedup (see Performance Optimization)

Example PyTorch installation with CUDA 12.4:

pip install torch --index-url https://download.pytorch.org/whl/cu124

Available Models

Model Alias Elements Description
aimnet2_wb97m_d3_X aimnet2 H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I wB97M-D3 (default)
aimnet2_b973c_d3_X aimnet2_b973c H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I B97-3c functional
aimnet2nse_X aimnet2nse H, C, N, O, F, S, Cl Open-shell chemistry
aimnet2-pd_X aimnet2pd H, C, N, O, F, P, S, Cl, Pd Palladium-containing systems

X = 0-3 for ensemble members. Ensemble averaging recommended for production use.

Installation

Basic Installation

Install from GitHub:

pip install git+https://github.com/isayevlab/aimnetcentral.git

Optional Features

AIMNet2 provides optional extras for different use cases:

ASE Calculator (for atomistic simulations with ASE):

pip install "aimnet[ase] @ git+https://github.com/isayevlab/aimnetcentral.git"

PySisyphus Calculator (for reaction path calculations):

pip install "aimnet[pysis] @ git+https://github.com/isayevlab/aimnetcentral.git"

Training (for model training and development):

pip install "aimnet[train] @ git+https://github.com/isayevlab/aimnetcentral.git"

All Features:

pip install "aimnet[ase,pysis,train] @ git+https://github.com/isayevlab/aimnetcentral.git"

Development Installation

For contributors, use uv for fast dependency management:

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

Quick Start

Basic Usage (Core)

from aimnet.calculators import AIMNet2Calculator # Load a pre-trained model calc = AIMNet2Calculator("aimnet2") # Prepare input data = { "coord": coordinates, # Nx3 array "numbers": atomic_numbers, # N array "charge": 0.0, } # Run inference results = calc(data, forces=True) print(results["energy"], results["forces"])

Output Data

The calculator returns a dictionary with the following keys:

Key Shape Description
energy (,) or (B,) Total energy in eV
charges (N,) or (B, N) Atomic partial charges in e
forces (N, 3) or (B, N, 3) Atomic forces in eV/A (if requested)
hessian (N, 3, N, 3) Second derivatives (if requested)
stress (3, 3) Stress tensor for PBC (if requested)

B = batch size, N = number of atoms

ASE Integration

With aimnet[ase] installed:

from ase.io import read from aimnet.calculators import AIMNet2ASE atoms = read("molecule.xyz") atoms.calc = AIMNet2ASE("aimnet2") energy = atoms.get_potential_energy() forces = atoms.get_forces()

Periodic Boundary Conditions

For periodic systems, provide a unit cell:

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

Long-Range Coulomb Methods

Configure electrostatic interactions for large or periodic systems:

# Damped-Shifted Force (DSF) - recommended for periodic systems calc.set_lrcoulomb_method("dsf", cutoff=15.0, dsf_alpha=0.2) # Ewald summation - for accurate periodic electrostatics calc.set_lrcoulomb_method("ewald", cutoff=15.0)

Performance Optimization

For molecular dynamics simulations, use compile_mode for ~5x speedup:

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

Requirements:

  • CUDA GPU required
  • Not compatible with periodic boundary conditions
  • Best for repeated inference on similar-sized systems

Training

With aimnet[train] installed:

aimnet train --config my_config.yaml --model aimnet2.yaml

Development

Common development tasks using make:

make check # Run linters and code quality checks make test # Run tests with coverage make docs # Build and serve documentation make build # Build distribution packages

Citation

If you use AIMNet2 in your research, please cite the appropriate paper:

AIMNet2 (main model):

@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.

About

AIMNet2: Fast and accurate machine-learned interatomic potential for molecular dynamics simulations

Resources

License

Contributing

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 5