Add DiffCSP as a new exploration engine (#251)

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


- **New Features**
- Introduced `variant_filter` for configuration filtering in exploration
tasks.
- Added support for the new `DiffCSP` exploration style, enhancing
exploration capabilities.
- New classes and methods for managing and executing `DiffCSP` tasks,
including `DiffCSPTaskGroup`.
  - CIF file conversion functionality added via the `DiffCSPGen` class.
- New class `DistanceConfFilter` for validating atomic configurations
based on distance criteria.
- Implemented the `RunRelax` class for structured execution of
relaxation tasks.

- **Bug Fixes**
- Enhanced configuration handling in multiple locations, addressing
previously unsupported filters.

- **Tests**
- Comprehensive test suites for `DistanceConfFilter`, `DiffCSPGen`,
`RunRelax`, and `PrepRunDiffCSP` to ensure functionality and
integration.

<!-- 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>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

dpgen2/entrypoint/args.py

@@ -355,17 +355,101 @@ def caly_args():
     ]
 
 
+def run_diffcsp_args():
+    doc_gen_tasks = "Number of DiffCSP generation tasks"
+    doc_gen_command = "Command for DiffCSP generation"
+    doc_relax_group_size = "Group size for relaxation"
+    return [
+        Argument(
+            "gen_tasks",
+            int,
+            optional=True,
+            default=1,
+            doc=doc_gen_tasks,
+        ),
+        Argument(
+            "gen_command",
+            str,
+            optional=False,
+            doc=doc_gen_command,
+        ),
+        Argument(
+            "relax_group_size",
+            int,
+            optional=True,
+            default=100,
+            doc=doc_relax_group_size,
+        ),
+    ]
+
+
+def diffcsp_args():
+    doc_config = "Configuration of DiffCSP 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_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>` "
+    )
+    doc_filters = "A list of configuration filters"
+
+    return [
+        Argument(
+            "config",
+            dict,
+            run_diffcsp_args(),
+            optional=False,
+            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("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument(
+            "filters",
+            list,
+            [],
+            [variant_filter()],
+            optional=True,
+            default=[],
+            doc=doc_filters,
+        ),
+    ]
+
+
 def variant_explore():
     doc = "The type of the exploration"
     doc_lmp = "The exploration by LAMMPS simulations"
     doc_calypso = "The exploration by CALYPSO structure prediction"
+    doc_diffcsp = "The exploration by DiffCSP"
     return Variant(
         "type",
         [
             Argument("lmp", dict, lmp_args(), doc=doc_lmp),
             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),
+            Argument("diffcsp", dict, diffcsp_args(), doc=doc_diffcsp),
         ],
         doc=doc,
     )

dpgen2/entrypoint/showkey.py

@@ -35,6 +35,6 @@ def showkey(
     all_step_keys = sum(folded_keys.values(), [])
     prt_str = print_keys_in_nice_format(
         all_step_keys,
-        ["run-train", "run-lmp", "run-fp"],
+        ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"],
     )
     print(prt_str)

dpgen2/entrypoint/submit.py

@@ -84,7 +84,9 @@
     LmpTemplateTaskGroup,
     NPTTaskGroup,
     caly_normalize,
+    diffcsp_normalize,
     make_calypso_task_group_from_config,
+    make_diffcsp_task_group_from_config,
     make_lmp_task_group_from_config,
     normalize_lmp_task_group_config,
 )
