test_aev.py 14.8 KB
Newer Older
1
2
3
4
5
import torch
import torchani
import unittest
import os
import pickle
6
7
import itertools
import ase
8
import ase.io
9
import math
10
import traceback
Gao, Xiang's avatar
Gao, Xiang committed
11
from common_aev_test import _TestAEVBase
12

13
14

path = os.path.dirname(os.path.realpath(__file__))
Gao, Xiang's avatar
Gao, Xiang committed
15
const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
16
17
18
N = 97


19
class TestAEVConstructor(torchani.testing.TestCase):
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
    # Test that checks that the friendly constructor
    # reproduces the values from ANI1x with the correct parameters
    def testCoverLinearly(self):
        consts = torchani.neurochem.Constants(const_file)
        aev_computer = torchani.AEVComputer(**consts)
        ani1x_values = {'radial_cutoff': 5.2,
                        'angular_cutoff': 3.5,
                        'radial_eta': 16.0,
                        'angular_eta': 8.0,
                        'radial_dist_divisions': 16,
                        'angular_dist_divisions': 4,
                        'zeta': 32.0,
                        'angle_sections': 8,
                        'num_species': 4}
        aev_computer_alt = torchani.AEVComputer.cover_linearly(**ani1x_values)
        constants = aev_computer.constants()
        constants_alt = aev_computer_alt.constants()
        for c, ca in zip(constants, constants_alt):
38
            self.assertEqual(c, ca)
39
40


41
class TestIsolated(torchani.testing.TestCase):
42
43
44
45
    # Tests that there is no error when atoms are separated
    # a distance greater than the cutoff radius from all other atoms
    # this can throw an IndexError for large distances or lone atoms
    def setUp(self):
Ignacio Pickering's avatar
Ignacio Pickering committed
46
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Gao, Xiang's avatar
Gao, Xiang committed
47
48
49
50
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts).to(self.device)
        self.species_to_tensor = consts.species_to_tensor
        self.rcr = self.aev_computer.Rcr
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
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
101
102
103
104
105
        self.rca = self.aev_computer.Rca

    def testCO2(self):
        species = self.species_to_tensor(['O', 'C', 'O']).to(self.device).unsqueeze(0)
        distances = [1.0, self.rca,
                     self.rca + 1e-4, self.rcr,
                     self.rcr + 1e-4, 2 * self.rcr]
        error = ()
        for dist in distances:
            coordinates = torch.tensor(
                [[[-dist, 0., 0.], [0., 0., 0.], [0., 0., dist]]],
                requires_grad=True, device=self.device)
            try:
                _, _ = self.aev_computer((species, coordinates))
            except IndexError:
                error = (traceback.format_exc(), dist)
            if error:
                self.fail(f'\n\n{error[0]}\nFailure at distance: {error[1]}\n'
                          f'Radial r_cut of aev_computer: {self.rcr}\n'
                          f'Angular r_cut of aev_computer: {self.rca}')

    def testH2(self):
        species = self.species_to_tensor(['H', 'H']).to(self.device).unsqueeze(0)
        distances = [1.0, self.rca,
                     self.rca + 1e-4, self.rcr,
                     self.rcr + 1e-4, 2 * self.rcr]
        error = ()
        for dist in distances:
            coordinates = torch.tensor(
                [[[0., 0., 0.], [0., 0., dist]]],
                requires_grad=True, device=self.device)
            try:
                _, _ = self.aev_computer((species, coordinates))
            except IndexError:
                error = (traceback.format_exc(), dist)
            if error:
                self.fail(f'\n\n{error[0]}\nFailure at distance: {error[1]}\n'
                          f'Radial r_cut of aev_computer: {self.rcr}\n'
                          f'Angular r_cut of aev_computer: {self.rca}')

    def testH(self):
        # Tests for failure on a single atom
        species = self.species_to_tensor(['H']).to(self.device).unsqueeze(0)
        error = ()
        coordinates = torch.tensor(
            [[[0., 0., 0.]]],
            requires_grad=True, device=self.device)
        try:
            _, _ = self.aev_computer((species, coordinates))
        except IndexError:
            error = (traceback.format_exc())
        if error:
            self.fail(f'\n\n{error}\nFailure on lone atom\n')


