Use PyTorch's TestCase to compare tensor (#533)

* Use PyTorch's TestCase to compare tensor

* save

* save

* save

* save

* save

* save

* save

* save

* save

* Update test_energies.py

* Update test_ase.py

* save

* flake8

tests/common_aev_test.py

-import unittest
 import torch
 import torchani
 import os
 
-tolerance = 1e-5
 
-
-class _TestAEVBase(unittest.TestCase):
+class _TestAEVBase(torchani.testing.TestCase):
 
     def setUp(self):
         path = os.path.dirname(os.path.realpath(__file__))
@@ -16,15 +13,11 @@ class _TestAEVBase(unittest.TestCase):
         self.radial_length = self.aev_computer.radial_length
         self.debug = False
 
-    def assertAEVEqual(self, expected_radial, expected_angular, aev, tolerance=tolerance):
+    def assertAEVEqual(self, expected_radial, expected_angular, aev):
         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)
+            print(torch.stack([expected_radial[0, aid, :], radial[0, aid, :]]))
+        self.assertEqual(expected_radial, radial)
+        self.assertEqual(expected_angular, angular)

tests/test_aev.py

@@ -16,7 +16,7 @@ const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3
 N = 97
 
 
-class TestAEVConstructor(unittest.TestCase):
+class TestAEVConstructor(torchani.testing.TestCase):
     # Test that checks that the friendly constructor
     # reproduces the values from ANI1x with the correct parameters
     def testCoverLinearly(self):
@@ -35,13 +35,10 @@ class TestAEVConstructor(unittest.TestCase):
         constants = aev_computer.constants()
         constants_alt = aev_computer_alt.constants()
         for c, ca in zip(constants, constants_alt):
-            if isinstance(c, torch.Tensor):
-                self.assertTrue(torch.isclose(c, ca).all())
-            else:
-                self.assertEqual(c, ca)
+            self.assertEqual(c, ca)
 
 
-class TestIsolated(unittest.TestCase):
+class TestIsolated(torchani.testing.TestCase):
     # 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
@@ -153,7 +150,7 @@ class TestAEVJIT(TestAEV):
         self.aev_computer = torch.jit.script(self.aev_computer)
 
 
-class TestPBCSeeEachOther(unittest.TestCase):
+class TestPBCSeeEachOther(torchani.testing.TestCase):
     def setUp(self):
         consts = torchani.neurochem.Constants(const_file)
         self.aev_computer = torchani.AEVComputer(**consts).to(torch.double)
@@ -175,7 +172,7 @@ class TestPBCSeeEachOther(unittest.TestCase):
         for _ in range(100):
             translation = torch.randn(3, dtype=torch.double)
             _, aev2 = self.aev_computer((species, coordinates + translation), cell=cell, pbc=pbc)
-            self.assertTrue(torch.allclose(aev, aev2))
+            self.assertEqual(aev, aev2)
 
     def testPBCConnersSeeEachOther(self):
         species = torch.tensor([[0, 0]])
@@ -257,7 +254,7 @@ class TestPBCSeeEachOther(unittest.TestCase):
         self.assertEqual(atom_index2.tolist(), [1])
 
 
-class TestAEVOnBoundary(unittest.TestCase):
+class TestAEVOnBoundary(torchani.testing.TestCase):
 
     def setUp(self):
         self.eps = 1e-9
@@ -280,13 +277,13 @@ class TestAEVOnBoundary(unittest.TestCase):
 
     def assertInCell(self, coordinates):
         coordinates_cell = coordinates @ self.inv_cell
-        self.assertTrue(torch.allclose(coordinates, coordinates_cell @ self.cell))
+        self.assertEqual(coordinates, coordinates_cell @ self.cell)
         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
-        self.assertTrue(torch.allclose(coordinates, coordinates_cell @ self.cell))
+        self.assertEqual(coordinates, coordinates_cell @ self.cell)
         in_cell = (coordinates_cell >= -self.eps) & (coordinates_cell <= 1 + self.eps)
         self.assertFalse(in_cell.all())
 
