Tests cleanup (#495)

* Remove test gradient from aev tests

* Add a dedicated test gradient for one model, and make it feasible

* add test_grad to unittests

* Add some comments and fix formatting

* flake8

* flake8

* remove unused var

* delete transform function, which was legacy from old pytorch probably

* Delete random_skip function, which was unused

* Allow testNeuroChem to use cuda if available, which makes it much faster

* small modification for consistency

* remove unused import

* simplify device call

* remove unused variable

* add main to test_grad

tests/common_aev_test.py

-import random
 import unittest
 import torch
 import torchani
@@ -14,12 +13,6 @@ class _TestAEVBase(unittest.TestCase):
         self.radial_length = self.aev_computer.radial_length
         self.debug = False
 
-    def transform(self, x):
-        return x
-
-    def random_skip(self, prob=0):
-        return random.random() < prob
-
     def assertAEVEqual(self, expected_radial, expected_angular, aev, tolerance=tolerance):
         radial = aev[..., :self.radial_length]
         angular = aev[..., self.radial_length:]

tests/test_aev.py

@@ -20,10 +20,7 @@ class TestIsolated(unittest.TestCase):
     # 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):
-        if torch.cuda.is_available():
-            self.device = 'cuda'
-        else:
-            self.device = 'cpu'
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         ani1x = torchani.models.ANI1x().to(self.device)
         self.aev_computer = ani1x.aev_computer
         self.species_to_tensor = ani1x.species_to_tensor
@@ -95,10 +92,6 @@ class TestAEV(_TestAEVBase):
                 species = torch.from_numpy(species)
                 expected_radial = torch.from_numpy(expected_radial)
                 expected_angular = torch.from_numpy(expected_angular)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                expected_radial = self.transform(expected_radial)
-                expected_angular = self.transform(expected_angular)
                 _, aev = self.aev_computer((species, coordinates))
                 self.assertAEVEqual(expected_radial, expected_angular, aev)
 
@@ -113,10 +106,6 @@ class TestAEV(_TestAEVBase):
                 species = torch.from_numpy(species)
                 radial = torch.from_numpy(radial)
                 angular = torch.from_numpy(angular)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                radial = self.transform(radial)
-                angular = self.transform(angular)
                 species_coordinates.append(torchani.utils.broadcast_first_dim(
                     {'species': species, 'coordinates': coordinates}))
                 radial_angular.append((radial, angular))

tests/test_aev_benzene_md.py

@@ -23,10 +23,6 @@ class TestAEVBenzeneMD(_TestAEVBase):
                 cell = torch.from_numpy(cell).float()
                 pbc = torch.from_numpy(pbc)
                 coordinates = torchani.utils.map2central(cell, coordinates, pbc)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                expected_radial = self.transform(expected_radial)
-                expected_angular = self.transform(expected_angular)
                 _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
                 self.assertAEVEqual(expected_radial, expected_angular, aev, 5e-5)
 

tests/test_aev_nist.py

@@ -14,8 +14,6 @@ class TestAEVNIST(_TestAEVBase):
         with open(datafile, 'rb') as f:
             data = pickle.load(f)
             for coordinates, species, radial, angular, _, _ in data:
-                if self.random_skip():
-                    continue
                 coordinates = torch.from_numpy(coordinates).to(torch.float)
                 species = torch.from_numpy(species)
                 radial = torch.from_numpy(radial).to(torch.float)

tests/test_aev_tripeptide_md.py

@@ -20,10 +20,6 @@ class TestAEVTripeptideMD(_TestAEVBase):
                 species = torch.from_numpy(species).unsqueeze(0)
                 expected_radial = torch.from_numpy(expected_radial).float().unsqueeze(0)
                 expected_angular = torch.from_numpy(expected_angular).float().unsqueeze(0)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                expected_radial = self.transform(expected_radial)
-                expected_angular = self.transform(expected_angular)
                 _, aev = self.aev_computer((species, coordinates))
                 self.assertAEVEqual(expected_radial, expected_angular, aev, tol)
 

tests/test_ase.py

@@ -10,7 +10,6 @@ import unittest
 import os
 
 path = os.path.dirname(os.path.realpath(__file__))
-N = 97
 tol = 5e-5
 
 

tests/test_energies.py

@@ -20,12 +20,6 @@ class TestEnergies(unittest.TestCase):
         self.nn = torchani.nn.Sequential(self.nnp, self.energy_shifter)
         self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
 
-    def random_skip(self):
-        return False
-
-    def transform(self, x):
-        return x
-
     def testIsomers(self):
         for i in range(N):
             datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
@@ -34,10 +28,7 @@ class TestEnergies(unittest.TestCase):
                 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)
