test_forces.py 2.42 KB
Newer Older
1
2
3
4
5
import torch
import torchani
import unittest
import os
import pickle
Jinze Xue's avatar
Jinze Xue committed
6
from torchani.testing import TestCase
7
8
9
10
11

path = os.path.dirname(os.path.realpath(__file__))
N = 97


Jinze Xue's avatar
Jinze Xue committed
12
class TestForce(TestCase):
13

14
    def setUp(self):
15
16
17
        model = torchani.models.ANI1x(model_index=0)
        self.aev_computer = model.aev_computer
        self.nnp = model.neural_networks
18
        self.model = torchani.nn.Sequential(self.aev_computer, self.nnp)
19

20
21
    def testIsomers(self):
        for i in range(N):
22
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
23
24
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, _, forces = pickle.load(f)
25
26
27
                coordinates = torch.from_numpy(coordinates)
                species = torch.from_numpy(species)
                forces = torch.from_numpy(forces)
Gao, Xiang's avatar
Gao, Xiang committed
28
                coordinates.requires_grad_(True)
29
30
31
                _, energies = self.model((species, coordinates))
                derivative = torch.autograd.grad(energies.sum(),
                                                 coordinates)[0]
32
                self.assertEqual(forces, -derivative)
33
34
35
36

    def testPadding(self):
        species_coordinates = []
        coordinates_forces = []
37
        for i in range(N):
38
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
39
40
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, _, forces = pickle.load(f)
41
42
43
                coordinates = torch.from_numpy(coordinates)
                species = torch.from_numpy(species)
                forces = torch.from_numpy(forces)
Gao, Xiang's avatar
Gao, Xiang committed
44
                coordinates.requires_grad_(True)
45
46
                species_coordinates.append(torchani.utils.broadcast_first_dim(
                    {'species': species, 'coordinates': coordinates}))
47
        species_coordinates = torchani.utils.pad_atomic_properties(
48
            species_coordinates)
49
        _, energies = self.model((species_coordinates['species'], species_coordinates['coordinates']))
50
51
52
53
        energies = energies.sum()
        for coordinates, forces in coordinates_forces:
            derivative = torch.autograd.grad(energies, coordinates,
                                             retain_graph=True)[0]
54
            self.assertEqual(forces, -derivative)
55
56


57
58
59
60
61
62
63
class TestForceJIT(TestForce):

    def setUp(self):
        super().setUp()
        self.model = torch.jit.script(self.model)


64
65
if __name__ == '__main__':
    unittest.main()