@@ -300,10 +297,10 @@ class TestAEVOnBoundary(unittest.TestCase):
             self.assertInCell(coordinates)
             _, aev = self.aev_computer((self.species, coordinates), cell=self.cell, pbc=self.pbc)
             self.assertGreater(aev.abs().max().item(), 0)
-            self.assertTrue(torch.allclose(aev, self.aev))
+            self.assertEqual(aev, self.aev)
 
 
-class TestAEVOnBenzenePBC(unittest.TestCase):
+class TestAEVOnBenzenePBC(torchani.testing.TestCase):
 
     def setUp(self):
         consts = torchani.neurochem.Constants(const_file)
@@ -319,7 +316,6 @@ class TestAEVOnBenzenePBC(unittest.TestCase):
         self.natoms = self.aev.shape[1]
 
     def testRepeat(self):
-        tolerance = 5e-6
         c1, c2, c3 = self.cell
         species2 = self.species.repeat(1, 4)
         coordinates2 = torch.cat([
@@ -332,7 +328,7 @@ class TestAEVOnBenzenePBC(unittest.TestCase):
         _, aev2 = self.aev_computer((species2, coordinates2), cell=cell2, pbc=self.pbc)
         for i in range(3):
             aev3 = aev2[:, i * self.natoms: (i + 1) * self.natoms, :]
-            self.assertTrue(torch.allclose(self.aev, aev3, atol=tolerance))
+            self.assertEqual(self.aev, aev3)
 
     def testManualMirror(self):
         c1, c2, c3 = self.cell
@@ -343,7 +339,7 @@ class TestAEVOnBenzenePBC(unittest.TestCase):
         ], dim=1)
         _, aev2 = self.aev_computer((species2, coordinates2))
         aev2 = aev2[:, :self.natoms, :]
-        self.assertTrue(torch.allclose(self.aev, aev2))
+        self.assertEqual(self.aev, aev2)
 
 
 if __name__ == '__main__':

tests/test_aev_benzene_md.py

@@ -24,7 +24,7 @@ class TestAEVBenzeneMD(_TestAEVBase):
                 pbc = torch.from_numpy(pbc)
                 coordinates = torchani.utils.map2central(cell, coordinates, pbc)
                 _, aev = self.aev_computer((species, coordinates), cell=cell, pbc=pbc)
-                self.assertAEVEqual(expected_radial, expected_angular, aev, 5e-5)
+                self.assertAEVEqual(expected_radial, expected_angular, aev)
 
 
 if __name__ == '__main__':

tests/test_aev_tripeptide_md.py

@@ -10,7 +10,6 @@ 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:
@@ -21,7 +20,7 @@ class TestAEVTripeptideMD(_TestAEVBase):
                 expected_radial = torch.from_numpy(expected_radial).float().unsqueeze(0)
                 expected_angular = torch.from_numpy(expected_angular).float().unsqueeze(0)
                 _, aev = self.aev_computer((species, coordinates))
-                self.assertAEVEqual(expected_radial, expected_angular, aev, tol)
+                self.assertAEVEqual(expected_radial, expected_angular, aev)
 
 
 if __name__ == '__main__':

tests/test_al.py

@@ -4,7 +4,7 @@ import math
 import unittest
 
 
-class TestALAtomic(unittest.TestCase):
+class TestALAtomic(torchani.testing.TestCase):
     def setUp(self):
         self.device = torch.device(
             'cuda' if torch.cuda.is_available() else 'cpu')
@@ -26,7 +26,7 @@ class TestALAtomic(unittest.TestCase):
     def testAverageAtomicEnergies(self):
         _, energies = self.model.atomic_energies(
             (self.species, self.coordinates))
