Torch interface

PyTorch is a very general deep-learning framework within the Python ecosystem and one of the most popular approches to training neural networks. TorchANI is a PyTorch-based implementation of the older ANI neural network potential for describing potential energy surfaces of molecular systems (and other properties). AIMNet2 is newer neural network potential, also implemented using PyTorch, capable of unprecedented accuracy for up to 14 chemical elements, support for charged systems and is capable of describing long-range interactions and dispersion. AIMNet2 paper: https://www.sciencedirect.com/org/science/article/pii/S2041652025006844

ASH features an interface allowing use of either TorchANI or AIMNet2 neural network models. The interface allows easy use of pre-trained neural network potentials. This allows the use of such neural network potentials as any other theory-level within ASH.

Energies and gradients can be requested (just like a regular QM or MM theory) and so a valid TorchTheory object can be used for single-point energies, geometry optimizations, numerical frequencies, surface scans, NEB, molecular dynamics etc. within ASH. Even hybrid ONIOM and QM/MM calculations are possible (with some limitations).

Periodic boundary conditions are also available but only for AIMNet2 potentials.

TorchTheory class:

class TorchTheory():
    def __init__(self, filename="torch.pt", model_name=None, model_object=None,
                model_file=None, printlevel=2, label="TorchTheory", numcores=1,
                platform=None, aimnet_mode="new",
                periodic=False, periodic_cell_vectors=None, periodic_cell_dimensions=None):

Keyword

Type

Default value

Details

filename

string

'torch.pt'

Filename used for PyTorch models when saving.

model_name

string

None

Name of pretrained model to use. Options: 'AIMNet2','ANI1x', 'ANI2x', 'ANI1ccx' (requires TorchANI or AIMNet2)

model_object

PyTorch NN model object

None

Read in a PyTorch model object directly.

model_file

Name of Pytorch model-file to read in

None

Read in a PyTorch model from a file.

platform

string

None

Type of platform to use when training/running. Options: 'cpu', 'cuda', 'mps' (Apple Silicon)

printlevel

integer

2

Printlevel

label

string

'TorchTheory'

Label of TorchTheory object

numcores

integer

1

Number of cores.

periodic

Boolean

False

Whether to use PBCs or not

periodic_cell_vectors

numpy array

None

Cell vectors as 3x3 numpy array in Angstrom.

periodic_cell_vectors

list

None

Cell dimensions as list of cell lengths and angles in Angstrom and degrees.

aimnet_mode

string

"new"

Whether to use new or old aimnet interface (newer recommended, old will be removed soon)

Torch/TorchANI/AIMNet2 installation

PyTorch

PyTorch needs to be installed on your system to use TorchTheory. The easiest way to install PyTorch is via pip:

pip install torch

TorchANI

To use pre-trained ANI neural network potentials you need to install TorchANI. See TorchANI documentation and TorchANI repository Torchani can be installed via pip:

pip install torchani

AIMNet2

To use AIMNet2 follow the installation instructions at https://isayevlab.github.io/aimnetcentral/#installation

pip install aimnet

AIMNet2 Examples

Basic AIMNet2 example

It is easy to use the AIMNet2 neural network potential with TorchTheory. The available models are: 'aimnet2', 'aimnet2_2025', 'aimnet2_b973c', 'aimnet2nse', 'aimnet2pd' and models are available for elements: 'H', 'C', 'N', 'O', 'F', 'Cl', 'S', 'Si', 'B', 'P', 'Br', 'As', 'I', 'Se'

from ash import *

#H2O fragment
frag = Fragment(databasefile="h2o.xyz", charge=0, mult=1)

# Create a TorchTheory object using the aimnet2_2025 neural network potential
theory = TorchTheory(model_name="aimnet2_2025", platform="cpu")

#Run a single-point energy+gradient calculation
#Optimizer,NumFreq, MolecularDynamics etc. should also work
result = Singlepoint(theory=theory, fragment=frag, Grad=True)

print(result.energy)
print(result.gradient)

AIMNet2 with PBCs

The AIMNet2 potentials can be used for periodic systems as well.

from ash import *

numcores=1
frag = Fragment(xyzfile="ammonia.xyz", charge=0, mult=1)

#Cell
cell_vectors = np.array([[5.01336,0.0,0.0],[0.0,5.01336,0.0],[0.0,0.0,5.01336]])

#Periodic AIMNet2 Torchtheory
theory = TorchTheory(model_name="aimnet2", numcores=numcores,
                periodic=True, periodic_cell_vectors=cell_vectors)

Optimizer(theory=theory, fragment=frag, coordsystem="hdlc")

TorchANI Examples

Basic TorchANI example

A pretrained ANI-based model using the TorchANI library can easily be used as well. The available models are: 'ANI1x', 'ANI1ccx', 'ANI2x' and they are available for elements: 'H', 'C', 'N', 'O'

from ash import *

#H2O fragment
frag = Fragment(databasefile="h2o.xyz", charge=0, mult=1)
# Create a TorchTheory object using the ANI2x neural network potential
theory = TorchTheory(model_name="ANI2x", platform="cpu")

#Run a single-point energy+gradient calculation
#Optimizer,NumFreq, MolecularDynamics etc. should also work
result = Singlepoint(theory=theory, fragment=frag, Grad=True)

print(result.energy)
print(result.gradient)

Loading a pretrained model from file

It's also possible to load a neural-network potential from file. Here we show an example of this by loading the ANI1x model from a file (Pytorch .pt format). This file was generated by first creating a TorchTheory object like above and then calling the save_model method.

from ash import *

#H2O fragment
frag = Fragment(databasefile="h2o.xyz", charge=0, mult=1)
# Create a TorchTheory object using the ANI1x neural network potential from a saved-file
theory = TorchTheory(model_file="savedANI1x.pt")
#Run a single-point energy+gradient calculation
result = Singlepoint(theory=theory, fragment=frag, Grad=True)

Defining a new PyTorch model from scratch

Not yet ready