Fix autograd for ASE calculator (#123)

Dockerfile

@@ -1,6 +1,6 @@
 FROM zasdfgbnm/pytorch-master
 RUN pacman -Sy --noconfirm python-sphinx python2-sphinx python-tqdm python2-tqdm python2-matplotlib python-matplotlib python-pillow python2-pillow flake8
-RUN pip install tensorboardX sphinx-gallery && pip2 install tensorboardX sphinx-gallery
+RUN pip install tensorboardX sphinx-gallery ase && pip2 install tensorboardX sphinx-gallery ase
 COPY . /torchani
 RUN cd torchani && pip install .
 RUN cd torchani && pip2 install .

tests/test_energies.py

@@ -14,10 +14,10 @@ class TestEnergies(unittest.TestCase):
     def setUp(self):
         self.tolerance = 5e-5
         builtins = torchani.neurochem.Builtins()
-        aev_computer = builtins.aev_computer
+        self.aev_computer = builtins.aev_computer
         nnp = builtins.models[0]
         shift_energy = builtins.energy_shifter
-        self.model = torch.nn.Sequential(aev_computer, nnp, shift_energy)
+        self.model = torch.nn.Sequential(self.aev_computer, nnp, shift_energy)
 
     def testIsomers(self):
         for i in range(N):
@@ -45,5 +45,12 @@ def testPadding(self):
         self.assertLess(max_diff, self.tolerance)
 
 
+class TestEnergiesASEComputer(TestEnergies):
+
+    def setUp(self):
+        super(TestEnergiesASEComputer, self).setUp()
+        self.aev_computer.neighborlist = torchani.ase.NeighborList()
+
+
 if __name__ == '__main__':
     unittest.main()

tests/test_forces.py

@@ -13,15 +13,21 @@ class TestForce(unittest.TestCase):
     def setUp(self):
         self.tolerance = 1e-5
         builtins = torchani.neurochem.Builtins()
-        aev_computer = builtins.aev_computer
+        self.aev_computer = builtins.aev_computer
         nnp = builtins.models[0]
-        self.model = torch.nn.Sequential(aev_computer, nnp)
+        self.model = torch.nn.Sequential(self.aev_computer, nnp)
+
+    def transform(self, x):
+        return x
 
     def testIsomers(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)
+                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(),
@@ -36,6 +42,9 @@ def testPadding(self):
             datafile = os.path.join(path, 'test_data/{}'.format(i))
             with open(datafile, 'rb') as f:
                 coordinates, species, _, _, _, forces = pickle.load(f)
+                coordinates = self.transform(coordinates)
+                species = self.transform(species)
+                forces = self.transform(forces)
                 coordinates.requires_grad_(True)
                 species_coordinates.append((species, coordinates))
                 coordinates_forces.append((coordinates, forces))
@@ -50,5 +59,16 @@ def testPadding(self):
             self.assertLess(max_diff, self.tolerance)
 
 
+class TestForceASEComputer(TestForce):
+
+    def setUp(self):
+        super(TestForceASEComputer, self).setUp()
+        self.aev_computer.neighborlist = torchani.ase.NeighborList()
+
+    def transform(self, x):
+        """To reduce the size of test cases for faster test speed"""
+        return x[:3, ...]
+
+
 if __name__ == '__main__':
     unittest.main()

torchani/ase.py

@@ -40,17 +40,18 @@ def __call__(self, species, coordinates, cutoff):
             atoms = s.shape[0]
             atoms_object = ase.Atoms(
                 'C'*atoms,  # chemical symbols are not important here
-                positions=c.numpy(),
+                positions=c.detach().numpy(),
                 pbc=self.pbc,
                 cell=self.cell)
-            idx1, idx2, d, D = ase.neighborlist.neighbor_list(
-                'ijdD', atoms_object, cutoff)
+            idx1, idx2 = ase.neighborlist.neighbor_list(
+                'ij', atoms_object, cutoff)
+            # NB: The absolute distance and distance vectors computed by
+            # `neighbor_list`can not be used since it does not preserve
+            # gradient information
             idx1 = torch.from_numpy(idx1).to(coordinates.device)
             idx2 = torch.from_numpy(idx2).to(coordinates.device)
-            d = torch.from_numpy(d).to(coordinates.device) \
-                .to(coordinates.dtype)
-            D = torch.from_numpy(D).to(coordinates.device) \
-                .to(coordinates.dtype)
+            D = c.index_select(0, idx2) - c.index_select(0, idx1)
+            d = D.norm(2, -1)
             neighbor_species1 = []
             neighbor_distances1 = []
             neighbor_vecs1 = []