fix: mode as a input

wangzyphysics · May 31, 2024 · a13fe8a · a13fe8a
1 parent 5830700
commit a13fe8a
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 29 deletions.
diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py
@@ -142,7 +142,6 @@ def make_concurrent_learning_op(
     train_style: str = "dp",
     explore_style: str = "lmp",
     fp_style: str = "vasp",
-    explore_config: dict = {},
     prep_train_config: dict = default_config,
     run_train_config: dict = default_config,
     prep_explore_config: dict = default_config,
@@ -155,7 +154,6 @@ def make_concurrent_learning_op(
     upload_python_packages: Optional[List[os.PathLike]] = None,
     valid_data: Optional[S3Artifact] = None,
 ):
-    expl_mode = explore_config.get("mode", "default")
     if train_style in ("dp", "dp-dist"):
         prep_run_train_op = PrepRunDPTrain(
             "prep-run-dp-train",
@@ -177,7 +175,8 @@ def make_concurrent_learning_op(
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
         )
-    elif explore_style == "calypso":
+    elif "calypso" in explore_style:
+        expl_mode = explore_style.split(":")[-1] if ":" in explore_style else "default"
         if expl_mode == "merge":
             caly_evo_step_op = CalyEvoStepMerge(
                 name="caly-evo-step",
@@ -199,14 +198,16 @@ def make_concurrent_learning_op(
                 upload_python_packages=upload_python_packages,
             )
         else:
-            raise KeyError(f"Unknown key: {expl_mode}, support `default` and `merge`.")
+            raise KeyError(
+                f"Unknown key: {explore_style}, support `calypso:default` and `calypso:merge`."
+            )
         prep_run_explore_op = PrepRunCaly(
             "prep-run-calypso",
             prep_caly_input_op=PrepCalyInput,
             caly_evo_step_op=caly_evo_step_op,
             prep_caly_model_devi_op=PrepCalyModelDevi,
             run_caly_model_devi_op=RunCalyModelDevi,
-            explore_config=explore_config,
+            expl_mode=expl_mode,
             prep_config=prep_explore_config,
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
@@ -510,7 +511,6 @@ def workflow_concurrent_learning(
         train_style,
         explore_style,
         fp_style,
-        explore_config=explore_config,
         prep_train_config=prep_train_config,
         run_train_config=run_train_config,
         prep_explore_config=prep_explore_config,

diff --git a/dpgen2/exploration/render/traj_render_lammps.py b/dpgen2/exploration/render/traj_render_lammps.py
@@ -45,15 +45,6 @@ def get_model_devi(
 
         return model_devi
 
-    def _load_one_model_devi_deprecated(self, fname, model_devi):
-        dd = np.loadtxt(fname)
-        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])
-
     def _load_one_model_devi(self, fname, model_devi):
         dd = np.loadtxt(fname)
         if len(np.shape(dd)) == 1:  # In case model-devi.out is 1-dimensional

diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py
@@ -57,7 +57,7 @@ def __init__(
         caly_evo_step_op: Union[OPTemplate, OP],
         prep_caly_model_devi_op: Type[OP],
         run_caly_model_devi_op: Type[OP],
-        explore_config: dict = {},
+        expl_mode: str = "default",
         prep_config: dict = normalize_step_dict({}),
         run_config: dict = normalize_step_dict({}),
         upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -115,7 +115,7 @@ def __init__(
             caly_evo_step_op,
             prep_caly_model_devi_op,
             run_caly_model_devi_op,
-            explore_config=explore_config,
+            expl_mode=expl_mode,
             prep_config=prep_config,
             run_config=run_config,
             upload_python_packages=upload_python_packages,
@@ -149,7 +149,7 @@ def _prep_run_caly(
     caly_evo_step_op: Union[OPTemplate, OP],
     prep_caly_model_devi_op: Type[OP],
     run_caly_model_devi_op: Type[OP],
-    explore_config: dict = {},
+    expl_mode: str = "default",
     prep_config: dict = normalize_step_dict({}),
     run_config: dict = normalize_step_dict({}),
     upload_python_packages: Optional[List[os.PathLike]] = None,
@@ -161,8 +161,6 @@ def _prep_run_caly(
     prep_executor = init_executor(prep_config.pop("executor"))
     run_executor = init_executor(run_config.pop("executor"))
     template_slice_config = run_config.pop("template_slice_config", {})
-    expl_mode = explore_config.get("mode", "default")
-    model_devi_group_size = explore_config.get("model_devi_group_size", 0)
 
     # prep caly input files
     prep_caly_input = Step(
@@ -255,7 +253,7 @@ def _prep_run_caly(
         ),
         parameters={
             "task_name": "prep-calypso-model-deviation",
-            "model_devi_group_size": model_devi_group_size,
+            "expl_config": prep_run_caly_steps.inputs.parameters["explore_config"],
         },
         artifacts={
             "traj_results": caly_evo_step.outputs.artifacts["traj_results"],

diff --git a/examples/calypso/input.test.json b/examples/calypso/input.test.json
@@ -60,8 +60,6 @@
             "_comment": "all"
         },
         "run_explore_config": {
-	    "mode": "default",
-	    "_mode": "merge",
             "template_config": {
                 "image": "registry.dp.tech/dptech/prod-11265/deepmdv3-dgpen2-calypso:v0.4",
                 "_comment": "all"
@@ -85,8 +83,7 @@
             },
             "template_slice_config": {
                 "group_size": 1,
-                "pool_size": 1,
-                "model_devi_group_size": 1
+                "pool_size": 1
             },
             "_comment": "all"
         },
@@ -413,10 +410,12 @@
         "_comment": "all"
     },
     "explore": {
-        "_type": "lmp",
-        "type": "calypso",
+        "type": "calypso:merge",
+        "_type": "calypso:default",
+        "_type": "calypso",
         "config": {
             "_command": "lmp -var restart 0",
+            "model_devi_group_size": 5,
             "run_calypso_command": "calypso.x",
             "run_opt_command": "python -u calypso_run_opt.py model.ckpt.pt"
         },

diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py
@@ -149,6 +149,7 @@ def test_caly_evo_step_merge_merge_mode(self):
             }
         )
         explore_config = {"mode": "merge", "model_devi_group_size": 30}
+        expl_mode = explore_config.get("mode")
         caly_evo_step_op = CalyEvoStepMerge(
             name="caly-evo-step",
             collect_run_caly=MockedCollRunCaly,
@@ -164,7 +165,7 @@ def test_caly_evo_step_merge_merge_mode(self):
             caly_evo_step_op,
             PrepCalyModelDevi,
             MockedRunCalyModelDevi,
-            explore_config=explore_config,
+            expl_mode=expl_mode,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,
@@ -204,6 +205,7 @@ def test_caly_evo_step_merge_default_mode(self):
             }
         )
         explore_config = {"mode": "default", "model_devi_group_size": 30}
+        expl_mode = explore_config.get("mode")
         caly_evo_step_op = CalyEvoStep(
             "caly-evo-run",
             MockedCollRunCaly,
@@ -219,7 +221,7 @@ def test_caly_evo_step_merge_default_mode(self):
             caly_evo_step_op,
             PrepCalyModelDevi,
             MockedRunCalyModelDevi,
-            explore_config=explore_config,
+            expl_mode=expl_mode,
             prep_config=prep_default_config,
             run_config=run_default_config,
             upload_python_packages=upload_python_packages,