-        self.assertTrue(energies.shape == self.coordinates.shape[:-1])
+        self.assertEqual(energies.shape, self.coordinates.shape[:-1])
         # energies of all hydrogens should be equal
         self.assertTrue((torch.isclose(
             energies[:, :-1],
@@ -57,11 +57,11 @@ class TestALQBC(TestALAtomic):
 
         # correctness of shape
         torch.set_printoptions(precision=15)
-        self.assertTrue(energies.shape[-1] == self.coordinates.shape[0])
-        self.assertTrue(energies.shape[0] == len(self.model.neural_networks))
-        self.assertTrue(
-            energies[0] == self.first_model((self.species,
-                                             self.coordinates)).energies)
+        self.assertEqual(energies.shape[-1], self.coordinates.shape[0])
+        self.assertEqual(energies.shape[0], len(self.model.neural_networks))
+        self.assertEqual(
+            energies[0], self.first_model((self.species,
+                                           self.coordinates)).energies)
         self.assertTrue(
             torch.isclose(
                 energies[0],
@@ -96,7 +96,7 @@ class TestALQBC(TestALAtomic):
         _, _, qbc = self.model.energies_qbcs((species, coordinates))
         std[0] = std[0] / math.sqrt(5)
         std[1] = std[1] / math.sqrt(4)
-        self.assertTrue(torch.isclose(std, qbc).all())
+        self.assertEqual(std, qbc)
 
 
 if __name__ == '__main__':

tests/test_ase.py

@@ -21,7 +21,7 @@ def get_numeric_force(atoms, eps):
     return fn
 
 
-class TestASE(unittest.TestCase):
+class TestASE(torchani.testing.TestCase):
 
     def setUp(self):
         self.model = torchani.models.ANI1x(model_index=0).double()
@@ -30,24 +30,19 @@ class TestASE(unittest.TestCase):
         # 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)
         dyn = Langevin(atoms, 5 * units.fs, 30000000 * units.kB, 0.002)
         dyn.run(100)
-        f = torch.from_numpy(atoms.get_forces())
+        f = atoms.get_forces()
         fn = get_numeric_force(atoms, 0.001)
-        df = (f - fn).abs().max()
-        avgf = f.abs().mean()
-        if avgf > 0:
-            self.assertLess(df / avgf, relative_tolerance)
+        self.assertEqual(f, fn, rtol=0.1, atol=0)
 
     def testWithNumericalStressWithPBCEnabled(self):
         # 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
@@ -66,10 +61,9 @@ class TestASE(unittest.TestCase):
         def test_stress():
             stress = benzene.get_stress()
             numerical_stress = calculator.calculate_numerical_stress(benzene)
-            diff = torch.from_numpy(stress - numerical_stress).abs().max().item()
-            self.assertLess(diff, tolerance)
-        dyn.attach(test_stress, interval=5)
-        dyn.run(20)
+            self.assertEqual(stress, numerical_stress)
+        dyn.attach(test_stress, interval=30)
+        dyn.run(120)
 
 
 class TestASEWithPTI(unittest.TestCase):

tests/test_cuaev.py

@@ -7,7 +7,7 @@ skipIfNoGPU = unittest.skipIf(not torch.cuda.is_available(),
 
 
 @unittest.skipIf(not torchani.has_cuaev, "only valid when cuaev is installed")
-class TestCUAEV(unittest.TestCase):
+class TestCUAEV(torchani.testing.TestCase):
 
     def testJIT(self):
         def f(coordinates, species, Rcr: float, Rca: float, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species: int):

tests/test_data.py

@@ -9,7 +9,7 @@ batch_size = 256
 ani1x_sae_dict = {'H': -0.60095298, 'C': -38.08316124, 'N': -54.7077577, 'O': -75.19446356}
 
 
-class TestData(unittest.TestCase):
+class TestData(torchani.testing.TestCase):
 
     def testTensorShape(self):
         ds = torchani.data.load(dataset_path).subtract_self_energies(ani1x_sae_dict).species_to_indices().shuffle().collate(batch_size).cache()
@@ -35,7 +35,7 @@ class TestData(unittest.TestCase):
     def testReEnter(self):
         # make sure that a dataset can be iterated multiple times
         ds = torchani.data.load(dataset_path)
-        for d in ds:
+        for _ in ds:
             pass
         entered = False
         for d in ds:
@@ -107,22 +107,19 @@ class TestData(unittest.TestCase):
         len(ds)
 
     def testSAE(self):
-        tolerance = 1e-5
         shifter = torchani.EnergyShifter(None)
         torchani.data.load(dataset_path).subtract_self_energies(shifter)
         true_self_energies = torch.tensor([-19.354171758844188,
                                            -19.354171758844046,
                                            -54.712238523648587,
                                            -75.162829556770987], dtype=torch.float64)
-        diff = torch.abs(true_self_energies - shifter.self_energies)
-        for e in diff:
-            self.assertLess(e, tolerance)
+        self.assertEqual(true_self_energies, shifter.self_energies)
 
     def testDataloader(self):
         shifter = torchani.EnergyShifter(None)
         dataset = list(torchani.data.load(dataset_path).subtract_self_energies(shifter).species_to_indices().shuffle())
         loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, collate_fn=torchani.data.collate_fn, num_workers=64)
-        for i in loader:
+        for _ in loader:
             pass
 
 

tests/test_energies.py

@@ -9,12 +9,11 @@ path = os.path.dirname(os.path.realpath(__file__))
 N = 97
 
 
-class TestEnergies(unittest.TestCase):
+class TestEnergies(torchani.testing.TestCase):
     # tests the predicions for a torchani.nn.Sequential(AEVComputer(),
     # ANIModel(), EnergyShifter()) against precomputed values
 
     def setUp(self):
-        self.tolerance = 5e-5
         model = torchani.models.ANI1x(model_index=0)
         self.aev_computer = model.aev_computer
         self.nnp = model.neural_networks
@@ -30,8 +29,7 @@ class TestEnergies(unittest.TestCase):
                 species = torch.from_numpy(species)
                 energies = torch.from_numpy(energies).to(torch.float)
                 energies_ = self.model((species, coordinates)).energies
-                max_diff = (energies - energies_).abs().max().item()
-                self.assertLess(max_diff, self.tolerance)
+                self.assertEqual(energies, energies_, exact_dtype=False)
 
     def testPadding(self):
         species_coordinates = []
@@ -50,8 +48,7 @@ class TestEnergies(unittest.TestCase):
             species_coordinates)
         energies = torch.cat(energies)
         energies_ = self.model((species_coordinates['species'], species_coordinates['coordinates'])).energies
-        max_diff = (energies - energies_).abs().max().item()
-        self.assertLess(max_diff, self.tolerance)
+        self.assertEqual(energies, energies_, exact_dtype=False)
 
 
 class TestEnergiesEnergyShifterJIT(TestEnergies):

tests/test_ensemble.py

@@ -8,10 +8,10 @@ path = os.path.dirname(os.path.realpath(__file__))
 N = 10
 
 
-class TestEnsemble(unittest.TestCase):
+class TestEnsemble(torchani.testing.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
@@ -25,10 +25,8 @@ class TestEnsemble(unittest.TestCase):
         energy2 = [m((species, coordinates))[1] for m in model_list]
         energy2 = sum(energy2) / len(model_list)
         force2 = torch.autograd.grad(energy2.sum(), coordinates)[0]
-        energy_diff = (energy1 - energy2).abs().max().item()
-        force_diff = (force1 - force2).abs().max().item()
-        self.assertLess(energy_diff, self.tol)
-        self.assertLess(force_diff, self.tol)
+        self.assertEqual(energy1, energy2)
+        self.assertEqual(force1, force2)
 
     def testGDB(self):
         for i in range(N):

tests/test_forces.py

@@ -8,10 +8,9 @@ path = os.path.dirname(os.path.realpath(__file__))
 N = 97
 
 
-class TestForce(unittest.TestCase):
+class TestForce(torchani.testing.TestCase):
 
     def setUp(self):
-        self.tolerance = 1e-5
         model = torchani.models.ANI1x(model_index=0)
         self.aev_computer = model.aev_computer
         self.nnp = model.neural_networks
@@ -29,8 +28,7 @@ class TestForce(unittest.TestCase):
                 _, 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)
+                self.assertEqual(forces, -derivative)
 
     def testPadding(self):
         species_coordinates = []
@@ -52,8 +50,7 @@ class TestForce(unittest.TestCase):
         for coordinates, forces in coordinates_forces:
             derivative = torch.autograd.grad(energies, coordinates,
                                              retain_graph=True)[0]
-            max_diff = (forces + derivative).abs().max().item()
-            self.assertLess(max_diff, self.tolerance)
+            self.assertEqual(forces, -derivative)
 
 
 class TestForceJIT(TestForce):

tests/test_grad.py

@@ -7,7 +7,7 @@ import pickle
 path = os.path.dirname(os.path.realpath(__file__))
 
 
-class TestGrad(unittest.TestCase):
+class TestGrad(torchani.testing.TestCase):
     # torch.autograd.gradcheck and torch.autograd.gradgradcheck verify that
     # the numerical and analytical gradient and hessian of a function
     # matches to within a given tolerance.

tests/test_jit_builtin_models.py

@@ -8,7 +8,7 @@ path = os.path.dirname(os.path.realpath(__file__))
 dspath = os.path.join(path, '../dataset/ani-1x/sample.h5')
 
 
-class TestBuiltinModelsJIT(unittest.TestCase):
+class TestBuiltinModelsJIT(torchani.testing.TestCase):
     # Tests if JIT compiled models have the same output energies
     # as eager (non JIT) models
 
@@ -22,7 +22,7 @@ class TestBuiltinModelsJIT(unittest.TestCase):
         input_ = (properties['species'], properties['coordinates'].float())
         _, e = model(input_)
         _, e2 = torch.jit.script(model)(input_)
-        self.assertTrue(torch.allclose(e, e2))
+        self.assertEqual(e, e2)
 
     def _test_ensemble(self, ensemble):
         self._test_model(ensemble)

tests/test_neurochem.py

@@ -9,7 +9,7 @@ iptpath = os.path.join(path, 'test_data/inputtrain.ipt')
 dspath = os.path.join(path, '../dataset/ani1-up_to_gdb4/ani_gdb_s01.h5')
 
 
-class TestNeuroChem(unittest.TestCase):
+class TestNeuroChem(torchani.testing.TestCase):
 
     def testNeuroChemTrainer(self):
         d = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

tests/test_padding.py

@@ -6,7 +6,7 @@ import torchani
 b = torchani.utils.broadcast_first_dim
 
 
-class TestPaddings(unittest.TestCase):
+class TestPaddings(torchani.testing.TestCase):
 
     def testVectorSpecies(self):
         species1 = torch.tensor([[0, 2, 3, 1]])
@@ -28,7 +28,7 @@ class TestPaddings(unittest.TestCase):
             [3, 2, 0, 1, 0],
             [3, 2, 0, 1, 0],
         ])
-        self.assertEqual((atomic_properties['species'] - expected_species).abs().max().item(), 0)
+        self.assertEqual(atomic_properties['species'], expected_species)
         self.assertEqual(atomic_properties['coordinates'].abs().max().item(), 0)
 
     def testTensorShape1NSpecies(self):
@@ -51,7 +51,7 @@ class TestPaddings(unittest.TestCase):
             [3, 2, 0, 1, 0],
             [3, 2, 0, 1, 0],
         ])