Gao, Xiang's avatar
Gao, Xiang committed
106
class TestAEV(_TestAEVBase):
107

108
109
    def testIsomers(self):
        for i in range(N):
110
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
111
112
113
            with open(datafile, 'rb') as f:
                coordinates, species, expected_radial, expected_angular, _, _ \
                    = pickle.load(f)
114
115
116
117
                coordinates = torch.from_numpy(coordinates)
                species = torch.from_numpy(species)
                expected_radial = torch.from_numpy(expected_radial)
                expected_angular = torch.from_numpy(expected_angular)
118
                _, aev = self.aev_computer((species, coordinates))
119
120
                self.assertAEVEqual(expected_radial, expected_angular, aev)

121
122
123
    def testPadding(self):
        species_coordinates = []
        radial_angular = []
124
        for i in range(N):
125
            datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
126
127
            with open(datafile, 'rb') as f:
                coordinates, species, radial, angular, _, _ = pickle.load(f)
128
129
130
131
                coordinates = torch.from_numpy(coordinates)
                species = torch.from_numpy(species)
                radial = torch.from_numpy(radial)
                angular = torch.from_numpy(angular)
132
133
                species_coordinates.append(torchani.utils.broadcast_first_dim(
                    {'species': species, 'coordinates': coordinates}))
134
                radial_angular.append((radial, angular))
135
        species_coordinates = torchani.utils.pad_atomic_properties(
136
            species_coordinates)
137
        _, aev = self.aev_computer((species_coordinates['species'], species_coordinates['coordinates']))
138
139
140
141
        start = 0
        for expected_radial, expected_angular in radial_angular:
            conformations = expected_radial.shape[0]
            atoms = expected_radial.shape[1]
142
            aev_ = aev[start:(start + conformations), 0:atoms]
143
            start += conformations
144
            self.assertAEVEqual(expected_radial, expected_angular, aev_)
145
146


147
148
149
150
151
152
class TestAEVJIT(TestAEV):
    def setUp(self):
        super().setUp()
        self.aev_computer = torch.jit.script(self.aev_computer)


153
class TestPBCSeeEachOther(torchani.testing.TestCase):
Gao, Xiang's avatar
Gao, Xiang committed
154
    def setUp(self):
Gao, Xiang's avatar
Gao, Xiang committed
155
156
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)
157
158
159
160
161
162
163
164
165
166
167

    def testTranslationalInvariancePBC(self):
        coordinates = torch.tensor(
            [[[0, 0, 0],
              [1, 0, 0],
              [0, 1, 0],
              [0, 0, 1],
              [0, 1, 1]]],
            dtype=torch.double, requires_grad=True)
        cell = torch.eye(3, dtype=torch.double) * 2
        species = torch.tensor([[1, 0, 0, 0, 0]], dtype=torch.long)
168
        pbc = torch.ones(3, dtype=torch.bool)
169

170
        _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
171
172
173

        for _ in range(100):
            translation = torch.randn(3, dtype=torch.double)
174
            _, aev2 = self.aev_computer((species, coordinates + translation), cell=cell, pbc=pbc)
175
            self.assertEqual(aev, aev2)
176
177
178
179

    def testPBCConnersSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        cell = torch.eye(3, dtype=torch.double) * 10
180
        pbc = torch.ones(3, dtype=torch.bool)
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
        allshifts = torchani.aev.compute_shifts(cell, pbc, 1)

        xyz1 = torch.tensor([0.1, 0.1, 0.1])
        xyz2s = [
            torch.tensor([9.9, 0.0, 0.0]),
            torch.tensor([0.0, 9.9, 0.0]),
            torch.tensor([0.0, 0.0, 9.9]),
            torch.tensor([9.9, 9.9, 0.0]),
            torch.tensor([0.0, 9.9, 9.9]),
            torch.tensor([9.9, 0.0, 9.9]),
            torch.tensor([9.9, 9.9, 9.9]),
        ]

        for xyz2 in xyz2s:
            coordinates = torch.stack([xyz1, xyz2]).to(torch.double).unsqueeze(0)
