aimnet2-nse

AIMNet2 NSE (Open-Shell)

AIMNet2 is a neural network interatomic potential for fast and accurate molecular simulations. This is a spin-polarized model for open-shell chemistry (radicals, triplet states, excited states). Accepts spin multiplicity as additional input.

Highlights

• Fast: Orders of magnitude faster than DFT, with optional torch.compile for ~5x additional speedup on GPU
• Accurate: Near-DFT accuracy for energies, forces, and charges
• Broad coverage: Supports 14 elements: H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
• Ensemble: 4 ensemble members for uncertainty estimation
• Features: Energy, forces, charges, Hessian, stress tensor, periodic boundary conditions

Installation

pip install "aimnet[hf]"

Quick Start

from aimnet.calculators import AIMNet2Calculator

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

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

# For open-shell species (e.g., methyl radical, doublet)
results = calc(
    {"coord": coords, "numbers": atomic_numbers, "charge": 0.0, "mult": 2.0},
    forces=True,
)

With ASE

from aimnet.calculators.aimnet2ase import AIMNet2ASE
from ase.build import molecule

atoms = molecule("H2O")
atoms.calc = AIMNet2ASE("isayevlab/aimnet2-nse", mult=2)

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

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.

Model Details

• Architecture: AIMNet2 (Atomic Environment Vectors + message-passing MLPs)
• Cutoff radius: 5.0 A
• DFT functional: wB97M
• Dispersion: D3BJ (externalized, handled by calculator)
• Coulomb: Short-range embedded + external long-range (DSF/Ewald)
• Format: safetensors (converted from PyTorch state dict)

Limitations

• Periodic systems are supported through the Python API / ASE interface, not through the Gradio demo
• Element coverage is limited to the 14 elements listed above; unsupported elements will raise an error

Citation

@article{anstine2025aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  year={2025},
  publisher={Royal Society of Chemistry},
  doi={10.1039/D4SC08572H}
}

License

MIT License