-        self.assertEqual((atomic_properties['species'] - expected_species).abs().max().item(), 0)
+        self.assertEqual(atomic_properties['species'], expected_species)
         self.assertEqual(atomic_properties['coordinates'].abs().max().item(), 0)
 
     def testTensorSpecies(self):
@@ -80,20 +80,17 @@ class TestPaddings(unittest.TestCase):
             [3, 2, 0, 1, 0],
             [3, 2, 0, 1, 0],
         ])
-        self.assertEqual((atomic_properties['species'] - expected_species).abs().max().item(), 0)
+        self.assertEqual(atomic_properties['species'], expected_species)
         self.assertEqual(atomic_properties['coordinates'].abs().max().item(), 0)
 
     def testPresentSpecies(self):
         species = torch.tensor([0, 1, 1, 0, 3, 7, -1, -1])
         present_species = torchani.utils.present_species(species)
         expected = torch.tensor([0, 1, 3, 7])
-        self.assertEqual((expected - present_species).abs().max().item(), 0)
+        self.assertEqual(expected, present_species)
 
 
-class TestStripRedundantPadding(unittest.TestCase):
-
-    def _assertTensorEqual(self, t1, t2):
-        self.assertEqual((t1 - t2).abs().max().item(), 0)
+class TestStripRedundantPadding(torchani.testing.TestCase):
 
     def testStripRestore(self):
         species1 = torch.randint(4, (5, 4), dtype=torch.long)
