Reduce test latency (#319)

* Move unittest check to GitHub Actions

Tests are parallelized

* Delete tests.yml

* Update setup.py

* Create test_requirements.txt

* Update unittest.yml

* Update unittest.yml

* Update README.md

* Delete install_dependencies.sh

* Delete install_dependencies_python2.sh

* Update unittest.yml

* split

* split

* fix

* fix

* -v

* more

* more

* fix

* fix

* flake8

* transform

* reduce testBenzeneMD

* version

* reduce number of epoches

* further reduce time

* Update docs.yml

examples/nnp_training.py

@@ -308,7 +308,7 @@ tensorboard = torch.utils.tensorboard.SummaryWriter()
 mse = torch.nn.MSELoss(reduction='none')
 
 print("training starting from epoch", AdamW_scheduler.last_epoch + 1)
-max_epochs = 200
+max_epochs = 10
 early_stopping_learning_rate = 1.0E-5
 best_model_checkpoint = 'best.pt'
 

examples/nnp_training_force.py

@@ -248,7 +248,9 @@ tensorboard = torch.utils.tensorboard.SummaryWriter()
 mse = torch.nn.MSELoss(reduction='none')
 
 print("training starting from epoch", AdamW_scheduler.last_epoch + 1)
-max_epochs = 20
+# We only train 3 epoches here in able to generate the docs quickly.
+# Real training should take much more than 3 epoches.
+max_epochs = 3
 early_stopping_learning_rate = 1.0E-5
 force_coefficient = 0.1  # controls the importance of energy loss vs force loss
 best_model_checkpoint = 'force-training-best.pt'

tests/common_aev_test.py

+import random
+import unittest
+import torch
+import torchani
+
+tolerance = 1e-5
+
+
+class _TestAEVBase(unittest.TestCase):
+
+    def setUp(self):
+        ani1x = torchani.models.ANI1x()
+        self.aev_computer = ani1x.aev_computer
+        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:]
+        radial_diff = expected_radial - radial
+        if self.debug:
+            aid = 1
+            print(torch.stack([expected_radial[0, aid, :], radial[0, aid, :], radial_diff.abs()[0, aid, :]], dim=1))
+        radial_max_error = torch.max(torch.abs(radial_diff)).item()
+        angular_diff = expected_angular - angular
+        angular_max_error = torch.max(torch.abs(angular_diff)).item()
+        self.assertLess(radial_max_error, tolerance)
+        self.assertLess(angular_max_error, tolerance)

tests/test_aev.py

@@ -3,18 +3,17 @@ import torchani
 import unittest
 import os
 import pickle
-import random
 import copy
 import itertools
 import ase
 import ase.io
 import math
 import traceback
+from common_aev_test import _TestAEVBase
 
 
 path = os.path.dirname(os.path.realpath(__file__))
 N = 97
-tolerance = 1e-5
 
 
 class TestIsolated(unittest.TestCase):
@@ -85,32 +84,7 @@ class TestIsolated(unittest.TestCase):
             self.fail(f'\n\n{error}\nFailure on lone atom\n')
 
 
-class TestAEV(unittest.TestCase):
-
-    def setUp(self):
-        ani1x = torchani.models.ANI1x()
-        self.aev_computer = ani1x.aev_computer
-        self.radial_length = self.aev_computer.radial_length
-        self.debug = False
-
-    def random_skip(self, prob=0):
-        return random.random() < prob
-
-    def transform(self, x):
-        return x
-
-    def assertAEVEqual(self, expected_radial, expected_angular, aev, tolerance=tolerance):
-        radial = aev[..., :self.radial_length]
-        angular = aev[..., self.radial_length:]
-        radial_diff = expected_radial - radial
-        if self.debug:
-            aid = 1
-            print(torch.stack([expected_radial[0, aid, :], radial[0, aid, :], radial_diff.abs()[0, aid, :]], dim=1))
-        radial_max_error = torch.max(torch.abs(radial_diff)).item()
-        angular_diff = expected_angular - angular
-        angular_max_error = torch.max(torch.abs(angular_diff)).item()
-        self.assertLess(radial_max_error, tolerance)
-        self.assertLess(angular_max_error, tolerance)
+class TestAEV(_TestAEVBase):
 
     def testIsomers(self):
         for i in range(N):
@@ -129,44 +103,6 @@ class TestAEV(unittest.TestCase):
                 _, aev = self.aev_computer((species, coordinates))
                 self.assertAEVEqual(expected_radial, expected_angular, aev)
 
-    def testBenzeneMD(self):
-        for i in range(10):
-            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, expected_radial, expected_angular, _, _, cell, pbc \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
-                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)
-                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)
-
-    def testTripeptideMD(self):
-        tol = 5e-6
-        for i in range(100):
-            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, expected_radial, expected_angular, _, _, _, _ \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
-                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)
-
     def testPadding(self):
         species_coordinates = []
         radial_angular = []
@@ -195,20 +131,6 @@ class TestAEV(unittest.TestCase):
             start += conformations
             self.assertAEVEqual(expected_radial, expected_angular, aev_)
 
