Fix caly evo step slice (deepmodeling#225)

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

- **New Features**
- Introduced new arguments and documentation for Calypso exploration
configuration.
- Enhanced trajectory parsing with additional filtering and selection
criteria.

- **Bug Fixes**
- Adjusted logic for handling `explore_style` parameter and mode
assignments.

- **Refactor**
- Renamed parameters and updated functions for consistency and clarity.
  - Removed unnecessary imports and adjusted variable names in tests.

- **Tests**
- Expanded test cases to include new configurations and additional
assertions.
- Updated test methods to handle new trajectory configurations and
verify outputs.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: zjgemi <[email protected]>
Co-authored-by: zjgemi <[email protected]>

dpgen2/entrypoint/args.py

@@ -198,16 +198,100 @@ def lmp_args():
     ]
 
 
+def run_expl_caly_conf_args():
+    doc_caly_model_devi_group_size = "group size for model deviation."
+    doc_run_calypso_command = "command of running calypso."
+    doc_caly_run_dp_opt_command = "command of running optimization with dp."
+    return [
+        Argument(
+            "model_devi_group_size",
+            int,
+            optional=True,
+            doc=doc_caly_model_devi_group_size,
+        ),
+        Argument(
+            "run_calypso_command",
+            str,
+            optional=True,
+            default="calypso.x",
+            doc=doc_run_calypso_command,
+        ),
+        Argument(
+            "run_opt_command",
+            str,
+            optional=True,
+            doc=doc_caly_run_dp_opt_command,
+        ),
+    ]
+
+
+def caly_args():
+    doc_config = "Configuration of calypso exploration"
+    doc_max_numb_iter = "Maximum number of iterations per stage"
+    doc_fatal_at_max = (
+        "Fatal when the number of iteration per stage reaches the `max_numb_iter`"
+    )
+    doc_output_nopbc = "Remove pbc of the output configurations"
+    doc_convergence = "The method of convergence check."
+    doc_configuration = "A list of initial configurations."
+    doc_stages = (
+        "The definition of exploration stages of type `List[List[ExplorationTaskGroup]`. "
+        "The outer list provides the enumeration of the exploration stages. "
+        "Then each stage is defined by a list of exploration task groups. "
+        "Each task group is described in :ref:`the task group definition<task_group_sec>` "
+    )
+
+    return [
+        Argument(
+            "config",
+            dict,
+            run_expl_caly_conf_args(),
+            optional=True,
+            default=RunLmp.normalize_config({}),
+            doc=doc_config,
+        ),
+        Argument(
+            "max_numb_iter", int, optional=True, default=10, doc=doc_max_numb_iter
+        ),
+        Argument(
+            "fatal_at_max", bool, optional=True, default=True, doc=doc_fatal_at_max
+        ),
+        Argument(
+            "output_nopbc", bool, optional=True, default=False, doc=doc_output_nopbc
+        ),
+        Argument(
+            "convergence",
+            dict,
+            [],
+            [variant_conv()],
+            optional=False,
+            doc=doc_convergence,
+        ),
+        Argument(
+            "configurations",
+            list,
+            [],
+            [variant_conf()],
+            optional=False,
+            repeat=True,
+            doc=doc_configuration,
+            alias=["configuration"],
+        ),
+        Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+    ]
+
+
 def variant_explore():
-    # TODO: add calypso_args
     doc = "The type of the exploration"
     doc_lmp = "The exploration by LAMMPS simulations"
     doc_calypso = "The exploration by CALYPSO structure prediction"
     return Variant(
         "type",
         [
             Argument("lmp", dict, lmp_args(), doc=doc_lmp),
-            Argument("calypso", dict, lmp_args(), doc=doc_calypso),
+            Argument("calypso", dict, caly_args(), doc=doc_calypso),
+            Argument("calypso:default", dict, caly_args(), doc=doc_calypso),
+            Argument("calypso:merge", dict, caly_args(), doc=doc_calypso),
         ],
         doc=doc,
     )

dpgen2/entrypoint/submit.py

@@ -154,7 +154,6 @@ def make_concurrent_learning_op(
     upload_python_packages: Optional[List[os.PathLike]] = None,
     valid_data: Optional[S3Artifact] = None,
 ):