196
197
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
198
199
200
201
202
            self.assertEqual(atom_index1.tolist(), [0])
            self.assertEqual(atom_index2.tolist(), [1])

    def testPBCSurfaceSeeEachOther(self):
        cell = torch.eye(3, dtype=torch.double) * 10
203
        pbc = torch.ones(3, dtype=torch.bool)
204
205
206
207
208
209
210
211
212
213
        allshifts = torchani.aev.compute_shifts(cell, pbc, 1)
        species = torch.tensor([[0, 0]])

        for i in range(3):
            xyz1 = torch.tensor([5.0, 5.0, 5.0], dtype=torch.double)
            xyz1[i] = 0.1
            xyz2 = xyz1.clone()
            xyz2[i] = 9.9

            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
214
215
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
216
217
218
219
220
            self.assertEqual(atom_index1.tolist(), [0])
            self.assertEqual(atom_index2.tolist(), [1])

    def testPBCEdgesSeeEachOther(self):
        cell = torch.eye(3, dtype=torch.double) * 10
221
        pbc = torch.ones(3, dtype=torch.bool)
222
223
224
225
226
227
228
229
230
231
        allshifts = torchani.aev.compute_shifts(cell, pbc, 1)
        species = torch.tensor([[0, 0]])

        for i, j in itertools.combinations(range(3), 2):
            xyz1 = torch.tensor([5.0, 5.0, 5.0], dtype=torch.double)
            xyz1[i] = 0.1
            xyz1[j] = 0.1
            for new_i, new_j in [[0.1, 9.9], [9.9, 0.1], [9.9, 9.9]]:
                xyz2 = xyz1.clone()
                xyz2[i] = new_i
232
                xyz2[j] = new_j
233
234

            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
235
236
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
237
238
239
240
241
242
243
            self.assertEqual(atom_index1.tolist(), [0])
            self.assertEqual(atom_index2.tolist(), [1])

    def testNonRectangularPBCConnersSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        cell = ase.geometry.cellpar_to_cell([10, 10, 10 * math.sqrt(2), 90, 45, 90])
        cell = torch.tensor(ase.geometry.complete_cell(cell), dtype=torch.double)
244
        pbc = torch.ones(3, dtype=torch.bool)
245
246
247
248
249
250
        allshifts = torchani.aev.compute_shifts(cell, pbc, 1)

        xyz1 = torch.tensor([0.1, 0.1, 0.05], dtype=torch.double)
        xyz2 = torch.tensor([10.0, 0.1, 0.1], dtype=torch.double)

        coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
251
252
        atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
        atom_index1, atom_index2 = atom_index12.unbind(0)
253
254
255
256
        self.assertEqual(atom_index1.tolist(), [0])
        self.assertEqual(atom_index2.tolist(), [1])


257
class TestAEVOnBoundary(torchani.testing.TestCase):
Gao, Xiang's avatar
Gao, Xiang committed
258

259
260
261
262
263
264
265
266
267
268
    def setUp(self):
        self.eps = 1e-9
        cell = ase.geometry.cellpar_to_cell([100, 100, 100 * math.sqrt(2), 90, 45, 90])
        self.cell = torch.tensor(ase.geometry.complete_cell(cell), dtype=torch.double)
        self.inv_cell = torch.inverse(self.cell)
        self.coordinates = torch.tensor([[[0.0, 0.0, 0.0],
                                          [1.0, -0.1, -0.1],
                                          [-0.1, 1.0, -0.1],
                                          [-0.1, -0.1, 1.0],
                                          [-1.0, -1.0, -1.0]]], dtype=torch.double)
269
        self.species = torch.tensor([[1, 0, 0, 0, 0]])
270
        self.pbc = torch.ones(3, dtype=torch.bool)
271
272
        self.v1, self.v2, self.v3 = self.cell
        self.center_coordinates = self.coordinates + 0.5 * (self.v1 + self.v2 + self.v3)