@@ -119,14 +116,14 @@ class TestStripRedundantPadding(unittest.TestCase):
             b({'species': species123[:5, ...], 'coordinates': coordinates123[:5, ...]}))
         species1_ = species_coordinates1_['species']
         coordinates1_ = species_coordinates1_['coordinates']
-        self._assertTensorEqual(species1_, species1)
-        self._assertTensorEqual(coordinates1_, coordinates1)
+        self.assertEqual(species1_, species1)
+        self.assertEqual(coordinates1_, coordinates1)
         species_coordinates12_ = torchani.utils.strip_redundant_padding(
             b({'species': species123[:7, ...], 'coordinates': coordinates123[:7, ...]}))
         species12_ = species_coordinates12_['species']
         coordinates12_ = species_coordinates12_['coordinates']
-        self._assertTensorEqual(species12_, species12)
-        self._assertTensorEqual(coordinates12_, coordinates12)
+        self.assertEqual(species12_, species12)
+        self.assertEqual(coordinates12_, coordinates12)
 
 
 if __name__ == '__main__':

tests/test_periodic_table_indexing.py

@@ -3,7 +3,7 @@ import torch
 import torchani
 
 
-class TestSpeciesConverter(unittest.TestCase):
+class TestSpeciesConverter(torchani.testing.TestCase):
 
     def setUp(self):
         self.c = torchani.SpeciesConverter(['H', 'C', 'N', 'O'])
