diff --git a/dpdata/system.py b/dpdata/system.py index 2614bc23b..2613166a7 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -776,9 +776,11 @@ 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"]), (int(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_predict.py b/tests/test_predict.py index 6ab00be36..9721323ec 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -106,7 +106,9 @@ def setUp(self): ) zero_driver = ZeroDriver() self.system_1 = ori_sys.predict(driver=zero_driver) - self.system_2 = ori_sys.minimize(driver=zero_driver, minimizer="ase") + self.system_2 = ori_sys.minimize( + driver=zero_driver, minimizer="ase", max_steps=100 + ) self.places = 6 self.e_places = 6 self.f_places = 6 @@ -123,7 +125,9 @@ def setUp(self): zero_driver = ZeroDriver() self.system_1 = list(multi_sys.predict(driver=zero_driver).systems.values())[0] self.system_2 = list( - multi_sys.minimize(driver=zero_driver, minimizer="ase").systems.values() + multi_sys.minimize( + driver=zero_driver, minimizer="ase", max_steps=100 + ).systems.values() )[0] self.places = 6 self.e_places = 6 diff --git a/tests/test_replicate.py b/tests/test_replicate.py index 3add2dc02..ded164c17 100644 --- 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 @@ -36,5 +37,32 @@ def setUp(self): 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()