@@ -97,15 +99,18 @@
 from dpgen2.op import (
     CollectData,
     CollRunCaly,
+    DiffCSPGen,
     PrepCalyDPOptim,
     PrepCalyInput,
     PrepCalyModelDevi,
     PrepDPTrain,
     PrepLmp,
+    PrepRelax,
     RunCalyDPOptim,
     RunCalyModelDevi,
     RunDPTrain,
     RunLmp,
+    RunRelax,
     SelectConfs,
 )
 from dpgen2.op.caly_evo_step_merge import (
@@ -114,6 +119,7 @@
 from dpgen2.superop import (
     ConcurrentLearningBlock,
     PrepRunCaly,
+    PrepRunDiffCSP,
     PrepRunDPTrain,
     PrepRunFp,
     PrepRunLmp,
@@ -223,6 +229,16 @@ def make_concurrent_learning_op(
             run_config=run_explore_config,
             upload_python_packages=upload_python_packages,
         )
+    elif explore_style == "diffcsp":
+        prep_run_explore_op = PrepRunDiffCSP(
+            "prep-run-diffcsp",
+            DiffCSPGen,
+            PrepRelax,
+            RunRelax,
+            prep_config=prep_explore_config,
+            run_config=run_explore_config,
+            upload_python_packages=upload_python_packages,
+        )
     else:
         raise RuntimeError(f"unknown explore_style {explore_style}")
 
@@ -269,12 +285,10 @@ def make_naive_exploration_scheduler(
 
     if explore_style == "lmp":
         return make_lmp_naive_exploration_scheduler(config)
-    elif "calypso" in explore_style:
-        return make_calypso_naive_exploration_scheduler(config)
+    elif "calypso" in explore_style or explore_style == "diffcsp":
+        return make_naive_exploration_scheduler_without_conf(config, explore_style)
     else:
-        raise KeyError(
-            f"Unknown key `{explore_style}`, Only support `lmp`, `calypso`, `calypso:merge` and `calypso:default`."
-        )
+        raise KeyError(f"Unknown explore_style `{explore_style}`")
 
 
 def get_conf_filters(config):
@@ -288,7 +302,7 @@ def get_conf_filters(config):
     return conf_filters
 
 
-def make_calypso_naive_exploration_scheduler(config):
+def make_naive_exploration_scheduler_without_conf(config, explore_style):
     model_devi_jobs = config["explore"]["stages"]
     fp_task_max = config["fp"]["task_max"]
     max_numb_iter = config["explore"]["max_numb_iter"]
@@ -300,6 +314,7 @@ def make_calypso_naive_exploration_scheduler(config):
     # report
     conv_style = convergence.pop("type")
     report = conv_styles[conv_style](**convergence)
+    # trajectory render, the format of the output trajs are assumed to be lammps/dump
     render = TrajRenderLammps(nopbc=output_nopbc)
     # selector
     selector = ConfSelectorFrames(
@@ -317,9 +332,16 @@ def make_calypso_naive_exploration_scheduler(config):
         # stage
         stage = ExplorationStage()
         for jj in job:
-            jconf = caly_normalize(jj)
-            # make task group
-            tgroup = make_calypso_task_group_from_config(jconf)
+            if "calypso" in explore_style:
+                jconf = caly_normalize(jj)
+                # make task group
+                tgroup = make_calypso_task_group_from_config(jconf)
+            elif explore_style == "diffcsp":
+                jconf = diffcsp_normalize(jj)
+                # make task group
+                tgroup = make_diffcsp_task_group_from_config(jconf)
+            else:
+                raise KeyError(f"Unknown explore_style `{explore_style}`")
             # add the list to task group
             tasks = tgroup.make_task()
             stage.add_task_group(tasks)
@@ -800,6 +822,12 @@ def get_superop(key):
         return re.sub("run-caly-model-devi-[0-9]*", "prep-run-explore", key)
     elif "caly-evo-step" in key:
         return re.sub("caly-evo-step-[0-9]*", "prep-run-explore", key)
+    elif "diffcsp-gen-" in key:
+        return re.sub("diffcsp-gen-[0-9]*", "prep-run-explore", key)
+    elif "prep-relax" in key:
+        return re.sub("prep-relax", "prep-run-explore", key)
+    elif "run-relax-" in key:
+        return re.sub("run-relax-[0-9]*", "prep-run-explore", key)
     return None
 
 
@@ -843,6 +871,9 @@ def get_resubmit_keys(
         "prep-run-explore",
         "prep-lmp",
         "run-lmp",
+        "diffcsp-gen",
+        "prep-relax",
+        "run-relax",
         "select-confs",
         "prep-run-fp",
         "prep-fp",
@@ -880,7 +911,7 @@ def get_resubmit_keys(
     )
     all_step_keys = sort_slice_ops(
         all_step_keys,
-        ["run-train", "run-lmp", "run-fp"],
+        ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"],
     )
     folded_keys = fold_keys(all_step_keys)
     return folded_keys
@@ -905,7 +936,7 @@ def resubmit_concurrent_learning(
     if list_steps:
         prt_str = print_keys_in_nice_format(
             all_step_keys,
-            ["run-train", "run-lmp", "run-fp"],
+            ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"],
         )
         print(prt_str)
 

dpgen2/exploration/task/__init__.py

@@ -7,19 +7,24 @@
 from .customized_lmp_template_task_group import (
     CustomizedLmpTemplateTaskGroup,
 )
+from .diffcsp_task_group import (
+    DiffCSPTaskGroup,
+)
 from .lmp_template_task_group import (
     LmpTemplateTaskGroup,
 )
 from .make_task_group_from_config import (
     caly_normalize,
     caly_task_group_args,
+    diffcsp_normalize,
 )
 from .make_task_group_from_config import (
     lmp_normalize as normalize_lmp_task_group_config,
 )
 from .make_task_group_from_config import (
     lmp_task_group_args,
     make_calypso_task_group_from_config,
+    make_diffcsp_task_group_from_config,
     make_lmp_task_group_from_config,
     variant_task_group,
 )

dpgen2/exploration/task/diffcsp_task_group.py

@@ -0,0 +1,47 @@
+from typing import (
+    Optional,
+)
+
+from .task import (
+    ExplorationTask,
+)
+from .task_group import (
+    ExplorationTaskGroup,
+)
+
+
+class DiffCSPTaskGroup(ExplorationTaskGroup):
+    def __init__(
+        self,
+        trj_freq: int = 10,
+        fmax: float = 1e-4,
+        steps: int = 200,
+        timeout: Optional[int] = None,
+    ):
+        super().__init__()
+        self.trj_freq = trj_freq
+        self.fmax = fmax
+        self.steps = steps
+        self.timeout = timeout
+
+    def make_task(self) -> "DiffCSPTaskGroup":
+        """
+        Make the DiffCSP task group.
+
+        Returns
+        -------
+        task_grp: DiffCSPTaskGroup
+            Return one DiffCSP task group.
+        """
+        # clear all existing tasks
+        self.clear()
+        self.add_task(self._make_diffcsp_task())
+        return self
+
+    def _make_diffcsp_task(self) -> ExplorationTask:
+        task = ExplorationTask()
+        task.trj_freq = self.trj_freq  # type: ignore
+        task.fmax = self.fmax  # type: ignore
+        task.steps = self.steps  # type: ignore
+        task.timeout = self.timeout  # type: ignore
+        return task