@@ -19,7 +19,7 @@ class TestSpeciesConverter(unittest.TestCase):
         ], dtype=torch.long)
         dummy_coordinates = torch.empty(2, 5, 3)
         output = self.c((input_, dummy_coordinates)).species
-        self.assertTrue(torch.allclose(output, expect))
+        self.assertEqual(output, expect)
 
 
 class TestSpeciesConverterJIT(TestSpeciesConverter):
@@ -29,7 +29,7 @@ class TestSpeciesConverterJIT(TestSpeciesConverter):
         self.c = torch.jit.script(self.c)
 
 
-class TestBuiltinEnsemblePeriodicTableIndex(unittest.TestCase):
+class TestBuiltinEnsemblePeriodicTableIndex(torchani.testing.TestCase):
 
     def setUp(self):
         self.model1 = torchani.models.ANI1x()
@@ -48,16 +48,16 @@ class TestBuiltinEnsemblePeriodicTableIndex(unittest.TestCase):
         energy2 = self.model2((self.species2, self.coordinates)).energies
         derivative1 = torch.autograd.grad(energy1.sum(), self.coordinates)[0]
         derivative2 = torch.autograd.grad(energy2.sum(), self.coordinates)[0]
-        self.assertTrue(torch.allclose(energy1, energy2))
-        self.assertTrue(torch.allclose(derivative1, derivative2))
+        self.assertEqual(energy1, energy2)
+        self.assertEqual(derivative1, derivative2)
 
     def testCH4Single(self):
         energy1 = self.model1[0]((self.species1, self.coordinates)).energies
         energy2 = self.model2[0]((self.species2, self.coordinates)).energies
         derivative1 = torch.autograd.grad(energy1.sum(), self.coordinates)[0]
         derivative2 = torch.autograd.grad(energy2.sum(), self.coordinates)[0]
