---
library_name: aimnet
license: mit
tags:
  - chemistry
  - computational-chemistry
  - molecular-dynamics
  - interatomic-potential
  - neural-network-potential
  - pytorch
  - aimnet2
pipeline_tag: other
---

# 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

```bash
pip install "aimnet[hf]"
```

## Quick Start

```python
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

```python
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.

| File | Description |
|------|-------------|
| `ensemble_0.safetensors` | Ensemble member 0 (default) |
| `ensemble_1.safetensors` | Ensemble member 1 |
| `ensemble_2.safetensors` | Ensemble member 2 |
| `ensemble_3.safetensors` | Ensemble member 3 |
| `config.json` | Shared model configuration and metadata |

## 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

```bibtex
@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