Merge pull request #162 from wanghan-iapcm/conf-mixed

support MultiSystem for file configuration

dpgen2/conf/file_conf.py

@@ -47,6 +47,15 @@ def __init__(
     def generate(
         self,
         type_map,
+    ) -> dpdata.MultiSystems:
+        if self.fmt in ["deepmd/npy/mixed"]:
+            return self.generate_mixed(type_map)
+        else:
+            return self.generate_std(type_map)
+
+    def generate_std(
+        self,
+        type_map,
     ) -> dpdata.MultiSystems:
         ms = dpdata.MultiSystems(type_map=type_map)
         for ff in self.files:
@@ -56,6 +65,22 @@ def generate(
             ms.append(ss)
         return ms
 
+    def generate_mixed(
+        self,
+        type_map,
+    ) -> dpdata.MultiSystems:
+        if len(self.files) > 1:
+            raise ValueError(
+                'the file format "deepmd/npy/mixed" is specified, '
+                "but more than one file is given, which is invalide "
+                "please provide one path that can be interpreted as "
+                "the dpdata.MultiSystems. "
+            )
+        assert "deepmd/npy/mixed" == self.fmt
+        ms = dpdata.MultiSystems(type_map=type_map)
+        ms.from_deepmd_npy_mixed(self.files[0], fmt="deepmd/npy/mixed", labeled=False)
+        return ms
+
     @staticmethod
     def args() -> List[Argument]:
         doc_files = "The paths to the configuration files. widecards are supported."

tests/conf/test_file_conf.py

@@ -159,6 +159,27 @@ def test_normalize_1(self):
         )
         self.assertEqual(out_data, expected_out_data)
 
+    def test_deepmd_mixed(self):
+        type_map = ["Cu", "Al", "Mg"]
+        ms = dpdata.MultiSystems(type_map=type_map)
+        ms.append(
+            dpdata.System(Path(self.prefix) / Path("poscar.foo.0"), fmt="vasp/poscar")
+        )
+        ms.append(
+            dpdata.System(Path(self.prefix) / Path("poscar.foo.1"), fmt="vasp/poscar")
+        )
+        ms.to("deepmd/npy/mixed", "test_mixed")
+        fcg = FileConfGenerator("test_mixed", fmt="deepmd/npy/mixed")
+        ms1 = fcg.generate(type_map)
+        self.assertEqual(len(ms), len(ms1))
+        for tt in ["atom_names", "atom_numbs", "atom_types"]:
+            self.assertEqual(ms[0][tt], ms1[0][tt])
+            self.assertEqual(ms[1][tt], ms1[1][tt])
+        for tt in ["cells", "coords"]:
+            np.testing.assert_almost_equal(ms[0][tt], ms1[0][tt])
+            np.testing.assert_almost_equal(ms[1][tt], ms1[1][tt])
+        shutil.rmtree("test_mixed")
+
 
 class TestFileConfGeneratorContent(unittest.TestCase):
     def test_list_1(self):