Torch/TorchANI interface

PyTorch is a very general and one of the most popular approches to training neural networks in Python. TorchANI is a PyTorch-based implementation of the ANI neural network potential for describing potential energy surfaces of molecular systems (and other properties).

ASH features a basic interface to PyTorch and TorchANI that allows easy use of pre-trained neural-network based on machine-learning models. This allows the use of such neural network potentials as theory-levels 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.

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, train=False):

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: 'ANI1x', 'ANI2x', 'ANI1ccx' (requires TorchANI)

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.

train

Boolean

False

Whether to create a TorchTheory object that will be used for training.

Torch/TorchANI installation

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

pip install torch

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

Examples

Basic TorchANI example

The easiest way to use a PyTorch-based neural network within ASH is to use a pretrained ANI-based model using the TorchANI library. 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