aimnet 0.1.0

AIMNet Machine Learned Interatomic Potential

• Github repository: https://github.com/isayevlab/aimnetcentral/
• Documentation https://isayevlab.github.io/aimnetcentral/

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	Elements	Description
aimnet2	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	wB97M-D3 (default)
aimnet2_b973c	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	B97-3c functional
aimnet2_2025	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	B97-3c + improved intermolecular interactions
aimnet2nse	H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I	Open-shell chemistry
aimnet2pd	H, B, C, N, O, F, Si, P, S, Cl, Se, Br, Pd, I	Palladium-containing systems

Each model has ensemble members (append _0 to _3). Ensemble averaging recommended for production use.

Installation

Basic Installation

Install from PyPI:

pip install aimnet

Optional Features

AIMNet2 provides optional extras for different use cases:

ASE Calculator (for atomistic simulations with ASE):

pip install "aimnet[ase]"

PySisyphus Calculator (for reaction path calculations):

pip install "aimnet[pysis]"

Training (for model training and development):

pip install "aimnet[train]"

All Features:

pip install "aimnet[ase,pysis,train]"

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", ewald_accuracy=1e-8)

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

Technical Details

Batching and Neighbor Lists

The AIMNet2Calculator automatically selects the optimal strategy based on system size (nb_threshold, default 120 atoms) and hardware:

1. Dense Mode (O(N²)): Used for small molecules on GPU. Input is kept in 3D batched format (B, N, 3). No neighbor list is computed; the model uses a fully connected graph for maximum parallelism.
2. Sparse Mode (O(N)): Used for large systems or CPU execution. Input is flattened to 2D (N_total, 3) with an adaptive neighbor list. This ensures linear memory scaling.

Adaptive Neighbor List

In sparse mode, AIMNet2 uses an AdaptiveNeighborList that automatically resizes its buffer to maintain efficient utilization (~75%) while preventing overflows.

• Format: The neighbor list is stored as a 2D integer matrix nbmat of shape (N_total, max_neighbors). Each row i contains the indices of atoms neighboring atom i.
• Padding: Rows with fewer neighbors than max_neighbors are padded with the index N_total (a dummy atom index).
• Buffer Management: The buffer size max_neighbors is always a multiple of 16 for memory alignment. It dynamically expands (by 1.5x) on overflow and shrinks if utilization drops significantly below the target, ensuring robust performance during MD simulations where density fluctuates.

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.