Faster tests (#8) (#537)

* modifications to make tests faster

* Now ani1x is not loaded for any aev_computer necessary

* flake*

* Avoid test warnings by loading from json and setting initial velocities to a small number

* Avoid loading from cif in ManualMirror and Repeat also

* Add comments to clarify test_energies and delete unused variable

* Fix inaccurate comment

* Delete unused variable in test_ensemble

* Add comments and clarifications to test_ase

* flake8
Co-authored-by: Ignacio Pickering <ign.pickering@gmail.com>

tests/common_aev_test.py

 import unittest
 import torch
 import torchani
+import os
 
 tolerance = 1e-5
 
@@ -8,8 +9,10 @@ tolerance = 1e-5
 class _TestAEVBase(unittest.TestCase):
 
     def setUp(self):
-        ani1x = torchani.models.ANI1x()
-        self.aev_computer = ani1x.aev_computer
+        path = os.path.dirname(os.path.realpath(__file__))
+        const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
+        consts = torchani.neurochem.Constants(const_file)
+        self.aev_computer = torchani.AEVComputer(**consts)
         self.radial_length = self.aev_computer.radial_length
         self.debug = False
 

tests/test_aev.py

@@ -12,6 +12,7 @@ from common_aev_test import _TestAEVBase
 
 
 path = os.path.dirname(os.path.realpath(__file__))
+const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params')  # noqa: E501
 N = 97
 
 
@@ -21,10 +22,10 @@ class TestIsolated(unittest.TestCase):
     # this can throw an IndexError for large distances or lone atoms
     def setUp(self):
         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
-        self.rcr = ani1x.aev_computer.Rcr
+        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
         self.rca = self.aev_computer.Rca
 
     def testCO2(self):
@@ -129,8 +130,8 @@ class TestAEVJIT(TestAEV):
 
 class TestPBCSeeEachOther(unittest.TestCase):
     def setUp(self):
-        self.ani1x = torchani.models.ANI1x()
-        self.aev_computer = self.ani1x.aev_computer.to(torch.double)
+        consts = torchani.neurochem.Constants(const_file)
+        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)
 
     def testTranslationalInvariancePBC(self):
         coordinates = torch.tensor(
@@ -247,8 +248,9 @@ class TestAEVOnBoundary(unittest.TestCase):
         self.pbc = torch.ones(3, dtype=torch.bool)
         self.v1, self.v2, self.v3 = self.cell
         self.center_coordinates = self.coordinates + 0.5 * (self.v1 + self.v2 + self.v3)
-        ani1x = torchani.models.ANI1x()
-        self.aev_computer = ani1x.aev_computer.to(torch.double)
+        consts = torchani.neurochem.Constants(const_file)
+        self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)
+
         _, self.aev = self.aev_computer((self.species, self.center_coordinates), cell=self.cell, pbc=self.pbc)
 
     def assertInCell(self, coordinates):
@@ -279,9 +281,9 @@ class TestAEVOnBoundary(unittest.TestCase):
 class TestAEVOnBenzenePBC(unittest.TestCase):
 
     def setUp(self):
-        ani1x = torchani.models.ANI1x()
-        self.aev_computer = ani1x.aev_computer
-        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.cif')
+        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')
         benzene = ase.io.read(filename)
         self.cell = torch.tensor(benzene.get_cell(complete=True)).float()
         self.pbc = torch.tensor(benzene.get_pbc(), dtype=torch.bool)

tests/test_aev_nist.py

@@ -13,7 +13,9 @@ class TestAEVNIST(_TestAEVBase):
         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:
+            # only use first 100 data points to make test take an
+            # acceptable time
+            for coordinates, species, radial, angular, _, _ in data[:100]:
                 coordinates = torch.from_numpy(coordinates).to(torch.float)
                 species = torch.from_numpy(species)
                 radial = torch.from_numpy(radial).to(torch.float)

tests/test_ase.py

@@ -4,13 +4,13 @@ from ase.md.nptberendsen import NPTBerendsen
 from ase import units
 from ase.io import read
 from ase.calculators.test import numeric_force
+import numpy as np
 import torch
 import torchani
 import unittest
 import os
 
 path = os.path.dirname(os.path.realpath(__file__))
-tol = 5e-5
 
 
 def get_numeric_force(atoms, eps):
@@ -27,6 +27,10 @@ class TestASE(unittest.TestCase):
         self.model = torchani.models.ANI1x(model_index=0).double()
 
     def testWithNumericalForceWithPBCEnabled(self):
+        # Run a Langevin thermostat dynamic for 100 steps and after the dynamic
+        # check once that the numerical and analytical force agree to a given
+        # relative tolerance
+        relative_tolerance = 0.1
         atoms = Diamond(symbol="C", pbc=True)
         calculator = self.model.ase()
         atoms.set_calculator(calculator)
@@ -37,11 +41,21 @@ class TestASE(unittest.TestCase):
         df = (f - fn).abs().max()
         avgf = f.abs().mean()
         if avgf > 0:
