add RunLmpHDF5

Signed-off-by: zjgemi <liuxin_zijian@163.com>

dpgen2/entrypoint/submit.py

@@ -110,6 +110,7 @@
     RunCalyModelDevi,
     RunDPTrain,
     RunLmp,
+    RunLmpHDF5,
     RunRelax,
     RunRelaxHDF5,
     SelectConfs,
@@ -187,7 +188,7 @@ def make_concurrent_learning_op(
         prep_run_explore_op = PrepRunLmp(
             "prep-run-lmp",
             PrepLmp,
-            RunLmp,
+            RunLmpHDF5 if explore_config["use_hdf5"] else RunLmp,  # type: ignore
             prep_config=prep_explore_config,
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,

dpgen2/op/__init__.py

@@ -36,6 +36,7 @@
 )
 from .run_lmp import (
     RunLmp,
+    RunLmpHDF5,
 )
 from .run_relax import (
     RunRelax,

dpgen2/op/run_lmp.py

@@ -14,6 +14,7 @@
     Tuple,
 )
 
+import numpy as np
 from dargs import (
     Argument,
     ArgumentEncoder,
@@ -26,6 +27,7 @@
     Artifact,
     BigParameter,
     FatalError,
+    HDF5Datasets,
     OPIOSign,
     TransientError,
 )
@@ -200,7 +202,7 @@ def execute(
         ret_dict = {
             "log": work_dir / lmp_log_name,
             "traj": work_dir / lmp_traj_name,
-            "model_devi": work_dir / lmp_model_devi_name,
+            "model_devi": self.get_model_devi(work_dir / lmp_model_devi_name),
         }
         plm_output = (
             {"plm_output": work_dir / plm_output_name}
@@ -213,13 +215,17 @@ def execute(
 
         return OPIO(ret_dict)
 
+    def get_model_devi(self, model_devi_file):
+        return model_devi_file
+
     @staticmethod
     def lmp_args():
         doc_lmp_cmd = "The command of LAMMPS"
         doc_teacher_model = "The teacher model in `Knowledge Distillation`"
         doc_shuffle_models = "Randomly pick a model from the group of models to drive theexploration MD simulation"
         doc_head = "Select a head from multitask"
         doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature"
+        doc_use_hdf5 = "Use HDF5 to store trajs and model_devis"
         return [
             Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd),
             Argument(
@@ -243,6 +249,13 @@ def lmp_args():
             Argument(
                 "model_frozen_head", str, optional=True, default=None, doc=doc_head
             ),
+            Argument(
+                "use_hdf5",
+                bool,
+                optional=True,
+                default=False,
+                doc=doc_use_hdf5,
+            ),
         ]
 
     @staticmethod
@@ -374,3 +387,15 @@ def merge_pimd_files():
             for model_devi_file in sorted(model_devi_files):
                 with open(model_devi_file, "r") as f2:
                     f.write(f2.read())
+
+
+class RunLmpHDF5(RunLmp):
+    @classmethod
+    def get_output_sign(cls):
+        output_sign = super().get_output_sign()
+        output_sign["traj"] = Artifact(HDF5Datasets)
+        output_sign["model_devi"] = Artifact(HDF5Datasets)
+        return output_sign
+
+    def get_model_devi(self, model_devi_file):
+        return np.loadtxt(model_devi_file)