diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..7c9978dc 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -212,6 +212,7 @@ def lmp_args(): "Each task group is described in :ref:`the task group definition` " ) doc_filters = "A list of configuration filters" + doc_lammps_input_file = "The template input file for LAMMPS simulation. Will pass it to dpdata to parse LAMMPS dump file in spin job." return [ Argument( @@ -259,6 +260,13 @@ def lmp_args(): default=[], doc=doc_filters, ), + Argument( + "lammps_input_file", + str, + optional=True, + default=None, + doc=doc_lammps_input_file, + ), ] diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 203478dc..4d8a1fad 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -81,6 +81,7 @@ CustomizedLmpTemplateTaskGroup, ExplorationStage, ExplorationTask, + LmpSpinTaskGroup, LmpTemplateTaskGroup, NPTTaskGroup, caly_normalize, @@ -318,7 +319,9 @@ def make_naive_exploration_scheduler_without_conf(config, explore_style): conv_style = convergence.pop("type") report = conv_styles[conv_style](**convergence) # trajectory render, the format of the output trajs are assumed to be lammps/dump - render = TrajRenderLammps(nopbc=output_nopbc) + render = TrajRenderLammps( + nopbc=output_nopbc, lammps_input_file=config["explore"]["lammps_input_file"] + ) # selector selector = ConfSelectorFrames( render, @@ -378,7 +381,11 @@ def make_lmp_naive_exploration_scheduler(config): # report conv_style = convergence.pop("type") report = conv_styles[conv_style](**convergence) - render = TrajRenderLammps(nopbc=output_nopbc, use_ele_temp=use_ele_temp) + render = TrajRenderLammps( + nopbc=output_nopbc, + use_ele_temp=use_ele_temp, + lammps_input_file=config["explore"]["lammps_input_file"], + ) # selector selector = ConfSelectorFrames( render, diff --git a/dpgen2/exploration/deviation/deviation_manager.py b/dpgen2/exploration/deviation/deviation_manager.py index cbc227c5..66ef1eea 100644 --- a/dpgen2/exploration/deviation/deviation_manager.py +++ b/dpgen2/exploration/deviation/deviation_manager.py @@ -19,6 +19,9 @@ class DeviManager(ABC): MAX_DEVI_F = "max_devi_f" MIN_DEVI_F = "min_devi_f" AVG_DEVI_F = "avg_devi_f" + MAX_DEVI_MF = "max_devi_mf" + MIN_DEVI_MF = "min_devi_mf" + AVG_DEVI_MF = "avg_devi_mf" def __init__(self) -> None: super().__init__() @@ -32,6 +35,9 @@ def _check_name(self, name: str): DeviManager.MAX_DEVI_F, DeviManager.MIN_DEVI_F, DeviManager.AVG_DEVI_F, + DeviManager.MAX_DEVI_MF, + DeviManager.MIN_DEVI_MF, + DeviManager.AVG_DEVI_MF, ), f"Error: unknown deviation name {name}" def add(self, name: str, deviation: np.ndarray) -> None: @@ -64,7 +70,9 @@ def get(self, name: str) -> List[Optional[np.ndarray]]: The name of the deviation. The name is restricted to (DeviManager.MAX_DEVI_V, DeviManager.MIN_DEVI_V, DeviManager.AVG_DEVI_V, DeviManager.MAX_DEVI_F, - DeviManager.MIN_DEVI_F, DeviManager.AVG_DEVI_F) + DeviManager.MIN_DEVI_F, DeviManager.AVG_DEVI_F, + DeviManager.MAX_DEVI_MF, DeviManager.MIN_DEVI_MF, + DeviManager.AVG_DEVI_MF) """ self._check_name(name) self._check_data() diff --git a/dpgen2/exploration/deviation/deviation_std.py b/dpgen2/exploration/deviation/deviation_std.py index f927b75b..0c41d00f 100644 --- a/dpgen2/exploration/deviation/deviation_std.py +++ b/dpgen2/exploration/deviation/deviation_std.py @@ -65,6 +65,9 @@ def _check_data(self) -> None: DeviManager.MAX_DEVI_F, DeviManager.MIN_DEVI_F, DeviManager.AVG_DEVI_F, + DeviManager.MAX_DEVI_MF, + DeviManager.MIN_DEVI_MF, + DeviManager.AVG_DEVI_MF, ) # check the length of model deviations frames = {} diff --git a/dpgen2/exploration/render/traj_render_lammps.py b/dpgen2/exploration/render/traj_render_lammps.py index 00b6a3de..477c7076 100644 --- a/dpgen2/exploration/render/traj_render_lammps.py +++ b/dpgen2/exploration/render/traj_render_lammps.py @@ -42,9 +42,14 @@ def __init__( self, nopbc: bool = False, use_ele_temp: int = 0, + lammps_input_file: str = None, # type: ignore ): self.nopbc = nopbc self.use_ele_temp = use_ele_temp + if lammps_input_file is not None: + self.lammps_input = Path(lammps_input_file).read_text() + else: + self.lammps_input = None def get_model_devi( self, @@ -74,6 +79,11 @@ def _load_one_model_devi(self, fname, model_devi): model_devi.add(DeviManager.MAX_DEVI_F, dd[:, 4]) # type: ignore model_devi.add(DeviManager.MIN_DEVI_F, dd[:, 5]) # type: ignore model_devi.add(DeviManager.AVG_DEVI_F, dd[:, 6]) # type: ignore + # assume the 7-9 columns are for MF + if dd.shape[1] >= 10: # type: ignore + model_devi.add(DeviManager.MAX_DEVI_MF, dd[:, 7]) # type: ignore + model_devi.add(DeviManager.MIN_DEVI_MF, dd[:, 8]) # type: ignore + model_devi.add(DeviManager.AVG_DEVI_MF, dd[:, 9]) # type: ignore def get_ele_temp(self, optional_outputs): ele_temp = [] @@ -117,13 +127,21 @@ def get_confs( traj_fmt = "lammps/dump" ms = dpdata.MultiSystems(type_map=type_map) + if self.lammps_input is not None: + lammps_input_file = "lammps_input.in" + Path(lammps_input_file).write_text(self.lammps_input) + else: + lammps_input_file = None for ii in range(ntraj): if len(id_selected[ii]) > 0: if isinstance(trajs[ii], HDF5Dataset): traj = StringIO(trajs[ii].get_data()) # type: ignore else: traj = trajs[ii] - ss = dpdata.System(traj, fmt=traj_fmt, type_map=type_map) + # for spin job, need to read input file to get the key of the spin data + ss = dpdata.System( + traj, fmt=traj_fmt, type_map=type_map, input_file=lammps_input_file + ) ss.nopbc = self.nopbc if ele_temp: self.set_ele_temp(ss, ele_temp[ii]) diff --git a/dpgen2/exploration/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index cd6087f2..e5bc3bab 100644 --- a/dpgen2/exploration/report/report_adaptive_lower.py +++ b/dpgen2/exploration/report/report_adaptive_lower.py @@ -27,14 +27,17 @@ class ExplorationReportAdaptiveLower(ExplorationReport): This report will treat a fixed number of frames that has force model deviation lower than `level_f_hi`, and virial model deviation - lower than `level_v_hi` as candidates. + lower than `level_v_hi` as candidates, and magnetic force model deviation + lower than `level_mf_hi` as candidates. The number of force frames is given by max(`numb_candi_f`, `rate_candi_f` * nframes) The number of virial frames is given by max(`numb_candi_v`, `rate_candi_v` * nframes) + The number of magnetic force frames is given by max(`numb_candi_mf`, `rate_candi_mf` * nframes) The lower force trust level will be set to the lowest force model deviation of the force frames. The lower virial trust level will be set to the lowest - virial model deviation of the virial frames + virial model deviation of the virial frames. The lower magnetic force trust + level will be set to the lowest magnetic force model deviation of the magnetic The exploration will be treat as converged if the differences in model deviations in the neighboring steps are less than `conv_tolerance` @@ -58,6 +61,14 @@ class ExplorationReportAdaptiveLower(ExplorationReport): rate_candi_v float The ratio of virial frames that has a model deviation lower than `level_v_hi` treated as candidate. + level_mf_hi float + The higher trust level of magnetic force model deviation + numb_candi_mf int + The number of magnetic force frames that has a model deviation lower than + `level_mf_hi` treated as candidate. + rate_candi_mf float + The ratio of magnetic force frames that has a model deviation lower than + `level_mf_hi` treated as candidate. n_checked_steps int The number of steps to check the convergence. conv_tolerance float @@ -79,21 +90,32 @@ def __init__( level_v_hi: Optional[float] = None, numb_candi_v: int = 0, rate_candi_v: float = 0.0, + level_mf_hi: Optional[float] = None, + numb_candi_mf: int = 0, + rate_candi_mf: float = 0.0, n_checked_steps: int = 2, conv_tolerance: float = 0.05, candi_sel_prob: str = "uniform", ): self.level_f_hi = level_f_hi self.level_v_hi = level_v_hi + self.level_mf_hi = level_mf_hi self.numb_candi_f = numb_candi_f self.rate_candi_f = rate_candi_f self.numb_candi_v = numb_candi_v self.rate_candi_v = rate_candi_v + self.numb_candi_mf = numb_candi_mf + self.rate_candi_mf = rate_candi_mf self.has_virial = self.level_v_hi is not None if not self.has_virial: self.level_v_hi = sys.float_info.max self.numb_candi_v = 0 self.rate_candi_v = 0.0 + self.has_mf = self.level_mf_hi is not None + if not self.has_mf: + self.level_mf_hi = sys.float_info.max + self.numb_candi_mf = 0 + self.rate_candi_mf = 0.0 self.n_checked_steps = n_checked_steps self.conv_tolerance = conv_tolerance self.model_devi = None @@ -123,6 +145,12 @@ def __init__( "v_hi", ) spaces += [10, 10] + if self.has_mf: + print_tuple += ( + "mf_lo", + "mf_hi", + ) + spaces += [10, 10] spaces += [8] self.fmt_str = " ".join([f"%{ii}s" for ii in spaces]) self.fmt_flt = "%.4f" @@ -145,13 +173,20 @@ def make_class_doc_link(key): rate_candi_f_link = make_class_doc_link("rate_candi_f") numb_candi_v_link = make_class_doc_link("numb_candi_v") rate_candi_v_link = make_class_doc_link("rate_candi_v") - numb_candi_s = f"{numb_candi_f_link} or {numb_candi_v_link}" - rate_candi_s = f"{rate_candi_f_link} or {rate_candi_v_link}" + numb_candi_mf_link = make_class_doc_link("numb_candi_mf") + rate_candi_mf_link = make_class_doc_link("rate_candi_mf") + numb_candi_s = ( + f"{numb_candi_f_link} or {numb_candi_v_link} or {numb_candi_mf_link}" + ) + rate_candi_s = ( + f"{rate_candi_f_link} or {rate_candi_v_link} or {rate_candi_mf_link}" + ) level_f_hi_link = make_class_doc_link("level_f_hi") level_v_hi_link = make_class_doc_link("level_v_hi") + level_mf_hi_link = make_class_doc_link("level_mf_hi") conv_tolerance_link = make_class_doc_link("conv_tolerance") n_checked_steps_link = make_class_doc_link("n_checked_steps") - return f"The method of adaptive adjust the lower trust levels. In each step of iterations, a number (set by {numb_candi_s}) or a ratio (set by {rate_candi_s}) of configurations with a model deviation lower than the higher trust level ({level_f_hi_link}, {level_v_hi_link}) are treated as candidates. The lowest model deviation of the candidates are treated as the lower trust level. If the lower trust level does not change significant (controlled by {conv_tolerance_link}) in {n_checked_steps_link}, the stage is treated as converged. " + return f"The method of adaptive adjust the lower trust levels. In each step of iterations, a number (set by {numb_candi_s}) or a ratio (set by {rate_candi_s}) of configurations with a model deviation lower than the higher trust level ({level_f_hi_link}, {level_v_hi_link}, {level_mf_hi_link}) are treated as candidates. The lowest model deviation of the candidates are treated as the lower trust level. If the lower trust level does not change significant (controlled by {conv_tolerance_link}) in {n_checked_steps_link}, the stage is treated as converged. " @staticmethod def args() -> List[Argument]: @@ -161,6 +196,9 @@ def args() -> List[Argument]: doc_level_v_hi = "The higher trust level of virial model deviation" doc_numb_candi_v = "The number of virial frames that has a model deviation lower than `level_v_hi` treated as candidate." doc_rate_candi_v = "The ratio of virial frames that has a model deviation lower than `level_v_hi` treated as candidate." + doc_level_mf_hi = "The higher trust level of magnetic force model deviation" + doc_numb_candi_mf = "The number of magnetic force frames that has a model deviation lower than `level_mf_hi` treated as candidate." + doc_rate_candi_mf = "The ratio of magnetic force frames that has a model deviation lower than `level_mf_hi` treated as candidate." doc_n_check_steps = "The number of steps to check the convergence." doc_conv_tolerance = "The convergence tolerance." doc_candi_sel_prob = ( @@ -191,6 +229,19 @@ def args() -> List[Argument]: Argument( "rate_candi_v", float, optional=True, default=0.0, doc=doc_rate_candi_v ), + Argument( + "level_mf_hi", float, optional=True, default=None, doc=doc_level_mf_hi + ), + Argument( + "numb_candi_mf", int, optional=True, default=0, doc=doc_numb_candi_mf + ), + Argument( + "rate_candi_mf", + float, + optional=True, + default=0.0, + doc=doc_rate_candi_mf, + ), Argument( "n_checked_steps", int, optional=True, default=2, doc=doc_n_check_steps ), @@ -222,6 +273,7 @@ def clear( self.model_devi = None self.md_f = [] self.md_v = [] + self.md_mf = [] def record( self, @@ -231,34 +283,48 @@ def record( self.ntraj += ntraj md_f = model_devi.get(DeviManager.MAX_DEVI_F) md_v = model_devi.get(DeviManager.MAX_DEVI_V) + md_mf = model_devi.get(DeviManager.MAX_DEVI_MF) self.md_f += md_f self.md_v += md_v + self.md_mf += md_mf # inits coll_f = [] coll_v = [] + coll_mf = [] # loop over trajs for ii in range(ntraj): - add_nframes, add_accur, add_failed, add_f, add_v = self._record_one_traj( - ii, md_f[ii], md_v[ii] - ) + ( + add_nframes, + add_accur, + add_failed, + add_f, + add_v, + add_mf, + ) = self._record_one_traj(ii, md_f[ii], md_v[ii], md_mf[ii]) self.nframes += add_nframes self.accur.update(add_accur) self.failed += add_failed coll_f += add_f coll_v += add_v + coll_mf += add_mf # sort coll_f.sort() coll_v.sort() + coll_mf.sort() assert len(coll_v) == len(coll_f) + assert len(coll_mf) == len(coll_f) # calcuate numbers numb_candi_f = max(self.numb_candi_f, int(self.rate_candi_f * len(coll_f))) numb_candi_v = max(self.numb_candi_v, int(self.rate_candi_v * len(coll_v))) + numb_candi_mf = max(self.numb_candi_mf, int(self.rate_candi_mf * len(coll_mf))) # adjust number of candidate if len(coll_f) < numb_candi_f: numb_candi_f = len(coll_f) if len(coll_v) < numb_candi_v: numb_candi_v = len(coll_v) + if len(coll_mf) < numb_candi_mf: + numb_candi_mf = len(coll_mf) # compute trust lo if numb_candi_v == 0: self.level_v_lo = self.level_v_hi @@ -266,6 +332,14 @@ def record( self.level_v_lo = coll_v[-numb_candi_v][0] if not self.has_virial: self.level_v_lo = None + + if numb_candi_mf == 0: + self.level_mf_lo = self.level_mf_hi + else: + self.level_mf_lo = coll_mf[-numb_candi_mf][0] + if not self.has_mf: + self.level_mf_lo = None + if numb_candi_f == 0: self.level_f_lo = self.level_f_hi else: @@ -275,6 +349,8 @@ def record( self.candi.add(tuple(coll_f[ii][1:])) for ii in range(len(coll_v) - numb_candi_v, len(coll_v)): self.candi.add(tuple(coll_v[ii][1:])) + for ii in range(len(coll_mf) - numb_candi_mf, len(coll_mf)): + self.candi.add(tuple(coll_mf[ii][1:])) # accurate set is substracted by the candidate set self.accur = self.accur - self.candi self.model_devi = model_devi @@ -288,12 +364,14 @@ def _record_one_traj( tt, md_f, md_v, + md_mf, ): """ Record one trajctory. tt: traj index md_f, md_v: model deviations of force and virial + md_mf: model deviations of magnetic force """ # check consistency if self.has_virial and md_v is None: @@ -301,9 +379,16 @@ def _record_one_traj( "report requires virial model deviation, but no virial " "model deviation is provided." ) + if self.has_mf and md_mf is None: + raise FatalError( + "report requires magnetic force model deviation, but no " + "magnetic force model deviation is provided." + ) # fake md_v as zeros if None is provided if md_v is None: md_v = np.zeros_like(md_f) + if md_mf is None: + md_mf = np.zeros_like(md_f) # loop over frames nframes = md_f.shape[0] assert nframes == md_v.shape[0] @@ -311,16 +396,22 @@ def _record_one_traj( accur = set() coll_f = [] coll_v = [] + coll_mf = [] for ii in range(nframes): - if md_f[ii] > self.level_f_hi or md_v[ii] > self.level_v_hi: + if ( + md_f[ii] > self.level_f_hi + or md_v[ii] > self.level_v_hi + or md_mf[ii] > self.level_mf_hi + ): failed.append((tt, ii)) else: coll_f.append([md_f[ii], tt, ii]) coll_v.append([md_v[ii], tt, ii]) + coll_mf.append([md_mf[ii], tt, ii]) # now accur takes all non-failed frames, # will be substracted by candidate later accur.add((tt, ii)) - return nframes, accur, failed, coll_f, coll_v + return nframes, accur, failed, coll_f, coll_v, coll_mf def _sequence_conv( self, @@ -348,6 +439,10 @@ def converged( all_level_v = [ii.level_v_lo for ii in reports] + [self.level_v_lo] all_level_v = all_level_v[-self.n_checked_steps :] conv = conv and self._sequence_conv(all_level_v) + if self.has_mf: + all_level_mf = [ii.level_mf_lo for ii in reports] + [self.level_mf_lo] + all_level_mf = all_level_mf[-self.n_checked_steps :] + conv = conv and self._sequence_conv(all_level_mf) return conv def failed_ratio( @@ -529,5 +624,10 @@ def print( fmt_flt % (self.level_v_lo), fmt_flt % (self.level_v_hi), ) + if self.has_mf: + print_tuple += ( + fmt_flt % (self.level_mf_lo), + fmt_flt % (self.level_mf_hi), + ) ret = " " + fmt_str % print_tuple return ret diff --git a/dpgen2/exploration/report/report_trust_levels_base.py b/dpgen2/exploration/report/report_trust_levels_base.py index 185ea5b1..b0e57701 100644 --- a/dpgen2/exploration/report/report_trust_levels_base.py +++ b/dpgen2/exploration/report/report_trust_levels_base.py @@ -31,15 +31,22 @@ def __init__( level_f_hi, level_v_lo=None, level_v_hi=None, + level_mf_lo=None, + level_mf_hi=None, conv_accuracy=0.9, ): self.level_f_lo = level_f_lo self.level_f_hi = level_f_hi self.level_v_lo = level_v_lo self.level_v_hi = level_v_hi + self.level_mf_lo = level_mf_lo + self.level_mf_hi = level_mf_hi self.conv_accuracy = conv_accuracy self.clear() self.v_level = (self.level_v_lo is not None) and (self.level_v_hi is not None) + self.mf_level = (self.level_mf_lo is not None) and ( + self.level_mf_hi is not None + ) self.model_devi = None print_tuple = ( @@ -59,6 +66,12 @@ def __init__( "v_hi", ) spaces += [10, 10] + if self.mf_level: + print_tuple += ( + "mf_lo", + "mf_hi", + ) + spaces += [10, 10] print_tuple += ("cvged",) spaces += [8] self.fmt_str = " ".join([f"%{ii}s" for ii in spaces]) @@ -75,6 +88,8 @@ def args() -> List[Argument]: doc_level_f_hi = "The higher trust level of force model deviation" doc_level_v_lo = "The lower trust level of virial model deviation" doc_level_v_hi = "The higher trust level of virial model deviation" + doc_level_mf_lo = "The lower trust level of magnetic force model deviation" + doc_level_mf_hi = "The higher trust level of magnetic force model deviation" doc_conv_accuracy = "If the ratio of accurate frames is larger than this value, the stage is converged" return [ Argument("level_f_lo", float, optional=False, doc=doc_level_f_lo), @@ -85,6 +100,20 @@ def args() -> List[Argument]: Argument( "level_v_hi", float, optional=True, default=None, doc=doc_level_v_hi ), + Argument( + "level_mf_lo", + float, + optional=True, + default=None, + doc=doc_level_mf_lo, + ), + Argument( + "level_mf_hi", + float, + optional=True, + default=None, + doc=doc_level_mf_hi, + ), Argument( "conv_accuracy", float, @@ -111,6 +140,7 @@ def record( ntraj = model_devi.ntraj md_f = model_devi.get(DeviManager.MAX_DEVI_F) md_v = model_devi.get(DeviManager.MAX_DEVI_V) + md_mf = model_devi.get(DeviManager.MAX_DEVI_MF) for ii in range(ntraj): id_f_cand, id_f_accu, id_f_fail = self._get_indexes( @@ -119,6 +149,9 @@ def record( id_v_cand, id_v_accu, id_v_fail = self._get_indexes( md_v[ii], self.level_v_lo, self.level_v_hi ) + id_mf_cand, id_mf_accu, id_mf_fail = self._get_indexes( + md_mf[ii], self.level_mf_lo, self.level_mf_hi + ) nframes, set_accu, set_cand, set_fail = self._record_one_traj( id_f_accu, id_f_cand, @@ -126,6 +159,9 @@ def record( id_v_accu, id_v_cand, id_v_fail, + id_mf_accu, + id_mf_cand, + id_mf_fail, ) # record self.traj_nframes.append(nframes) @@ -170,6 +206,9 @@ def _record_one_traj( id_v_accu, id_v_cand, id_v_fail, + id_mf_accu, + id_mf_cand, + id_mf_fail, ): """ Record one trajctory. inputs are the indexes of candidate, accurate and failed frames. @@ -180,27 +219,47 @@ def _record_one_traj( if novirial: assert id_v_accu is None assert id_v_fail is None + nomagforce = id_mf_cand is None + if nomagforce: + assert id_mf_accu is None + assert id_mf_fail is None nframes = np.size(np.concatenate((id_f_cand, id_f_accu, id_f_fail))) if (not novirial) and nframes != np.size( np.concatenate((id_v_cand, id_v_accu, id_v_fail)) ): raise FatalError("number of frames by virial ") + if (not nomagforce) and nframes != np.size( + np.concatenate((id_mf_cand, id_mf_accu, id_mf_fail)) + ): + raise FatalError("number of frames by magnetic force ") # nframes # to sets + set_full = set([ii for ii in range(nframes)]) set_f_accu = set(id_f_accu) set_f_cand = set(id_f_cand) set_f_fail = set(id_f_fail) - set_v_accu = set([ii for ii in range(nframes)]) if novirial else set(id_v_accu) + set_v_accu = set_full if novirial else set(id_v_accu) set_v_cand = set([]) if novirial else set(id_v_cand) set_v_fail = set([]) if novirial else set(id_v_fail) + set_mf_accu = set_full if nomagforce else set(id_mf_accu) + set_mf_cand = set([]) if nomagforce else set(id_mf_cand) + set_mf_fail = set([]) if nomagforce else set(id_mf_fail) + + # check consistency + assert set_full == set_f_accu | set_f_cand | set_f_fail + for accu, cand, fail in [ + [set_f_accu, set_f_cand, set_f_fail], + [set_v_accu, set_v_cand, set_v_fail], + [set_mf_accu, set_mf_cand, set_mf_fail], + ]: + assert 0 == len(accu & cand) + assert 0 == len(accu & fail) + assert 0 == len(cand & fail) + # accu, cand, fail - set_accu = set_f_accu & set_v_accu - set_cand = ( - (set_f_cand & set_v_accu) - | (set_f_cand & set_v_cand) - | (set_f_accu & set_v_cand) - ) - set_fail = set_f_fail | set_v_fail + set_accu = set_f_accu & set_v_accu & set_mf_accu + set_fail = set_f_fail | set_v_fail | set_mf_fail + set_cand = set_full - set_accu - set_fail # check size assert nframes == len(set_accu | set_cand | set_fail) assert 0 == len(set_accu & set_cand) @@ -271,6 +330,11 @@ def print( fmt_flt % (self.level_v_lo), fmt_flt % (self.level_v_hi), ) + if self.mf_level: + print_tuple += ( + fmt_flt % (self.level_mf_lo), + fmt_flt % (self.level_mf_hi), + ) print_tuple += (str(self.converged()),) ret = " " + fmt_str % print_tuple return ret diff --git a/dpgen2/exploration/report/report_trust_levels_max.py b/dpgen2/exploration/report/report_trust_levels_max.py index e847c1e6..bc4fd610 100644 --- a/dpgen2/exploration/report/report_trust_levels_max.py +++ b/dpgen2/exploration/report/report_trust_levels_max.py @@ -109,7 +109,9 @@ def make_class_doc_link(key): level_f_hi_link = make_class_doc_link("level_f_hi") level_v_hi_link = make_class_doc_link("level_v_hi") + level_mf_hi_link = make_class_doc_link("level_mf_hi") level_f_lo_link = make_class_doc_link("level_f_lo") level_v_lo_link = make_class_doc_link("level_v_lo") + level_mf_lo_link = make_class_doc_link("level_mf_lo") conv_accuracy_link = make_class_doc_link("conv_accuracy") - return f"The configurations with force model deviation between {level_f_lo_link}, {level_f_hi_link} or virial model deviation between {level_v_lo_link} and {level_v_hi_link} are treated as candidates (The virial model deviation check is optional). The configurations with maximal model deviation in the candidates are sent for FP calculations. If the ratio of accurate (below {level_f_lo_link} and {level_v_lo_link}) is higher then {conv_accuracy_link}, the stage is treated as converged." + return f"The configurations with force model deviation between {level_f_lo_link}, {level_f_hi_link} or virial model deviation between {level_v_lo_link} and {level_v_hi_link}, or magnetic force model deviation between {level_mf_lo_link} and {level_mf_hi_link}, are treated as candidates (The virial model deviation check is optional). The configurations with maximal model deviation in the candidates are sent for FP calculations. If the ratio of accurate (below {level_f_lo_link} and {level_v_lo_link}) is higher then {conv_accuracy_link}, the stage is treated as converged." diff --git a/dpgen2/exploration/task/__init__.py b/dpgen2/exploration/task/__init__.py index 534a8828..b60b0bce 100644 --- a/dpgen2/exploration/task/__init__.py +++ b/dpgen2/exploration/task/__init__.py @@ -10,6 +10,9 @@ from .diffcsp_task_group import ( DiffCSPTaskGroup, ) +from .lmp_spin_task_group import ( + LmpSpinTaskGroup, +) from .lmp_template_task_group import ( LmpTemplateTaskGroup, ) diff --git a/dpgen2/exploration/task/lmp_spin_task_group.py b/dpgen2/exploration/task/lmp_spin_task_group.py new file mode 100644 index 00000000..c900d522 --- /dev/null +++ b/dpgen2/exploration/task/lmp_spin_task_group.py @@ -0,0 +1,112 @@ +import itertools +import random +from pathlib import ( + Path, +) +from typing import ( + List, + Optional, +) + +from dpgen2.constants import ( + lmp_conf_name, + lmp_input_name, + lmp_traj_name, + model_name_pattern, + plm_input_name, + plm_output_name, +) + +from .conf_sampling_task_group import ( + ConfSamplingTaskGroup, +) +from .lmp import ( + make_lmp_input, +) +from .task import ( + ExplorationTask, +) + + +class LmpSpinTaskGroup(ConfSamplingTaskGroup): + def __init__( + self, + ): + super().__init__() + self.lmp_set = False + self.plm_set = False + + def set_lmp( + self, + numb_models: int, + lmp_template_fname: str, + plm_template_fname: Optional[str] = None, + revisions: dict = {}, + ) -> None: + self.lmp_template = Path(lmp_template_fname).read_text().split("\n") + self.revisions = revisions + self.lmp_set = True + self.model_list = sorted([model_name_pattern % ii for ii in range(numb_models)]) + if plm_template_fname is not None: + self.plm_template = Path(plm_template_fname).read_text().split("\n") + self.plm_set = True + + def make_task( + self, + ) -> "LmpSpinTaskGroup": + if not self.conf_set: + raise RuntimeError("confs are not set") + if not self.lmp_set: + raise RuntimeError("Lammps SPIN template and revisions are not set") + # clear all existing tasks + self.clear() + confs = self._sample_confs() + templates = [self.lmp_template] + conts = self.make_cont(templates, self.revisions) + nconts = len(conts[0]) + for cc, ii in itertools.product(confs, range(nconts)): # type: ignore + self.add_task(self._make_lmp_task(cc, conts[0][ii])) + return self + + def make_cont( + self, + templates: list, + revisions: dict, + ): + keys = revisions.keys() + prod_vv = [revisions[kk] for kk in keys] + ntemplate = len(templates) + ret = [[] for ii in range(ntemplate)] + for vv in itertools.product(*prod_vv): + for ii in range(ntemplate): + tt = templates[ii].copy() + ret[ii].append("\n".join(revise_by_keys(tt, keys, vv))) + return ret + + def _make_lmp_task( + self, + conf: str, + lmp_cont: str, + plm_cont: Optional[str] = None, + ) -> ExplorationTask: + task = ExplorationTask() + task.add_file( + lmp_conf_name, + conf, + ).add_file( + lmp_input_name, + lmp_cont, + ) + if plm_cont is not None: + task.add_file( + plm_input_name, + plm_cont, + ) + return task + + +def revise_by_keys(lmp_lines, keys, values): + for kk, vv in zip(keys, values): # type: ignore + for ii in range(len(lmp_lines)): + lmp_lines[ii] = lmp_lines[ii].replace(kk, str(vv)) + return lmp_lines diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index 3b793c58..78cc148e 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -19,6 +19,9 @@ from dpgen2.exploration.task.customized_lmp_template_task_group import ( CustomizedLmpTemplateTaskGroup, ) +from dpgen2.exploration.task.lmp_spin_task_group import ( + LmpSpinTaskGroup, +) from dpgen2.exploration.task.lmp_template_task_group import ( LmpTemplateTaskGroup, ) @@ -177,6 +180,45 @@ def lmp_template_task_group_args(): ] +def lmp_spin_task_group_args(): + doc_lmp_template_fname = "The file name of lammps input template" + doc_plm_template_fname = "The file name of plumed input template" + doc_revisions = "The revisions. Should be a dict providing the key - list of desired values pair. Key is the word to be replaced in the templates, and it may appear in both the lammps and plumed input templates. All values in the value list will be enmerated." + + return [ + Argument("conf_idx", list, optional=False, doc=doc_conf_idx, alias=["sys_idx"]), + Argument( + "n_sample", + int, + optional=True, + default=None, + doc=doc_n_sample, + ), + Argument( + "lmp_template_fname", + str, + optional=False, + doc=doc_lmp_template_fname, + alias=["lmp_template", "lmp"], + ), + Argument( + "plm_template_fname", + str, + optional=True, + default=None, + doc=doc_plm_template_fname, + alias=["plm_template", "plm"], + ), + Argument( + "revisions", + dict, + optional=True, + default={}, + doc=doc_revisions, + ), + ] + + def customized_lmp_template_task_group_args(): doc_input_lmp_tmpl_name = "The file name of lammps input template" doc_input_plm_tmpl_name = "The file name of plumed input template" @@ -303,6 +345,12 @@ def variant_task_group(): lmp_template_task_group_args(), doc=doc_lmp_template, ), + Argument( + "lmp-spin", + dict, + lmp_spin_task_group_args(), + doc=doc_lmp_template, + ), Argument( "customized-lmp-template", dict, @@ -635,6 +683,15 @@ def make_lmp_task_group_from_config( lmp_template, **config, ) + elif config["type"] == "lmp-spin": + tgroup = LmpSpinTaskGroup() + config.pop("type") + lmp_template = config.pop("lmp_template_fname") + tgroup.set_lmp( + numb_models, + lmp_template, + **config, + ) elif config["type"] == "customized-lmp-template": tgroup = CustomizedLmpTemplateTaskGroup() config.pop("type") diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 5a9782f4..d7cb656f 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -423,12 +423,22 @@ def write_other_to_input_script( for v in odict["loss_dict"].values(): if isinstance(v, dict): v["start_pref_e"] = config["init_model_start_pref_e"] - v["start_pref_f"] = config["init_model_start_pref_f"] v["start_pref_v"] = config["init_model_start_pref_v"] + # for spin job, the key is "start_pref_fr" and "start_pref_fm" + if config.get("spin", False): + v["start_pref_fr"] = config["init_model_start_pref_f"] + v["start_pref_fm"] = config["init_model_start_pref_fm"] + else: + v["start_pref_f"] = config["init_model_start_pref_f"] + else: odict["loss"]["start_pref_e"] = config["init_model_start_pref_e"] - odict["loss"]["start_pref_f"] = config["init_model_start_pref_f"] odict["loss"]["start_pref_v"] = config["init_model_start_pref_v"] + if config.get("spin", False): + odict["loss"]["start_pref_fr"] = config["init_model_start_pref_f"] + odict["loss"]["start_pref_fm"] = config["init_model_start_pref_fm"] + else: + odict["loss"]["start_pref_f"] = config["init_model_start_pref_f"] if major_version == "1": odict["training"]["stop_batch"] = config["init_model_numb_steps"] elif major_version == "2": @@ -517,6 +527,10 @@ def training_args(): doc_init_model_start_pref_f = ( "The start force prefactor in loss when init-model" ) + doc_init_model_start_pref_fm = ( + "The start magnetic force prefactor in loss when init-model for spin job" + ) + doc_spin = "If is a spin job" doc_init_model_start_pref_v = ( "The start virial prefactor in loss when init-model" ) @@ -587,6 +601,20 @@ def training_args(): default=100, doc=doc_init_model_start_pref_f, ), + Argument( + "init_model_start_pref_fm", + float, + optional=True, + default=100, + doc=doc_init_model_start_pref_fm, + ), + Argument( + "spin", + bool, + optional=True, + default=False, + doc=doc_spin, + ), Argument( "init_model_start_pref_v", float, diff --git a/tests/exploration/test_devi_manager.py b/tests/exploration/test_devi_manager.py index a0a29e68..027cf8e9 100644 --- a/tests/exploration/test_devi_manager.py +++ b/tests/exploration/test_devi_manager.py @@ -31,11 +31,13 @@ def test_success(self): ) ) self.assertEqual(model_devi.get(DeviManager.MAX_DEVI_V), [None, None]) + self.assertEqual(model_devi.get(DeviManager.MIN_DEVI_MF), [None, None]) model_devi.clear() self.assertEqual(model_devi.ntraj, 0) self.assertEqual(model_devi.get(DeviManager.MAX_DEVI_F), []) self.assertEqual(model_devi.get(DeviManager.MAX_DEVI_V), []) + self.assertEqual(model_devi.get(DeviManager.MIN_DEVI_MF), []) def test_add_invalid_name(self): model_devi = DeviManagerStd() @@ -72,6 +74,7 @@ def test_devi_manager_std_check_data(self): model_devi.add(DeviManager.MAX_DEVI_F, np.array([1, 2, 3])) model_devi.add(DeviManager.MAX_DEVI_F, np.array([4, 5, 6])) model_devi.add(DeviManager.MAX_DEVI_V, np.array([4, 5, 6])) + model_devi.add(DeviManager.MAX_DEVI_MF, np.array([7, 8, 9])) self.assertEqual(model_devi.ntraj, 2) @@ -82,6 +85,13 @@ def test_devi_manager_std_check_data(self): DeviManager.MAX_DEVI_V, ) + self.assertRaisesRegex( + AssertionError, + "Error: the number of model deviation", + model_devi.get, + DeviManager.MAX_DEVI_MF, + ) + model_devi = DeviManagerStd() model_devi.add(DeviManager.MAX_DEVI_V, np.array([1, 2, 3])) @@ -109,3 +119,21 @@ def test_devi_manager_std_check_data(self): model_devi.get, DeviManager.MAX_DEVI_V, ) + + model_devi = DeviManagerStd() + model_devi.add(DeviManager.MAX_DEVI_F, np.array([1, 2, 3])) + model_devi.add(DeviManager.MAX_DEVI_F, np.array([4, 5, 6])) + model_devi.add(DeviManager.MAX_DEVI_MF, np.array([1, 2, 3])) + model_devi.add(DeviManager.MAX_DEVI_MF, np.array([4, 5])) + self.assertRaisesRegex( + AssertionError, + f"Error: the number of frames in", + model_devi.get, + DeviManager.MAX_DEVI_F, + ) + self.assertRaisesRegex( + AssertionError, + f"Error: the number of frames in", + model_devi.get, + DeviManager.MAX_DEVI_MF, + ) diff --git a/tests/exploration/test_lmp_spin_task_group.py b/tests/exploration/test_lmp_spin_task_group.py new file mode 100644 index 00000000..721889c0 --- /dev/null +++ b/tests/exploration/test_lmp_spin_task_group.py @@ -0,0 +1,218 @@ +import itertools +import os +import textwrap +import unittest +from pathlib import ( + Path, +) +from typing import ( + List, + Set, +) + +import numpy as np + +try: + from exploration.context import ( + dpgen2, + ) +except ModuleNotFoundError: + # case of upload everything to argo, no context needed + pass +from unittest.mock import ( + Mock, + patch, +) + +from dpgen2.constants import ( + lmp_conf_name, + lmp_input_name, + plm_input_name, +) +from dpgen2.exploration.task import ( + ExplorationStage, + LmpSpinTaskGroup, +) + +in_lmp_template = textwrap.dedent( + """variable NSTEPS equal V_NSTEPS +variable THERMO_FREQ equal 10 +variable DUMP_FREQ equal 10 +variable TEMP equal V_TEMP +variable PRES equal 0.0 +variable TAU_T equal 0.100000 +variable TAU_P equal 0.500000 +variable MASS equal V_MASS + +units metal +boundary p p p +atom_style atomic + +neighbor 1.0 bin + +box tilt large +read_data conf.lmp +change_box all triclinic +mass 1 27.000000 +mass 2 24.000000 + +pair_style deepspin model.000.pth model.001.pth model.002.pth model.003.pth out_freq ${THERMO_FREQ} +pair_coeff * * + +thermo_style custom step temp pe ke etotal press vol lx ly lz xy xz yz +thermo ${THERMO_FREQ} + +dump dpgen_dump + +velocity all create ${TEMP} 826513 +fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P} mass ${MASS} + +timestep 0.002000 +run + +compute mag all spin +compute pe all pe +compute ke all ke +compute temp all temp +compute spin all property/atom sp spx spy spz fmx fmy fmz fx fy fz + +thermo 10 +thermo_style custom step time v_magnorm temp v_tmag press etotal ke pe v_emag econserve + +dump dpgen_dump all custom 10 traj.dump id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10] +dump_modify dpgen_dump sort id + +run ${NSTEPS} +""" +) + +expected_lmp_template = textwrap.dedent( + """variable NSTEPS equal V_NSTEPS +variable THERMO_FREQ equal 10 +variable DUMP_FREQ equal 10 +variable TEMP equal V_TEMP +variable PRES equal 0.0 +variable TAU_T equal 0.100000 +variable TAU_P equal 0.500000 +variable MASS equal V_MASS + +units metal +boundary p p p +atom_style atomic + +neighbor 1.0 bin + +box tilt large +read_data conf.lmp +change_box all triclinic +mass 1 27.000000 +mass 2 24.000000 + +pair_style deepspin model.000.pth model.001.pth model.002.pth model.003.pth out_freq ${THERMO_FREQ} +pair_coeff * * + +thermo_style custom step temp pe ke etotal press vol lx ly lz xy xz yz +thermo ${THERMO_FREQ} + +dump dpgen_dump + +velocity all create ${TEMP} 826513 +fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P} mass ${MASS} + +timestep 0.002000 +run + +compute mag all spin +compute pe all pe +compute ke all ke +compute temp all temp +compute spin all property/atom sp spx spy spz fmx fmy fmz fx fy fz + +thermo 10 +thermo_style custom step time v_magnorm temp v_tmag press etotal ke pe v_emag econserve + +dump dpgen_dump all custom 10 traj.dump id type x y z c_spin[1] c_spin[2] c_spin[3] c_spin[4] c_spin[5] c_spin[6] c_spin[7] c_spin[8] c_spin[9] c_spin[10] +dump_modify dpgen_dump sort id + +run ${NSTEPS} +""" +) + + +class TestLmpSpinTaskGroup(unittest.TestCase): + def setUp(self): + self.lmp_template_fname = Path("lmp.template") + self.lmp_template_fname.write_text(in_lmp_template) + self.numb_models = 8 + self.confs = ["foo", "bar"] + self.lmp_rev_mat = { + "V_NSTEPS": [1000], + "V_TEMP": [50, 100], + "V_MASS": [0.01, 0.1], + } + self.rev_empty = {} + + def tearDown(self): + os.remove(self.lmp_template_fname) + + def test_lmp(self): + task_group = LmpSpinTaskGroup() + task_group.set_conf(self.confs) + task_group.set_lmp( + self.numb_models, self.lmp_template_fname, revisions=self.lmp_rev_mat + ) + task_group.make_task() + ngroup = len(task_group) + self.assertEqual( + ngroup, + len(self.confs) + * len(self.lmp_rev_mat["V_NSTEPS"]) + * len(self.lmp_rev_mat["V_TEMP"]) + * len(self.lmp_rev_mat["V_MASS"]), + ) + + idx = 0 + for cc, ii, jj, kk in itertools.product( + range(len(self.confs)), + range(len(self.lmp_rev_mat["V_NSTEPS"])), + range(len(self.lmp_rev_mat["V_TEMP"])), + range(len(self.lmp_rev_mat["V_MASS"])), + ): + ee = expected_lmp_template.split("\n") + ee[0] = ee[0].replace("V_NSTEPS", str(self.lmp_rev_mat["V_NSTEPS"][ii])) + ee[3] = ee[3].replace("V_TEMP", str(self.lmp_rev_mat["V_TEMP"][jj])) + ee[7] = ee[7].replace("V_MASS", str(self.lmp_rev_mat["V_MASS"][kk])) + self.assertEqual( + task_group[idx].files()[lmp_conf_name], + self.confs[cc], + ) + self.assertEqual( + task_group[idx].files()[lmp_input_name].split("\n"), + ee, + ) + idx += 1 + + def test_lmp_empty(self): + task_group = LmpSpinTaskGroup() + task_group.set_conf(self.confs) + task_group.set_lmp( + self.numb_models, self.lmp_template_fname, revisions=self.rev_empty + ) + task_group.make_task() + ngroup = len(task_group) + self.assertEqual( + ngroup, + len(self.confs), + ) + idx = 0 + for cc in range(len(self.confs)): + ee = expected_lmp_template.split("\n") + self.assertEqual( + task_group[idx].files()[lmp_conf_name], + self.confs[cc], + ) + self.assertEqual( + task_group[idx].files()[lmp_input_name].split("\n"), + ee, + ) + idx += 1 diff --git a/tests/exploration/test_report_adaptive_lower.py b/tests/exploration/test_report_adaptive_lower.py index f5c13354..1bd1384c 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -296,6 +296,147 @@ class MockedReport: self.assertEqual(ter.accurate_ratio(), 9.0 / 18.0) self.assertEqual(ter.failed_ratio(), 6.0 / 18.0) + def test_mf(self): + model_devi = DeviManagerStd() + model_devi.add( + DeviManager.MAX_DEVI_MF, + np.array([0.90, 0.10, 0.91, 0.11, 0.50, 0.53, 0.51, 0.52, 0.92]), + ) + model_devi.add( + DeviManager.MAX_DEVI_MF, + np.array([0.40, 0.20, 0.80, 0.81, 0.82, 0.21, 0.41, 0.22, 0.42]), + ) + + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.40, 0.20, 0.21, 0.80, 0.81, 0.53, 0.22, 0.82, 0.42]), + ) + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.50, 0.90, 0.91, 0.92, 0.51, 0.52, 0.10, 0.11, 0.12]), + ) + + expected_fail_ = [[0, 2, 8], [2, 3, 4]] + expected_fail = set() + for idx, ii in enumerate(expected_fail_): + for jj in ii: + expected_fail.add((idx, jj)) + expected_cand = set([(0, 6), (0, 7), (0, 5)]) + expected_accu = set( + [(0, 1), (0, 3), (0, 4), (1, 0), (1, 1), (1, 5), (1, 6), (1, 7), (1, 8)] + ) + + ter = ExplorationReportAdaptiveLower( + level_f_hi=1.0, + numb_candi_f=0, + rate_candi_f=0.000, + level_mf_hi=0.76, + numb_candi_mf=3, + rate_candi_mf=0.001, + n_checked_steps=3, + conv_tolerance=0.001, + ) + ter.record(model_devi) + self.assertEqual(ter.candi, expected_cand) + self.assertEqual(ter.accur, expected_accu) + self.assertEqual(set(ter.failed), expected_fail) + + class MockedReport: + level_f_lo = 0 + level_v_lo = 0 + level_mf_lo = 0 + + mr = MockedReport() + mr.level_f_lo = 1.0 + mr.level_v_lo = 0.51 + mr.level_mf_lo = 0.51 + self.assertFalse(ter.converged([mr])) + self.assertTrue(ter.converged([mr, mr])) + self.assertTrue(ter.converged([mr, mr, mr])) + + picked = ter.get_candidate_ids(2) + npicked = 0 + self.assertEqual(len(picked), 2) + for ii in range(2): + for jj in picked[ii]: + self.assertTrue((ii, jj) in expected_cand) + npicked += 1 + self.assertEqual(npicked, 2) + self.assertEqual(ter.candidate_ratio(), 3.0 / 18.0) + self.assertEqual(ter.accurate_ratio(), 9.0 / 18.0) + self.assertEqual(ter.failed_ratio(), 6.0 / 18.0) + + #### need modify + def test_v_mf(self): + model_devi = DeviManagerStd() + model_devi.add( + DeviManager.MAX_DEVI_V, + np.array([0.90, 0.10, 0.91, 0.11, 0.50, 0.53, 0.51, 0.52, 0.92]), + ) + model_devi.add( + DeviManager.MAX_DEVI_V, + np.array([0.40, 0.20, 0.80, 0.81, 0.82, 0.21, 0.41, 0.22, 0.42]), + ) + + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.40, 0.20, 0.21, 0.80, 0.81, 0.53, 0.22, 0.82, 0.42]), + ) + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.50, 0.90, 0.91, 0.92, 0.51, 0.52, 0.10, 0.11, 0.12]), + ) + + expected_fail_ = [[0, 2, 8], [2, 3, 4]] + expected_fail = set() + for idx, ii in enumerate(expected_fail_): + for jj in ii: + expected_fail.add((idx, jj)) + expected_cand = set([(0, 6), (0, 7), (0, 5)]) + expected_accu = set( + [(0, 1), (0, 3), (0, 4), (1, 0), (1, 1), (1, 5), (1, 6), (1, 7), (1, 8)] + ) + + ter = ExplorationReportAdaptiveLower( + level_f_hi=1.0, + numb_candi_f=0, + rate_candi_f=0.000, + level_v_hi=0.76, + numb_candi_v=3, + rate_candi_v=0.001, + n_checked_steps=3, + conv_tolerance=0.001, + ) + ter.record(model_devi) + self.assertEqual(ter.candi, expected_cand) + self.assertEqual(ter.accur, expected_accu) + self.assertEqual(set(ter.failed), expected_fail) + + class MockedReport: + level_f_lo = 0 + level_v_lo = 0 + level_mf_lo = 0 + + mr = MockedReport() + mr.level_f_lo = 1.0 + mr.level_v_lo = 0.51 + mr.level_mf_lo = 0.51 + self.assertFalse(ter.converged([mr])) + self.assertTrue(ter.converged([mr, mr])) + self.assertTrue(ter.converged([mr, mr, mr])) + + picked = ter.get_candidate_ids(2) + npicked = 0 + self.assertEqual(len(picked), 2) + for ii in range(2): + for jj in picked[ii]: + self.assertTrue((ii, jj) in expected_cand) + npicked += 1 + self.assertEqual(npicked, 2) + self.assertEqual(ter.candidate_ratio(), 3.0 / 18.0) + self.assertEqual(ter.accurate_ratio(), 9.0 / 18.0) + self.assertEqual(ter.failed_ratio(), 6.0 / 18.0) + def test_args(self): input_dict = { "level_f_hi": 1.0, @@ -315,4 +456,6 @@ def test_args(self): self.assertEqual(data["n_checked_steps"], 2) self.assertAlmostEqual(data["conv_tolerance"], 0.01) self.assertAlmostEqual(data["candi_sel_prob"], "uniform") + self.assertTrue(data["level_mf_hi"] is None) + self.assertEqual(data["numb_candi_mf"], 0) ExplorationReportAdaptiveLower(*data) diff --git a/tests/exploration/test_report_trust_levels.py b/tests/exploration/test_report_trust_levels.py index 7f6c79a9..ba40d69b 100644 --- a/tests/exploration/test_report_trust_levels.py +++ b/tests/exploration/test_report_trust_levels.py @@ -37,6 +37,9 @@ def test_exploration_report_trust_levels(self): self.fv_selection_test(ExplorationReportTrustLevelsRandom) self.fv_selection_test(ExplorationReportTrustLevelsMax) + self.mf_selection_test(ExplorationReportTrustLevelsRandom) + self.mf_selection_test(ExplorationReportTrustLevelsMax) + self.f_selection_test(ExplorationReportTrustLevelsRandom) self.f_selection_test(ExplorationReportTrustLevelsMax) @@ -97,6 +100,53 @@ def fv_selection_test(self, exploration_report: ExplorationReportTrustLevels): self.assertEqual(ter.accurate_ratio(), 2.0 / 18.0) self.assertEqual(ter.failed_ratio(), 10.0 / 18.0) + def mf_selection_test(self, exploration_report: ExplorationReportTrustLevels): + model_devi = DeviManagerStd() + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.90, 0.10, 0.91, 0.11, 0.50, 0.12, 0.51, 0.52, 0.92]), + ) + model_devi.add( + DeviManager.MAX_DEVI_F, + np.array([0.40, 0.20, 0.80, 0.81, 0.82, 0.21, 0.41, 0.22, 0.42]), + ) + model_devi.add( + DeviManager.MAX_DEVI_MF, + np.array([0.40, 0.20, 0.21, 0.80, 0.81, 0.41, 0.22, 0.82, 0.42]), + ) + model_devi.add( + DeviManager.MAX_DEVI_MF, + np.array([0.50, 0.90, 0.91, 0.92, 0.51, 0.52, 0.10, 0.11, 0.12]), + ) + + expected_accu = [[1], [7]] + expected_cand = [[6, 5], [8, 6, 0, 5]] + expected_fail = [[0, 2, 3, 4, 7, 8], [1, 2, 3, 4]] + expected_accu = [set(ii) for ii in expected_accu] + expected_cand = [set(ii) for ii in expected_cand] + expected_fail = [set(ii) for ii in expected_fail] + all_cand_sel = [(0, 6), (0, 5), (1, 8), (1, 6), (1, 0), (1, 5)] + + ter = exploration_report( + 0.3, 0.6, level_mf_lo=0.3, level_mf_hi=0.6, conv_accuracy=0.9 + ) + ter.record(model_devi) + self.assertEqual(ter.traj_cand, expected_cand) + self.assertEqual(ter.traj_accu, expected_accu) + self.assertEqual(ter.traj_fail, expected_fail) + + picked = ter.get_candidate_ids(2) + npicked = 0 + self.assertEqual(len(picked), 2) + for ii in range(2): + for jj in picked[ii]: + self.assertTrue(jj in expected_cand[ii]) + npicked += 1 + self.assertEqual(npicked, 2) + self.assertEqual(ter.candidate_ratio(), 6.0 / 18.0) + self.assertEqual(ter.accurate_ratio(), 2.0 / 18.0) + self.assertEqual(ter.failed_ratio(), 10.0 / 18.0) + def f_selection_test(self, exploration_report: ExplorationReportTrustLevels): model_devi = DeviManagerStd() model_devi.add( diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 7649b520..09dad868 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -89,7 +89,19 @@ def setUp(self): "init_model_start_pref_f": 100, "init_model_start_pref_v": 0.0, } + self.config_spin = { + "spin": True, + "init_model_policy": "no", + "init_model_old_ratio": 0.9, + "init_model_numb_steps": 400000, + "init_model_start_lr": 1e-4, + "init_model_start_pref_e": 0.1, + "init_model_start_pref_f": 100, + "init_model_start_pref_fm": 100, + "init_model_start_pref_v": 0.0, + } self.config = RunDPTrain.normalize_config(self.config) + self.config_spin = RunDPTrain.normalize_config(self.config_spin) self.old_data_size = ( self.init_nframs_0 + self.init_nframs_1 + sum(self.nframes_0) @@ -169,6 +181,35 @@ def setUp(self): }, } + self.expected_init_model_odict_v2_spin = { + "training": { + "training_data": { + "systems": [ + "init/data-0", + "init/data-1", + "data-0/foo3", + "data-0/foo4", + "data-1/foo2", + "data-1/foo3", + "data-1/foo5", + ], + "batch_size": "auto", + "auto_prob": "prob_sys_size; 0:4:0.9; 4:7:0.1", + }, + "disp_file": "lcurve.out", + "numb_steps": 400000, + }, + "learning_rate": { + "start_lr": 1e-4, + }, + "loss": { + "start_pref_e": 0.1, + "start_pref_v": 0.0, + "start_pref_fr": 100, + "start_pref_fm": 100, + }, + } + self.idict_v1 = { "training": { "systems": [], @@ -548,6 +589,75 @@ def test_exec_v2_init_model(self, mocked_run): jdata = json.load(fp) self.assertDictEqual(jdata, self.expected_init_model_odict_v2) + @patch("dpgen2.op.run_dp_train.run_command") + def test_exec_v2_init_model_spin(self, mocked_run): + mocked_run.side_effect = [(0, "foo\n", ""), (0, "bar\n", "")] + + config = self.config_spin.copy() + config["init_model_policy"] = "yes" + + task_path = self.task_path + Path(task_path).mkdir(exist_ok=True) + idict_v2 = self.idict_v2.copy() + idict_v2["loss"] = { + "start_pref_e": 0.1, + "start_pref_v": 1, + "start_pref_fr": 1, + "start_pref_fm": 1, + } + with open(Path(task_path) / train_script_name, "w") as fp: + json.dump(idict_v2, fp, indent=4) + task_name = self.task_name + work_dir = Path(task_name) + + ptrain = RunDPTrain() + out = ptrain.execute( + OPIO( + { + "config": config, + "task_name": task_name, + "task_path": Path(task_path), + "init_model": Path(self.init_model), + "init_data": [Path(ii) for ii in self.init_data], + "iter_data": [Path(ii) for ii in self.iter_data], + } + ) + ) + self.assertEqual(out["script"], work_dir / train_script_name) + self.assertEqual(out["model"], work_dir / "frozen_model.pb") + self.assertEqual(out["lcurve"], work_dir / "lcurve.out") + self.assertEqual(out["log"], work_dir / "train.log") + + calls = [ + call( + [ + "dp", + "train", + "--init-frz-model", + str(self.init_model), + train_script_name, + ] + ), + call(["dp", "freeze", "-o", "frozen_model.pb"]), + ] + mocked_run.assert_has_calls(calls) + + self.assertTrue(work_dir.is_dir()) + self.assertTrue(out["log"].is_file()) + self.assertEqual( + out["log"].read_text(), + "#=================== train std out ===================\n" + "foo\n" + "#=================== train std err ===================\n" + "#=================== freeze std out ===================\n" + "bar\n" + "#=================== freeze std err ===================\n", + ) + with open(out["script"]) as fp: + jdata = json.load(fp) + print(jdata) + self.assertDictEqual(jdata, self.expected_init_model_odict_v2_spin) + @patch("dpgen2.op.run_dp_train.run_command") def test_exec_v2_train_error(self, mocked_run): mocked_run.side_effect = [(1, "", "foo\n"), (0, "bar\n", "")]