Test non-rectangular PBC (#139)

tests/test_ase.py

@@ -7,6 +7,7 @@
 import unittest
 import numpy
 import itertools
+import math
 
 
 def get_numeric_force(atoms, eps):
@@ -158,6 +159,24 @@ def testPBCEdgesSeeEachOther(self):
 
             self.assertTensorEqual(neighbor_coordinate, xyz2)
 
+    def testNonRectangularPBCConnersSeeEachOther(self):
+        species = torch.tensor([[0, 0]])
+        neighborlist = torchani.ase.NeighborList(
+            cell=[10, 10, 10 * math.sqrt(2), 90, 45, 90], pbc=True)
+
+        xyz1 = torch.tensor([0.1, 0.1, 0.05])
+        xyz2 = torch.tensor([10.0, 0.1, 0.1])
+        mirror = torch.tensor([0.0, 0.1, 0.1])
+
+        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
+        for i in range(3):
+            if mirror[i] > 5:
+                mirror[i] -= 10
+        self.assertTensorEqual(neighbor_coordinate, mirror)
+
 
 if __name__ == '__main__':
     unittest.main()