Test boundary see each other (#128)

tests/test_ase.py

@@ -6,6 +6,7 @@ import torch
 import torchani
 import unittest
 import numpy
+import itertools
 
 
 def get_numeric_force(atoms, eps):
@@ -86,6 +87,77 @@ class TestASE(unittest.TestCase):
             atoms.set_positions(positions + translation)
             self.assertEqual(e, atoms.get_potential_energy())
 
+    def assertTensorEqual(self, a, b):
+        self.assertLess((a - b).abs().max().item(), 1e-6)
+
+    def testPBCConnersSeeEachOther(self):
+        species = torch.tensor([[0, 0]])
+        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)
+
+        xyz1 = torch.tensor([0.1, 0.1, 0.1])
+        xyz2s = [
+            torch.tensor([9.9, 0.0, 0.0]),
+            torch.tensor([0.0, 9.9, 0.0]),
+            torch.tensor([0.0, 0.0, 9.9]),
+            torch.tensor([9.9, 9.9, 0.0]),
+            torch.tensor([0.0, 9.9, 9.9]),
+            torch.tensor([9.9, 0.0, 9.9]),
+            torch.tensor([9.9, 9.9, 9.9]),
+        ]
+
+        for xyz2 in xyz2s:
+            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
+            s, _, D = neighborlist(species, coordinates, 1)
+            self.assertListEqual(list(s.shape), [1, 2, 1])
+            neighbor_coordinate = D[0][0].squeeze() + xyz1
+            mirror = xyz2
+            for i in range(3):
+                if mirror[i] > 5:
+                    mirror[i] -= 10
+            self.assertTensorEqual(neighbor_coordinate, mirror)
+
+    def testPBCSurfaceSeeEachOther(self):
+        species = torch.tensor([[0, 0]])
+        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)
+
+        for i in range(3):
+            xyz1 = torch.tensor([5.0, 5.0, 5.0])
+            xyz1[i] = 0.1
+            xyz2 = xyz1.clone()
+            xyz2[i] = 9.9
+
+            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
+            s, _, D = neighborlist(species, coordinates, 1)
+            self.assertListEqual(list(s.shape), [1, 2, 1])
+            neighbor_coordinate = D[0][0].squeeze() + xyz1
+            xyz2[i] = -0.1
+            self.assertTensorEqual(neighbor_coordinate, xyz2)
+
+    def testPBCEdgesSeeEachOther(self):
+        species = torch.tensor([[0, 0]])
+        neighborlist = torchani.ase.NeighborList(cell=[10, 10, 10], pbc=True)
+
+        for i, j in itertools.combinations(range(3), 2):
+            xyz1 = torch.tensor([5.0, 5.0, 5.0])
+            xyz1[i] = 0.1
+            xyz1[j] = 0.1
+            for new_i, new_j in [[0.1, 9.9], [9.9, 0.1], [9.9, 9.9]]:
+                xyz2 = xyz1.clone()
+                xyz2[i] = new_i
+                xyz2[j] = new_i
+
+            coordinates = torch.stack([xyz1, xyz2]).unsqueeze(0)
+            s, _, D = neighborlist(species, coordinates, 1)
+            self.assertListEqual(list(s.shape), [1, 2, 1])
+            neighbor_coordinate = D[0][0].squeeze() + xyz1
+
+            if xyz2[i] > 5:
+                xyz2[i] = -0.1
+            if xyz2[j] > 5:
+                xyz2[j] = -0.1
+
+            self.assertTensorEqual(neighbor_coordinate, xyz2)
+
 
 if __name__ == '__main__':
     unittest.main()