test_energies.py 4.33 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
import torch
import torchani
import unittest
import os
import pickle


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


12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class TestCorrectInput(torchani.testing.TestCase):

    def setUp(self):
        self.model = torchani.models.ANI1x(model_index=0, periodic_table_index=False)
        self.converter = torchani.nn.SpeciesConverter(['H', 'C', 'N', 'O'])
        self.aev_computer = self.model.aev_computer
        self.ani_model = self.model.neural_networks

    def testUnknownSpecies(self):
        # unsupported atomic number raises a value error
        self.assertRaises(ValueError, self.converter, (torch.tensor([[1, 1, 7, 10]]), torch.zeros((1, 4, 3))))
        # larger index than supported by the model raises a value error
        self.assertRaises(ValueError, self.model, (torch.tensor([[0, 1, 2, 4]]), torch.zeros((1, 4, 3))))

    def testIncorrectShape(self):
        # non matching shapes between species and coordinates
        self.assertRaises(AssertionError, self.model, (torch.tensor([[0, 1, 2, 3]]), torch.zeros((1, 3, 3))))
        self.assertRaises(AssertionError, self.aev_computer, (torch.tensor([[0, 1, 2, 3]]), torch.zeros((1, 3, 3))))
        self.assertRaises(AssertionError, self.ani_model, (torch.tensor([[0, 1, 2, 3]]), torch.zeros((1, 3, 384))))
        self.assertRaises(AssertionError, self.model, (torch.tensor([[0, 1, 2, 3]]), torch.zeros((1, 4, 4))))
        self.assertRaises(AssertionError, self.model, (torch.tensor([0, 1, 2, 3]), torch.zeros((4, 3))))


35
class TestEnergies(torchani.testing.TestCase):
Gao, Xiang's avatar
Gao, Xiang committed
36
37
    # tests the predicions for a torchani.nn.Sequential(AEVComputer(),
    # ANIModel(), EnergyShifter()) against precomputed values
38

39
    def setUp(self):
40
41
42
43
        model = torchani.models.ANI1x(model_index=0)
        self.aev_computer = model.aev_computer
        self.nnp = model.neural_networks
        self.energy_shifter = model.energy_shifter
44
        self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
45

46
    def testIsomers(self):
47
        for i in range(N):
48
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
49
50
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, energies, _ = pickle.load(f)
51
52
53
                coordinates = torch.from_numpy(coordinates).to(torch.float)
                species = torch.from_numpy(species)
                energies = torch.from_numpy(energies).to(torch.float)
Ignacio Pickering's avatar
Ignacio Pickering committed
54
                energies_ = self.model((species, coordinates)).energies
55
                self.assertEqual(energies, energies_, exact_dtype=False)
56
57
58
59
60

    def testPadding(self):
        species_coordinates = []
        energies = []
        for i in range(N):
61
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
62
63
            with open(datafile, 'rb') as f:
                coordinates, species, _, _, e, _ = pickle.load(f)
64
65
66
                coordinates = torch.from_numpy(coordinates).to(torch.float)
                species = torch.from_numpy(species)
                e = torch.from_numpy(e).to(torch.float)
67
68
                species_coordinates.append(
                    torchani.utils.broadcast_first_dim({'species': species, 'coordinates': coordinates}))
69
                energies.append(e)
70
        species_coordinates = torchani.utils.pad_atomic_properties(
71
72
            species_coordinates)
        energies = torch.cat(energies)
Ignacio Pickering's avatar
Ignacio Pickering committed
73
        energies_ = self.model((species_coordinates['species'], species_coordinates['coordinates'])).energies
74
        self.assertEqual(energies, energies_, exact_dtype=False)
75
76


77
class TestEnergiesEnergyShifterJIT(TestEnergies):
Gao, Xiang's avatar
Gao, Xiang committed
78
    # only JIT compile the energy shifter and repeat all tests
79

80
81
82
    def setUp(self):
        super().setUp()
        self.energy_shifter = torch.jit.script(self.energy_shifter)
83
        self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
84
85


86
class TestEnergiesANIModelJIT(TestEnergies):
Gao, Xiang's avatar
Gao, Xiang committed
87
    # only JIT compile the ANI nnp ANIModel and repeat all tests
88
89
90
91
92
93
94
95

    def setUp(self):
        super().setUp()
        self.nnp = torch.jit.script(self.nnp)
        self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)


class TestEnergiesJIT(TestEnergies):
Gao, Xiang's avatar
Gao, Xiang committed
96
    # JIT compile the whole model and repeat all tests
97
98
99
100
101
102

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


103
104
if __name__ == '__main__':
    unittest.main()