support one_h5 for simplify model_devi step (#1185)

For `dp model_devi`, large number of small NumPy files are still quite
expensive on HPCs.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

dpgen/generator/arginfo.py

@@ -92,9 +92,7 @@ def training_args() -> List[Argument]:
     )
     doc_model_devi_activation_func = "The activation function in the model. The shape of list should be (N_models, 2), where 2 represents the embedding and fitting network. This option will override default parameters."
     doc_srtab_file_path = "The path of the table for the short-range pairwise interaction which is needed when using DP-ZBL potential"
-    doc_one_h5 = (
-        "Before training, all of the training data will be merged into one HDF5 file."
-    )
+    doc_one_h5 = "When using DeePMD-kit, all of the input data will be merged into one HDF5 file."
 
     return [
         Argument("numb_models", int, optional=False, doc=doc_numb_models),

dpgen/simplify/simplify.py

@@ -218,17 +218,26 @@ def run_model_devi(iter_index, jdata, mdata):
     # models
     commands = []
     detail_file_name = detail_file_name_prefix
+    system_file_name = rest_data_name + ".old"
+    if jdata.get("one_h5", False):
+        # convert system to one h5 file
+        system_path = os.path.join(work_path, system_file_name)
+        dpdata.MultiSystems(type_map=jdata["type_map"]).from_deepmd_npy(
+            system_path,
+            labeled=False,
+        ).to_deepmd_hdf5(system_path + ".hdf5")
+        system_file_name += ".hdf5"
     command = "{dp} model-devi -m {model} -s {system} -o {detail_file}".format(
         dp=mdata.get("model_devi_command", "dp"),
         model=" ".join(task_model_list),
-        system=rest_data_name + ".old",
+        system=system_file_name,
         detail_file=detail_file_name,
     )
     commands = [command]
     # submit
     model_devi_group_size = mdata.get("model_devi_group_size", 1)
 
-    forward_files = [rest_data_name + ".old"]
+    forward_files = [system_file_name]
     backward_files = [detail_file_name]
 
     api_version = mdata.get("api_version", "1.0")

tests/simplify/__init__.py

@@ -0,0 +1 @@
+

tests/simplify/test_get_multi_system.py

@@ -1,10 +1,14 @@
 import os
 import shutil
+import sys
 import unittest
 
 import dpdata
 import numpy as np
-from context import dpgen
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+__package__ = "simplify"
+from .context import dpgen
 
 
 class TestGetMultiSystem(unittest.TestCase):
@@ -44,7 +48,7 @@ def test_get_multi_system(self):
                 with self.subTest(list_data=list_data, labeled=labeled):
                     ms = dpgen.simplify.simplify.get_multi_system(
                         self.data if list_data else self.data[0],
-                        {"labeled": labeled},
+                        {"labeled": labeled, "type_map": ["H"]},
                     )
                     assert isinstance(ms, dpdata.MultiSystems)
                     for ss in ms.systems.values():

tests/simplify/test_run_model_devi.py

@@ -0,0 +1,86 @@
+import os
+import shutil
+import sys
+import tempfile
+import textwrap
+import unittest
+from pathlib import Path
+
+import dpdata
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+__package__ = "simplify"
+from .context import dpgen
+
+
+class TestOneH5(unittest.TestCase):
+    def setUp(self):
+        work_path = Path("iter.000000") / "01.model_devi"
+        work_path.mkdir(parents=True, exist_ok=True)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with open(Path(tmpdir) / "test.xyz", "w") as f:
+                f.write(
+                    textwrap.dedent(
+                        """\
+                    2
+
+                    H 0.0 0.0 0.0
+                    H 0.0 0.0 1.0
+                    """
+                    )
+                )
+            dpdata.MultiSystems(
+                dpdata.System(Path(tmpdir) / "test.xyz", fmt="xyz")
+            ).to_deepmd_npy(
+                work_path / (dpgen.simplify.simplify.rest_data_name + ".old")
+            )
+
+    def tearDown(self) -> None:
+        shutil.rmtree("iter.000000")
+
+    def test_npy(self):
+        jdata = {
+            "type_map": ["H"],
+        }
+        with tempfile.TemporaryDirectory() as remote_root:
+            mdata = {
+                "model_devi_command": (
+                    f"test -d {dpgen.simplify.simplify.rest_data_name}.old"
+                    f"&& touch {dpgen.simplify.simplify.detail_file_name_prefix}"
+                    "&& echo dp"
+                ),
+                "model_devi_machine": {
+                    "context_type": "LocalContext",
+                    "batch_type": "shell",
+                    "local_root": "./",
+                    "remote_root": remote_root,
+                },
+                "model_devi_resources": {
+                    "group_size": 1,
+                },
+            }
+            dpgen.simplify.simplify.run_model_devi(0, jdata=jdata, mdata=mdata)
+
+    def test_one_h5(self):
+        jdata = {
+            "type_map": ["H"],
+            "one_h5": True,
+        }
+        with tempfile.TemporaryDirectory() as remote_root:
+            mdata = {
+                "model_devi_command": (
+                    f"test -f {dpgen.simplify.simplify.rest_data_name}.old.hdf5"
+                    f"&& touch {dpgen.simplify.simplify.detail_file_name_prefix}"
+                    "&& echo dp"
+                ),
+                "model_devi_machine": {
+                    "context_type": "LocalContext",
+                    "batch_type": "shell",
+                    "local_root": "./",
+                    "remote_root": remote_root,
+                },
+                "model_devi_resources": {
+                    "group_size": 1,
+                },
+            }
+            dpgen.simplify.simplify.run_model_devi(0, jdata=jdata, mdata=mdata)