-        self.assertTrue(torch.allclose(energy1, energy2))
-        self.assertTrue(torch.allclose(derivative1, derivative2))
+        self.assertEqual(energy1, energy2)
+        self.assertEqual(derivative1, derivative2)
 
 
 if __name__ == '__main__':

tests/test_structure_optim.py

@@ -11,7 +11,7 @@ from ase import Atoms
 path = os.path.dirname(os.path.realpath(__file__))
 
 
-class TestStructureOptimization(unittest.TestCase):
+class TestStructureOptimization(torchani.testing.TestCase):
 
     def setUp(self):
         self.tolerance = 1e-6
@@ -31,10 +31,7 @@ class TestStructureOptimization(unittest.TestCase):
                 opt = BFGS(atoms)
                 opt.run()
                 coordinates = atoms.get_positions()
-                coordinates = torch.from_numpy(coordinates)
-                distances = (old_coordinates - coordinates).norm(dim=1)
-                rmse = distances.mean()
-                self.assertLess(rmse, self.tolerance)
+                self.assertEqual(old_coordinates, coordinates)
 
 
 if __name__ == '__main__':

tests/test_utils.py

@@ -3,7 +3,7 @@ import torch
 import torchani
 
 
-class TestUtils(unittest.TestCase):
+class TestUtils(torchani.testing.TestCase):
 
     def testChemicalSymbolsToInts(self):
         str2i = torchani.utils.ChemicalSymbolsToInts(['A', 'B', 'C', 'D', 'E', 'F'])

tests/test_vibrational.py

@@ -13,7 +13,7 @@ path = os.path.dirname(os.path.realpath(__file__))
 path = os.path.join(path, '../dataset/xyz_files/H2O.xyz')
 
 
-class TestVibrational(unittest.TestCase):
+class TestVibrational(torchani.testing.TestCase):
 
     def testVibrationalWavenumbers(self):
         model = torchani.models.ANI1x().double()
@@ -44,8 +44,7 @@ class TestVibrational(unittest.TestCase):
         freq2, modes2, _, _ = torchani.utils.vibrational_analysis(masses[species], hessian)
         freq2 = freq2[6:].float()
         modes2 = modes2[6:]
-        ratio = freq2 / freq
-        self.assertLess((ratio - 1).abs().max(), 0.02)
+        self.assertEqual(freq, freq2, atol=0, rtol=0.02, exact_dtype=False)
 
         diff1 = (modes - modes2).abs().max(dim=-1).values.max(dim=-1).values
         diff2 = (modes + modes2).abs().max(dim=-1).values.max(dim=-1).values

torchani/__init__.py

@@ -39,6 +39,7 @@ from . import units
 from pkg_resources import get_distribution, DistributionNotFound
 import warnings
 import importlib_metadata
+from . import testing
 
 has_cuaev = 'torchani.cuaev' in importlib_metadata.metadata(__package__).get_all('Provides')
 
@@ -55,7 +56,7 @@ except DistributionNotFound:
     pass
 
 __all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble', 'SpeciesConverter',
-           'utils', 'neurochem', 'models', 'units', 'has_cuaev']
+           'utils', 'neurochem', 'models', 'units', 'has_cuaev', 'testing']
 
 try:
     from . import ase  # noqa: F401