diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a96c7d1d..6271a38c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,14 +17,11 @@ repos: # Python - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.1.3 + rev: v0.16.0 hooks: + - id: ruff-check + args: [--fix] - id: ruff-format -- repo: https://github.com/PyCQA/isort - rev: 5.12.0 - hooks: - - id: isort - files: \.py$ # numpydoc - repo: https://github.com/Carreau/velin rev: 0.0.12 diff --git a/docs/conf.py b/docs/conf.py index b68a1726..6abe342d 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -19,7 +19,7 @@ # -- Project information ----------------------------------------------------- project = "DPGEN2" -copyright = "2022-%d, DeepModeling" % date.today().year +copyright = f"2022-{date.today().year}, DeepModeling" author = "DeepModeling" diff --git a/docs/exploration.md b/docs/exploration.md index 9e6216ed..3647dd32 100644 --- a/docs/exploration.md +++ b/docs/exploration.md @@ -26,11 +26,11 @@ class StageScheduler(ABC): @abstractmethod def plan_next_iteration( - self, - hist_reports : List[ExplorationReport], - report : ExplorationReport, - confs : List[Path], - ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector] : + self, + hist_reports: List[ExplorationReport], + report: ExplorationReport, + confs: List[Path], + ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector]: """ Make the plan for the next iteration of the stage. @@ -66,7 +66,7 @@ One may check more details on the [exploratin task group](#exploration-task-grou DPGEN2 defines a python class `ExplorationTask` to manage all necessry files needed to run a exploration task. It can be used as the example provided in the doc string. ```python -class ExplorationTask(): +class ExplorationTask: """Define the files needed by an exploration task. Examples @@ -74,7 +74,7 @@ class ExplorationTask(): >>> # this example dumps all files needed by the task. >>> files = exploration_task.files() ... for file_name, file_content in files.items(): - ... with open(file_name, 'w') as fp: + ... with open(file_name, "w") as fp: ... fp.write(file_content) """ @@ -94,8 +94,8 @@ class ExplorationTaskGroup(Sequence): ... def add_group( - self, - group : 'ExplorationTaskGroup', + self, + group: "ExplorationTaskGroup", ): """Add another group to the group.""" ... @@ -167,7 +167,9 @@ Then we define the exploration arguments schema in the dpgen2 input configuratio ```python def xxx_task_group_args(): - doc_xxx_task_grp = "XXX exploration tasks. dpgen2 will generate the XXX input script" + doc_xxx_task_grp = ( + "XXX exploration tasks. dpgen2 will generate the XXX input script" + ) return Argument( "task_group", dict, @@ -177,12 +179,14 @@ def xxx_task_group_args(): doc=doc_xxx_task_grp, ) + def xxx_normalize(config): args = xxx_task_group_args() config = args.normalize_value(config, trim_pattern="_*") args.check_value(config, strict=False) return config + def make_xxx_task_group_from_config(config): config = xxx_normalize(config) tgroup = XXXTaskGroup(**config) @@ -197,6 +201,7 @@ def xxx_args(): # Argument(...), ] + def variant_explore(): # ... doc_xxx = "The exploration by XXX" diff --git a/docs/operator.md b/docs/operator.md index 970330e8..b43237a9 100644 --- a/docs/operator.md +++ b/docs/operator.md @@ -92,31 +92,34 @@ The `prep_train` prepares a list of paths, each of which contains all necessary The `run_train` slices the list of paths, and assign each item in the list to a DeePMD-kit task. The task is executed by `run_train_op`. This is a very nice feature of `dflow`, because the developer only needs to implement how one DeePMD-kit task is executed, and then all the items in the task list will be executed [in parallel](https://github.com/dptech-corp/dflow/blob/master/README.md#315-produce-parallel-steps-using-loop). See the following code to see how it works ```python - run_train = Step( - 'run-train', - template=PythonOPTemplate( - run_train_op, - image=run_train_image, - slices = Slices( - "int('{{item}}')", - input_parameter = ["task_name"], - input_artifact = ["task_path", "init_model"], - output_artifact = ["model", "lcurve", "log", "script"], - ), +run_train = Step( + "run-train", + template=PythonOPTemplate( + run_train_op, + image=run_train_image, + slices=Slices( + "int('{{item}}')", + input_parameter=["task_name"], + input_artifact=["task_path", "init_model"], + output_artifact=["model", "lcurve", "log", "script"], ), - parameters={ - "config" : train_steps.inputs.parameters["train_config"], - "task_name" : prep_train.outputs.parameters["task_names"], - }, - artifacts={ - 'task_path' : prep_train.outputs.artifacts['task_paths'], - "init_model" : train_steps.inputs.artifacts['init_models'], - "init_data": train_steps.inputs.artifacts['init_data'], - "iter_data": train_steps.inputs.artifacts['iter_data'], - }, - with_sequence=argo_sequence(argo_len(prep_train.outputs.parameters["task_names"]), format=train_index_pattern), - key = step_keys['run-train'], - ) + ), + parameters={ + "config": train_steps.inputs.parameters["train_config"], + "task_name": prep_train.outputs.parameters["task_names"], + }, + artifacts={ + "task_path": prep_train.outputs.artifacts["task_paths"], + "init_model": train_steps.inputs.artifacts["init_models"], + "init_data": train_steps.inputs.artifacts["init_data"], + "iter_data": train_steps.inputs.artifacts["iter_data"], + }, + with_sequence=argo_sequence( + argo_len(prep_train.outputs.parameters["task_names"]), + format=train_index_pattern, + ), + key=step_keys["run-train"], +) ``` The input parameter `"task_names"` and artifacts `"task_paths"` and `"init_model"` are sliced and supplied to each DeePMD-kit task. The output artifacts of the tasks (`"model"`, `"lcurve"`, `"log"` and `"script"`) are stacked in the same order as the input lists. These lists are assigned as the outputs of `train_steps` by ```python @@ -144,26 +147,32 @@ from dflow.python import ( OPIOSign, Artifact, ) + + class RunDPTrain(OP): @classmethod def get_input_sign(cls): - return OPIOSign({ - "config" : dict, - "task_name" : str, - "task_path" : Artifact(Path), - "init_model" : Artifact(Path), - "init_data" : Artifact(List[Path]), - "iter_data" : Artifact(List[Path]), - }) + return OPIOSign( + { + "config": dict, + "task_name": str, + "task_path": Artifact(Path), + "init_model": Artifact(Path), + "init_data": Artifact(List[Path]), + "iter_data": Artifact(List[Path]), + } + ) @classmethod def get_output_sign(cls): - return OPIOSign({ - "script" : Artifact(Path), - "model" : Artifact(Path), - "lcurve" : Artifact(Path), - "log" : Artifact(Path), - }) + return OPIOSign( + { + "script": Artifact(Path), + "model": Artifact(Path), + "lcurve": Artifact(Path), + "log": Artifact(Path), + } + ) ``` All items not defined as `Artifact` are treated as parameters of the `OP`. The concept of parameter and artifact are explained in the [dflow document](https://github.com/dptech-corp/dflow/blob/master/README.md#Parametersandartifacts). To be short, the artifacts can be `pathlib.Path` or a list of `pathlib.Path`. The artifacts are passed by the file system. Other data structures are treated as parameters, they are passed as variables encoded in `str`. Therefore, a large amout of information should be stored in artifacts, otherwise they can be considered as parameters. @@ -171,47 +180,51 @@ All items not defined as `Artifact` are treated as parameters of the `OP`. The c The operation of the `OP` is implemented in method `execute`, and are run in docker containers. Again taking the `execute` method of `RunDPTrain` as an example ```python - @OP.exec_sign_check - def execute( - self, - ip : OPIO, - ) -> OPIO: - ... - task_name = ip['task_name'] - task_path = ip['task_path'] - init_model = ip['init_model'] - init_data = ip['init_data'] - iter_data = ip['iter_data'] - ... - work_dir = Path(task_name) - ... - # here copy all files in task_path to work_dir - ... - with set_directory(work_dir): - fplog = open('train.log', 'w') - def clean_before_quit(): - fplog.close() - # train model - command = ['dp', 'train', train_script_name] - ret, out, err = run_command(command) - if ret != 0: - clean_before_quit() - raise FatalError('dp train failed') - fplog.write(out) - # freeze model - ret, out, err = run_command(['dp', 'freeze', '-o', 'frozen_model.pb']) - if ret != 0: - clean_before_quit() - raise FatalError('dp freeze failed') - fplog.write(out) +@OP.exec_sign_check +def execute( + self, + ip: OPIO, +) -> OPIO: + ... + task_name = ip["task_name"] + task_path = ip["task_path"] + init_model = ip["init_model"] + init_data = ip["init_data"] + iter_data = ip["iter_data"] + ... + work_dir = Path(task_name) + ... + # here copy all files in task_path to work_dir + ... + with set_directory(work_dir): + fplog = open("train.log", "w") + + def clean_before_quit(): + fplog.close() + + # train model + command = ["dp", "train", train_script_name] + ret, out, err = run_command(command) + if ret != 0: clean_before_quit() - - return OPIO({ - "script" : work_dir / train_script_name, - "model" : work_dir / "frozen_model.pb", - "lcurve" : work_dir / "lcurve.out", - "log" : work_dir / "train.log", - }) + raise FatalError("dp train failed") + fplog.write(out) + # freeze model + ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"]) + if ret != 0: + clean_before_quit() + raise FatalError("dp freeze failed") + fplog.write(out) + clean_before_quit() + + return OPIO( + { + "script": work_dir / train_script_name, + "model": work_dir / "frozen_model.pb", + "lcurve": work_dir / "lcurve.out", + "log": work_dir / "train.log", + } + ) ``` The inputs and outputs variables are recorded in data structure `dflow.python.OPIO`, which is initialized by a Python dict. The keys in the input/output `dict`, and the types of the input/output variables will be checked against their signatures by decorator `OP.exec_sign_check`. If any key or type does not match, an exception will be raised. diff --git a/dpgen2/conf/alloy_conf.py b/dpgen2/conf/alloy_conf.py index a885e623..adddd2fe 100644 --- a/dpgen2/conf/alloy_conf.py +++ b/dpgen2/conf/alloy_conf.py @@ -4,9 +4,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, Union, ) @@ -53,9 +51,9 @@ class AlloyConfGenerator(ConfGenerator): def __init__( self, numb_confs, - lattice: Union[dpdata.System, Tuple[str, float]], - replicate: Union[List[int], Tuple[int, int, int], int, None] = None, - concentration: Union[List[List[float]], List[float], None] = None, + lattice: Union[dpdata.System, tuple[str, float]], + replicate: Union[list[int], tuple[int, int, int], int, None] = None, + concentration: Union[list[list[float]], list[float], None] = None, cell_pert_frac: float = 0.0, atom_pert_dist: float = 0.0, ): @@ -104,7 +102,7 @@ def doc() -> str: return f"Generate alloys with {make_link('a certain lattice or user proided structure', 'explore[lmp]/configurations[alloy]/lattice')}, the elements randomly occuping the lattice with {make_link('user provided probability', 'explore[lmp]/configurations[alloy]/concentration')} ." @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: from dpgen2.entrypoint.args import ( make_link, ) @@ -162,9 +160,9 @@ class AlloyConf: def __init__( self, - lattice: Union[dpdata.System, Tuple[str, float]], - type_map: List[str], - replicate: Union[List[int], Tuple[int, int, int], int, None] = None, + lattice: Union[dpdata.System, tuple[str, float]], + type_map: list[str], + replicate: Union[list[int], tuple[int, int, int], int, None] = None, ) -> None: # init sys if not isinstance(lattice, dpdata.System): @@ -184,18 +182,18 @@ def __init__( sys.data["atom_numbs"] = [0] * self.ntypes sys.data["atom_numbs"][0] = self.natoms sys.data["atom_types"] = np.array([0] * self.natoms, dtype=int) - self.type_population = [ii for ii in range(self.ntypes)] + self.type_population = list(range(self.ntypes)) # record sys self.sys = sys def generate_file_content( self, numb_confs, - concentration: Union[List[List[float]], List[float], None] = None, + concentration: Union[list[list[float]], list[float], None] = None, cell_pert_frac: float = 0.0, atom_pert_dist: float = 0.0, fmt: str = "lammps/lmp", - ) -> List[str]: + ) -> list[str]: """ Parameters ---------- @@ -232,10 +230,10 @@ def generate_file_content( def generate_systems( self, numb_confs, - concentration: Union[List[List[float]], List[float], None] = None, + concentration: Union[list[list[float]], list[float], None] = None, cell_pert_frac: float = 0.0, atom_pert_dist: float = 0.0, - ) -> List[dpdata.System]: + ) -> list[dpdata.System]: """ Parameters ---------- @@ -265,7 +263,7 @@ def generate_systems( def _generate_one_sys( self, - concentration: Union[List[List[float]], List[float], None] = None, + concentration: Union[list[list[float]], list[float], None] = None, cell_pert_frac: float = 0.0, atom_pert_dist: float = 0.0, ) -> dpdata.System: @@ -351,11 +349,11 @@ def gen_doc(*, make_anchor=True, make_link=True, **kwargs): def generate_alloy_conf_file_content( - lattice: Union[dpdata.System, Tuple[str, float]], - type_map: List[str], + lattice: Union[dpdata.System, tuple[str, float]], + type_map: list[str], numb_confs, - replicate: Union[List[int], Tuple[int, int, int], int, None] = None, - concentration: Union[List[List[float]], List[float], None] = None, + replicate: Union[list[int], tuple[int, int, int], int, None] = None, + concentration: Union[list[list[float]], list[float], None] = None, cell_pert_frac: float = 0.0, atom_pert_dist: float = 0.0, fmt: str = "lammps/lmp", diff --git a/dpgen2/conf/conf_generator.py b/dpgen2/conf/conf_generator.py index e5e3b7b0..a30f362e 100644 --- a/dpgen2/conf/conf_generator.py +++ b/dpgen2/conf/conf_generator.py @@ -6,10 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - Dict, - List, -) import dargs import dpdata @@ -40,8 +36,8 @@ def get_file_content( self, type_map, fmt="lammps/lmp", - ) -> List[str]: - r"""Get the file content of configurations + ) -> list[str]: + r"""Get the file content of configurations. Parameters ---------- @@ -67,15 +63,15 @@ def get_file_content( @staticmethod @abstractmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: pass @classmethod def normalize_config( cls, - data: Dict = {}, + data: dict = {}, strict: bool = True, - ) -> Dict: + ) -> dict: r"""Normalized the argument. Parameters diff --git a/dpgen2/conf/file_conf.py b/dpgen2/conf/file_conf.py index e8ae2443..ca2f80ce 100644 --- a/dpgen2/conf/file_conf.py +++ b/dpgen2/conf/file_conf.py @@ -4,9 +4,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, Union, ) @@ -24,7 +22,7 @@ class FileConfGenerator(ConfGenerator): def __init__( self, - files: Union[str, List[str]], + files: Union[str, list[str]], fmt: str = "auto", prefix: Optional[str] = None, remove_pbc: Optional[bool] = False, @@ -86,7 +84,7 @@ def doc() -> str: return "Generate alloys from user provided file(s). The file(s) are assume to be load by `dpdata`." @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: doc_files = "The paths to the configuration files. widecards are supported." doc_prefix = "The prefix of file paths." doc_fmt = "The format (dpdata accepted formats) of the files." diff --git a/dpgen2/conf/unit_cells.py b/dpgen2/conf/unit_cells.py index ad392023..dcc891bf 100644 --- a/dpgen2/conf/unit_cells.py +++ b/dpgen2/conf/unit_cells.py @@ -40,16 +40,16 @@ def gen_box(self): def poscar_unit(self, latt): box = self.gen_box() ret = "" - ret += "BCC : a = %f \n" % latt - ret += "%.16f\n" % (latt) - ret += "%.16f %.16f %.16f\n" % (box[0][0], box[0][1], box[0][2]) - ret += "%.16f %.16f %.16f\n" % (box[1][0], box[1][1], box[1][2]) - ret += "%.16f %.16f %.16f\n" % (box[2][0], box[2][1], box[2][2]) + ret += f"BCC : a = {latt:f} \n" + ret += f"{latt:.16f}\n" + ret += f"{box[0][0]:.16f} {box[0][1]:.16f} {box[0][2]:.16f}\n" + ret += f"{box[1][0]:.16f} {box[1][1]:.16f} {box[1][2]:.16f}\n" + ret += f"{box[2][0]:.16f} {box[2][1]:.16f} {box[2][2]:.16f}\n" ret += "Type\n" - ret += "%d\n" % self.numb_atoms() + ret += f"{self.numb_atoms():d}\n" ret += "Direct\n" - ret += "%.16f %.16f %.16f\n" % (0.0, 0.0, 0.0) - ret += "%.16f %.16f %.16f\n" % (0.5, 0.5, 0.5) + ret += f"{0.0:.16f} {0.0:.16f} {0.0:.16f}\n" + ret += f"{0.5:.16f} {0.5:.16f} {0.5:.16f}\n" return ret @@ -63,18 +63,18 @@ def gen_box(self): def poscar_unit(self, latt): box = self.gen_box() ret = "" - ret += "FCC : a = %f \n" % latt - ret += "%.16f\n" % (latt) - ret += "%.16f %.16f %.16f\n" % (box[0][0], box[0][1], box[0][2]) - ret += "%.16f %.16f %.16f\n" % (box[1][0], box[1][1], box[1][2]) - ret += "%.16f %.16f %.16f\n" % (box[2][0], box[2][1], box[2][2]) + ret += f"FCC : a = {latt:f} \n" + ret += f"{latt:.16f}\n" + ret += f"{box[0][0]:.16f} {box[0][1]:.16f} {box[0][2]:.16f}\n" + ret += f"{box[1][0]:.16f} {box[1][1]:.16f} {box[1][2]:.16f}\n" + ret += f"{box[2][0]:.16f} {box[2][1]:.16f} {box[2][2]:.16f}\n" ret += "Type\n" - ret += "%d\n" % self.numb_atoms() + ret += f"{self.numb_atoms():d}\n" ret += "Direct\n" - ret += "%.16f %.16f %.16f\n" % (0.0, 0.0, 0.0) - ret += "%.16f %.16f %.16f\n" % (0.5, 0.5, 0.0) - ret += "%.16f %.16f %.16f\n" % (0.5, 0.0, 0.5) - ret += "%.16f %.16f %.16f\n" % (0.0, 0.5, 0.5) + ret += f"{0.0:.16f} {0.0:.16f} {0.0:.16f}\n" + ret += f"{0.5:.16f} {0.5:.16f} {0.0:.16f}\n" + ret += f"{0.5:.16f} {0.0:.16f} {0.5:.16f}\n" + ret += f"{0.0:.16f} {0.5:.16f} {0.5:.16f}\n" return ret @@ -91,16 +91,16 @@ def gen_box(self): def poscar_unit(self, latt): box = self.gen_box() ret = "" - ret += "HCP : a = %f / sqrt(2)\n" % latt + ret += f"HCP : a = {latt:f} / sqrt(2)\n" ret += "%.16f\n" % (latt / np.sqrt(2)) - ret += "%.16f %.16f %.16f\n" % (box[0][0], box[0][1], box[0][2]) - ret += "%.16f %.16f %.16f\n" % (box[1][0], box[1][1], box[1][2]) - ret += "%.16f %.16f %.16f\n" % (box[2][0], box[2][1], box[2][2]) + ret += f"{box[0][0]:.16f} {box[0][1]:.16f} {box[0][2]:.16f}\n" + ret += f"{box[1][0]:.16f} {box[1][1]:.16f} {box[1][2]:.16f}\n" + ret += f"{box[2][0]:.16f} {box[2][1]:.16f} {box[2][2]:.16f}\n" ret += "Type\n" - ret += "%d\n" % self.numb_atoms() + ret += f"{self.numb_atoms():d}\n" ret += "Direct\n" - ret += "%.16f %.16f %.16f\n" % (0, 0, 0) - ret += "%.16f %.16f %.16f\n" % (1.0 / 3, 1.0 / 3, 1.0 / 2) + ret += f"{0:.16f} {0:.16f} {0:.16f}\n" + ret += f"{1.0 / 3:.16f} {1.0 / 3:.16f} {1.0 / 2:.16f}\n" return ret @@ -114,15 +114,15 @@ def gen_box(self): def poscar_unit(self, latt): box = self.gen_box() ret = "" - ret += "SC : a = %f \n" % latt - ret += "%.16f\n" % (latt) - ret += "%.16f %.16f %.16f\n" % (box[0][0], box[0][1], box[0][2]) - ret += "%.16f %.16f %.16f\n" % (box[1][0], box[1][1], box[1][2]) - ret += "%.16f %.16f %.16f\n" % (box[2][0], box[2][1], box[2][2]) + ret += f"SC : a = {latt:f} \n" + ret += f"{latt:.16f}\n" + ret += f"{box[0][0]:.16f} {box[0][1]:.16f} {box[0][2]:.16f}\n" + ret += f"{box[1][0]:.16f} {box[1][1]:.16f} {box[1][2]:.16f}\n" + ret += f"{box[2][0]:.16f} {box[2][1]:.16f} {box[2][2]:.16f}\n" ret += "Type\n" - ret += "%d\n" % self.numb_atoms() + ret += f"{self.numb_atoms():d}\n" ret += "Direct\n" - ret += "%.16f %.16f %.16f\n" % (0.0, 0.0, 0.0) + ret += f"{0.0:.16f} {0.0:.16f} {0.0:.16f}\n" return ret @@ -142,21 +142,17 @@ def poscar_unit(self, latt): box = self.gen_box() ret = "" ret += "DIAMOND\n" - ret += "%.16f\n" % (latt) - ret += "%.16f %.16f %.16f\n" % (box[0][0], box[0][1], box[0][2]) - ret += "%.16f %.16f %.16f\n" % (box[1][0], box[1][1], box[1][2]) - ret += "%.16f %.16f %.16f\n" % (box[2][0], box[2][1], box[2][2]) + ret += f"{latt:.16f}\n" + ret += f"{box[0][0]:.16f} {box[0][1]:.16f} {box[0][2]:.16f}\n" + ret += f"{box[1][0]:.16f} {box[1][1]:.16f} {box[1][2]:.16f}\n" + ret += f"{box[2][0]:.16f} {box[2][1]:.16f} {box[2][2]:.16f}\n" ret += "Type\n" - ret += "%d\n" % self.numb_atoms() + ret += f"{self.numb_atoms():d}\n" ret += "Direct\n" - ret += "%.16f %.16f %.16f\n" % ( - 0.12500000000000, - 0.12500000000000, - 0.12500000000000, + ret += ( + f"{0.12500000000000:.16f} {0.12500000000000:.16f} {0.12500000000000:.16f}\n" ) - ret += "%.16f %.16f %.16f\n" % ( - 0.87500000000000, - 0.87500000000000, - 0.87500000000000, + ret += ( + f"{0.87500000000000:.16f} {0.87500000000000:.16f} {0.87500000000000:.16f}\n" ) return ret diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..dd6e8e91 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -1,7 +1,4 @@ import textwrap -from typing import ( - List, -) import dargs from dargs import ( @@ -61,7 +58,7 @@ def dp_dist_train_args(): doc=doc_config, ), Argument( - "template_script", [List[str], str], optional=False, doc=doc_template_script + "template_script", [list[str], str], optional=False, doc=doc_template_script ), Argument("student_model_path", str, optional=True, doc=dock_student_model_path), Argument( @@ -100,11 +97,11 @@ def dp_train_args(): ), Argument("numb_models", int, optional=True, default=4, doc=doc_numb_models), Argument( - "template_script", [List[str], str], optional=False, doc=doc_template_script + "template_script", [list[str], str], optional=False, doc=doc_template_script ), Argument( "init_models_paths", - List[str], + list[str], optional=True, default=None, doc=doc_init_models_paths, @@ -186,7 +183,7 @@ def variant_filter(): kk, dict, conf_filter_styles[kk].args(), - doc="Configuration filter of type %s" % kk, + doc=f"Configuration filter of type {kk}", ) ) return Variant( @@ -249,7 +246,7 @@ def lmp_args(): doc=doc_configuration, alias=["configuration"], ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list[list[dict]], optional=False, doc=doc_stages), Argument( "filters", list, @@ -342,7 +339,7 @@ def caly_args(): doc=doc_configuration, alias=["configuration"], ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list[list[dict]], optional=False, doc=doc_stages), Argument( "filters", list, @@ -432,7 +429,7 @@ def diffcsp_args(): optional=False, doc=doc_convergence, ), - Argument("stages", List[List[dict]], optional=False, doc=doc_stages), + Argument("stages", list[list[dict]], optional=False, doc=doc_stages), Argument( "filters", list, @@ -545,8 +542,8 @@ def input_args(): doc_multi_valid_data_uri = "The URI of validation data for multitask" return [ - Argument("type_map", List[str], optional=False, doc=doc_type_map), - Argument("mass_map", List[float], optional=False, doc=doc_mass_map), + Argument("type_map", list[str], optional=False, doc=doc_type_map), + Argument("mass_map", list[float], optional=False, doc=doc_mass_map), Argument( "init_data_prefix", str, @@ -560,7 +557,7 @@ def input_args(): ), Argument( "init_data_sys", - [List[str], str], + [list[str], str], optional=True, default=None, doc=doc_init_sys, @@ -609,7 +606,7 @@ def input_args(): ), Argument( "valid_data_sys", - [List[str], str], + [list[str], str], optional=True, default=None, doc=doc_valid_sys, @@ -840,7 +837,7 @@ def submit_args(default_step_config=normalize_step_dict({})): ), Argument( "upload_python_packages", - [List[str], str], + [list[str], str], optional=True, default=None, doc=doc_upload_python_packages, diff --git a/dpgen2/entrypoint/common.py b/dpgen2/entrypoint/common.py index 0d0af9e8..7a954e43 100644 --- a/dpgen2/entrypoint/common.py +++ b/dpgen2/entrypoint/common.py @@ -3,8 +3,6 @@ Path, ) from typing import ( - Dict, - List, Optional, Union, ) @@ -38,7 +36,7 @@ def global_config_workflow( bohrium_config_from_dict(wf_config["bohrium_config"]) -def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: +def expand_sys_str(root_dir: Union[str, Path]) -> list[str]: root_dir = Path(root_dir) matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()] if (root_dir / "type.raw").is_file(): @@ -46,7 +44,7 @@ def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: return matches -def expand_idx(in_list) -> List[int]: +def expand_idx(in_list) -> list[int]: ret = [] for ii in in_list: if isinstance(ii, int): @@ -64,5 +62,5 @@ def expand_idx(in_list) -> List[int]: ret += [int(range_str[0])] else: raise RuntimeError("not expected range string", step_str[0]) - ret = sorted(list(set(ret))) + ret = sorted(set(ret)) return ret diff --git a/dpgen2/entrypoint/download.py b/dpgen2/entrypoint/download.py index 7f095039..3c4364e0 100644 --- a/dpgen2/entrypoint/download.py +++ b/dpgen2/entrypoint/download.py @@ -1,7 +1,5 @@ import logging from typing import ( - Dict, - List, Optional, Union, ) @@ -25,9 +23,9 @@ def download_by_def( workflow_id, - wf_config: Dict = {}, - iterations: Optional[List[int]] = None, - step_defs: Optional[List[str]] = None, + wf_config: dict = {}, + iterations: Optional[list[int]] = None, + step_defs: Optional[list[str]] = None, prefix: Optional[str] = None, chk_pnt: bool = False, ): @@ -42,8 +40,8 @@ def download_by_def( def download( workflow_id, - wf_config: Optional[Dict] = {}, - wf_keys: Optional[List] = None, + wf_config: Optional[dict] = {}, + wf_keys: Optional[list] = None, prefix: Optional[str] = None, chk_pnt: bool = False, ): diff --git a/dpgen2/entrypoint/main.py b/dpgen2/entrypoint/main.py index 00ee55e3..744aa61f 100644 --- a/dpgen2/entrypoint/main.py +++ b/dpgen2/entrypoint/main.py @@ -3,7 +3,6 @@ import logging import textwrap from typing import ( - List, Optional, ) @@ -287,13 +286,13 @@ def main_parser() -> argparse.ArgumentParser: "-v", "--version", action="version", - version="DPGEN v%s" % __version__, + version=f"DPGEN v{__version__}", ) return parser -def parse_args(args: Optional[List[str]] = None): +def parse_args(args: Optional[list[str]] = None): """DPGEN2 commandline options argument parsing. Parameters @@ -358,7 +357,7 @@ def main(): config = json.load(fp) wfid = args.ID if args.list_supported is not None and args.list_supported: - print(print_op_download_setting()) + pass elif args.keys is not None: download( wfid, diff --git a/dpgen2/entrypoint/showkey.py b/dpgen2/entrypoint/showkey.py index 498533bf..be9b1913 100644 --- a/dpgen2/entrypoint/showkey.py +++ b/dpgen2/entrypoint/showkey.py @@ -29,4 +29,3 @@ def showkey( all_step_keys, ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], ) - print(prt_str) diff --git a/dpgen2/entrypoint/status.py b/dpgen2/entrypoint/status.py index 1698f518..d90efda2 100644 --- a/dpgen2/entrypoint/status.py +++ b/dpgen2/entrypoint/status.py @@ -1,6 +1,5 @@ import logging from typing import ( - Dict, Optional, ) @@ -20,7 +19,7 @@ def status( workflow_id, - wf_config: Optional[Dict] = {}, + wf_config: Optional[dict] = {}, ): wf_config = normalize_args(wf_config) @@ -34,6 +33,5 @@ def status( if scheduler is not None: ptr_str = scheduler.print_convergence() - print(ptr_str) else: logging.warn("no scheduler is finished") diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 9e07374f..536f83a4 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -9,8 +9,6 @@ Path, ) from typing import ( - Dict, - List, Optional, ) @@ -159,9 +157,9 @@ def make_concurrent_learning_op( select_confs_config: dict = default_config, collect_data_config: dict = default_config, cl_step_config: dict = default_config, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, valid_data: Optional[S3Artifact] = None, - train_optional_files: Optional[List[str]] = None, + train_optional_files: Optional[list[str]] = None, explore_config: Optional[dict] = None, ): if train_style in ("dp", "dp-dist"): @@ -463,7 +461,7 @@ def get_systems_from_data(data, data_prefix=None): def workflow_concurrent_learning( - config: Dict, + config: dict, ) -> Step: default_config = config["default_step_config"] @@ -505,7 +503,7 @@ def workflow_concurrent_learning( if upload_python_packages is not None and isinstance(upload_python_packages, str): upload_python_packages = [upload_python_packages] if upload_python_packages is not None: - _upload_python_packages: List[os.PathLike] = [ + _upload_python_packages: list[os.PathLike] = [ Path(ii) for ii in upload_python_packages ] upload_python_packages = _upload_python_packages @@ -564,9 +562,9 @@ def workflow_concurrent_learning( "teacher_model_path" in explore_config and explore_config["teacher_model_path"] is not None ): - assert os.path.exists( - explore_config["teacher_model_path"] - ), f"No such file: {explore_config['teacher_model_path']}" + assert os.path.exists(explore_config["teacher_model_path"]), ( + f"No such file: {explore_config['teacher_model_path']}" + ) explore_config["teacher_model_path"] = BinaryFileInput( explore_config["teacher_model_path"] ) @@ -579,12 +577,12 @@ def workflow_concurrent_learning( fp_config["run"] = config["fp"]["run_config"] fp_config["extra_output_files"] = config["fp"]["extra_output_files"] if fp_style == "deepmd": - assert ( - "teacher_model_path" in fp_config["run"] - ), "Cannot find 'teacher_model_path' in config['fp']['run_config'] when fp_style == 'deepmd'" - assert os.path.exists( - fp_config["run"]["teacher_model_path"] - ), f"No such file: {fp_config['run']['teacher_model_path']}" + assert "teacher_model_path" in fp_config["run"], ( + "Cannot find 'teacher_model_path' in config['fp']['run_config'] when fp_style == 'deepmd'" + ) + assert os.path.exists(fp_config["run"]["teacher_model_path"]), ( + f"No such file: {fp_config['run']['teacher_model_path']}" + ) fp_config["run"]["teacher_model_path"] = BinaryFileInput( fp_config["run"]["teacher_model_path"] ) @@ -670,13 +668,13 @@ def get_scheduler_ids( if get_subkey(ii.key, 1) == "scheduler": scheduler_ids.append(idx) scheduler_keys = [reuse_step[ii].key for ii in scheduler_ids] - assert ( - sorted(scheduler_keys) == scheduler_keys - ), "The scheduler keys are not properly sorted" + assert sorted(scheduler_keys) == scheduler_keys, ( + "The scheduler keys are not properly sorted" + ) if len(scheduler_ids) == 0: logging.warning( - "No scheduler found in the workflow, " "does not do any replacement." + "No scheduler found in the workflow, does not do any replacement." ) return scheduler_ids @@ -735,7 +733,7 @@ def copy_scheduler_plans( def submit_concurrent_learning( wf_config, - reuse_step: Optional[List[ArgoStep]] = None, + reuse_step: Optional[list[ArgoStep]] = None, replace_scheduler: bool = False, no_submission: bool = False, ): @@ -867,7 +865,6 @@ def resubmit_concurrent_learning( all_step_keys, ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], ) - print(prt_str) if reuse is None: return None diff --git a/dpgen2/entrypoint/watch.py b/dpgen2/entrypoint/watch.py index e8654039..69044a4d 100644 --- a/dpgen2/entrypoint/watch.py +++ b/dpgen2/entrypoint/watch.py @@ -1,8 +1,6 @@ import logging import time from typing import ( - Dict, - List, Optional, ) @@ -31,9 +29,9 @@ def update_finished_steps( wf, - finished_keys: Optional[List[str]] = None, + finished_keys: Optional[list[str]] = None, download: Optional[bool] = False, - watching_keys: Optional[List[str]] = None, + watching_keys: Optional[list[str]] = None, prefix: Optional[str] = None, chk_pnt: bool = False, ): @@ -58,7 +56,7 @@ def update_finished_steps( def watch( workflow_id, wf_config: Optional[dict] = None, - watching_keys: Optional[List] = default_watching_keys, + watching_keys: Optional[list] = default_watching_keys, frequency: float = 600.0, download: bool = False, prefix: Optional[str] = None, diff --git a/dpgen2/exploration/deviation/deviation_manager.py b/dpgen2/exploration/deviation/deviation_manager.py index cbc227c5..ff9c1763 100644 --- a/dpgen2/exploration/deviation/deviation_manager.py +++ b/dpgen2/exploration/deviation/deviation_manager.py @@ -3,7 +3,6 @@ abstractmethod, ) from typing import ( - List, Optional, ) @@ -55,7 +54,7 @@ def add(self, name: str, deviation: np.ndarray) -> None: def _add(self, name: str, deviation: np.ndarray) -> None: pass - def get(self, name: str) -> List[Optional[np.ndarray]]: + def get(self, name: str) -> list[Optional[np.ndarray]]: r"""Gat a model deviation from this manager. Parameters @@ -71,7 +70,7 @@ def get(self, name: str) -> List[Optional[np.ndarray]]: return self._get(name) @abstractmethod - def _get(self, name: str) -> List[Optional[np.ndarray]]: + def _get(self, name: str) -> list[Optional[np.ndarray]]: pass @abstractmethod @@ -81,5 +80,5 @@ def clear(self) -> None: @abstractmethod def _check_data(self) -> None: - r"""Check if data is valid""" + r"""Check if data is valid.""" pass diff --git a/dpgen2/exploration/deviation/deviation_std.py b/dpgen2/exploration/deviation/deviation_std.py index f927b75b..54be0c02 100644 --- a/dpgen2/exploration/deviation/deviation_std.py +++ b/dpgen2/exploration/deviation/deviation_std.py @@ -2,8 +2,6 @@ defaultdict, ) from typing import ( - Dict, - List, Optional, ) @@ -33,9 +31,9 @@ def __init__(self): self._data = defaultdict(list) def _add(self, name: str, deviation: np.ndarray) -> None: - assert isinstance( - deviation, np.ndarray - ), f"Error: deviation(type: {type(deviation)}) is not a np.ndarray" + assert isinstance(deviation, np.ndarray), ( + f"Error: deviation(type: {type(deviation)}) is not a np.ndarray" + ) assert len(deviation.shape) == 1, ( f"Error: deviation(shape: {deviation.shape}) is not a " + f"one-dimensional array" @@ -44,7 +42,7 @@ def _add(self, name: str, deviation: np.ndarray) -> None: self._data[name].append(deviation) self.ntraj = max(self.ntraj, len(self._data[name])) - def _get(self, name: str) -> List[Optional[np.ndarray]]: + def _get(self, name: str) -> list[Optional[np.ndarray]]: if self.ntraj == 0: return [] elif len(self._data[name]) == 0: @@ -57,7 +55,7 @@ def clear(self) -> None: return None def _check_data(self) -> None: - r"""Check if data is valid""" + r"""Check if data is valid.""" model_devi_names = ( DeviManager.MAX_DEVI_V, DeviManager.MIN_DEVI_V, @@ -86,9 +84,9 @@ def _check_data(self) -> None: frames.pop(name) # check if "max_devi_f" exists - assert ( - len(self._data[DeviManager.MAX_DEVI_F]) == self.ntraj - ), f"Error: cannot find model deviation {DeviManager.MAX_DEVI_F}" + assert len(self._data[DeviManager.MAX_DEVI_F]) == self.ntraj, ( + f"Error: cannot find model deviation {DeviManager.MAX_DEVI_F}" + ) # check if the length of the arrays corresponding to the same # trajectory has the same number of frames diff --git a/dpgen2/exploration/render/traj_render.py b/dpgen2/exploration/render/traj_render.py index 5c9f0c41..dd6e5ccd 100644 --- a/dpgen2/exploration/render/traj_render.py +++ b/dpgen2/exploration/render/traj_render.py @@ -7,9 +7,7 @@ ) from typing import ( TYPE_CHECKING, - List, Optional, - Tuple, Union, ) @@ -33,7 +31,7 @@ class TrajRender(ABC): @abstractmethod def get_model_devi( self, - files: Union[List[Path], List[HDF5Dataset]], + files: Union[list[Path], list[HDF5Dataset]], ) -> DeviManager: r"""Get model deviations from recording files. @@ -51,11 +49,11 @@ def get_model_devi( @abstractmethod def get_confs( self, - traj: Union[List[Path], List[HDF5Dataset]], - id_selected: List[List[int]], - type_map: Optional[List[str]] = None, + traj: Union[list[Path], list[HDF5Dataset]], + id_selected: list[list[int]], + type_map: Optional[list[str]] = None, conf_filters: Optional["ConfFilters"] = None, - optional_outputs: Optional[List[Path]] = None, + optional_outputs: Optional[list[Path]] = None, ) -> dpdata.MultiSystems: r"""Get configurations from trajectory by selection. diff --git a/dpgen2/exploration/render/traj_render_lammps.py b/dpgen2/exploration/render/traj_render_lammps.py index 00b6a3de..aea8b8e8 100644 --- a/dpgen2/exploration/render/traj_render_lammps.py +++ b/dpgen2/exploration/render/traj_render_lammps.py @@ -7,9 +7,7 @@ ) from typing import ( TYPE_CHECKING, - List, Optional, - Tuple, Union, ) @@ -48,7 +46,7 @@ def __init__( def get_model_devi( self, - files: Union[List[Path], List[HDF5Dataset]], + files: Union[list[Path], list[HDF5Dataset]], ) -> DeviManager: ntraj = len(files) @@ -78,7 +76,7 @@ def _load_one_model_devi(self, fname, model_devi): def get_ele_temp(self, optional_outputs): ele_temp = [] for ii in range(len(optional_outputs)): - with open(optional_outputs[ii], "r") as f: + with open(optional_outputs[ii]) as f: data = json.load(f) if self.use_ele_temp: ele_temp.append(data["ele_temp"]) @@ -103,11 +101,11 @@ def set_ele_temp(self, system, ele_temp): def get_confs( self, - trajs: Union[List[Path], List[HDF5Dataset]], - id_selected: List[List[int]], - type_map: Optional[List[str]] = None, + trajs: Union[list[Path], list[HDF5Dataset]], + id_selected: list[list[int]], + type_map: Optional[list[str]] = None, conf_filters: Optional["ConfFilters"] = None, - optional_outputs: Optional[List[Path]] = None, + optional_outputs: Optional[list[Path]] = None, ) -> dpdata.MultiSystems: ntraj = len(trajs) ele_temp = None diff --git a/dpgen2/exploration/report/report.py b/dpgen2/exploration/report/report.py index d1c43fe6..132acf0a 100644 --- a/dpgen2/exploration/report/report.py +++ b/dpgen2/exploration/report/report.py @@ -3,9 +3,7 @@ abstractmethod, ) from typing import ( - List, Optional, - Tuple, ) import numpy as np @@ -18,7 +16,7 @@ class ExplorationReport(ABC): @abstractmethod def clear(self): - r"""Clear the report""" + r"""Clear the report.""" pass @abstractmethod @@ -26,7 +24,7 @@ def record( self, model_devi: DeviManager, ): - r"""Record the model deviations of the trajectories + r"""Record the model deviations of the trajectories. Parameters ---------- @@ -62,15 +60,15 @@ def converged( @abstractmethod def no_candidate(self) -> bool: - r"""If no candidate configuration is found""" + r"""If no candidate configuration is found.""" pass @abstractmethod def get_candidate_ids( self, max_nframes: Optional[int] = None, - ) -> List[List[int]]: - r"""Get indexes of candidate configurations + ) -> list[list[int]]: + r"""Get indexes of candidate configurations. Parameters ---------- @@ -88,7 +86,7 @@ def get_candidate_ids( @abstractmethod def print_header(self) -> str: - r"""Print the header of report""" + r"""Print the header of report.""" pass @abstractmethod @@ -98,5 +96,5 @@ def print( idx_in_stage: int, iter_idx: int, ) -> str: - r"""Print the report""" + r"""Print the report.""" pass diff --git a/dpgen2/exploration/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index cd2989f3..c2d2e0b2 100644 --- a/dpgen2/exploration/report/report_adaptive_lower.py +++ b/dpgen2/exploration/report/report_adaptive_lower.py @@ -1,9 +1,7 @@ import random import sys from typing import ( - List, Optional, - Tuple, ) import numpy as np @@ -154,7 +152,7 @@ def make_class_doc_link(key): 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. " @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: doc_level_f_hi = "The higher trust level of force model deviation" doc_numb_candi_f = "The number of force frames that has a model deviation lower than `level_f_hi` treated as candidate." doc_rate_candi_f = "The ratio of force frames that has a model deviation lower than `level_f_hi` treated as candidate." @@ -375,7 +373,7 @@ def get_candidate_ids( self, max_nframes: Optional[int] = None, clear: bool = True, - ) -> List[List[int]]: + ) -> list[list[int]]: ntraj = self.ntraj id_cand = self._get_candidates(max_nframes) id_cand_list = [[] for ii in range(ntraj)] @@ -389,7 +387,7 @@ def get_candidate_ids( def _get_candidates( self, max_nframes: Optional[int] = None, - ) -> List[Tuple[int, int]]: + ) -> list[tuple[int, int]]: if self.candi_sel_prob == "uniform": return self._get_candidates_uniform(max_nframes) elif self.candi_sel_prob == "inv_pop_f": @@ -400,7 +398,7 @@ def _get_candidates( def _get_candidates_uniform( self, max_nframes: Optional[int] = None, - ) -> List[Tuple[int, int]]: + ) -> list[tuple[int, int]]: """ Get candidates. If number of candidates is larger than `max_nframes`, then randomly pick `max_nframes` frames from the candidates. @@ -426,7 +424,7 @@ def _get_candidates_uniform( def _get_candidates_inv_pop_f( self, max_nframes: Optional[int] = None, - ) -> List[Tuple[int, int]]: + ) -> list[tuple[int, int]]: """ Get candidates. If number of candidates is larger than `max_nframes`, then randomly pick `max_nframes` frames from the candidates. @@ -459,7 +457,7 @@ def _get_candidates_inv_pop_f( def _choice_prob_inv_pop_f( self, - candi: List, + candi: list, ): """Compute the probability of candi frames according to the inverse population in the model deviation statistics. @@ -492,9 +490,7 @@ def _histo_idx( self, devi_f: float, ) -> int: - """ - return the index in histogram given a force model deviation. - """ + """Return the index in histogram given a force model deviation.""" dh = (self.level_f_hi - self.level_f_lo) / self.nhist hist_idx = int((devi_f - self.level_f_lo) / dh) if hist_idx < 0: @@ -504,7 +500,7 @@ def _histo_idx( return hist_idx def print_header(self) -> str: - r"""Print the header of report""" + r"""Print the header of report.""" return self.header_str def print( @@ -513,7 +509,7 @@ def print( idx_in_stage: int, iter_idx: int, ) -> str: - r"""Print the report""" + r"""Print the report.""" fmt_str = self.fmt_str fmt_flt = self.fmt_flt print_tuple = ( diff --git a/dpgen2/exploration/report/report_trust_levels_base.py b/dpgen2/exploration/report/report_trust_levels_base.py index 185ea5b1..3f23702c 100644 --- a/dpgen2/exploration/report/report_trust_levels_base.py +++ b/dpgen2/exploration/report/report_trust_levels_base.py @@ -3,9 +3,7 @@ abstractmethod, ) from typing import ( - List, Optional, - Tuple, ) import numpy as np @@ -70,7 +68,7 @@ def __init__( self._candidate_ratio = None @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: doc_level_f_lo = "The lower trust level of force model deviation" doc_level_f_hi = "The higher trust level of force model deviation" doc_level_v_lo = "The lower trust level of virial model deviation" @@ -171,10 +169,7 @@ def _record_one_traj( id_v_cand, id_v_fail, ): - """ - Record one trajctory. inputs are the indexes of candidate, accurate and failed frames. - - """ + """Record one trajctory. inputs are the indexes of candidate, accurate and failed frames.""" # check consistency novirial = id_v_cand is None if novirial: @@ -190,9 +185,9 @@ def _record_one_traj( 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_cand = set([]) if novirial else set(id_v_cand) - set_v_fail = set([]) if novirial else set(id_v_fail) + set_v_accu = set(range(nframes)) 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) # accu, cand, fail set_accu = set_f_accu & set_v_accu set_cand = ( @@ -211,7 +206,7 @@ def _record_one_traj( @abstractmethod def converged( self, - reports: Optional[List[ExplorationReport]] = None, + reports: Optional[list[ExplorationReport]] = None, ) -> bool: pass @@ -240,11 +235,11 @@ def no_candidate(self) -> bool: def get_candidate_ids( self, max_nframes: Optional[int] = None, - ) -> List[List[int]]: + ) -> list[list[int]]: pass def print_header(self) -> str: - r"""Print the header of report""" + r"""Print the header of report.""" return self.header_str def print( @@ -253,7 +248,7 @@ def print( idx_in_stage: int, iter_idx: int, ) -> str: - r"""Print the report""" + r"""Print the report.""" fmt_str = self.fmt_str fmt_flt = self.fmt_flt print_tuple = ( diff --git a/dpgen2/exploration/report/report_trust_levels_max.py b/dpgen2/exploration/report/report_trust_levels_max.py index e847c1e6..76fe0b06 100644 --- a/dpgen2/exploration/report/report_trust_levels_max.py +++ b/dpgen2/exploration/report/report_trust_levels_max.py @@ -1,8 +1,6 @@ import random from typing import ( - List, Optional, - Tuple, ) import numpy as np @@ -27,7 +25,7 @@ class ExplorationReportTrustLevelsMax(ExplorationReportTrustLevels): def converged( self, - reports: Optional[List[ExplorationReport]] = None, + reports: Optional[list[ExplorationReport]] = None, ) -> bool: r"""Check if the exploration is converged. @@ -49,7 +47,7 @@ def get_candidate_ids( self, max_nframes: Optional[int] = None, clear: bool = True, - ) -> List[List[int]]: + ) -> list[list[int]]: ntraj = len(self.traj_nframes) id_cand = self._get_candidates(max_nframes) id_cand_list = [[] for ii in range(ntraj)] @@ -63,7 +61,7 @@ def get_candidate_ids( def _get_candidates( self, max_nframes: Optional[int] = None, - ) -> List[Tuple[int, int]]: + ) -> list[tuple[int, int]]: """ Get candidates. If number of candidates is larger than `max_nframes`, then select `max_nframes` frames with the largest `max_devi_f` from diff --git a/dpgen2/exploration/report/report_trust_levels_random.py b/dpgen2/exploration/report/report_trust_levels_random.py index 540ad0d2..6e725465 100644 --- a/dpgen2/exploration/report/report_trust_levels_random.py +++ b/dpgen2/exploration/report/report_trust_levels_random.py @@ -1,8 +1,6 @@ import random from typing import ( - List, Optional, - Tuple, ) import numpy as np @@ -27,7 +25,7 @@ class ExplorationReportTrustLevelsRandom(ExplorationReportTrustLevels): def converged( self, - reports: Optional[List[ExplorationReport]] = None, + reports: Optional[list[ExplorationReport]] = None, ) -> bool: r"""Check if the exploration is converged. @@ -49,7 +47,7 @@ def get_candidate_ids( self, max_nframes: Optional[int] = None, clear: bool = True, - ) -> List[List[int]]: + ) -> list[list[int]]: ntraj = len(self.traj_nframes) id_cand = self._get_candidates(max_nframes) id_cand_list = [[] for ii in range(ntraj)] @@ -63,7 +61,7 @@ def get_candidate_ids( def _get_candidates( self, max_nframes: Optional[int] = None, - ) -> List[Tuple[int, int]]: + ) -> list[tuple[int, int]]: """ Get candidates. If number of candidates is larger than `max_nframes`, then randomly pick `max_nframes` frames from the candidates. diff --git a/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py b/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py index 8ab8662f..0922caec 100644 --- a/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py +++ b/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py @@ -2,9 +2,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, Union, ) @@ -71,8 +69,8 @@ def reached_max_iteration(self): def plan_next_iteration( self, report: Optional[ExplorationReport] = None, - trajs: Optional[Union[List[Path], List[HDF5Dataset]]] = None, - ) -> Tuple[bool, Optional[BaseExplorationTaskGroup], Optional[ConfSelector]]: + trajs: Optional[Union[list[Path], list[HDF5Dataset]]] = None, + ) -> tuple[bool, Optional[BaseExplorationTaskGroup], Optional[ConfSelector]]: if self.complete(): raise FatalError("Cannot plan because the stage has completed.") if report is None: diff --git a/dpgen2/exploration/scheduler/scheduler.py b/dpgen2/exploration/scheduler/scheduler.py index ff55fa23..7bdd3920 100644 --- a/dpgen2/exploration/scheduler/scheduler.py +++ b/dpgen2/exploration/scheduler/scheduler.py @@ -2,9 +2,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, Union, ) @@ -33,10 +31,7 @@ class ExplorationScheduler: - """ - The exploration scheduler. - - """ + """The exploration scheduler.""" def __init__( self, @@ -91,17 +86,11 @@ def get_iteration(self): return tot_iter def complete(self): - """ - Tell if all stages are converged. - - """ + """Tell if all stages are converged.""" return self.complete_ def force_stage_complete(self): - """ - Force complete the current stage - - """ + """Force complete the current stage.""" self.stage_schedulers[self.cur_stage].force_complete() self.cur_stage += 1 if self.cur_stage < len(self.stage_schedulers): @@ -114,8 +103,8 @@ def force_stage_complete(self): def plan_next_iteration( self, report: Optional[ExplorationReport] = None, - trajs: Optional[Union[List[Path], List[HDF5Dataset]]] = None, - ) -> Tuple[bool, Optional[ExplorationTaskGroup], Optional[ConfSelector]]: + trajs: Optional[Union[list[Path], list[HDF5Dataset]]] = None, + ) -> tuple[bool, Optional[ExplorationTaskGroup], Optional[ConfSelector]]: """ Make the plan for the next DPGEN iteration. @@ -136,7 +125,6 @@ def plan_next_iteration( The configuration selector for the next iteration. Should be `None` if converged. """ - try: stg_complete, expl_task_grp, conf_selector = self.stage_schedulers[ self.cur_stage @@ -145,7 +133,7 @@ def plan_next_iteration( trajs, ) except FatalError as e: - raise FatalError(f"stage {self.cur_stage}: " + str(e)) + raise FatalError(f"stage {self.cur_stage}: {e}") from e if stg_complete: self.cur_stage += 1 @@ -164,10 +152,7 @@ def plan_next_iteration( return stg_complete, expl_task_grp, conf_selector def get_stage_of_iterations(self): - """ - Get the stage index and the index in the stage of iterations. - - """ + """Get the stage index and the index in the stage of iterations.""" stages = self.stage_schedulers n_stage_iters = [] for ii in range(self.get_stage() + 1): @@ -196,7 +181,7 @@ def get_stage_of_iterations(self): def get_convergence_ratio(self): """ - Get the accurate, candidate and failed ratios of the iterations + Get the accurate, candidate and failed ratios of the iterations. Returns ------- diff --git a/dpgen2/exploration/scheduler/stage_scheduler.py b/dpgen2/exploration/scheduler/stage_scheduler.py index 18fe5593..0015f06c 100644 --- a/dpgen2/exploration/scheduler/stage_scheduler.py +++ b/dpgen2/exploration/scheduler/stage_scheduler.py @@ -6,8 +6,6 @@ Path, ) from typing import ( - List, - Tuple, Union, ) @@ -28,14 +26,12 @@ class StageScheduler(ABC): - """ - The scheduler for an exploration stage. - """ + """The scheduler for an exploration stage.""" @abstractmethod def converged(self) -> bool: """ - Tell if the stage is converged + Tell if the stage is converged. Returns ------- @@ -47,7 +43,7 @@ def converged(self) -> bool: @abstractmethod def complete(self) -> bool: """ - Tell if the stage is complete + Tell if the stage is complete. Returns ------- @@ -58,16 +54,13 @@ def complete(self) -> bool: @abstractmethod def force_complete(self): - """ - For complete the stage - - """ + """For complete the stage.""" pass @abstractmethod def next_iteration(self) -> int: """ - Return the index of the next iteration + Return the index of the next iteration. Returns ------- @@ -77,9 +70,9 @@ def next_iteration(self) -> int: pass @abstractmethod - def get_reports(self) -> List[ExplorationReport]: + def get_reports(self) -> list[ExplorationReport]: """ - Return all exploration reports + Return all exploration reports. Returns ------- @@ -92,8 +85,8 @@ def get_reports(self) -> List[ExplorationReport]: def plan_next_iteration( self, report: ExplorationReport, - trajs: Union[List[Path], List[HDF5Dataset]], - ) -> Tuple[bool, ExplorationTaskGroup, ConfSelector]: + trajs: Union[list[Path], list[HDF5Dataset]], + ) -> tuple[bool, ExplorationTaskGroup, ConfSelector]: """ Make the plan for the next iteration of the stage. diff --git a/dpgen2/exploration/selector/conf_filter.py b/dpgen2/exploration/selector/conf_filter.py index 3fca483f..1fce69a9 100644 --- a/dpgen2/exploration/selector/conf_filter.py +++ b/dpgen2/exploration/selector/conf_filter.py @@ -6,9 +6,6 @@ ABC, abstractmethod, ) -from typing import ( - List, -) import dpdata import numpy as np @@ -37,8 +34,8 @@ def check( def batched_check( self, - frames: List[dpdata.System], - ) -> List[bool]: + frames: list[dpdata.System], + ) -> list[bool]: """Check if a list of configurations are valid. Parameters diff --git a/dpgen2/exploration/selector/conf_selector.py b/dpgen2/exploration/selector/conf_selector.py index f24a7d31..26d64e67 100644 --- a/dpgen2/exploration/selector/conf_selector.py +++ b/dpgen2/exploration/selector/conf_selector.py @@ -6,10 +6,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Tuple, Union, ) @@ -33,9 +30,9 @@ class ConfSelector(ABC): @abstractmethod def select( self, - trajs: Union[List[Path], List[HDF5Dataset]], - model_devis: Union[List[Path], List[HDF5Dataset]], - type_map: Optional[List[str]] = None, - optional_outputs: Optional[List[Path]] = None, - ) -> Tuple[List[Path], ExplorationReport]: + trajs: Union[list[Path], list[HDF5Dataset]], + model_devis: Union[list[Path], list[HDF5Dataset]], + type_map: Optional[list[str]] = None, + optional_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..b1f420bb 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -6,9 +6,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, Union, ) @@ -34,7 +32,8 @@ class ConfSelectorFrames(ConfSelector): """Select frames from trajectories as confs. - Parameters: + Parameters + ---------- trust_level: TrustLevel The trust level conf_filter: ConfFilters @@ -56,12 +55,12 @@ def __init__( def select( self, - trajs: Union[List[Path], List[HDF5Dataset]], - model_devis: Union[List[Path], List[HDF5Dataset]], - type_map: Optional[List[str]] = None, - optional_outputs: Optional[List[Path]] = None, - ) -> Tuple[List[Path], ExplorationReport]: - """Select configurations + trajs: Union[list[Path], list[HDF5Dataset]], + model_devis: Union[list[Path], list[HDF5Dataset]], + type_map: Optional[list[str]] = None, + optional_outputs: Optional[list[Path]] = None, + ) -> tuple[list[Path], ExplorationReport]: + """Select configurations. Parameters ---------- diff --git a/dpgen2/exploration/selector/distance_conf_filter.py b/dpgen2/exploration/selector/distance_conf_filter.py index 7748e70d..a4106daf 100644 --- a/dpgen2/exploration/selector/distance_conf_filter.py +++ b/dpgen2/exploration/selector/distance_conf_filter.py @@ -5,9 +5,6 @@ from copy import ( deepcopy, ) -from typing import ( - List, -) import dargs import dpdata @@ -204,7 +201,6 @@ def check( for a in A: for b in B: if a < b: - print(f"Lattice length {a:.3f} is less than safe distance {b:.3f} ") return False num_atoms = len(coords) @@ -225,7 +221,7 @@ def check( def batched_check( self, - frames: List[dpdata.System], + frames: list[dpdata.System], ): if self.max_workers == 1: return list(map(self.check, frames)) @@ -234,7 +230,7 @@ def batched_check( return list(executor.map(self.check, frames)) @staticmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `ConfFilter`. Returns @@ -242,7 +238,6 @@ def args() -> List[dargs.Argument]: arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of the `ConfFilter`. """ - doc_max_workers = ( "The maximum number of processes used to filter configurations, " + "None represents as many as the processors of the machine, and 1 for serial" @@ -311,7 +306,7 @@ def check( def batched_check( self, - frames: List[dpdata.System], + frames: list[dpdata.System], ): if self.max_workers == 1: return list(map(self.check, frames)) @@ -320,7 +315,7 @@ def batched_check( return list(executor.map(self.check, frames)) @staticmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `ConfFilter`. Returns @@ -328,7 +323,6 @@ def args() -> List[dargs.Argument]: arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of the `ConfFilter`. """ - doc_max_workers = ( "The maximum number of processes used to filter configurations, " + "None represents as many as the processors of the machine, and 1 for serial" @@ -383,13 +377,13 @@ def check( c = cell[2][2] # type: ignore if check_multiples(a, b, c, self.length_ratio): - logging.warning("One side is %s larger than another" % self.length_ratio) + logging.warning(f"One side is {self.length_ratio} larger than another") return False return True def batched_check( self, - frames: List[dpdata.System], + frames: list[dpdata.System], ): if self.max_workers == 1: return list(map(self.check, frames)) @@ -398,7 +392,7 @@ def batched_check( return list(executor.map(self.check, frames)) @staticmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `ConfFilter`. Returns @@ -406,7 +400,6 @@ def args() -> List[dargs.Argument]: arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of the `ConfFilter`. """ - doc_max_workers = ( "The maximum number of processes used to filter configurations, " + "None represents as many as the processors of the machine, and 1 for serial" diff --git a/dpgen2/exploration/task/__init__.py b/dpgen2/exploration/task/__init__.py index 534a8828..b020ac69 100644 --- a/dpgen2/exploration/task/__init__.py +++ b/dpgen2/exploration/task/__init__.py @@ -17,17 +17,15 @@ caly_normalize, caly_task_group_args, diffcsp_normalize, -) -from .make_task_group_from_config import ( - lmp_normalize as normalize_lmp_task_group_config, -) -from .make_task_group_from_config import ( lmp_task_group_args, make_calypso_task_group_from_config, make_diffcsp_task_group_from_config, make_lmp_task_group_from_config, variant_task_group, ) +from .make_task_group_from_config import ( + lmp_normalize as normalize_lmp_task_group_config, +) from .npt_task_group import ( NPTTaskGroup, ) diff --git a/dpgen2/exploration/task/caly_task_group.py b/dpgen2/exploration/task/caly_task_group.py index 14594f38..86f2ac0a 100644 --- a/dpgen2/exploration/task/caly_task_group.py +++ b/dpgen2/exploration/task/caly_task_group.py @@ -1,9 +1,6 @@ import copy import logging import random -from typing import ( - List, -) import numpy as np @@ -87,7 +84,7 @@ def set_params( pop_size: int = 30, max_step: int = 5, system_name: str = "CALYPSO", - numb_of_formula: List[int] = [1, 1], + numb_of_formula: list[int] = [1, 1], pressure: float = 0.001, fmax: float = 0.01, volume: float = 0, @@ -103,26 +100,24 @@ def set_params( pick_step: int = 1, parallel: bool = False, split: bool = True, - spec_space_group: List[int] = [2, 230], + spec_space_group: list[int] = [2, 230], vsc: bool = True, - ctrl_range: List[List[int]] = [[1, 10]], + ctrl_range: list[list[int]] = [[1, 10]], max_numb_atoms: int = 100, opt_step: int = 1000, ): - """ - Set calypso parameters - """ + """Set calypso parameters.""" self.numb_of_species = numb_of_species self.numb_of_atoms = numb_of_atoms if isinstance(name_of_atoms, list) and all( - [isinstance(i, list) for i in name_of_atoms] + 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)): + if any((set(s) - overlap) == 0 for s in 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." ) diff --git a/dpgen2/exploration/task/calypso/caly_input.py b/dpgen2/exploration/task/calypso/caly_input.py index 66f2b511..e30168d2 100644 --- a/dpgen2/exploration/task/calypso/caly_input.py +++ b/dpgen2/exploration/task/calypso/caly_input.py @@ -1,6 +1,5 @@ import random from typing import ( - List, Optional, ) @@ -239,14 +238,14 @@ def check(): def make_calypso_input( numb_of_species: int, - name_of_atoms: List[str], + name_of_atoms: list[str], atomic_number, - numb_of_atoms: List[int], + numb_of_atoms: list[int], distance_of_ions, pop_size: int = 30, max_step: int = 5, system_name: str = "CALYPSO", - numb_of_formula: List[int] = [1, 1], + numb_of_formula: list[int] = [1, 1], pressure: float = 0.001, # KBar fmax: float = 0.01, volume: float = 0, @@ -262,9 +261,9 @@ def make_calypso_input( pick_step: int = 1, parallel: bool = False, split: bool = True, - spec_space_group: List[int] = [2, 230], + spec_space_group: list[int] = [2, 230], vsc: bool = False, - ctrl_range: List[List[int]] = [[1, 10]], + ctrl_range: list[list[int]] = [[1, 10]], max_numb_atoms: int = 100, **kwargs, ): diff --git a/dpgen2/exploration/task/conf_sampling_task_group.py b/dpgen2/exploration/task/conf_sampling_task_group.py index 4c5ee0c3..a63757f1 100644 --- a/dpgen2/exploration/task/conf_sampling_task_group.py +++ b/dpgen2/exploration/task/conf_sampling_task_group.py @@ -1,7 +1,6 @@ import itertools import random from typing import ( - List, Optional, ) @@ -28,12 +27,12 @@ def __init__( def set_conf( self, - conf_list: List[str], + conf_list: list[str], n_sample: Optional[int] = None, random_sample: bool = False, ): """ - Set the configurations of exploration + Set the configurations of exploration. Parameters ---------- diff --git a/dpgen2/exploration/task/customized_lmp_template_task_group.py b/dpgen2/exploration/task/customized_lmp_template_task_group.py index d7022516..4c56f9f9 100644 --- a/dpgen2/exploration/task/customized_lmp_template_task_group.py +++ b/dpgen2/exploration/task/customized_lmp_template_task_group.py @@ -6,7 +6,6 @@ Path, ) from typing import ( - List, Optional, Union, ) @@ -51,14 +50,14 @@ def __init__( def set_lmp( self, numb_models: int, - custom_shell_commands: List[str], + custom_shell_commands: list[str], revisions: dict = {}, traj_freq: int = 10, input_lmp_conf_name: str = lmp_conf_name, input_lmp_tmpl_name: str = lmp_input_name, input_plm_tmpl_name: Optional[str] = None, - input_extra_files: List[str] = [], - output_dir_pattern: Union[str, List[str]] = "*", + input_extra_files: list[str] = [], + output_dir_pattern: Union[str, list[str]] = "*", output_lmp_conf_name: str = lmp_conf_name, output_lmp_tmpl_name: str = lmp_input_name, output_plm_tmpl_name: Optional[str] = None, diff --git a/dpgen2/exploration/task/lmp/lmp_input.py b/dpgen2/exploration/task/lmp/lmp_input.py index ad3cef89..80d159b4 100644 --- a/dpgen2/exploration/task/lmp/lmp_input.py +++ b/dpgen2/exploration/task/lmp/lmp_input.py @@ -1,6 +1,5 @@ import random from typing import ( - List, Optional, ) @@ -31,12 +30,12 @@ def _sample_sphere(): def make_lmp_input( conf_file: str, ensemble: str, - graphs: List[str], + graphs: list[str], nsteps: int, dt: float, neidelay: Optional[int], trj_freq: int, - mass_map: List[float], + mass_map: list[float], temp: float, tau_t: float = 0.1, pres: Optional[float] = None, @@ -65,19 +64,19 @@ def make_lmp_input( ) if "npt" in ensemble and pres is None: raise RuntimeError("the pressre should be provided for npt ensemble") - ret = "variable NSTEPS equal %d\n" % nsteps - ret += "variable THERMO_FREQ equal %d\n" % trj_freq - ret += "variable DUMP_FREQ equal %d\n" % trj_freq - ret += "variable TEMP equal %f\n" % temp + ret = f"variable NSTEPS equal {nsteps:d}\n" + ret += f"variable THERMO_FREQ equal {trj_freq:d}\n" + ret += f"variable DUMP_FREQ equal {trj_freq:d}\n" + ret += f"variable TEMP equal {temp:f}\n" if ele_temp_f is not None: - ret += "variable ELE_TEMP equal %f\n" % ele_temp_f + ret += f"variable ELE_TEMP equal {ele_temp_f:f}\n" if ele_temp_a is not None: - ret += "variable ELE_TEMP equal %f\n" % ele_temp_a + ret += f"variable ELE_TEMP equal {ele_temp_a:f}\n" if pres is not None: - ret += "variable PRES equal %f\n" % pres - ret += "variable TAU_T equal %f\n" % tau_t + ret += f"variable PRES equal {pres:f}\n" + ret += f"variable TAU_T equal {tau_t:f}\n" if pres is not None: - ret += "variable TAU_P equal %f\n" % tau_p + ret += f"variable TAU_P equal {tau_p:f}\n" ret += "\n" ret += "units metal\n" if nopbc: @@ -88,16 +87,13 @@ def make_lmp_input( ret += "\n" ret += "neighbor 1.0 bin\n" if neidelay is not None: - ret += "neigh_modify delay %d\n" % neidelay + ret += f"neigh_modify delay {neidelay:d}\n" ret += "\n" ret += "box tilt large\n" - ret += ( - 'if "${restart} > 0" then "read_restart dpgen.restart.*" else "read_data %s"\n' - % conf_file - ) + ret += f'if "${{restart}} > 0" then "read_restart dpgen.restart.*" else "read_data {conf_file}"\n' ret += "change_box all triclinic\n" for jj in range(len(mass_map)): - ret += "mass %d %f\n" % (jj + 1, mass_map[jj]) + ret += f"mass {jj + 1:d} {mass_map[jj]:f}\n" graph_list = "" for ii in graphs: graph_list += ii + " " @@ -108,28 +104,21 @@ def make_lmp_input( ) if Version(deepmd_version) < Version("1"): # 0.x - ret += "pair_style deepmd %s ${THERMO_FREQ} %s\n" % ( - graph_list, - model_devi_file_name, - ) + ret += f"pair_style deepmd {graph_list} ${{THERMO_FREQ}} {model_devi_file_name}\n" else: # 1.x keywords = "" if use_clusters: keywords += "atomic " if relative_f_epsilon is not None: - keywords += "relative %s " % relative_f_epsilon + keywords += f"relative {relative_f_epsilon} " if relative_v_epsilon is not None: - keywords += "relative_v %s " % relative_v_epsilon + keywords += f"relative_v {relative_v_epsilon} " if ele_temp_f is not None: keywords += "fparam ${ELE_TEMP}" if ele_temp_a is not None: keywords += "aparam ${ELE_TEMP}" - ret += "pair_style deepmd %s out_freq ${THERMO_FREQ} out_file %s %s\n" % ( - graph_list, - model_devi_file_name, - keywords, - ) + ret += f"pair_style deepmd {graph_list} out_freq ${{THERMO_FREQ}} out_file {model_devi_file_name} {keywords}\n" ret += "pair_coeff * *\n" ret += "\n" ret += "thermo_style custom step temp pe ke etotal press vol lx ly lz xy xz yz\n" @@ -140,15 +129,13 @@ def make_lmp_input( lmp_traj_file_name = ( lmp_pimd_traj_name % pimd_bead if pimd_bead is not None else lmp_traj_name ) - ret += ( - "dump 1 all custom ${DUMP_FREQ} %s id type x y z fx fy fz\n" - % lmp_traj_file_name - ) + ret += f"dump 1 all custom ${{DUMP_FREQ}} {lmp_traj_file_name} id type x y z fx fy fz\n" ret += "restart 10000 dpgen.restart\n" ret += "\n" if pka_e is None: - ret += 'if "${restart} == 0" then "velocity all create ${TEMP} %d"' % ( - random.randrange(max_seed - 1) + 1 + ret += ( + 'if "${restart} == 0" then "velocity all create ${TEMP} ' + f'{random.randrange(max_seed - 1) + 1:d}"' ) else: sys = dpdata.System(conf_file, fmt="lammps/lmp") @@ -162,21 +149,16 @@ def make_lmp_input( / (0.5 * pka_mass * 1e-3 / pc.Avogadro * (pc.angstrom / pc.pico) ** 2) ) # type: ignore pka_vn = np.sqrt(pka_vn) - print(pka_vn) pka_vec = _sample_sphere() pka_vec *= pka_vn ret += "group first id 1\n" - ret += 'if "${restart} == 0" then "velocity first set %f %f %f"\n' % ( - pka_vec[0], - pka_vec[1], - pka_vec[2], - ) + ret += f'if "${{restart}} == 0" then "velocity first set {pka_vec[0]:f} {pka_vec[1]:f} {pka_vec[2]:f}"\n' ret += "fix 2 all momentum 1 linear 1 1 1\n" ret += "\n" if ensemble.split("-")[0] == "npt": assert pres is not None if nopbc: - raise RuntimeError("ensemble %s is conflicting with nopbc" % ensemble) + raise RuntimeError(f"ensemble {ensemble} is conflicting with nopbc") if ensemble == "npt" or ensemble == "npt-i" or ensemble == "npt-iso": ret += "fix 1 all npt temp ${TEMP} ${TEMP} ${TAU_T} iso ${PRES} ${PRES} ${TAU_P}\n" elif ensemble == "npt-a" or ensemble == "npt-aniso": @@ -193,6 +175,6 @@ def make_lmp_input( ret += "velocity all zero linear\n" ret += "fix fm all momentum 1 linear 1 1 1\n" ret += "\n" - ret += "timestep %f\n" % dt + ret += f"timestep {dt:f}\n" ret += "run ${NSTEPS} upto\n" return ret diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index d1f8e9fc..4b3edb28 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/dpgen2/exploration/task/lmp_template_task_group.py @@ -4,7 +4,6 @@ Path, ) from typing import ( - List, Optional, ) @@ -48,7 +47,7 @@ def set_lmp( traj_freq: int = 10, extra_pair_style_args: str = "", pimd_bead: Optional[str] = None, - input_extra_files: Optional[List[str]] = None, + input_extra_files: Optional[list[str]] = None, ) -> None: self.lmp_template = Path(lmp_template_fname).read_text().split("\n") self.revisions = revisions @@ -161,9 +160,9 @@ def find_only_one_key(lmp_lines, key): if len(words) >= nkey and words[:nkey] == key: found.append(idx) if len(found) > 1: - raise RuntimeError("found %d keywords %s" % (len(found), key)) + raise RuntimeError(f"found {len(found)} keywords {key}") if len(found) == 0: - raise RuntimeError("failed to find keyword %s" % (key)) + raise RuntimeError(f"failed to find keyword {key}") return found[0] @@ -184,11 +183,9 @@ def revise_lmp_input_model( if pimd_bead is not None else lmp_model_devi_name ) - lmp_lines[idx] = "pair_style deepmd %s out_freq %d out_file %s%s" % ( - graph_list, - trj_freq, - model_devi_file_name, - extra_pair_style_args, + lmp_lines[idx] = ( + f"pair_style deepmd {graph_list} out_freq {trj_freq:d} " + f"out_file {model_devi_file_name}{extra_pair_style_args}" ) return lmp_lines @@ -198,17 +195,16 @@ def revise_lmp_input_dump(lmp_lines, trj_freq, pimd_bead=None): lmp_traj_file_name = ( lmp_pimd_traj_name % pimd_bead if pimd_bead is not None else lmp_traj_name ) - lmp_lines[ - idx - ] = f"dump dpgen_dump all custom {trj_freq} {lmp_traj_file_name} id type x y z" + lmp_lines[idx] = ( + f"dump dpgen_dump all custom {trj_freq} {lmp_traj_file_name} id type x y z" + ) return lmp_lines def revise_lmp_input_plm(lmp_lines, in_plm, out_plm="output.plumed"): idx = find_only_one_key(lmp_lines, ["fix", "dpgen_plm"]) - lmp_lines[idx] = "fix dpgen_plm all plumed plumedfile %s outfile %s" % ( - in_plm, - out_plm, + lmp_lines[idx] = ( + f"fix dpgen_plm all plumed plumedfile {in_plm} outfile {out_plm}" ) 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 2859ac8f..05f5b12e 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -670,4 +670,4 @@ def make_lmp_task_group_from_config( if __name__ == "__main__": - print(lmp_normalize({"type": "lmp-md"})) + pass diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py index ada5a15a..404d85dd 100644 --- a/dpgen2/exploration/task/npt_task_group.py +++ b/dpgen2/exploration/task/npt_task_group.py @@ -4,7 +4,6 @@ Path, ) from typing import ( - List, Optional, ) @@ -36,8 +35,8 @@ def set_md( self, numb_models, mass_map, - temps: List[float], - press: Optional[List[float]] = None, + temps: list[float], + press: Optional[list[float]] = None, ens: str = "npt", dt: float = 0.001, nsteps: int = 1000, @@ -53,11 +52,9 @@ def set_md( ele_temp_f: Optional[float] = None, ele_temp_a: Optional[float] = None, pimd_bead: Optional[str] = None, - input_extra_files: Optional[List[str]] = None, + input_extra_files: Optional[list[str]] = None, ): - """ - Set MD parameters - """ + """Set MD parameters.""" self.graphs = [model_name_pattern % ii for ii in range(numb_models)] self.mass_map = mass_map self.temps = temps diff --git a/dpgen2/exploration/task/stage.py b/dpgen2/exploration/task/stage.py index 17ab150e..c6e9f1a5 100644 --- a/dpgen2/exploration/task/stage.py +++ b/dpgen2/exploration/task/stage.py @@ -2,9 +2,6 @@ ABC, abstractmethod, ) -from typing import ( - List, -) from dpgen2.constants import ( lmp_conf_name, @@ -22,19 +19,13 @@ class ExplorationStage: - """ - The exploration stage. - - """ + """The exploration stage.""" def __init__(self): self.clear() def clear(self): - """ - Clear all exploration group. - - """ + """Clear all exploration group.""" self.explor_groups = [] def add_task_group( @@ -42,7 +33,7 @@ def add_task_group( grp: ExplorationTaskGroup, ): """ - Add an exploration group + Add an exploration group. Parameters ---------- @@ -67,7 +58,6 @@ def make_task( added to the stage. """ - lmp_task_grp = BaseExplorationTaskGroup() for ii in self.explor_groups: # lmp_task_grp.add_group(ii.make_task()) diff --git a/dpgen2/exploration/task/task.py b/dpgen2/exploration/task/task.py index 24f7d02f..f1e9b87b 100644 --- a/dpgen2/exploration/task/task.py +++ b/dpgen2/exploration/task/task.py @@ -2,11 +2,6 @@ from collections.abc import ( Sequence, ) -from typing import ( - Dict, - List, - Tuple, -) class ExplorationTask: @@ -17,7 +12,7 @@ class ExplorationTask: >>> # this example dumps all files needed by the task. >>> files = exploration_task.files() ... for file_name, file_content in files.items(): - ... with open(file_name, 'w') as fp: + ... with open(file_name, "w") as fp: ... fp.write(file_content) """ @@ -32,7 +27,7 @@ def add_file( fname: str, fcont: str, ): - """Add file to the task + """Add file to the task. Parameters ---------- @@ -45,7 +40,7 @@ def add_file( self._files[fname] = fcont return self - def files(self) -> Dict: + def files(self) -> dict: """Get all files for the task. Returns diff --git a/dpgen2/exploration/task/task_group.py b/dpgen2/exploration/task/task_group.py index c603f80b..a2fb8e99 100644 --- a/dpgen2/exploration/task/task_group.py +++ b/dpgen2/exploration/task/task_group.py @@ -5,11 +5,6 @@ from collections.abc import ( Sequence, ) -from typing import ( - Dict, - List, - Tuple, -) from .task import ( ExplorationTask, @@ -24,19 +19,19 @@ def __init__(self): self.clear() def __getitem__(self, ii: int) -> ExplorationTask: - """Get the `ii`th task""" + """Get the `ii`th task.""" return self.task_list[ii] def __len__(self) -> int: - """Get the number of tasks in the group""" + """Get the number of tasks in the group.""" return len(self.task_list) def clear(self) -> None: self._task_list = [] @property - def task_list(self) -> List[ExplorationTask]: - """Get the `list` of `ExplorationTask`""" + def task_list(self) -> list[ExplorationTask]: + """Get the `list` of `ExplorationTask`.""" return self._task_list def add_task(self, task: ExplorationTask): @@ -110,4 +105,3 @@ def task_list(self): grp = FooTaskGroup(3) for ii in grp: fcs = ii.files() - print(fcs) diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py index 190a1090..fea346b2 100644 --- a/dpgen2/flow/dpgen_loop.py +++ b/dpgen2/flow/dpgen_loop.py @@ -7,7 +7,6 @@ Path, ) from typing import ( - List, Optional, Union, ) @@ -61,6 +60,7 @@ dump_object_to_file, load_object_from_file, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.step_config import ( init_executor, ) @@ -93,7 +93,7 @@ def get_input_sign(cls): { "exploration_scheduler": BigParameter(ExplorationScheduler), "exploration_report": BigParameter(ExplorationReport), - "trajs": Artifact(Union[List[Path], HDF5Datasets]), + "trajs": Artifact(Union[DflowList[Path], HDF5Datasets]), } ) @@ -169,7 +169,7 @@ def __init__( name: str, block_op: ConcurrentLearningBlock, step_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): self._input_parameters = { "block_id": InputParameter(), @@ -214,7 +214,7 @@ def __init__( self.step_keys = {} for ii in self._my_keys: self.step_keys[ii] = "--".join( - ["%s" % self.inputs.parameters["block_id"], ii] + ["{}".format(self.inputs.parameters["block_id"]), ii] ) self = _loop( @@ -253,7 +253,7 @@ def __init__( name: str, block_op: ConcurrentLearningBlock, step_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): self.loop = ConcurrentLearningLoop( name + "-loop", @@ -347,7 +347,7 @@ def _loop( name: str, block_op: OPTemplate, step_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): step_config = deepcopy(step_config) step_template_config = step_config.pop("template_config") @@ -449,7 +449,7 @@ def _loop( "init_data": steps.inputs.artifacts["init_data"], "iter_data": block_step.outputs.artifacts["iter_data"], }, - when="%s == false" % (scheduler_step.outputs.parameters["converged"]), + when="{} == false".format(scheduler_step.outputs.parameters["converged"]), ) steps.add(next_step) @@ -481,7 +481,7 @@ def _dpgen( loop_op, loop_key, step_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): step_config = deepcopy(step_config) step_template_config = step_config.pop("template_config") @@ -553,7 +553,7 @@ def _dpgen( "init_data": steps.inputs.artifacts["init_data"], "iter_data": steps.inputs.artifacts["iter_data"], }, - key="--".join(["%s" % id_step.outputs.parameters["block_id"], loop_key]), + key="--".join(["{}".format(id_step.outputs.parameters["block_id"]), loop_key]), ) steps.add(loop_step) diff --git a/dpgen2/fp/abacus.py b/dpgen2/fp/abacus.py index 28769b01..a80bc80c 100644 --- a/dpgen2/fp/abacus.py +++ b/dpgen2/fp/abacus.py @@ -1,9 +1,6 @@ from pathlib import ( Path, ) -from typing import ( - List, -) import dpdata from dargs import ( @@ -17,6 +14,8 @@ OPIOSign, ) +from dpgen2.utils.dflow_types import DflowList + try: from fpop.abacus import ( AbacusInputs, @@ -78,8 +77,8 @@ def get_input_sign(cls): return OPIOSign( { "config": BigParameter(dict), - "type_map": List[str], - "confs": Artifact(List[Path]), + "type_map": list[str], + "confs": Artifact(DflowList[Path]), } ) @@ -87,8 +86,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": BigParameter(List[str]), - "task_paths": Artifact(List[Path]), + "task_names": BigParameter(list[str]), + "task_paths": Artifact(DflowList[Path]), } ) @@ -115,7 +114,7 @@ def execute( s["atom_types"][i] = atom_names.index(s["atom_names"][t]) # type: ignore https://github.com/microsoft/pyright/issues/5620 s.data["atom_numbs"] = atom_numbs s.data["atom_names"] = atom_names - target = "output/%s" % system + target = f"output/{system}" s.to("deepmd/npy", target) confs.append(Path(target)) else: @@ -132,12 +131,7 @@ def execute( return op.execute(op_in) # type: ignore in the case of not importing fpop -from typing import ( - Tuple, -) - - -def get_suffix_calculation(INPUT: List[str]) -> Tuple[str, str]: +def get_suffix_calculation(INPUT: list[str]) -> tuple[str, str]: suffix = "ABACUS" calculation = "scf" for iline in INPUT: @@ -166,7 +160,7 @@ def get_output_sign(cls): { "log": Artifact(Path), "labeled_data": Artifact(Path), - "extra_outputs": Artifact(List[Path]), + "extra_outputs": Artifact(DflowList[Path]), } ) @@ -189,7 +183,7 @@ def execute( workdir = op_out["backward_dir"].parent # convert the output to deepmd/npy format - with open("%s/INPUT" % workdir, "r") as f: + with open(f"{workdir}/INPUT") as f: INPUT = f.readlines() _, calculation = get_suffix_calculation(INPUT) if calculation == "scf": @@ -199,7 +193,7 @@ def execute( elif calculation in ["relax", "cell-relax"]: sys = dpdata.LabeledSystem(str(workdir), fmt="abacus/relax") else: - raise ValueError("Type of calculation %s not supported" % calculation) + raise ValueError(f"Type of calculation {calculation} not supported") out_name = fp_default_out_data_name sys.to("deepmd/npy", workdir / out_name) diff --git a/dpgen2/fp/cp2k.py b/dpgen2/fp/cp2k.py index 67ed5d55..d28db70e 100644 --- a/dpgen2/fp/cp2k.py +++ b/dpgen2/fp/cp2k.py @@ -3,7 +3,6 @@ Path, ) from typing import ( - List, Optional, ) @@ -19,6 +18,8 @@ OPIOSign, ) +from dpgen2.utils.dflow_types import DflowList + try: from fpop.cp2k import ( Cp2kInputs, @@ -48,8 +49,8 @@ def get_input_sign(cls): return OPIOSign( { "config": BigParameter(dict), - "type_map": List[str], - "confs": Artifact(List[Path]), + "type_map": list[str], + "confs": Artifact(DflowList[Path]), } ) @@ -57,8 +58,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": BigParameter(List[str]), - "task_paths": Artifact(List[Path]), + "task_names": BigParameter(list[str]), + "task_paths": Artifact(DflowList[Path]), } ) @@ -85,7 +86,7 @@ def execute( s["atom_types"][i] = atom_names.index(s["atom_names"][t]) # type: ignore https://github.com/microsoft/pyright/issues/5620 s.data["atom_numbs"] = atom_numbs s.data["atom_names"] = atom_names - target = "output/%s" % system + target = f"output/{system}" s.to("deepmd/npy", target) confs.append(Path(target)) else: @@ -102,7 +103,7 @@ def execute( return op.execute(op_in) # type: ignore in the case of not importing fpop -def get_run_type(lines: List[str]) -> Optional[str]: +def get_run_type(lines: list[str]) -> Optional[str]: for line in lines: if "RUN_TYPE" in line: return line.split()[-1] @@ -126,7 +127,7 @@ def get_output_sign(cls): { "log": Artifact(Path), "labeled_data": Artifact(Path), - "extra_outputs": Artifact(List[Path]), + "extra_outputs": Artifact(DflowList[Path]), } ) @@ -152,7 +153,7 @@ def execute( file_path = os.path.join(str(workdir), "output.log") # convert the output to deepmd/npy format - with open(workdir / "input.inp", "r") as f: + with open(workdir / "input.inp") as f: lines = f.readlines() # 获取 RUN_TYPE diff --git a/dpgen2/fp/deepmd.py b/dpgen2/fp/deepmd.py index 43fb200a..fb555cf4 100644 --- a/dpgen2/fp/deepmd.py +++ b/dpgen2/fp/deepmd.py @@ -1,13 +1,12 @@ """Prep and Run Gaussian tasks.""" + import os from pathlib import ( Path, ) from typing import ( Any, - List, Optional, - Tuple, ) import dpdata @@ -48,7 +47,7 @@ class DeepmdInputs: @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: return [] def __init__(self, **kwargs: Any): @@ -74,7 +73,7 @@ def prep_task( class RunDeepmd(RunFp): - def input_files(self) -> List[str]: + def input_files(self) -> list[str]: r"""The mandatory input files to run a Deepmd task. Returns @@ -85,7 +84,7 @@ def input_files(self) -> List[str]: """ return [deepmd_input_path] - def optional_input_files(self) -> List[str]: + def optional_input_files(self) -> list[str]: r"""The optional input files to run a Deepmd task. Returns @@ -101,8 +100,8 @@ def run_task( teacher_model_path: BinaryFileInput, out: str, log: str, - ) -> Tuple[str, str]: - r"""Defines how one FP task runs + ) -> tuple[str, str]: + r"""Defines how one FP task runs. Parameters ---------- @@ -197,7 +196,7 @@ def _dp_infer(self, dp, type_map_teacher, out_name): ss.to("deepmd/npy", out_name) @staticmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `run_task` method. Returns @@ -205,7 +204,6 @@ def args() -> List[dargs.Argument]: arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of `run_task` method. """ - doc_deepmd_teacher_model = ( "The path of teacher model, which can be loaded by deepmd.infer.DeepPot" ) diff --git a/dpgen2/fp/gaussian.py b/dpgen2/fp/gaussian.py index b6aba200..a06815d8 100644 --- a/dpgen2/fp/gaussian.py +++ b/dpgen2/fp/gaussian.py @@ -1,10 +1,9 @@ """Prep and Run Gaussian tasks.""" + import logging from typing import ( Any, - List, Optional, - Tuple, ) import dpdata @@ -38,7 +37,7 @@ class GaussianInputs: @staticmethod - def args() -> List[Argument]: + def args() -> list[Argument]: r"""The arguments of the GaussianInputs class.""" doc_keywords = "Gaussian keywords, e.g. force b3lyp/6-31g**. If a list, run multiple steps." doc_multiplicity = ( @@ -104,12 +103,11 @@ def prep_task( inputs : GaussianInputs The GaussianInputs object handels all other input files of the task. """ - conf_frame.to("gaussian/gjf", gaussian_input_name, **inputs.data) class RunGaussian(RunFp): - def input_files(self) -> List[str]: + def input_files(self) -> list[str]: r"""The mandatory input files to run a Gaussian task. Returns @@ -120,7 +118,7 @@ def input_files(self) -> List[str]: """ return [gaussian_input_name] - def optional_input_files(self) -> List[str]: + def optional_input_files(self) -> list[str]: r"""The optional input files to run a Gaussian task. Returns @@ -136,8 +134,8 @@ def run_task( command: str, out: str, post_command: Optional[str] = None, - ) -> Tuple[str, str]: - r"""Defines how one FP task runs + ) -> tuple[str, str]: + r"""Defines how one FP task runs. Parameters ---------- @@ -195,7 +193,7 @@ def run_task( return out_name, gaussian_output_name @staticmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `run_task` method. Returns @@ -203,7 +201,6 @@ def args() -> List[dargs.Argument]: arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of `run_task` method. """ - doc_gaussian_cmd = "The command of Gaussian" doc_gaussian_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`." doc_post_command = "The command after Gaussian" diff --git a/dpgen2/fp/prep_fp.py b/dpgen2/fp/prep_fp.py index 962be031..ce2f9fbb 100644 --- a/dpgen2/fp/prep_fp.py +++ b/dpgen2/fp/prep_fp.py @@ -7,8 +7,6 @@ ) from typing import ( Any, - List, - Tuple, ) import dpdata @@ -27,6 +25,7 @@ set_directory, setup_ele_temp, ) +from dpgen2.utils.dflow_types import DflowList class PrepFp(OP, ABC): @@ -45,8 +44,8 @@ def get_input_sign(cls): return OPIOSign( { "config": BigParameter(dict), - "type_map": List[str], - "confs": Artifact(List[Path]), + "type_map": list[str], + "confs": Artifact(DflowList[Path]), } ) @@ -54,8 +53,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": BigParameter(List[str]), - "task_paths": Artifact(List[Path]), + "task_names": BigParameter(list[str]), + "task_paths": Artifact(DflowList[Path]), } ) @@ -100,7 +99,6 @@ def execute( - `task_names`: (`List[str]`) The name of tasks. Will be used as the identities of the tasks. The names of different tasks are different. - `task_paths`: (`Artifact(List[Path])`) The parepared working paths of the tasks. Contains all input files needed to start the FP. The order fo the Paths should be consistent with `op["task_names"]` """ - inputs = ip["config"]["inputs"] confs = ip["confs"] type_map = ip["type_map"] @@ -137,7 +135,7 @@ def _exec_one_frame( idx, inputs, conf_frame: dpdata.System, - ) -> Tuple[str, Path]: + ) -> tuple[str, Path]: task_name = fp_task_pattern % idx task_path = Path(task_name) with set_directory(task_path): diff --git a/dpgen2/fp/run_fp.py b/dpgen2/fp/run_fp.py index 254936d0..5df83535 100644 --- a/dpgen2/fp/run_fp.py +++ b/dpgen2/fp/run_fp.py @@ -6,11 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - Dict, - List, - Tuple, -) import dargs from dflow.python import ( @@ -26,6 +21,7 @@ from dpgen2.utils.chdir import ( set_directory, ) +from dpgen2.utils.dflow_types import DflowList class RunFp(OP, ABC): @@ -55,12 +51,12 @@ def get_output_sign(cls): { "log": Artifact(Path), "labeled_data": Artifact(Path), - "extra_outputs": Artifact(List[Path]), + "extra_outputs": Artifact(DflowList[Path]), } ) @abstractmethod - def input_files(self) -> List[str]: + def input_files(self) -> list[str]: r"""The mandatory input files to run a FP task. Returns @@ -72,7 +68,7 @@ def input_files(self) -> List[str]: pass @abstractmethod - def optional_input_files(self) -> List[str]: + def optional_input_files(self) -> list[str]: r"""The optional input files to run a FP task. Returns @@ -87,8 +83,8 @@ def optional_input_files(self) -> List[str]: def run_task( self, **kwargs, - ) -> Tuple[str, str]: - r"""Defines how one FP task runs + ) -> tuple[str, str]: + r"""Defines how one FP task runs. Parameters ---------- @@ -107,7 +103,7 @@ def run_task( @staticmethod @abstractmethod - def args() -> List[dargs.Argument]: + def args() -> list[dargs.Argument]: r"""The argument definition of the `run_task` method. Returns @@ -118,7 +114,7 @@ def args() -> List[dargs.Argument]: pass @classmethod - def normalize_config(cls, data: Dict = {}, strict: bool = True) -> Dict: + def normalize_config(cls, data: dict = {}, strict: bool = True) -> dict: r"""Normalized the argument. Parameters diff --git a/dpgen2/fp/vasp.py b/dpgen2/fp/vasp.py index b8b12b28..a1f32cce 100644 --- a/dpgen2/fp/vasp.py +++ b/dpgen2/fp/vasp.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) import dpdata import numpy as np @@ -64,7 +60,8 @@ def clean_lines(string_list, remove_empty_lines=True): remove_empty_lines: Set to True to skip lines which are empty after stripping. - Returns: + Returns + ------- List of clean strings with no whitespaces. """ for s in string_list: @@ -134,7 +131,6 @@ def prep_task( vasp_inputs : VaspInputs The VaspInputs object handels all other input files of the task. """ - conf_frame.to("vasp/poscar", vasp_conf_name) incar = vasp_inputs.incar_template incar = self.set_ele_temp(conf_frame, incar) @@ -147,7 +143,7 @@ def prep_task( class RunVasp(RunFp): - def input_files(self) -> List[str]: + def input_files(self) -> list[str]: r"""The mandatory input files to run a vasp task. Returns @@ -158,7 +154,7 @@ def input_files(self) -> List[str]: """ return [vasp_conf_name, vasp_input_name, vasp_pot_name, vasp_kp_name] - def optional_input_files(self) -> List[str]: + def optional_input_files(self) -> list[str]: r"""The optional input files to run a vasp task. Returns @@ -171,7 +167,7 @@ def optional_input_files(self) -> List[str]: def set_ele_temp(self, system): if os.path.exists("job.json"): - with open("job.json", "r") as f: + with open("job.json") as f: data = json.load(f) if "use_ele_temp" in data and "ele_temp" in data: if data["use_ele_temp"] == 1: @@ -188,8 +184,8 @@ def run_task( command: str, out: str, log: str, - ) -> Tuple[str, str]: - r"""Defines how one FP task runs + ) -> tuple[str, str]: + r"""Defines how one FP task runs. Parameters ---------- @@ -207,7 +203,6 @@ def run_task( log_name: str The file name of the log. """ - log_name = log out_name = out # run vasp @@ -235,7 +230,6 @@ def args(): arguments: List[dargs.Argument] List of dargs.Argument defines the arguments of `run_task` method. """ - doc_vasp_cmd = "The command of VASP" doc_vasp_log = "The log file name of VASP" doc_vasp_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`." diff --git a/dpgen2/fp/vasp_input.py b/dpgen2/fp/vasp_input.py index 2b60b18e..87e63fb1 100644 --- a/dpgen2/fp/vasp_input.py +++ b/dpgen2/fp/vasp_input.py @@ -2,11 +2,7 @@ Path, ) from typing import ( - Dict, - List, Optional, - Set, - Tuple, Union, ) @@ -22,9 +18,9 @@ class VaspInputs: def __init__( self, - kspacing: Union[float, List[float]], + kspacing: Union[float, list[float]], incar: str, - pp_files: Dict[str, str], + pp_files: dict[str, str], kgamma: bool = True, ): """ @@ -66,7 +62,7 @@ def incar_from_file( def potcars_from_file( self, - dict_fnames: Dict[str, str], + dict_fnames: dict[str, str], ): self._potcars = {} for kk, vv in dict_fnames.items(): @@ -127,7 +123,7 @@ def _make_vasp_kp_gamma(kpoints): ret += "Automatic mesh\n" ret += "0\n" ret += "Gamma\n" - ret += "%d %d %d\n" % (kpoints[0], kpoints[1], kpoints[2]) + ret += f"{kpoints[0]:d} {kpoints[1]:d} {kpoints[2]:d}\n" ret += "0 0 0\n" return ret @@ -137,7 +133,7 @@ def _make_vasp_kp_mp(kpoints): ret += "K-Points\n" ret += "0\n" ret += "Monkhorst Pack\n" - ret += "%d %d %d\n" % (kpoints[0], kpoints[1], kpoints[2]) + ret += f"{kpoints[0]:d} {kpoints[1]:d} {kpoints[2]:d}\n" ret += "0 0 0\n" return ret diff --git a/dpgen2/op/caly_evo_step_merge.py b/dpgen2/op/caly_evo_step_merge.py index 8abda6dd..71bb25fc 100644 --- a/dpgen2/op/caly_evo_step_merge.py +++ b/dpgen2/op/caly_evo_step_merge.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow import ( Step, @@ -44,6 +40,7 @@ BinaryFileInput, set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -70,7 +67,7 @@ def get_input_sign(cls): "caly_check_opt_file": Artifact(Path), "results": Artifact(Path, optional=True), "step": Artifact(Path, optional=True), - "opt_results_dir": Artifact(List[Path], optional=True), + "opt_results_dir": Artifact(DflowList[Path], optional=True), "qhull_input": Artifact(Path, optional=True), } ) @@ -79,7 +76,7 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "traj_results": Artifact(List[Path]), + "traj_results": Artifact(DflowList[Path]), } ) @@ -117,9 +114,9 @@ def execute( output_sign = self.get_output_sign() for k in step.outputs.artifacts: path_list = download_artifact(step.outputs.artifacts[k]) - if output_sign[k].type == List[Path]: + if output_sign[k].type == list[Path]: if not isinstance(path_list, list) or any( - [p is not None and not isinstance(p, str) for p in path_list] + p is not None and not isinstance(p, str) for p in path_list ): path_list = list(flatten(path_list).values()) out[k] = [Path(p) for p in path_list] diff --git a/dpgen2/op/collect_data.py b/dpgen2/op/collect_data.py index 68f397ff..a4689aaa 100644 --- a/dpgen2/op/collect_data.py +++ b/dpgen2/op/collect_data.py @@ -3,11 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, - Tuple, -) import dpdata from dflow.python import ( @@ -22,6 +17,7 @@ from dpgen2.utils import ( setup_ele_temp, ) +from dpgen2.utils.dflow_types import DflowList class CollectData(OP): @@ -44,13 +40,13 @@ def get_input_sign(cls): return OPIOSign( { "name": str, - "type_map": List[str], + "type_map": list[str], "optional_parameter": Parameter( dict, default=CollectData.default_optional_parameter, ), - "labeled_data": Artifact(List[Path]), - "iter_data": Artifact(List[Path]), + "labeled_data": Artifact(DflowList[Path]), + "iter_data": Artifact(DflowList[Path]), } ) @@ -58,7 +54,7 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "iter_data": Artifact(List[Path]), + "iter_data": Artifact(DflowList[Path]), } ) diff --git a/dpgen2/op/collect_run_caly.py b/dpgen2/op/collect_run_caly.py index 4b6148f6..240860e6 100644 --- a/dpgen2/op/collect_run_caly.py +++ b/dpgen2/op/collect_run_caly.py @@ -8,10 +8,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Tuple, ) from dargs import ( @@ -38,6 +35,7 @@ BinaryFileInput, set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -67,7 +65,7 @@ def get_input_sign(cls): type=Path, optional=True ), # dir named results for evo "opt_results_dir": Artifact( - type=List[Path], optional=True + type=DflowList[Path], optional=True ), # dir contains POSCAR* CONTCAR* OUTCAR* "qhull_input": Artifact(type=Path, optional=True), # for vsc } @@ -246,7 +244,7 @@ def prep_last_calypso_file(step, results, opt_results_dir, qhull_input, vsc): def get_value_from_inputdat(filename): max_step = 0 vsc = False - with open(filename, "r") as f: + with open(filename) as f: lines = f.readlines() for line in lines: if "MaxStep" in line: diff --git a/dpgen2/op/diffcsp_gen.py b/dpgen2/op/diffcsp_gen.py index a7ed5ef1..12aee6af 100644 --- a/dpgen2/op/diffcsp_gen.py +++ b/dpgen2/op/diffcsp_gen.py @@ -3,9 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - List, -) from dflow.python import ( OP, @@ -14,6 +11,8 @@ OPIOSign, ) +from dpgen2.utils.dflow_types import DflowList + def convert_pt_to_cif(input_file, output_dir): import numpy as np @@ -55,7 +54,7 @@ def convert_pt_to_cif(input_file, output_dir): lattice, atom_type, frac_coord, coords_are_cartesian=False ) - filename = "%s.cif" % i + filename = f"{i}.cif" file_path = os.path.join(output_dir, filename) structure.to(filename=file_path) now_atom += atom_num @@ -75,7 +74,7 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "cifs": Artifact(List[Path]), + "cifs": Artifact(DflowList[Path]), } ) @@ -88,12 +87,12 @@ def execute( args = cmd.split() try: i = args.index("--model_path") - except ValueError: - raise RuntimeError("Path of DiffCSP model not provided.") + except ValueError as exc: + raise RuntimeError("Path of DiffCSP model not provided.") from exc model_path = args[i + 1] subprocess.run(cmd, shell=True, check=True) result_file = os.path.join(model_path, "eval_gen.pt") - task_dir = "diffcsp.%s" % ip["task_id"] + task_dir = f"diffcsp.{ip['task_id']}" convert_pt_to_cif(result_file, task_dir) return OPIO( { diff --git a/dpgen2/op/md_settings.py b/dpgen2/op/md_settings.py index 6916e3b8..949500bd 100644 --- a/dpgen2/op/md_settings.py +++ b/dpgen2/op/md_settings.py @@ -1,6 +1,5 @@ import json from typing import ( - List, Optional, ) @@ -12,8 +11,8 @@ def __init__( dt: float, nsteps: int, trj_freq: int, - temps: Optional[List[float]] = None, - press: Optional[List[float]] = None, + temps: Optional[list[float]] = None, + press: Optional[list[float]] = None, tau_t: float = 0.1, tau_p: float = 0.5, pka_e: Optional[float] = None, diff --git a/dpgen2/op/prep_caly_dp_optim.py b/dpgen2/op/prep_caly_dp_optim.py index d2e4d8b0..9dee911e 100644 --- a/dpgen2/op/prep_caly_dp_optim.py +++ b/dpgen2/op/prep_caly_dp_optim.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, @@ -33,6 +29,7 @@ BinaryFileInput, set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -69,8 +66,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": Parameter(List[str]), - "task_dirs": Artifact(List[Path]), + "task_names": Parameter(list[str]), + "task_dirs": Artifact(DflowList[Path]), "caly_run_opt_file": Artifact(Path), # from prep_caly_input "caly_check_opt_file": Artifact(Path), # from prep_caly_input } diff --git a/dpgen2/op/prep_caly_input.py b/dpgen2/op/prep_caly_input.py index e3da359a..09f03ae0 100644 --- a/dpgen2/op/prep_caly_input.py +++ b/dpgen2/op/prep_caly_input.py @@ -3,10 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, @@ -31,6 +27,7 @@ from dpgen2.utils import ( set_directory, ) +from dpgen2.utils.dflow_types import DflowList vsc_keys = { "VSC": "F", @@ -319,10 +316,10 @@ def get_output_sign(cls): return OPIOSign( { "ntasks": Parameter(int), - "task_names": BigParameter(List[str]), # task dir names - "input_dat_files": Artifact(List[Path]), # `input.dat`s - "caly_run_opt_files": Artifact(List[Path]), - "caly_check_opt_files": Artifact(List[Path]), + "task_names": BigParameter(list[str]), # task dir names + "input_dat_files": Artifact(DflowList[Path]), # `input.dat`s + "caly_run_opt_files": Artifact(DflowList[Path]), + "caly_check_opt_files": Artifact(DflowList[Path]), } ) @@ -349,7 +346,6 @@ def execute( - `caly_run_opt_files`: (`Artifact(List[Path])`) - `caly_check_opt_files`: (`Artifact(List[Path])`) """ - cc = 0 task_paths = [] input_dat_files = [] diff --git a/dpgen2/op/prep_caly_model_devi.py b/dpgen2/op/prep_caly_model_devi.py index 3b070959..e87ad4a1 100644 --- a/dpgen2/op/prep_caly_model_devi.py +++ b/dpgen2/op/prep_caly_model_devi.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, @@ -33,6 +29,7 @@ BinaryFileInput, set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -49,7 +46,7 @@ def get_input_sign(cls): { "task_name": Parameter(str), "config": BigParameter(dict), - "traj_results": Artifact(List[Path]), + "traj_results": Artifact(DflowList[Path]), } ) @@ -57,8 +54,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_name_list": Parameter(List[str]), - "grouped_traj_list": Artifact(List[Path]), + "task_name_list": Parameter(list[str]), + "grouped_traj_list": Artifact(DflowList[Path]), } ) diff --git a/dpgen2/op/prep_dp_train.py b/dpgen2/op/prep_dp_train.py index 20fe58c2..73b96171 100644 --- a/dpgen2/op/prep_dp_train.py +++ b/dpgen2/op/prep_dp_train.py @@ -5,8 +5,6 @@ Path, ) from typing import ( - List, - Tuple, Union, ) @@ -22,6 +20,7 @@ train_script_name, train_task_pattern, ) +from dpgen2.utils.dflow_types import DflowList class PrepDPTrain(OP): @@ -38,7 +37,7 @@ class PrepDPTrain(OP): def get_input_sign(cls): return OPIOSign( { - "template_script": BigParameter(Union[dict, List[dict]]), + "template_script": BigParameter(Union[dict, list[dict]]), "numb_models": int, } ) @@ -47,8 +46,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": BigParameter(List[str]), - "task_paths": Artifact(List[Path]), + "task_names": BigParameter(list[str]), + "task_paths": Artifact(DflowList[Path]), } ) diff --git a/dpgen2/op/prep_lmp.py b/dpgen2/op/prep_lmp.py index e1b5c026..9efc0b55 100644 --- a/dpgen2/op/prep_lmp.py +++ b/dpgen2/op/prep_lmp.py @@ -3,10 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, @@ -23,6 +19,7 @@ BaseExplorationTaskGroup, ExplorationTaskGroup, ) +from dpgen2.utils.dflow_types import DflowList class PrepLmp(OP): @@ -48,8 +45,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "task_names": BigParameter(List[str]), - "task_paths": Artifact(List[Path]), + "task_names": BigParameter(list[str]), + "task_paths": Artifact(DflowList[Path]), } ) @@ -74,7 +71,6 @@ def execute( - `task_names`: (`List[str]`) The name of tasks. Will be used as the identities of the tasks. The names of different tasks are different. - `task_paths`: (`Artifact(List[Path])`) The parepared working paths of the tasks. Contains all input files needed to start the LAMMPS simulation. The order fo the Paths should be consistent with `op["task_names"]` """ - lmp_task_grp = ip["lmp_task_grp"] cc = 0 task_paths = [] diff --git a/dpgen2/op/prep_relax.py b/dpgen2/op/prep_relax.py index 1ee2869a..87348457 100644 --- a/dpgen2/op/prep_relax.py +++ b/dpgen2/op/prep_relax.py @@ -2,9 +2,6 @@ from pathlib import ( Path, ) -from typing import ( - List, -) from dflow.python import ( OP, @@ -13,6 +10,8 @@ OPIOSign, ) +from dpgen2.utils.dflow_types import DflowList + class PrepRelax(OP): @classmethod @@ -20,7 +19,7 @@ def get_input_sign(cls): return OPIOSign( { "expl_config": dict, - "cifs": Artifact(List[Path]), + "cifs": Artifact(DflowList[Path]), } ) @@ -29,7 +28,7 @@ def get_output_sign(cls): return OPIOSign( { "ntasks": int, - "task_paths": Artifact(List[Path]), + "task_paths": Artifact(DflowList[Path]), } ) @@ -44,10 +43,10 @@ def execute( ntasks = int(ncifs / group_size) task_paths = [] for i in range(ntasks): - task_dir = Path("task.%06d" % i) + task_dir = Path(f"task.{i:06d}") task_dir.mkdir(exist_ok=True) for j in range(group_size * i, min(group_size * (i + 1), ncifs)): - os.symlink(ip["cifs"][j], task_dir / ("%s.cif" % j)) + os.symlink(ip["cifs"][j], task_dir / (f"{j}.cif")) task_paths.append(task_dir) return OPIO( { diff --git a/dpgen2/op/run_caly_dp_optim.py b/dpgen2/op/run_caly_dp_optim.py index 639c97d5..e00297e6 100644 --- a/dpgen2/op/run_caly_dp_optim.py +++ b/dpgen2/op/run_caly_dp_optim.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, diff --git a/dpgen2/op/run_caly_model_devi.py b/dpgen2/op/run_caly_model_devi.py index 9e191cf2..aaf0aa8b 100644 --- a/dpgen2/op/run_caly_model_devi.py +++ b/dpgen2/op/run_caly_model_devi.py @@ -5,7 +5,6 @@ Path, ) from typing import ( - List, Union, ) @@ -22,6 +21,7 @@ from dpgen2.utils import ( set_directory, ) +from dpgen2.utils.dflow_types import DflowList class RunCalyModelDevi(OP): @@ -36,10 +36,10 @@ class RunCalyModelDevi(OP): def get_input_sign(cls): return OPIOSign( { - "type_map": Parameter(List[str]), + "type_map": Parameter(list[str]), "task_name": Parameter(str), - "traj_dirs": Artifact(List[Path]), - "models": Artifact(List[Path]), + "traj_dirs": Artifact(DflowList[Path]), + "models": Artifact(DflowList[Path]), } ) @@ -48,8 +48,8 @@ def get_output_sign(cls): return OPIOSign( { "task_name": Parameter(str), - "traj": Artifact(List[Path]), - "model_devi": Artifact(List[Path]), + "traj": Artifact(DflowList[Path]), + "model_devi": Artifact(DflowList[Path]), } ) @@ -128,9 +128,9 @@ def execute( traj_str = dump_str_dict[key] model_devis = devis_dict[key] - assert len(traj_str) == len( - model_devis - ), "The length of traj_str and model_devis should be same." + assert len(traj_str) == len(model_devis), ( + "The length of traj_str and model_devis should be same." + ) for idx in range(len(model_devis)): traj_str[idx] = traj_str[idx] % idx model_devis[idx][0] = idx @@ -158,7 +158,7 @@ def execute( def atoms2lmpdump(atoms, struc_idx, type_map, ignore=False): - """down triangle cell can be obtained from + """Down triangle cell can be obtained from cell params: a, b, c, alpha, beta, gamma. cell = cellpar_to_cell([a, b, c, alpha, beta, gamma]) lx, ly, lz = cell[0][0], cell[1][1], cell[2][2] @@ -169,7 +169,7 @@ def atoms2lmpdump(atoms, struc_idx, type_map, ignore=False): ylo_bound = ylo + MIN(0.0,yz) yhi_bound = yhi + MAX(0.0,yz) zlo_bound = zlo - zhi_bound = zhi + zhi_bound = zhi. ref: https://docs.lammps.org/Howto_triclinic.html """ @@ -208,19 +208,19 @@ def atoms2lmpdump(atoms, struc_idx, type_map, ignore=False): zhi_bound = zhi dump_str += "ITEM: BOX BOUNDS xy xz yz pp pp pp\n" - dump_str += "%20.10f %20.10f %20.10f\n" % (xlo_bound, xhi_bound, xy) - dump_str += "%20.10f %20.10f %20.10f\n" % (ylo_bound, yhi_bound, xz) - dump_str += "%20.10f %20.10f %20.10f\n" % (zlo_bound, zhi_bound, yz) + dump_str += f"{xlo_bound:20.10f} {xhi_bound:20.10f} {xy:20.10f}\n" + dump_str += f"{ylo_bound:20.10f} {yhi_bound:20.10f} {xz:20.10f}\n" + dump_str += f"{zlo_bound:20.10f} {zhi_bound:20.10f} {yz:20.10f}\n" dump_str += "ITEM: ATOMS id type x y z fx fy fz\n" for idx, atom in enumerate(new_atoms): type_id = type_map.index(atom.symbol) + 1 # type: ignore - dump_str += "%5d %5d" % (idx + 1, type_id) - dump_str += "%20.10f %20.10f %20.10f" % ( - atom.position[0], # type: ignore - atom.position[1], # type: ignore - atom.position[2], # type: ignore + dump_str += f"{idx + 1:5d} {type_id:5d}" + dump_str += ( + f"{atom.position[0]:20.10f} " # type: ignore + f"{atom.position[1]:20.10f} " # type: ignore + f"{atom.position[2]:20.10f}" # type: ignore ) - dump_str += "%20.10f %20.10f %20.10f\n" % (0, 0, 0) + dump_str += f"{0:20.10f} {0:20.10f} {0:20.10f}\n" # dump_str = dump_str.strip("\n") return dump_str @@ -266,7 +266,7 @@ def parse_traj(traj_file): "H": 0.813, } - trajs: List[Atoms] = read(traj_file, index=":", format="traj") # type: ignore + trajs: list[Atoms] = read(traj_file, index=":", format="traj") # type: ignore dthresh = 0.72 numb_traj = len(trajs) assert numb_traj >= 1, "traj file is broken." @@ -279,7 +279,7 @@ def parse_traj(traj_file): dis_mtx[row, col] = np.nan is_reasonable = np.nanmin(dis_mtx) > dthresh - selected_traj: Union[List[Atoms], None] = None + selected_traj: Union[DflowList[Atoms], None] = None if is_reasonable: if len(trajs) >= 20: selected_traj = [trajs[iii] for iii in [4, 9, -10, -5, -1]] @@ -326,14 +326,12 @@ def parse_traj(traj_file): def write_model_devi_out(devi: np.ndarray, fname: Union[str, Path], header: str = ""): assert devi.shape[1] == 8 - header = "%s\n%10s" % (header, "step") + header = f"{header}\n{'step':>10}" for item in "vf": - header += "%19s%19s%19s" % ( - f"max_devi_{item}", - f"min_devi_{item}", - f"avg_devi_{item}", + header += ( + f"{'max_devi_' + item:>19}{'min_devi_' + item:>19}{'avg_devi_' + item:>19}" ) - header += "%19s" % "devi_e" + header += f"{'devi_e':>19}" with open(fname, "ab") as fp: np.savetxt( fp, diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index c0cf2d4d..a5f4f6f2 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -9,10 +9,7 @@ Path, ) from typing import ( - Dict, - List, Optional, - Tuple, Union, ) @@ -42,6 +39,7 @@ from dpgen2.utils.chdir import ( set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -62,8 +60,8 @@ def _make_train_command( if impl == "tensorflow" and os.path.isfile("checkpoint"): checkpoint = "model.ckpt" elif impl == "pytorch" and len(glob.glob("model.ckpt-[0-9]*.pt")) > 0: - checkpoint = "model.ckpt-%s.pt" % max( - [int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")] + checkpoint = "model.ckpt-{}.pt".format( + max([int(f[11:-3]) for f in glob.glob("model.ckpt-[0-9]*.pt")]) ) else: checkpoint = None @@ -130,9 +128,9 @@ def get_input_sign(cls): "task_path": Artifact(Path), "init_model": Artifact(Path, optional=True), "init_data": Artifact(NestedDict[Path]), - "iter_data": Artifact(List[Path]), + "iter_data": Artifact(DflowList[Path]), "valid_data": Artifact(NestedDict[Path], optional=True), - "optional_files": Artifact(List[Path], optional=True), + "optional_files": Artifact(DflowList[Path], optional=True), } ) @@ -375,11 +373,11 @@ def clean_before_quit(): def write_data_to_input_script( idict: dict, config, - init_data: Union[List[Path], Dict[str, List[Path]]], - iter_data: List[Path], + init_data: Union[DflowList[Path], dict[str, list[Path]]], + iter_data: list[Path], auto_prob_str: str = "prob_sys_size", major_version: str = "1", - valid_data: Optional[Union[List[Path], Dict[str, List[Path]]]] = None, + valid_data: Optional[Union[DflowList[Path], dict[str, list[Path]]]] = None, ): odict = idict.copy() if config["multitask"]: @@ -706,12 +704,12 @@ def _expand_all_multi_sys_to_sys(list_multi_sys): return all_sys_dirs -def split_valid(systems: List[str], valid_ratio: float): +def split_valid(systems: list[str], valid_ratio: float): train_systems = [] valid_systems = [] for system in systems: d = dpdata.MultiSystems() - mixed_type = len(glob.glob("%s/*/real_atom_types.npy" % system)) > 0 + mixed_type = len(glob.glob(f"{system}/*/real_atom_types.npy")) > 0 if mixed_type: d.load_systems_from_file(system, fmt="deepmd/npy/mixed") else: @@ -735,11 +733,11 @@ def split_valid(systems: List[str], valid_ratio: float): target = "train_data/" + system if mixed_type: # The multisystem is loaded from one dir, thus we can safely keep one dir - train_multi_systems.to_deepmd_npy_mixed("%s.tmp" % target) # type: ignore - fs = os.listdir("%s.tmp" % target) + train_multi_systems.to_deepmd_npy_mixed(f"{target}.tmp") # type: ignore + fs = os.listdir(f"{target}.tmp") assert len(fs) == 1 - os.rename(os.path.join("%s.tmp" % target, fs[0]), target) - os.rmdir("%s.tmp" % target) + os.rename(os.path.join(f"{target}.tmp", fs[0]), target) + os.rmdir(f"{target}.tmp") else: train_multi_systems[0].to_deepmd_npy(target) # type: ignore train_systems.append(os.path.abspath(target)) @@ -748,11 +746,11 @@ def split_valid(systems: List[str], valid_ratio: float): target = "valid_data/" + system if mixed_type: # The multisystem is loaded from one dir, thus we can safely keep one dir - valid_multi_systems.to_deepmd_npy_mixed("%s.tmp" % target) # type: ignore - fs = os.listdir("%s.tmp" % target) + valid_multi_systems.to_deepmd_npy_mixed(f"{target}.tmp") # type: ignore + fs = os.listdir(f"{target}.tmp") assert len(fs) == 1 - os.rename(os.path.join("%s.tmp" % target, fs[0]), target) - os.rmdir("%s.tmp" % target) + os.rename(os.path.join(f"{target}.tmp", fs[0]), target) + os.rmdir(f"{target}.tmp") else: valid_multi_systems[0].to_deepmd_npy(target) # type: ignore valid_systems.append(os.path.abspath(target)) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 60cd9305..9b05fea8 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -8,10 +8,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Tuple, ) import numpy as np @@ -47,6 +44,7 @@ BinaryFileInput, set_directory, ) +from dpgen2.utils.dflow_types import DflowList from dpgen2.utils.run_command import ( run_command, ) @@ -70,7 +68,7 @@ def get_input_sign(cls): "config": BigParameter(dict), "task_name": BigParameter(str), "task_path": Artifact(Path), - "models": Artifact(List[Path]), + "models": Artifact(DflowList[Path]), } ) @@ -83,7 +81,7 @@ def get_output_sign(cls): "model_devi": Artifact(Path), "plm_output": Artifact(Path, optional=True), "optional_output": Artifact(Path, optional=True), - "extra_outputs": Artifact(List[Path]), + "extra_outputs": Artifact(DflowList[Path]), } ) @@ -132,9 +130,9 @@ def execute( work_dir = Path(task_name) if teacher_model is not None: - assert ( - len(model_files) == 1 - ), "One model is enough in knowledge distillation" + assert len(model_files) == 1, ( + "One model is enough in knowledge distillation" + ) ext = os.path.splitext(teacher_model.file_name)[-1] teacher_model_file = "teacher_model" + ext teacher_model.save_as_file(teacher_model_file) @@ -158,7 +156,7 @@ def execute( freeze_model(mm, mname, config.get("model_frozen_head")) else: raise RuntimeError( - "Model file with extension '%s' is not supported" % ext + f"Model file with extension '{ext}' is not supported" ) model_names.append(mname) @@ -283,7 +281,7 @@ def normalize_config(data={}): config_args = RunLmp.lmp_args -def set_models(lmp_input_name: str, model_names: List[str]): +def set_models(lmp_input_name: str, model_names: list[str]): with open(lmp_input_name, encoding="utf8") as f: lmp_input_lines = f.readlines() @@ -306,7 +304,7 @@ def set_models(lmp_input_name: str, model_names: List[str]): break if match_first == -1: raise RuntimeError( - f"cannot file model pattern {pattern} in line " f" {lmp_input_lines[idx]}" + f"cannot file model pattern {pattern} in line {lmp_input_lines[idx]}" ) if match_last == -1: raise RuntimeError(f"last matching index should not be -1, terribly wrong ") @@ -331,10 +329,10 @@ def find_only_one_key(lmp_lines, key, raise_not_found=True): if len(words) >= nkey and words[:nkey] == key: found.append(idx) if len(found) > 1: - raise RuntimeError("found %d keywords %s" % (len(found), key)) + raise RuntimeError(f"found {len(found)} keywords {key}") if len(found) == 0: if raise_not_found: - raise RuntimeError("failed to find keyword %s" % (key)) + raise RuntimeError(f"failed to find keyword {key}") else: return None return found[0] @@ -363,10 +361,10 @@ def get_ele_temp(lmp_log_name): def freeze_model(input_model, frozen_model, head=None): - freeze_args = "-o %s" % frozen_model + freeze_args = f"-o {frozen_model}" if head is not None: - freeze_args += " --head %s" % head - freeze_cmd = "dp --pt freeze -c %s %s" % (input_model, freeze_args) + freeze_args += f" --head {head}" + freeze_cmd = f"dp --pt freeze -c {input_model} {freeze_args}" ret, out, err = run_command(freeze_cmd, shell=True) if ret != 0: logging.error( @@ -392,13 +390,13 @@ def merge_pimd_files(): if len(traj_files) > 0: with open(lmp_traj_name, "w") as f: for traj_file in sorted(traj_files): - with open(traj_file, "r") as f2: + with open(traj_file) as f2: f.write(f2.read()) model_devi_files = glob.glob("model_devi.*.out") if len(model_devi_files) > 0: with open(lmp_model_devi_name, "w") as f: for model_devi_file in sorted(model_devi_files): - with open(model_devi_file, "r") as f2: + with open(model_devi_file) as f2: f.write(f2.read()) diff --git a/dpgen2/op/run_relax.py b/dpgen2/op/run_relax.py index 672275d8..1c4056f9 100644 --- a/dpgen2/op/run_relax.py +++ b/dpgen2/op/run_relax.py @@ -3,9 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - List, -) from dargs import ( Argument, @@ -25,6 +22,7 @@ from dpgen2.exploration.task import ( DiffCSPTaskGroup, ) +from dpgen2.utils.dflow_types import DflowList from .run_caly_model_devi import ( atoms2lmpdump, @@ -42,7 +40,7 @@ def get_input_sign(cls): "diffcsp_task_grp": BigParameter(DiffCSPTaskGroup), "expl_config": dict, "task_path": Artifact(Path), - "models": Artifact(List[Path]), + "models": Artifact(DflowList[Path]), } ) @@ -50,8 +48,8 @@ def get_input_sign(cls): def get_output_sign(cls): return OPIOSign( { - "trajs": Artifact(List[Path]), - "model_devis": Artifact(List[Path]), + "trajs": Artifact(DflowList[Path]), + "model_devis": Artifact(DflowList[Path]), } ) @@ -62,14 +60,14 @@ def write_traj(self, dump_str, traj_file): def write_model_devi(self, devi, model_devi_file): import numpy as np - header = "%10s%19s%19s%19s%19s%19s%19s" % ( - "step", - "max_devi_v", - "min_devi_v", - "avg_devi_v", - "max_devi_f", - "min_devi_f", - "avg_devi_f", + header = ( + f"{'step':>10}" + f"{'max_devi_v':>19}" + f"{'min_devi_v':>19}" + f"{'avg_devi_v':>19}" + f"{'max_devi_f':>19}" + f"{'min_devi_f':>19}" + f"{'avg_devi_f':>19}" ) np.savetxt( model_devi_file, @@ -193,14 +191,14 @@ def execute( ) forces_list[j] = forces virial_list[j] = virial / len(atype) - traj_file = ip["task_path"] / ("traj.%s.dump" % fname) + traj_file = ip["task_path"] / (f"traj.{fname}.dump") traj_file = self.write_traj(dump_str, traj_file) trajs.append(traj_file) devi = [np.array(step_list)] devi += list(calc_model_devi_v(np.array(virial_list))) devi += list(calc_model_devi_f(np.array(forces_list))) devi = np.vstack(devi).T - model_devi_file = ip["task_path"] / ("model_devi.%s.out" % fname) + model_devi_file = ip["task_path"] / (f"model_devi.{fname}.out") model_devi_file = self.write_model_devi(devi, model_devi_file) model_devis.append(model_devi_file) return OPIO( diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index e8ba891d..1f25aadd 100644 --- a/dpgen2/op/select_confs.py +++ b/dpgen2/op/select_confs.py @@ -4,9 +4,6 @@ Path, ) from typing import ( - List, - Set, - Tuple, Union, ) @@ -26,6 +23,7 @@ from dpgen2.exploration.selector import ( ConfSelector, ) +from dpgen2.utils.dflow_types import DflowList class SelectConfs(OP): @@ -36,10 +34,10 @@ def get_input_sign(cls): return OPIOSign( { "conf_selector": ConfSelector, - "type_map": List[str], - "trajs": Artifact(Union[List[Path], HDF5Datasets]), - "model_devis": Artifact(Union[List[Path], HDF5Datasets]), - "optional_outputs": Artifact(List[Path], optional=True), + "type_map": list[str], + "trajs": Artifact(Union[DflowList[Path], HDF5Datasets]), + "model_devis": Artifact(Union[DflowList[Path], HDF5Datasets]), + "optional_outputs": Artifact(DflowList[Path], optional=True), } ) @@ -48,7 +46,7 @@ def get_output_sign(cls): return OPIOSign( { "report": BigParameter(ExplorationReport), - "confs": Artifact(List[Path]), + "confs": Artifact(DflowList[Path]), } ) @@ -77,7 +75,6 @@ def execute( - `conf`: (`Artifact(List[Path])`) The selected configurations. """ - conf_selector = ip["conf_selector"] type_map = ip["type_map"] @@ -111,11 +108,11 @@ def validate_trajs( ntrajs = len(trajs) if ntrajs != len(model_devis): raise FatalError( - "length of trajs list is not equal to the " "model_devis list" + "length of trajs list is not equal to the model_devis list" ) if optional_outputs and ntrajs != len(optional_outputs): raise FatalError( - "length of trajs list is not equal to the " "optional_output list" + "length of trajs list is not equal to the optional_output list" ) rett = [] retm = [] diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py index 0e39ab38..a66823b0 100644 --- a/dpgen2/superop/block.py +++ b/dpgen2/superop/block.py @@ -7,11 +7,7 @@ ) from typing import ( Any, - Dict, - List, Optional, - Set, - Type, Union, ) @@ -90,12 +86,12 @@ def __init__( name: str, prep_run_dp_train_op: PrepRunDPTrain, prep_run_explore_op: Union[PrepRunLmp, PrepRunCaly, PrepRunDiffCSP], - select_confs_op: Type[OP], + select_confs_op: type[OP], prep_run_fp_op: PrepRunFp, - collect_data_op: Type[OP], + collect_data_op: type[OP], select_confs_config: dict = normalize_step_dict({}), collect_data_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): self._input_parameters = { "block_id": InputParameter(), @@ -148,7 +144,7 @@ def __init__( self.step_keys = {} for ii in self._my_keys: self.step_keys[ii] = "--".join( - ["%s" % self.inputs.parameters["block_id"], ii] + ["{}".format(self.inputs.parameters["block_id"]), ii] ) self = _block_cl( @@ -188,16 +184,16 @@ def keys(self): def _block_cl( block_steps: Steps, - step_keys: Dict[str, Any], + step_keys: dict[str, Any], name: str, prep_run_dp_train_op: OPTemplate, prep_run_explore_op: OPTemplate, - select_confs_op: Type[OP], + select_confs_op: type[OP], prep_run_fp_op: OPTemplate, - collect_data_op: Type[OP], + collect_data_op: type[OP], select_confs_config: dict = normalize_step_dict({}), collect_data_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): select_confs_config = deepcopy(select_confs_config) collect_data_config = deepcopy(collect_data_config) @@ -228,7 +224,7 @@ def _block_cl( "iter_data": block_steps.inputs.artifacts["iter_data"], }, key="--".join( - ["%s" % block_steps.inputs.parameters["block_id"], "prep-run-train"] + ["{}".format(block_steps.inputs.parameters["block_id"]), "prep-run-train"] ), ) block_steps.add(prep_run_dp_train) @@ -246,7 +242,7 @@ def _block_cl( "models": prep_run_dp_train.outputs.artifacts["models"], }, key="--".join( - ["%s" % block_steps.inputs.parameters["block_id"], "prep-run-explore"] + ["{}".format(block_steps.inputs.parameters["block_id"]), "prep-run-explore"] ), ) block_steps.add(prep_run_explore) @@ -288,7 +284,7 @@ def _block_cl( "confs": select_confs.outputs.artifacts["confs"], }, key="--".join( - ["%s" % block_steps.inputs.parameters["block_id"], "prep-run-fp"] + ["{}".format(block_steps.inputs.parameters["block_id"]), "prep-run-fp"] ), ) block_steps.add(prep_run_fp) diff --git a/dpgen2/superop/caly_evo_step.py b/dpgen2/superop/caly_evo_step.py index cd2be501..a5063f41 100644 --- a/dpgen2/superop/caly_evo_step.py +++ b/dpgen2/superop/caly_evo_step.py @@ -6,10 +6,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Type, ) from dflow import ( @@ -51,13 +48,13 @@ class CalyEvoStep(Steps): def __init__( self, name: str, - collect_run_caly: Type[OP], - prep_dp_optim: Type[OP], - run_dp_optim: Type[OP], + collect_run_caly: type[OP], + prep_dp_optim: type[OP], + run_dp_optim: type[OP], expl_mode: str = "default", prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): self.expl_mode = expl_mode self._input_parameters = { @@ -132,12 +129,12 @@ def keys(self): def _caly_evo_step( caly_evo_step_steps, step_keys, - collect_run_calypso_op: Type[OP], - prep_dp_optim_op: Type[OP], - run_dp_optim_op: Type[OP], + collect_run_calypso_op: type[OP], + prep_dp_optim_op: type[OP], + run_dp_optim_op: type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) @@ -181,8 +178,7 @@ def wise_executor(expl_mode, origin_executor_config): "opt_results_dir": caly_evo_step_steps.inputs.artifacts["opt_results_dir"], "qhull_input": caly_evo_step_steps.inputs.artifacts["qhull_input"], }, - key="%s--collect-run-calypso-%s-%s" - % ( + key="{}--collect-run-calypso-{}-{}".format( caly_evo_step_steps.inputs.parameters["block_id"], caly_evo_step_steps.inputs.parameters["iter_num"], caly_evo_step_steps.inputs.parameters["cnt_num"], @@ -215,8 +211,7 @@ def wise_executor(expl_mode, origin_executor_config): "caly_check_opt_file" ], }, - key="%s--prep-dp-optim-%s-%s" - % ( + key="{}--prep-dp-optim-{}-{}".format( caly_evo_step_steps.inputs.parameters["block_id"], caly_evo_step_steps.inputs.parameters["iter_num"], caly_evo_step_steps.inputs.parameters["cnt_num"], @@ -248,8 +243,7 @@ def wise_executor(expl_mode, origin_executor_config): artifacts={ "task_dir": prep_dp_optim.outputs.artifacts["task_dirs"], }, - key="%s--run-dp-optim-%s-%s-{{item}}" - % ( + key="{}--run-dp-optim-{}-{}-{{{{item}}}}".format( caly_evo_step_steps.inputs.parameters["block_id"], caly_evo_step_steps.inputs.parameters["iter_num"], caly_evo_step_steps.inputs.parameters["cnt_num"], @@ -284,7 +278,7 @@ def wise_executor(expl_mode, origin_executor_config): "caly_check_opt_file" ], }, - when="%s == false" % (collect_run_calypso.outputs.parameters["finished"]), + when="{} == false".format(collect_run_calypso.outputs.parameters["finished"]), ) caly_evo_step_steps.add(next_step) diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py index daa48143..9748029d 100644 --- a/dpgen2/superop/prep_run_calypso.py +++ b/dpgen2/superop/prep_run_calypso.py @@ -7,10 +7,7 @@ ) from typing import ( Any, - Dict, - List, Optional, - Type, Union, ) @@ -53,14 +50,14 @@ class PrepRunCaly(Steps): def __init__( self, name: str, - prep_caly_input_op: Type[OP], + prep_caly_input_op: type[OP], caly_evo_step_op: Union[OPTemplate, OP], - prep_caly_model_devi_op: Type[OP], - run_caly_model_devi_op: Type[OP], + prep_caly_model_devi_op: type[OP], + run_caly_model_devi_op: type[OP], expl_mode: str = "default", prep_config: Optional[dict] = None, run_config: Optional[dict] = None, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = normalize_step_dict({}) if prep_config is None else prep_config run_config = normalize_step_dict({}) if run_config is None else run_config @@ -102,13 +99,21 @@ def __init__( ] self.step_keys = {} ii = "prep-caly-input" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "caly-evo-step-{{item}}" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "prep-caly-model-devi" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "run-caly-model-devi" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) self = _prep_run_caly( self, @@ -146,15 +151,15 @@ def keys(self): def _prep_run_caly( prep_run_caly_steps: Steps, - step_keys: Dict[str, Any], - prep_caly_input_op: Type[OP], + step_keys: dict[str, Any], + prep_caly_input_op: type[OP], caly_evo_step_op: Union[OPTemplate, OP], - prep_caly_model_devi_op: Type[OP], - run_caly_model_devi_op: Type[OP], + prep_caly_model_devi_op: type[OP], + run_caly_model_devi_op: type[OP], expl_mode: str = "default", prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) @@ -260,8 +265,9 @@ def _prep_run_caly( artifacts={ "traj_results": caly_evo_step.outputs.artifacts["traj_results"], }, - key="%s--prep-caly-model-devi" - % (prep_run_caly_steps.inputs.parameters["block_id"],), + key="{}--prep-caly-model-devi".format( + prep_run_caly_steps.inputs.parameters["block_id"] + ), executor=prep_executor, ) prep_run_caly_steps.add(prep_caly_model_devi) @@ -287,8 +293,9 @@ def _prep_run_caly( "traj_dirs": prep_caly_model_devi.outputs.artifacts["grouped_traj_list"], "models": prep_run_caly_steps.inputs.artifacts["models"], }, - key="%s--run-caly-model-devi-{{item}}" - % (prep_run_caly_steps.inputs.parameters["block_id"],), + key="{}--run-caly-model-devi-{{{{item}}}}".format( + prep_run_caly_steps.inputs.parameters["block_id"] + ), executor=run_executor, **prep_config, ) diff --git a/dpgen2/superop/prep_run_diffcsp.py b/dpgen2/superop/prep_run_diffcsp.py index c44851fe..cf1223ac 100644 --- a/dpgen2/superop/prep_run_diffcsp.py +++ b/dpgen2/superop/prep_run_diffcsp.py @@ -7,10 +7,7 @@ ) from typing import ( Any, - Dict, - List, Optional, - Type, Union, ) @@ -50,12 +47,12 @@ class PrepRunDiffCSP(Steps): def __init__( self, name: str, - diffcsp_gen_op: Type[OP], - prep_relax_op: Type[OP], - run_relax_op: Type[OP], + diffcsp_gen_op: type[OP], + prep_relax_op: type[OP], + run_relax_op: type[OP], prep_config: Optional[dict] = None, run_config: Optional[dict] = None, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = normalize_step_dict({}) if prep_config is None else prep_config run_config = normalize_step_dict({}) if run_config is None else run_config @@ -120,12 +117,12 @@ def keys(self): def _prep_run_diffcsp( prep_run_diffcsp_steps: Steps, - diffcsp_gen_op: Type[OP], - prep_relax_op: Type[OP], - run_relax_op: Type[OP], + diffcsp_gen_op: type[OP], + prep_relax_op: type[OP], + run_relax_op: type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) @@ -157,7 +154,7 @@ def _prep_run_diffcsp( "task_id": "{{item}}", "config": expl_config, }, - key="%s--diffcsp-gen-{{item}}" % block_id, + key=f"{block_id}--diffcsp-gen-{{{{item}}}}", executor=prep_executor, with_sequence=argo_sequence(expl_config["gen_tasks"], format="%06d"), # type: ignore ) @@ -176,7 +173,7 @@ def _prep_run_diffcsp( artifacts={ "cifs": diffcsp_gen.outputs.artifacts["cifs"], }, - key="%s--prep-relax" % block_id, + key=f"{block_id}--prep-relax", executor=prep_executor, ) prep_run_diffcsp_steps.add(prep_relax) @@ -202,7 +199,7 @@ def _prep_run_diffcsp( "models": models, "task_path": prep_relax.outputs.artifacts["task_paths"], }, - key="%s--run-relax-{{item}}" % block_id, + key=f"{block_id}--run-relax-{{{{item}}}}", executor=run_executor, with_sequence=argo_sequence( prep_relax.outputs.parameters["ntasks"], format="%06d" diff --git a/dpgen2/superop/prep_run_dp_train.py b/dpgen2/superop/prep_run_dp_train.py index 0fd988e4..12421a65 100644 --- a/dpgen2/superop/prep_run_dp_train.py +++ b/dpgen2/superop/prep_run_dp_train.py @@ -7,10 +7,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Type, ) from dflow import ( @@ -58,13 +55,13 @@ class PrepRunDPTrain(Steps): def __init__( self, name: str, - prep_train_op: Type[OP], - run_train_op: Type[RunDPTrain], + prep_train_op: type[OP], + run_train_op: type[RunDPTrain], prep_config: Optional[dict] = None, run_config: Optional[dict] = None, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, valid_data: Optional[S3Artifact] = None, - optional_files: Optional[List[str]] = None, + optional_files: Optional[list[str]] = None, ): prep_config = normalize_step_dict({}) if prep_config is None else prep_config run_config = normalize_step_dict({}) if run_config is None else run_config @@ -107,10 +104,12 @@ def __init__( self._keys = ["prep-train", "run-train"] self.step_keys = {} ii = "prep-train" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "run-train" self.step_keys[ii] = "--".join( - ["%s" % self.inputs.parameters["block_id"], ii + "-{{item}}"] + ["{}".format(self.inputs.parameters["block_id"]), ii + "-{{item}}"] ) self = _prep_run_dp_train( @@ -149,13 +148,13 @@ def keys(self): def _prep_run_dp_train( train_steps, step_keys, - prep_train_op: Type[OP], - run_train_op: Type[RunDPTrain], + prep_train_op: type[OP], + run_train_op: type[RunDPTrain], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, valid_data: Optional[S3Artifact] = None, - optional_files: Optional[List[str]] = None, + optional_files: Optional[list[str]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) diff --git a/dpgen2/superop/prep_run_fp.py b/dpgen2/superop/prep_run_fp.py index ba659c6d..493a1cd0 100644 --- a/dpgen2/superop/prep_run_fp.py +++ b/dpgen2/superop/prep_run_fp.py @@ -6,10 +6,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Type, ) from dflow import ( @@ -50,11 +47,11 @@ class PrepRunFp(Steps): def __init__( self, name: str, - prep_op: Type[OP], - run_op: Type[OP], + prep_op: type[OP], + run_op: type[OP], prep_config: Optional[dict] = None, run_config: Optional[dict] = None, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = normalize_step_dict({}) if prep_config is None else prep_config run_config = normalize_step_dict({}) if run_config is None else run_config @@ -88,10 +85,12 @@ def __init__( self._keys = ["prep-fp", "run-fp"] self.step_keys = {} ii = "prep-fp" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "run-fp" self.step_keys[ii] = "--".join( - ["%s" % self.inputs.parameters["block_id"], ii + "-{{item}}"] + ["{}".format(self.inputs.parameters["block_id"]), ii + "-{{item}}"] ) self = _prep_run_fp( @@ -128,11 +127,11 @@ def keys(self): def _prep_run_fp( prep_run_steps, step_keys, - prep_op: Type[OP], - run_op: Type[OP], + prep_op: type[OP], + run_op: type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 3e0a0a0f..ad42935f 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -6,10 +6,7 @@ Path, ) from typing import ( - List, Optional, - Set, - Type, ) from dflow import ( @@ -50,11 +47,11 @@ class PrepRunLmp(Steps): def __init__( self, name: str, - prep_op: Type[OP], - run_op: Type[OP], + prep_op: type[OP], + run_op: type[OP], prep_config: Optional[dict] = None, run_config: Optional[dict] = None, - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = normalize_step_dict({}) if prep_config is None else prep_config run_config = normalize_step_dict({}) if run_config is None else run_config @@ -94,10 +91,12 @@ def __init__( self._keys = ["prep-lmp", "run-lmp"] self.step_keys = {} ii = "prep-lmp" - self.step_keys[ii] = "--".join(["%s" % self.inputs.parameters["block_id"], ii]) + self.step_keys[ii] = "--".join( + ["{}".format(self.inputs.parameters["block_id"]), ii] + ) ii = "run-lmp" self.step_keys[ii] = "--".join( - ["%s" % self.inputs.parameters["block_id"], ii + "-{{item}}"] + ["{}".format(self.inputs.parameters["block_id"]), ii + "-{{item}}"] ) self = _prep_run_lmp( @@ -134,11 +133,11 @@ def keys(self): def _prep_run_lmp( prep_run_steps, step_keys, - prep_op: Type[OP], - run_op: Type[OP], + prep_op: type[OP], + run_op: type[OP], prep_config: dict = normalize_step_dict({}), run_config: dict = normalize_step_dict({}), - upload_python_packages: Optional[List[os.PathLike]] = None, + upload_python_packages: Optional[list[os.PathLike]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) diff --git a/dpgen2/utils/__init__.py b/dpgen2/utils/__init__.py index b7824f57..6341e765 100644 --- a/dpgen2/utils/__init__.py +++ b/dpgen2/utils/__init__.py @@ -40,8 +40,6 @@ from .step_config import gen_doc as gen_doc_step_dict from .step_config import ( init_executor, -) -from .step_config import normalize as normalize_step_dict -from .step_config import ( step_conf_args, ) +from .step_config import normalize as normalize_step_dict diff --git a/dpgen2/utils/artifact_uri.py b/dpgen2/utils/artifact_uri.py index 4f7ec035..e1822563 100644 --- a/dpgen2/utils/artifact_uri.py +++ b/dpgen2/utils/artifact_uri.py @@ -11,13 +11,13 @@ def get_artifact_from_uri(uri): elif uri.startswith("oss://"): return S3Artifact(uri[6:]) else: - raise ValueError("Unrecognized scheme of URI: %s" % uri) + raise ValueError(f"Unrecognized scheme of URI: {uri}") def upload_artifact_and_print_uri(files, name): art = upload_artifact(files) if s3_config["repo_type"] == "s3" and hasattr(art, "key"): - print("%s has been uploaded to s3://%s" % (name, art.key)) + pass elif s3_config["repo_type"] == "oss" and hasattr(art, "key"): - print("%s has been uploaded to oss://%s" % (name, art.key)) + pass return art diff --git a/dpgen2/utils/binary_file_input.py b/dpgen2/utils/binary_file_input.py index db099a00..c665c584 100644 --- a/dpgen2/utils/binary_file_input.py +++ b/dpgen2/utils/binary_file_input.py @@ -1,4 +1,5 @@ -"""Binary file inputs""" +"""Binary file inputs.""" + import os import warnings from pathlib import ( @@ -6,9 +7,7 @@ ) from typing import ( Any, - List, Optional, - Tuple, Union, ) @@ -30,9 +29,9 @@ def __init__(self, path: Union[str, Path], ext: Optional[str] = None) -> None: self.ext = ext if self.ext: - assert ( - os.path.splitext(path)[-1] == self.ext - ), f'File extension mismatch, require "{ext}", current "{os.path.splitext(path)[-1]}", file path: {str(path)}' + assert os.path.splitext(path)[-1] == self.ext, ( + f'File extension mismatch, require "{ext}", current "{os.path.splitext(path)[-1]}", file path: {str(path)}' + ) self.file_name = os.path.basename(path) with open(path, "rb") as f: @@ -42,7 +41,7 @@ def save_as_file(self, path: Union[str, Path]) -> None: if self.ext and os.path.splitext(path)[-1] != self.ext: warnings.warn( f'warning: file extension mismatch! Extension of input file is "{self.ext}",' - + f"current extension is \"{str(path).split('.')[-1]}\"" + + f'current extension is "{str(path).split(".")[-1]}"' ) with open(path, "wb") as file: diff --git a/dpgen2/utils/chdir.py b/dpgen2/utils/chdir.py index 3bc0b5c0..5be879e5 100644 --- a/dpgen2/utils/chdir.py +++ b/dpgen2/utils/chdir.py @@ -33,7 +33,7 @@ def set_directory(path: Path): Examples -------- >>> with set_directory("some_path"): - ... do_something() + ... do_something() """ cwd = Path().absolute() path.mkdir(exist_ok=True, parents=True) diff --git a/dpgen2/utils/dflow_config.py b/dpgen2/utils/dflow_config.py index e6a7b9ba..f7e2b489 100644 --- a/dpgen2/utils/dflow_config.py +++ b/dpgen2/utils/dflow_config.py @@ -19,7 +19,7 @@ def dflow_config_lower( dflow_config, ): dflow_s3_config = {} - keys = [kk for kk in dflow_config.keys()] + keys = list(dflow_config.keys()) for kk in keys: if kk[:3] == "s3_": dflow_s3_config[kk[3:]] = dflow_config.pop(kk) @@ -40,7 +40,7 @@ def dflow_config( config_data, ): """ - set the dflow config by `config_data` + Set the dflow config by `config_data`. the keys starting with "s3_" will be treated as s3_config keys, other keys are treated as config keys. @@ -53,9 +53,6 @@ def dflow_config( def dflow_s3_config( config_data, ): - """ - set the s3 config by `config_data` - - """ + """Set the s3 config by `config_data`.""" if config_data is not None: dflow_s3_config_lower(config_data) diff --git a/dpgen2/utils/dflow_query.py b/dpgen2/utils/dflow_query.py index 933c8cc6..bf7c4245 100644 --- a/dpgen2/utils/dflow_query.py +++ b/dpgen2/utils/dflow_query.py @@ -2,7 +2,6 @@ import re from typing import ( Any, - List, Optional, ) @@ -23,12 +22,10 @@ def get_iteration( def matched_step_key( - all_keys: List[str], - step_keys: Optional[List[str]] = None, + all_keys: list[str], + step_keys: Optional[list[str]] = None, ): - """ - returns the keys in `all_keys` that matches any of the `step_keys` - """ + """Returns the keys in `all_keys` that matches any of the `step_keys`.""" if step_keys is None: return all_keys ret = [] @@ -46,11 +43,9 @@ def matched_step_key( def get_last_scheduler( wf: Any, - keys: List[str], + keys: list[str], ): - """ - get the output Scheduler of the last successful iteration - """ + """Get the output Scheduler of the last successful iteration.""" outputs = wf.query_global_outputs() if ( outputs is not None @@ -80,11 +75,9 @@ def get_last_scheduler( def get_all_schedulers( wf: Any, - keys: List[str], + keys: list[str], ): - """ - get the output Scheduler of the all the iterations - """ + """Get the output Scheduler of the all the iterations.""" scheduler_keys = sorted(matched_step_key(keys, ["scheduler"])) if len(scheduler_keys) == 0: return None @@ -97,21 +90,17 @@ def get_all_schedulers( def get_last_iteration( - keys: List[str], + keys: list[str], ): - """ - get the index of the last iteraction from a list of step keys. - """ + """Get the index of the last iteraction from a list of step keys.""" return int(sorted([get_subkey(ii, 0) for ii in keys])[-1].split("-")[1]) def find_slice_ranges( - keys: List[str], + keys: list[str], sliced_subkey: str, ): - """ - find range of sliced OPs that matches the pattern 'iter-[0-9]*--{sliced_subkey}-[0-9]*' - """ + """Find range of sliced OPs that matches the pattern 'iter-[0-9]*--{sliced_subkey}-[0-9]*'.""" found_range = [] tmp_range = [] status = "not-found" @@ -139,12 +128,10 @@ def _sort_slice_ops(keys, sliced_subkey): def sort_slice_ops( - keys: List[str], - sliced_subkey: List[str], + keys: list[str], + sliced_subkey: list[str], ): - """ - sort the keys of the sliced ops. the keys of the sliced ops contains sliced_subkey - """ + """Sort the keys of the sliced ops. the keys of the sliced ops contains sliced_subkey.""" if isinstance(sliced_subkey, str): sliced_subkey = [sliced_subkey] for ii in sliced_subkey: @@ -153,8 +140,8 @@ def sort_slice_ops( def print_keys_in_nice_format( - keys: List[str], - sliced_subkey: List[str], + keys: list[str], + sliced_subkey: list[str], idx_fmt_len: int = 8, ): keys = sort_slice_ops(keys, sliced_subkey) @@ -165,9 +152,9 @@ def print_keys_in_nice_format( slice_0 = [ii[0] for ii in slice_range] slice_1 = [ii[1] for ii in slice_range] - normal_fmt = f"%{idx_fmt_len*2+4}d" + normal_fmt = f"%{idx_fmt_len * 2 + 4}d" range_fmt = f"%d -> %d" - range_s_fmt = f"%{idx_fmt_len*2+4}s" + range_s_fmt = f"%{idx_fmt_len * 2 + 4}s" idx = 0 ret = [] diff --git a/dpgen2/utils/dflow_types.py b/dpgen2/utils/dflow_types.py new file mode 100644 index 00000000..55790c35 --- /dev/null +++ b/dpgen2/utils/dflow_types.py @@ -0,0 +1,8 @@ +"""Compatibility aliases for types inspected by dflow at runtime.""" + +import typing + +# dflow validates artifact schema types against ``typing.List`` objects. The +# equivalent PEP 585 ``list`` form is rejected by current pydflow releases, so +# keep this runtime-only alias until dflow accepts both representations. +DflowList = typing.List # noqa: UP006 diff --git a/dpgen2/utils/download_dpgen2_artifacts.py b/dpgen2/utils/download_dpgen2_artifacts.py index 67c2aaf0..d82f3c32 100644 --- a/dpgen2/utils/download_dpgen2_artifacts.py +++ b/dpgen2/utils/download_dpgen2_artifacts.py @@ -4,7 +4,6 @@ Path, ) from typing import ( - List, Optional, ) @@ -102,17 +101,16 @@ def download_dpgen2_artifacts( chk_pnt: bool = False, ): """ - download the artifacts of a step. + Download the artifacts of a step. the key should be of format 'iter-xxxxxx--subkey-of-step-xxxxxx' the input and output artifacts will be downloaded to prefix/iter-xxxxxx/key-of-step/inputs/ and - prefix/iter-xxxxxx/key-of-step/outputs/ + prefix/iter-xxxxxx/key-of-step/outputs/. the downloaded input and output artifacts of steps are defined by `op_download_setting` """ - iteration = get_iteration(key) subkey = get_subkey(key) mypath = Path(iteration) @@ -158,8 +156,8 @@ def download_dpgen2_artifacts( def download_dpgen2_artifacts_by_def( wf: Workflow, - iterations: Optional[List[int]] = None, - step_defs: Optional[List[str]] = None, + iterations: Optional[list[int]] = None, + step_defs: Optional[list[str]] = None, prefix: Optional[str] = None, chk_pnt: bool = False, ): @@ -273,7 +271,7 @@ def _get_all_iterations(step_keys): if ii.startswith("iter-"): ii = int(ii.split("-")[1]) ret.append(ii) - ret = sorted(list(set(ret))) + ret = sorted(set(ret)) return ret @@ -283,7 +281,7 @@ def _get_all_queried_steps(wf_step_keys, dld_items): ret.append(ii.split(global_step_def_split)[0]) ret = set(ret) ret = ret.intersection(set(wf_step_keys)) - return sorted(list(ret)) + return sorted(ret) def _get_dld_items( diff --git a/dpgen2/utils/obj_artifact.py b/dpgen2/utils/obj_artifact.py index 05ab85eb..9d1e712c 100644 --- a/dpgen2/utils/obj_artifact.py +++ b/dpgen2/utils/obj_artifact.py @@ -8,10 +8,7 @@ def dump_object_to_file( obj, fname, ): - """ - pickle dump object to a file - - """ + """Pickle dump object to a file.""" with open(fname, "wb") as fp: pickle.dump(obj, fp) return Path(fname) @@ -20,10 +17,7 @@ def dump_object_to_file( def load_object_from_file( fname, ): - """ - pickle load object from a file - - """ + """Pickle load object from a file.""" with open(fname, "rb") as fp: obj = pickle.load(fp) return obj diff --git a/dpgen2/utils/run_command.py b/dpgen2/utils/run_command.py index 2d5c5764..390ea5e3 100644 --- a/dpgen2/utils/run_command.py +++ b/dpgen2/utils/run_command.py @@ -1,7 +1,5 @@ import os from typing import ( - List, - Tuple, Union, ) @@ -12,9 +10,9 @@ def run_command( - cmd: Union[str, List[str]], + cmd: Union[str, list[str]], shell: bool = False, -) -> Tuple[int, str, str]: +) -> tuple[int, str, str]: interactive = False if config["mode"] == "debug" else True return dflow_run_command( cmd, raise_error=False, try_bash=shell, interactive=interactive diff --git a/dpgen2/utils/step_config.py b/dpgen2/utils/step_config.py index 591f474e..258c67ac 100644 --- a/dpgen2/utils/step_config.py +++ b/dpgen2/utils/step_config.py @@ -18,7 +18,7 @@ def dispatcher_args(): - """free style dispatcher args""" + """Free style dispatcher args.""" return [] diff --git a/pyproject.toml b/pyproject.toml index ccf004f9..54a7a3c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,6 +81,33 @@ exclude = [ "dpgen2/_version.py", ] -[tool.isort] -profile = "black" -force_grid_wrap = 1 +[tool.ruff.format] +docstring-code-format = true + +[tool.ruff.lint] +select = [ + "D", + "UP", + "C4", + "TID251", + "TID253", + "T20", + "B904", + "N804", + "N805", + "I", +] +ignore = [ + "D100", + "D101", + "D102", + "D103", + "D104", + "D105", + "D205", + "D401", + "D404", +] + +[tool.ruff.lint.pydocstyle] +convention = "numpy" diff --git a/tests/entrypoint/test_workflow.py b/tests/entrypoint/test_workflow.py index 5b48bcca..81ed3d91 100644 --- a/tests/entrypoint/test_workflow.py +++ b/tests/entrypoint/test_workflow.py @@ -3,9 +3,11 @@ import shutil import textwrap import unittest +from unittest import ( + mock, +) import dflow -import mock from dflow import ( Workflow, ) diff --git a/tests/exploration/test_conf_filter.py b/tests/exploration/test_conf_filter.py index a0c36ba1..3a0fb343 100644 --- a/tests/exploration/test_conf_filter.py +++ b/tests/exploration/test_conf_filter.py @@ -1,14 +1,14 @@ import os import unittest +from unittest.mock import ( + patch, +) import dpdata import numpy as np from fake_data_set import ( fake_system, ) -from mock import ( - patch, -) # isort: off from .context import ( diff --git a/tests/exploration/test_customized_lmp_templ_task_group.py b/tests/exploration/test_customized_lmp_templ_task_group.py index 57462759..ab02dd8c 100644 --- a/tests/exploration/test_customized_lmp_templ_task_group.py +++ b/tests/exploration/test_customized_lmp_templ_task_group.py @@ -6,10 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/exploration/test_exploration_group.py b/tests/exploration/test_exploration_group.py index f71ce8f4..d60c70fb 100644 --- a/tests/exploration/test_exploration_group.py +++ b/tests/exploration/test_exploration_group.py @@ -4,10 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/exploration/test_exploration_scheduler.py b/tests/exploration/test_exploration_scheduler.py index 324904b0..a2c75863 100644 --- a/tests/exploration/test_exploration_scheduler.py +++ b/tests/exploration/test_exploration_scheduler.py @@ -4,10 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/exploration/test_lmp_templ_task_group.py b/tests/exploration/test_lmp_templ_task_group.py index 6211137d..91442674 100644 --- a/tests/exploration/test_lmp_templ_task_group.py +++ b/tests/exploration/test_lmp_templ_task_group.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/exploration/test_make_task_group_from_config.py b/tests/exploration/test_make_task_group_from_config.py index f9fe93be..988b5aa8 100644 --- a/tests/exploration/test_make_task_group_from_config.py +++ b/tests/exploration/test_make_task_group_from_config.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/exploration/test_report_adaptive_lower.py b/tests/exploration/test_report_adaptive_lower.py index b5f123ca..14c8a27e 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -4,8 +4,10 @@ from collections import ( Counter, ) +from unittest import ( + mock, +) -import mock import numpy as np from dargs import ( Argument, @@ -52,8 +54,8 @@ def test_fv(self): for idx, ii in enumerate(expected_fail_): for jj in ii: expected_fail.add((idx, jj)) - expected_cand = set([(0, 5), (0, 6), (1, 8), (1, 0), (1, 5)]) - expected_accu = set([(0, 1), (1, 6), (1, 7)]) + expected_cand = {(0, 5), (0, 6), (1, 8), (1, 0), (1, 5)} + expected_accu = {(0, 1), (1, 6), (1, 7)} ter = ExplorationReportAdaptiveLower( level_f_hi=0.7, @@ -135,10 +137,18 @@ def test_f(self): 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)] - ) + expected_cand = {(0, 6), (0, 7), (0, 5)} + expected_accu = { + (0, 1), + (0, 3), + (0, 4), + (1, 0), + (1, 1), + (1, 5), + (1, 6), + (1, 7), + (1, 8), + } ter = ExplorationReportAdaptiveLower( level_f_hi=0.7, @@ -186,7 +196,7 @@ def test_f_inv_pop(self): [(0, 6), (0, 7), (0, 5)] + [(0, 1), (0, 3), (0, 4), (1, 0), (1, 1), (1, 5), (1, 6), (1, 7), (1, 8)] ) - expected_accu = set([]) + expected_accu = set() ter = ExplorationReportAdaptiveLower( level_f_hi=0.7, @@ -255,10 +265,18 @@ def test_v(self): 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)] - ) + expected_cand = {(0, 6), (0, 7), (0, 5)} + expected_accu = { + (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, diff --git a/tests/fp/data.vasp.kp.gf/make_kp_test.py b/tests/fp/data.vasp.kp.gf/make_kp_test.py index dddf4584..7d365bcd 100644 --- a/tests/fp/data.vasp.kp.gf/make_kp_test.py +++ b/tests/fp/data.vasp.kp.gf/make_kp_test.py @@ -22,5 +22,5 @@ def make_one(out_dir): ntest = 30 for ii in range(ntest): - out_dir = "test.%03d" % ii + out_dir = f"test.{ii:03d}" make_one(out_dir) diff --git a/tests/fp/test_abacus.py b/tests/fp/test_abacus.py index 32d6c979..d4008626 100644 --- a/tests/fp/test_abacus.py +++ b/tests/fp/test_abacus.py @@ -40,8 +40,9 @@ def test_abacus(self): data_path / "INPUT", {"Na": data_path / "Na_ONCV_PBE-1.0.upf"} ), "run": { - "command": "cp -r %s OUT.ABACUS && cat %s" - % (data_path / "OUT.ABACUS", data_path / "log"), + "command": "cp -r {} OUT.ABACUS && cat {}".format( + data_path / "OUT.ABACUS", data_path / "log" + ), }, "extra_output_files": [], } diff --git a/tests/fp/test_cp2k.py b/tests/fp/test_cp2k.py index 22c87b28..8f12400c 100644 --- a/tests/fp/test_cp2k.py +++ b/tests/fp/test_cp2k.py @@ -45,8 +45,9 @@ def test_cp2k(self): fp_config = { "inputs": FpOpCp2kInputs(data_path / "input.inp"), "run": { - "command": "cp -r %s output.log && cat %s" - % (data_path / "output.log", data_path / "output.log"), + "command": "cp -r {} output.log && cat {}".format( + data_path / "output.log", data_path / "output.log" + ), }, "extra_output_files": [], } diff --git a/tests/fp/test_run_vasp.py b/tests/fp/test_run_vasp.py index 7202d6ce..7e3638a9 100644 --- a/tests/fp/test_run_vasp.py +++ b/tests/fp/test_run_vasp.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index 6e6cfadf..8d60f524 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -18,9 +18,7 @@ Path, ) from typing import ( - List, Optional, - Tuple, ) try: @@ -366,9 +364,9 @@ def execute( task_id = int(ip["task_name"].split(".")[1]) assert task_path.is_dir() assert ip["task_name"] in str(ip["task_path"]) - assert ( - len(models) == mocked_numb_models - ), f"{len(models)} == {mocked_numb_models}" + assert len(models) == mocked_numb_models, ( + f"{len(models)} == {mocked_numb_models}" + ) for ii in range(mocked_numb_models): assert ip["models"][ii].is_file() assert "model" in str(ip["models"][ii]) @@ -863,11 +861,11 @@ def __init__( def select( self, - trajs: List[Path], - model_devis: List[Path], - type_map: List[str] = None, - optional_outputs: Optional[List[Path]] = None, - ) -> Tuple[List[Path], ExplorationReport]: + trajs: list[Path], + model_devis: list[Path], + type_map: list[str] = None, + optional_outputs: Optional[list[Path]] = None, + ) -> tuple[list[Path], ExplorationReport]: confs = [] if len(trajs) == mocked_numb_lmp_tasks: # get output from prep_run_lmp @@ -988,7 +986,7 @@ def execute( Path("step").write_text("2") else: step_num = Path("step").read_text().strip() - Path("step").write_text(f"{int(step_num)+1}") + Path("step").write_text(f"{int(step_num) + 1}") if qhull_input is None: Path("test_qconvex.in").write_text("") diff --git a/tests/op/test_collect_data.py b/tests/op/test_collect_data.py index 2e558193..8e78a59b 100644 --- a/tests/op/test_collect_data.py +++ b/tests/op/test_collect_data.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import dpdata import numpy as np @@ -18,11 +25,6 @@ fake_multi_sys, fake_system, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( diff --git a/tests/op/test_collect_run_caly.py b/tests/op/test_collect_run_caly.py index 22ec37db..610aa875 100644 --- a/tests/op/test_collect_run_caly.py +++ b/tests/op/test_collect_run_caly.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( diff --git a/tests/op/test_prep_caly_dp_optim.py b/tests/op/test_prep_caly_dp_optim.py index d1ffd724..4f5743de 100644 --- a/tests/op/test_prep_caly_dp_optim.py +++ b/tests/op/test_prep_caly_dp_optim.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( @@ -67,7 +69,7 @@ def setUp(self): self.template_slice_config = {"group_size": 3} self.group_size = self.template_slice_config["group_size"] - grouped_poscar_list = [i for i in range(0, nposcar, self.group_size)] + grouped_poscar_list = list(range(0, nposcar, self.group_size)) self.ngrouped = len(grouped_poscar_list) self.ref_task_dirs = [] for i in range(0, self.ngrouped): diff --git a/tests/op/test_prep_caly_input.py b/tests/op/test_prep_caly_input.py index 0fac49c4..1f06ba19 100644 --- a/tests/op/test_prep_caly_input.py +++ b/tests/op/test_prep_caly_input.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from dpgen2.constants import ( diff --git a/tests/op/test_prep_caly_model_devi.py b/tests/op/test_prep_caly_model_devi.py index 6cc6337f..8253d808 100644 --- a/tests/op/test_prep_caly_model_devi.py +++ b/tests/op/test_prep_caly_model_devi.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( diff --git a/tests/op/test_prep_dp_train.py b/tests/op/test_prep_dp_train.py index a380e221..8fb3d532 100644 --- a/tests/op/test_prep_dp_train.py +++ b/tests/op/test_prep_dp_train.py @@ -4,6 +4,9 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) import numpy as np from dflow.python import ( @@ -12,9 +15,6 @@ Artifact, OPIOSign, ) -from mock import ( - mock, -) # isort: off from .context import ( diff --git a/tests/op/test_prep_relax.py b/tests/op/test_prep_relax.py index 83dcc2b7..45c3de1d 100644 --- a/tests/op/test_prep_relax.py +++ b/tests/op/test_prep_relax.py @@ -18,7 +18,7 @@ class TestPrepRelax(unittest.TestCase): def testPrepRelax(self): cifs = [] for i in range(4): - p = Path("%i.cif" % i) + p = Path(f"{i:d}.cif") p.write_text("Mocked cif.") cifs.append(p) op_in = OPIO( @@ -34,13 +34,13 @@ def testPrepRelax(self): self.assertEqual(op_out["ntasks"], 2) self.assertEqual(len(op_out["task_paths"]), 2) for i, task_path in enumerate(op_out["task_paths"]): - self.assertEqual(str(task_path), "task.%06d" % i) + self.assertEqual(str(task_path), f"task.{i:06d}") self.assertEqual(len(list(task_path.iterdir())), 2) def tearDown(self): for i in range(2): - if os.path.isdir("task.%06d" % i): - shutil.rmtree("task.%06d" % i) + if os.path.isdir(f"task.{i:06d}"): + shutil.rmtree(f"task.{i:06d}") for i in range(4): - if os.path.isfile("%s.cif" % i): - os.remove("%s.cif" % i) + if os.path.isfile(f"{i}.cif"): + os.remove(f"{i}.cif") diff --git a/tests/op/test_run_caly_dp_optim.py b/tests/op/test_run_caly_dp_optim.py index ee6afd43..0105ef15 100644 --- a/tests/op/test_run_caly_dp_optim.py +++ b/tests/op/test_run_caly_dp_optim.py @@ -4,6 +4,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import numpy as np from dflow.python import ( @@ -13,11 +20,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 45ba950c..d2838d79 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -6,6 +6,10 @@ from pathlib import ( Path, ) +from unittest.mock import ( + call, + patch, +) import dpdata import numpy as np @@ -21,10 +25,6 @@ fake_multi_sys, fake_system, ) -from mock import ( - call, - patch, -) # isort: off from .context import ( @@ -76,7 +76,7 @@ def setUp(self): ss_0.to_deepmd_npy("init/data-0") ss_1.to_deepmd_npy("init/data-1") self.init_data = [Path("init/data-0"), Path("init/data-1")] - self.init_data = sorted(list(self.init_data)) + self.init_data = sorted(self.init_data) self.init_model = Path("bar.pb") @@ -868,7 +868,7 @@ def setUp(self): ss_0.to_deepmd_npy("init/data-0") ss_1.to_deepmd_npy("init/data-1") self.init_data = [Path("init/data-0"), Path("init/data-1")] - self.init_data = sorted(list(self.init_data)) + self.init_data = sorted(self.init_data) self.init_model = Path("bar.pb") diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 650fd82e..d3caebdd 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -5,6 +5,13 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) +from unittest.mock import ( + call, + patch, +) import dpdata import numpy as np @@ -15,11 +22,6 @@ OPIOSign, TransientError, ) -from mock import ( - call, - mock, - patch, -) # isort: off from .context import ( @@ -249,7 +251,7 @@ def test_success(self, mocked_run): ) # The number of models have to be 2 in knowledge distillation - self.assertEqual(len(list((work_dir.glob("*.pb")))), 2) + self.assertEqual(len(list(work_dir.glob("*.pb"))), 2) def swap_element(arg): @@ -316,7 +318,7 @@ def tearDown(self): class TestMergePIMDFiles(unittest.TestCase): def test_merge_pimd_files(self): for i in range(1, 3): - with open("traj.%s.dump" % i, "w") as f: + with open(f"traj.{i}.dump", "w") as f: f.write( """ITEM: TIMESTEP 0 @@ -345,7 +347,7 @@ def test_merge_pimd_files(self): """ ) for i in range(1, 3): - with open("model_devi.%s.out" % i, "w") as f: + with open(f"model_devi.{i}.out", "w") as f: f.write( """# step max_devi_v min_devi_v avg_devi_v max_devi_f min_devi_f avg_devi_f 0 9.023897e-17 3.548771e-17 5.237314e-17 8.196123e-16 1.225653e-16 3.941002e-16 diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py index fa0571a7..bfdc4cf0 100644 --- a/tests/test_block_cl.py +++ b/tests/test_block_cl.py @@ -7,10 +7,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np @@ -45,6 +41,10 @@ except ModuleNotFoundError: # case of upload everything to argo, no context needed pass +from unittest.mock import ( + patch, +) + from context import ( default_host, default_image, @@ -52,9 +52,6 @@ skip_ut_with_dflow_reason, upload_python_packages, ) -from mock import ( - patch, -) from mocked_ops import ( MockedCollectData, MockedCollectDataCheckOptParam, diff --git a/tests/test_caly_evo_step.py b/tests/test_caly_evo_step.py index 2aff3614..51dafecd 100644 --- a/tests/test_caly_evo_step.py +++ b/tests/test_caly_evo_step.py @@ -7,10 +7,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np diff --git a/tests/test_collect_data.py b/tests/test_collect_data.py index 95b7e4c9..5f10e1db 100644 --- a/tests/test_collect_data.py +++ b/tests/test_collect_data.py @@ -6,10 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np @@ -108,8 +104,8 @@ def test(self): @unittest.skipIf(skip_ut_with_dflow, skip_ut_with_dflow_reason) class TestMockedCollectDataArgo(unittest.TestCase): def setUp(self): - self.iter_data = set(("foo/iter0", "bar/iter1")) - self.iter_data = set([Path(ii) for ii in self.iter_data]) + self.iter_data = {"foo/iter0", "bar/iter1"} + self.iter_data = {Path(ii) for ii in self.iter_data} self.name = "outdata" self.labeled_data = ["d0", "d1"] self.labeled_data = [Path(ii) for ii in self.labeled_data] diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py index 49b4c873..5c830891 100644 --- a/tests/test_dpgen_loop.py +++ b/tests/test_dpgen_loop.py @@ -8,10 +8,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np diff --git a/tests/test_merge_caly_evo_step.py b/tests/test_merge_caly_evo_step.py index 3f22ba56..57c54a36 100644 --- a/tests/test_merge_caly_evo_step.py +++ b/tests/test_merge_caly_evo_step.py @@ -7,10 +7,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py index e949410d..adcc1b61 100644 --- a/tests/test_prep_run_caly.py +++ b/tests/test_prep_run_caly.py @@ -7,10 +7,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np diff --git a/tests/test_prep_run_diffcsp.py b/tests/test_prep_run_diffcsp.py index d5569d17..5f4da22c 100644 --- a/tests/test_prep_run_diffcsp.py +++ b/tests/test_prep_run_diffcsp.py @@ -36,10 +36,10 @@ def execute( self, ip: OPIO, ) -> OPIO: - task_dir = Path("diffcsp.%s" % ip["task_id"]) + task_dir = Path("diffcsp.{}".format(ip["task_id"])) task_dir.mkdir(exist_ok=True) for i in range(2): - fpath = task_dir / ("%s.cif" % i) + fpath = task_dir / (f"{i}.cif") fpath.write_text("Mocked cif.") return OPIO( { @@ -60,10 +60,10 @@ def execute( model_devis = [] for cif in cifs: name = cif[:-4] - traj = ip["task_path"] / ("traj.%s.dump" % name) + traj = ip["task_path"] / (f"traj.{name}.dump") traj.write_text("Mocked traj.") trajs.append(traj) - model_devi = ip["task_path"] / ("model_devi.%s.out" % name) + model_devi = ip["task_path"] / (f"model_devi.{name}.out") model_devi.write_text("Mocked model_devi.") model_devis.append(model_devi) return OPIO( diff --git a/tests/test_prep_run_dp_labeling.py b/tests/test_prep_run_dp_labeling.py index 67256f86..4a3882aa 100644 --- a/tests/test_prep_run_dp_labeling.py +++ b/tests/test_prep_run_dp_labeling.py @@ -5,6 +5,10 @@ from pathlib import ( Path, ) +from unittest.mock import ( + Mock, + patch, +) import dpdata import numpy as np @@ -14,10 +18,6 @@ from dflow.python import ( FatalError, ) -from mock import ( - Mock, - patch, -) from dpgen2.fp.deepmd import ( PrepDeepmd, diff --git a/tests/test_prep_run_dp_train.py b/tests/test_prep_run_dp_train.py index 536ca4b9..4f235d19 100644 --- a/tests/test_prep_run_dp_train.py +++ b/tests/test_prep_run_dp_train.py @@ -6,10 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np from dflow import ( @@ -99,7 +95,7 @@ def _check_log( lines[1 + ii].split(" "), [ "data", - str(revised_fname(Path(path) / sorted(list(init_data))[ii])), + str(revised_fname(Path(path) / sorted(init_data)[ii])), "OK", ], ) @@ -108,7 +104,7 @@ def _check_log( lines[3 + ii].split(" "), [ "data", - str(revised_fname(Path(path) / sorted(list(iter_data))[ii])), + str(revised_fname(Path(path) / sorted(iter_data)[ii])), "OK", ], ) diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py index 3b350240..f60f8689 100644 --- a/tests/test_prep_run_lmp.py +++ b/tests/test_prep_run_lmp.py @@ -7,10 +7,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np @@ -167,7 +163,7 @@ def setUp(self): def check_run_lmp_output( self, task_name: str, - models: List[Path], + models: list[Path], ): cwd = os.getcwd() os.chdir(task_name) @@ -246,7 +242,7 @@ def tearDown(self): def check_run_lmp_output( self, task_name: str, - models: List[Path], + models: list[Path], ): cwd = os.getcwd() os.chdir(task_name) diff --git a/tests/test_prep_run_vasp.py b/tests/test_prep_run_vasp.py index 77030bff..b4ab332c 100644 --- a/tests/test_prep_run_vasp.py +++ b/tests/test_prep_run_vasp.py @@ -6,10 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import jsonpickle import numpy as np diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py index 491d42f7..216d8697 100644 --- a/tests/test_select_confs.py +++ b/tests/test_select_confs.py @@ -6,11 +6,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, - Tuple, -) import jsonpickle import numpy as np diff --git a/tests/utils/test_dflow_query.py b/tests/utils/test_dflow_query.py index dad71643..a6df2f6e 100644 --- a/tests/utils/test_dflow_query.py +++ b/tests/utils/test_dflow_query.py @@ -4,10 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) import numpy as np diff --git a/tests/utils/test_dl_dpgen2_arti.py b/tests/utils/test_dl_dpgen2_arti.py index 037a0d80..0d9a5377 100644 --- a/tests/utils/test_dl_dpgen2_arti.py +++ b/tests/utils/test_dl_dpgen2_arti.py @@ -7,10 +7,12 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) import dflow import dpdata -import mock import numpy as np # isort: off diff --git a/tests/utils/test_dl_dpgen2_arti_by_def.py b/tests/utils/test_dl_dpgen2_arti_by_def.py index e6a30a32..23fe4742 100644 --- a/tests/utils/test_dl_dpgen2_arti_by_def.py +++ b/tests/utils/test_dl_dpgen2_arti_by_def.py @@ -8,10 +8,12 @@ from pathlib import ( Path, ) +from unittest import ( + mock, +) import dflow import dpdata -import mock import numpy as np # isort: off