From bba3e9f28eed4d97f649ff79812e31034b3c8d61 Mon Sep 17 00:00:00 2001
From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Date: Thu, 6 Jun 2024 00:25:35 +0800
Subject: [PATCH] fix: the replicate will fail if the atom types of system is
 not sorted (#667)

- add UT to detect the issue.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Bug Fixes**
- Improved the logic for replicating atom types in the system, ensuring
more accurate replication behavior.

- **Tests**
- Added new tests to verify the replication of atom types, enhancing
test coverage and reliability.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 dpdata/system.py        |  6 ++++--
 tests/test_predict.py   |  8 ++++++--
 tests/test_replicate.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

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()