README.md

---
library_name: aimnet
license: mit
tags:
  - chemistry
  - computational-chemistry
  - reactions
  - transition-states
  - reaction-pathways
  - neb
  - interatomic-potential
  - neural-network-potential
  - pytorch
  - aimnet2
pipeline_tag: other
---

# AIMNet2-rxn

AIMNet2-rxn is a neural network interatomic potential specialized for **generalized organic reaction modeling** — energies, forces, charges, and electrostatic moments along reaction coordinates at near-ωB97M-V/def2-TZVPP accuracy. Trained on ~4.7M reaction-relevant geometries to handle bond breaking/forming, transition states, and minimum-energy pathways out of the box, without system-specific fine-tuning.

## Highlights

- **Reaction-aware:** ~1 kcal mol⁻¹ across IRC profiles, ~1.6 kcal mol⁻¹ RMSD on reaction energies, ~2 kcal mol⁻¹ RMSD on diverse-reaction barrier heights, ~3 kcal mol⁻¹ RMSD on the Grambow benchmark — all without retraining or fine-tuning. (Numbers from the paper; see below.)
- **High-throughput pathway search:** Enables minimum-energy-pathway searches at ~10⁶ reactions/day via batched NEB on a single GPU — roughly 10⁶× the throughput of serial NEB-DFT at the reference level. Per-call speedup is closer to 10³–10⁴×; the 10⁶× figure is end-to-end batched NEB throughput, not per-energy-evaluation.
- **TS- and NEB-ready:** Designed for transition-state optimization and (batched) nudged-elastic-band minimum-energy-pathway searches.
- **Ensemble:** 4 ensemble members for uncertainty estimation.
- **Outputs:** energy, forces (autograd), atomic charges, molecular dipole, molecular quadrupole.

## Element coverage

**H, C, N, O only.** The model was trained on closed-shell organics. Unsupported elements (F, P, S, …) are not validated by the calculator and will silently produce undefined predictions — read `implemented_species` from this repo's `config.json` (`[1, 6, 7, 8]`) and check input atomic numbers yourself before running on novel systems. For broader element coverage on equilibrium structures, use [`isayevlab/aimnet2-wb97m-d3`](https://huggingface.co/isayevlab/aimnet2-wb97m-d3) or related families.

## Energy convention

AIMNet2-rxn outputs are on a **learned per-element shifted electronic-energy scale** (eV). They are designed to be accurate for energy *differences* between geometries that share atomic composition — barrier heights, reaction energies, IRC profiles, conformer energies. They are **not** absolute ωB97M-V/def2-TZVPP electronic energies and **not** formal atomization energies.

Practical implications:
- Methane sits near a few eV on this scale, **not** the −1100 eV scale used by `aimnet2-wb97m-d3` and the other AIMNet2 families. Do **not** mix or compare energies across families.
- Single-point energies in isolation are not physically meaningful. Always interpret as a difference vs. another geometry from this same model.

## 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-rxn")

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

### With ASE — minimum energy path via NEB

```python
from aimnet.calculators.aimnet2ase import AIMNet2ASE
from ase.mep import NEB
from ase.optimize import LBFGS

# images: list[ase.Atoms] interpolated between reactant and product
calc_factory = lambda: AIMNet2ASE("isayevlab/aimnet2-rxn")
for img in images:
    img.calc = calc_factory()

neb = NEB(images, climb=True)
LBFGS(neb).run(fmax=0.05)
```

For batched-NEB (BNEB) on millions of pathways, see the AIMNet2-rxn paper for the recommended workflow with `pysisyphus`.

### Transition-state and Hessian workflows

For TS optimization or any workflow that requires a Hessian (e.g. P-RFO, dimer, IRC, vibrational analysis), **set `compile_model=False`** when constructing the calculator:

```python
calc = AIMNet2Calculator("isayevlab/aimnet2-rxn", compile_model=False)
```

`torch.compile` (Dynamo) is known to hang on the double-backward path through GELU activations that Hessian autograd requires. The default avoids the issue, but explicit `compile_model=True` for speed should be reserved for pure energy/force evaluations.