-            self.assertLess(df / avgf, 0.1)
+            self.assertLess(df / avgf, relative_tolerance)
 
     def testWithNumericalStressWithPBCEnabled(self):
-        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.cif')
+        # Run NPT dynamics for some steps and periodically check that the
+        # numerical and analytical stresses agree up to a given
+        # absolute difference
+        tolerance = 1e-5
+        filename = os.path.join(path, '../tools/generate-unit-test-expect/others/Benzene.json')
         benzene = read(filename)
+        # set velocities to a very small value to avoid division by zero
+        # warning due to initial zero temperature.
+        #
+        # Note that there are 4 benzene molecules, thus, 48 atoms in
+        # Benzene.json
+        benzene.set_velocities(np.full((48, 3), 1e-15))
         calculator = self.model.ase()
         benzene.set_calculator(calculator)
         dyn = NPTBerendsen(benzene, timestep=0.1 * units.fs,
@@ -53,12 +67,15 @@ class TestASE(unittest.TestCase):
             stress = benzene.get_stress()
             numerical_stress = calculator.calculate_numerical_stress(benzene)
             diff = torch.from_numpy(stress - numerical_stress).abs().max().item()
-            self.assertLess(diff, tol)
-        dyn.attach(test_stress, interval=30)
-        dyn.run(120)
+            self.assertLess(diff, tolerance)
+        dyn.attach(test_stress, interval=5)
+        dyn.run(20)
 
 
 class TestASEWithPTI(unittest.TestCase):
+    # Tests that the values obtained by wrapping a BuiltinModel or
+    # BuiltinEnsemble with a calculator are the same with and without
+    # periodic_table_index
 
     def setUp(self):
         self.model_pti = torchani.models.ANI1x(periodic_table_index=True).double()

tests/test_data.py

@@ -6,15 +6,13 @@ import unittest
 path = os.path.dirname(os.path.realpath(__file__))
 dataset_path = os.path.join(path, '../dataset/ani-1x/sample.h5')
 batch_size = 256
-ani1x = torchani.models.ANI1x()
-sae_dict = ani1x.sae_dict
-aev_computer = ani1x.aev_computer
+ani1x_sae_dict = {'H': -0.60095298, 'C': -38.08316124, 'N': -54.7077577, 'O': -75.19446356}
 
 
 class TestData(unittest.TestCase):
 
     def testTensorShape(self):
-        ds = torchani.data.load(dataset_path).subtract_self_energies(sae_dict).species_to_indices().shuffle().collate(batch_size).cache()
+        ds = torchani.data.load(dataset_path).subtract_self_energies(ani1x_sae_dict).species_to_indices().shuffle().collate(batch_size).cache()
         for d in ds:
             species = d['species']
             coordinates = d['coordinates']
@@ -28,7 +26,7 @@ class TestData(unittest.TestCase):
             self.assertEqual(coordinates.shape[0], energies.shape[0])
 
     def testNoUnnecessaryPadding(self):
-        ds = torchani.data.load(dataset_path).subtract_self_energies(sae_dict).species_to_indices().shuffle().collate(batch_size).cache()
+        ds = torchani.data.load(dataset_path).subtract_self_energies(ani1x_sae_dict).species_to_indices().shuffle().collate(batch_size).cache()
         for d in ds:
             species = d['species']
             non_padding = (species >= 0)[:, -1].nonzero()
@@ -44,7 +42,7 @@ class TestData(unittest.TestCase):
             entered = True
         self.assertTrue(entered)
 
-        ds = ds.subtract_self_energies(sae_dict)
+        ds = ds.subtract_self_energies(ani1x_sae_dict)
         entered = False
         for d in ds:
             entered = True

tests/test_energies.py

@@ -10,6 +10,8 @@ N = 97
 
 
 class TestEnergies(unittest.TestCase):
+    # tests the predicions for a torchani.nn.Sequential(AEVComputer(),
+    # ANIModel(), EnergyShifter()) against precomputed values
 
     def setUp(self):
         self.tolerance = 5e-5
@@ -17,7 +19,6 @@ class TestEnergies(unittest.TestCase):
         self.aev_computer = model.aev_computer
         self.nnp = model.neural_networks
         self.energy_shifter = model.energy_shifter
-        self.nn = torchani.nn.Sequential(self.nnp, self.energy_shifter)
         self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
 
     def testIsomers(self):
@@ -54,24 +55,25 @@ class TestEnergies(unittest.TestCase):
 
 
 class TestEnergiesEnergyShifterJIT(TestEnergies):
+    # only JIT compile the energy shifter and repeat all tests
 
     def setUp(self):
         super().setUp()
         self.energy_shifter = torch.jit.script(self.energy_shifter)
-        self.nn = torchani.nn.Sequential(self.nnp, self.energy_shifter)
         self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
 
 
 class TestEnergiesANIModelJIT(TestEnergies):
+    # only JIT compile the ANI nnp ANIModel and repeat all tests
 
     def setUp(self):
         super().setUp()
         self.nnp = torch.jit.script(self.nnp)
-        self.nn = torchani.nn.Sequential(self.nnp, self.energy_shifter)
         self.model = torchani.nn.Sequential(self.aev_computer, self.nnp, self.energy_shifter)
 
 
 class TestEnergiesJIT(TestEnergies):
+    # JIT compile the whole model and repeat all tests
 
     def setUp(self):
         super().setUp()

tests/test_ensemble.py

@@ -12,7 +12,6 @@ class TestEnsemble(unittest.TestCase):
 
     def setUp(self):
         self.tol = 1e-5
-        self.conformations = 20
         ani1x = torchani.models.ANI1x()
         self.aev_computer = ani1x.aev_computer
         self.model_iterator = ani1x.neural_networks

tests/test_jit_builtin_models.py

@@ -9,10 +9,13 @@ dspath = os.path.join(path, '../dataset/ani-1x/sample.h5')
 
 
 class TestBuiltinModelsJIT(unittest.TestCase):
+    # Tests if JIT compiled models have the same output energies
+    # as eager (non JIT) models
 
     def setUp(self):
-        self.ani1ccx = torchani.models.ANI1ccx()
-        self.ds = torchani.data.load(dspath).subtract_self_energies(self.ani1ccx.sae_dict).species_to_indices().shuffle().collate(256).cache()
+        # in general self energies should be subtracted, and shuffle should be
+        # performed, but for these tests this is not important
+        self.ds = torchani.data.load(dspath).species_to_indices().collate(256).cache()
 
     def _test_model(self, model):
         properties = next(iter(self.ds))

tools/generate-unit-test-expect/others/Benzene.json

+{"1": {
+ "cell": {"array": {"__ndarray__": [[3, 3], "float64", [7.287, 0.0, 0.0, 0.0, 9.2, 0.0, 0.0, 0.0, 6.688]]}, "pbc": {"__ndarray__": [[3], "bool", [true, true, true]]}, "__ase_objtype__": "cell"},
+ "ctime": 20.492114349528013,
+ "mtime": 20.492114349528013,
+ "numbers": {"__ndarray__": [[48], "int64", [6, 6, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1, 6, 6, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1, 6, 6, 6, 6, 6, 6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1]]},
+ "pbc": {"__ndarray__": [[3], "bool", [true, true, true]]},
+ "positions": {"__ndarray__": [[48, 3], "float64", [6.8956881, 1.3109999999999997, 0.0648736, 4.034811899999999, 7.888999999999999, 3.4088736, 0.3913119, 5.911, 3.2791264, 3.2521881, 3.289, 6.623126399999999, 0.3913119, 7.888999999999999, 6.623126399999999, 3.2521881, 1.3109999999999997, 3.2791264, 6.8956881, 3.289, 3.4088736, 4.034811899999999, 5.911, 0.0648736, 6.667605, 2.2632, 0.227392, 4.262894999999999, 6.9368, 3.571392, 0.619395, 6.863199999999999, 3.116608, 3.024105, 2.3367999999999998, 6.460608, 0.619395, 6.9368, 6.460608, 3.024105, 2.2632, 3.116608, 6.667605, 2.3367999999999998, 3.571392, 4.262894999999999, 6.863199999999999, 0.227392, 0.612108, 0.85008, 0.9182624, 3.031392, 8.34992, 4.2622624, 6.674892, 5.45008, 2.4257376, 4.255608, 3.74992, 5.7697376, 6.674892, 8.34992, 5.7697376, 4.255608, 0.85008, 2.4257376, 0.612108, 3.74992, 4.2622624, 3.031392, 5.45008, 0.9182624, 1.02018, 1.4351999999999998, 1.464672, 2.62332, 7.764799999999999, 4.808672, 6.26682, 6.0352, 1.879328, 4.66368, 3.1647999999999996, 5.2233279999999995, 6.26682, 7.764799999999999, 5.2233279999999995, 4.66368, 1.4351999999999998, 1.879328, 1.02018, 3.1647999999999996, 4.808672, 2.62332, 6.0352, 1.464672, 6.3083559000000005, 0.47931999999999997, 5.862032, 4.6221441, 8.72068, 2.518032, 0.9786441, 5.07932, 4.169967999999999, 2.6648559, 4.12068, 0.8259679999999999, 0.9786441, 8.72068, 0.8259679999999999, 2.6648559, 0.47931999999999997, 4.169967999999999, 6.3083559000000005, 4.12068, 2.518032, 4.6221441, 5.07932, 5.862032, 5.68386, 0.736, 5.323648, 5.24664, 8.464, 1.9796480000000003, 1.60314, 5.335999999999999, 4.708352, 2.04036, 3.8639999999999994, 1.3643519999999998, 1.60314, 8.464, 1.3643519999999998, 2.04036, 0.736, 4.708352, 5.68386, 3.8639999999999994, 1.9796480000000003, 5.24664, 5.335999999999999, 5.323648]]},
+ "unique_id": "518b9a7c9d2506abec880303a9cf7ce8",
+ "user": "ignacio"},
+"ids": [1],
+"nextid": 2}