-
Notifications
You must be signed in to change notification settings - Fork 2
/
ani_grad-2x.py
50 lines (31 loc) · 1.05 KB
/
ani_grad-2x.py
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
# calculate energy and gradient with ANI
import torch
import torchani
import math
import numpy as np
import sys
device = torch.device('cpu')
inp_f = "mol.tmp"
with open(inp_f,"r") as f:
natom = int(f.readline())
l1 = []
l3 = []
for i in range(natom):
l0 = f.readline().split()
l2 = l0[1:4]
x = float(l2[0]) * 0.529177
y = float(l2[1]) * 0.529177
z = float(l2[2]) * 0.529177
l1.append([x,y,z])
l3.append(int(l0[0]))
model = torchani.models.ANI2x(periodic_table_index=True).to(device).double()
species = torch.tensor(np.array(l3), device=device, dtype=torch.long).unsqueeze(0)
coordinates = torch.from_numpy(np.array(l1)).unsqueeze(0).requires_grad_(True)
masses = torchani.utils.get_atomic_masses(species)
energies = model((species, coordinates)).energies
E1 = energies.item()
print("%20.12F%20.12F%20.12F%20.12F"% (E1,0,0,0))
derivative = torch.autograd.grad(energies.sum(), coordinates)[0]
G1 = derivative.numpy()[0] * 0.529177
for i in G1:
print("%20.12F%20.12F%20.12F"% tuple(i))