Gao, Xiang's avatar
Gao, Xiang committed
273
274
275
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)

276
        _, self.aev = self.aev_computer((self.species, self.center_coordinates), cell=self.cell, pbc=self.pbc)
277
278
279

    def assertInCell(self, coordinates):
        coordinates_cell = coordinates @ self.inv_cell
280
        self.assertEqual(coordinates, coordinates_cell @ self.cell)
281
282
283
284
285
        in_cell = (coordinates_cell >= -self.eps) & (coordinates_cell <= 1 + self.eps)
        self.assertTrue(in_cell.all())

    def assertNotInCell(self, coordinates):
        coordinates_cell = coordinates @ self.inv_cell
286
        self.assertEqual(coordinates, coordinates_cell @ self.cell)
287
288
289
290
291
292
293
294
295
296
297
        in_cell = (coordinates_cell >= -self.eps) & (coordinates_cell <= 1 + self.eps)
        self.assertFalse(in_cell.all())

    def testCornerSurfaceAndEdge(self):
        for i, j, k in itertools.product([0, 0.5, 1], repeat=3):
            if i == 0.5 and j == 0.5 and k == 0.5:
                continue
            coordinates = self.coordinates + i * self.v1 + j * self.v2 + k * self.v3
            self.assertNotInCell(coordinates)
            coordinates = torchani.utils.map2central(self.cell, coordinates, self.pbc)
            self.assertInCell(coordinates)
298
            _, aev = self.aev_computer((self.species, coordinates), cell=self.cell, pbc=self.pbc)
299
            self.assertGreater(aev.abs().max().item(), 0)
300
            self.assertEqual(aev, self.aev)
301

Gao, Xiang's avatar
Gao, Xiang committed
302

303
class TestAEVOnBenzenePBC(torchani.testing.TestCase):
304
305

    def setUp(self):
Gao, Xiang's avatar
Gao, Xiang committed
306
307
308
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts)
        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.json')
309
310
        benzene = ase.io.read(filename)
        self.cell = torch.tensor(benzene.get_cell(complete=True)).float()
311
        self.pbc = torch.tensor(benzene.get_pbc(), dtype=torch.bool)
312
313
314
        species_to_tensor = torchani.utils.ChemicalSymbolsToInts(['H', 'C', 'N', 'O'])
        self.species = species_to_tensor(benzene.get_chemical_symbols()).unsqueeze(0)
        self.coordinates = torch.tensor(benzene.get_positions()).unsqueeze(0).float()
315
        _, self.aev = self.aev_computer((self.species, self.coordinates), cell=self.cell, pbc=self.pbc)
316
        self.natoms = self.aev.shape[1]
317
318
319
320
321
322
323
324
325
326
327

    def testRepeat(self):
        c1, c2, c3 = self.cell
        species2 = self.species.repeat(1, 4)
        coordinates2 = torch.cat([
            self.coordinates,
            self.coordinates + c1,
            self.coordinates + 2 * c1,
            self.coordinates + 3 * c1,
        ], dim=1)
        cell2 = torch.stack([4 * c1, c2, c3])
328
        _, aev2 = self.aev_computer((species2, coordinates2), cell=cell2, pbc=self.pbc)
329
330
        for i in range(3):
            aev3 = aev2[:, i * self.natoms: (i + 1) * self.natoms, :]
331
            self.assertEqual(self.aev, aev3)
332
333
334
335
336
337
338
339
340
341

    def testManualMirror(self):
        c1, c2, c3 = self.cell
        species2 = self.species.repeat(1, 3 ** 3)
        coordinates2 = torch.cat([
            self.coordinates + i * c1 + j * c2 + k * c3
            for i, j, k in itertools.product([0, -1, 1], repeat=3)
        ], dim=1)
        _, aev2 = self.aev_computer((species2, coordinates2))
        aev2 = aev2[:, :self.natoms, :]
342
        self.assertEqual(self.aev, aev2)
343
344


345
346
if __name__ == '__main__':
    unittest.main()