test_aev.py 15.1 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
Jinze Xue's avatar
Jinze Xue committed
12
from torchani.testing import TestCase
13

14
15

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


Jinze Xue's avatar
Jinze Xue committed
20
class TestAEVConstructor(TestCase):
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
    # 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):
39
            self.assertEqual(c, ca)
40
41


Jinze Xue's avatar
Jinze Xue committed
42
class TestIsolated(TestCase):
43
44
45
46
    # 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
47
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Gao, Xiang's avatar
Gao, Xiang committed
48
49
50
51
        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
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
106
        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
107
class TestAEV(_TestAEVBase):
108

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

122
123
124
125
126
127
128
    def testNoNan(self):
        # AEV should not output NaN even when coordinates are superimposed
        coordinates = torch.ones(1, 3, 3, dtype=torch.float)
        species = torch.zeros(1, 3, dtype=torch.long)
        _, aev = self.aev_computer((species, coordinates))
        self.assertFalse(torch.isnan(aev).any())

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


155
156
157
158
159
160
class TestAEVJIT(TestAEV):
    def setUp(self):
        super().setUp()
        self.aev_computer = torch.jit.script(self.aev_computer)


Jinze Xue's avatar
Jinze Xue committed
161
class TestPBCSeeEachOther(TestCase):
Gao, Xiang's avatar
Gao, Xiang committed
162
    def setUp(self):
Gao, Xiang's avatar
Gao, Xiang committed
163
164
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)
165
166
167
168
169
170
171
172
173
174
175

    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)
176
        pbc = torch.ones(3, dtype=torch.bool)
177

178
        _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
179
180
181

        for _ in range(100):
            translation = torch.randn(3, dtype=torch.double)
182
            _, aev2 = self.aev_computer((species, coordinates + translation), cell=cell, pbc=pbc)
183
            self.assertEqual(aev, aev2)
184
185
186
187

    def testPBCConnersSeeEachOther(self):
        species = torch.tensor([[0, 0]])
        cell = torch.eye(3, dtype=torch.double) * 10
188
        pbc = torch.ones(3, dtype=torch.bool)
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
        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)
204
205
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
206
207
208
209
210
            self.assertEqual(atom_index1.tolist(), [0])
            self.assertEqual(atom_index2.tolist(), [1])

    def testPBCSurfaceSeeEachOther(self):
        cell = torch.eye(3, dtype=torch.double) * 10
211
        pbc = torch.ones(3, dtype=torch.bool)
212
213
214
215
216
217
218
219
220
221
        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)
222
223
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
224
225
226
227
228
            self.assertEqual(atom_index1.tolist(), [0])
            self.assertEqual(atom_index2.tolist(), [1])

    def testPBCEdgesSeeEachOther(self):
        cell = torch.eye(3, dtype=torch.double) * 10
229
        pbc = torch.ones(3, dtype=torch.bool)
230
231
232
233
234
235
236
237
238
239
        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
240
                xyz2[j] = new_j
241
242

            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
243
244
            atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
            atom_index1, atom_index2 = atom_index12.unbind(0)
245
246
247
248
249
250
251
            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)
252
        pbc = torch.ones(3, dtype=torch.bool)
253
254
255
256
257
258
        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)
259
260
        atom_index12, _ = torchani.aev.neighbor_pairs(species == -1, coordinates, cell, allshifts, 1)
        atom_index1, atom_index2 = atom_index12.unbind(0)
261
262
263
264
        self.assertEqual(atom_index1.tolist(), [0])
        self.assertEqual(atom_index2.tolist(), [1])


Jinze Xue's avatar
Jinze Xue committed
265
class TestAEVOnBoundary(TestCase):
Gao, Xiang's avatar
Gao, Xiang committed
266

267
268
269
270
271
272
273
274
275
276
    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)
277
        self.species = torch.tensor([[1, 0, 0, 0, 0]])
278
        self.pbc = torch.ones(3, dtype=torch.bool)
279
280
        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
281
282
283
        consts = torchani.neurochem.Constants(const_file)
        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)

284
        _, self.aev = self.aev_computer((self.species, self.center_coordinates), cell=self.cell, pbc=self.pbc)
285
286
287

    def assertInCell(self, coordinates):
        coordinates_cell = coordinates @ self.inv_cell
288
        self.assertEqual(coordinates, coordinates_cell @ self.cell)
289
290
291
292
293
        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
294
        self.assertEqual(coordinates, coordinates_cell @ self.cell)
295
296
297
298
299
300
301
302
303
304
305
        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)
306
            _, aev = self.aev_computer((self.species, coordinates), cell=self.cell, pbc=self.pbc)
307
            self.assertGreater(aev.abs().max().item(), 0)
308
            self.assertEqual(aev, self.aev)
309

Gao, Xiang's avatar
Gao, Xiang committed
310

Jinze Xue's avatar
Jinze Xue committed
311
class TestAEVOnBenzenePBC(TestCase):
312
313

    def setUp(self):
Gao, Xiang's avatar
Gao, Xiang committed
314
315
316
        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')
317
318
        benzene = ase.io.read(filename)
        self.cell = torch.tensor(benzene.get_cell(complete=True)).float()
319
        self.pbc = torch.tensor(benzene.get_pbc(), dtype=torch.bool)
320
321
322
        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()
323
        _, self.aev = self.aev_computer((self.species, self.coordinates), cell=self.cell, pbc=self.pbc)
324
        self.natoms = self.aev.shape[1]
325
326
327
328
329
330
331
332
333
334
335

    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])
336
        _, aev2 = self.aev_computer((species2, coordinates2), cell=cell2, pbc=self.pbc)
337
338
        for i in range(3):
            aev3 = aev2[:, i * self.natoms: (i + 1) * self.natoms, :]
339
            self.assertEqual(self.aev, aev3)
340
341
342
343
344
345
346
347
348
349

    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, :]
350
        self.assertEqual(self.aev, aev2)
351
352


353
354
if __name__ == '__main__':
    unittest.main()