test_energies.py 3.52 KB
Newer Older
1
2
3
4
5
import torch
import torchani
import unittest
import os
import pickle
6
7
import math
import random
8
9
10
11
12
13
14
15


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


class TestEnergies(unittest.TestCase):

16
    def setUp(self):
17
        self.tolerance = 5e-5
18
        builtins = torchani.neurochem.Builtins()
19
        self.aev_computer = builtins.aev_computer
20
21
        nnp = builtins.models[0]
        shift_energy = builtins.energy_shifter
22
        self.model = torch.nn.Sequential(self.aev_computer, nnp, shift_energy)
23

24
25
26
27
28
29
    def random_skip(self):
        return False

    def transform(self, x):
        return x

30
    def testIsomers(self):
31
        for i in range(N):
32
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
33
34
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, energies, _ = pickle.load(f)
35
36
37
38
39
40
                coordinates = torch.from_numpy(coordinates).to(torch.float)
                species = torch.from_numpy(species)
                energies = torch.from_numpy(energies).to(torch.float)
                coordinates = self.transform(coordinates)
                species = self.transform(species)
                energies = self.transform(energies)
41
42
43
44
45
46
47
48
                _, energies_ = self.model((species, coordinates))
                max_diff = (energies - energies_).abs().max().item()
                self.assertLess(max_diff, self.tolerance)

    def testPadding(self):
        species_coordinates = []
        energies = []
        for i in range(N):
49
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
50
51
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, e, _ = pickle.load(f)
52
53
54
55
56
57
                coordinates = torch.from_numpy(coordinates).to(torch.float)
                species = torch.from_numpy(species)
                e = torch.from_numpy(e).to(torch.float)
                coordinates = self.transform(coordinates)
                species = self.transform(species)
                e = self.transform(e)
58
                species_coordinates.append((species, coordinates))
59
                energies.append(e)
60
        species, coordinates = torchani.utils.pad_coordinates(
61
62
63
64
65
            species_coordinates)
        energies = torch.cat(energies)
        _, energies_ = self.model((species, coordinates))
        max_diff = (energies - energies_).abs().max().item()
        self.assertLess(max_diff, self.tolerance)
66

67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    def testNIST(self):
        datafile = os.path.join(path, 'test_data/NIST/all')
        with open(datafile, 'rb') as f:
            data = pickle.load(f)
            for coordinates, species, _, _, e, _ in data:
                if self.random_skip():
                    continue
                coordinates = torch.from_numpy(coordinates).to(torch.float)
                species = torch.from_numpy(species)
                energies = torch.from_numpy(e).to(torch.float)
                _, energies_ = self.model((species, coordinates))
                natoms = coordinates.shape[1]
                max_diff = (energies - energies_).abs().max().item()
                self.assertLess(max_diff / math.sqrt(natoms), self.tolerance)

82

83
84
85
86
87
class TestEnergiesASEComputer(TestEnergies):

    def setUp(self):
        super(TestEnergiesASEComputer, self).setUp()

88
89
90
91
92
93
94
95
    def transform(self, x):
        """To reduce the size of test cases for faster test speed"""
        return x[:2, ...]

    def random_skip(self):
        """To reduce the size of test cases for faster test speed"""
        return random.random() < 0.95

96

97
98
if __name__ == '__main__':
    unittest.main()