rewrite unittests to completely remove dependency on NeuroChem (#2)

rewrite unittests to completely remove dependency on NeuroChem

tests/test_energies.py

+import torch
+import numpy
+import torchani
+import unittest
+import os
+import pickle
+
+
+path = os.path.dirname(os.path.realpath(__file__))
+N = 97
+
+
+class TestEnergies(unittest.TestCase):
+
+    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
+        self.tolerance = 5e-5
+        self.aev_computer = torchani.SortedAEV(
+            dtype=dtype, device=torch.device('cpu'))
+        self.nnp = torchani.ModelOnAEV(self.aev_computer, from_nc=None)
+
+    def _test_molecule(self, coordinates, species, energies):
+        shift_energy = torchani.EnergyShifter(torchani.buildin_sae_file)
+        energies_ = self.nnp(coordinates, species).squeeze()
+        energies_ = shift_energy.add_sae(energies_, species)
+        max_diff = (energies - energies_).abs().max().item()
+        self.assertLess(max_diff, self.tolerance)
+
+    def testGDB(self):
+        for i in range(N):
+            datafile = os.path.join(path, 'test_data/{}'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, energies, _ = pickle.load(f)
+                self._test_molecule(coordinates, species, energies)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_forces.py

+import torch
+import numpy
+import torchani
+import unittest
+import os
+import pickle
+
+path = os.path.dirname(os.path.realpath(__file__))
+N = 97
+
+
+class TestForce(unittest.TestCase):
+
+    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
+        self.tolerance = 1e-5
+        self.aev_computer = torchani.SortedAEV(
+            dtype=dtype, device=torch.device('cpu'))
+        self.nnp = torchani.ModelOnAEV(
+            self.aev_computer, derivative=True, from_nc=None)
+
+    def _test_molecule(self, coordinates, species, forces):
+        _, derivative = self.nnp(coordinates, species)
+        max_diff = (forces + derivative).abs().max().item()
+        self.assertLess(max_diff, self.tolerance)
+
+    def testGDB(self):
+        for i in range(N):
+            datafile = os.path.join(path, 'test_data/{}'.format(i))
+            with open(datafile, 'rb') as f:
+                coordinates, species, _, _, _, forces = pickle.load(f)
+                self._test_molecule(coordinates, species, forces)
+
+
+if __name__ == '__main__':
+    unittest.main()

torchani/neurochem_aev.py → tools/generate-unit-test-expect.py

 import torch
+import os
 import ase
 import pyNeuroChem
 import ase_interface
 import numpy
-from .aev_base import AEVComputer
-from . import buildin_const_file, buildin_sae_file, buildin_network_dir, default_dtype, default_device
+import torchani
+import pickle
+from torchani import buildin_const_file, buildin_sae_file, buildin_network_dir, default_dtype, default_device
+import torchani.pyanitools
 
+path = os.path.dirname(os.path.realpath(__file__))
+conv_au_ev = 27.21138505
 
-class NeuroChemAEV (AEVComputer):
-    """The AEV computer that dump out AEV from pyNeuroChem
 
-    Attributes
-    ----------
-    sae_file : str
-        The name of the original file that stores self atomic energies.
-    network_dir : str
-        The name ending with '/' of the directory that stores networks in NeuroChem's format.
-    nc : pyNeuroChem.molecule
-        The internal object of pyNeuroChem which can be used to dump out AEVs, energies, forces,
-        activations, etc.
-    """
+class NeuroChem (torchani.aev_base.AEVComputer):
 
     def __init__(self, dtype=default_dtype, device=default_device, const_file=buildin_const_file, sae_file=buildin_sae_file, network_dir=buildin_network_dir):
-        super(NeuroChemAEV, self).__init__(False, dtype, device, const_file)
+        super(NeuroChem, self).__init__(False, dtype, device, const_file)
         self.sae_file = sae_file
         self.network_dir = network_dir
-        self.nc = pyNeuroChem.molecule(const_file, sae_file, network_dir, 0)
+        self.nc = pyNeuroChem.molecule(
+            self.const_file, self.sae_file, self.network_dir, 0)
 
     def _get_radial_part(self, fullaev):
-        """Get the radial part of AEV from the full AEV
-
-        Parameters
-        ----------
-        fullaev : pytorch tensor of `dtype`
-            The full AEV in shape (conformations, atoms, `radial_length()+angular_length()`).
-
-        Returns
-        -------
-        pytorch tensor of `dtype`
-            The radial AEV in shape(conformations, atoms, `radial_length()`)
-        """
         radial_size = self.radial_length
         return fullaev[:, :, :radial_size]
 
     def _get_angular_part(self, fullaev):
-        """Get the angular part of AEV from the full AEV
-
-        Parameters
-        ----------
-        fullaev : pytorch tensor of `dtype`
-            The full AEV in shape (conformations, atoms, `radial_length()+angular_length()`).
-
-        Returns
-        -------
-        pytorch tensor of `dtype`
-            The radial AEV in shape (conformations, atoms, `angular_length()`)
-        """
         radial_size = self.radial_length
         return fullaev[:, :, radial_size:]
 
-    def _compute_neurochem_aevs_per_conformation(self, coordinates, species):
-        """Get the full AEV for a single conformation
-
-        Parameters
-        ----------
-        coordinates : pytorch tensor of `dtype`
-            The xyz coordinates in shape (atoms, 3).
-        species : list of str
-            The list specifying the species of each atom. The length of the
-            list must be the same as the number of atoms.
-
-        Returns
-        -------
-        pytorch tensor of `dtype`
-            The full AEV for all atoms in shape (atoms, `radial_length()+angular_length()`)
-        """
+    def _per_conformation(self, coordinates, species):
         atoms = coordinates.shape[0]
         mol = ase.Atoms(''.join(species), positions=coordinates)
         mol.set_calculator(ase_interface.ANI(False))
         mol.calc.setnc(self.nc)
-        _ = mol.get_potential_energy()
-        aevs = [self.nc.atomicenvironments(j) for j in range(atoms)]
+        energy = mol.get_potential_energy() / conv_au_ev
+        aevs = [mol.calc.nc.atomicenvironments(j) for j in range(atoms)]
+        force = mol.get_forces() / conv_au_ev
         aevs = numpy.stack(aevs)
-        return aevs
+        return aevs, energy, force
 
-    def __call__(self, coordinates, species):
+    def forward(self, coordinates, species):
         conformations = coordinates.shape[0]
-        aevs = [self._compute_neurochem_aevs_per_conformation(
+        results = [self._per_conformation(
             coordinates[i], species) for i in range(conformations)]
+        aevs, energies, forces = zip(*results)
         aevs = torch.from_numpy(numpy.stack(aevs)).type(
             self.dtype).to(self.device)
-        return self._get_radial_part(aevs), self._get_angular_part(aevs)
+        energies = torch.from_numpy(numpy.stack(energies)).type(
+            self.dtype).to(self.device)
+        forces = torch.from_numpy(numpy.stack(forces)).type(
+            self.dtype).to(self.device)
+        return self._get_radial_part(aevs), self._get_angular_part(aevs), energies, forces
+
+
+aev = torchani.SortedAEV(device=torch.device('cpu'))
+ncaev = NeuroChem(device=torch.device('cpu'))
+mol_count = 0
+
+for i in [1, 2, 3, 4]:
+    data_file = os.path.join(
+        path, '../tests/dataset/ani_gdb_s0{}.h5'.format(i))
+    adl = torchani.pyanitools.anidataloader(data_file)
+    for data in adl:
+        coordinates = data['coordinates'][:10, :]
+        coordinates = torch.from_numpy(coordinates).type(aev.dtype)
+        species = data['species']
+        coordinates, species = aev.sort_by_species(coordinates, species)
+        smiles = ''.join(data['smiles'])
+        radial, angular, energies, forces = ncaev(coordinates, species)
+        pickleobj = (coordinates, species, radial, angular, energies, forces)
+        dumpfile = os.path.join(
+            path, '../tests/test_data/{}'.format(mol_count))
+        with open(dumpfile, 'wb') as f:
+            pickle.dump(pickleobj, f)
+        mol_count += 1

torchani/__init__.py

@@ -13,9 +13,6 @@ buildin_network_dir = pkg_resources.resource_filename(
 buildin_model_prefix = pkg_resources.resource_filename(
     __name__, 'resources/ani-1x_dft_x8ens/train')
 
-buildin_dataset_dir = pkg_resources.resource_filename(
-    __name__, 'resources/')
-
 default_dtype = torch.float32
 default_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -27,10 +24,3 @@ import logging
 __all__ = ['SortedAEV', 'EnergyShifter', 'ModelOnAEV', 'PerSpeciesFromNeuroChem', 'data',
            'buildin_const_file', 'buildin_sae_file', 'buildin_network_dir', 'buildin_dataset_dir',
            'default_dtype', 'default_device']
-
-try:
-    from .neurochem_aev import NeuroChemAEV
-    __all__.append('NeuroChemAEV')
-except ImportError:
-    logging.log(logging.WARNING,
-                'Unable to import NeuroChemAEV, please check your pyNeuroChem installation.')

torchani/aev.py

@@ -92,7 +92,8 @@ class SortedAEV(AEVComputer):
         torch.Tensor
             A tensor of shape (..., neighbors, `radial_sublength`) storing the subAEVs.
         """
-        distances = distances.unsqueeze(-1).unsqueeze(-1) #TODO: allow unsqueeze to insert multiple dimensions
+        distances = distances.unsqueeze(
+            -1).unsqueeze(-1)  # TODO: allow unsqueeze to insert multiple dimensions
         fc = _cutoff_cosine(distances, self.Rcr)
         # Note that in the equation in the paper there is no 0.25 coefficient, but in NeuroChem there is such a coefficient. We choose to be consistent with NeuroChem instead of the paper here.
         ret = 0.25 * torch.exp(-self.EtaR * (distances - self.ShfR)**2) * fc
@@ -118,9 +119,9 @@ class SortedAEV(AEVComputer):
             Tensor of shape (..., pairs, `angular_sublength`) storing the subAEVs.
         """
         vectors1 = vectors1.unsqueeze(
-            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) #TODO: allow unsqueeze to plug in multiple dims
+            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)  # TODO: allow unsqueeze to plug in multiple dims
         vectors2 = vectors2.unsqueeze(
-            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) #TODO: allow unsqueeze to plug in multiple dims
+            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)  # TODO: allow unsqueeze to plug in multiple dims
         distances1 = vectors1.norm(2, dim=-5)
         distances2 = vectors2.norm(2, dim=-5)
 
@@ -179,7 +180,8 @@ class SortedAEV(AEVComputer):
         distances, indices = distances.sort(-1)
 
         min_distances, _ = distances.flatten(end_dim=1).min(0)
-        inRcr = (min_distances <= self.Rcr).nonzero().flatten()[1:]  #TODO: can we use something like find_first?
+        inRcr = (min_distances <= self.Rcr).nonzero().flatten()[
+            1:]  # TODO: can we use something like find_first?
         inRca = (min_distances <= self.Rca).nonzero().flatten()[1:]
 
         distances = distances.index_select(-1, inRcr)
@@ -188,9 +190,12 @@ class SortedAEV(AEVComputer):
 
         indices_a = indices.index_select(-1, inRca)
         new_shape = list(indices_a.shape) + [3]
-        _indices_a = indices_a.unsqueeze(-1).expand(*new_shape)  #TODO: can we add something like expand_dim(dim=0, repeat=3)
-        vec = vec.gather(-2, _indices_a)  #TODO: can we make gather broadcast??
-        vec = self.combinations(vec, -2)  #TODO: can we move combinations to ATen?
+        # TODO: can we add something like expand_dim(dim=0, repeat=3)
+        _indices_a = indices_a.unsqueeze(-1).expand(*new_shape)
+        # TODO: can we make gather broadcast??
+        vec = vec.gather(-2, _indices_a)
+        # TODO: can we move combinations to ATen?
+        vec = self.combinations(vec, -2)
         angular_terms = self.angular_subaev_terms(
             *vec) if vec is not None else None
 
@@ -227,7 +232,6 @@ class SortedAEV(AEVComputer):
                   torch.arange(len(self.species), device=self.device))
         return mask_r
 
-
     def compute_mask_a(self, species_a, present_species):
         """Partition indices according to their species, angular part
 
@@ -299,12 +303,12 @@ class SortedAEV(AEVComputer):
         radial_aevs = present_radial_aevs.flatten(start_dim=2)
 
         # assemble angular subaev
-        rev_indices = {present_species[i].item(): i  #TODO: can we use find_first?
+        rev_indices = {present_species[i].item(): i  # TODO: can we use find_first?
                        for i in range(len(present_species))}
         """Tensor of shape (conformations, atoms, present species, present species, angular_length)"""
         angular_aevs = []
-        zero_angular_subaev = torch.zeros(  #TODO: can we make stack and cat broadcast?
-            conformations, atoms, self.angular_sublength, dtype=self.dtype, device=self.device)  #TODO: can we make torch.zeros, torch.ones typeless and deviceless?
+        zero_angular_subaev = torch.zeros(  # TODO: can we make stack and cat broadcast?
+            conformations, atoms, self.angular_sublength, dtype=self.dtype, device=self.device)  # TODO: can we make torch.zeros, torch.ones typeless and deviceless?
         for s1, s2 in itertools.combinations_with_replacement(range(len(self.species)), 2):
             # TODO: can we remove this if pytorch support 0 size tensors?
             if s1 in rev_indices and s2 in rev_indices and mask_a is not None:

torchani/aev_base.py

@@ -6,7 +6,7 @@ from .benchmarked import BenchmarkedModule
 
 class AEVComputer(BenchmarkedModule):
     __constants__ = ['Rcr', 'Rca', 'dtype', 'device', 'radial_sublength',
-        'radial_length', 'angular_sublength', 'angular_length', 'aev_length']
+                     'radial_length', 'angular_sublength', 'angular_length', 'aev_length']
 
     """Base class of various implementations of AEV computer
 

torchani/tests/test_aev.py

-import torch
-import numpy
-import torchani
-import unittest
-import pyanitools
-import os
-import logging
-
-
-class TestAEV(unittest.TestCase):
-
-    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
-        self.aev = torchani.SortedAEV(dtype=dtype, device=device)
-        self.ncaev = torchani.NeuroChemAEV(dtype=dtype, device=device)
-        self.tolerance = 1e-5
-        self.logger = logging.getLogger('smiles')
-
-    def _test_molecule(self, coordinates, species):
-        coordinates = torch.from_numpy(coordinates).type(
-            self.aev.dtype).to(self.aev.device)
-        coordinates, species = self.aev.sort_by_species(coordinates, species)
-        radial_neurochem, angular_neurochem = self.ncaev(coordinates, species)
-        radial_torchani, angular_torchani = self.aev(coordinates, species)
-        radial_diff = radial_neurochem - radial_torchani
-        radial_max_error = torch.max(torch.abs(radial_diff))
-        angular_diff = angular_neurochem - angular_torchani
-        angular_max_error = torch.max(torch.abs(angular_diff))
-        if radial_max_error > self.tolerance:
-            print('radial aev for', species)
-            for i in range(len(species)):
-                r1 = radial_neurochem[0, i, :]
-                r2 = radial_torchani[0, i, :]
-                max_err = torch.max(torch.abs(r1 - r2))
-                if max_err > self.tolerance:
-                    print('atom', i, 'species', species[i], 'radial aevs:')
-                    print(torch.stack([r1, r2, r1-r2], dim=1))
-                    break
-        if angular_max_error > self.tolerance:
-            print('angular aev for', species)
-            for i in range(len(species)):
-                r1 = angular_neurochem[0, i, :]
-                r2 = angular_torchani[0, i, :]
-                max_err = torch.max(torch.abs(r1 - r2))
-                if max_err > self.tolerance:
-                    print('atom', i, 'species', species[i], 'angular aevs:')
-                    print(torch.stack([r1, r2, r1-r2], dim=1))
-                    break
-        self.assertLess(radial_max_error, self.tolerance)
-        self.assertLess(angular_max_error, self.tolerance)
-
-    def _test_datafile(self, number):
-        data_file = os.path.join(
-            torchani.buildin_dataset_dir, 'ani_gdb_s0{}.h5'.format(number))
-        adl = pyanitools.anidataloader(data_file)
-        for data in adl:
-            coordinates = data['coordinates'][:10, :]
-            species = data['species']
-            smiles = ''.join(data['smiles'])
-            self._test_molecule(coordinates, species)
-            self.logger.info('Test pass: ' + smiles)
-
-    def testGDB01(self):
-        self._test_datafile(1)
-
-    def testGDB02(self):
-        self._test_datafile(2)
-
-    def testGDB03(self):
-        self._test_datafile(3)
-
-    def testGDB04(self):
-        self._test_datafile(4)
-
-    def testCH4(self):
-        # return
-        coordinates = numpy.array([[[0.03192167,  0.00638559,  0.01301679],
-                                    [-0.83140486,  0.39370209, -0.26395324],
-                                    [-0.66518241, -0.84461308,  0.20759389],
-                                    [0.45554739,  0.54289633,  0.81170881],
-                                    [0.66091919, -0.16799635, -0.91037834]]], numpy.float32)
-        species = ['C', 'H', 'H', 'H', 'H']
-        self._test_molecule(coordinates, species)
-
-
-if __name__ == '__main__':
-    unittest.main()

torchani/tests/test_force_neurochem.py

-import torch
-import numpy
-import torchani
-import unittest
-import ase
-import pyNeuroChem
-import ase_interface
-import pyanitools
-import os
-import logging
-
-
-class TestForceNeuroChem(unittest.TestCase):
-
-    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
-        self.tolerance = 1e-5
-        self.logger = logging.getLogger('smiles')
-        self.ncaev = torchani.NeuroChemAEV(dtype=dtype, device=device)
-        self.aev_computer = torchani.SortedAEV(dtype=dtype, device=device)
-        self.model1 = torchani.ModelOnAEV(
-            self.aev_computer, derivative=True, device=device, from_nc=None)
-        self.model2 = torchani.ModelOnAEV(
-            self.aev_computer, derivative=True, device=device, from_nc=None, ensemble=1)
-        self.logger = logging.getLogger('smiles')
-
-    def _test_molecule(self, coordinates, species):
-        _, force1 = self.model1(coordinates, species)
-        _, force2 = self.model1(coordinates, species)
-        conformations = coordinates.shape[0]
-        for i in range(conformations):
-            c = coordinates[i]
-            mol = ase.Atoms(''.join(species), positions=c)
-            mol.set_calculator(ase_interface.ANI(False))
-            mol.calc.setnc(self.ncaev.nc)
-            _ = mol.get_potential_energy()
-            force_nc = self.ncaev.nc.force()
-            force_nc = torch.from_numpy(force_nc).type(
-                self.aev_computer.dtype).to(self.aev_computer.device)
-            max_diff1 = torch.max(force_nc+force1[i])
-            max_diff2 = torch.max(force_nc+force2[i])
-            max_diff = max(max_diff1, max_diff2)
-            self.assertLess(max_diff, self.tolerance)
-
-    def testCH4(self):
-        coordinates = torch.tensor([[[0.03192167,  0.00638559,  0.01301679],
-                                     [-0.83140486,  0.39370209, -0.26395324],
-                                     [-0.66518241, -0.84461308,  0.20759389],
-                                     [0.45554739,   0.54289633,  0.81170881],
-                                     [0.66091919,  -0.16799635, -0.91037834]],
-                                    [[0,  0,  0],
-                                     [0,  0,  1],
-                                     [1,  0,  0],
-                                     [0,  1,  0],
-                                     [-1, -1, -1]],
-                                    ], dtype=self.aev_computer.dtype, device=self.aev_computer.device)
-        species = ['C', 'H', 'H', 'H', 'H']
-        self._test_molecule(coordinates, species)
-
-    def _test_by_file(self, number):
-        data_file = os.path.join(
-            torchani.buildin_dataset_dir, 'ani_gdb_s0{}.h5'.format(number))
-        adl = pyanitools.anidataloader(data_file)
-        for data in adl:
-            coordinates = data['coordinates'][:10, :]
-            coordinates = torch.from_numpy(coordinates).type(
-                self.aev_computer.dtype).to(self.aev_computer.device)
-            species = data['species']
-            smiles = ''.join(data['smiles'])
-            self._test_molecule(coordinates, species)
-            self.logger.info('Test pass: ' + smiles)
-
-    def testGDB01(self):
-        self._test_by_file(1)
-
-    def testGDB02(self):
-        self._test_by_file(2)
-
-    def testGDB03(self):
-        self._test_by_file(3)
-
-    def testGDB04(self):
-        self._test_by_file(4)
-
-
-if __name__ == '__main__':
-    unittest.main()

torchani/tests/test_inference.py

-import torch
-import numpy
-import torchani
-import unittest
-import os
-import logging
-import pyanitools
-import ase
-import pyNeuroChem
-import ase_interface
-
-
-class TestInference(unittest.TestCase):
-
-    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
-        self.tolerance = 1e-5
-        self.ncaev = torchani.NeuroChemAEV(dtype=dtype, device=device)
-        self.nn = torchani.ModelOnAEV(
-            self.ncaev, from_nc=self.ncaev.network_dir)
-        self.nn1 = torchani.ModelOnAEV(self.ncaev, from_nc=None)
-        self.nn2 = torchani.ModelOnAEV(
-            self.ncaev, from_nc=torchani.buildin_model_prefix, ensemble=1)
-        self.logger = logging.getLogger('smiles')
-        self.shift_energy = torchani.EnergyShifter(self.ncaev.sae_file)
-
-    def _get_neurochem_energies(self, coordinates, species):
-        conformations = coordinates.shape[0]
-        nc_energies = []
-        for i in range(conformations):
-            c = coordinates[i]
-            mol = ase.Atoms(''.join(species), positions=c)
-            mol.set_calculator(ase_interface.ANI(False))
-            mol.calc.setnc(self.ncaev.nc)
-            _ = mol.get_potential_energy()
-            e = self.ncaev.nc.energy()[0]
-            nc_energies.append(e)
-        nc_energies = torch.DoubleTensor(nc_energies)
-        return nc_energies.type(self.ncaev.dtype).to(self.ncaev.device)
-
-    def _test_molecule_energy(self, coordinates, species):
-        energies = self._get_neurochem_energies(coordinates, species)
-        energies = self.shift_energy.subtract_sae(energies, species)
-        coordinates = torch.from_numpy(coordinates).type(
-            self.ncaev.dtype).to(self.ncaev.device)
-        pred_energies1 = self.nn1(coordinates, species).squeeze()
-        pred_energies2 = self.nn2(coordinates, species).squeeze()
-        maxdiff1 = torch.max(torch.abs(pred_energies1 - energies)).item()
-        maxdiff2 = torch.max(torch.abs(pred_energies2 - energies)).item()
-        maxdiff = max(maxdiff1, maxdiff2)
-        maxdiff_per_atom = maxdiff / len(species)
-        self.assertLess(maxdiff_per_atom, self.tolerance)
-
-    def _test_activations(self, coordinates, species):
-        conformations = coordinates.shape[0]
-        atoms = coordinates.shape[1]
-        radial_aev, angular_aev = self.nn.aev_computer(coordinates, species)
-        aev = torch.cat([radial_aev, angular_aev], dim=2)
-        for i in range(conformations):
-            for j in range(atoms):
-                model_X = getattr(self.nn, 'model_' + species[j])
-                layers = model_X.layers
-                for layer in range(layers):
-                    # get activation from NeuroChem
-                    c = coordinates[i]
-                    mol = ase.Atoms(''.join(species), positions=c)
-                    mol.set_calculator(ase_interface.ANI(False))
-                    mol.calc.setnc(self.ncaev.nc)
-                    _ = mol.get_potential_energy()
-                    nca = self.ncaev.nc.activations(j, layer, 0)
-                    nca = torch.from_numpy(nca).type(
-                        self.ncaev.dtype).to(self.ncaev.device)
-                    # get activation from ModelOnAEV
-                    atom_aev = aev[:, j, :]
-                    a = model_X.get_activations(atom_aev, layer)
-                    a = a[i].view(-1)
-                    # compute diff
-                    maxdiff = torch.max(torch.abs(nca - a)).item()
-                    self.assertLess(maxdiff, self.tolerance)
-
-    def _test_by_file(self, number):
-        data_file = os.path.join(
-            torchani.buildin_dataset_dir, 'ani_gdb_s0{}.h5'.format(number))
-        adl = pyanitools.anidataloader(data_file)
-        for data in adl:
-            coordinates = data['coordinates'][:10, :]
-            species = data['species']
-            smiles = ''.join(data['smiles'])
-            self._test_activations(coordinates, species)
-            self._test_molecule_energy(coordinates, species)
-            self.logger.info('Test pass: ' + smiles)
-
-    def testGDB01(self):
-        self._test_by_file(1)
-
-    def testGDB02(self):
-        self._test_by_file(2)
-
-    def testGDB03(self):
-        self._test_by_file(3)
-
-    def testGDB04(self):
-        self._test_by_file(4)
-
-
-if __name__ == '__main__':
-    unittest.main()

torchani/tests/test_neurochem_loader.py

-import torch
-import numpy
-import torchani
-import unittest
-import logging
-
-
-class TestNeuroChemLoader(unittest.TestCase):
-
-    def setUp(self, dtype=torchani.default_dtype, device=torchani.default_device):
-        self.tolerance = 1e-5
-        self.ncaev = torchani.NeuroChemAEV(dtype=dtype, device=device)
-        self.logger = logging.getLogger('species')
-
-    def testLoader(self):
-        nn = torchani.ModelOnAEV(
-            self.ncaev, from_nc=self.ncaev.network_dir)
-        for i in range(len(self.ncaev.species)):
-            s = self.ncaev.species[i]
-            model_X = getattr(nn, 'model_' + s)
-            self.logger.info(s)
-            for j in range(model_X.layers):
-                linear = getattr(model_X, 'layer{}'.format(j))
-                ncparams = self.ncaev.nc.getntwkparams(i, j)
-                ncw = ncparams['weights']
-                ncw = torch.from_numpy(ncw).type(
-                    self.ncaev.dtype).to(self.ncaev.device)
-                ncb = numpy.transpose(ncparams['biases'])
-                ncb = torch.from_numpy(ncb).type(
-                    self.ncaev.dtype).to(self.ncaev.device)
-                max_wdiff = torch.max(
-                    torch.abs(ncw - linear.weight.data)).item()
-                max_bdiff = torch.max(torch.abs(ncb - linear.bias.data)).item()
-                self.assertEqual(max_bdiff, 0.0)
-                self.assertEqual(max_wdiff, 0.0)
-
-
-if __name__ == '__main__':
-    unittest.main()

torchani/tests/test_onnx.py

-import torchani
-import unittest
-import torch
-import cntk
-import tempfile
-import os
-import numpy
-
-
-class TestONNX(unittest.TestCase):
-
-    def setUp(self):
-        self.tolerance = 1e-5
-
-    def testONNX(self):  # not ready yet
-        return
-
-        # molecule structure: CH2OH
-        species = ['C', 'H', 'H', 'O', 'H']
-        coordinates = [
-            [0, 0, 0],  # C
-            [0, 0, 1],  # H
-            [1, 0, 0],  # H
-            [0, 1, 0],  # O
-            [0, 1, 1],  # H
-        ]
-
-        # compute aev using pytorch
-        aev_computer = torchani.AEV()
-        coordinates = torch.FloatTensor(coordinates)
-        coordinates = coordinates.unsqueeze(0)
-        radial_aev, angular_aev = aev_computer(coordinates, species)
-        aev = torch.cat([radial_aev, angular_aev], dim=2).numpy()
-
-        # temp directory storing exported networks
-        tmpdir = tempfile.TemporaryDirectory()
-        tmpdirname = tmpdir.name
-
-        ####################################################
-        # Step 1: use pytorch to export all graphs into ONNX
-        ####################################################
-
-        # TODO: exporting AEV to ONNX is not supported yet,
-        # due to lack of operators in ONNX. Add this support
-        # when ONNX support this operation.
-
-        aev_computer.export_radial_subaev_onnx(
-            os.path.join(tmpdirname, 'radial.onnx'))
-
-        # Export neural network potential to ONNX
-        model = torchani.ModelOnAEV(aev_computer, from_nc=None)
-        model.export_onnx(tmpdirname)
-
-        #####################################
-        # Step 2: import from ONNX using CNTK
-        #####################################
-        networks = {}
-        for s in aev_computer.species:
-            nn_onnx = os.path.join(tmpdirname, '{}.proto'.format(s))
-            networks[s] = cntk.Function.load(
-                nn_onnx, format=cntk.ModelFormat.ONNX)
-
-        ###################################
-        # Step 3: compute energy using CNTX
-        ###################################
-        energy1 = 0
-        for i in range(len(species)):
-            atomic_aev = aev[:, i, :]
-            network = networks[species[i]]
-            atomic_energy = network(atomic_aev)[0, 0, 0]
-            energy1 += atomic_energy
-
-        ###############################################
-        # Test only: check the CNTK result with pytorch
-        ###############################################
-        energy2 = model(coordinates, species).squeeze().item()
-        self.assertLessEqual(abs(energy1 - energy2), self.tolerance)
-
-
-if __name__ == '__main__':
-    unittest.main()