← all repositories
aiqm/torchani

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

A PyTorch library for training ANI-style neural network potentials, recently overhauled with breaking changes and C++/CUDA extensions.

545 stars Python Domain AppsML Frameworks
torchani
Velocity · 7d
+0.2
★ / day
Trend
steady
star history

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.

heatdrop uses Google Analytics to see which pages get read — nothing else. Your call. How we handle data.