## 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) with `dipole` and `quadrupole` heads
- **Cutoff radius:** 5.0 Å
- **DFT reference:** **ωB97M-V / def2-TZVPP** (range-separated hybrid meta-GGA with VV10 nonlocal correlation; Mardirossian & Head-Gordon, *J. Chem. Phys.* 144, 214110, 2016)
- **Dispersion:** Supplied by ωB97M-V's VV10 nonlocal correlation kernel, included in every training-set energy and force. No separate D3/D4 correction is added at inference (`needs_dispersion: false` in `config.json`); doing so would double-count the dispersion already baked into the reference data. AIMNet2's external D3 path remains available for other family members whose reference functional does not carry its own dispersion (e.g. `aimnet2-wb97m-d3`).
- **Coulomb:** Short-range embedded (≤4.6 Å) + external long-range (≥4.6 Å, applied by the calculator via `LRCoulomb`). The 4.6 Å crossover is **physically frozen** — it is the SR/LR cancellation point used during training. Do not alter `coulomb_sr_rc` at inference.
- **Training set:** 4,685,165 geometries combining the AIMNet2 base set with the RGD-1 reaction database, all recomputed at ωB97M-V/def2-TZVPP
- **Format:** safetensors (converted from JIT TorchScript via reconstructed YAML)

## Intended use

Mechanism enumeration, transition-state searches, NEB / BNEB, reaction-coordinate energy profiles, high-throughput reactivity screening on closed-shell organic systems.

## Calculator-enforced safeguards (with `aimnet>=NEXT_VERSION`)

The calculator validates inputs and refuses to silently produce undefined output:

- **Element scope** — passing atomic numbers outside `[1, 6, 7, 8]` raises `ValueError` (bypass with `validate_species=False`).
- **Net charge** — `charge != 0` raises `ValueError` (zwitterions are OK as long as the net charge is zero; for net-charged species use [`isayevlab/aimnet2-wb97m-d3`](https://huggingface.co/isayevlab/aimnet2-wb97m-d3)).
- **Hessian + `torch.compile`** — combining both raises `RuntimeError` (Dynamo and GELU double-backward hang); reconstruct the calculator with `compile_model=False` for TS / IRC / vibrational analysis.
- **Coulomb cutoff** — `set_lrcoulomb_method(...)` warns if you change the SR/LR crossover from 4.6 Å (physically frozen during training).
- **Cross-family mixing** — constructing calculators from two AIMNet2 families in one process emits a one-time `UserWarning` about energy-scale incompatibility.

`afv.weight` rows for elements outside `[1, 6, 7, 8]` are NaN-padded at the safetensors level, so even with `validate_species=False` the model produces NaN (not plausible-looking nonsense) for unsupported elements. Two layers of defense.

## Limitations

- **Element scope:** H, C, N, O only — passing other elements is not validated by the calculator and yields undefined output.
- **Closed-shell only:** Open-shell species (radicals, doublets) are not handled by AIMNet2-rxn. For open-shell reactions and radical chemistry, use [`isayevlab/aimnet2-nse`](https://huggingface.co/isayevlab/aimnet2-nse) — a fully reactive model with parametrization focused on open-shell systems and radicals.
- **Net-charged species (anions, cations):** Out of scope for this release; future work per the paper. **Net-neutral zwitterions are within the training distribution and supported** — local charge separation in a globally neutral molecule (e.g. amino acid neutral forms, proton-transfer intermediates) is a deliberate part of training.
- **Transition metals & heteroatoms beyond H/C/N/O:** Not supported.
- **Equilibrium MD on large organics:** While accurate near equilibrium, the wb97m-d3 family is generally preferred when reaction information is not needed.
- **Hessian/TS workflows under `torch.compile`:** Set `compile_model=False` (see Quick Start). `torch.compile` and Hessian autograd through GELU activations are known to interact poorly.

## Citation

```bibtex
@misc{anstine2025aimnet2rxn,
  title  = {AIMNet2-rxn: A Machine Learned Potential for Generalized Reaction Modeling on a Millions-of-Pathways Scale},
  author = {Anstine, Dylan M. and Zhao, Qiyuan and Zubatiuk, Roman and Zhang, Shuhao and Singla, Veerupaksh and Nikitin, Filipp and Savoie, Brett M. and Isayev, Olexandr},
  year   = {2025},
  doi    = {10.26434/chemrxiv-2025-hpdmg},
  note   = {ChemRxiv preprint}
}
```

## License

MIT License