-                _, energies_ = self.model((species, coordinates))
+                energies_ = self.model((species, coordinates)).energies
                 max_diff = (energies - energies_).abs().max().item()
                 self.assertLess(max_diff, self.tolerance)
 
@@ -51,16 +42,13 @@ class TestEnergies(unittest.TestCase):
                 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)
                 species_coordinates.append(
                     torchani.utils.broadcast_first_dim({'species': species, 'coordinates': coordinates}))
                 energies.append(e)
         species_coordinates = torchani.utils.pad_atomic_properties(
             species_coordinates)
         energies = torch.cat(energies)
-        _, energies_ = self.model((species_coordinates['species'], species_coordinates['coordinates']))
+        energies_ = self.model((species_coordinates['species'], species_coordinates['coordinates'])).energies
         max_diff = (energies - energies_).abs().max().item()
         self.assertLess(max_diff, self.tolerance)
 

tests/test_forces.py

@@ -17,12 +17,6 @@ class TestForce(unittest.TestCase):
         self.nnp = model.neural_networks
         self.model = torchani.nn.Sequential(self.aev_computer, self.nnp)
 
-    def random_skip(self):
-        return False
-
-    def transform(self, x):
-        return x
-
     def testIsomers(self):
         for i in range(N):
             datafile = os.path.join(path, 'test_data/ANI1_subset/{}'.format(i))
@@ -31,9 +25,6 @@ class TestForce(unittest.TestCase):
                 coordinates = torch.from_numpy(coordinates)
                 species = torch.from_numpy(species)
                 forces = torch.from_numpy(forces)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                forces = self.transform(forces)
                 coordinates.requires_grad_(True)
                 _, energies = self.model((species, coordinates))
                 derivative = torch.autograd.grad(energies.sum(),
@@ -51,9 +42,6 @@ class TestForce(unittest.TestCase):
                 coordinates = torch.from_numpy(coordinates)
                 species = torch.from_numpy(species)
                 forces = torch.from_numpy(forces)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                forces = self.transform(forces)
                 coordinates.requires_grad_(True)
                 species_coordinates.append(torchani.utils.broadcast_first_dim(
                     {'species': species, 'coordinates': coordinates}))

tests/test_grad.py

@@ -54,3 +54,7 @@ class TestGrad(unittest.TestCase):
             torch.autograd.gradgradcheck(lambda x: self.model((species, x)).energies,
                                          coordinates,
                                          nondet_tol=1e-13)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_neurochem.py

@@ -12,7 +12,7 @@ dspath = os.path.join(path, '../dataset/ani1-up_to_gdb4/ani_gdb_s01.h5')
 class TestNeuroChem(unittest.TestCase):
 
     def testNeuroChemTrainer(self):
-        d = torch.device('cpu')
+        d = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         trainer = torchani.neurochem.Trainer(iptpath, d, True, os.path.join(path, 'runs'))
 
         # test if loader construct correct model

tests/test_periodic_table_indexing.py

@@ -40,7 +40,7 @@ class TestBuiltinEnsemblePeriodicTableIndex(unittest.TestCase):
                                           [0.45554739, 0.54289633, 0.81170881],
                                           [0.66091919, -0.16799635, -0.91037834]]],
                                         requires_grad=True)
-        self.species1 = self.model1.species_to_tensor('CHHHH').unsqueeze(0)
+        self.species1 = self.model1.species_to_tensor(['C', 'H', 'H', 'H', 'H']).unsqueeze(0)
         self.species2 = torch.tensor([[6, 1, 1, 1, 1]])
 
     def testCH4Ensemble(self):