Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
load_from_neurochem.py 4.36 KB
Newer Older
Gao, Xiang's avatar
Add ANI model zoo (#169)  
Gao, Xiang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50 
# -*- coding: utf-8 -*-
"""
Construct Model From NeuroChem Files
====================================

This tutorial illustrates how to manually load model from `NeuroChem files`_.

.. _NeuroChem files:
    https://github.com/isayev/ASE_ANI/tree/master/ani_models

"""

###############################################################################
# To begin with, let's first import the modules we will use:
import os
import torch
import torchani
import ase


###############################################################################
# Now let's read constants from constant file and construct AEV computer.
try:
    path = os.path.dirname(os.path.realpath(__file__))
except NameError:
    path = os.getcwd()
const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
consts = torchani.neurochem.Constants(const_file)
aev_computer = torchani.AEVComputer(**consts)

###############################################################################
# Now let's read self energies and construct energy shifter.
sae_file = os.path.join(path, '../torchani/resources/ani-1x_8x/sae_linfit.dat')  # noqa: E501
energy_shifter = torchani.neurochem.load_sae(sae_file)

###############################################################################
# Now let's read a whole ensemble of models.
model_prefix = os.path.join(path, '../torchani/resources/ani-1x_8x/train')  # noqa: E501
ensemble = torchani.neurochem.load_model_ensemble(consts.species, model_prefix, 8)  # noqa: E501

###############################################################################
# Or alternatively a single model.
model_dir = os.path.join(path, '../torchani/resources/ani-1x_8x/train0/networks')  # noqa: E501
model = torchani.neurochem.load_model(consts.species, model_dir)

###############################################################################
# You can create the pipeline of computing energies:
# (Coordinates) -[AEVComputer]-> (AEV) -[Neural Network]->
# (Raw energies) -[EnergyShifter]-> (Final energies)
# From using either the ensemble or a single model:
Farhad Ramezanghorbani's avatar
Use a modified Sequential class with type annotation (#308)  
Farhad Ramezanghorbani committed
51
52 
nnp1 = torchani.nn.Sequential(aev_computer, ensemble, energy_shifter)
nnp2 = torchani.nn.Sequential(aev_computer, model, energy_shifter)
Gao, Xiang's avatar
Add ANI model zoo (#169)  
Gao, Xiang committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66 
print(nnp1)
print(nnp2)

###############################################################################
# You can also create an ASE calculator using the ensemble or single model:
calculator1 = torchani.ase.Calculator(consts.species, aev_computer,
                                      ensemble, energy_shifter)
calculator2 = torchani.ase.Calculator(consts.species, aev_computer,
                                      model, energy_shifter)
print(calculator1)
print(calculator1)

###############################################################################
# Now let's define a methane molecule
Gao, Xiang's avatar
Config flake8 and fix may linter errors (#198)  
Gao, Xiang committed
67
68
69
70
71 
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
Add ANI model zoo (#169)  
Gao, Xiang committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100 
                           requires_grad=True)
species = consts.species_to_tensor('CHHHH').unsqueeze(0)
methane = ase.Atoms('CHHHH', positions=coordinates.squeeze().detach().numpy())

###############################################################################
# Now let's compute energies using the ensemble directly:
_, energy = nnp1((species, coordinates))
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
print('Energy:', energy.item())
print('Force:', force.squeeze())

###############################################################################
# And using the ASE interface of the ensemble:
methane.set_calculator(calculator1)
print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
print('Force:', methane.get_forces() / ase.units.Hartree)

###############################################################################
# We can do the same thing with the single model:
_, energy = nnp2((species, coordinates))
derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
force = -derivative
print('Energy:', energy.item())
print('Force:', force.squeeze())

methane.set_calculator(calculator2)
print('Energy:', methane.get_potential_energy() / ase.units.Hartree)
print('Force:', methane.get_forces() / ase.units.Hartree)