From a13fe8a49d22e035fefc3928fd8fd98cab8f9491 Mon Sep 17 00:00:00 2001 From: zhenyu wang Date: Fri, 31 May 2024 20:33:46 +0800 Subject: [PATCH] fix: mode as a input --- dpgen2/entrypoint/submit.py | 12 ++++++------ dpgen2/exploration/render/traj_render_lammps.py | 9 --------- dpgen2/superop/prep_run_calypso.py | 10 ++++------ examples/calypso/input.test.json | 11 +++++------ tests/test_prep_run_caly.py | 6 ++++-- 5 files changed, 19 insertions(+), 29 deletions(-) diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 0dad5aab..8ce882ca 100644 --- 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 index cfdefabd..5d4cb56b 100644 --- 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 index eb71643d..8723ba32 100644 --- 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 index 432d28f4..c93e658c 100644 --- 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 index ac2a4574..3320d001 100644 --- 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,