-    expl_mode = run_explore_config.get("mode", "default")
     if train_style in ("dp", "dp-dist"):
         prep_run_train_op = PrepRunDPTrain(
             "prep-run-dp-train",
@@ -176,13 +175,15 @@ 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",
                 collect_run_caly=CollRunCaly,
                 prep_dp_optim=PrepCalyDPOptim,
                 run_dp_optim=RunCalyDPOptim,
+                expl_mode=expl_mode,
                 prep_config=prep_explore_config,
                 run_config=run_explore_config,
                 upload_python_packages=None,
@@ -193,18 +194,22 @@ def make_concurrent_learning_op(
                 collect_run_caly=CollRunCaly,
                 prep_dp_optim=PrepCalyDPOptim,
                 run_dp_optim=RunCalyDPOptim,
+                expl_mode=expl_mode,
                 prep_config=prep_explore_config,
                 run_config=run_explore_config,
                 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,
+            expl_mode=expl_mode,
             prep_config=prep_explore_config,
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
@@ -255,8 +260,12 @@ def make_naive_exploration_scheduler(
 
     if explore_style == "lmp":
         return make_lmp_naive_exploration_scheduler(config)
-    elif explore_style == "calypso":
+    elif "calypso" in explore_style:
         return make_calypso_naive_exploration_scheduler(config)
+    else:
+        raise KeyError(
+            f"Unknown key `{explore_style}`, Only support `lmp`, `calypso`, `calypso:merge` and `calypso:default`."
+        )
 
 
 def make_calypso_naive_exploration_scheduler(config):
@@ -453,6 +462,8 @@ def workflow_concurrent_learning(
 ) -> Tuple[Step, Optional[Step]]:
     default_config = config["default_step_config"]
 
+    train_config = config["train"]["config"]
+    explore_config = config["explore"]["config"]
     train_style = config["train"]["type"]
     explore_style = config["explore"]["type"]
     fp_style = config["fp"]["type"]
@@ -527,8 +538,7 @@ def workflow_concurrent_learning(
         template_script = [json.loads(Path(ii).read_text()) for ii in template_script_]
     else:
         template_script = json.loads(Path(template_script_).read_text())
-    train_config = config["train"]["config"]
-    explore_config = config["explore"]["config"]
+
     if (
         "teacher_model_path" in explore_config
         and explore_config["teacher_model_path"] is not None

dpgen2/exploration/render/traj_render_lammps.py

@@ -47,6 +47,9 @@ def get_model_devi(
 
     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
+            dd = dd.reshape((1, len(dd)))
+
         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])

dpgen2/exploration/task/__init__.py

@@ -21,7 +21,6 @@
     lmp_task_group_args,
     make_calypso_task_group_from_config,
     make_lmp_task_group_from_config,
-    make_task_group_from_config,
     variant_task_group,
 )
 from .npt_task_group import (

dpgen2/exploration/task/caly_task_group.py

@@ -107,6 +107,7 @@ def set_params(
         vsc: bool = True,
         ctrl_range: List[List[int]] = [[1, 10]],
         max_numb_atoms: int = 100,
+        opt_step: int = 1000,
     ):
         """
         Set calypso parameters
@@ -185,6 +186,7 @@ def set_params(
         self.vsc = vsc
         self.ctrl_range = ctrl_range
         self.max_numb_atoms = max_numb_atoms
+        self.opt_step = opt_step
 
         self.caly_set = True
 
@@ -234,6 +236,7 @@ def _make_caly_task(self) -> ExplorationTask:
             self.vsc,
             self.ctrl_range,
             self.max_numb_atoms,
+            opt_step=self.opt_step,
         )
         task = ExplorationTask()
         task.add_file(calypso_input_file, input_file_str)

dpgen2/exploration/task/calypso/caly_input.py

@@ -98,12 +98,11 @@ def Write_Outcar(outcar, element, ele, volume, lat, pos, ene, force, stress, pst
     enthalpy = ene + pstress * volume / 1602.17733
     f.write('enthalpy is  TOTEN    = %20.6f %20.6f\\n' % (enthalpy, enthalpy/na))
 
-def run_opt(fmax, stress):
+def run_opt(fmax, stress, opt_step):
     # Using the ASE&DP to Optimize Configures
 
     calc = DP(model=sys.argv[1])    # init the model before iteration
 
-    Opt_Step = 1000
     start = time.time()
     # pstress kbar
     pstress = stress
@@ -118,7 +117,7 @@ def run_opt(fmax, stress):
         to_be_opti.calc = calc
         ucf = UnitCellFilter(to_be_opti, scalar_pressure=aim_stress)
         opt = LBFGS(ucf,trajectory=poscar.strip("POSCAR_") + '.traj')
-        opt.run(fmax=fmax,steps=Opt_Step)
+        opt.run(fmax=fmax,steps=opt_step)
         atoms_lat = to_be_opti.cell
         atoms_pos = to_be_opti.positions
         atoms_force = to_be_opti.get_forces()
@@ -138,7 +137,7 @@ def run_opt(fmax, stress):
 
 calypso_run_opt_str_end = """
 if __name__ == '__main__':
-    run_opt(fmax=%.3f, stress=%.3f)
+    run_opt(fmax=%.3f, stress=%.3f, opt_step=%.3f)
 """
 
 calypso_check_opt_str = """#!/usr/bin/env python3
@@ -248,7 +247,7 @@ def make_calypso_input(
     max_step: int = 5,
     system_name: str = "CALYPSO",
     numb_of_formula: List[int] = [1, 1],
-    pressure: float = 0.001,
+    pressure: float = 0.001,  # KBar
     fmax: float = 0.01,
     volume: float = 0,
     ialgo: int = 2,
@@ -333,7 +332,12 @@ def make_calypso_input(
     file_str += vsc_ctrl_range + "\n"
     file_str += "@End\n"
 
-    run_opt_str = calypso_run_opt_str + calypso_run_opt_str_end % (fmax, pressure)
+    opt_step = kwargs.get("opt_step", 1000)
+    run_opt_str = calypso_run_opt_str + calypso_run_opt_str_end % (
+        fmax,
+        pressure,
+        opt_step,
+    )
     check_opt_str = calypso_check_opt_str
 
     return file_str, run_opt_str, check_opt_str

dpgen2/exploration/task/make_task_group_from_config.py

@@ -351,6 +351,13 @@ def caly_task_grp_args():
             default=0.01,
             doc="the converge criterion. The force on all individual atoms should be less than `fmax`.",
         ),
+        Argument(
+            "opt_step",
+            float,
+            optional=True,
+            default=1000,
+            doc="the converge criterion. The force on all individual atoms should be less than `fmax`.",
+        ),
         Argument(
             "volume",
             float,
@@ -492,18 +499,6 @@ def config_strip_confidx(
     return cc
 
 
-def make_task_group_from_config(
-    numb_models,
-    mass_map,
-    config,
-):
-    explore_type = config["explore"]["type"]
-    if explore_type == "lmp":
-        return make_lmp_task_group_from_config(numb_models, mass_map, config)
-    elif explore_type == "calypso":
-        return make_calypso_task_group_from_config(config)
-
-
 def make_calypso_task_group_from_config(config):
     config.pop("type", None)
     config = caly_normalize(config)

dpgen2/op/prep_caly_model_devi.py

@@ -48,7 +48,7 @@ def get_input_sign(cls):
         return OPIOSign(
             {
                 "task_name": Parameter(str),
-                "template_slice_config": Parameter(dict),
+                "config": BigParameter(dict),
                 "traj_results": Artifact(List[Path]),
             }
         )
@@ -74,7 +74,7 @@ def execute(
         ip : dict
             Input dict with components:
             - `task_name` : (`str`)
-            - `template_slice_config` : (`dict`)
+            - `config` : (`BigParameter(dict)`)
             - `traj_results` : (`Path`)
 
         Returns
@@ -93,9 +93,8 @@ def execute(
             for traj_dir in traj_results_dir
             for traj in Path(traj_dir).rglob("*.traj")
         ]
-        group_size = ip["template_slice_config"].get(
-            "model_devi_group_size", len(trajs)
-        )
+        expl_config = ip["config"]
+        group_size = expl_config.get("model_devi_group_size", len(trajs))
 
         with set_directory(work_dir):
             grouped_trajs_list = [

dpgen2/op/run_caly_dp_optim.py

@@ -102,7 +102,7 @@ def execute(
 
         config = ip["config"] if ip["config"] is not None else {}
         command = config.get(
-            f"run_opt_command", "python -u calypso_run_opt.py {model_name}"
+            "run_opt_command", "python -u calypso_run_opt.py model.ckpt.pt"
         )
 
         work_dir = Path(ip["task_name"])