import dpdata
import numpy as np
from collections import Counter
from typing import (
List,
Optional,
Tuple,
)
from pathlib import Path
from . import (
ConfSelector,
ConfFilters,
)
from dpgen2.exploration.report import ExplorationReport
from dpgen2.exploration.selector import TrustLevel
from dpgen2.exploration.render import TrajRender
[docs]class ConfSelectorFrames(ConfSelector):
"""Select frames from trajectories as confs.
Parameters:
trust_level: TrustLevel
The trust level
conf_filter: ConfFilters
The configuration filter
"""
def __init__(
self,
traj_render: TrajRender,
report: ExplorationReport,
max_numb_sel : Optional[int] = None,
conf_filters : Optional[ConfFilters] = None,
):
self.max_numb_sel = max_numb_sel
self.conf_filters = conf_filters
self.traj_render = traj_render
self.report = report
[docs] def select (
self,
trajs : List[Path],
model_devis : List[Path],
type_map : Optional[List[str]] = None,
) -> Tuple[List[ Path ], ExplorationReport]:
"""Select configurations
Parameters
----------
trajs : List[Path]
A `list` of `Path` to trajectory files generated by LAMMPS
model_devis : List[Path]
A `list` of `Path` to model deviation files generated by LAMMPS.
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
type_map : List[str]
The `type_map` of the systems
Returns
-------
confs : List[Path]
The selected confgurations, stored in a folder in deepmd/npy format, can be parsed as dpdata.MultiSystems. The `list` only has one item.
report : ExplorationReport
The exploration report recoding the status of the exploration.
"""
ntraj = len(trajs)
assert(ntraj == len(model_devis))
mdf, mdv = self.traj_render.get_model_devi(model_devis)
self.report.clear()
self.report.record(mdf, mdv)
id_cand_list = self.report.get_candidate_ids(self.max_numb_sel)
ms = self.traj_render.get_confs(trajs, id_cand_list, type_map, self.conf_filters)
out_path = Path('confs')
out_path.mkdir(exist_ok=True)
ms.to_deepmd_npy(out_path)
return [out_path], self.report