energy_force.py 2.45 KB
Newer Older
Gao, Xiang's avatar
Gao, Xiang committed
1
2
# -*- coding: utf-8 -*-
"""
Gao, Xiang's avatar
Gao, Xiang committed
3
4
Computing Energy and Force Using Models Inside Model Zoo
========================================================
Gao, Xiang's avatar
Gao, Xiang committed
5

Gao, Xiang's avatar
Gao, Xiang committed
6
7
TorchANI has a model zoo trained by NeuroChem. These models are shipped with
TorchANI and can be used directly.
Gao, Xiang's avatar
Gao, Xiang committed
8
9
10
11
"""

###############################################################################
# To begin with, let's first import the modules we will use:
Xiang Gao's avatar
Xiang Gao committed
12
13
14
import torch
import torchani

Gao, Xiang's avatar
Gao, Xiang committed
15
16
###############################################################################
# Let's now manually specify the device we want TorchANI to run:
Xiang Gao's avatar
Xiang Gao committed
17
18
device = torch.device('cpu')

Gao, Xiang's avatar
Gao, Xiang committed
19
###############################################################################
Gao, Xiang's avatar
Gao, Xiang committed
20
21
22
23
24
25
# 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.
model = torchani.models.ANI1ccx()
Gao, Xiang's avatar
Gao, Xiang committed
26
27
28
29
30
31
32
33
34

###############################################################################
# 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``.
Xiang Gao's avatar
Xiang Gao committed
35
36
37
38
39
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]]],
Gao, Xiang's avatar
Gao, Xiang committed
40
                           requires_grad=True, device=device)
Gao, Xiang's avatar
Gao, Xiang committed
41
species = model.species_to_tensor('CHHHH').to(device).unsqueeze(0)
Xiang Gao's avatar
Xiang Gao committed
42

Gao, Xiang's avatar
Gao, Xiang committed
43
44
###############################################################################
# Now let's compute energy and force:
45
_, energy = model((species, coordinates))
Gao, Xiang's avatar
Gao, Xiang committed
46
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
Xiang Gao's avatar
Xiang Gao committed
47
48
force = -derivative

Gao, Xiang's avatar
Gao, Xiang committed
49
50
###############################################################################
# And print to see the result:
Xiang Gao's avatar
Xiang Gao committed
51
print('Energy:', energy.item())
Gao, Xiang's avatar
Gao, Xiang committed
52
print('Force:', force.squeeze())