fix: the replicate will fail if the atom types of system is not sorted

wanghan-iapcm · Jun 5, 2024 · 2e1e993 · 2e1e993
1 parent 199afc1
commit 2e1e993
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/dpdata/system.py b/dpdata/system.py
@@ -776,9 +776,9 @@ def replicate(self, ncopy: list[int] | tuple[int, int, int]):
         tmp.data["atom_numbs"] = list(
             np.array(np.copy(data["atom_numbs"])) * np.prod(ncopy)
         )
-        tmp.data["atom_types"] = np.sort(
-            np.tile(np.copy(data["atom_types"]), np.prod(ncopy).item()), kind="stable"
-        )
+        tmp.data["atom_types"] = np.tile(np.copy(data["atom_types"]), (np.prod(ncopy),) + (1,))
+        tmp.data["atom_types"] = np.transpose(tmp.data["atom_types"]).reshape([-1])
+
         tmp.data["cells"] = np.copy(data["cells"])
         for ii in range(3):
             tmp.data["cells"][:, ii, :] *= ncopy[ii]

diff --git a/tests/test_replicate.py b/tests/test_replicate.py
@@ -2,6 +2,7 @@
 
 import unittest
 
+import numpy as np
 from comp_sys import CompSys, IsPBC
 from context import dpdata
 
@@ -35,6 +36,24 @@ def setUp(self):
         self.system_2 = dpdata.System("poscars/POSCAR.SiC", fmt="vasp/poscar")
         self.places = 6
 
+class TestReplicateTriclinicBox(unittest.TestCase, CompSys, IsPBC):
+    def setUp(self):
+        self.system_1 = dpdata.System()
+        self.system_1.data["atom_names"] = ["foo", "bar"]
+        self.system_1.data["atom_types"] = np.array([1, 0], dtype=int)
+        self.system_1.data["atom_numbs"] = [1, 1]
+        self.system_1.data["cells"] = np.array([10, 0, 0, 0, 10, 0, 0, 0, 10], dtype=float).reshape(1,3,3)
+        self.system_1.data["coords"] = np.array([0, 0, 0, 0, 0, 1], dtype=float).reshape(1,2,3)
+        self.system_1 = self.system_1.replicate([2,1,1])
+
+        self.system_2 = dpdata.System()
+        self.system_2.data["atom_names"] = ["foo", "bar"]
+        self.system_2.data["atom_types"] = np.array([1, 1, 0, 0], dtype=int)
+        self.system_2.data["atom_numbs"] = [2, 2]
+        self.system_2.data["cells"] = np.array([20, 0, 0, 0, 10, 0, 0, 0, 10], dtype=float).reshape(1,3,3)
+        self.system_2.data["coords"] = np.array([0, 0, 0, 10, 0, 0, 0, 0, 1, 10, 0, 1], dtype=float).reshape(1,4,3)
+        self.places = 6
+
 
 if __name__ == "__main__":
     unittest.main()