From be2c7b6df30a31e2df1e680b1af0af9e25084e3c Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Sun, 16 Aug 2026 14:31:36 +0800 Subject: [PATCH 1/6] feat(cv-filter): filter candidates by PLUMED CV regions --- docs/input.md | 42 +++++ dpgen2/entrypoint/args.py | 13 ++ dpgen2/entrypoint/submit.py | 6 + dpgen2/exploration/report/report.py | 7 + .../report/report_adaptive_lower.py | 14 ++ .../report/report_trust_levels_base.py | 12 ++ dpgen2/exploration/selector/__init__.py | 3 + dpgen2/exploration/selector/conf_selector.py | 1 + .../selector/conf_selector_frame.py | 23 +++ .../exploration/selector/plumed_cv_filter.py | 144 ++++++++++++++++++ dpgen2/op/run_lmp.py | 20 ++- dpgen2/op/select_confs.py | 33 +++- dpgen2/superop/block.py | 3 + tests/exploration/test_conf_selector_frame.py | 27 +++- tests/exploration/test_plumed_cv_filter.py | 88 +++++++++++ .../exploration/test_report_adaptive_lower.py | 15 ++ tests/exploration/test_report_trust_levels.py | 8 + tests/op/test_run_lmp.py | 25 +++ tests/test_select_confs.py | 20 +++ 19 files changed, 495 insertions(+), 9 deletions(-) create mode 100644 dpgen2/exploration/selector/plumed_cv_filter.py create mode 100644 tests/exploration/test_plumed_cv_filter.py diff --git a/docs/input.md b/docs/input.md index 6fc03bd1..382711ff 100644 --- a/docs/input.md +++ b/docs/input.md @@ -131,6 +131,48 @@ The {dargs:argument}`"stages"` defines the exploration stag The {dargs:argument}`"n_sample"` tells the number of confgiruations randomly sampled from the set picked by {dargs:argument}`"conf_idx"` from {dargs:argument}`"configurations"` for each exploration task. All configurations has the equal possibility to be sampled. The default value of `"n_sample"` is `null`, in this case all picked configurations are sampled. In the example, we have 3 samples for stage 0 task group 0 and 2 thermodynamic states (NVT, T=50 and 100K), then the task group has 3x2=6 NVT DPMD tasks. +#### PLUMED CV candidate filtering + +LAMMPS exploration candidates can be restricted to a union of named PLUMED CV +regions before the configured random or maximum-model-deviation selection: + +```json +"explore": { + "config": { + "plm_output_file": "COLVAR" + }, + "cv_filter": { + "regions": [ + {"d": [0.08, 0.12]}, + {"v": [1.8, 2.2]} + ] + } +} +``` + +Each region is a mapping from a field in the PLUMED `#! FIELDS` header to a +lower-inclusive, upper-exclusive interval. Conditions within a region are +combined by AND and regions are combined by OR. The PLUMED input must write the +selected fields with the same stride as `trj_freq`, for example: + +```plumed +LOAD FILE=/absolute/path/ReactiveVoronoi.so +d: DISTANCE ATOMS=1,2 +v: VORONOI_COORDINATION ... +PRINT ARG=d,v STRIDE=10 FILE=COLVAR +``` + +`LOAD` is only needed for CVs that are not built into the active PLUMED. Build +such a plugin with that same PLUMED installation (for example, `plumed mklib +ReactiveVoronoi.cpp`); shared libraries from a different compiler or PLUMED +build may be ABI-incompatible. + +The region bounds use the units written to `COLVAR` (PLUMED's default length +unit is nm). The filter fails if the file, field, finite values, strictly +increasing `time`, or row-to-trajectory alignment is invalid. Model-deviation +trust levels are applied first, followed by the CV regions and then the existing +candidate limit and selection policy. + ### FP diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..ddc1610a 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -20,6 +20,7 @@ conv_styles, ) from dpgen2.exploration.selector import ( + PlumedCVFilter, conf_filter_styles, ) from dpgen2.fp import ( @@ -212,6 +213,10 @@ def lmp_args(): "Each task group is described in :ref:`the task group definition` " ) doc_filters = "A list of configuration filters" + doc_cv_filter = ( + "Optional PLUMED CV candidate filter. The regions key is a list of " + "field-to-[lower, upper] mappings; regions are combined by OR." + ) return [ Argument( @@ -259,6 +264,14 @@ def lmp_args(): default=[], doc=doc_filters, ), + Argument( + "cv_filter", + dict, + PlumedCVFilter.args(), + optional=True, + default=None, + doc=doc_cv_filter, + ), ] diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 9e07374f..597c7f88 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -68,6 +68,7 @@ from dpgen2.exploration.selector import ( ConfFilters, ConfSelectorFrames, + PlumedCVFilter, conf_filter_styles, ) from dpgen2.exploration.task import ( @@ -366,6 +367,10 @@ def make_lmp_naive_exploration_scheduler(config): convergence = config["explore"]["convergence"] output_nopbc = config["explore"]["output_nopbc"] conf_filters = get_conf_filters(config["explore"]["filters"]) + cv_filter_config = config["explore"]["cv_filter"] + cv_filter = ( + PlumedCVFilter(**cv_filter_config) if cv_filter_config is not None else None + ) use_ele_temp = config["inputs"]["use_ele_temp"] scheduler = ExplorationScheduler() # report @@ -378,6 +383,7 @@ def make_lmp_naive_exploration_scheduler(config): report, fp_task_max, conf_filters, + cv_filter, ) sys_configs_lmp = [] diff --git a/dpgen2/exploration/report/report.py b/dpgen2/exploration/report/report.py index d1c43fe6..066d75f8 100644 --- a/dpgen2/exploration/report/report.py +++ b/dpgen2/exploration/report/report.py @@ -65,6 +65,13 @@ def no_candidate(self) -> bool: r"""If no candidate configuration is found""" pass + def restrict_candidate_ids( + self, + allowed_ids: List[List[int]], + ) -> None: + r"""Restrict model-deviation candidates before sampling.""" + raise NotImplementedError + @abstractmethod def get_candidate_ids( self, diff --git a/dpgen2/exploration/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index cd2989f3..42e20b63 100644 --- a/dpgen2/exploration/report/report_adaptive_lower.py +++ b/dpgen2/exploration/report/report_adaptive_lower.py @@ -371,6 +371,20 @@ def candidate_ratio( def no_candidate(self) -> bool: return self._no_candidate + def restrict_candidate_ids( + self, + allowed_ids: List[List[int]], + ) -> None: + if len(allowed_ids) != self.ntraj: + raise FatalError("candidate filter and trajectories have different lengths") + allowed = { + (traj_idx, frame_idx) + for traj_idx, frame_ids in enumerate(allowed_ids) + for frame_idx in frame_ids + } + self.candi &= allowed + self._no_candidate = len(self.candi) == 0 + def get_candidate_ids( self, max_nframes: Optional[int] = None, diff --git a/dpgen2/exploration/report/report_trust_levels_base.py b/dpgen2/exploration/report/report_trust_levels_base.py index 185ea5b1..3e28d4ef 100644 --- a/dpgen2/exploration/report/report_trust_levels_base.py +++ b/dpgen2/exploration/report/report_trust_levels_base.py @@ -236,6 +236,18 @@ def candidate_ratio( def no_candidate(self) -> bool: return self._no_candidate + def restrict_candidate_ids( + self, + allowed_ids: List[List[int]], + ) -> None: + if len(allowed_ids) != len(self.traj_cand): + raise FatalError("candidate filter and trajectories have different lengths") + self.traj_cand = [ + candidates & set(allowed) + for candidates, allowed in zip(self.traj_cand, allowed_ids) + ] + self._no_candidate = sum(len(candidates) for candidates in self.traj_cand) == 0 + @abstractmethod def get_candidate_ids( self, diff --git a/dpgen2/exploration/selector/__init__.py b/dpgen2/exploration/selector/__init__.py index cfed094f..ce015a2d 100644 --- a/dpgen2/exploration/selector/__init__.py +++ b/dpgen2/exploration/selector/__init__.py @@ -13,6 +13,9 @@ BoxSkewnessConfFilter, DistanceConfFilter, ) +from .plumed_cv_filter import ( + PlumedCVFilter, +) conf_filter_styles = { "distance": DistanceConfFilter, diff --git a/dpgen2/exploration/selector/conf_selector.py b/dpgen2/exploration/selector/conf_selector.py index f24a7d31..22b7be64 100644 --- a/dpgen2/exploration/selector/conf_selector.py +++ b/dpgen2/exploration/selector/conf_selector.py @@ -37,5 +37,6 @@ def select( model_devis: Union[List[Path], List[HDF5Dataset]], type_map: Optional[List[str]] = None, optional_outputs: Optional[List[Path]] = None, + plm_outputs: Optional[List[Path]] = None, ) -> Tuple[List[Path], ExplorationReport]: pass diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index fc116f88..35de3b81 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -14,10 +14,16 @@ import dpdata import numpy as np +from dflow.python import ( + FatalError, +) from dflow.python.opio import ( HDF5Dataset, ) +from dpgen2.exploration.deviation import ( + DeviManager, +) from dpgen2.exploration.render import ( TrajRender, ) @@ -29,6 +35,9 @@ ConfFilters, ConfSelector, ) +from .plumed_cv_filter import ( + PlumedCVFilter, +) class ConfSelectorFrames(ConfSelector): @@ -48,11 +57,13 @@ def __init__( report: ExplorationReport, max_numb_sel: Optional[int] = None, conf_filters: Optional[ConfFilters] = None, + plumed_cv_filter: Optional[PlumedCVFilter] = None, ): self.max_numb_sel = max_numb_sel self.conf_filters = conf_filters self.traj_render = traj_render self.report = report + self.plumed_cv_filter = plumed_cv_filter def select( self, @@ -60,6 +71,7 @@ def select( model_devis: Union[List[Path], List[HDF5Dataset]], type_map: Optional[List[str]] = None, optional_outputs: Optional[List[Path]] = None, + plm_outputs: Optional[List[Path]] = None, ) -> Tuple[List[Path], ExplorationReport]: """Select configurations @@ -92,6 +104,17 @@ def select( self.report.clear() self.report.record(md_model_devi) + if self.plumed_cv_filter is not None: + if plm_outputs is None or any(output is None for output in plm_outputs): + raise FatalError( + "PLUMED CV filtering requires one output per trajectory" + ) + md_f = md_model_devi.get(DeviManager.MAX_DEVI_F) + allowed_ids = self.plumed_cv_filter.get_selected_ids( + plm_outputs, + [len(values) for values in md_f], + ) + self.report.restrict_candidate_ids(allowed_ids) id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) ms = self.traj_render.get_confs( diff --git a/dpgen2/exploration/selector/plumed_cv_filter.py b/dpgen2/exploration/selector/plumed_cv_filter.py new file mode 100644 index 00000000..865a0ca2 --- /dev/null +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -0,0 +1,144 @@ +from pathlib import ( + Path, +) +from typing import ( + Dict, + List, + Sequence, +) + +import numpy as np +from dargs import ( + Argument, +) +from dflow.python import ( + FatalError, +) + + +class PlumedCVFilter: + """Select frames in a union of PLUMED CV regions. + + Each region maps field names from ``#! FIELDS`` to a lower-inclusive, + upper-exclusive interval. A frame is selected when it matches every + interval in any one region. + """ + + @staticmethod + def args() -> List[Argument]: + return [ + Argument( + "regions", + list, + optional=False, + doc=( + "A list of PLUMED field-to-[lower, upper] mappings. " + "Fields within a region are ANDed; regions are ORed." + ), + ) + ] + + def __init__(self, regions: List[Dict[str, Sequence[float]]]): + if ( + not isinstance(regions, list) + or not regions + or any(not isinstance(region, dict) or not region for region in regions) + ): + raise ValueError("PLUMED CV regions must be a non-empty list of dicts") + self.regions = [] + for region in regions: + normalized = {} + for field, bounds in region.items(): + if not isinstance(field, str) or not field: + raise ValueError("PLUMED field names must be non-empty strings") + if not isinstance(bounds, (list, tuple)) or len(bounds) != 2: + raise ValueError(f"invalid interval for PLUMED field {field!r}") + try: + normalized[field] = (float(bounds[0]), float(bounds[1])) + except (TypeError, ValueError) as exc: + raise ValueError( + f"invalid interval for PLUMED field {field!r}" + ) from exc + if normalized[field][0] >= normalized[field][1]: + raise ValueError(f"invalid interval for PLUMED field {field!r}") + if not np.all(np.isfinite(normalized[field])): + raise ValueError(f"non-finite interval for PLUMED field {field!r}") + self.regions.append(normalized) + + def get_selected_ids( + self, + files: List[Path], + nframes: List[int], + ) -> List[List[int]]: + if len(files) != len(nframes): + raise FatalError("PLUMED outputs and trajectories have different lengths") + selected = [] + for file, expected_nframes in zip(files, nframes): + fields, values = self._read(file) + if len(values) != expected_nframes: + raise FatalError( + f"PLUMED output {file} has {len(values)} rows, expected " + f"{expected_nframes}; PRINT STRIDE must match the trajectory stride" + ) + field_idx = {field: idx for idx, field in enumerate(fields)} + keep = np.zeros(expected_nframes, dtype=bool) + for region in self.regions: + in_region = np.ones(expected_nframes, dtype=bool) + for field, (lower, upper) in region.items(): + if field not in field_idx: + raise FatalError( + f"PLUMED field {field!r} is missing from {file}" + ) + column = values[:, field_idx[field]] + in_region &= (column >= lower) & (column < upper) + keep |= in_region + selected.append(np.flatnonzero(keep).tolist()) + return selected + + @staticmethod + def _read(file: Path): + fields = None + rows = [] + try: + with open(file, encoding="utf8") as handle: + for line_number, line in enumerate(handle, 1): + words = line.split() + if not words: + continue + if words[:2] == ["#!", "FIELDS"]: + new_fields = words[2:] + if fields is not None and fields != new_fields: + raise FatalError( + f"inconsistent PLUMED FIELDS headers in {file}" + ) + fields = new_fields + elif words[0].startswith("#"): + continue + else: + if fields is None: + raise FatalError( + f"PLUMED numeric row precedes FIELDS header in {file}" + ) + try: + rows.append([float(value) for value in words]) + except ValueError as exc: + raise FatalError( + f"invalid PLUMED numeric row {line_number} in {file}" + ) from exc + except OSError as exc: + raise FatalError(f"cannot read PLUMED output {file}: {exc}") from exc + if fields is None: + raise FatalError(f"PLUMED FIELDS header is missing from {file}") + if not fields or len(fields) != len(set(fields)): + raise FatalError(f"PLUMED FIELDS must be non-empty and unique in {file}") + if any(len(row) != len(fields) for row in rows): + raise FatalError(f"PLUMED row width does not match FIELDS in {file}") + values = np.asarray(rows, dtype=float).reshape((-1, len(fields))) + if not np.all(np.isfinite(values)): + raise FatalError(f"non-finite PLUMED values in {file}") + if "time" not in fields: + raise FatalError(f"PLUMED time field is missing from {file}") + time = values[:, fields.index("time")] + if len(time) > 1 and np.any(np.diff(time) <= 0): + raise FatalError(f"PLUMED time must be strictly increasing in {file}") + return fields, values diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 60cd9305..885c750a 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -122,6 +122,7 @@ def execute( command = config["command"] teacher_model: Optional[BinaryFileInput] = config["teacher_model_path"] shuffle_models: Optional[bool] = config["shuffle_models"] + plm_output_file = config["plm_output_file"] task_name = ip["task_name"] task_path = ip["task_path"] models = ip["models"] @@ -206,8 +207,8 @@ def execute( "model_devi": self.get_model_devi(work_dir / lmp_model_devi_name), } plm_output = ( - {"plm_output": work_dir / plm_output_name} - if (work_dir / plm_output_name).is_file() + {"plm_output": work_dir / plm_output_file} + if (work_dir / plm_output_file).is_file() else {} ) ret_dict.update(plm_output) @@ -231,6 +232,10 @@ def lmp_args(): doc_head = "Select a head from multitask" doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature" doc_use_hdf5 = "Use HDF5 to store trajs and model_devis" + doc_plm_output_file = ( + "PLUMED output artifact to collect. Set this to the FILE used by " + "PLUMED PRINT when filtering candidates by CV." + ) doc_extra_output_files = "Extra output file names, support wildcards" return [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), @@ -262,6 +267,13 @@ def lmp_args(): default=False, doc=doc_use_hdf5, ), + Argument( + "plm_output_file", + str, + optional=True, + default=plm_output_name, + doc=doc_plm_output_file, + ), Argument( "extra_output_files", list, @@ -277,6 +289,10 @@ def normalize_config(data={}): base = Argument("base", dict, ta) data = base.normalize_value(data, trim_pattern="_*") base.check_value(data, strict=True) + if data["plm_output_file"] in {"", ".", ".."} or ( + Path(data["plm_output_file"]).name != data["plm_output_file"] + ): + raise ValueError("plm_output_file must be a file name, not a path") return data diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index e8ba891d..65e0b5a8 100644 --- a/dpgen2/op/select_confs.py +++ b/dpgen2/op/select_confs.py @@ -40,6 +40,7 @@ def get_input_sign(cls): "trajs": Artifact(Union[List[Path], HDF5Datasets]), "model_devis": Artifact(Union[List[Path], HDF5Datasets]), "optional_outputs": Artifact(List[Path], optional=True), + "plm_outputs": Artifact(List[Path], optional=True), } ) @@ -84,16 +85,19 @@ def execute( trajs = ip["trajs"] model_devis = ip["model_devis"] optional_outputs = ip["optional_outputs"] + plm_outputs = ip["plm_outputs"] + plm_outputs = SelectConfs.validate_plm_outputs(trajs, model_devis, plm_outputs) trajs, model_devis, optional_outputs = SelectConfs.validate_trajs( trajs, model_devis, optional_outputs ) - confs, report = conf_selector.select( - trajs, - model_devis, - type_map=type_map, - optional_outputs=optional_outputs, - ) + select_kwargs = { + "type_map": type_map, + "optional_outputs": optional_outputs, + } + if plm_outputs is not None: + select_kwargs["plm_outputs"] = plm_outputs + confs, report = conf_selector.select(trajs, model_devis, **select_kwargs) return OPIO( { @@ -144,3 +148,20 @@ def validate_trajs( else: raise FatalError(f"trajs frame is {tt} while model_devis frame is {mm}") return rett, retm, reto + + @staticmethod + def validate_plm_outputs(trajs, model_devis, plm_outputs=None): + if plm_outputs is None: + return None + if len(trajs) != len(plm_outputs): + raise FatalError("length of trajs list is not equal to the plm_output list") + ret = [] + for traj, model_devi, plm_output in zip(trajs, model_devis, plm_outputs): + if traj is None and model_devi is None: + if plm_output is not None: + raise FatalError( + f"trajs frame is {traj} while plm_outputs frame is {plm_output}" + ) + elif traj is not None and model_devi is not None: + ret.append(plm_output) + return ret diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py index 0e39ab38..916df329 100644 --- a/dpgen2/superop/block.py +++ b/dpgen2/superop/block.py @@ -269,6 +269,9 @@ def _block_cl( "optional_outputs": prep_run_explore.outputs.artifacts["optional_outputs"] if "optional_outputs" in prep_run_explore.outputs.artifacts else None, + "plm_outputs": prep_run_explore.outputs.artifacts["plm_output"] + if "plm_output" in prep_run_explore.outputs.artifacts + else None, }, key=step_keys["select-confs"], executor=select_confs_executor, diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py index 52b30248..e62ae7d4 100644 --- a/tests/exploration/test_conf_selector_frame.py +++ b/tests/exploration/test_conf_selector_frame.py @@ -17,10 +17,12 @@ TrajRenderLammps, ) from dpgen2.exploration.report import ( + ExplorationReportTrustLevelsMax, ExplorationReportTrustLevelsRandom, ) from dpgen2.exploration.selector import ( ConfSelectorFrames, + PlumedCVFilter, ) # isort: on @@ -85,13 +87,36 @@ def setUp(self): self.type_map = ["O", "H"] def tearDown(self): - for ii in ["foo.dump", "bar.dump", "foo.md", "bar.md"]: + for ii in ["foo.dump", "bar.dump", "foo.md", "bar.md", "foo.cv", "bar.cv"]: if Path(ii).is_file(): os.remove(ii) for ii in ["confs"]: if Path(ii).is_dir(): shutil.rmtree(ii) + def test_plumed_filter_precedes_max_selection(self): + plm_outputs = [Path("foo.cv"), Path("bar.cv")] + for output in plm_outputs: + output.write_text( + "#! FIELDS time cv\n0.0 0.5\n1.0 0.5\n2.0 1.5\n" + ) + conf_selector = ConfSelectorFrames( + TrajRenderLammps(), + ExplorationReportTrustLevelsMax(0.1, 0.5), + max_numb_sel=1, + plumed_cv_filter=PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]), + ) + confs, _ = conf_selector.select( + self.trajs, + self.model_devis, + self.type_map, + plm_outputs=plm_outputs, + ) + ms = dpdata.MultiSystems(type_map=self.type_map) + ms.from_deepmd_npy(confs[0], labeled=False) + self.assertEqual(ms[0].get_nframes(), 1) + self.assertAlmostEqual(ms[0]["coords"][0][0][1], 3.87, places=2) + def test_f_0(self): report = ExplorationReportTrustLevelsRandom(0.1, 0.5, conv_accuracy=0.9) traj_render = TrajRenderLammps() diff --git a/tests/exploration/test_plumed_cv_filter.py b/tests/exploration/test_plumed_cv_filter.py new file mode 100644 index 00000000..e4cfb3d3 --- /dev/null +++ b/tests/exploration/test_plumed_cv_filter.py @@ -0,0 +1,88 @@ +import tempfile +import unittest +from pathlib import ( + Path, +) + +from dargs import ( + Argument, +) +from dflow.python import ( + FatalError, +) + +from dpgen2.exploration.selector import ( + PlumedCVFilter, +) + + +class TestPlumedCVFilter(unittest.TestCase): + def test_config_schema(self): + schema = Argument("cv_filter", dict, PlumedCVFilter.args()) + config = schema.normalize_value( + {"regions": [{"distance": [0.8, 1.2]}]} + ) + schema.check_value(config, strict=True) + self.assertEqual(config["regions"][0]["distance"], [0.8, 1.2]) + + def test_union_of_regions(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time d1 d2\n" "0.0 0.5 2.0\n" "1.0 0.5 4.0\n" "2.0 2.5 9.0\n" + ) + cv_filter = PlumedCVFilter( + regions=[{"d1": [0.0, 1.0], "d2": [1.0, 3.0]}, {"d1": [2.0, 3.0]}] + ) + self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 2]]) + + def test_intervals_are_lower_inclusive_and_upper_exclusive(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv\n0.0 1.0\n1.0 2.0\n2.0 3.0\n" + ) + cv_filter = PlumedCVFilter(regions=[{"cv": [1.0, 3.0]}]) + self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 1]]) + + def test_alignment_fails_closed(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text("#! FIELDS time cv\n0.0 0.5\n") + with self.assertRaises(FatalError): + PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + [output], [2] + ) + + def test_invalid_config_and_file_fail_closed(self): + invalid_regions = [ + [], + [{}], + [{"cv": [1.0, 1.0]}], + [{"cv": [0.0, float("inf")]}], + [{"cv": "01"}], + ] + for regions in invalid_regions: + with self.subTest(regions=regions), self.assertRaises(ValueError): + PlumedCVFilter(regions=regions) + + invalid_outputs = [ + "0.0 0.5\n#! FIELDS time cv\n", + "#! FIELDS time cv\n0.0 nan\n", + "#! FIELDS time cv\n0.0\n", + "#! FIELDS time cv\n#! FIELDS time other\n0.0 0.5\n", + "#! FIELDS cv\n0.5\n", + "#! FIELDS time cv\n1.0 0.5\n0.0 0.5\n", + ] + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(FatalError): + PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + [Path(tmpdir) / "missing"], [1] + ) + for index, content in enumerate(invalid_outputs): + output = Path(tmpdir) / f"COLVAR.{index}" + output.write_text(content) + with self.subTest(content=content), self.assertRaises(FatalError): + PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + [output], [1] + ) diff --git a/tests/exploration/test_report_adaptive_lower.py b/tests/exploration/test_report_adaptive_lower.py index b5f123ca..78a1426d 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -27,6 +27,21 @@ class TestTrajsExplorationReport(unittest.TestCase): + def test_candidate_filter(self): + model_devi = DeviManagerStd() + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.10, 0.20, 0.30]), + ) + report = ExplorationReportAdaptiveLower( + level_f_hi=1.0, + numb_candi_f=2, + rate_candi_f=0.0, + ) + report.record(model_devi) + report.restrict_candidate_ids([[0, 1]]) + self.assertEqual(report.get_candidate_ids(), [[1]]) + def test_fv(self): model_devi = DeviManagerStd() model_devi.add( diff --git a/tests/exploration/test_report_trust_levels.py b/tests/exploration/test_report_trust_levels.py index 7f6c79a9..fb2c6a0f 100644 --- a/tests/exploration/test_report_trust_levels.py +++ b/tests/exploration/test_report_trust_levels.py @@ -217,6 +217,14 @@ def test_max_selection(self): npicked += 1 self.assertEqual(npicked, 2) + def test_max_selection_after_candidate_filter(self): + model_devi = DeviManagerStd() + model_devi.add(DeviManager.MAX_DEVI_F, np.array([0.40, 0.50, 0.55])) + ter = ExplorationReportTrustLevelsMax(0.3, 0.6) + ter.record(model_devi) + ter.restrict_candidate_ids([[0, 1]]) + self.assertEqual(ter.get_candidate_ids(1), [[1]]) + def test_random_selection_convergence(self): # case 1 model_devi = DeviManagerStd() diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 650fd82e..b10d0a42 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -67,6 +67,31 @@ def tearDown(self): if Path(self.task_name).is_dir(): shutil.rmtree(self.task_name) + def test_plm_output_file_config(self): + config = RunLmp.normalize_config({"plm_output_file": "COLVAR"}) + self.assertEqual(config["plm_output_file"], "COLVAR") + with self.assertRaises(ValueError): + RunLmp.normalize_config({"plm_output_file": "outputs/COLVAR"}) + for invalid_name in ["", ".", ".."]: + with self.subTest(invalid_name=invalid_name), self.assertRaises(ValueError): + RunLmp.normalize_config({"plm_output_file": invalid_name}) + + @patch("dpgen2.op.run_lmp.run_command") + def test_plm_output_file_collection(self, mocked_run): + mocked_run.return_value = (0, "", "") + (self.task_path / "COLVAR").write_text("#! FIELDS time cv\n0.0 0.5\n") + out = RunLmp().execute( + OPIO( + { + "config": {"plm_output_file": "COLVAR"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + self.assertEqual(out["plm_output"], Path(self.task_name) / "COLVAR") + @patch("dpgen2.op.run_lmp.run_command") def test_success(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", "")] diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py index 491d42f7..9722546e 100644 --- a/tests/test_select_confs.py +++ b/tests/test_select_confs.py @@ -165,3 +165,23 @@ def test_validate_trajs(self): trajs, model_devis, optional_outputs = SelectConfs.validate_trajs( trajs, model_devis, optional_outputs ) + + def test_validate_plm_outputs(self): + trajs = ["foo", None, "bar"] + model_devis = ["foo.md", None, "bar.md"] + self.assertEqual( + SelectConfs.validate_plm_outputs( + trajs, model_devis, ["foo.cv", None, "bar.cv"] + ), + ["foo.cv", "bar.cv"], + ) + self.assertEqual( + SelectConfs.validate_plm_outputs(trajs, model_devis), + None, + ) + with self.assertRaises(FatalError): + SelectConfs.validate_plm_outputs(trajs, model_devis, ["foo.cv"]) + with self.assertRaises(FatalError): + SelectConfs.validate_plm_outputs( + trajs, model_devis, ["foo.cv", "unexpected.cv", "bar.cv"] + ) From 38a33a89e91f44b5a63909f5d4969f2fd977316a Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Sun, 16 Aug 2026 14:31:36 +0800 Subject: [PATCH 2/6] feat(cv-filter): add configurable candidate sampling --- docs/input.md | 28 ++ dpgen2/entrypoint/args.py | 4 +- dpgen2/exploration/report/report.py | 3 + .../selector/conf_selector_frame.py | 26 +- .../exploration/selector/plumed_cv_filter.py | 294 +++++++++++++++++- tests/exploration/test_conf_selector_frame.py | 34 +- tests/exploration/test_plumed_cv_filter.py | 160 +++++++++- 7 files changed, 518 insertions(+), 31 deletions(-) diff --git a/docs/input.md b/docs/input.md index 382711ff..50ba6bfa 100644 --- a/docs/input.md +++ b/docs/input.md @@ -173,6 +173,34 @@ increasing `time`, or row-to-trajectory alignment is invalid. Model-deviation trust levels are applied first, followed by the CV regions and then the existing candidate limit and selection policy. +To make the final candidates cover a primary CV interval instead of clustering +where the trajectory spends most of its time, add an optional sampling policy: + +```json +"cv_filter": { + "regions": [ + {"d": [0.08, 0.12]}, + {"d": [0.16, 0.24], "v": [1.8, 2.2]} + ], + "sampling": { + "mode": "uniform", + "field": "d", + "n_bins": 10, + "within_bin": "max_deviation", + "seed": 20260815 + } +} +``` + +`uniform` divides the configured interval for `field` in each region into +equal-width bins. Regions receive a balanced share of the FP task limit and +the selected non-empty bins span the available interval. `within_bin` is +either `random` or `max_deviation`; `seed` makes random choices reproducible. +Every region must bound the primary `field`; its other CVs remain AND +constraints. Use `{"mode": "random", "seed": 20260815}` for reproducible +random selection after the CV filter. If `sampling` is omitted, DPGEN2 keeps +the convergence report's existing random or maximum-deviation selection. + ### FP diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index ddc1610a..c6f5abe6 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -215,7 +215,9 @@ def lmp_args(): doc_filters = "A list of configuration filters" doc_cv_filter = ( "Optional PLUMED CV candidate filter. The regions key is a list of " - "field-to-[lower, upper] mappings; regions are combined by OR." + "field-to-[lower, upper] mappings; regions are combined by OR. An " + "optional sampling key selects candidates randomly or uniformly along " + "one primary CV." ) return [ diff --git a/dpgen2/exploration/report/report.py b/dpgen2/exploration/report/report.py index 066d75f8..3c21b45a 100644 --- a/dpgen2/exploration/report/report.py +++ b/dpgen2/exploration/report/report.py @@ -76,6 +76,7 @@ def restrict_candidate_ids( def get_candidate_ids( self, max_nframes: Optional[int] = None, + clear: bool = True, ) -> List[List[int]]: r"""Get indexes of candidate configurations @@ -83,6 +84,8 @@ def get_candidate_ids( ---------- max_nframes The maximal number of frames of candidates. + clear + Clear frame-level report data after selection. Returns ------- diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index 35de3b81..59ee50c5 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -104,18 +104,32 @@ def select( self.report.clear() self.report.record(md_model_devi) + id_cand_list = None if self.plumed_cv_filter is not None: if plm_outputs is None or any(output is None for output in plm_outputs): raise FatalError( "PLUMED CV filtering requires one output per trajectory" ) md_f = md_model_devi.get(DeviManager.MAX_DEVI_F) - allowed_ids = self.plumed_cv_filter.get_selected_ids( - plm_outputs, - [len(values) for values in md_f], - ) - self.report.restrict_candidate_ids(allowed_ids) - id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) + if self.plumed_cv_filter.sampling is None: + allowed_ids = self.plumed_cv_filter.get_selected_ids( + plm_outputs, + [len(values) for values in md_f], + ) + self.report.restrict_candidate_ids(allowed_ids) + else: + candidate_ids = self.report.get_candidate_ids(None, clear=False) + sampled_ids = self.plumed_cv_filter.select_candidate_ids( + plm_outputs, + [len(values) for values in md_f], + candidate_ids, + self.max_numb_sel, + md_f, + ) + self.report.restrict_candidate_ids(sampled_ids) + id_cand_list = self.report.get_candidate_ids() + if id_cand_list is None: + id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) ms = self.traj_render.get_confs( trajs, diff --git a/dpgen2/exploration/selector/plumed_cv_filter.py b/dpgen2/exploration/selector/plumed_cv_filter.py index 865a0ca2..99f25caf 100644 --- a/dpgen2/exploration/selector/plumed_cv_filter.py +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -4,7 +4,9 @@ from typing import ( Dict, List, + Optional, Sequence, + Tuple, ) import numpy as np @@ -35,10 +37,32 @@ def args() -> List[Argument]: "A list of PLUMED field-to-[lower, upper] mappings. " "Fields within a region are ANDed; regions are ORed." ), - ) + ), + Argument( + "sampling", + dict, + [ + Argument("mode", str, optional=False), + Argument("field", str, optional=True, default=None), + Argument("n_bins", int, optional=True, default=10), + Argument("within_bin", str, optional=True, default="random"), + Argument("seed", int, optional=True, default=0), + ], + optional=True, + default=None, + doc=( + "Optional final candidate sampling. mode is random or uniform; " + "uniform uses field and n_bins, then random or max_deviation " + "within each bin." + ), + ), ] - def __init__(self, regions: List[Dict[str, Sequence[float]]]): + def __init__( + self, + regions: List[Dict[str, Sequence[float]]], + sampling: Optional[Dict] = None, + ): if ( not isinstance(regions, list) or not regions @@ -64,15 +88,250 @@ def __init__(self, regions: List[Dict[str, Sequence[float]]]): if not np.all(np.isfinite(normalized[field])): raise ValueError(f"non-finite interval for PLUMED field {field!r}") self.regions.append(normalized) + self.sampling = self._normalize_sampling(sampling) + + def _normalize_sampling(self, sampling: Optional[Dict]): + if sampling is None: + return None + if not isinstance(sampling, dict): + raise ValueError("PLUMED CV sampling must be a dict") + mode = sampling.get("mode") + if mode not in {"random", "uniform"}: + raise ValueError("PLUMED CV sampling mode must be random or uniform") + seed = sampling.get("seed", 0) + if not isinstance(seed, int) or isinstance(seed, bool) or seed < 0: + raise ValueError("PLUMED CV sampling seed must be a non-negative integer") + normalized = {"mode": mode, "seed": seed} + if mode == "uniform": + field = sampling.get("field") + n_bins = sampling.get("n_bins", 10) + within_bin = sampling.get("within_bin", "random") + if not isinstance(field, str) or not field: + raise ValueError("uniform PLUMED CV sampling requires a field") + if any(field not in region for region in self.regions): + raise ValueError( + f"uniform sampling field {field!r} must bound every region" + ) + if not isinstance(n_bins, int) or isinstance(n_bins, bool) or n_bins <= 0: + raise ValueError("PLUMED CV sampling n_bins must be positive") + if within_bin not in {"random", "max_deviation"}: + raise ValueError("PLUMED CV within_bin must be random or max_deviation") + normalized.update( + {"field": field, "n_bins": n_bins, "within_bin": within_bin} + ) + return normalized def get_selected_ids( self, files: List[Path], nframes: List[int], ) -> List[List[int]]: + outputs = self._load_outputs(files, nframes) + selected = [] + for fields, values in outputs: + masks = self._region_masks(fields, values) + selected.append(np.flatnonzero(np.logical_or.reduce(masks)).tolist()) + return selected + + def select_candidate_ids( + self, + files: List[Path], + nframes: List[int], + candidate_ids: List[List[int]], + max_nframes: Optional[int], + max_devi_f: Optional[List[np.ndarray]] = None, + ) -> List[List[int]]: + """Filter model-deviation candidates and apply configured sampling.""" + if self.sampling is None: + raise ValueError("PLUMED CV sampling is not configured") + if len(candidate_ids) != len(files): + raise FatalError("candidate IDs and trajectories have different lengths") + if max_nframes is not None and max_nframes < 0: + raise ValueError("max_nframes must be non-negative") + + outputs = self._load_outputs(files, nframes) + masks_by_traj = [ + self._region_masks(fields, values) for fields, values in outputs + ] + candidates = [] + for traj_idx, frame_ids in enumerate(candidate_ids): + for frame_idx in sorted(set(frame_ids)): + if ( + not isinstance(frame_idx, (int, np.integer)) + or frame_idx < 0 + or frame_idx >= nframes[traj_idx] + ): + raise FatalError("candidate frame index is outside the trajectory") + if any(mask[frame_idx] for mask in masks_by_traj[traj_idx]): + candidates.append((traj_idx, int(frame_idx))) + + limit = ( + len(candidates) + if max_nframes is None + else min(max_nframes, len(candidates)) + ) + if limit == len(candidates): + return self._group_candidates(candidates, len(files)) + rng = np.random.default_rng(self.sampling["seed"]) + if self.sampling["mode"] == "random": + picked = [ + candidates[ii] for ii in rng.choice(len(candidates), limit, False) + ] + else: + picked = self._sample_uniform( + candidates, + outputs, + masks_by_traj, + limit, + max_devi_f, + rng, + ) + return self._group_candidates(picked, len(files)) + + def _sample_uniform( + self, + candidates: List[Tuple[int, int]], + outputs, + masks_by_traj, + limit: int, + max_devi_f: Optional[List[np.ndarray]], + rng, + ) -> List[Tuple[int, int]]: + field = self.sampling["field"] + n_bins = self.sampling["n_bins"] + within_bin = self.sampling["within_bin"] + if within_bin == "max_deviation" and max_devi_f is None: + raise FatalError("max_deviation sampling requires force model deviations") + + buckets = [dict() for _ in self.regions] + for candidate in candidates: + traj_idx, frame_idx = candidate + fields, values = outputs[traj_idx] + value = values[frame_idx, fields.index(field)] + for region_idx, (region, mask) in enumerate( + zip(self.regions, masks_by_traj[traj_idx]) + ): + if not mask[frame_idx]: + continue + lower, upper = region[field] + bin_idx = min( + int((value - lower) / (upper - lower) * n_bins), n_bins - 1 + ) + buckets[region_idx].setdefault(bin_idx, []).append(candidate) + + capacities = [ + len({item for items in region.values() for item in items}) + for region in buckets + ] + region_quotas = self._balanced_quotas(capacities, limit) + picked = [] + picked_set = set() + for region, quota in zip(buckets, region_quotas): + available = { + bin_idx: [item for item in items if item not in picked_set] + for bin_idx, items in region.items() + } + available = {key: value for key, value in available.items() if value} + bin_quotas = self._bin_quotas(available, quota) + for bin_idx, count in bin_quotas.items(): + chosen = self._pick_within_bin( + available[bin_idx], count, within_bin, max_devi_f, rng + ) + picked.extend(chosen) + picked_set.update(chosen) + + if len(picked) < limit: + remaining = [item for item in candidates if item not in picked_set] + picked.extend( + self._pick_within_bin( + remaining, limit - len(picked), within_bin, max_devi_f, rng + ) + ) + return picked[:limit] + + @staticmethod + def _balanced_quotas(capacities: List[int], total: int) -> List[int]: + quotas = [0] * len(capacities) + active = [idx for idx, capacity in enumerate(capacities) if capacity] + if total <= 0 or not active: + return quotas + if total <= len(active): + if total == 1: + selected = [active[len(active) // 2]] + else: + selected = [ + active[round(idx * (len(active) - 1) / (total - 1))] + for idx in range(total) + ] + for idx in selected: + quotas[idx] = 1 + return quotas + for idx in active: + quotas[idx] = 1 + total -= len(active) + while total: + progressed = False + for idx, capacity in enumerate(capacities): + if quotas[idx] < capacity: + quotas[idx] += 1 + total -= 1 + progressed = True + if not total: + break + if not progressed: + break + return quotas + + @classmethod + def _bin_quotas(cls, buckets, total: int): + bins = sorted(buckets) + if not bins or total <= 0: + return {} + if total <= len(bins): + if total == 1: + selected_bins = [bins[len(bins) // 2]] + else: + selected_bins = [ + bins[round(idx * (len(bins) - 1) / (total - 1))] + for idx in range(total) + ] + return {bin_idx: 1 for bin_idx in selected_bins} + quotas = [1] * len(bins) + extra = cls._balanced_quotas( + [len(buckets[bin_idx]) - 1 for bin_idx in bins], total - len(bins) + ) + return { + bin_idx: quota + increment + for bin_idx, quota, increment in zip(bins, quotas, extra) + } + + @staticmethod + def _pick_within_bin(candidates, count, mode, max_devi_f, rng): + if count <= 0: + return [] + if mode == "random": + order = rng.permutation(len(candidates))[:count] + return [candidates[ii] for ii in order] + return sorted( + candidates, + key=lambda item: ( + -max_devi_f[item[0]][item[1]], + item[0], + item[1], + ), + )[:count] + + @staticmethod + def _group_candidates(candidates, ntraj): + grouped = [[] for _ in range(ntraj)] + for traj_idx, frame_idx in sorted(candidates): + grouped[traj_idx].append(frame_idx) + return grouped + + def _load_outputs(self, files: List[Path], nframes: List[int]): if len(files) != len(nframes): raise FatalError("PLUMED outputs and trajectories have different lengths") - selected = [] + outputs = [] for file, expected_nframes in zip(files, nframes): fields, values = self._read(file) if len(values) != expected_nframes: @@ -80,20 +339,21 @@ def get_selected_ids( f"PLUMED output {file} has {len(values)} rows, expected " f"{expected_nframes}; PRINT STRIDE must match the trajectory stride" ) - field_idx = {field: idx for idx, field in enumerate(fields)} - keep = np.zeros(expected_nframes, dtype=bool) - for region in self.regions: - in_region = np.ones(expected_nframes, dtype=bool) - for field, (lower, upper) in region.items(): - if field not in field_idx: - raise FatalError( - f"PLUMED field {field!r} is missing from {file}" - ) - column = values[:, field_idx[field]] - in_region &= (column >= lower) & (column < upper) - keep |= in_region - selected.append(np.flatnonzero(keep).tolist()) - return selected + outputs.append((fields, values)) + return outputs + + def _region_masks(self, fields, values): + field_idx = {field: idx for idx, field in enumerate(fields)} + masks = [] + for region in self.regions: + in_region = np.ones(len(values), dtype=bool) + for field, (lower, upper) in region.items(): + if field not in field_idx: + raise FatalError(f"PLUMED field {field!r} is missing") + column = values[:, field_idx[field]] + in_region &= (column >= lower) & (column < upper) + masks.append(in_region) + return masks @staticmethod def _read(file: Path): diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py index e62ae7d4..6b655ed1 100644 --- a/tests/exploration/test_conf_selector_frame.py +++ b/tests/exploration/test_conf_selector_frame.py @@ -97,9 +97,7 @@ def tearDown(self): def test_plumed_filter_precedes_max_selection(self): plm_outputs = [Path("foo.cv"), Path("bar.cv")] for output in plm_outputs: - output.write_text( - "#! FIELDS time cv\n0.0 0.5\n1.0 0.5\n2.0 1.5\n" - ) + output.write_text("#! FIELDS time cv\n0.0 0.5\n1.0 0.5\n2.0 1.5\n") conf_selector = ConfSelectorFrames( TrajRenderLammps(), ExplorationReportTrustLevelsMax(0.1, 0.5), @@ -117,6 +115,36 @@ def test_plumed_filter_precedes_max_selection(self): self.assertEqual(ms[0].get_nframes(), 1) self.assertAlmostEqual(ms[0]["coords"][0][0][1], 3.87, places=2) + def test_plumed_uniform_sampling_is_final_selection(self): + plm_outputs = [Path("foo.cv"), Path("bar.cv")] + for output in plm_outputs: + output.write_text("#! FIELDS time cv\n0.0 0.05\n1.0 0.45\n2.0 0.95\n") + conf_selector = ConfSelectorFrames( + TrajRenderLammps(), + ExplorationReportTrustLevelsMax(0.1, 0.5), + max_numb_sel=2, + plumed_cv_filter=PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + sampling={ + "mode": "uniform", + "field": "cv", + "n_bins": 10, + "within_bin": "max_deviation", + }, + ), + ) + confs, _ = conf_selector.select( + self.trajs, + self.model_devis, + self.type_map, + plm_outputs=plm_outputs, + ) + ms = dpdata.MultiSystems(type_map=self.type_map) + ms.from_deepmd_npy(confs[0], labeled=False) + self.assertEqual(ms[0].get_nframes(), 2) + self.assertAlmostEqual(ms[0]["coords"][0][0][1], 2.87, places=2) + self.assertAlmostEqual(ms[0]["coords"][1][0][1], 4.87, places=2) + def test_f_0(self): report = ExplorationReportTrustLevelsRandom(0.1, 0.5, conv_accuracy=0.9) traj_render = TrajRenderLammps() diff --git a/tests/exploration/test_plumed_cv_filter.py b/tests/exploration/test_plumed_cv_filter.py index e4cfb3d3..2af9e5a7 100644 --- a/tests/exploration/test_plumed_cv_filter.py +++ b/tests/exploration/test_plumed_cv_filter.py @@ -4,6 +4,7 @@ Path, ) +import numpy as np from dargs import ( Argument, ) @@ -20,10 +21,19 @@ class TestPlumedCVFilter(unittest.TestCase): def test_config_schema(self): schema = Argument("cv_filter", dict, PlumedCVFilter.args()) config = schema.normalize_value( - {"regions": [{"distance": [0.8, 1.2]}]} + { + "regions": [{"distance": [0.8, 1.2]}], + "sampling": { + "mode": "uniform", + "field": "distance", + "n_bins": 8, + }, + } ) schema.check_value(config, strict=True) self.assertEqual(config["regions"][0]["distance"], [0.8, 1.2]) + self.assertEqual(config["sampling"]["within_bin"], "random") + self.assertEqual(config["sampling"]["seed"], 0) def test_union_of_regions(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -39,12 +49,141 @@ def test_union_of_regions(self): def test_intervals_are_lower_inclusive_and_upper_exclusive(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" - output.write_text( - "#! FIELDS time cv\n0.0 1.0\n1.0 2.0\n2.0 3.0\n" - ) + output.write_text("#! FIELDS time cv\n0.0 1.0\n1.0 2.0\n2.0 3.0\n") cv_filter = PlumedCVFilter(regions=[{"cv": [1.0, 3.0]}]) self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 1]]) + def test_random_sampling_is_reproducible(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv\n" + + "".join(f"{ii}.0 {ii / 20:.3f}\n" for ii in range(20)) + ) + kwargs = { + "regions": [{"cv": [0.25, 0.75]}], + "sampling": {"mode": "random", "seed": 17}, + } + selected = PlumedCVFilter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + repeated = PlumedCVFilter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + self.assertEqual(selected, repeated) + self.assertEqual(len(selected[0]), 5) + self.assertTrue(all(5 <= frame < 15 for frame in selected[0])) + + def test_uniform_sampling_spans_nonempty_bins(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + values = [0.05, 0.06] + [ii / 10 + 0.05 for ii in range(1, 10)] + output.write_text( + "#! FIELDS time cv\n" + + "".join(f"{ii}.0 {value:.3f}\n" for ii, value in enumerate(values)) + ) + deviations = np.arange(len(values), dtype=float) + deviations[0] = 100.0 + deviations[1] = 200.0 + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + sampling={ + "mode": "uniform", + "field": "cv", + "n_bins": 10, + "within_bin": "max_deviation", + "seed": 3, + }, + ) + selected = cv_filter.select_candidate_ids( + [output], [len(values)], [list(range(len(values)))], 3, [deviations] + ) + self.assertEqual(selected, [[1, 5, 10]]) + + def test_uniform_random_within_bin_is_reproducible(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv\n" + + "".join(f"{ii}.0 {ii / 20:.3f}\n" for ii in range(20)) + ) + kwargs = { + "regions": [{"cv": [0.0, 1.0]}], + "sampling": { + "mode": "uniform", + "field": "cv", + "n_bins": 5, + "within_bin": "random", + "seed": 29, + }, + } + selected = PlumedCVFilter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + repeated = PlumedCVFilter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + self.assertEqual(selected, repeated) + self.assertEqual(len(selected[0]), 5) + + def test_uniform_sampling_balances_regions_and_honors_and_conditions(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv gate\n" + "0.0 0.1 0.5\n" + "1.0 0.9 0.5\n" + "2.0 2.1 0.5\n" + "3.0 2.9 0.5\n" + "4.0 0.5 1.5\n" + ) + cv_filter = PlumedCVFilter( + regions=[ + {"cv": [0.0, 1.0], "gate": [0.0, 1.0]}, + {"cv": [2.0, 3.0], "gate": [0.0, 1.0]}, + ], + sampling={ + "mode": "uniform", + "field": "cv", + "n_bins": 10, + "within_bin": "random", + "seed": 9, + }, + ) + selected = cv_filter.select_candidate_ids( + [output], [5], [list(range(5))], 2 + )[0] + self.assertEqual(len(selected), 2) + self.assertTrue(any(frame in {0, 1} for frame in selected)) + self.assertTrue(any(frame in {2, 3} for frame in selected)) + self.assertNotIn(4, selected) + + def test_uniform_sampling_handles_empty_bins(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text("#! FIELDS time cv\n0.0 0.05\n1.0 0.45\n2.0 0.95\n") + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + sampling={"mode": "uniform", "field": "cv", "n_bins": 10}, + ) + self.assertEqual( + cv_filter.select_candidate_ids([output], [3], [list(range(3))], 2), + [[0, 2]], + ) + + def test_uniform_sampling_spreads_a_small_limit_across_regions(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text("#! FIELDS time cv\n0.0 0.5\n1.0 2.5\n2.0 4.5\n") + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}, {"cv": [2.0, 3.0]}, {"cv": [4.0, 5.0]}], + sampling={"mode": "uniform", "field": "cv", "n_bins": 4}, + ) + self.assertEqual( + cv_filter.select_candidate_ids([output], [3], [list(range(3))], 2), + [[0, 2]], + ) + def test_alignment_fails_closed(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" @@ -66,6 +205,19 @@ def test_invalid_config_and_file_fail_closed(self): with self.subTest(regions=regions), self.assertRaises(ValueError): PlumedCVFilter(regions=regions) + invalid_sampling = [ + {"mode": "weighted"}, + {"mode": "uniform"}, + {"mode": "uniform", "field": "other"}, + {"mode": "uniform", "field": "cv", "n_bins": 0}, + {"mode": "uniform", "field": "cv", "within_bin": "first"}, + {"mode": "random", "seed": True}, + {"mode": "random", "seed": -1}, + ] + for sampling in invalid_sampling: + with self.subTest(sampling=sampling), self.assertRaises(ValueError): + PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}], sampling=sampling) + invalid_outputs = [ "0.0 0.5\n#! FIELDS time cv\n", "#! FIELDS time cv\n0.0 nan\n", From ff311987854c62c625605cbc837f12099cec0190 Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Sun, 16 Aug 2026 14:31:36 +0800 Subject: [PATCH 3/6] feat(cv-filter): add uniform CV coverage and audit --- docs/input.md | 73 +- dpgen2/entrypoint/args.py | 9 +- .../selector/conf_selector_frame.py | 22 +- .../exploration/selector/plumed_cv_filter.py | 654 +++++++++++++++--- tests/exploration/test_conf_selector_frame.py | 20 +- tests/exploration/test_plumed_cv_filter.py | 185 +++++ 6 files changed, 848 insertions(+), 115 deletions(-) diff --git a/docs/input.md b/docs/input.md index 50ba6bfa..629846a3 100644 --- a/docs/input.md +++ b/docs/input.md @@ -134,7 +134,8 @@ The {dargs:argument}`"n_sample"` tells the number o #### PLUMED CV candidate filtering LAMMPS exploration candidates can be restricted to a union of named PLUMED CV -regions before the configured random or maximum-model-deviation selection: +regions after the model-deviation trust window and before final CV-space +coverage or an explicitly configured selection policy: ```json "explore": { @@ -145,7 +146,8 @@ regions before the configured random or maximum-model-deviation selection: "regions": [ {"d": [0.08, 0.12]}, {"v": [1.8, 2.2]} - ] + ], + "sampling": {"mode": "report"} } } ``` @@ -173,8 +175,12 @@ increasing `time`, or row-to-trajectory alignment is invalid. Model-deviation trust levels are applied first, followed by the CV regions and then the existing candidate limit and selection policy. -To make the final candidates cover a primary CV interval instead of clustering -where the trajectory spends most of its time, add an optional sampling policy: +By default, DPGEN2 prevents candidates from clustering where the trajectory +spends most of its time. If all regions share one CV field, it uses 10 +equal-width bins along that CV. If they share two CV fields, it uses a 10 by 10 +grid. In both cases it covers separated non-empty bins or cells and selects the +largest force model deviation within each one. An explicit policy can override +these defaults: ```json "cv_filter": { @@ -198,8 +204,63 @@ the selected non-empty bins span the available interval. `within_bin` is either `random` or `max_deviation`; `seed` makes random choices reproducible. Every region must bound the primary `field`; its other CVs remain AND constraints. Use `{"mode": "random", "seed": 20260815}` for reproducible -random selection after the CV filter. If `sampling` is omitted, DPGEN2 keeps -the convergence report's existing random or maximum-deviation selection. +candidate-frame random selection after the CV filter; this follows the +trajectory's CV density and can therefore cluster in a highly populated CV +region. Use `{"mode": "report"}` to retain the +convergence report's original random or maximum-deviation selection. If regions +do not share exactly one or two CV fields, sampling must be specified because +DPGEN2 cannot infer an unambiguous coverage space. + +For two-CV coverage and auditable reaction windows, regions may also have +names and weights: + +```json +"cv_filter": { + "regions": [ + { + "name": "incipient_contact", + "conditions": { + "iondistance": [0.2, 2.0], + "ionization": [0.2, 2.0] + } + }, + { + "name": "separated", + "conditions": { + "iondistance": [2.0, 10.0], + "ionization": [0.2, 2.0] + }, + "weight": 1.0 + } + ], + "sampling": { + "mode": "grid", + "grid": {"iondistance": 8, "ionization": 4}, + "within_bin": "max_deviation", + "seed": 20260815, + "min_frame_gap": 5 + }, + "time_alignment": { + "start": 0.0, + "step": 0.01, + "atol": 1e-8 + } +} +``` + +`grid` currently requires exactly two CV fields, both bounded by every region. +It allocates the candidate limit across regions (using `weight` when supplied), +covers separated non-empty cells before adding extra frames, and then uses +`within_bin` inside each cell. `min_frame_gap` is a minimum frame-index +separation within each trajectory. A spacing constraint may leave the result +underfilled; DPGEN2 reports this instead of silently relaxing the constraint. +`time_alignment` additionally verifies `time = start + frame * step` with the +configured absolute tolerance. + +When a CV filter is active, the selected DeepMD data directory contains +`cv_selection.csv` and `cv_selection_summary.json`. They record trajectory and +frame IDs, time, maximum force model deviation, CV values, matching regions, +grid cells, population counts, spacing rejections, and any underfilled quota. ### FP diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index c6f5abe6..2ac2b107 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -214,10 +214,11 @@ def lmp_args(): ) doc_filters = "A list of configuration filters" doc_cv_filter = ( - "Optional PLUMED CV candidate filter. The regions key is a list of " - "field-to-[lower, upper] mappings; regions are combined by OR. An " - "optional sampling key selects candidates randomly or uniformly along " - "one primary CV." + "Optional PLUMED CV candidate filter. A region may be a field-to-" + "[lower, upper] mapping or a named conditions mapping; regions are " + "combined by OR. By default, one or two common CVs are covered uniformly. " + "Sampling also supports explicit random, uniform, grid, or report modes, " + "with optional frame spacing and time alignment." ) return [ diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index 59ee50c5..da99fe59 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -84,6 +84,8 @@ def select( Format: each line has 7 numbers they are used as # frame_id md_v_max md_v_min md_v_mean md_f_max md_f_min md_f_mean where `md` stands for model deviation, v for virial and f for force + DeePMD outputs may append an eighth ``devi_e`` column, which does not + change the first seven columns consumed here. type_map : List[str] The `type_map` of the systems optional_outputs : List[Path] @@ -105,12 +107,14 @@ def select( self.report.clear() self.report.record(md_model_devi) id_cand_list = None + cv_audit = None if self.plumed_cv_filter is not None: if plm_outputs is None or any(output is None for output in plm_outputs): raise FatalError( "PLUMED CV filtering requires one output per trajectory" ) md_f = md_model_devi.get(DeviManager.MAX_DEVI_F) + candidate_ids = self.report.get_candidate_ids(None, clear=False) if self.plumed_cv_filter.sampling is None: allowed_ids = self.plumed_cv_filter.get_selected_ids( plm_outputs, @@ -118,18 +122,30 @@ def select( ) self.report.restrict_candidate_ids(allowed_ids) else: - candidate_ids = self.report.get_candidate_ids(None, clear=False) - sampled_ids = self.plumed_cv_filter.select_candidate_ids( + ( + sampled_ids, + records, + summary, + ) = self.plumed_cv_filter.select_candidate_ids_with_audit( plm_outputs, [len(values) for values in md_f], candidate_ids, self.max_numb_sel, md_f, ) + cv_audit = (records, summary) self.report.restrict_candidate_ids(sampled_ids) id_cand_list = self.report.get_candidate_ids() if id_cand_list is None: id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) + if self.plumed_cv_filter is not None and cv_audit is None: + cv_audit = self.plumed_cv_filter.audit_candidate_ids( + plm_outputs, + [len(values) for values in md_f], + candidate_ids, + id_cand_list, + md_f, + ) ms = self.traj_render.get_confs( trajs, @@ -142,5 +158,7 @@ def select( out_path = Path("confs") out_path.mkdir(exist_ok=True) ms.to_deepmd_npy(out_path) # type: ignore + if cv_audit is not None: + self.plumed_cv_filter.write_audit(out_path, *cv_audit) return [out_path], copy.deepcopy(self.report) diff --git a/dpgen2/exploration/selector/plumed_cv_filter.py b/dpgen2/exploration/selector/plumed_cv_filter.py index 99f25caf..563a8329 100644 --- a/dpgen2/exploration/selector/plumed_cv_filter.py +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -1,3 +1,5 @@ +import csv +import json from pathlib import ( Path, ) @@ -5,7 +7,6 @@ Dict, List, Optional, - Sequence, Tuple, ) @@ -22,8 +23,9 @@ class PlumedCVFilter: """Select frames in a union of PLUMED CV regions. Each region maps field names from ``#! FIELDS`` to a lower-inclusive, - upper-exclusive interval. A frame is selected when it matches every - interval in any one region. + upper-exclusive interval. Fields within a region are ANDed; regions are + ORed. A region may use the legacy bare mapping or the named form + ``{"name": ..., "conditions": {...}}``. """ @staticmethod @@ -34,8 +36,9 @@ def args() -> List[Argument]: list, optional=False, doc=( - "A list of PLUMED field-to-[lower, upper] mappings. " - "Fields within a region are ANDed; regions are ORed." + "A list of PLUMED field-to-[lower, upper] mappings or named " + "regions with name, conditions, and optional weight. Fields " + "within a region are ANDed; regions are ORed." ), ), Argument( @@ -45,23 +48,38 @@ def args() -> List[Argument]: Argument("mode", str, optional=False), Argument("field", str, optional=True, default=None), Argument("n_bins", int, optional=True, default=10), + Argument("grid", dict, optional=True, default=None), Argument("within_bin", str, optional=True, default="random"), Argument("seed", int, optional=True, default=0), + Argument("min_frame_gap", int, optional=True, default=0), ], optional=True, default=None, doc=( - "Optional final candidate sampling. mode is random or uniform; " - "uniform uses field and n_bins, then random or max_deviation " - "within each bin." + "Optional final candidate sampling. mode is random, uniform, " + "grid, or report. If omitted, one or two common CV fields are " + "covered uniformly. report keeps the report's selection policy." ), ), + Argument( + "time_alignment", + dict, + [ + Argument("start", float, optional=True, default=0.0), + Argument("step", float, optional=False), + Argument("atol", float, optional=True, default=1e-8), + ], + optional=True, + default=None, + doc="Optional expected PLUMED time = start + frame * step.", + ), ] def __init__( self, - regions: List[Dict[str, Sequence[float]]], + regions: List[Dict], sampling: Optional[Dict] = None, + time_alignment: Optional[Dict] = None, ): if ( not isinstance(regions, list) @@ -69,10 +87,39 @@ def __init__( or any(not isinstance(region, dict) or not region for region in regions) ): raise ValueError("PLUMED CV regions must be a non-empty list of dicts") + self.regions = [] - for region in regions: + self.region_names = [] + self.region_weights = [] + for region_idx, region in enumerate(regions): + if "conditions" in region: + unknown = set(region) - {"name", "conditions", "weight"} + if unknown: + raise ValueError( + f"unknown named PLUMED region keys: {sorted(unknown)}" + ) + conditions = region["conditions"] + name = region.get("name", f"region_{region_idx}") + weight = region.get("weight", 1.0) + else: + conditions = region + name = f"region_{region_idx}" + weight = 1.0 + if not isinstance(name, str) or not name: + raise ValueError("PLUMED region names must be non-empty strings") + if name in self.region_names: + raise ValueError(f"duplicate PLUMED region name {name!r}") + try: + weight = float(weight) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid weight for PLUMED region {name!r}") from exc + if not np.isfinite(weight) or weight <= 0: + raise ValueError(f"invalid weight for PLUMED region {name!r}") + if not isinstance(conditions, dict) or not conditions: + raise ValueError(f"PLUMED region {name!r} conditions must be a dict") + normalized = {} - for field, bounds in region.items(): + for field, bounds in conditions.items(): if not isinstance(field, str) or not field: raise ValueError("PLUMED field names must be non-empty strings") if not isinstance(bounds, (list, tuple)) or len(bounds) != 2: @@ -88,7 +135,36 @@ def __init__( if not np.all(np.isfinite(normalized[field])): raise ValueError(f"non-finite interval for PLUMED field {field!r}") self.regions.append(normalized) + self.region_names.append(name) + self.region_weights.append(weight) + + self.cv_fields = sorted({field for region in self.regions for field in region}) + self.sampling_inferred = sampling is None + if sampling is None: + sampling = self._default_sampling() self.sampling = self._normalize_sampling(sampling) + self.time_alignment = self._normalize_time_alignment(time_alignment) + + def _default_sampling(self): + common_fields = sorted( + set.intersection(*(set(region) for region in self.regions)) + ) + if len(common_fields) == 1: + return { + "mode": "uniform", + "field": common_fields[0], + "n_bins": 10, + "within_bin": "max_deviation", + } + if len(common_fields) == 2: + return { + "mode": "grid", + "grid": {field: 10 for field in common_fields}, + "within_bin": "max_deviation", + } + raise ValueError( + "sampling must be explicit unless regions share exactly one or two CV fields" + ) def _normalize_sampling(self, sampling: Optional[Dict]): if sampling is None: @@ -96,31 +172,81 @@ def _normalize_sampling(self, sampling: Optional[Dict]): if not isinstance(sampling, dict): raise ValueError("PLUMED CV sampling must be a dict") mode = sampling.get("mode") - if mode not in {"random", "uniform"}: - raise ValueError("PLUMED CV sampling mode must be random or uniform") + if mode not in {"random", "uniform", "grid", "report"}: + raise ValueError( + "PLUMED CV sampling mode must be random, uniform, grid, or report" + ) + if mode == "report": + return None seed = sampling.get("seed", 0) if not isinstance(seed, int) or isinstance(seed, bool) or seed < 0: raise ValueError("PLUMED CV sampling seed must be a non-negative integer") - normalized = {"mode": mode, "seed": seed} + min_frame_gap = sampling.get("min_frame_gap", 0) + if ( + not isinstance(min_frame_gap, int) + or isinstance(min_frame_gap, bool) + or min_frame_gap < 0 + ): + raise ValueError("PLUMED CV min_frame_gap must be non-negative") + normalized = { + "mode": mode, + "seed": seed, + "min_frame_gap": min_frame_gap, + "within_bin": "random", + } + if mode == "random": + return normalized + + within_bin = sampling.get("within_bin", "random") + if within_bin not in {"random", "max_deviation"}: + raise ValueError("PLUMED CV within_bin must be random or max_deviation") if mode == "uniform": field = sampling.get("field") n_bins = sampling.get("n_bins", 10) - within_bin = sampling.get("within_bin", "random") if not isinstance(field, str) or not field: raise ValueError("uniform PLUMED CV sampling requires a field") - if any(field not in region for region in self.regions): - raise ValueError( - f"uniform sampling field {field!r} must bound every region" - ) if not isinstance(n_bins, int) or isinstance(n_bins, bool) or n_bins <= 0: raise ValueError("PLUMED CV sampling n_bins must be positive") - if within_bin not in {"random", "max_deviation"}: - raise ValueError("PLUMED CV within_bin must be random or max_deviation") - normalized.update( - {"field": field, "n_bins": n_bins, "within_bin": within_bin} - ) + grid = {field: n_bins} + normalized.update({"field": field, "n_bins": n_bins}) + else: + grid = sampling.get("grid") + if not isinstance(grid, dict) or len(grid) != 2: + raise ValueError("grid PLUMED CV sampling requires exactly two fields") + if any( + not isinstance(field, str) + or not field + or not isinstance(n_bins, int) + or isinstance(n_bins, bool) + or n_bins <= 0 + for field, n_bins in grid.items() + ): + raise ValueError("PLUMED CV grid fields and bin counts are invalid") + grid = dict(grid) + if any(field not in region for region in self.regions for field in grid): + raise ValueError("every grid field must bound every PLUMED CV region") + normalized.update({"grid": grid, "within_bin": within_bin}) return normalized + @staticmethod + def _normalize_time_alignment(time_alignment: Optional[Dict]): + if time_alignment is None: + return None + if not isinstance(time_alignment, dict): + raise ValueError("PLUMED CV time_alignment must be a dict") + unknown = set(time_alignment) - {"start", "step", "atol"} + if unknown: + raise ValueError(f"unknown PLUMED time_alignment keys: {sorted(unknown)}") + try: + start = float(time_alignment.get("start", 0.0)) + step = float(time_alignment["step"]) + atol = float(time_alignment.get("atol", 1e-8)) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("invalid PLUMED CV time_alignment") from exc + if not np.all(np.isfinite([start, step, atol])) or step <= 0 or atol < 0: + raise ValueError("invalid PLUMED CV time_alignment") + return {"start": start, "step": step, "atol": atol} + def get_selected_ids( self, files: List[Path], @@ -142,6 +268,20 @@ def select_candidate_ids( max_devi_f: Optional[List[np.ndarray]] = None, ) -> List[List[int]]: """Filter model-deviation candidates and apply configured sampling.""" + selected, _, _ = self.select_candidate_ids_with_audit( + files, nframes, candidate_ids, max_nframes, max_devi_f + ) + return selected + + def select_candidate_ids_with_audit( + self, + files: List[Path], + nframes: List[int], + candidate_ids: List[List[int]], + max_nframes: Optional[int], + max_devi_f: Optional[List[np.ndarray]] = None, + ): + """Return sampled IDs plus selected-frame records and aggregate counts.""" if self.sampling is None: raise ValueError("PLUMED CV sampling is not configured") if len(candidate_ids) != len(files): @@ -153,32 +293,27 @@ def select_candidate_ids( masks_by_traj = [ self._region_masks(fields, values) for fields, values in outputs ] - candidates = [] - for traj_idx, frame_ids in enumerate(candidate_ids): - for frame_idx in sorted(set(frame_ids)): - if ( - not isinstance(frame_idx, (int, np.integer)) - or frame_idx < 0 - or frame_idx >= nframes[traj_idx] - ): - raise FatalError("candidate frame index is outside the trajectory") - if any(mask[frame_idx] for mask in masks_by_traj[traj_idx]): - candidates.append((traj_idx, int(frame_idx))) - + candidates = self._eligible_candidates(candidate_ids, nframes, masks_by_traj) limit = ( len(candidates) if max_nframes is None else min(max_nframes, len(candidates)) ) - if limit == len(candidates): - return self._group_candidates(candidates, len(files)) rng = np.random.default_rng(self.sampling["seed"]) + rejected_by_gap = set() + picked = [] if self.sampling["mode"] == "random": - picked = [ - candidates[ii] for ii in rng.choice(len(candidates), limit, False) - ] + self._pick_candidates( + candidates, + limit, + "random", + max_devi_f, + rng, + picked, + rejected_by_gap, + ) else: - picked = self._sample_uniform( + picked, rejected_by_gap = self._sample_grid( candidates, outputs, masks_by_traj, @@ -186,9 +321,69 @@ def select_candidate_ids( max_devi_f, rng, ) - return self._group_candidates(picked, len(files)) + records, summary = self._make_audit( + candidate_ids, + candidates, + picked, + outputs, + masks_by_traj, + max_devi_f, + limit, + rejected_by_gap, + self.sampling["mode"], + ) + return self._group_candidates(picked, len(files)), records, summary + + def audit_candidate_ids( + self, + files: List[Path], + nframes: List[int], + candidate_ids: List[List[int]], + selected_ids: List[List[int]], + max_devi_f: Optional[List[np.ndarray]] = None, + ): + """Audit filtering followed by the report's existing sampling policy.""" + outputs = self._load_outputs(files, nframes) + masks_by_traj = [ + self._region_masks(fields, values) for fields, values in outputs + ] + candidates = self._eligible_candidates(candidate_ids, nframes, masks_by_traj) + eligible_set = set(candidates) + picked = [] + for traj_idx, frame_ids in enumerate(selected_ids): + for frame_idx in sorted(set(frame_ids)): + candidate = (traj_idx, int(frame_idx)) + if candidate not in eligible_set: + raise FatalError("selected frame is outside PLUMED CV candidates") + picked.append(candidate) + records, summary = self._make_audit( + candidate_ids, + candidates, + picked, + outputs, + masks_by_traj, + max_devi_f, + len(picked), + set(), + "report", + ) + return records, summary - def _sample_uniform( + def _eligible_candidates(self, candidate_ids, nframes, masks_by_traj): + candidates = [] + for traj_idx, frame_ids in enumerate(candidate_ids): + for frame_idx in sorted(set(frame_ids)): + if ( + not isinstance(frame_idx, (int, np.integer)) + or frame_idx < 0 + or frame_idx >= nframes[traj_idx] + ): + raise FatalError("candidate frame index is outside the trajectory") + if any(mask[frame_idx] for mask in masks_by_traj[traj_idx]): + candidates.append((traj_idx, int(frame_idx))) + return candidates + + def _sample_grid( self, candidates: List[Tuple[int, int]], outputs, @@ -196,9 +391,7 @@ def _sample_uniform( limit: int, max_devi_f: Optional[List[np.ndarray]], rng, - ) -> List[Tuple[int, int]]: - field = self.sampling["field"] - n_bins = self.sampling["n_bins"] + ): within_bin = self.sampling["within_bin"] if within_bin == "max_deviation" and max_devi_f is None: raise FatalError("max_deviation sampling requires force model deviations") @@ -207,47 +400,130 @@ def _sample_uniform( for candidate in candidates: traj_idx, frame_idx = candidate fields, values = outputs[traj_idx] - value = values[frame_idx, fields.index(field)] - for region_idx, (region, mask) in enumerate( - zip(self.regions, masks_by_traj[traj_idx]) - ): + for region_idx, mask in enumerate(masks_by_traj[traj_idx]): if not mask[frame_idx]: continue - lower, upper = region[field] - bin_idx = min( - int((value - lower) / (upper - lower) * n_bins), n_bins - 1 - ) - buckets[region_idx].setdefault(bin_idx, []).append(candidate) + cell = self._cell_key(region_idx, fields, values[frame_idx]) + buckets[region_idx].setdefault(cell, []).append(candidate) capacities = [ len({item for items in region.values() for item in items}) for region in buckets ] - region_quotas = self._balanced_quotas(capacities, limit) + region_quotas = self._weighted_quotas(capacities, self.region_weights, limit) picked = [] picked_set = set() + rejected_by_gap = set() + grid_sizes = tuple(self.sampling["grid"].values()) for region, quota in zip(buckets, region_quotas): available = { - bin_idx: [item for item in items if item not in picked_set] - for bin_idx, items in region.items() + cell: [item for item in items if item not in picked_set] + for cell, items in region.items() } - available = {key: value for key, value in available.items() if value} - bin_quotas = self._bin_quotas(available, quota) - for bin_idx, count in bin_quotas.items(): - chosen = self._pick_within_bin( - available[bin_idx], count, within_bin, max_devi_f, rng + available = {cell: items for cell, items in available.items() if items} + cell_quotas = self._cell_quotas(available, quota, grid_sizes) + for cell, count in cell_quotas.items(): + before = len(picked) + self._pick_candidates( + available[cell], + count, + within_bin, + max_devi_f, + rng, + picked, + rejected_by_gap, ) - picked.extend(chosen) - picked_set.update(chosen) + picked_set.update(picked[before:]) if len(picked) < limit: remaining = [item for item in candidates if item not in picked_set] - picked.extend( - self._pick_within_bin( - remaining, limit - len(picked), within_bin, max_devi_f, rng + self._pick_candidates( + remaining, + limit - len(picked), + within_bin, + max_devi_f, + rng, + picked, + rejected_by_gap, + ) + return picked[:limit], rejected_by_gap + + def _pick_candidates( + self, + candidates, + count, + mode, + max_devi_f, + rng, + picked, + rejected_by_gap, + ): + if count <= 0: + return + initially_picked = len(picked) + picked_set = set(picked) + candidates = list( + dict.fromkeys(item for item in candidates if item not in picked_set) + ) + if mode == "random": + order = rng.permutation(len(candidates)) + ordered = [candidates[ii] for ii in order] + else: + if max_devi_f is None: + raise FatalError( + "max_deviation sampling requires force model deviations" ) + ordered = sorted( + candidates, + key=lambda item: ( + -self._deviation_value(max_devi_f, item), + item[0], + item[1], + ), ) - return picked[:limit] + min_frame_gap = self.sampling["min_frame_gap"] + for candidate in ordered: + if len(picked) - initially_picked >= count: + break + if min_frame_gap and any( + candidate[0] == other[0] + and abs(candidate[1] - other[1]) < min_frame_gap + for other in picked + ): + rejected_by_gap.add(candidate) + continue + picked.append(candidate) + + @staticmethod + def _deviation_value(max_devi_f, candidate): + traj_idx, frame_idx = candidate + try: + value = float(max_devi_f[traj_idx][frame_idx]) + except (IndexError, TypeError, ValueError) as exc: + raise FatalError( + "force model deviations do not match trajectories" + ) from exc + if not np.isfinite(value): + raise FatalError("force model deviations must be finite") + return value + + @classmethod + def _weighted_quotas(cls, capacities, weights, total): + if not capacities or len(set(weights)) == 1: + return cls._balanced_quotas(capacities, total) + quotas = [0] * len(capacities) + for _ in range(total): + available = [ + idx for idx, capacity in enumerate(capacities) if quotas[idx] < capacity + ] + if not available: + break + idx = max( + available, + key=lambda ii: (weights[ii] / (quotas[ii] + 1), -ii), + ) + quotas[idx] += 1 + return quotas @staticmethod def _balanced_quotas(capacities: List[int], total: int) -> List[int]: @@ -283,43 +559,61 @@ def _balanced_quotas(capacities: List[int], total: int) -> List[int]: return quotas @classmethod - def _bin_quotas(cls, buckets, total: int): - bins = sorted(buckets) - if not bins or total <= 0: + def _cell_quotas(cls, buckets, total, grid_sizes): + cells = sorted(buckets) + if not cells or total <= 0: return {} - if total <= len(bins): - if total == 1: - selected_bins = [bins[len(bins) // 2]] - else: - selected_bins = [ - bins[round(idx * (len(bins) - 1) / (total - 1))] - for idx in range(total) - ] - return {bin_idx: 1 for bin_idx in selected_bins} - quotas = [1] * len(bins) + if total <= len(cells): + return {cell: 1 for cell in cls._spread_cells(cells, total, grid_sizes)} extra = cls._balanced_quotas( - [len(buckets[bin_idx]) - 1 for bin_idx in bins], total - len(bins) + [len(buckets[cell]) - 1 for cell in cells], total - len(cells) ) - return { - bin_idx: quota + increment - for bin_idx, quota, increment in zip(bins, quotas, extra) - } + return {cell: 1 + increment for cell, increment in zip(cells, extra)} @staticmethod - def _pick_within_bin(candidates, count, mode, max_devi_f, rng): - if count <= 0: - return [] - if mode == "random": - order = rng.permutation(len(candidates))[:count] - return [candidates[ii] for ii in order] - return sorted( - candidates, - key=lambda item: ( - -max_devi_f[item[0]][item[1]], - item[0], - item[1], - ), - )[:count] + def _spread_cells(cells, total, grid_sizes): + if total == 1: + center = tuple((size - 1) / 2 for size in grid_sizes) + return [ + min( + cells, + key=lambda cell: sum( + ((value - middle) / max(size - 1, 1)) ** 2 + for value, middle, size in zip(cell, center, grid_sizes) + ), + ) + ] + chosen = [cells[0]] + while len(chosen) < total: + best = None + best_distance = -1.0 + for cell in cells: + if cell in chosen: + continue + distance = min( + sum( + ((left - right) / max(size - 1, 1)) ** 2 + for left, right, size in zip(cell, other, grid_sizes) + ) + for other in chosen + ) + if distance > best_distance: + best = cell + best_distance = distance + chosen.append(best) + return chosen + + def _cell_key(self, region_idx, fields, row): + field_idx = {field: idx for idx, field in enumerate(fields)} + region = self.regions[region_idx] + cell = [] + for field, n_bins in self.sampling["grid"].items(): + lower, upper = region[field] + value = row[field_idx[field]] + cell.append( + min(int((value - lower) / (upper - lower) * n_bins), n_bins - 1) + ) + return tuple(cell) @staticmethod def _group_candidates(candidates, ntraj): @@ -328,6 +622,147 @@ def _group_candidates(candidates, ntraj): grouped[traj_idx].append(frame_idx) return grouped + def _make_audit( + self, + candidate_ids, + candidates, + picked, + outputs, + masks_by_traj, + max_devi_f, + requested, + rejected_by_gap, + mode, + ): + grid = None if self.sampling is None else self.sampling.get("grid") + within_bin = ( + "" if self.sampling is None else self.sampling.get("within_bin", "") + ) + seed = "" if self.sampling is None else self.sampling["seed"] + records = [] + for traj_idx, frame_idx in sorted(picked): + fields, values = outputs[traj_idx] + row = values[frame_idx] + field_idx = {field: idx for idx, field in enumerate(fields)} + matched = [ + idx + for idx, mask in enumerate(masks_by_traj[traj_idx]) + if mask[frame_idx] + ] + cells = [] + if grid is not None: + for region_idx in matched: + cell = self._cell_key(region_idx, fields, row) + labels = ",".join( + f"{field}={value}" for field, value in zip(grid, cell) + ) + cells.append(f"{self.region_names[region_idx]}:{labels}") + record = { + "traj_idx": traj_idx, + "frame_idx": frame_idx, + "time": float(row[field_idx["time"]]), + "max_devi_f": ( + "" + if max_devi_f is None + else self._deviation_value(max_devi_f, (traj_idx, frame_idx)) + ), + "region_names": ";".join(self.region_names[idx] for idx in matched), + "cell_or_bin": ";".join(cells), + "sampling_mode": mode, + "within_bin": within_bin, + "seed": seed, + } + for field in self.cv_fields: + record[f"cv_{field}"] = float(row[field_idx[field]]) + records.append(record) + + eligible_set = set(candidates) + picked_set = set(picked) + per_region = {} + for region_idx, name in enumerate(self.region_names): + eligible_region = { + candidate + for candidate in eligible_set + if masks_by_traj[candidate[0]][region_idx][candidate[1]] + } + selected_region = eligible_region & picked_set + eligible_cells = set() + selected_cells = set() + if grid is not None: + for candidate in eligible_region: + fields, values = outputs[candidate[0]] + cell = self._cell_key(region_idx, fields, values[candidate[1]]) + eligible_cells.add(cell) + if candidate in selected_region: + selected_cells.add(cell) + per_region[name] = { + "eligible": len(eligible_region), + "selected": len(selected_region), + "nonempty_cells": len(eligible_cells), + "selected_cells": len(selected_cells), + } + + trust_candidates = sum(len(set(frame_ids)) for frame_ids in candidate_ids) + summary = { + "trust_candidates": trust_candidates, + "cv_eligible_candidates": len(eligible_set), + "requested": requested, + "selected": len(picked_set), + "underfilled_quota": max(requested - len(picked_set), 0), + "min_frame_gap_rejects": len(rejected_by_gap), + "sampling_mode": mode, + "sampling_inferred": self.sampling_inferred, + "within_bin": within_bin, + "seed": seed, + "min_frame_gap": ( + 0 if self.sampling is None else self.sampling["min_frame_gap"] + ), + "interval_semantics": "lower-inclusive, upper-exclusive", + "cv_fields": self.cv_fields, + "time_alignment": self.time_alignment, + "per_region": per_region, + "regions": [ + { + "name": name, + "weight": weight, + "conditions": { + field: list(bounds) for field, bounds in region.items() + }, + } + for name, weight, region in zip( + self.region_names, self.region_weights, self.regions + ) + ], + } + return records, summary + + @staticmethod + def write_audit(out_path: Path, records, summary): + out_path.mkdir(exist_ok=True) + base_fields = [ + "traj_idx", + "frame_idx", + "time", + "max_devi_f", + "region_names", + "cell_or_bin", + "sampling_mode", + "within_bin", + "seed", + ] + fieldnames = base_fields + [f"cv_{field}" for field in summary["cv_fields"]] + with (out_path / "cv_selection.csv").open( + "w", newline="", encoding="utf8" + ) as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(records) + with (out_path / "cv_selection_summary.json").open( + "w", encoding="utf8" + ) as handle: + json.dump(summary, handle, indent=2, sort_keys=True) + handle.write("\n") + def _load_outputs(self, files: List[Path], nframes: List[int]): if len(files) != len(nframes): raise FatalError("PLUMED outputs and trajectories have different lengths") @@ -339,6 +774,21 @@ def _load_outputs(self, files: List[Path], nframes: List[int]): f"PLUMED output {file} has {len(values)} rows, expected " f"{expected_nframes}; PRINT STRIDE must match the trajectory stride" ) + if self.time_alignment is not None: + time = values[:, fields.index("time")] + expected_time = ( + self.time_alignment["start"] + + np.arange(expected_nframes) * self.time_alignment["step"] + ) + if not np.allclose( + time, + expected_time, + rtol=0.0, + atol=self.time_alignment["atol"], + ): + raise FatalError( + f"PLUMED time in {file} does not match configured frame times" + ) outputs.append((fields, values)) return outputs diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py index 6b655ed1..7d299afd 100644 --- a/tests/exploration/test_conf_selector_frame.py +++ b/tests/exploration/test_conf_selector_frame.py @@ -1,3 +1,5 @@ +import csv +import json import os import shutil import textwrap @@ -102,7 +104,9 @@ def test_plumed_filter_precedes_max_selection(self): TrajRenderLammps(), ExplorationReportTrustLevelsMax(0.1, 0.5), max_numb_sel=1, - plumed_cv_filter=PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]), + plumed_cv_filter=PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], sampling={"mode": "report"} + ), ) confs, _ = conf_selector.select( self.trajs, @@ -144,6 +148,20 @@ def test_plumed_uniform_sampling_is_final_selection(self): self.assertEqual(ms[0].get_nframes(), 2) self.assertAlmostEqual(ms[0]["coords"][0][0][1], 2.87, places=2) self.assertAlmostEqual(ms[0]["coords"][1][0][1], 4.87, places=2) + audit_csv = Path("confs/cv_selection.csv") + audit_json = Path("confs/cv_selection_summary.json") + self.assertTrue(audit_csv.is_file()) + self.assertTrue(audit_json.is_file()) + with audit_csv.open(newline="") as handle: + rows = list(csv.DictReader(handle)) + summary = json.loads(audit_json.read_text()) + self.assertEqual(len(rows), 2) + self.assertEqual({int(row["frame_idx"]) for row in rows}, {0, 2}) + self.assertTrue(all(row["max_devi_f"] for row in rows)) + self.assertTrue(all(row["cv_cv"] for row in rows)) + self.assertEqual(summary["trust_candidates"], 6) + self.assertEqual(summary["cv_eligible_candidates"], 6) + self.assertEqual(summary["selected"], 2) def test_f_0(self): report = ExplorationReportTrustLevelsRandom(0.1, 0.5, conv_accuracy=0.9) diff --git a/tests/exploration/test_plumed_cv_filter.py b/tests/exploration/test_plumed_cv_filter.py index 2af9e5a7..315554c8 100644 --- a/tests/exploration/test_plumed_cv_filter.py +++ b/tests/exploration/test_plumed_cv_filter.py @@ -35,6 +35,30 @@ def test_config_schema(self): self.assertEqual(config["sampling"]["within_bin"], "random") self.assertEqual(config["sampling"]["seed"], 0) + grid_config = schema.normalize_value( + { + "regions": [ + { + "name": "contact", + "conditions": { + "iondistance": [0.2, 2.0], + "ionization": [0.25, 2.0], + }, + } + ], + "sampling": { + "mode": "grid", + "grid": {"iondistance": 4, "ionization": 2}, + "within_bin": "max_deviation", + "min_frame_gap": 5, + }, + "time_alignment": {"start": 0.0, "step": 0.01}, + } + ) + schema.check_value(grid_config, strict=True) + self.assertEqual(grid_config["sampling"]["min_frame_gap"], 5) + self.assertEqual(grid_config["time_alignment"]["atol"], 1e-8) + def test_union_of_regions(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" @@ -74,6 +98,42 @@ def test_random_sampling_is_reproducible(self): self.assertEqual(len(selected[0]), 5) self.assertTrue(all(5 <= frame < 15 for frame in selected[0])) + def test_default_sampling_covers_one_cv_uniformly(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + values = [0.01] + [0.8] * 10 + [0.99] + output.write_text( + "#! FIELDS time cv\n" + + "".join(f"{ii}.0 {value}\n" for ii, value in enumerate(values)) + ) + deviations = np.arange(len(values), dtype=float) + cv_filter = PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]) + self.assertEqual(cv_filter.sampling["mode"], "uniform") + self.assertEqual(cv_filter.sampling["grid"], {"cv": 10}) + self.assertEqual(cv_filter.sampling["within_bin"], "max_deviation") + selected = cv_filter.select_candidate_ids( + [output], [len(values)], [list(range(len(values)))], 3, [deviations] + )[0] + self.assertEqual(selected, [0, 10, 11]) + + def test_default_sampling_covers_two_cv_grid(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv1 cv2\n" + "0.0 0.1 0.1\n" + "1.0 0.9 0.1\n" + "2.0 0.1 0.9\n" + "3.0 0.9 0.9\n" + ) + cv_filter = PlumedCVFilter(regions=[{"cv1": [0.0, 1.0], "cv2": [0.0, 1.0]}]) + self.assertEqual(cv_filter.sampling["mode"], "grid") + self.assertEqual(cv_filter.sampling["grid"], {"cv1": 10, "cv2": 10}) + selected = cv_filter.select_candidate_ids( + [output], [4], [list(range(4))], 4, [np.arange(4.0)] + ) + self.assertEqual(selected, [[0, 1, 2, 3]]) + def test_uniform_sampling_spans_nonempty_bins(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" @@ -184,6 +244,111 @@ def test_uniform_sampling_spreads_a_small_limit_across_regions(self): [[0, 2]], ) + def test_named_region_grid_sampling_covers_two_cvs_and_records_audit(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time iondistance ionization\n" + "0.00 0.25 0.25\n" + "0.01 0.30 0.30\n" + "0.02 0.75 0.25\n" + "0.03 0.25 0.75\n" + "0.04 0.75 0.75\n" + "0.05 0.75 1.25\n" + ) + deviations = np.asarray([1.0, 9.0, 8.0, 7.0, 6.0, 100.0]) + cv_filter = PlumedCVFilter( + regions=[ + { + "name": "ion_pair", + "conditions": { + "iondistance": [0.2, 1.0], + "ionization": [0.2, 1.0], + }, + } + ], + sampling={ + "mode": "grid", + "grid": {"iondistance": 2, "ionization": 2}, + "within_bin": "max_deviation", + "seed": 20260815, + }, + time_alignment={"start": 0.0, "step": 0.01}, + ) + selected, records, summary = cv_filter.select_candidate_ids_with_audit( + [output], [6], [list(range(6))], 4, [deviations] + ) + self.assertEqual(selected, [[1, 2, 3, 4]]) + self.assertEqual({row["region_names"] for row in records}, {"ion_pair"}) + self.assertEqual(len({row["cell_or_bin"] for row in records}), 4) + self.assertEqual(summary["trust_candidates"], 6) + self.assertEqual(summary["cv_eligible_candidates"], 5) + self.assertEqual(summary["selected"], 4) + self.assertEqual(summary["per_region"]["ion_pair"]["nonempty_cells"], 4) + + def test_min_frame_gap_reports_underfilled_quota(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv\n" + "".join(f"{ii}.0 0.5\n" for ii in range(10)) + ) + deviations = np.arange(10, 0, -1, dtype=float) + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + sampling={ + "mode": "uniform", + "field": "cv", + "n_bins": 1, + "within_bin": "max_deviation", + "min_frame_gap": 3, + }, + ) + selected, _, summary = cv_filter.select_candidate_ids_with_audit( + [output], [10], [list(range(10))], 5, [deviations] + ) + self.assertEqual(selected, [[0, 3, 6, 9]]) + self.assertEqual(summary["underfilled_quota"], 1) + self.assertEqual(summary["min_frame_gap_rejects"], 6) + + def test_named_region_weights_allocate_grid_quota(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time cv\n" + + "".join( + f"{ii}.0 {value}\n" + for ii, value in enumerate([0.1, 0.3, 0.6, 0.9, 2.1, 2.3, 2.6, 2.9]) + ) + ) + cv_filter = PlumedCVFilter( + regions=[ + {"name": "low", "conditions": {"cv": [0.0, 1.0]}}, + { + "name": "high", + "conditions": {"cv": [2.0, 3.0]}, + "weight": 3.0, + }, + ], + sampling={"mode": "uniform", "field": "cv", "n_bins": 4}, + ) + _, records, summary = cv_filter.select_candidate_ids_with_audit( + [output], [8], [list(range(8))], 4 + ) + self.assertEqual(summary["per_region"]["low"]["selected"], 1) + self.assertEqual(summary["per_region"]["high"]["selected"], 3) + self.assertEqual(len(records), 4) + + def test_explicit_time_alignment_fails_closed(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text("#! FIELDS time cv\n0.000 0.5\n0.010 0.5\n0.021 0.5\n") + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + time_alignment={"start": 0.0, "step": 0.01, "atol": 1e-8}, + ) + with self.assertRaises(FatalError): + cv_filter.get_selected_ids([output], [3]) + def test_alignment_fails_closed(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" @@ -211,13 +376,33 @@ def test_invalid_config_and_file_fail_closed(self): {"mode": "uniform", "field": "other"}, {"mode": "uniform", "field": "cv", "n_bins": 0}, {"mode": "uniform", "field": "cv", "within_bin": "first"}, + {"mode": "grid", "grid": {"cv": 2}}, + {"mode": "grid", "grid": {"cv": 2, "other": 2}}, {"mode": "random", "seed": True}, {"mode": "random", "seed": -1}, + {"mode": "random", "min_frame_gap": True}, ] for sampling in invalid_sampling: with self.subTest(sampling=sampling), self.assertRaises(ValueError): PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}], sampling=sampling) + with self.assertRaises(ValueError): + PlumedCVFilter(regions=[{"cv1": [0.0, 1.0]}, {"cv2": [0.0, 1.0]}]) + + for time_alignment in [ + {}, + {"step": 0.0}, + {"step": 1.0, "atol": -1.0}, + {"step": 1.0, "unknown": 2.0}, + ]: + with self.subTest(time_alignment=time_alignment), self.assertRaises( + ValueError + ): + PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + time_alignment=time_alignment, + ) + invalid_outputs = [ "0.0 0.5\n#! FIELDS time cv\n", "#! FIELDS time cv\n0.0 nan\n", From 7fd89755c1dc2ca016132901138c0370a3fa5784 Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Sun, 16 Aug 2026 14:31:36 +0800 Subject: [PATCH 4/6] fix(cv-filter): finalize generic labels and typing --- docs/input.md | 43 +++++++++++++++++++ dpgen2/entrypoint/args.py | 8 ++-- .../selector/conf_selector_frame.py | 32 +++++++++----- .../exploration/selector/plumed_cv_filter.py | 31 ++++++++----- tests/exploration/test_plumed_cv_filter.py | 9 ++++ 5 files changed, 98 insertions(+), 25 deletions(-) diff --git a/docs/input.md b/docs/input.md index 629846a3..c821a7a1 100644 --- a/docs/input.md +++ b/docs/input.md @@ -164,6 +164,49 @@ v: VORONOI_COORDINATION ... PRINT ARG=d,v STRIDE=10 FILE=COLVAR ``` +The condition keys are exact labels from `#! FIELDS`; DPGEN2 has no reserved +CV names such as `iondistance` or `ionization`. Labels are matched by name, not +by their column position, so reordering `PRINT ARG=d,v` to `PRINT ARG=v,d` +does not change a correctly named filter. Semantic PLUMED labels are therefore +safer than numeric column references. + +For a single CV, keep only that condition. Omitting `sampling` then selects +uniformly across 10 equal-width bins by default: + +```json +"cv_filter": { + "regions": [ + { + "name": "target_window", + "conditions": {"reaction_coordinate": [0.2, 2.0]} + } + ] +} +``` + +Represent disjoint intervals as separate named regions. For example, +`[0.2, 2.0)` and `[3.0, 4.0)` of the same CV are a union because regions are +ORed: + +```json +"cv_filter": { + "regions": [ + { + "name": "segment_1", + "conditions": {"reaction_coordinate": [0.2, 2.0]} + }, + { + "name": "segment_2", + "conditions": {"reaction_coordinate": [3.0, 4.0]} + } + ] +} +``` + +To apply another CV as an AND constraint, repeat it inside each segment's +`conditions`. Keeping segments as regions also gives every interval its own +name, population, selected count, and audit provenance. + `LOAD` is only needed for CVs that are not built into the active PLUMED. Build such a plugin with that same PLUMED installation (for example, `plumed mklib ReactiveVoronoi.cpp`); shared libraries from a different compiler or PLUMED diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index 2ac2b107..de6fb5de 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -216,9 +216,11 @@ def lmp_args(): doc_cv_filter = ( "Optional PLUMED CV candidate filter. A region may be a field-to-" "[lower, upper] mapping or a named conditions mapping; regions are " - "combined by OR. By default, one or two common CVs are covered uniformly. " - "Sampling also supports explicit random, uniform, grid, or report modes, " - "with optional frame spacing and time alignment." + "combined by OR. Field names are exact PLUMED FIELDS labels and are " + "matched independently of column order. By default, one or two common " + "CVs are covered uniformly. Sampling also supports explicit random, " + "uniform, grid, or report modes, with optional frame spacing and time " + "alignment." ) return [ diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index da99fe59..42d615d3 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -10,6 +10,7 @@ Optional, Tuple, Union, + cast, ) import dpdata @@ -108,16 +109,21 @@ def select( self.report.record(md_model_devi) id_cand_list = None cv_audit = None - if self.plumed_cv_filter is not None: + plumed_cv_filter = self.plumed_cv_filter + plm_files = None + md_f = None + candidate_ids = None + if plumed_cv_filter is not None: if plm_outputs is None or any(output is None for output in plm_outputs): raise FatalError( "PLUMED CV filtering requires one output per trajectory" ) - md_f = md_model_devi.get(DeviManager.MAX_DEVI_F) + plm_files = [output for output in plm_outputs if output is not None] + md_f = cast(List[np.ndarray], md_model_devi.get(DeviManager.MAX_DEVI_F)) candidate_ids = self.report.get_candidate_ids(None, clear=False) - if self.plumed_cv_filter.sampling is None: - allowed_ids = self.plumed_cv_filter.get_selected_ids( - plm_outputs, + if plumed_cv_filter.sampling is None: + allowed_ids = plumed_cv_filter.get_selected_ids( + plm_files, [len(values) for values in md_f], ) self.report.restrict_candidate_ids(allowed_ids) @@ -126,8 +132,8 @@ def select( sampled_ids, records, summary, - ) = self.plumed_cv_filter.select_candidate_ids_with_audit( - plm_outputs, + ) = plumed_cv_filter.select_candidate_ids_with_audit( + plm_files, [len(values) for values in md_f], candidate_ids, self.max_numb_sel, @@ -138,9 +144,12 @@ def select( id_cand_list = self.report.get_candidate_ids() if id_cand_list is None: id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) - if self.plumed_cv_filter is not None and cv_audit is None: - cv_audit = self.plumed_cv_filter.audit_candidate_ids( - plm_outputs, + if plumed_cv_filter is not None and cv_audit is None: + assert plm_files is not None + assert md_f is not None + assert candidate_ids is not None + cv_audit = plumed_cv_filter.audit_candidate_ids( + plm_files, [len(values) for values in md_f], candidate_ids, id_cand_list, @@ -159,6 +168,7 @@ def select( out_path.mkdir(exist_ok=True) ms.to_deepmd_npy(out_path) # type: ignore if cv_audit is not None: - self.plumed_cv_filter.write_audit(out_path, *cv_audit) + assert plumed_cv_filter is not None + plumed_cv_filter.write_audit(out_path, *cv_audit) return [out_path], copy.deepcopy(self.report) diff --git a/dpgen2/exploration/selector/plumed_cv_filter.py b/dpgen2/exploration/selector/plumed_cv_filter.py index 563a8329..47ecfaec 100644 --- a/dpgen2/exploration/selector/plumed_cv_filter.py +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -22,10 +22,11 @@ class PlumedCVFilter: """Select frames in a union of PLUMED CV regions. - Each region maps field names from ``#! FIELDS`` to a lower-inclusive, - upper-exclusive interval. Fields within a region are ANDed; regions are - ORed. A region may use the legacy bare mapping or the named form - ``{"name": ..., "conditions": {...}}``. + Each region maps exact field labels from ``#! FIELDS`` to a + lower-inclusive, upper-exclusive interval. Labels are matched by name, not + column position, and are otherwise arbitrary. Fields within a region are + ANDed; regions are ORed. A region may use the legacy bare mapping or the + named form ``{"name": ..., "conditions": {...}}``. """ @staticmethod @@ -36,9 +37,11 @@ def args() -> List[Argument]: list, optional=False, doc=( - "A list of PLUMED field-to-[lower, upper] mappings or named " - "regions with name, conditions, and optional weight. Fields " - "within a region are ANDed; regions are ORed." + "A list of exact PLUMED FIELDS-label-to-[lower, upper] " + "mappings or named regions with name, conditions, and " + "optional weight. Labels are matched by name, not column " + "position. Fields within a region are ANDed; regions are " + "ORed." ), ), Argument( @@ -392,7 +395,9 @@ def _sample_grid( max_devi_f: Optional[List[np.ndarray]], rng, ): - within_bin = self.sampling["within_bin"] + sampling = self.sampling + assert sampling is not None + within_bin = sampling["within_bin"] if within_bin == "max_deviation" and max_devi_f is None: raise FatalError("max_deviation sampling requires force model deviations") @@ -414,7 +419,7 @@ def _sample_grid( picked = [] picked_set = set() rejected_by_gap = set() - grid_sizes = tuple(self.sampling["grid"].values()) + grid_sizes = tuple(sampling["grid"].values()) for region, quota in zip(buckets, region_quotas): available = { cell: [item for item in items if item not in picked_set] @@ -458,6 +463,8 @@ def _pick_candidates( picked, rejected_by_gap, ): + sampling = self.sampling + assert sampling is not None if count <= 0: return initially_picked = len(picked) @@ -481,7 +488,7 @@ def _pick_candidates( item[1], ), ) - min_frame_gap = self.sampling["min_frame_gap"] + min_frame_gap = sampling["min_frame_gap"] for candidate in ordered: if len(picked) - initially_picked >= count: break @@ -604,10 +611,12 @@ def _spread_cells(cells, total, grid_sizes): return chosen def _cell_key(self, region_idx, fields, row): + sampling = self.sampling + assert sampling is not None field_idx = {field: idx for idx, field in enumerate(fields)} region = self.regions[region_idx] cell = [] - for field, n_bins in self.sampling["grid"].items(): + for field, n_bins in sampling["grid"].items(): lower, upper = region[field] value = row[field_idx[field]] cell.append( diff --git a/tests/exploration/test_plumed_cv_filter.py b/tests/exploration/test_plumed_cv_filter.py index 315554c8..5bfa8405 100644 --- a/tests/exploration/test_plumed_cv_filter.py +++ b/tests/exploration/test_plumed_cv_filter.py @@ -70,6 +70,15 @@ def test_union_of_regions(self): ) self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 2]]) + def test_field_names_are_header_based_not_positional(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text( + "#! FIELDS time second_cv first_cv\n" "0.0 9.0 0.5\n" "1.0 0.5 9.0\n" + ) + cv_filter = PlumedCVFilter(regions=[{"first_cv": [0.0, 1.0]}]) + self.assertEqual(cv_filter.get_selected_ids([output], [2]), [[0]]) + def test_intervals_are_lower_inclusive_and_upper_exclusive(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" From 26f9c7608f3ed6642ca855cf22ed528463c2b394 Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Sun, 16 Aug 2026 17:17:15 +0800 Subject: [PATCH 5/6] fix(cv-filter): prevent stale PLUMED output reuse --- dpgen2/op/run_lmp.py | 15 ++++++++++++++ tests/op/test_run_lmp.py | 45 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 885c750a..3de6f9b9 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -129,6 +129,11 @@ def execute( # input_files = [lmp_conf_name, lmp_input_name] # input_files = [(Path(task_path) / ii).resolve() for ii in input_files] input_files = [ii.resolve() for ii in Path(task_path).iterdir()] + if plm_output_file in {ii.name for ii in input_files}: + raise FatalError( + f"PLUMED output file {plm_output_file!r} collides with a staged " + "LAMMPS input file" + ) model_files = [Path(ii).resolve() for ii in models] work_dir = Path(task_name) @@ -168,6 +173,16 @@ def execute( set_models(lmp_input_name, model_names) + # A retried task may reuse its working directory. Remove an output + # from an earlier attempt so it cannot be collected as fresh data. + plm_output_path = Path(plm_output_file) + if plm_output_path.is_file() or plm_output_path.is_symlink(): + plm_output_path.unlink() + elif plm_output_path.exists(): + raise FatalError( + f"PLUMED output path {plm_output_file!r} is not a file" + ) + # run lmp command = " ".join([command, "-i", lmp_input_name, "-log", lmp_log_name]) ret, out, err = run_command(command, shell=True) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index b10d0a42..0edf7c1d 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -12,6 +12,7 @@ OP, OPIO, Artifact, + FatalError, OPIOSign, TransientError, ) @@ -78,8 +79,11 @@ def test_plm_output_file_config(self): @patch("dpgen2.op.run_lmp.run_command") def test_plm_output_file_collection(self, mocked_run): - mocked_run.return_value = (0, "", "") - (self.task_path / "COLVAR").write_text("#! FIELDS time cv\n0.0 0.5\n") + def run_with_plumed_output(*args, **kwargs): + Path("COLVAR").write_text("#! FIELDS time cv\n0.0 0.5\n") + return 0, "", "" + + mocked_run.side_effect = run_with_plumed_output out = RunLmp().execute( OPIO( { @@ -92,6 +96,43 @@ def test_plm_output_file_collection(self, mocked_run): ) self.assertEqual(out["plm_output"], Path(self.task_name) / "COLVAR") + @patch("dpgen2.op.run_lmp.run_command") + def test_plm_output_file_rejects_staged_input(self, mocked_run): + (self.task_path / "COLVAR").write_text("stale input") + with self.assertRaisesRegex(FatalError, "collides with a staged"): + RunLmp().execute( + OPIO( + { + "config": {"plm_output_file": "COLVAR"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + mocked_run.assert_not_called() + + @patch("dpgen2.op.run_lmp.run_command") + def test_plm_output_file_does_not_reuse_stale_output(self, mocked_run): + mocked_run.return_value = (0, "", "") + Path(self.task_name).mkdir() + stale_output = Path(self.task_name) / "COLVAR" + stale_output.write_text("stale output") + + out = RunLmp().execute( + OPIO( + { + "config": {"plm_output_file": "COLVAR"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + + self.assertIsNone(out["plm_output"]) + self.assertFalse(stale_output.exists()) + @patch("dpgen2.op.run_lmp.run_command") def test_success(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", "")] From b528c1c51cbeaa8f5488f7cf6f14431c81e49351 Mon Sep 17 00:00:00 2001 From: Pengchao_Zhang Date: Tue, 1 Sep 2026 11:07:16 +0800 Subject: [PATCH 6/6] fix(cv-filter): address upstream review findings --- docs/input.md | 24 ++- dpgen2/entrypoint/args.py | 4 +- .../report/report_adaptive_lower.py | 5 +- .../report/report_trust_levels_base.py | 12 +- .../selector/conf_selector_frame.py | 15 +- .../exploration/selector/plumed_cv_filter.py | 172 +++++++++++++----- dpgen2/op/run_lmp.py | 57 ++++-- dpgen2/op/select_confs.py | 14 +- dpgen2/utils/download_dpgen2_artifacts.py | 1 + tests/entrypoint/test_submit_args.py | 9 + tests/exploration/test_conf_selector_frame.py | 49 ++++- tests/exploration/test_plumed_cv_filter.py | 86 ++++++--- .../exploration/test_report_adaptive_lower.py | 21 +++ tests/exploration/test_report_trust_levels.py | 14 ++ tests/mocked_ops.py | 1 + tests/op/test_run_lmp.py | 56 +++++- tests/utils/test_dl_dpgen2_arti.py | 15 ++ 17 files changed, 425 insertions(+), 130 deletions(-) diff --git a/docs/input.md b/docs/input.md index c821a7a1..0a008019 100644 --- a/docs/input.md +++ b/docs/input.md @@ -147,7 +147,8 @@ coverage or an explicitly configured selection policy: {"d": [0.08, 0.12]}, {"v": [1.8, 2.2]} ], - "sampling": {"mode": "report"} + "sampling": {"mode": "report"}, + "time_alignment": {"start": 0.0, "step": 0.01} } } ``` @@ -164,6 +165,11 @@ v: VORONOI_COORDINATION ... PRINT ARG=d,v STRIDE=10 FILE=COLVAR ``` +`plm_output_file` defaults to `COLVAR`; set it explicitly when `PRINT FILE` +uses another name. `time_alignment` is required because equal row counts alone +cannot detect a phase offset between the trajectory and COLVAR. Set `start` and +`step` to the times of trajectory frame 0 and one frame interval, respectively. + The condition keys are exact labels from `#! FIELDS`; DPGEN2 has no reserved CV names such as `iondistance` or `ionization`. Labels are matched by name, not by their column position, so reordering `PRINT ARG=d,v` to `PRINT ARG=v,d` @@ -180,7 +186,8 @@ uniformly across 10 equal-width bins by default: "name": "target_window", "conditions": {"reaction_coordinate": [0.2, 2.0]} } - ] + ], + "time_alignment": {"start": 0.0, "step": 0.01} } ``` @@ -199,7 +206,8 @@ ORed: "name": "segment_2", "conditions": {"reaction_coordinate": [3.0, 4.0]} } - ] + ], + "time_alignment": {"start": 0.0, "step": 0.01} } ``` @@ -237,7 +245,8 @@ these defaults: "n_bins": 10, "within_bin": "max_deviation", "seed": 20260815 - } + }, + "time_alignment": {"start": 0.0, "step": 0.01} } ``` @@ -286,7 +295,7 @@ names and weights: "time_alignment": { "start": 0.0, "step": 0.01, - "atol": 1e-8 + "atol": 1e-6 } } ``` @@ -297,8 +306,9 @@ covers separated non-empty cells before adding extra frames, and then uses `within_bin` inside each cell. `min_frame_gap` is a minimum frame-index separation within each trajectory. A spacing constraint may leave the result underfilled; DPGEN2 reports this instead of silently relaxing the constraint. -`time_alignment` additionally verifies `time = start + frame * step` with the -configured absolute tolerance. +`time_alignment` verifies `time = start + frame * step` with an absolute +tolerance of `1e-6` by default. Increase `atol` explicitly when a coarse PLUMED +`FMT` rounds time more strongly. When a CV filter is active, the selected DeepMD data directory contains `cv_selection.csv` and `cv_selection_summary.json`. They record trajectory and diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index de6fb5de..e96ed142 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -219,8 +219,8 @@ def lmp_args(): "combined by OR. Field names are exact PLUMED FIELDS labels and are " "matched independently of column order. By default, one or two common " "CVs are covered uniformly. Sampling also supports explicit random, " - "uniform, grid, or report modes, with optional frame spacing and time " - "alignment." + "uniform, grid, or report modes, with optional frame spacing. Explicit " + "time alignment is required to bind COLVAR rows to trajectory frames." ) return [ diff --git a/dpgen2/exploration/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index 42e20b63..5521bbad 100644 --- a/dpgen2/exploration/report/report_adaptive_lower.py +++ b/dpgen2/exploration/report/report_adaptive_lower.py @@ -279,6 +279,9 @@ def record( self.accur = self.accur - self.candi self.model_devi = model_devi self._no_candidate = len(self.candi) == 0 + self._update_ratios() + + def _update_ratios(self) -> None: self._failed_ratio = float(len(self.failed)) / float(self.nframes) self._accurate_ratio = float(len(self.accur)) / float(self.nframes) self._candidate_ratio = float(len(self.candi)) / float(self.nframes) @@ -383,7 +386,7 @@ def restrict_candidate_ids( for frame_idx in frame_ids } self.candi &= allowed - self._no_candidate = len(self.candi) == 0 + self._update_ratios() def get_candidate_ids( self, diff --git a/dpgen2/exploration/report/report_trust_levels_base.py b/dpgen2/exploration/report/report_trust_levels_base.py index 3e28d4ef..79d9e263 100644 --- a/dpgen2/exploration/report/report_trust_levels_base.py +++ b/dpgen2/exploration/report/report_trust_levels_base.py @@ -138,14 +138,18 @@ def record( assert len(self.traj_fail) == ntraj self.model_devi = model_devi self._no_candidate = sum([len(ii) for ii in self.traj_cand]) == 0 + self._update_ratios() + + def _update_ratios(self) -> None: + nframes = float(sum(self.traj_nframes)) self._failed_ratio = float(sum([len(ii) for ii in self.traj_fail])) / float( - sum(self.traj_nframes) + nframes ) self._accurate_ratio = float(sum([len(ii) for ii in self.traj_accu])) / float( - sum(self.traj_nframes) + nframes ) self._candidate_ratio = float(sum([len(ii) for ii in self.traj_cand])) / float( - sum(self.traj_nframes) + nframes ) def _get_indexes( @@ -246,7 +250,7 @@ def restrict_candidate_ids( candidates & set(allowed) for candidates, allowed in zip(self.traj_cand, allowed_ids) ] - self._no_candidate = sum(len(candidates) for candidates in self.traj_cand) == 0 + self._update_ratios() @abstractmethod def get_candidate_ids( diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index 42d615d3..17dd65b0 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -111,20 +111,30 @@ def select( cv_audit = None plumed_cv_filter = self.plumed_cv_filter plm_files = None + loaded_plm_outputs = None md_f = None candidate_ids = None if plumed_cv_filter is not None: - if plm_outputs is None or any(output is None for output in plm_outputs): + if ( + plm_outputs is None + or len(plm_outputs) != ntraj + or any(output is None for output in plm_outputs) + ): raise FatalError( "PLUMED CV filtering requires one output per trajectory" ) - plm_files = [output for output in plm_outputs if output is not None] + plm_files = cast(List[Path], list(plm_outputs)) md_f = cast(List[np.ndarray], md_model_devi.get(DeviManager.MAX_DEVI_F)) candidate_ids = self.report.get_candidate_ids(None, clear=False) if plumed_cv_filter.sampling is None: + loaded_plm_outputs = plumed_cv_filter.load_outputs( + plm_files, + [len(values) for values in md_f], + ) allowed_ids = plumed_cv_filter.get_selected_ids( plm_files, [len(values) for values in md_f], + loaded_outputs=loaded_plm_outputs, ) self.report.restrict_candidate_ids(allowed_ids) else: @@ -154,6 +164,7 @@ def select( candidate_ids, id_cand_list, md_f, + loaded_outputs=loaded_plm_outputs, ) ms = self.traj_render.get_confs( diff --git a/dpgen2/exploration/selector/plumed_cv_filter.py b/dpgen2/exploration/selector/plumed_cv_filter.py index 47ecfaec..9949f330 100644 --- a/dpgen2/exploration/selector/plumed_cv_filter.py +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -18,6 +18,8 @@ FatalError, ) +PlumedOutputs = List[Tuple[List[str], np.ndarray]] + class PlumedCVFilter: """Select frames in a union of PLUMED CV regions. @@ -48,13 +50,54 @@ def args() -> List[Argument]: "sampling", dict, [ - Argument("mode", str, optional=False), - Argument("field", str, optional=True, default=None), - Argument("n_bins", int, optional=True, default=10), - Argument("grid", dict, optional=True, default=None), - Argument("within_bin", str, optional=True, default="random"), - Argument("seed", int, optional=True, default=0), - Argument("min_frame_gap", int, optional=True, default=0), + Argument( + "mode", + str, + optional=False, + doc="Selection mode: random, uniform, grid, or report.", + ), + Argument( + "field", + str, + optional=True, + default=None, + doc="CV field covered by one-dimensional uniform sampling.", + ), + Argument( + "n_bins", + int, + optional=True, + default=10, + doc="Number of equal-width bins for uniform sampling.", + ), + Argument( + "grid", + dict, + optional=True, + default=None, + doc="Two CV field-to-bin-count mappings for grid sampling.", + ), + Argument( + "within_bin", + str, + optional=True, + default="random", + doc="Frame policy within a bin: random or max_deviation.", + ), + Argument( + "seed", + int, + optional=True, + default=0, + doc="Non-negative random seed used by sampling.", + ), + Argument( + "min_frame_gap", + int, + optional=True, + default=0, + doc="Minimum frame-index separation within a trajectory.", + ), ], optional=True, default=None, @@ -68,13 +111,29 @@ def args() -> List[Argument]: "time_alignment", dict, [ - Argument("start", float, optional=True, default=0.0), - Argument("step", float, optional=False), - Argument("atol", float, optional=True, default=1e-8), + Argument( + "start", + float, + optional=True, + default=0.0, + doc="Expected PLUMED time of trajectory frame zero.", + ), + Argument( + "step", + float, + optional=False, + doc="Expected PLUMED time interval between trajectory frames.", + ), + Argument( + "atol", + float, + optional=True, + default=1e-6, + doc="Absolute tolerance for PLUMED time alignment.", + ), ], - optional=True, - default=None, - doc="Optional expected PLUMED time = start + frame * step.", + optional=False, + doc="Required expected PLUMED time = start + frame * step.", ), ] @@ -232,9 +291,9 @@ def _normalize_sampling(self, sampling: Optional[Dict]): return normalized @staticmethod - def _normalize_time_alignment(time_alignment: Optional[Dict]): + def _normalize_time_alignment(time_alignment: Optional[Dict]) -> Dict: if time_alignment is None: - return None + raise ValueError("PLUMED CV time_alignment is required") if not isinstance(time_alignment, dict): raise ValueError("PLUMED CV time_alignment must be a dict") unknown = set(time_alignment) - {"start", "step", "atol"} @@ -243,7 +302,7 @@ def _normalize_time_alignment(time_alignment: Optional[Dict]): try: start = float(time_alignment.get("start", 0.0)) step = float(time_alignment["step"]) - atol = float(time_alignment.get("atol", 1e-8)) + atol = float(time_alignment.get("atol", 1e-6)) except (KeyError, TypeError, ValueError) as exc: raise ValueError("invalid PLUMED CV time_alignment") from exc if not np.all(np.isfinite([start, step, atol])) or step <= 0 or atol < 0: @@ -254,14 +313,23 @@ def get_selected_ids( self, files: List[Path], nframes: List[int], + loaded_outputs: Optional[PlumedOutputs] = None, ) -> List[List[int]]: - outputs = self._load_outputs(files, nframes) + outputs = ( + self.load_outputs(files, nframes) + if loaded_outputs is None + else loaded_outputs + ) selected = [] for fields, values in outputs: masks = self._region_masks(fields, values) selected.append(np.flatnonzero(np.logical_or.reduce(masks)).tolist()) return selected + def load_outputs(self, files: List[Path], nframes: List[int]) -> PlumedOutputs: + """Parse and validate PLUMED outputs once for a selection pass.""" + return self._load_outputs(files, nframes) + def select_candidate_ids( self, files: List[Path], @@ -292,7 +360,7 @@ def select_candidate_ids_with_audit( if max_nframes is not None and max_nframes < 0: raise ValueError("max_nframes must be non-negative") - outputs = self._load_outputs(files, nframes) + outputs = self.load_outputs(files, nframes) masks_by_traj = [ self._region_masks(fields, values) for fields, values in outputs ] @@ -344,9 +412,14 @@ def audit_candidate_ids( candidate_ids: List[List[int]], selected_ids: List[List[int]], max_devi_f: Optional[List[np.ndarray]] = None, + loaded_outputs: Optional[PlumedOutputs] = None, ): """Audit filtering followed by the report's existing sampling policy.""" - outputs = self._load_outputs(files, nframes) + outputs = ( + self.load_outputs(files, nframes) + if loaded_outputs is None + else loaded_outputs + ) masks_by_traj = [ self._region_masks(fields, values) for fields, values in outputs ] @@ -590,24 +663,26 @@ def _spread_cells(cells, total, grid_sizes): ), ) ] + + def squared_distance(left_cell, right_cell): + return sum( + ((left - right) / max(size - 1, 1)) ** 2 + for left, right, size in zip(left_cell, right_cell, grid_sizes) + ) + chosen = [cells[0]] + chosen_set = {cells[0]} + nearest = {cell: squared_distance(cell, cells[0]) for cell in cells[1:]} while len(chosen) < total: - best = None - best_distance = -1.0 - for cell in cells: - if cell in chosen: - continue - distance = min( - sum( - ((left - right) / max(size - 1, 1)) ** 2 - for left, right, size in zip(cell, other, grid_sizes) - ) - for other in chosen - ) - if distance > best_distance: - best = cell - best_distance = distance + best = max( + (cell for cell in cells if cell not in chosen_set), + key=lambda cell: nearest[cell], + ) chosen.append(best) + chosen_set.add(best) + for cell in cells: + if cell not in chosen_set: + nearest[cell] = min(nearest[cell], squared_distance(cell, best)) return chosen def _cell_key(self, region_idx, fields, row): @@ -772,7 +847,7 @@ def write_audit(out_path: Path, records, summary): json.dump(summary, handle, indent=2, sort_keys=True) handle.write("\n") - def _load_outputs(self, files: List[Path], nframes: List[int]): + def _load_outputs(self, files: List[Path], nframes: List[int]) -> PlumedOutputs: if len(files) != len(nframes): raise FatalError("PLUMED outputs and trajectories have different lengths") outputs = [] @@ -783,21 +858,20 @@ def _load_outputs(self, files: List[Path], nframes: List[int]): f"PLUMED output {file} has {len(values)} rows, expected " f"{expected_nframes}; PRINT STRIDE must match the trajectory stride" ) - if self.time_alignment is not None: - time = values[:, fields.index("time")] - expected_time = ( - self.time_alignment["start"] - + np.arange(expected_nframes) * self.time_alignment["step"] + time = values[:, fields.index("time")] + expected_time = ( + self.time_alignment["start"] + + np.arange(expected_nframes) * self.time_alignment["step"] + ) + if not np.allclose( + time, + expected_time, + rtol=0.0, + atol=self.time_alignment["atol"], + ): + raise FatalError( + f"PLUMED time in {file} does not match configured frame times" ) - if not np.allclose( - time, - expected_time, - rtol=0.0, - atol=self.time_alignment["atol"], - ): - raise FatalError( - f"PLUMED time in {file} does not match configured frame times" - ) outputs.append((fields, values)) return outputs diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 3de6f9b9..4bdde01e 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -21,6 +21,9 @@ Variant, dargs, ) +from dargs.dargs import ( + ArgumentError, +) from dflow.python import ( OP, OPIO, @@ -118,7 +121,10 @@ def execute( On the failure of LAMMPS execution. Handle different failure cases? e.g. loss atoms. """ config = ip["config"] if ip["config"] is not None else {} - config = RunLmp.normalize_config(config) + try: + config = RunLmp.normalize_config(config) + except ArgumentError as exc: + raise FatalError(f"invalid LAMMPS configuration: {exc}") from exc command = config["command"] teacher_model: Optional[BinaryFileInput] = config["teacher_model_path"] shuffle_models: Optional[bool] = config["shuffle_models"] @@ -146,7 +152,33 @@ def execute( teacher_model.save_as_file(teacher_model_file) model_files = [Path(teacher_model_file).resolve()] + model_files + generated_names = { + lmp_log_name, + lmp_model_devi_name, + lmp_traj_name, + plm_output_name, + "job.json", + } + for idx in range(len(model_files)): + generated_names.add(model_name_pattern % idx) + generated_names.add(pytorch_model_name_pattern % idx) + if plm_output_file in generated_names: + raise FatalError( + f"PLUMED output file {plm_output_file!r} collides with a generated " + "LAMMPS, PLUMED, or model file" + ) + with set_directory(work_dir): + # Remove a pre-existing output before creating any task links. This + # prevents stale CV data from surviving a retried task. + plm_output_path = Path(plm_output_file) + if plm_output_path.is_file() or plm_output_path.is_symlink(): + plm_output_path.unlink() + elif plm_output_path.exists(): + raise FatalError( + f"PLUMED output path {plm_output_file!r} is not a file" + ) + # link input files for ii in input_files: iname = ii.name @@ -173,16 +205,6 @@ def execute( set_models(lmp_input_name, model_names) - # A retried task may reuse its working directory. Remove an output - # from an earlier attempt so it cannot be collected as fresh data. - plm_output_path = Path(plm_output_file) - if plm_output_path.is_file() or plm_output_path.is_symlink(): - plm_output_path.unlink() - elif plm_output_path.exists(): - raise FatalError( - f"PLUMED output path {plm_output_file!r} is not a file" - ) - # run lmp command = " ".join([command, "-i", lmp_input_name, "-log", lmp_log_name]) ret, out, err = run_command(command, shell=True) @@ -248,8 +270,8 @@ def lmp_args(): doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature" doc_use_hdf5 = "Use HDF5 to store trajs and model_devis" doc_plm_output_file = ( - "PLUMED output artifact to collect. Set this to the FILE used by " - "PLUMED PRINT when filtering candidates by CV." + "PLUMED CV output artifact to collect. It must match the FILE used " + "by PLUMED PRINT and defaults to COLVAR." ) doc_extra_output_files = "Extra output file names, support wildcards" return [ @@ -286,7 +308,10 @@ def lmp_args(): "plm_output_file", str, optional=True, - default=plm_output_name, + default="COLVAR", + extra_check=lambda value: value not in {"", ".", ".."} + and Path(value).name == value, + extra_check_errmsg="must be a file name, not a path", doc=doc_plm_output_file, ), Argument( @@ -304,10 +329,6 @@ def normalize_config(data={}): base = Argument("base", dict, ta) data = base.normalize_value(data, trim_pattern="_*") base.check_value(data, strict=True) - if data["plm_output_file"] in {"", ".", ".."} or ( - Path(data["plm_output_file"]).name != data["plm_output_file"] - ): - raise ValueError("plm_output_file must be a file name, not a path") return data diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index 65e0b5a8..e43f4012 100644 --- a/dpgen2/op/select_confs.py +++ b/dpgen2/op/select_confs.py @@ -91,13 +91,13 @@ def execute( trajs, model_devis, optional_outputs ) - select_kwargs = { - "type_map": type_map, - "optional_outputs": optional_outputs, - } - if plm_outputs is not None: - select_kwargs["plm_outputs"] = plm_outputs - confs, report = conf_selector.select(trajs, model_devis, **select_kwargs) + confs, report = conf_selector.select( + trajs, + model_devis, + type_map=type_map, + optional_outputs=optional_outputs, + plm_outputs=plm_outputs, + ) return OPIO( { diff --git a/dpgen2/utils/download_dpgen2_artifacts.py b/dpgen2/utils/download_dpgen2_artifacts.py index 67c2aaf0..c308aed9 100644 --- a/dpgen2/utils/download_dpgen2_artifacts.py +++ b/dpgen2/utils/download_dpgen2_artifacts.py @@ -64,6 +64,7 @@ def add_output( .add_output("logs") .add_output("trajs") .add_output("model_devis") + .add_output("plm_output") .add_output("extra_outputs"), "prep-run-fp": DownloadDefinition() .add_input("confs") diff --git a/tests/entrypoint/test_submit_args.py b/tests/entrypoint/test_submit_args.py index ec7e4582..7b0423de 100644 --- a/tests/entrypoint/test_submit_args.py +++ b/tests/entrypoint/test_submit_args.py @@ -11,6 +11,9 @@ import dpdata import numpy as np +from dargs.dargs import ( + ArgumentError, +) # isort: off from .context import ( @@ -155,6 +158,12 @@ def test_bohrium(self): }, ) + def test_plm_output_file_path_is_rejected_at_submit(self): + data = json.loads(new_str) + data["explore"]["config"]["plm_output_file"] = "outputs/COLVAR" + with self.assertRaisesRegex(ArgumentError, "must be a file name"): + normalize(data) + old_str = textwrap.dedent( """ diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py index 7d299afd..29cbfc18 100644 --- a/tests/exploration/test_conf_selector_frame.py +++ b/tests/exploration/test_conf_selector_frame.py @@ -7,9 +7,15 @@ from pathlib import ( Path, ) +from unittest.mock import ( + patch, +) import dpdata import numpy as np +from dflow.python import ( + FatalError, +) # isort: off from .context import ( @@ -100,25 +106,49 @@ def test_plumed_filter_precedes_max_selection(self): plm_outputs = [Path("foo.cv"), Path("bar.cv")] for output in plm_outputs: output.write_text("#! FIELDS time cv\n0.0 0.5\n1.0 0.5\n2.0 1.5\n") + cv_filter = PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + sampling={"mode": "report"}, + time_alignment={"start": 0.0, "step": 1.0}, + ) conf_selector = ConfSelectorFrames( TrajRenderLammps(), ExplorationReportTrustLevelsMax(0.1, 0.5), max_numb_sel=1, - plumed_cv_filter=PlumedCVFilter( - regions=[{"cv": [0.0, 1.0]}], sampling={"mode": "report"} - ), - ) - confs, _ = conf_selector.select( - self.trajs, - self.model_devis, - self.type_map, - plm_outputs=plm_outputs, + plumed_cv_filter=cv_filter, ) + with patch.object( + cv_filter, "_load_outputs", wraps=cv_filter._load_outputs + ) as mocked_load: + confs, _ = conf_selector.select( + self.trajs, + self.model_devis, + self.type_map, + plm_outputs=plm_outputs, + ) + mocked_load.assert_called_once() ms = dpdata.MultiSystems(type_map=self.type_map) ms.from_deepmd_npy(confs[0], labeled=False) self.assertEqual(ms[0].get_nframes(), 1) self.assertAlmostEqual(ms[0]["coords"][0][0][1], 3.87, places=2) + def test_plumed_filter_requires_one_output_per_trajectory(self): + conf_selector = ConfSelectorFrames( + TrajRenderLammps(), + ExplorationReportTrustLevelsMax(0.1, 0.5), + plumed_cv_filter=PlumedCVFilter( + regions=[{"cv": [0.0, 1.0]}], + time_alignment={"start": 0.0, "step": 1.0}, + ), + ) + with self.assertRaisesRegex(FatalError, "one output per trajectory"): + conf_selector.select( + self.trajs, + self.model_devis, + self.type_map, + plm_outputs=[Path("foo.cv")], + ) + def test_plumed_uniform_sampling_is_final_selection(self): plm_outputs = [Path("foo.cv"), Path("bar.cv")] for output in plm_outputs: @@ -135,6 +165,7 @@ def test_plumed_uniform_sampling_is_final_selection(self): "n_bins": 10, "within_bin": "max_deviation", }, + time_alignment={"start": 0.0, "step": 1.0}, ), ) confs, _ = conf_selector.select( diff --git a/tests/exploration/test_plumed_cv_filter.py b/tests/exploration/test_plumed_cv_filter.py index 5bfa8405..87211910 100644 --- a/tests/exploration/test_plumed_cv_filter.py +++ b/tests/exploration/test_plumed_cv_filter.py @@ -8,6 +8,9 @@ from dargs import ( Argument, ) +from dargs.dargs import ( + ArgumentKeyError, +) from dflow.python import ( FatalError, ) @@ -17,6 +20,11 @@ ) +def make_filter(**kwargs): + kwargs.setdefault("time_alignment", {"start": 0.0, "step": 1.0}) + return PlumedCVFilter(**kwargs) + + class TestPlumedCVFilter(unittest.TestCase): def test_config_schema(self): schema = Argument("cv_filter", dict, PlumedCVFilter.args()) @@ -28,6 +36,7 @@ def test_config_schema(self): "field": "distance", "n_bins": 8, }, + "time_alignment": {"start": 0.0, "step": 1.0}, } ) schema.check_value(config, strict=True) @@ -57,7 +66,23 @@ def test_config_schema(self): ) schema.check_value(grid_config, strict=True) self.assertEqual(grid_config["sampling"]["min_frame_gap"], 5) - self.assertEqual(grid_config["time_alignment"]["atol"], 1e-8) + self.assertEqual(grid_config["time_alignment"]["atol"], 1e-6) + + def test_time_alignment_is_required(self): + schema = Argument("cv_filter", dict, PlumedCVFilter.args()) + config = schema.normalize_value({"regions": [{"cv": [0.0, 1.0]}]}) + with self.assertRaisesRegex(ArgumentKeyError, "time_alignment"): + schema.check_value(config, strict=True) + + with self.assertRaisesRegex(ValueError, "time_alignment is required"): + PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]) + + def test_spread_cells_preserves_deterministic_coverage_order(self): + cells = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 2)] + self.assertEqual( + PlumedCVFilter._spread_cells(cells, 5, (3, 3)), + [(0, 0), (2, 2), (1, 1), (0, 1), (1, 0)], + ) def test_union_of_regions(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -65,7 +90,7 @@ def test_union_of_regions(self): output.write_text( "#! FIELDS time d1 d2\n" "0.0 0.5 2.0\n" "1.0 0.5 4.0\n" "2.0 2.5 9.0\n" ) - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"d1": [0.0, 1.0], "d2": [1.0, 3.0]}, {"d1": [2.0, 3.0]}] ) self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 2]]) @@ -76,14 +101,14 @@ def test_field_names_are_header_based_not_positional(self): output.write_text( "#! FIELDS time second_cv first_cv\n" "0.0 9.0 0.5\n" "1.0 0.5 9.0\n" ) - cv_filter = PlumedCVFilter(regions=[{"first_cv": [0.0, 1.0]}]) + cv_filter = make_filter(regions=[{"first_cv": [0.0, 1.0]}]) self.assertEqual(cv_filter.get_selected_ids([output], [2]), [[0]]) def test_intervals_are_lower_inclusive_and_upper_exclusive(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" output.write_text("#! FIELDS time cv\n0.0 1.0\n1.0 2.0\n2.0 3.0\n") - cv_filter = PlumedCVFilter(regions=[{"cv": [1.0, 3.0]}]) + cv_filter = make_filter(regions=[{"cv": [1.0, 3.0]}]) self.assertEqual(cv_filter.get_selected_ids([output], [3]), [[0, 1]]) def test_random_sampling_is_reproducible(self): @@ -97,10 +122,10 @@ def test_random_sampling_is_reproducible(self): "regions": [{"cv": [0.25, 0.75]}], "sampling": {"mode": "random", "seed": 17}, } - selected = PlumedCVFilter(**kwargs).select_candidate_ids( + selected = make_filter(**kwargs).select_candidate_ids( [output], [20], [list(range(20))], 5 ) - repeated = PlumedCVFilter(**kwargs).select_candidate_ids( + repeated = make_filter(**kwargs).select_candidate_ids( [output], [20], [list(range(20))], 5 ) self.assertEqual(selected, repeated) @@ -116,7 +141,7 @@ def test_default_sampling_covers_one_cv_uniformly(self): + "".join(f"{ii}.0 {value}\n" for ii, value in enumerate(values)) ) deviations = np.arange(len(values), dtype=float) - cv_filter = PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]) + cv_filter = make_filter(regions=[{"cv": [0.0, 1.0]}]) self.assertEqual(cv_filter.sampling["mode"], "uniform") self.assertEqual(cv_filter.sampling["grid"], {"cv": 10}) self.assertEqual(cv_filter.sampling["within_bin"], "max_deviation") @@ -135,7 +160,7 @@ def test_default_sampling_covers_two_cv_grid(self): "2.0 0.1 0.9\n" "3.0 0.9 0.9\n" ) - cv_filter = PlumedCVFilter(regions=[{"cv1": [0.0, 1.0], "cv2": [0.0, 1.0]}]) + cv_filter = make_filter(regions=[{"cv1": [0.0, 1.0], "cv2": [0.0, 1.0]}]) self.assertEqual(cv_filter.sampling["mode"], "grid") self.assertEqual(cv_filter.sampling["grid"], {"cv1": 10, "cv2": 10}) selected = cv_filter.select_candidate_ids( @@ -154,7 +179,7 @@ def test_uniform_sampling_spans_nonempty_bins(self): deviations = np.arange(len(values), dtype=float) deviations[0] = 100.0 deviations[1] = 200.0 - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"cv": [0.0, 1.0]}], sampling={ "mode": "uniform", @@ -186,10 +211,10 @@ def test_uniform_random_within_bin_is_reproducible(self): "seed": 29, }, } - selected = PlumedCVFilter(**kwargs).select_candidate_ids( + selected = make_filter(**kwargs).select_candidate_ids( [output], [20], [list(range(20))], 5 ) - repeated = PlumedCVFilter(**kwargs).select_candidate_ids( + repeated = make_filter(**kwargs).select_candidate_ids( [output], [20], [list(range(20))], 5 ) self.assertEqual(selected, repeated) @@ -206,7 +231,7 @@ def test_uniform_sampling_balances_regions_and_honors_and_conditions(self): "3.0 2.9 0.5\n" "4.0 0.5 1.5\n" ) - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[ {"cv": [0.0, 1.0], "gate": [0.0, 1.0]}, {"cv": [2.0, 3.0], "gate": [0.0, 1.0]}, @@ -231,7 +256,7 @@ def test_uniform_sampling_handles_empty_bins(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" output.write_text("#! FIELDS time cv\n0.0 0.05\n1.0 0.45\n2.0 0.95\n") - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"cv": [0.0, 1.0]}], sampling={"mode": "uniform", "field": "cv", "n_bins": 10}, ) @@ -244,7 +269,7 @@ def test_uniform_sampling_spreads_a_small_limit_across_regions(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" output.write_text("#! FIELDS time cv\n0.0 0.5\n1.0 2.5\n2.0 4.5\n") - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"cv": [0.0, 1.0]}, {"cv": [2.0, 3.0]}, {"cv": [4.0, 5.0]}], sampling={"mode": "uniform", "field": "cv", "n_bins": 4}, ) @@ -266,7 +291,7 @@ def test_named_region_grid_sampling_covers_two_cvs_and_records_audit(self): "0.05 0.75 1.25\n" ) deviations = np.asarray([1.0, 9.0, 8.0, 7.0, 6.0, 100.0]) - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[ { "name": "ion_pair", @@ -302,7 +327,7 @@ def test_min_frame_gap_reports_underfilled_quota(self): "#! FIELDS time cv\n" + "".join(f"{ii}.0 0.5\n" for ii in range(10)) ) deviations = np.arange(10, 0, -1, dtype=float) - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"cv": [0.0, 1.0]}], sampling={ "mode": "uniform", @@ -329,7 +354,7 @@ def test_named_region_weights_allocate_grid_quota(self): for ii, value in enumerate([0.1, 0.3, 0.6, 0.9, 2.1, 2.3, 2.6, 2.9]) ) ) - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[ {"name": "low", "conditions": {"cv": [0.0, 1.0]}}, { @@ -351,19 +376,30 @@ def test_explicit_time_alignment_fails_closed(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" output.write_text("#! FIELDS time cv\n0.000 0.5\n0.010 0.5\n0.021 0.5\n") - cv_filter = PlumedCVFilter( + cv_filter = make_filter( regions=[{"cv": [0.0, 1.0]}], time_alignment={"start": 0.0, "step": 0.01, "atol": 1e-8}, ) with self.assertRaises(FatalError): cv_filter.get_selected_ids([output], [3]) + def test_equal_row_count_with_phase_offset_fails_closed(self): + with tempfile.TemporaryDirectory() as tmpdir: + output = Path(tmpdir) / "COLVAR" + output.write_text("#! FIELDS time cv\n10.0 0.5\n20.0 0.5\n30.0 0.5\n") + cv_filter = make_filter( + regions=[{"cv": [0.0, 1.0]}], + time_alignment={"start": 0.0, "step": 10.0}, + ) + with self.assertRaisesRegex(FatalError, "configured frame times"): + cv_filter.get_selected_ids([output], [3]) + def test_alignment_fails_closed(self): with tempfile.TemporaryDirectory() as tmpdir: output = Path(tmpdir) / "COLVAR" output.write_text("#! FIELDS time cv\n0.0 0.5\n") with self.assertRaises(FatalError): - PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + make_filter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( [output], [2] ) @@ -377,7 +413,7 @@ def test_invalid_config_and_file_fail_closed(self): ] for regions in invalid_regions: with self.subTest(regions=regions), self.assertRaises(ValueError): - PlumedCVFilter(regions=regions) + make_filter(regions=regions) invalid_sampling = [ {"mode": "weighted"}, @@ -393,10 +429,10 @@ def test_invalid_config_and_file_fail_closed(self): ] for sampling in invalid_sampling: with self.subTest(sampling=sampling), self.assertRaises(ValueError): - PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}], sampling=sampling) + make_filter(regions=[{"cv": [0.0, 1.0]}], sampling=sampling) with self.assertRaises(ValueError): - PlumedCVFilter(regions=[{"cv1": [0.0, 1.0]}, {"cv2": [0.0, 1.0]}]) + make_filter(regions=[{"cv1": [0.0, 1.0]}, {"cv2": [0.0, 1.0]}]) for time_alignment in [ {}, @@ -407,7 +443,7 @@ def test_invalid_config_and_file_fail_closed(self): with self.subTest(time_alignment=time_alignment), self.assertRaises( ValueError ): - PlumedCVFilter( + make_filter( regions=[{"cv": [0.0, 1.0]}], time_alignment=time_alignment, ) @@ -422,13 +458,13 @@ def test_invalid_config_and_file_fail_closed(self): ] with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(FatalError): - PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + make_filter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( [Path(tmpdir) / "missing"], [1] ) for index, content in enumerate(invalid_outputs): output = Path(tmpdir) / f"COLVAR.{index}" output.write_text(content) with self.subTest(content=content), self.assertRaises(FatalError): - PlumedCVFilter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( + make_filter(regions=[{"cv": [0.0, 1.0]}]).get_selected_ids( [output], [1] ) diff --git a/tests/exploration/test_report_adaptive_lower.py b/tests/exploration/test_report_adaptive_lower.py index 78a1426d..876b3fde 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -42,6 +42,27 @@ def test_candidate_filter(self): report.restrict_candidate_ids([[0, 1]]) self.assertEqual(report.get_candidate_ids(), [[1]]) + def test_empty_cv_filter_preserves_trust_candidate_state(self): + model_devi = DeviManagerStd() + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.10, 0.20, 0.30]), + ) + report = ExplorationReportAdaptiveLower( + level_f_hi=1.0, + numb_candi_f=2, + rate_candi_f=0.0, + ) + report.record(model_devi) + + self.assertFalse(report.no_candidate()) + report.restrict_candidate_ids([[]]) + + self.assertFalse(report.no_candidate()) + self.assertEqual(report.candidate_ratio(), 0.0) + self.assertEqual(report.accurate_ratio(), 1.0 / 3.0) + self.assertEqual(report.failed_ratio(), 0.0) + def test_fv(self): model_devi = DeviManagerStd() model_devi.add( diff --git a/tests/exploration/test_report_trust_levels.py b/tests/exploration/test_report_trust_levels.py index fb2c6a0f..d059222c 100644 --- a/tests/exploration/test_report_trust_levels.py +++ b/tests/exploration/test_report_trust_levels.py @@ -225,6 +225,20 @@ def test_max_selection_after_candidate_filter(self): ter.restrict_candidate_ids([[0, 1]]) self.assertEqual(ter.get_candidate_ids(1), [[1]]) + def test_empty_cv_filter_preserves_trust_candidate_state(self): + model_devi = DeviManagerStd() + model_devi.add(DeviManager.MAX_DEVI_F, np.array([0.40, 0.20])) + report = ExplorationReportTrustLevelsMax(0.3, 0.6) + report.record(model_devi) + + self.assertFalse(report.no_candidate()) + report.restrict_candidate_ids([[]]) + + self.assertFalse(report.no_candidate()) + self.assertEqual(report.candidate_ratio(), 0.0) + self.assertEqual(report.accurate_ratio(), 0.5) + self.assertEqual(report.failed_ratio(), 0.0) + def test_random_selection_convergence(self): # case 1 model_devi = DeviManagerStd() diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index 6e6cfadf..cd28346a 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -867,6 +867,7 @@ def select( model_devis: List[Path], type_map: List[str] = None, optional_outputs: Optional[List[Path]] = None, + plm_outputs: Optional[List[Path]] = None, ) -> Tuple[List[Path], ExplorationReport]: confs = [] if len(trajs) == mocked_numb_lmp_tasks: diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 0edf7c1d..9e5843b9 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -8,6 +8,9 @@ import dpdata import numpy as np +from dargs.dargs import ( + ArgumentError, +) from dflow.python import ( OP, OPIO, @@ -69,32 +72,50 @@ def tearDown(self): shutil.rmtree(self.task_name) def test_plm_output_file_config(self): + self.assertEqual(RunLmp.normalize_config({})["plm_output_file"], "COLVAR") config = RunLmp.normalize_config({"plm_output_file": "COLVAR"}) self.assertEqual(config["plm_output_file"], "COLVAR") - with self.assertRaises(ValueError): + with self.assertRaises(ArgumentError): RunLmp.normalize_config({"plm_output_file": "outputs/COLVAR"}) for invalid_name in ["", ".", ".."]: - with self.subTest(invalid_name=invalid_name), self.assertRaises(ValueError): + with self.subTest(invalid_name=invalid_name), self.assertRaises( + ArgumentError + ): RunLmp.normalize_config({"plm_output_file": invalid_name}) + @patch("dpgen2.op.run_lmp.run_command") + def test_plm_output_file_invalid_config_is_fatal(self, mocked_run): + with self.assertRaisesRegex(FatalError, "invalid LAMMPS configuration"): + RunLmp().execute( + OPIO( + { + "config": {"plm_output_file": "outputs/COLVAR"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + mocked_run.assert_not_called() + @patch("dpgen2.op.run_lmp.run_command") def test_plm_output_file_collection(self, mocked_run): def run_with_plumed_output(*args, **kwargs): - Path("COLVAR").write_text("#! FIELDS time cv\n0.0 0.5\n") + Path("PLUMED_OUT").write_text("#! FIELDS time cv\n0.0 0.5\n") return 0, "", "" mocked_run.side_effect = run_with_plumed_output out = RunLmp().execute( OPIO( { - "config": {"plm_output_file": "COLVAR"}, + "config": {"plm_output_file": "PLUMED_OUT"}, "task_name": self.task_name, "task_path": self.task_path, "models": self.models, } ) ) - self.assertEqual(out["plm_output"], Path(self.task_name) / "COLVAR") + self.assertEqual(out["plm_output"], Path(self.task_name) / "PLUMED_OUT") @patch("dpgen2.op.run_lmp.run_command") def test_plm_output_file_rejects_staged_input(self, mocked_run): @@ -114,11 +135,16 @@ def test_plm_output_file_rejects_staged_input(self, mocked_run): @patch("dpgen2.op.run_lmp.run_command") def test_plm_output_file_does_not_reuse_stale_output(self, mocked_run): - mocked_run.return_value = (0, "", "") Path(self.task_name).mkdir() stale_output = Path(self.task_name) / "COLVAR" stale_output.write_text("stale output") + def run_without_stale_output(*args, **kwargs): + self.assertFalse(Path("COLVAR").exists()) + return 0, "", "" + + mocked_run.side_effect = run_without_stale_output + out = RunLmp().execute( OPIO( { @@ -133,6 +159,24 @@ def test_plm_output_file_does_not_reuse_stale_output(self, mocked_run): self.assertIsNone(out["plm_output"]) self.assertFalse(stale_output.exists()) + @patch("dpgen2.op.run_lmp.run_command") + def test_plm_output_file_rejects_generated_name(self, mocked_run): + for output_name in ["output.plumed", model_name_pattern % 0, lmp_log_name]: + with self.subTest(output_name=output_name), self.assertRaisesRegex( + FatalError, "collides with a generated" + ): + RunLmp().execute( + OPIO( + { + "config": {"plm_output_file": output_name}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + mocked_run.assert_not_called() + @patch("dpgen2.op.run_lmp.run_command") def test_success(self, mocked_run): mocked_run.side_effect = [(0, "foo\n", "")] diff --git a/tests/utils/test_dl_dpgen2_arti.py b/tests/utils/test_dl_dpgen2_arti.py index 037a0d80..6b193369 100644 --- a/tests/utils/test_dl_dpgen2_arti.py +++ b/tests/utils/test_dl_dpgen2_arti.py @@ -123,6 +123,11 @@ def test_lmp_download(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-plm_output", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), mock.call( "arti-extra_outputs", path=Path("iter-000001/prep-run-explore/outputs"), @@ -258,6 +263,11 @@ def test_update_finished_steps_exist_steps(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-plm_output", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), mock.call( "arti-extra_outputs", path=Path("iter-000001/prep-run-explore/outputs"), @@ -325,6 +335,11 @@ def test_update_finished_steps_none_steps(self, mocked_dl): path=Path("iter-000001/prep-run-explore/outputs"), skip_exists=True, ), + mock.call( + "arti-plm_output", + path=Path("iter-000001/prep-run-explore/outputs"), + skip_exists=True, + ), mock.call( "arti-extra_outputs", path=Path("iter-000001/prep-run-explore/outputs"),