From 2e1e993aa1dfbc8456360b5c69a11c33b9861f59 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 5 Jun 2024 22:56:16 +0800 Subject: [PATCH] fix: the replicate will fail if the atom types of system is not sorted --- dpdata/system.py | 6 +++--- tests/test_replicate.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/dpdata/system.py b/dpdata/system.py index 2614bc23..2de98368 100644 --- 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 index 3add2dc0..bdd8c5bb 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 @@ -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()