From 07df321d710b82409e1d65a25c342cf3ebb3dfca Mon Sep 17 00:00:00 2001 From: zhenyu <76582286+wangzyphysics@users.noreply.github.com> Date: Tue, 30 Apr 2024 20:47:41 +0800 Subject: [PATCH] Calypso speedup by refactorizing model deviation. (#217) 1. add a slice when running op `run_caly_model_devi` 2. support `distanceofion` in dict format, e.g. distanceofion = {"Mg": 0.8, "Al": 0.9}, the value denotes the radius of each element, and the mindistance of different element is equal to the sum of value of element Mg and Al and multipy 0.7. 3. support providing a list of list format `name_of_element`, e.g. name_of_element = [["Mg", "Al", "Cu", "Li"], ["La", "Ce", "Ca", "Li"], ["H"]], code will randomly choose one of element in each list to set up a input file of calypso. ## Summary by CodeRabbit - **New Features** - Introduced a new operation `PrepCalyModelDevi` for organizing and preparing model deviations. - Enhanced the concurrent learning operation to include the preparation of Calypso model deviations. - Updated the Calypso task generation logic to handle different types of input parameters more flexibly. - **Bug Fixes** - Improved error handling and logging in the Calypso operation to better manage failures. - **Tests** - Added new tests to validate the creation and handling of Calypso tasks with various configurations. - Implemented unit tests for the new `PrepCalyModelDevi` class to ensure its functionality. - **Refactor** - Reorganized logic and removed outdated code in several operations to streamline processes and improve clarity. --- dpgen2/entrypoint/submit.py | 7 +- dpgen2/exploration/task/caly_task_group.py | 110 +++++++++++++++- dpgen2/exploration/task/calypso/caly_input.py | 2 +- .../task/make_task_group_from_config.py | 6 +- dpgen2/op/__init__.py | 3 + dpgen2/op/collect_run_caly.py | 47 ++++--- dpgen2/op/prep_caly_dp_optim.py | 5 +- dpgen2/op/prep_caly_model_devi.py | 122 ++++++++++++++++++ dpgen2/superop/prep_run_calypso.py | 45 +++++-- tests/exploration/test_exploration_group.py | 54 ++++++-- tests/mocked_ops.py | 54 ++++---- tests/op/test_prep_caly_model_devi.py | 119 +++++++++++++++++ tests/op/test_run_caly_model_devi.py | 9 +- tests/test_prep_run_caly.py | 6 +- 14 files changed, 503 insertions(+), 86 deletions(-) create mode 100644 dpgen2/op/prep_caly_model_devi.py create mode 100644 tests/op/test_prep_caly_model_devi.py diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 5d81b041..b11682b7 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -94,6 +94,7 @@ CollRunCaly, PrepCalyDPOptim, PrepCalyInput, + PrepCalyModelDevi, PrepDPTrain, PrepLmp, RunCalyDPOptim, @@ -183,6 +184,7 @@ def make_concurrent_learning_op( "prep-run-calypso", prep_caly_input_op=PrepCalyInput, caly_evo_step_op=caly_evo_step_op, + prep_caly_model_devi_op=PrepCalyModelDevi, run_caly_model_devi_op=RunCalyModelDevi, prep_config=prep_explore_config, run_config=run_explore_config, @@ -803,8 +805,10 @@ def get_superop(key): return re.sub("prep-dp-optim-[0-9]*-[0-9]*", "prep-run-explore", key) elif "run-dp-optim-" in key: return re.sub("run-dp-optim-[0-9]*-[0-9]*-[0-9]*", "prep-run-explore", key) + elif "prep-caly-model-devi" in key: + return key.replace("prep-caly-model-devi", "prep-run-explore") elif "run-caly-model-devi" in key: - return key.replace("run-caly-model-devi", "prep-run-explore") + return re.sub("run-caly-model-devi-[0-9]*", "prep-run-explore", key) return None @@ -849,6 +853,7 @@ def get_resubmit_keys( "collect-run-calypso", "prep-dp-optim", "run-dp-optim", + "prep-caly-model-devi", "run-caly-model-devi", "prep-run-explore", "prep-lmp", diff --git a/dpgen2/exploration/task/caly_task_group.py b/dpgen2/exploration/task/caly_task_group.py index 4a0b2c25..6eef4827 100644 --- a/dpgen2/exploration/task/caly_task_group.py +++ b/dpgen2/exploration/task/caly_task_group.py @@ -1,7 +1,12 @@ +import copy +import logging +import random from typing import ( List, ) +import numpy as np + from dpgen2.constants import ( calypso_check_opt_file, calypso_input_file, @@ -18,6 +23,55 @@ ExplorationTaskGroup, ) +atomic_symbols = ( + 'X', # placeholder + 'H', 'He', + 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', + 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', + 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', + 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', + 'Cs', 'Ba', + 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', + 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', + 'Fr', 'Ra', + 'Ac','Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', + 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og', +) # fmt: skip +atomic_number_map = {key: value for value, key in enumerate(atomic_symbols)} +# Covalent radii from: +# +# Covalent radii revisited, +# Beatriz Cordero, Verónica Gómez, Ana E. Platero-Prats, Marc Revés, +# Jorge Echeverría, Eduard Cremades, Flavia Barragán and Santiago Alvarez, +# Dalton Trans., 2008, 2832-2838 DOI:10.1039/B801115J +UNKN = 0.2 +covalent_radii = [ + # X, placeholder + UNKN, + # H He + 0.31, 0.28, + # Li Be B C N O F Ne + 1.28, 0.96, 0.84, 0.76, 0.71, 0.66, 0.57, 0.58, + # Na Mg Al Si P S Cl Ar + 1.66, 1.41, 1.21, 1.11, 1.07, 1.05, 1.02, 1.06, + # K Ca Sc Ti V Cr Mn Fe Co Ni Cu Zn Ga Ge As Se Br Kr + 2.03, 1.76, 1.70, 1.60, 1.53, 1.39, 1.39, 1.32, 1.26, 1.24, 1.32, 1.22, 1.22, 1.20, 1.19, 1.20, 1.20, 1.16, + # Rb Sr Y Zr Nb Mo Tc Ru Rh Pd Au Cd In Sn Sb Te I Xe + 2.20, 1.95, 1.90, 1.75, 1.64, 1.54, 1.47, 1.46, 1.42, 1.39, 1.45, 1.44, 1.42, 1.39, 1.39, 1.38, 1.39, 1.40, + # Cs Ba + 2.44, 2.15, + # La Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Lu + 2.07, 2.04, 2.03, 2.01, 1.99, 1.98, 1.98, 1.96, 1.94, 1.92, 1.92, 1.89, 1.90, 1.87, 1.87, + # Hf Ta W Re Os Ir Pt Au Hg Tl Pb Bi Po At Rn + 1.75, 1.70, 1.62, 1.51, 1.44, 1.41, 1.36, 1.36, 1.32, 1.45, 1.46, 1.48, 1.40, 1.50, 1.50, + # Fr Ra + 2.60, 2.21, + # Ac Th Pa U Np Pu Am Cm Bk Cf Es Fm Md No Lr + 2.15, 2.06, 2.00, 1.96, 1.90, 1.87, 1.80, 1.69, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, + # Rf Db Sg Bh Hs Mt Ds Rg Cn Nh Fl Mc Lv Ts Og + UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, UNKN, +] # fmt: skip + class CalyTaskGroup(ExplorationTaskGroup): def __init__(self): @@ -27,9 +81,9 @@ def set_params( self, numb_of_species, name_of_atoms, - atomic_number, numb_of_atoms, - distance_of_ions, + distance_of_ions=None, + atomic_number=None, pop_size: int = 30, max_step: int = 5, system_name: str = "CALYPSO", @@ -58,10 +112,56 @@ def set_params( Set calypso parameters """ self.numb_of_species = numb_of_species - self.name_of_atoms = name_of_atoms - self.atomic_number = atomic_number self.numb_of_atoms = numb_of_atoms - self.distance_of_ions = distance_of_ions + + if isinstance(name_of_atoms, list) and all( + [isinstance(i, list) for i in name_of_atoms] + ): + overlap = set(name_of_atoms[0]) + for temp in name_of_atoms[1:]: + overlap = overlap & set(temp) + + if any(map(lambda s: (set(s) - overlap) == 0, name_of_atoms)): + raise ValueError( + f"Any sub-list should not equal with intersection, e.g. [[A,B,C], [B,C], [C]] is not allowed." + ) + + while True: + choice = [] + for _atoms in name_of_atoms: + value = random.choice(_atoms) + logging.info( + f"randomly choose {value} from {_atoms}, already choose: {choice}" + ) + if value in choice: + break + choice.append(value) + else: + break + self.name_of_atoms = choice + logging.info(f"The final choice is {self.name_of_atoms}") + self.atomic_number = [atomic_symbols.index(i) for i in self.name_of_atoms] + else: + self.name_of_atoms = name_of_atoms + self.atomic_number = atomic_number + + if isinstance(distance_of_ions, dict): + updated_table = copy.deepcopy(covalent_radii) + for key, value in distance_of_ions.items(): + updated_table[atomic_number_map[key]] = value + + temp_distance_mtx = np.zeros((numb_of_species, numb_of_species)) + for i in range(numb_of_species): + for j in range(numb_of_species): + temp_distance_mtx[i][j] = round( + updated_table[atomic_number_map[self.name_of_atoms[i]]] * 0.7 + + updated_table[atomic_number_map[self.name_of_atoms[j]]] * 0.7, + 2, + ) + self.distance_of_ions = temp_distance_mtx + else: + self.distance_of_ions = distance_of_ions + self.pop_size = pop_size self.max_step = max_step self.system_name = system_name diff --git a/dpgen2/exploration/task/calypso/caly_input.py b/dpgen2/exploration/task/calypso/caly_input.py index 28b668a6..f12eccbd 100644 --- a/dpgen2/exploration/task/calypso/caly_input.py +++ b/dpgen2/exploration/task/calypso/caly_input.py @@ -241,7 +241,7 @@ def check(): def make_calypso_input( numb_of_species: int, name_of_atoms: List[str], - atomic_number: List[int], + atomic_number, numb_of_atoms: List[int], distance_of_ions, pop_size: int = 30, diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index 7705cbf6..584a980a 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -294,7 +294,7 @@ def caly_task_grp_args(): Argument( "atomic_number", list, - optional=False, + optional=True, doc="atomic number of each element.", ), Argument( @@ -305,8 +305,8 @@ def caly_task_grp_args(): ), Argument( "distance_of_ions", - list, - optional=False, + [list, dict], + optional=True, doc="the distance matrix between different elements.", ), Argument( diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index 573b8a07..ff4c8a8d 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -10,6 +10,9 @@ from .prep_caly_input import ( PrepCalyInput, ) +from .prep_caly_model_devi import ( + PrepCalyModelDevi, +) from .prep_dp_train import ( PrepDPTrain, ) diff --git a/dpgen2/op/collect_run_caly.py b/dpgen2/op/collect_run_caly.py index 80332c5c..d5738310 100644 --- a/dpgen2/op/collect_run_caly.py +++ b/dpgen2/op/collect_run_caly.py @@ -141,11 +141,6 @@ def execute( results = ( ip["results"].resolve() if ip["results"] is not None else ip["results"] ) - # opt_results_dir = ( - # ip["opt_results_dir"].resolve() - # if ip["opt_results_dir"] is not None - # else ip["opt_results_dir"] - # ) opt_results_dir = [] if ip["opt_results_dir"] is not None: for temp in ip["opt_results_dir"]: @@ -162,26 +157,30 @@ def execute( prep_last_calypso_file(step, results, opt_results_dir, qhull_input, vsc) # copy input.dat Path(input_file.name).symlink_to(input_file) - # run calypso - command = " ".join([command, ">", calypso_log_name]) - ret, out, err = run_command(command, shell=True) - if ret != 0: - logging.error( - "".join( - ( - "calypso failed\n", - "command was: ", - command, - "out msg: ", - out, - "\n", - "err msg: ", - err, - "\n", + + finished = "true" if int(cnt_num) == int(max_step) else "false" + + if finished == "false": + # run calypso + command = " ".join([command, ">", calypso_log_name]) + ret, out, err = run_command(command, shell=True) + if ret != 0: + logging.error( + "".join( + ( + "calypso failed\n", + "command was: ", + command, + "out msg: ", + out, + "\n", + "err msg: ", + err, + "\n", + ) ) ) - ) - raise TransientError("calypso failed") + raise TransientError("calypso failed") poscar_dir = Path("poscar_dir") poscar_dir.mkdir(parents=True, exist_ok=True) @@ -190,7 +189,7 @@ def execute( shutil.copyfile(poscar, target) step = Path("step").read_text().strip() - finished = "true" if int(cnt_num) == int(max_step) else "false" + # finished = "true" if int(cnt_num) == int(max_step) else "false" if not Path("test_qconvex.in").exists(): Path("test_qconvex.in").write_text("") diff --git a/dpgen2/op/prep_caly_dp_optim.py b/dpgen2/op/prep_caly_dp_optim.py index c982456e..81e69255 100644 --- a/dpgen2/op/prep_caly_dp_optim.py +++ b/dpgen2/op/prep_caly_dp_optim.py @@ -106,8 +106,6 @@ def execute( - `caly_check_opt_file` : (`Path`) """ - group_size = ip["template_slice_config"]["group_size"] - finished = ip["finished"] work_dir = Path(ip["task_name"]) @@ -119,6 +117,9 @@ def execute( caly_check_opt_file = _caly_check_opt_file.resolve() poscar_list = [poscar.resolve() for poscar in poscar_dir.rglob("POSCAR_*")] poscar_list = sorted(poscar_list, key=lambda x: int(x.name.strip("POSCAR_"))) + + group_size = ip["template_slice_config"].get("group_size", len(poscar_list)) + model_name = "frozen_model.pb" model_list = [model.resolve() for model in models_dir.rglob(model_name)] if len(model_list) == 0: diff --git a/dpgen2/op/prep_caly_model_devi.py b/dpgen2/op/prep_caly_model_devi.py new file mode 100644 index 00000000..c3ac8755 --- /dev/null +++ b/dpgen2/op/prep_caly_model_devi.py @@ -0,0 +1,122 @@ +import json +import logging +import pickle +import shutil +from pathlib import ( + Path, +) +from typing import ( + List, + Tuple, +) + +from dflow.python import ( + OP, + OPIO, + Artifact, + BigParameter, + OPIOSign, + Parameter, + TransientError, +) + +from dpgen2.constants import ( + calypso_check_opt_file, + calypso_opt_dir_name, + calypso_run_opt_file, + model_name_pattern, +) +from dpgen2.exploration.task import ( + ExplorationTaskGroup, +) +from dpgen2.utils import ( + BinaryFileInput, + set_directory, +) +from dpgen2.utils.run_command import ( + run_command, +) + + +class PrepCalyModelDevi(OP): + """Prepare the working directories and input file according to slices information + for making model deviation. + """ + + @classmethod + def get_input_sign(cls): + return OPIOSign( + { + "task_name": Parameter(str), + "template_slice_config": Parameter(dict), + "traj_results": Artifact(List[Path]), + } + ) + + @classmethod + def get_output_sign(cls): + return OPIOSign( + { + "task_name_list": Parameter(List[str]), + "grouped_traj_list": Artifact(List[Path]), + } + ) + + @OP.exec_sign_check + def execute( + self, + ip: OPIO, + ) -> OPIO: + """Execute the OP. + + Parameters + ---------- + ip : dict + Input dict with components: + - `task_name` : (`str`) + - `template_slice_config` : (`dict`) + - `traj_results` : (`Path`) + + Returns + ------- + op : dict + Output dict with components: + + - `task_name_list`: (`List[str]`) + - `grouped_traj_list`: (`Artifact(List[Path])`) + + """ + work_dir = Path(ip["task_name"]) + traj_results_dir = [Path(dir_name).resolve() for dir_name in ip["traj_results"]] + trajs = [ + traj.resolve() + for traj_dir in traj_results_dir + for traj in Path(traj_dir).rglob("*.traj") + ] + group_size = ip["template_slice_config"].get( + "model_devi_group_size", len(trajs) + ) + + with set_directory(work_dir): + grouped_trajs_list = [ + trajs[i : i + group_size] for i in range(0, len(trajs), group_size) + ] + + traj_cnt = 0 + task_dirs = [] + for idx, grouped_trajs in enumerate(grouped_trajs_list): + trajs_path = Path(f"trajs_part_{idx}") + task_dirs.append(work_dir / trajs_path) + with set_directory(trajs_path): + for traj in grouped_trajs: + Path(f"{traj_cnt}.{traj.name}").symlink_to(traj) + traj_cnt += 1 + + task_names = [str(task_dir) for task_dir in task_dirs] + + return OPIO( + { + "task_name_list": task_names, + "grouped_traj_list": task_dirs, + } + ) diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py index ae3bb2bb..647e5bb2 100644 --- a/dpgen2/superop/prep_run_calypso.py +++ b/dpgen2/superop/prep_run_calypso.py @@ -58,6 +58,7 @@ def __init__( name: str, prep_caly_input_op: Type[OP], caly_evo_step_op: OPTemplate, + prep_caly_model_devi_op: Type[OP], run_caly_model_devi_op: Type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), @@ -96,6 +97,7 @@ def __init__( self._keys = [ "prep-caly-input", "caly-evo-step-{{item}}", + "prep-caly-model-devi", "run-caly-model-devi", ] self.step_keys = {} @@ -103,6 +105,8 @@ def __init__( self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) ii = "caly-evo-step-{{item}}" self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + ii = "prep-caly-model-devi" + self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) ii = "run-caly-model-devi" self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) @@ -111,6 +115,7 @@ def __init__( self.step_keys, prep_caly_input_op, caly_evo_step_op, + prep_caly_model_devi_op, run_caly_model_devi_op, prep_config=prep_config, run_config=run_config, @@ -143,6 +148,7 @@ def _prep_run_caly( step_keys: Dict[str, Any], prep_caly_input_op: Type[OP], caly_evo_step_op: OPTemplate, + prep_caly_model_devi_op: Type[OP], run_caly_model_devi_op: Type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), @@ -219,32 +225,55 @@ def _prep_run_caly( ) prep_run_caly_steps.add(caly_evo_step) + # prep_caly_model_devi + prep_caly_model_devi = Step( + "prep-caly-model-devi", + template=PythonOPTemplate( + prep_caly_model_devi_op, + python_packages=upload_python_packages, + **run_template_config, + ), + parameters={ + "task_name": "prep-calypso-model-deviation", + "template_slice_config": template_slice_config, + }, + artifacts={ + "traj_results": caly_evo_step.outputs.artifacts["traj_results"], + }, + key="%s--prep-caly-model-devi" + % (prep_run_caly_steps.inputs.parameters["block_id"],), + executor=prep_executor, + ) + prep_run_caly_steps.add(prep_caly_model_devi) + # run model devi run_caly_model_devi = Step( "run-caly-model-devi", template=PythonOPTemplate( run_caly_model_devi_op, + slices=Slices( + input_parameter=["task_name"], + input_artifact=["traj_dirs"], + output_artifact=["traj", "model_devi"], + ), python_packages=upload_python_packages, - **prep_template_config, + **run_template_config, ), parameters={ "type_map": prep_run_caly_steps.inputs.parameters["type_map"], - "task_name": "run-calypso-model-devi", + "task_name": prep_caly_model_devi.outputs.parameters["task_name_list"], }, artifacts={ - "traj_dirs": caly_evo_step.outputs.artifacts["traj_results"], + "traj_dirs": prep_caly_model_devi.outputs.artifacts["grouped_traj_list"], "models": prep_run_caly_steps.inputs.artifacts["models"], }, - key=step_keys["run-caly-model-devi"], + key="%s--run-caly-model-devi-{{item}}" + % (prep_run_caly_steps.inputs.parameters["block_id"],), executor=run_executor, **prep_config, ) prep_run_caly_steps.add(run_caly_model_devi) - # prep_run_caly_steps.outputs.parameters[ - # "task_names" - # ].value_from_parameter = prep_caly_input.outputs.parameters["task_names"], - # run_caly_model_devi.outputs.parameters["task_name"] prep_run_caly_steps.outputs.artifacts[ "trajs" ]._from = run_caly_model_devi.outputs.artifacts["traj"] diff --git a/tests/exploration/test_exploration_group.py b/tests/exploration/test_exploration_group.py index 6f542ed1..f42b13c0 100644 --- a/tests/exploration/test_exploration_group.py +++ b/tests/exploration/test_exploration_group.py @@ -623,8 +623,32 @@ def setUp(self): "fmax": 1, } ) + self.config_random = caly_normalize( + { + "name_of_atoms": [ + ["Li", "La", "Mg", "Al"], + ["La", "Be", "Mg", "Ca"], + ["H"], + ], + "numb_of_atoms": [10, 10, 10], + "numb_of_species": 3, + "distance_of_ions": [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], + } + ) + self.config_random_dict = caly_normalize( + { + "name_of_atoms": [ + ["Li", "La", "Mg", "Al"], + ["La", "Be", "Mg", "Ca"], + ["H"], + ], + "numb_of_atoms": [10, 10, 10], + "numb_of_species": 3, + "distance_of_ions": {"Li": 1.0, "La": 1.1, "H": 0.5}, + } + ) - def test(self): + def test_00_make_calypso_task(self): tgroup_1 = make_calypso_task_group_from_config(self.config_1) tgroup_2 = make_calypso_task_group_from_config(self.config_2) @@ -639,13 +663,21 @@ def test(self): self.maxDiff = None ii = 0 self.assertEqual(task_group[ii].files()[calypso_input_file], ref_input) - # self.assertEqual( - # task_group[ii].files()[calypso_run_opt_file], - # ref_run_opt, - # ) - # ii += 1 - # self.assertEqual(task_group[ii].files()[lmp_conf_name], "foo") - # self.assertEqual( - # task_group[ii].files()[lmp_input_name], - # in_template_npt % (100.0, 10.0), - # ) + + def test_01_make_random_calypso_task(self): + tgroup_1 = make_calypso_task_group_from_config(self.config_random) + + stage = ExplorationStage() + stage.add_task_group(tgroup_1) + + task_group = stage.make_task() + self.assertEqual(len(task_group), 1) + + def test_02_make_random_dict_calypso_task(self): + tgroup_1 = make_calypso_task_group_from_config(self.config_random_dict) + + stage = ExplorationStage() + stage.add_task_group(tgroup_1) + + task_group = stage.make_task() + self.assertEqual(len(task_group), 1) diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index 89f570bf..97cf8ada 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -983,11 +983,7 @@ def execute( results = ( ip["results"].resolve() if ip["results"] is not None else ip["results"] ) - # opt_results_dir = ( - # ip["opt_results_dir"].resolve() - # if ip["opt_results_dir"] is not None - # else ip["opt_results_dir"] - # ) + opt_results_dir = [] if ip["opt_results_dir"] is not None: for temp in ip["opt_results_dir"]: @@ -1008,38 +1004,40 @@ def execute( Path(results.name).symlink_to(results) # Path(opt_results_dir.name).symlink_to(opt_results_dir) - for i in range(5): - Path(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") - - if step is None: - Path("step").write_text("2") - else: - step_num = Path("step").read_text().strip() - Path("step").write_text(f"{int(step_num)+1}") - - if qhull_input is None: - Path("test_qconvex.in").write_text("") - - step_num = int(Path("step").read_text().strip()) - - if results is None: - Path("results").mkdir(parents=True, exist_ok=True) - for i in range(1, step_num): + finished = "true" if int(cnt_num) == int(max_step) else "false" + if finished == "false": + for i in range(5): + Path(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + + if step is None: + Path("step").write_text("2") + else: + step_num = Path("step").read_text().strip() + Path("step").write_text(f"{int(step_num)+1}") + + if qhull_input is None: + Path("test_qconvex.in").write_text("") + + step_num = int(Path("step").read_text().strip()) + + if results is None: + Path("results").mkdir(parents=True, exist_ok=True) + for i in range(1, step_num): + Path(f"results/pso_ini_{i}").write_text(f"pso_ini_{i}") + Path(f"results/pso_opt_{i}").write_text(f"pso_opt_{i}") + Path(f"results/pso_sor_{i}").write_text(f"pso_sor_{i}") + else: + i = step_num Path(f"results/pso_ini_{i}").write_text(f"pso_ini_{i}") Path(f"results/pso_opt_{i}").write_text(f"pso_opt_{i}") Path(f"results/pso_sor_{i}").write_text(f"pso_sor_{i}") - else: - i = step_num - Path(f"results/pso_ini_{i}").write_text(f"pso_ini_{i}") - Path(f"results/pso_opt_{i}").write_text(f"pso_opt_{i}") - Path(f"results/pso_sor_{i}").write_text(f"pso_sor_{i}") poscar_dir = Path("poscar_dir") poscar_dir.mkdir(parents=True, exist_ok=True) for poscar in Path().glob("POSCAR_*"): target = poscar_dir.joinpath(poscar.name) shutil.copyfile(poscar, target) - finished = "true" if int(cnt_num) == int(max_step) else "false" + # finished = "true" if int(cnt_num) == int(max_step) else "false" os.chdir(cwd) ret_dict = { diff --git a/tests/op/test_prep_caly_model_devi.py b/tests/op/test_prep_caly_model_devi.py new file mode 100644 index 00000000..fae0a0b4 --- /dev/null +++ b/tests/op/test_prep_caly_model_devi.py @@ -0,0 +1,119 @@ +import os +import shutil +import unittest +from pathlib import ( + Path, +) + +import numpy as np +from dflow.python import ( + OP, + OPIO, + Artifact, + OPIOSign, + TransientError, +) +from mock import ( + call, + mock, + patch, +) + +# isort: off +from .context import ( + dpgen2, +) +from dpgen2.constants import ( + calypso_task_pattern, + model_name_pattern, + calypso_run_opt_file, + calypso_check_opt_file, +) +from dpgen2.op import PrepCalyModelDevi +from dpgen2.utils import ( + BinaryFileInput, +) + +# isort: on + + +class TestPrepCalyModelDevi(unittest.TestCase): + def setUp(self): + self.run_dir_name = "run_dir" + + self.ref_dir = Path("ref_dir") + self.ref_dir.mkdir(parents=True, exist_ok=True) + + self.ref_traj_results = [] + + ntrajs_dir = 5 + ntrajs_per_dir = 2 + for dir_index in range(ntrajs_dir): + dir_name = self.ref_dir.joinpath(f"traj_dir_{dir_index}") + dir_name.mkdir(parents=True, exist_ok=True) + self.ref_traj_results.append(dir_name) + for traj_index in range(ntrajs_per_dir): + dir_name.joinpath(f"{dir_index}.{traj_index}.traj").write_text( + f"trajectory.{dir_index}.{traj_index}" + ) + + self.group_size = 5 + self.ngroup = ntrajs_dir * ntrajs_per_dir / self.group_size + self.template_slice_config_1 = {"model_devi_group_size": self.group_size} + self.template_slice_config_2 = {} + + def tearDown(self): + shutil.rmtree(self.ref_dir) + shutil.rmtree(self.run_dir_name) + + def test_00_success(self): + op = PrepCalyModelDevi() + out = op.execute( + OPIO( + { + "task_name": self.run_dir_name, + "template_slice_config": self.template_slice_config_1, + "traj_results": self.ref_traj_results, + } + ) + ) + # check output length + self.assertEqual(len(out["task_name_list"]), self.ngroup) + self.assertEqual(len(out["grouped_traj_list"]), self.ngroup) + # check filename + self.assertEqual(out["task_name_list"][0], "run_dir/trajs_part_0") + self.assertEqual(out["task_name_list"][1], "run_dir/trajs_part_1") + # check file exists + self.assertTrue(Path(out["grouped_traj_list"][0]).exists()) + self.assertTrue(Path(out["grouped_traj_list"][1]).exists()) + + traj_list = list(Path(out["grouped_traj_list"][0]).rglob("*traj")) + # check traj number + self.assertEqual(len(traj_list), 5) + # check traj file name + # self.assertTrue(Path("run_dir/trajs_part_0/0.0.0.traj") in traj_list) + + def test_01_success(self): + op = PrepCalyModelDevi() + out = op.execute( + OPIO( + { + "task_name": self.run_dir_name, + "template_slice_config": self.template_slice_config_2, + "traj_results": self.ref_traj_results, + } + ) + ) + # check output length + self.assertEqual(len(out["task_name_list"]), 1) + self.assertEqual(len(out["grouped_traj_list"]), 1) + # check filename + self.assertEqual(out["task_name_list"][0], "run_dir/trajs_part_0") + # check file exists + self.assertTrue(Path(out["grouped_traj_list"][0]).exists()) + + traj_list = list(Path(out["grouped_traj_list"][0]).rglob("*traj")) + # check traj number + self.assertEqual(len(traj_list), 10) + # check traj file name + # self.assertTrue(Path("run_dir/trajs_part_0/0.0.0.traj") in traj_list) diff --git a/tests/op/test_run_caly_model_devi.py b/tests/op/test_run_caly_model_devi.py index 6d45b403..9040814c 100644 --- a/tests/op/test_run_caly_model_devi.py +++ b/tests/op/test_run_caly_model_devi.py @@ -1,6 +1,9 @@ import os import shutil import unittest +from ast import ( + Slice, +) from pathlib import ( Path, ) @@ -17,11 +20,15 @@ from ase.io import ( write, ) +from dflow import ( + Step, +) from dflow.python import ( OP, OPIO, Artifact, OPIOSign, + PythonOPTemplate, TransientError, ) @@ -131,8 +138,6 @@ def test_01_atoms2lmpdump(self): dump_str = atoms2lmpdump(self.atoms_normal_2, 1, self.type_map) self.assertEqual(dump_str, self.ref_dump_str) - # @patch("dpgen2.op.run_caly_model_devi.RunCalyModelDevi.import_deepmd_package.calc_model_devi") - # @patch("dpgen2.op.run_caly_model_devi.RunCalyModelDevi.import_deepmd_package.DP") @unittest.skipIf(x == 1, "deepmd package not exists.") @patch("deepmd.infer.calc_model_devi") @patch("deepmd.infer.DeepPot") diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py index ee7d06a7..d555d418 100644 --- a/tests/test_prep_run_caly.py +++ b/tests/test_prep_run_caly.py @@ -75,6 +75,9 @@ from dpgen2.op.prep_caly_input import ( PrepCalyInput, ) +from dpgen2.op.prep_caly_model_devi import ( + PrepCalyModelDevi, +) from dpgen2.op.run_caly_model_devi import ( RunCalyModelDevi, ) @@ -101,6 +104,7 @@ "template_slice_config": { "group_size": 2, "pool_size": 1, + "model_devi_group_size": 30, }, } ) @@ -143,7 +147,6 @@ def tearDown(self): shutil.rmtree(self.work_dir, ignore_errors=True) for i in Path().glob("prep-run-caly-step*"): shutil.rmtree(i, ignore_errors=True) - # shutil.rmtree("upload", ignore_errors=True) def test(self): caly_evo_step_op = CalyEvoStep( @@ -159,6 +162,7 @@ def test(self): "prep-run-calypso", PrepCalyInput, caly_evo_step_op, + PrepCalyModelDevi, MockedRunCalyModelDevi, prep_config=prep_default_config, run_config=run_default_config,