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