diff --git a/docs/input.md b/docs/input.md index 6fc03bd1..0a008019 100644 --- a/docs/input.md +++ b/docs/input.md @@ -131,6 +131,190 @@ 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 after the model-deviation trust window and before final CV-space +coverage or an explicitly configured selection policy: + +```json +"explore": { + "config": { + "plm_output_file": "COLVAR" + }, + "cv_filter": { + "regions": [ + {"d": [0.08, 0.12]}, + {"v": [1.8, 2.2]} + ], + "sampling": {"mode": "report"}, + "time_alignment": {"start": 0.0, "step": 0.01} + } +} +``` + +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 +``` + +`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` +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]} + } + ], + "time_alignment": {"start": 0.0, "step": 0.01} +} +``` + +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]} + } + ], + "time_alignment": {"start": 0.0, "step": 0.01} +} +``` + +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 +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. + +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": { + "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 + }, + "time_alignment": {"start": 0.0, "step": 0.01} +} +``` + +`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 +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-6 + } +} +``` + +`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` 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 +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 df11ff7f..e96ed142 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,15 @@ 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. A region may be a field-to-" + "[lower, upper] mapping or a named conditions mapping; regions are " + "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. Explicit " + "time alignment is required to bind COLVAR rows to trajectory frames." + ) return [ Argument( @@ -259,6 +269,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..3c21b45a 100644 --- a/dpgen2/exploration/report/report.py +++ b/dpgen2/exploration/report/report.py @@ -65,10 +65,18 @@ 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, max_nframes: Optional[int] = None, + clear: bool = True, ) -> List[List[int]]: r"""Get indexes of candidate configurations @@ -76,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/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index cd2989f3..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) @@ -371,6 +374,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._update_ratios() + 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..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( @@ -236,6 +240,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._update_ratios() + @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..17dd65b0 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -10,14 +10,21 @@ Optional, Tuple, Union, + cast, ) 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 +36,9 @@ ConfFilters, ConfSelector, ) +from .plumed_cv_filter import ( + PlumedCVFilter, +) class ConfSelectorFrames(ConfSelector): @@ -48,11 +58,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 +72,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 @@ -72,6 +85,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] @@ -92,7 +107,65 @@ def select( self.report.clear() self.report.record(md_model_devi) - id_cand_list = self.report.get_candidate_ids(self.max_numb_sel) + id_cand_list = None + 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 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 = 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: + ( + sampled_ids, + records, + summary, + ) = plumed_cv_filter.select_candidate_ids_with_audit( + plm_files, + [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 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, + md_f, + loaded_outputs=loaded_plm_outputs, + ) ms = self.traj_render.get_confs( trajs, @@ -105,5 +178,8 @@ 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: + 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 new file mode 100644 index 00000000..9949f330 --- /dev/null +++ b/dpgen2/exploration/selector/plumed_cv_filter.py @@ -0,0 +1,937 @@ +import csv +import json +from pathlib import ( + Path, +) +from typing import ( + Dict, + List, + Optional, + Tuple, +) + +import numpy as np +from dargs import ( + Argument, +) +from dflow.python import ( + FatalError, +) + +PlumedOutputs = List[Tuple[List[str], np.ndarray]] + + +class PlumedCVFilter: + """Select frames in a union of PLUMED CV regions. + + 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 + def args() -> List[Argument]: + return [ + Argument( + "regions", + list, + optional=False, + doc=( + "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( + "sampling", + dict, + [ + 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, + doc=( + "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, + 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=False, + doc="Required expected PLUMED time = start + frame * step.", + ), + ] + + def __init__( + self, + regions: List[Dict], + sampling: Optional[Dict] = None, + time_alignment: Optional[Dict] = None, + ): + 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 = [] + 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 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: + 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) + 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: + 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", "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") + 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) + if not isinstance(field, str) or not field: + raise ValueError("uniform PLUMED CV sampling requires a field") + 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") + 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]) -> Dict: + if time_alignment is 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"} + 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-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: + raise ValueError("invalid PLUMED CV time_alignment") + return {"start": start, "step": step, "atol": atol} + + 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) + 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], + 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.""" + 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): + 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 = self._eligible_candidates(candidate_ids, nframes, masks_by_traj) + limit = ( + len(candidates) + if max_nframes is None + else min(max_nframes, len(candidates)) + ) + rng = np.random.default_rng(self.sampling["seed"]) + rejected_by_gap = set() + picked = [] + if self.sampling["mode"] == "random": + self._pick_candidates( + candidates, + limit, + "random", + max_devi_f, + rng, + picked, + rejected_by_gap, + ) + else: + picked, rejected_by_gap = self._sample_grid( + candidates, + outputs, + masks_by_traj, + limit, + max_devi_f, + rng, + ) + 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, + loaded_outputs: Optional[PlumedOutputs] = None, + ): + """Audit filtering followed by the report's existing sampling policy.""" + 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 + ] + 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 _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, + masks_by_traj, + limit: int, + max_devi_f: Optional[List[np.ndarray]], + rng, + ): + 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") + + buckets = [dict() for _ in self.regions] + for candidate in candidates: + traj_idx, frame_idx = candidate + fields, values = outputs[traj_idx] + for region_idx, mask in enumerate(masks_by_traj[traj_idx]): + if not mask[frame_idx]: + continue + 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._weighted_quotas(capacities, self.region_weights, limit) + picked = [] + picked_set = set() + rejected_by_gap = set() + 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] + for cell, items in region.items() + } + 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_set.update(picked[before:]) + + if len(picked) < limit: + remaining = [item for item in candidates if item not in picked_set] + 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, + ): + sampling = self.sampling + assert sampling is not None + 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], + ), + ) + min_frame_gap = 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]: + 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 _cell_quotas(cls, buckets, total, grid_sizes): + cells = sorted(buckets) + if not cells or total <= 0: + return {} + if total <= len(cells): + return {cell: 1 for cell in cls._spread_cells(cells, total, grid_sizes)} + extra = cls._balanced_quotas( + [len(buckets[cell]) - 1 for cell in cells], total - len(cells) + ) + return {cell: 1 + increment for cell, increment in zip(cells, extra)} + + @staticmethod + 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) + ), + ) + ] + + 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 = 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): + 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 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): + grouped = [[] for _ in range(ntraj)] + for traj_idx, frame_idx in sorted(candidates): + 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]) -> PlumedOutputs: + if len(files) != len(nframes): + raise FatalError("PLUMED outputs and trajectories have different lengths") + outputs = [] + 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" + ) + 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 + + 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): + 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..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,16 +121,25 @@ 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"] + plm_output_file = config["plm_output_file"] task_name = ip["task_name"] task_path = ip["task_path"] models = ip["models"] # 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) @@ -140,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 @@ -206,8 +244,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 +269,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 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 [ Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd), @@ -262,6 +304,16 @@ def lmp_args(): default=False, doc=doc_use_hdf5, ), + Argument( + "plm_output_file", + str, + optional=True, + 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( "extra_output_files", list, diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index e8ba891d..e43f4012 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,6 +85,8 @@ 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 ) @@ -93,6 +96,7 @@ def execute( model_devis, type_map=type_map, optional_outputs=optional_outputs, + plm_outputs=plm_outputs, ) 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/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 52b30248..29cbfc18 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 @@ -5,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 ( @@ -17,10 +25,12 @@ TrajRenderLammps, ) from dpgen2.exploration.report import ( + ExplorationReportTrustLevelsMax, ExplorationReportTrustLevelsRandom, ) from dpgen2.exploration.selector import ( ConfSelectorFrames, + PlumedCVFilter, ) # isort: on @@ -85,13 +95,105 @@ 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") + 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=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: + 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", + }, + time_alignment={"start": 0.0, "step": 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(), 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) 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..87211910 --- /dev/null +++ b/tests/exploration/test_plumed_cv_filter.py @@ -0,0 +1,470 @@ +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np +from dargs import ( + Argument, +) +from dargs.dargs import ( + ArgumentKeyError, +) +from dflow.python import ( + FatalError, +) + +from dpgen2.exploration.selector import ( + PlumedCVFilter, +) + + +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()) + config = schema.normalize_value( + { + "regions": [{"distance": [0.8, 1.2]}], + "sampling": { + "mode": "uniform", + "field": "distance", + "n_bins": 8, + }, + "time_alignment": {"start": 0.0, "step": 1.0}, + } + ) + 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) + + 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-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: + 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 = 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]]) + + 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 = 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 = 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): + 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 = make_filter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + repeated = make_filter(**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_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 = 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") + 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 = 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( + [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" + 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 = make_filter( + 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 = make_filter(**kwargs).select_candidate_ids( + [output], [20], [list(range(20))], 5 + ) + repeated = make_filter(**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 = make_filter( + 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 = make_filter( + 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 = 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}, + ) + self.assertEqual( + cv_filter.select_candidate_ids([output], [3], [list(range(3))], 2), + [[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 = make_filter( + 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 = make_filter( + 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 = make_filter( + 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 = 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): + make_filter(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): + make_filter(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": "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): + make_filter(regions=[{"cv": [0.0, 1.0]}], sampling=sampling) + + with self.assertRaises(ValueError): + make_filter(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 + ): + make_filter( + 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", + "#! 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): + 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): + 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 b5f123ca..876b3fde 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -27,6 +27,42 @@ 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_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 7f6c79a9..d059222c 100644 --- a/tests/exploration/test_report_trust_levels.py +++ b/tests/exploration/test_report_trust_levels.py @@ -217,6 +217,28 @@ 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_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 650fd82e..9e5843b9 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -8,10 +8,14 @@ import dpdata import numpy as np +from dargs.dargs import ( + ArgumentError, +) from dflow.python import ( OP, OPIO, Artifact, + FatalError, OPIOSign, TransientError, ) @@ -67,6 +71,112 @@ def tearDown(self): if Path(self.task_name).is_dir(): 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(ArgumentError): + RunLmp.normalize_config({"plm_output_file": "outputs/COLVAR"}) + for invalid_name in ["", ".", ".."]: + 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("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": "PLUMED_OUT"}, + "task_name": self.task_name, + "task_path": self.task_path, + "models": self.models, + } + ) + ) + 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): + (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): + 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( + { + "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_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/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"] + ) 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"),