-    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, 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)
-                angular = torch.from_numpy(angular).to(torch.float)
-                _, aev = self.aev_computer((species, coordinates))
-                self.assertAEVEqual(radial, angular, aev)
-
     @unittest.skipIf(not torch.cuda.is_available(), "Too slow on CPU")
     def testGradient(self):
         """Test validity of autodiff by comparing analytical and numerical

tests/test_aev_benzene_md.py

+import os
+import torch
+import pickle
+import torchani
+import unittest
+from common_aev_test import _TestAEVBase
+
+path = os.path.dirname(os.path.realpath(__file__))
+
+
+class TestAEVBenzeneMD(_TestAEVBase):
+
+    def testBenzeneMD(self):
+        for i in [2, 8]:
+            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, expected_radial, expected_angular, _, _, cell, pbc \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
+                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)
+                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)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_aev_nist.py

+import os
+import torch
+import pickle
+import unittest
+from common_aev_test import _TestAEVBase
+
+path = os.path.dirname(os.path.realpath(__file__))
+
+
+class TestAEVNIST(_TestAEVBase):
+
+    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, 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)
+                angular = torch.from_numpy(angular).to(torch.float)
+                _, aev = self.aev_computer((species, coordinates))
+                self.assertAEVEqual(radial, angular, aev)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_aev_tripeptide_md.py

+import os
+import torch
+import pickle
+import unittest
+from common_aev_test import _TestAEVBase
+
+path = os.path.dirname(os.path.realpath(__file__))
+
+
+class TestAEVTripeptideMD(_TestAEVBase):
+
+    def testTripeptideMD(self):
+        tol = 5e-6
+        for i in range(100):
+            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, expected_radial, expected_angular, _, _, _, _ \
+                    = pickle.load(f)
+                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
+                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)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_energies.py

@@ -3,7 +3,6 @@ import torchani
 import unittest
 import os
 import pickle
-import math
 
 
 path = os.path.dirname(os.path.realpath(__file__))
@@ -42,42 +41,6 @@ class TestEnergies(unittest.TestCase):
                 max_diff = (energies - energies_).abs().max().item()
                 self.assertLess(max_diff, self.tolerance)
 
-    def testBenzeneMD(self):
-        tolerance = 1e-5
-        for i in range(10):
-            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, _, _, energies, _, cell, pbc \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
-                species = torch.from_numpy(species).unsqueeze(0)
-                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)
-                energies = self.transform(energies)
-                _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
-                _, energies_ = self.nn((species, aev))
-                max_diff = (energies - energies_).abs().max().item()
-                self.assertLess(max_diff, tolerance)
-
-    def testTripeptideMD(self):
-        tolerance = 2e-4
-        for i in range(100):
-            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, _, _, energies, _, _, _ \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0)
-                species = torch.from_numpy(species).unsqueeze(0)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                energies = self.transform(energies)
-                _, energies_ = self.model((species, coordinates))
-                max_diff = (energies - energies_).abs().max().item()
-                self.assertLess(max_diff, tolerance)
-
     def testPadding(self):
         species_coordinates = []
         energies = []
@@ -100,21 +63,6 @@ class TestEnergies(unittest.TestCase):
         max_diff = (energies - energies_).abs().max().item()
         self.assertLess(max_diff, self.tolerance)
 
-    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)
-
 
 class TestEnergiesEnergyShifterJIT(TestEnergies):
     def setUp(self):

tests/test_forces.py

@@ -66,65 +66,6 @@ class TestForce(unittest.TestCase):
             max_diff = (forces + derivative).abs().max().item()
             self.assertLess(max_diff, self.tolerance)
 
-    def testBenzeneMD(self):
-        tolerance = 1e-5
-        for i in range(10):
-            datafile = os.path.join(path, 'test_data/benzene-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, _, _, _, forces, cell, pbc \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0).requires_grad_(True)
-                species = torch.from_numpy(species).unsqueeze(0)
-                cell = torch.from_numpy(cell).float()
-                pbc = torch.from_numpy(pbc)
-                forces = torch.from_numpy(forces)
-                coordinates = torchani.utils.map2central(cell, coordinates, pbc)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                forces = self.transform(forces)
-                _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
-                _, energies_ = self.nnp((species, aev))
-                derivative = torch.autograd.grad(energies_.sum(),
-                                                 coordinates)[0]
-                max_diff = (forces + derivative).abs().max().item()
-                self.assertLess(max_diff, tolerance)
-
-    def testTripeptideMD(self):
-        tolerance = 2e-6
-        for i in range(100):
-            datafile = os.path.join(path, 'test_data/tripeptide-md/{}.dat'.format(i))
-            with open(datafile, 'rb') as f:
-                coordinates, species, _, _, _, forces, _, _ \
-                    = pickle.load(f)
-                coordinates = torch.from_numpy(coordinates).float().unsqueeze(0).requires_grad_(True)
-                species = torch.from_numpy(species).unsqueeze(0)
-                forces = torch.from_numpy(forces)
-                coordinates = self.transform(coordinates)
-                species = self.transform(species)
-                forces = self.transform(forces)
-                _, energies_ = self.model((species, coordinates))
-                derivative = torch.autograd.grad(energies_.sum(),
-                                                 coordinates)[0]
-                max_diff = (forces + derivative).abs().max().item()
-                self.assertLess(max_diff, tolerance)
-
-    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, _, _, _, forces in data:
-                if self.random_skip():
-                    continue
-                coordinates = torch.from_numpy(coordinates).to(torch.float) \
-                                   .requires_grad_(True)
-                species = torch.from_numpy(species)
-                forces = torch.from_numpy(forces).to(torch.float)
-                _, energies = self.model((species, coordinates))
-                derivative = torch.autograd.grad(energies.sum(),
-                                                 coordinates)[0]
-                max_diff = (forces + derivative).abs().max().item()
-                self.assertLess(max_diff, self.tolerance)
-
 
 if __name__ == '__main__':
     unittest.main()