add distance conf filter

dpgen2/entrypoint/args.py

@@ -19,6 +19,9 @@
 from dpgen2.exploration.report import (
     conv_styles,
 )
+from dpgen2.exploration.selector import (
+    conf_filter_styles,
+)
 from dpgen2.fp import (
     fp_styles,
 )
@@ -174,6 +177,25 @@ def variant_conf():
     )
 
 
+def variant_filter():
+    doc = "the type of the configuration filter."
+    var_list = []
+    for kk in conf_filter_styles.keys():
+        var_list.append(
+            Argument(
+                kk,
+                dict,
+                conf_filter_styles[kk].args(),
+                doc="Configuration filter of type %s" % kk,
+            )
+        )
+    return Variant(
+        "type",
+        var_list,
+        doc=doc,
+    )
+
+
 def lmp_args():
     doc_config = "Configuration of lmp exploration"
     doc_max_numb_iter = "Maximum number of iterations per stage"
@@ -189,6 +211,7 @@ def lmp_args():
         "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(
@@ -227,6 +250,15 @@ def lmp_args():
             alias=["configuration"],
         ),
         Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument(
+            "filters",
+            list,
+            [],
+            [variant_filter()],
+            optional=True,
+            default=[],
+            doc=doc_filters,
+        ),
     ]
 
 
@@ -272,6 +304,7 @@ def caly_args():
         "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(
@@ -310,6 +343,15 @@ def caly_args():
             alias=["configuration"],
         ),
         Argument("stages", List[List[dict]], optional=False, doc=doc_stages),
+        Argument(
+            "filters",
+            list,
+            [],
+            [variant_filter()],
+            optional=True,
+            default=[],
+            doc=doc_filters,
+        ),
     ]
 
 

dpgen2/entrypoint/submit.py

@@ -5,6 +5,9 @@
 import os
 import pickle
 import re
+from copy import (
+    deepcopy,
+)
 from pathlib import (
     Path,
 )
@@ -70,7 +73,9 @@
     ExplorationScheduler,
 )
 from dpgen2.exploration.selector import (
+    ConfFilters,
     ConfSelectorFrames,
+    conf_filter_styles,
 )
 from dpgen2.exploration.task import (
     CustomizedLmpTemplateTaskGroup,
@@ -272,13 +277,25 @@
         )
 
 
+def get_conf_filters(config):
+    conf_filters = None
+    if len(config) > 0:
+        conf_filters = ConfFilters()
+        for c in config:
+            c = deepcopy(c)
+            conf_filter = conf_filter_styles[c.pop("type")](**c)
+            conf_filters.add(conf_filter)
+    return conf_filters
+
+
 def make_calypso_naive_exploration_scheduler(config):
     model_devi_jobs = config["explore"]["stages"]
     fp_task_max = config["fp"]["task_max"]
     max_numb_iter = config["explore"]["max_numb_iter"]
     fatal_at_max = config["explore"]["fatal_at_max"]
     convergence = config["explore"]["convergence"]
     output_nopbc = config["explore"]["output_nopbc"]
+    conf_filters = get_conf_filters(config["explore"]["filters"])
     scheduler = ExplorationScheduler()
     # report
     conv_style = convergence.pop("type")
@@ -289,6 +306,7 @@
         render,
         report,
         fp_task_max,
+        conf_filters,
     )
 
     for job_ in model_devi_jobs:
@@ -329,6 +347,7 @@
     fatal_at_max = config["explore"]["fatal_at_max"]
     convergence = config["explore"]["convergence"]
     output_nopbc = config["explore"]["output_nopbc"]
+    conf_filters = get_conf_filters(config["explore"]["filters"])
     scheduler = ExplorationScheduler()
     # report
     conv_style = convergence.pop("type")
@@ -339,6 +358,7 @@
         render,
         report,
         fp_task_max,
+        conf_filters,
     )
 
     sys_configs_lmp = []

dpgen2/exploration/render/traj_render_lammps.py

@@ -64,7 +64,6 @@
         type_map: Optional[List[str]] = None,
         conf_filters: Optional["ConfFilters"] = None,
     ) -> dpdata.MultiSystems:
-        del conf_filters  # by far does not support conf filters
         ntraj = len(trajs)
         traj_fmt = "lammps/dump"
         ms = dpdata.MultiSystems(type_map=type_map)
@@ -74,4 +73,11 @@
                 ss.nopbc = self.nopbc
                 ss = ss.sub_system(id_selected[ii])
                 ms.append(ss)
+        if conf_filters is not None:
+            ms2 = dpdata.MultiSystems(type_map=type_map)
+            for s in ms:
+                s2 = conf_filters.check(s)
+                if len(s2) > 0:
+                    ms2.append(s2)
+            ms = ms2
         return ms

dpgen2/exploration/selector/__init__.py

@@ -8,3 +8,10 @@
 from .conf_selector_frame import (
     ConfSelectorFrames,
 )
+from .distance_conf_filter import (
+    DistanceConfFilter,
+)
+
+conf_filter_styles = {
+    "distance": DistanceConfFilter,
+}

dpgen2/exploration/selector/conf_filter.py

@@ -67,7 +67,7 @@ def check(
                     ff.check(
                         conf["coords"][ii],
                         conf["cells"][ii],
-                        conf["atom_types"],
+                        np.array([conf["atom_names"][t] for t in conf["atom_types"]]),
                         conf.nopbc,
                     )
                     for ii in range(conf.get_nframes())