TorchANI

Accurate Neural Network Potentials in PyTorch

Status: Active Development

TorchANI is a PyTorch implementation of ANI (Accurate NeurAl networK engINe for Molecular Energies) neural network potentials. It provides fast, GPU-accelerated molecular energy and force calculations with near-quantum accuracy.

Features

  • High Performance: GPU acceleration with PyTorch for 100,000× speedup over DFT
  • Pre-trained Models: ANI-1x, ANI-1ccx, and ANI-2x potentials ready to use
  • Easy Integration: Compatible with ASE, RDKit, and major MD engines
  • Extensible: Train custom potentials on your own data
  • Open Source: MIT licensed, actively maintained

Installation

Via pip

pip install torchani

Via conda

conda install -c conda-forge torchani

Quick Start

import torch
import torchani

# Load pre-trained ANI-2x model
model = torchani.models.ANI2x(periodic_table_index=True)

# Define a molecule (coordinates in Angstroms)
coordinates = torch.tensor([[[0.0, 0.0, 0.0],
                             [0.0, 0.0, 1.1]]])
species = torch.tensor([[1, 1]])  # H2

# Calculate energy and forces
energy = model((species, coordinates)).energies
forces = -torch.autograd.grad(energy.sum(), coordinates)[0]

print(f"Energy: {energy.item()} Hartree")
print(f"Forces:\n{forces}")

Applications

  • Molecular dynamics simulations
  • Geometry optimization
  • Transition state searches
  • Conformational sampling
  • Drug binding calculations

Performance

  • Accuracy: Within 1 kcal/mol of CCSD(T) for organic molecules
  • Speed: ~10 ms per molecule on GPU
  • Coverage: CHNO atoms (ANI-1x/1ccx), CHNOSFCl atoms (ANI-2x)

Citation

If you use TorchANI in your research, please cite:

@article{gao2020torchani,
  title={TorchANI: A Free and Open Source PyTorch-Based Deep Learning Implementation of the ANI Neural Network Potentials},
  author={Gao, Xiang and Ramezanghorbani, Farhad and Isayev, Olexandr and Smith, Justin S and Roitberg, Adrian E},
  journal={Journal of Chemical Information and Modeling},
  year={2020}
}

Installation

pip install torchani
conda install -c conda-forge torchani