Neural networks that feel atomic forces, now with a v2 rewrite

What it does

TorchANI trains and runs “ANI-style” neural network interatomic potentials — essentially machine-learned approximations of quantum mechanical forces between atoms. You feed it molecular geometries; it predicts energies and forces fast enough to drive molecular dynamics simulations. The project is maintained by the Roitberg group and has been around long enough to need a migration guide for its own 2.0 rewrite.

The interesting bit

The library ships with custom C++ and CUDA extensions for descriptor computation and network inference, which you build post-install with an ani build-extensions command. It also exposes a command-line interface (ani --help) and plugs into Amber for mixed quantum/classical (ML/MM) simulations via a separate interface project.

Key highlights

• Requires PyTorch ≥ 2.0; tested against PyTorch 2.8 and CUDA 12.8
• Custom CUDA/C++ extensions for GPU acceleration; CPU-only runs are explicitly “degraded”
• CLI utility ani included for common operations
• Migration guide provided for 1.x users; legacy state dicts accessible via .legacy_state_dict()
• Conda package exists but is unmaintained; pip is the recommended install path even inside conda envs

Caveats

• Untested on AMD GPUs (ROCm/HIP) and on Apple Metal Performance Shaders
• macOS installation requires manual tweaks to environment.yaml; no CUDA support
• Conda packaging is effectively deprecated and CI-gated behind specific branch name strings

Verdict

Computational chemists and molecular dynamics hackers already using PyTorch should look here; if you’re not simulating atoms, this is specialized tooling with a nontrivial install path.