Note
Go to the end to download the full example code
Computing Energy and Force Using Models Inside Model Zoo¶
TorchANI has a model zoo trained by NeuroChem. These models are shipped with TorchANI and can be used directly.
To begin with, let’s first import the modules we will use:
import torch
import torchani
Let’s now manually specify the device we want TorchANI to run:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Let’s now load the built-in ANI-1ccx models. The builtin ANI-1ccx contains 8 models trained with diffrent initialization. Predicting the energy and force using the average of the 8 models outperform using a single model, so it is always recommended to use an ensemble, unless the speed of computation is an issue in your application.
The periodic_table_index
arguments tells TorchANI to use element index
in periodic table to index species. If not specified, you need to use
0, 1, 2, 3, … to index species
model = torchani.models.ANI2x(periodic_table_index=True).to(device)
/opt/hostedtoolcache/Python/3.8.17/x64/lib/python3.8/site-packages/torchani-2.2.3-py3.8.egg/torchani/resources/
Now let’s define the coordinate and species. If you just want to compute the
energy and force for a single structure like in this example, you need to
make the coordinate tensor has shape (1, Na, 3)
and species has shape
(1, Na)
, where Na
is the number of atoms in the molecule, the
preceding 1
in the shape is here to support batch processing like in
training. If you have N
different structures to compute, then make it
N
.
Note
The coordinates are in Angstrom, and the energies you get are in Hartree
coordinates = torch.tensor([[[0.03192167, 0.00638559, 0.01301679],
[-0.83140486, 0.39370209, -0.26395324],
[-0.66518241, -0.84461308, 0.20759389],
[0.45554739, 0.54289633, 0.81170881],
[0.66091919, -0.16799635, -0.91037834]]],
requires_grad=True, device=device)
# In periodic table, C = 6 and H = 1
species = torch.tensor([[6, 1, 1, 1, 1]], device=device)
Now let’s compute energy and force:
energy = model((species, coordinates)).energies
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
And print to see the result:
Energy: -40.45979070350399
Force: tensor([[ 0.048, -0.130, -0.055],
[-0.135, 0.158, -0.078],
[ 0.080, -0.039, 0.039],
[ 0.025, 0.008, 0.043],
[-0.018, 0.004, 0.051]])
you can also get the atomic energies (WARNING: these have no physical meaning) by calling:
_, atomic_energies = model.atomic_energies((species, coordinates))
this gives you the average (shifted) energies over all models of the ensemble by default, with the same shape as the coordinates. Dummy atoms, if present, will have an energy of zero
print('Average Atomic energies, for species 6 1 1 1 1', atomic_energies)
Average Atomic energies, for species 6 1 1 1 1 tensor([[-38.080, -0.581, -0.591, -0.605, -0.604]],
grad_fn=<MeanBackward1>)
you can also access model specific atomic energies
_, atomic_energies = model.atomic_energies((species, coordinates), average=False)
print('Atomic energies of first model, for species 6 1 1 1 1', atomic_energies[0, :, :])
Atomic energies of first model, for species 6 1 1 1 1 tensor([[-38.084, -0.580, -0.590, -0.603, -0.603]],
grad_fn=<SliceBackward0>)
Total running time of the script: (0 minutes 5.452 seconds)