fix sort_by_species (#57)

tests/test_energies.py

@@ -21,6 +21,12 @@ class TestEnergies(unittest.TestCase):
                                          nnp, shift_energy)
 
     def _test_molecule(self, coordinates, species, energies):
+        # generate a random permute
+        atoms = len(species)
+        randperm = torch.randperm(atoms)
+        coordinates = coordinates.index_select(1, randperm)
+        species = [species[i] for i in randperm.tolist()]
+
         _, energies_ = self.model((species, coordinates))
         max_diff = (energies - energies_.squeeze()).abs().max().item()
         self.assertLess(max_diff, self.tolerance)

tests/test_forces.py

@@ -18,6 +18,13 @@ class TestForce(unittest.TestCase):
         self.model = torch.nn.Sequential(prepare, aev_computer, nnp)
 
     def _test_molecule(self, coordinates, species, forces):
+        # generate a random permute
+        atoms = len(species)
+        randperm = torch.randperm(atoms)
+        coordinates = coordinates.index_select(1, randperm)
+        forces = forces.index_select(1, randperm)
+        species = [species[i] for i in randperm.tolist()]
+
         coordinates = torch.tensor(coordinates, requires_grad=True)
         _, energies = self.model((species, coordinates))
         derivative = torch.autograd.grad(energies.sum(), coordinates)[0]