Add RunLmpHDF5 (#267)

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

- **New Features**
- Introduced dynamic selection of exploration operations based on
configuration settings, enhancing configurability.
- Added a new `RunLmpHDF5` option for improved data handling and output
types.
- **Improvements**
- Enhanced the `RunLmp` class with new methods for better processing of
model deviation files and output data storage options.
- Improved type safety and error handling in the `TrajRenderLammps`
class, including better management of temperature settings.
- **Dependency Updates**
- Updated the version requirement for the `pydflow` dependency to
`>=1.8.95`.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: zjgemi <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.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/exploration/render/traj_render_lammps.py

@@ -63,15 +63,17 @@ def _load_one_model_devi(self, fname, model_devi):
             dd = fname.get_data()
         else:
             dd = np.loadtxt(fname)
-        if len(np.shape(dd)) == 1:  # In case model-devi.out is 1-dimensional
-            dd = dd.reshape((1, len(dd)))
+        if (
+            len(np.shape(dd)) == 1  # type: ignore
+        ):  # In case model-devi.out is 1-dimensional
+            dd = dd.reshape((1, len(dd)))  # type: ignore
 
-        model_devi.add(DeviManager.MAX_DEVI_V, dd[:, 1])
-        model_devi.add(DeviManager.MIN_DEVI_V, dd[:, 2])
-        model_devi.add(DeviManager.AVG_DEVI_V, dd[:, 3])
-        model_devi.add(DeviManager.MAX_DEVI_F, dd[:, 4])
-        model_devi.add(DeviManager.MIN_DEVI_F, dd[:, 5])
-        model_devi.add(DeviManager.AVG_DEVI_F, dd[:, 6])
+        model_devi.add(DeviManager.MAX_DEVI_V, dd[:, 1])  # type: ignore
+        model_devi.add(DeviManager.MIN_DEVI_V, dd[:, 2])  # type: ignore
+        model_devi.add(DeviManager.AVG_DEVI_V, dd[:, 3])  # type: ignore
+        model_devi.add(DeviManager.MAX_DEVI_F, dd[:, 4])  # type: ignore
+        model_devi.add(DeviManager.MIN_DEVI_F, dd[:, 5])  # type: ignore
+        model_devi.add(DeviManager.AVG_DEVI_F, dd[:, 6])  # type: ignore
 
     def get_ele_temp(self, optional_outputs):
         ele_temp = []

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)

pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
 dependencies = [
 	     'numpy',
 	     'dpdata>=0.2.20',
-	     'pydflow>=1.8.88',
+	     'pydflow>=1.8.95',
 	     'dargs>=0.3.1',
 	     'scipy',
 	     'lbg',