diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 50a42b5c..536522cc 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -23,7 +23,7 @@ jobs: pip install -e .[test] pip install mock coverage pytest - name: Test - run: SKIP_UT_WITH_DFLOW=0 DFLOW_DEBUG=1 coverage run --source=./dpgen2 -m unittest -v -f && coverage report + run: SKIP_UT_WITH_DFLOW=0 DFLOW_DEBUG=1 coverage run --source=./dpgen2 -m unittest -v && coverage report - uses: codecov/codecov-action@v5 env: CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} diff --git a/.gitignore b/.gitignore index f471ce9e..864269e7 100644 --- a/.gitignore +++ b/.gitignore @@ -49,6 +49,7 @@ tests/caly-* tests/coll-* tests/dpgen-* tests/upload +upload/ tests/init_data/ tests/incar tests/potcar diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a96c7d1d..4c713eb4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,8 +17,10 @@ repos: # Python - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.1.3 + rev: v0.12.12 hooks: + - id: ruff + args: [--fix] - id: ruff-format - repo: https://github.com/PyCQA/isort rev: 5.12.0 diff --git a/docs/conf.py b/docs/conf.py index b68a1726..c91b989d 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -13,13 +13,14 @@ import os import sys from datetime import ( - date, + datetime, + timezone, ) # -- Project information ----------------------------------------------------- project = "DPGEN2" -copyright = "2022-%d, DeepModeling" % date.today().year +copyright = "2022-%d, DeepModeling" % datetime.now(timezone.utc).year author = "DeepModeling" diff --git a/dpgen2/__init__.py b/dpgen2/__init__.py index 4217e612..c8a20d76 100644 --- a/dpgen2/__init__.py +++ b/dpgen2/__init__.py @@ -1,6 +1,4 @@ try: from ._version import version as __version__ except ImportError: - from .__about__ import ( - __version__, - ) + from .__about__ import __version__ as __version__ diff --git a/dpgen2/conf/__init__.py b/dpgen2/conf/__init__.py index 059469c7..77d35f86 100644 --- a/dpgen2/conf/__init__.py +++ b/dpgen2/conf/__init__.py @@ -12,3 +12,10 @@ "alloy": AlloyConfGenerator, "file": FileConfGenerator, } + +__all__ = [ + "AlloyConfGenerator", + "ConfGenerator", + "FileConfGenerator", + "conf_styles", +] diff --git a/dpgen2/conf/alloy_conf.py b/dpgen2/conf/alloy_conf.py index a885e623..59c4048f 100644 --- a/dpgen2/conf/alloy_conf.py +++ b/dpgen2/conf/alloy_conf.py @@ -5,7 +5,6 @@ ) from typing import ( List, - Optional, Tuple, Union, ) @@ -14,7 +13,6 @@ import numpy as np from dargs import ( Argument, - Variant, ) from .conf_generator import ( diff --git a/dpgen2/conf/file_conf.py b/dpgen2/conf/file_conf.py index e8ae2443..270f3d67 100644 --- a/dpgen2/conf/file_conf.py +++ b/dpgen2/conf/file_conf.py @@ -1,19 +1,16 @@ import glob -import os from pathlib import ( Path, ) from typing import ( List, Optional, - Tuple, Union, ) import dpdata from dargs import ( Argument, - Variant, ) from .conf_generator import ( diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index df11ff7f..32e8af4e 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -9,13 +9,9 @@ Variant, ) -import dpgen2 from dpgen2.conf import ( conf_styles, ) -from dpgen2.constants import ( - default_image, -) from dpgen2.exploration.report import ( conv_styles, ) @@ -83,7 +79,6 @@ def dp_dist_train_args(): def dp_train_args(): doc_numb_models = "Number of models trained for evaluating the model deviation" - doc_config = "Configuration of training" doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. " doc_init_models_paths = "the paths to initial models" doc_init_models_uri = "The URI of initial models" diff --git a/dpgen2/entrypoint/common.py b/dpgen2/entrypoint/common.py index 0d0af9e8..85c2388d 100644 --- a/dpgen2/entrypoint/common.py +++ b/dpgen2/entrypoint/common.py @@ -3,9 +3,7 @@ Path, ) from typing import ( - Dict, List, - Optional, Union, ) @@ -13,14 +11,8 @@ from dpgen2.utils import ( bohrium_config_from_dict, - dump_object_to_file, - load_object_from_file, - matched_step_key, - print_keys_in_nice_format, - sort_slice_ops, workflow_config_from_dict, ) -from dpgen2.utils.step_config import normalize as normalize_step_dict def global_config_workflow( diff --git a/dpgen2/entrypoint/download.py b/dpgen2/entrypoint/download.py index 7f095039..b67ad37e 100644 --- a/dpgen2/entrypoint/download.py +++ b/dpgen2/entrypoint/download.py @@ -3,7 +3,6 @@ Dict, List, Optional, - Union, ) from dflow import ( @@ -14,9 +13,6 @@ from dpgen2.entrypoint.common import ( global_config_workflow, ) -from dpgen2.utils.dflow_query import ( - matched_step_key, -) from dpgen2.utils.download_dpgen2_artifacts import ( download_dpgen2_artifacts, download_dpgen2_artifacts_by_def, diff --git a/dpgen2/entrypoint/main.py b/dpgen2/entrypoint/main.py index 08493917..2c8c0250 100644 --- a/dpgen2/entrypoint/main.py +++ b/dpgen2/entrypoint/main.py @@ -1,22 +1,12 @@ import argparse import json import logging -import os import textwrap from typing import ( List, Optional, ) -import dflow -from dflow import ( - Step, - Steps, - Workflow, - download_artifact, - upload_artifact, -) - from dpgen2 import ( __version__, ) @@ -41,11 +31,8 @@ status, ) from .submit import ( - make_concurrent_learning_op, - make_naive_exploration_scheduler, resubmit_concurrent_learning, submit_concurrent_learning, - workflow_concurrent_learning, ) from .watch import ( default_watching_keys, @@ -327,7 +314,7 @@ def main(): logging.basicConfig(level=logging.INFO) args = parse_args() - dict_args = vars(args) + vars(args) if args.command == "submit": with open(args.CONFIG) as fp: diff --git a/dpgen2/entrypoint/showkey.py b/dpgen2/entrypoint/showkey.py index 36055426..291dce3a 100644 --- a/dpgen2/entrypoint/showkey.py +++ b/dpgen2/entrypoint/showkey.py @@ -1,11 +1,6 @@ -import glob -import os -import pickle -from pathlib import ( - Path, -) +import functools +import operator -import dpdata from dflow import ( Workflow, ) @@ -32,7 +27,7 @@ def showkey( wf = Workflow(id=wf_id) folded_keys = get_resubmit_keys(wf) - all_step_keys = sum(folded_keys.values(), []) + all_step_keys = functools.reduce(operator.iadd, folded_keys.values(), []) prt_str = print_keys_in_nice_format( all_step_keys, ["run-train", "run-lmp", "run-fp", "diffcsp-gen", "run-relax"], diff --git a/dpgen2/entrypoint/status.py b/dpgen2/entrypoint/status.py index d6164647..7ae02f25 100644 --- a/dpgen2/entrypoint/status.py +++ b/dpgen2/entrypoint/status.py @@ -1,9 +1,7 @@ import logging from typing import ( Dict, - List, Optional, - Union, ) from dflow import ( @@ -15,7 +13,6 @@ global_config_workflow, ) from dpgen2.utils.dflow_query import ( - get_all_schedulers, get_last_scheduler, ) diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 203478dc..e52abc82 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -1,10 +1,9 @@ import copy -import glob +import functools import json import logging +import operator import os -import pickle -import re from copy import ( deepcopy, ) @@ -15,44 +14,20 @@ Dict, List, Optional, - Tuple, - Type, - Union, ) -import dpdata from dflow import ( ArgoStep, - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, S3Artifact, Step, - Steps, Workflow, - argo_range, - download_artifact, upload_artifact, ) -from dflow.python import ( - OP, - OPIO, - Artifact, - FatalError, - OPIOSign, - PythonOPTemplate, - TransientError, - upload_packages, -) from dpgen2.conf import ( conf_styles, ) from dpgen2.constants import ( - default_host, default_image, ) from dpgen2.entrypoint.args import normalize as normalize_args @@ -65,7 +40,6 @@ TrajRenderLammps, ) from dpgen2.exploration.report import ( - ExplorationReportTrustLevelsRandom, conv_styles, ) from dpgen2.exploration.scheduler import ( @@ -78,11 +52,7 @@ conf_filter_styles, ) from dpgen2.exploration.task import ( - CustomizedLmpTemplateTaskGroup, ExplorationStage, - ExplorationTask, - LmpTemplateTaskGroup, - NPTTaskGroup, caly_normalize, diffcsp_normalize, make_calypso_task_group_from_config, @@ -131,16 +101,12 @@ ) from dpgen2.utils import ( BinaryFileInput, - bohrium_config_from_dict, - dump_object_to_file, get_artifact_from_uri, get_subkey, - load_object_from_file, matched_step_key, print_keys_in_nice_format, sort_slice_ops, upload_artifact_and_print_uri, - workflow_config_from_dict, ) from dpgen2.utils.step_config import normalize as normalize_step_dict @@ -465,14 +431,14 @@ def get_systems_from_data(data, data_prefix=None): assert isinstance(data, list) if data_prefix is not None: data = [os.path.join(data_prefix, ii) for ii in data] - data = sum([expand_sys_str(ii) for ii in data], []) + data = functools.reduce(operator.iadd, [expand_sys_str(ii) for ii in data], []) return data def workflow_concurrent_learning( config: Dict, ) -> Step: - default_config = config["default_step_config"] + config["default_step_config"] train_config = config["train"]["config"] explore_config = config["explore"]["config"] @@ -571,9 +537,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"] ) @@ -586,12 +552,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"] - ), f"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"] ) @@ -677,13 +643,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 @@ -891,7 +857,7 @@ def resubmit_concurrent_learning( # reuse the super OP iif all steps within it are reused if set(v) == set(folded_keys[k]): reused_folded_keys[k] = [k] - reused_keys = sum(reused_folded_keys.values(), []) + reused_keys = functools.reduce(operator.iadd, reused_folded_keys.values(), []) reuse_step = old_wf.query_step(key=reused_keys, sort_by_generation=True) wf = submit_concurrent_learning( diff --git a/dpgen2/entrypoint/watch.py b/dpgen2/entrypoint/watch.py index 891be817..7df13fc1 100644 --- a/dpgen2/entrypoint/watch.py +++ b/dpgen2/entrypoint/watch.py @@ -4,7 +4,6 @@ Dict, List, Optional, - Union, ) from dflow import ( @@ -43,15 +42,15 @@ def update_finished_steps( if finished_keys is not None: diff_keys = [] for kk in wf_keys: - if not (kk in finished_keys): + if kk not in finished_keys: diff_keys.append(kk) else: diff_keys = wf_keys for kk in diff_keys: - logging.info(f'steps {kk.ljust(50,"-")} finished') + logging.info(f"steps {kk.ljust(50, '-')} finished") if download: download_dpgen2_artifacts(wf, kk, prefix=prefix, chk_pnt=chk_pnt) - logging.info(f'steps {kk.ljust(50,"-")} downloaded') + logging.info(f"steps {kk.ljust(50, '-')} downloaded") finished_keys = wf_keys return finished_keys diff --git a/dpgen2/entrypoint/workflow.py b/dpgen2/entrypoint/workflow.py index 8cc4cc04..4b3e2012 100644 --- a/dpgen2/entrypoint/workflow.py +++ b/dpgen2/entrypoint/workflow.py @@ -1,7 +1,4 @@ import argparse -import json -import logging -import os from typing import ( Optional, ) diff --git a/dpgen2/exploration/deviation/__init__.py b/dpgen2/exploration/deviation/__init__.py index 537c2f94..8640e3b7 100644 --- a/dpgen2/exploration/deviation/__init__.py +++ b/dpgen2/exploration/deviation/__init__.py @@ -1,6 +1,2 @@ -from .deviation_manager import ( - DeviManager, -) -from .deviation_std import ( - DeviManagerStd, -) +from .deviation_manager import DeviManager as DeviManager +from .deviation_std import DeviManagerStd as DeviManagerStd diff --git a/dpgen2/exploration/deviation/deviation_std.py b/dpgen2/exploration/deviation/deviation_std.py index f927b75b..c574629e 100644 --- a/dpgen2/exploration/deviation/deviation_std.py +++ b/dpgen2/exploration/deviation/deviation_std.py @@ -2,7 +2,6 @@ defaultdict, ) from typing import ( - Dict, List, Optional, ) @@ -33,12 +32,12 @@ 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" + + "one-dimensional array" ) self._data[name].append(deviation) @@ -73,7 +72,7 @@ def _check_data(self) -> None: assert len(self._data[name]) == self.ntraj, ( f"Error: the number of model deviation {name} " + f"({len(self._data[name])}) and trajectory files ({self.ntraj}) " - + f"are not equal." + + "are not equal." ) for idx, ndarray in enumerate(self._data[name]): assert isinstance(ndarray, np.ndarray), ( @@ -86,9 +85,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/__init__.py b/dpgen2/exploration/render/__init__.py index e926c484..36e6ddce 100644 --- a/dpgen2/exploration/render/__init__.py +++ b/dpgen2/exploration/render/__init__.py @@ -4,3 +4,8 @@ from .traj_render_lammps import ( TrajRenderLammps, ) + +__all__ = [ + "TrajRender", + "TrajRenderLammps", +] diff --git a/dpgen2/exploration/render/traj_render.py b/dpgen2/exploration/render/traj_render.py index 5c9f0c41..236b6efd 100644 --- a/dpgen2/exploration/render/traj_render.py +++ b/dpgen2/exploration/render/traj_render.py @@ -9,12 +9,10 @@ TYPE_CHECKING, List, Optional, - Tuple, Union, ) import dpdata -import numpy as np from dflow.python.opio import ( HDF5Dataset, ) diff --git a/dpgen2/exploration/render/traj_render_lammps.py b/dpgen2/exploration/render/traj_render_lammps.py index 00b6a3de..e2afc6a9 100644 --- a/dpgen2/exploration/render/traj_render_lammps.py +++ b/dpgen2/exploration/render/traj_render_lammps.py @@ -9,7 +9,6 @@ TYPE_CHECKING, List, Optional, - Tuple, Union, ) diff --git a/dpgen2/exploration/report/__init__.py b/dpgen2/exploration/report/__init__.py index 282068c2..dc28350f 100644 --- a/dpgen2/exploration/report/__init__.py +++ b/dpgen2/exploration/report/__init__.py @@ -16,3 +16,11 @@ "fixed-levels-max-select": ExplorationReportTrustLevelsMax, "adaptive-lower": ExplorationReportAdaptiveLower, } + +__all__ = [ + "ExplorationReport", + "ExplorationReportAdaptiveLower", + "ExplorationReportTrustLevelsMax", + "ExplorationReportTrustLevelsRandom", + "conv_styles", +] diff --git a/dpgen2/exploration/report/report.py b/dpgen2/exploration/report/report.py index d1c43fe6..62bce76d 100644 --- a/dpgen2/exploration/report/report.py +++ b/dpgen2/exploration/report/report.py @@ -5,11 +5,8 @@ from typing import ( List, Optional, - Tuple, ) -import numpy as np - from ..deviation import ( DeviManager, ) diff --git a/dpgen2/exploration/report/report_adaptive_lower.py b/dpgen2/exploration/report/report_adaptive_lower.py index cd2989f3..848da3dd 100644 --- a/dpgen2/exploration/report/report_adaptive_lower.py +++ b/dpgen2/exploration/report/report_adaptive_lower.py @@ -446,7 +446,8 @@ def _get_candidates_inv_pop_f( self.candi_picked = [(ii[0], ii[1]) for ii in self.candi] if max_nframes is not None and max_nframes < len(self.candi_picked): prob = self._choice_prob_inv_pop_f(self.candi_picked) - indices = np.random.choice( + rng = np.random.default_rng() + indices = rng.choice( len(self.candi_picked), size=max_nframes, replace=False, diff --git a/dpgen2/exploration/report/report_trust_levels_base.py b/dpgen2/exploration/report/report_trust_levels_base.py index 185ea5b1..87f1e732 100644 --- a/dpgen2/exploration/report/report_trust_levels_base.py +++ b/dpgen2/exploration/report/report_trust_levels_base.py @@ -1,11 +1,9 @@ -import random from abc import ( abstractmethod, ) from typing import ( List, Optional, - Tuple, ) import numpy as np diff --git a/dpgen2/exploration/report/report_trust_levels_max.py b/dpgen2/exploration/report/report_trust_levels_max.py index e847c1e6..ab63f598 100644 --- a/dpgen2/exploration/report/report_trust_levels_max.py +++ b/dpgen2/exploration/report/report_trust_levels_max.py @@ -1,18 +1,9 @@ -import random from typing import ( List, Optional, Tuple, ) -import numpy as np -from dargs import ( - Argument, -) -from dflow.python import ( - FatalError, -) - from ..deviation import ( DeviManager, ) diff --git a/dpgen2/exploration/report/report_trust_levels_random.py b/dpgen2/exploration/report/report_trust_levels_random.py index 540ad0d2..5733a19f 100644 --- a/dpgen2/exploration/report/report_trust_levels_random.py +++ b/dpgen2/exploration/report/report_trust_levels_random.py @@ -5,17 +5,6 @@ Tuple, ) -import numpy as np -from dargs import ( - Argument, -) -from dflow.python import ( - FatalError, -) - -from ..deviation import ( - DeviManager, -) from . import ( ExplorationReport, ) diff --git a/dpgen2/exploration/scheduler/__init__.py b/dpgen2/exploration/scheduler/__init__.py index 5c05e006..29fc1f86 100644 --- a/dpgen2/exploration/scheduler/__init__.py +++ b/dpgen2/exploration/scheduler/__init__.py @@ -7,3 +7,9 @@ from .stage_scheduler import ( StageScheduler, ) + +__all__ = [ + "ConvergenceCheckStageScheduler", + "ExplorationScheduler", + "StageScheduler", +] diff --git a/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py b/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py index 8ab8662f..04c5c9ff 100644 --- a/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py +++ b/dpgen2/exploration/scheduler/convergence_check_stage_scheduler.py @@ -24,7 +24,6 @@ from dpgen2.exploration.task import ( BaseExplorationTaskGroup, ExplorationStage, - ExplorationTaskGroup, ) from .stage_scheduler import ( diff --git a/dpgen2/exploration/scheduler/scheduler.py b/dpgen2/exploration/scheduler/scheduler.py index ff55fa23..d079ddf6 100644 --- a/dpgen2/exploration/scheduler/scheduler.py +++ b/dpgen2/exploration/scheduler/scheduler.py @@ -23,7 +23,6 @@ ConfSelector, ) from dpgen2.exploration.task import ( - ExplorationStage, ExplorationTaskGroup, ) @@ -252,8 +251,8 @@ def print_last_iteration(self, print_header=False): ) if self.complete(): - ret.append(f"# All stages converged") - return "\n".join(ret + [""]) + ret.append("# All stages converged") + return "\n".join([*ret, ""]) def print_convergence(self): ret = [] @@ -288,5 +287,5 @@ def print_convergence(self): _summary = self._print_prev_summary(prev_stg_idx) assert _summary is not None ret.append(_summary) - ret.append(f"# All stages converged") - return "\n".join(ret + [""]) + ret.append("# All stages converged") + return "\n".join([*ret, ""]) diff --git a/dpgen2/exploration/scheduler/stage_scheduler.py b/dpgen2/exploration/scheduler/stage_scheduler.py index 18fe5593..41ea237d 100644 --- a/dpgen2/exploration/scheduler/stage_scheduler.py +++ b/dpgen2/exploration/scheduler/stage_scheduler.py @@ -22,7 +22,6 @@ ConfSelector, ) from dpgen2.exploration.task import ( - ExplorationStage, ExplorationTaskGroup, ) diff --git a/dpgen2/exploration/selector/__init__.py b/dpgen2/exploration/selector/__init__.py index cfed094f..1d3326a3 100644 --- a/dpgen2/exploration/selector/__init__.py +++ b/dpgen2/exploration/selector/__init__.py @@ -19,3 +19,14 @@ "box_skewness": BoxSkewnessConfFilter, "box_length": BoxLengthFilter, } + +__all__ = [ + "BoxLengthFilter", + "BoxSkewnessConfFilter", + "ConfFilter", + "ConfFilters", + "ConfSelector", + "ConfSelectorFrames", + "DistanceConfFilter", + "conf_filter_styles", +] diff --git a/dpgen2/exploration/selector/conf_filter.py b/dpgen2/exploration/selector/conf_filter.py index 3fca483f..dedf3f7b 100644 --- a/dpgen2/exploration/selector/conf_filter.py +++ b/dpgen2/exploration/selector/conf_filter.py @@ -11,7 +11,6 @@ ) import dpdata -import numpy as np class ConfFilter(ABC): diff --git a/dpgen2/exploration/selector/conf_selector.py b/dpgen2/exploration/selector/conf_selector.py index f24a7d31..b8d69941 100644 --- a/dpgen2/exploration/selector/conf_selector.py +++ b/dpgen2/exploration/selector/conf_selector.py @@ -8,12 +8,10 @@ from typing import ( List, Optional, - Set, Tuple, Union, ) -import dpdata from dflow.python.opio import ( HDF5Dataset, ) @@ -22,10 +20,6 @@ ExplorationReport, ) -from . import ( - ConfFilters, -) - class ConfSelector(ABC): """Select configurations from trajectory and model deviation files.""" diff --git a/dpgen2/exploration/selector/conf_selector_frame.py b/dpgen2/exploration/selector/conf_selector_frame.py index fc116f88..752fa173 100644 --- a/dpgen2/exploration/selector/conf_selector_frame.py +++ b/dpgen2/exploration/selector/conf_selector_frame.py @@ -1,7 +1,4 @@ import copy -from collections import ( - Counter, -) from pathlib import ( Path, ) @@ -12,8 +9,6 @@ Union, ) -import dpdata -import numpy as np from dflow.python.opio import ( HDF5Dataset, ) diff --git a/dpgen2/exploration/selector/distance_conf_filter.py b/dpgen2/exploration/selector/distance_conf_filter.py index 7748e70d..4778c765 100644 --- a/dpgen2/exploration/selector/distance_conf_filter.py +++ b/dpgen2/exploration/selector/distance_conf_filter.py @@ -148,9 +148,6 @@ def check( from ase import ( Atoms, ) - from ase.build import ( - make_supercell, - ) safe_dist = deepcopy(safe_dist_dict) safe_dist.update(self.custom_safe_dist) diff --git a/dpgen2/exploration/task/__init__.py b/dpgen2/exploration/task/__init__.py index 534a8828..d083ddff 100644 --- a/dpgen2/exploration/task/__init__.py +++ b/dpgen2/exploration/task/__init__.py @@ -41,3 +41,25 @@ BaseExplorationTaskGroup, ExplorationTaskGroup, ) + +__all__ = [ + "BaseExplorationTaskGroup", + "CalyTaskGroup", + "ConfSamplingTaskGroup", + "CustomizedLmpTemplateTaskGroup", + "DiffCSPTaskGroup", + "ExplorationStage", + "ExplorationTask", + "ExplorationTaskGroup", + "LmpTemplateTaskGroup", + "NPTTaskGroup", + "caly_normalize", + "caly_task_group_args", + "diffcsp_normalize", + "lmp_task_group_args", + "make_calypso_task_group_from_config", + "make_diffcsp_task_group_from_config", + "make_lmp_task_group_from_config", + "normalize_lmp_task_group_config", + "variant_task_group", +] diff --git a/dpgen2/exploration/task/caly_task_group.py b/dpgen2/exploration/task/caly_task_group.py index 14594f38..b9cdfec0 100644 --- a/dpgen2/exploration/task/caly_task_group.py +++ b/dpgen2/exploration/task/caly_task_group.py @@ -124,7 +124,7 @@ def set_params( if any(map(lambda s: (set(s) - overlap) == 0, name_of_atoms)): raise ValueError( - f"Any sub-list should not equal with intersection, e.g. [[A,B,C], [B,C], [C]] is not allowed." + "Any sub-list should not equal with intersection, e.g. [[A,B,C], [B,C], [C]] is not allowed." ) while True: diff --git a/dpgen2/exploration/task/calypso/__init__.py b/dpgen2/exploration/task/calypso/__init__.py index 61f3cdcd..33d15515 100644 --- a/dpgen2/exploration/task/calypso/__init__.py +++ b/dpgen2/exploration/task/calypso/__init__.py @@ -4,3 +4,10 @@ calypso_run_opt_str_end, make_calypso_input, ) + +__all__ = [ + "calypso_check_opt_str", + "calypso_run_opt_str", + "calypso_run_opt_str_end", + "make_calypso_input", +] diff --git a/dpgen2/exploration/task/calypso/caly_input.py b/dpgen2/exploration/task/calypso/caly_input.py index 66f2b511..7b31530a 100644 --- a/dpgen2/exploration/task/calypso/caly_input.py +++ b/dpgen2/exploration/task/calypso/caly_input.py @@ -1,19 +1,8 @@ -import random from typing import ( List, - Optional, ) -import dpdata import numpy as np -import scipy.constants as pc -from packaging.version import ( - Version, -) - -from dpgen2.constants import ( - lmp_traj_name, -) calypso_run_opt_str = """#!/usr/bin/env python3 @@ -322,9 +311,9 @@ def make_calypso_input( file_str = "" for key, value in necessary_keys.items(): - file_str += f"{key} = {str(value)}\n" + file_str += f"{key} = {value!s}\n" for key, value in default_key_value.items(): - file_str += f"{key} = {str(value)}\n" + file_str += f"{key} = {value!s}\n" file_str += "@DistanceOfIon\n" file_str += distance_of_ions_str + "\n" file_str += "@End\n" diff --git a/dpgen2/exploration/task/conf_sampling_task_group.py b/dpgen2/exploration/task/conf_sampling_task_group.py index 4c5ee0c3..e933e5d6 100644 --- a/dpgen2/exploration/task/conf_sampling_task_group.py +++ b/dpgen2/exploration/task/conf_sampling_task_group.py @@ -1,19 +1,9 @@ -import itertools import random from typing import ( List, Optional, ) -from dpgen2.constants import ( - lmp_conf_name, - lmp_input_name, - model_name_pattern, -) - -from .task import ( - ExplorationTask, -) from .task_group import ( ExplorationTaskGroup, ) diff --git a/dpgen2/exploration/task/customized_lmp_template_task_group.py b/dpgen2/exploration/task/customized_lmp_template_task_group.py index d7022516..39e895da 100644 --- a/dpgen2/exploration/task/customized_lmp_template_task_group.py +++ b/dpgen2/exploration/task/customized_lmp_template_task_group.py @@ -18,8 +18,6 @@ from dpgen2.constants import ( lmp_conf_name, lmp_input_name, - model_name_pattern, - plm_input_name, ) from dpgen2.utils import ( run_command, @@ -29,15 +27,9 @@ from .conf_sampling_task_group import ( ConfSamplingTaskGroup, ) -from .lmp import ( - make_lmp_input, -) from .lmp_template_task_group import ( LmpTemplateTaskGroup, ) -from .task import ( - ExplorationTask, -) class CustomizedLmpTemplateTaskGroup(ConfSamplingTaskGroup): diff --git a/dpgen2/exploration/task/lmp/__init__.py b/dpgen2/exploration/task/lmp/__init__.py index ceaa824a..bd66130e 100644 --- a/dpgen2/exploration/task/lmp/__init__.py +++ b/dpgen2/exploration/task/lmp/__init__.py @@ -1,3 +1,7 @@ from .lmp_input import ( make_lmp_input, ) + +__all__ = [ + "make_lmp_input", +] diff --git a/dpgen2/exploration/task/lmp/lmp_input.py b/dpgen2/exploration/task/lmp/lmp_input.py index c2a22b60..eb11caa7 100644 --- a/dpgen2/exploration/task/lmp/lmp_input.py +++ b/dpgen2/exploration/task/lmp/lmp_input.py @@ -20,8 +20,9 @@ def _sample_sphere(): + rng = np.random.default_rng() while True: - vv = np.array([np.random.normal(), np.random.normal(), np.random.normal()]) + vv = np.array([rng.normal(), rng.normal(), rng.normal()]) vn = np.linalg.norm(vv) if vn < 0.2: continue diff --git a/dpgen2/exploration/task/lmp_template_task_group.py b/dpgen2/exploration/task/lmp_template_task_group.py index d1f8e9fc..e4b29cf9 100644 --- a/dpgen2/exploration/task/lmp_template_task_group.py +++ b/dpgen2/exploration/task/lmp_template_task_group.py @@ -1,5 +1,4 @@ import itertools -import random from pathlib import ( Path, ) @@ -23,9 +22,6 @@ from .conf_sampling_task_group import ( ConfSamplingTaskGroup, ) -from .lmp import ( - make_lmp_input, -) from .task import ( ExplorationTask, ) @@ -198,9 +194,9 @@ 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 diff --git a/dpgen2/exploration/task/make_task_group_from_config.py b/dpgen2/exploration/task/make_task_group_from_config.py index 2859ac8f..77bd2eea 100644 --- a/dpgen2/exploration/task/make_task_group_from_config.py +++ b/dpgen2/exploration/task/make_task_group_from_config.py @@ -1,4 +1,3 @@ -import dargs from dargs import ( Argument, Variant, @@ -7,7 +6,6 @@ from dpgen2.constants import ( lmp_conf_name, lmp_input_name, - model_name_pattern, plm_input_name, ) from dpgen2.exploration.task import ( diff --git a/dpgen2/exploration/task/npt_task_group.py b/dpgen2/exploration/task/npt_task_group.py index ada5a15a..5ff96e24 100644 --- a/dpgen2/exploration/task/npt_task_group.py +++ b/dpgen2/exploration/task/npt_task_group.py @@ -1,5 +1,4 @@ import itertools -import random from pathlib import ( Path, ) diff --git a/dpgen2/exploration/task/stage.py b/dpgen2/exploration/task/stage.py index 17ab150e..e2f78602 100644 --- a/dpgen2/exploration/task/stage.py +++ b/dpgen2/exploration/task/stage.py @@ -1,20 +1,3 @@ -from abc import ( - ABC, - abstractmethod, -) -from typing import ( - List, -) - -from dpgen2.constants import ( - lmp_conf_name, - lmp_input_name, - model_name_pattern, -) - -from .task import ( - ExplorationTask, -) from .task_group import ( BaseExplorationTaskGroup, ExplorationTaskGroup, diff --git a/dpgen2/exploration/task/task.py b/dpgen2/exploration/task/task.py index 24f7d02f..ca30f343 100644 --- a/dpgen2/exploration/task/task.py +++ b/dpgen2/exploration/task/task.py @@ -1,11 +1,5 @@ -import os -from collections.abc import ( - Sequence, -) from typing import ( Dict, - List, - Tuple, ) @@ -17,7 +11,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) """ diff --git a/dpgen2/exploration/task/task_group.py b/dpgen2/exploration/task/task_group.py index c603f80b..7f8adc7e 100644 --- a/dpgen2/exploration/task/task_group.py +++ b/dpgen2/exploration/task/task_group.py @@ -6,9 +6,7 @@ Sequence, ) from typing import ( - Dict, List, - Tuple, ) from .task import ( diff --git a/dpgen2/flow/__init__.py b/dpgen2/flow/__init__.py index 98c490c2..2a77f3ab 100644 --- a/dpgen2/flow/__init__.py +++ b/dpgen2/flow/__init__.py @@ -1,3 +1,7 @@ from .dpgen_loop import ( ConcurrentLearning, ) + +__all__ = [ + "ConcurrentLearning", +] diff --git a/dpgen2/flow/dpgen_loop.py b/dpgen2/flow/dpgen_loop.py index 190a1090..17ab597e 100644 --- a/dpgen2/flow/dpgen_loop.py +++ b/dpgen2/flow/dpgen_loop.py @@ -1,5 +1,4 @@ import os -import pickle from copy import ( deepcopy, ) @@ -12,7 +11,6 @@ Union, ) -import jsonpickle from dflow import ( InputArtifact, InputParameter, @@ -23,13 +21,7 @@ Outputs, Step, Steps, - Workflow, - argo_len, - argo_range, - argo_sequence, - download_artifact, if_expression, - upload_artifact, ) from dflow.python import ( OP, @@ -39,7 +31,6 @@ HDF5Datasets, OPIOSign, PythonOPTemplate, - Slices, ) from dpgen2.exploration.report import ( @@ -57,10 +48,6 @@ from dpgen2.superop.block import ( ConcurrentLearningBlock, ) -from dpgen2.utils import ( - dump_object_to_file, - load_object_from_file, -) from dpgen2.utils.step_config import ( init_executor, ) @@ -153,7 +140,7 @@ def execute( ) -> OPIO: scheduler = ip["exploration_scheduler"] - stage = scheduler.get_stage() + scheduler.get_stage() iteration = scheduler.get_iteration() return OPIO( @@ -338,7 +325,7 @@ def init_keys(self): @property def loop_keys(self): - return [self.loop_key] + self.loop.keys + return [self.loop_key, *self.loop.keys] def _loop( @@ -456,17 +443,17 @@ def _loop( steps.outputs.parameters[ "exploration_scheduler" ].value_from_expression = if_expression( - _if=(scheduler_step.outputs.parameters["converged"] == True), + _if=(scheduler_step.outputs.parameters["converged"]), _then=scheduler_step.outputs.parameters["exploration_scheduler"], _else=next_step.outputs.parameters["exploration_scheduler"], ) steps.outputs.artifacts["models"].from_expression = if_expression( - _if=(scheduler_step.outputs.parameters["converged"] == True), + _if=(scheduler_step.outputs.parameters["converged"]), _then=block_step.outputs.artifacts["models"], _else=next_step.outputs.artifacts["models"], ) steps.outputs.artifacts["iter_data"].from_expression = if_expression( - _if=(scheduler_step.outputs.parameters["converged"] == True), + _if=(scheduler_step.outputs.parameters["converged"]), _then=block_step.outputs.artifacts["iter_data"], _else=next_step.outputs.artifacts["iter_data"], ) diff --git a/dpgen2/fp/abacus.py b/dpgen2/fp/abacus.py index 28769b01..2940dd18 100644 --- a/dpgen2/fp/abacus.py +++ b/dpgen2/fp/abacus.py @@ -3,6 +3,7 @@ ) from typing import ( List, + Tuple, ) import dpdata @@ -132,11 +133,6 @@ 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]: suffix = "ABACUS" calculation = "scf" diff --git a/dpgen2/fp/deepmd.py b/dpgen2/fp/deepmd.py index 43fb200a..5b894124 100644 --- a/dpgen2/fp/deepmd.py +++ b/dpgen2/fp/deepmd.py @@ -1,4 +1,5 @@ """Prep and Run Gaussian tasks.""" + import os from pathlib import ( Path, @@ -6,7 +7,6 @@ from typing import ( Any, List, - Optional, Tuple, ) @@ -18,7 +18,6 @@ ) from dflow.python import ( FatalError, - TransientError, ) from dpgen2.constants import ( diff --git a/dpgen2/fp/gaussian.py b/dpgen2/fp/gaussian.py index b6aba200..72ab6a74 100644 --- a/dpgen2/fp/gaussian.py +++ b/dpgen2/fp/gaussian.py @@ -1,4 +1,5 @@ """Prep and Run Gaussian tasks.""" + import logging from typing import ( Any, diff --git a/dpgen2/fp/prep_fp.py b/dpgen2/fp/prep_fp.py index a43f582a..962be031 100644 --- a/dpgen2/fp/prep_fp.py +++ b/dpgen2/fp/prep_fp.py @@ -1,5 +1,3 @@ -import json -import os from abc import ( ABC, abstractmethod, @@ -9,11 +7,8 @@ ) from typing import ( Any, - Dict, List, - Set, Tuple, - Union, ) import dpdata diff --git a/dpgen2/fp/run_fp.py b/dpgen2/fp/run_fp.py index 64e0c3a8..a18c45f6 100644 --- a/dpgen2/fp/run_fp.py +++ b/dpgen2/fp/run_fp.py @@ -1,4 +1,3 @@ -import json import os from abc import ( ABC, @@ -10,12 +9,10 @@ from typing import ( Dict, List, - Set, Tuple, ) import dargs -import dpdata from dflow.python import ( OP, OPIO, @@ -23,7 +20,6 @@ BigParameter, FatalError, OPIOSign, - TransientError, ) from dpgen2.utils.chdir import ( diff --git a/dpgen2/fp/vasp.py b/dpgen2/fp/vasp.py index 8d8cebbe..5fc15394 100644 --- a/dpgen2/fp/vasp.py +++ b/dpgen2/fp/vasp.py @@ -6,29 +6,16 @@ Path, ) from typing import ( - Dict, List, - Optional, - Set, Tuple, - Union, ) import dpdata import numpy as np from dargs import ( Argument, - ArgumentEncoder, - Variant, - dargs, ) from dflow.python import ( - OP, - OPIO, - Artifact, - BigParameter, - FatalError, - OPIOSign, TransientError, ) @@ -49,7 +36,6 @@ ) from .vasp_input import ( VaspInputs, - make_kspacing_kpoints, ) # global static variables diff --git a/dpgen2/fp/vasp_input.py b/dpgen2/fp/vasp_input.py index 8bbad820..81b6cdd5 100644 --- a/dpgen2/fp/vasp_input.py +++ b/dpgen2/fp/vasp_input.py @@ -4,19 +4,12 @@ from typing import ( Dict, List, - Optional, - Set, - Tuple, Union, ) -import dpdata import numpy as np from dargs import ( Argument, - ArgumentEncoder, - Variant, - dargs, ) diff --git a/dpgen2/op/__init__.py b/dpgen2/op/__init__.py index f4fec3a2..f28bed81 100644 --- a/dpgen2/op/__init__.py +++ b/dpgen2/op/__init__.py @@ -45,3 +45,23 @@ from .select_confs import ( SelectConfs, ) + +__all__ = [ + "CollRunCaly", + "CollectData", + "DiffCSPGen", + "PrepCalyDPOptim", + "PrepCalyInput", + "PrepCalyModelDevi", + "PrepDPTrain", + "PrepLmp", + "PrepRelax", + "RunCalyDPOptim", + "RunCalyModelDevi", + "RunDPTrain", + "RunLmp", + "RunLmpHDF5", + "RunRelax", + "RunRelaxHDF5", + "SelectConfs", +] diff --git a/dpgen2/op/caly_evo_step_merge.py b/dpgen2/op/caly_evo_step_merge.py index 8abda6dd..0520e808 100644 --- a/dpgen2/op/caly_evo_step_merge.py +++ b/dpgen2/op/caly_evo_step_merge.py @@ -1,13 +1,8 @@ -import json -import logging -import pickle -import shutil from pathlib import ( Path, ) from typing import ( List, - Tuple, ) from dflow import ( @@ -24,29 +19,14 @@ OPIOSign, Parameter, Slices, - TransientError, ) from dflow.utils import ( flatten, ) -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_run_opt_file, -) -from dpgen2.exploration.task import ( - ExplorationTaskGroup, -) from dpgen2.superop.caly_evo_step import ( CalyEvoStep, ) -from dpgen2.utils import ( - BinaryFileInput, - set_directory, -) -from dpgen2.utils.run_command import ( - run_command, -) class CalyEvoStepMerge(OP): diff --git a/dpgen2/op/collect_data.py b/dpgen2/op/collect_data.py index 68f397ff..ebffb0ea 100644 --- a/dpgen2/op/collect_data.py +++ b/dpgen2/op/collect_data.py @@ -1,12 +1,9 @@ -import json -import os from pathlib import ( Path, ) from typing import ( + ClassVar, List, - Set, - Tuple, ) import dpdata @@ -14,7 +11,6 @@ OP, OPIO, Artifact, - BigParameter, OPIOSign, Parameter, ) @@ -35,7 +31,7 @@ class CollectData(OP): """ - default_optional_parameter = { + default_optional_parameter: ClassVar = { "mixed_type": False, } diff --git a/dpgen2/op/collect_run_caly.py b/dpgen2/op/collect_run_caly.py index 4b6148f6..a8bd93f7 100644 --- a/dpgen2/op/collect_run_caly.py +++ b/dpgen2/op/collect_run_caly.py @@ -1,31 +1,20 @@ -import json import logging -import os -import random -import re import shutil from pathlib import ( Path, ) from typing import ( List, - Optional, - Set, - Tuple, ) from dargs import ( Argument, - ArgumentEncoder, - Variant, - dargs, ) from dflow.python import ( OP, OPIO, Artifact, BigParameter, - FatalError, OPIOSign, Parameter, TransientError, @@ -35,7 +24,6 @@ calypso_log_name, ) from dpgen2.utils import ( - BinaryFileInput, set_directory, ) from dpgen2.utils.run_command import ( diff --git a/dpgen2/op/prep_caly_dp_optim.py b/dpgen2/op/prep_caly_dp_optim.py index d2e4d8b0..aaf48acc 100644 --- a/dpgen2/op/prep_caly_dp_optim.py +++ b/dpgen2/op/prep_caly_dp_optim.py @@ -1,13 +1,8 @@ -import json -import logging -import pickle -import shutil from pathlib import ( Path, ) from typing import ( List, - Tuple, ) from dflow.python import ( @@ -17,25 +12,11 @@ BigParameter, OPIOSign, Parameter, - TransientError, ) -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_opt_dir_name, - calypso_run_opt_file, - model_name_pattern, -) -from dpgen2.exploration.task import ( - ExplorationTaskGroup, -) from dpgen2.utils import ( - BinaryFileInput, set_directory, ) -from dpgen2.utils.run_command import ( - run_command, -) class PrepCalyDPOptim(OP): diff --git a/dpgen2/op/prep_caly_input.py b/dpgen2/op/prep_caly_input.py index e3da359a..f5080506 100644 --- a/dpgen2/op/prep_caly_input.py +++ b/dpgen2/op/prep_caly_input.py @@ -1,11 +1,8 @@ -import json -import pickle from pathlib import ( Path, ) from typing import ( List, - Tuple, ) from dflow.python import ( @@ -22,14 +19,9 @@ calypso_input_file, calypso_run_opt_file, calypso_task_pattern, - model_name_pattern, ) from dpgen2.exploration.task import ( BaseExplorationTaskGroup, - ExplorationTaskGroup, -) -from dpgen2.utils import ( - set_directory, ) vsc_keys = { diff --git a/dpgen2/op/prep_caly_model_devi.py b/dpgen2/op/prep_caly_model_devi.py index 3b070959..2a4b46e5 100644 --- a/dpgen2/op/prep_caly_model_devi.py +++ b/dpgen2/op/prep_caly_model_devi.py @@ -1,13 +1,8 @@ -import json -import logging -import pickle -import shutil from pathlib import ( Path, ) from typing import ( List, - Tuple, ) from dflow.python import ( @@ -17,25 +12,11 @@ BigParameter, OPIOSign, Parameter, - TransientError, ) -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_opt_dir_name, - calypso_run_opt_file, - model_name_pattern, -) -from dpgen2.exploration.task import ( - ExplorationTaskGroup, -) from dpgen2.utils import ( - BinaryFileInput, set_directory, ) -from dpgen2.utils.run_command import ( - run_command, -) class PrepCalyModelDevi(OP): diff --git a/dpgen2/op/prep_dp_train.py b/dpgen2/op/prep_dp_train.py index 20fe58c2..4f497b59 100644 --- a/dpgen2/op/prep_dp_train.py +++ b/dpgen2/op/prep_dp_train.py @@ -6,7 +6,6 @@ ) from typing import ( List, - Tuple, Union, ) @@ -79,7 +78,7 @@ def execute( template = ip["template_script"] numb_models = ip["numb_models"] osubdirs = [] - if type(template) != list: + if not isinstance(template, list): template = [template for ii in range(numb_models)] else: if not (len(template) == numb_models): diff --git a/dpgen2/op/prep_lmp.py b/dpgen2/op/prep_lmp.py index e1b5c026..7a5be422 100644 --- a/dpgen2/op/prep_lmp.py +++ b/dpgen2/op/prep_lmp.py @@ -1,11 +1,8 @@ -import json -import pickle from pathlib import ( Path, ) from typing import ( List, - Tuple, ) from dflow.python import ( @@ -21,7 +18,6 @@ ) from dpgen2.exploration.task import ( BaseExplorationTaskGroup, - ExplorationTaskGroup, ) diff --git a/dpgen2/op/run_caly_dp_optim.py b/dpgen2/op/run_caly_dp_optim.py index 639c97d5..6d2de21e 100644 --- a/dpgen2/op/run_caly_dp_optim.py +++ b/dpgen2/op/run_caly_dp_optim.py @@ -1,14 +1,8 @@ -import json import logging -import pickle import shutil from pathlib import ( Path, ) -from typing import ( - List, - Tuple, -) from dflow.python import ( OP, @@ -20,15 +14,7 @@ TransientError, ) -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_run_opt_file, -) -from dpgen2.exploration.task import ( - ExplorationTaskGroup, -) from dpgen2.utils import ( - BinaryFileInput, set_directory, ) from dpgen2.utils.run_command import ( diff --git a/dpgen2/op/run_caly_model_devi.py b/dpgen2/op/run_caly_model_devi.py index 9e191cf2..d78201cd 100644 --- a/dpgen2/op/run_caly_model_devi.py +++ b/dpgen2/op/run_caly_model_devi.py @@ -14,7 +14,6 @@ OP, OPIO, Artifact, - BigParameter, OPIOSign, Parameter, ) @@ -107,7 +106,7 @@ def execute( if atoms_list is None: continue for atoms in atoms_list: - natoms = len(atoms) + len(atoms) dump_str = atoms2lmpdump(atoms, tcount, type_map, ignore=True) dump_str_dict[tcount].append(dump_str) @@ -128,9 +127,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 @@ -284,9 +283,8 @@ def parse_traj(traj_file): if len(trajs) >= 20: selected_traj = [trajs[iii] for iii in [4, 9, -10, -5, -1]] elif 5 <= len(trajs) < 20: - selected_traj = [ - trajs[np.random.randint(3, len(trajs) - 1)] for _ in range(4) - ] + rng = np.random.default_rng() + selected_traj = [trajs[rng.integers(3, len(trajs) - 1)] for _ in range(4)] selected_traj.append(trajs[-1]) elif 3 <= len(trajs) < 5: selected_traj = [trajs[round((len(trajs) - 1) / 2)]] diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 5a9782f4..c36b263a 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -1,7 +1,9 @@ +import functools import glob import json import logging import math +import operator import os import random import shutil @@ -9,19 +11,16 @@ Path, ) from typing import ( + ClassVar, Dict, List, Optional, - Tuple, Union, ) import dpdata from dargs import ( Argument, - ArgumentEncoder, - Variant, - dargs, ) from dflow.python import ( OP, @@ -32,12 +31,10 @@ NestedDict, OPIOSign, Parameter, - TransientError, ) from dpgen2.constants import ( train_script_name, - train_task_pattern, ) from dpgen2.utils.chdir import ( set_directory, @@ -69,7 +66,7 @@ def _make_train_command( checkpoint = None # case of restart if checkpoint is not None: - command = dp_command + ["train", "--restart", checkpoint, train_script_name] + command = [*dp_command, "train", "--restart", checkpoint, train_script_name] return command # case of init model and finetune assert checkpoint is None @@ -79,25 +76,18 @@ def _make_train_command( ) if case_init_model: init_flag = "--init-frz-model" if impl == "tensorflow" else "--init-model" - command = dp_command + [ + command = [*dp_command, "train", init_flag, str(init_model), train_script_name] + elif case_finetune: + command = [ + *dp_command, "train", - init_flag, - str(init_model), train_script_name, + "--finetune", + str(init_model), + *finetune_args.split(), ] - elif case_finetune: - command = ( - dp_command - + [ - "train", - train_script_name, - "--finetune", - str(init_model), - ] - + finetune_args.split() - ) else: - command = dp_command + ["train", train_script_name] + command = [*dp_command, "train", train_script_name] command += train_args.split() return command @@ -112,7 +102,7 @@ class RunDPTrain(OP): """ - default_optional_parameter = { + default_optional_parameter: ClassVar = { "mixed_type": False, "finetune_mode": "no", } @@ -236,7 +226,7 @@ def execute( len_init = len(init_data) numb_old = len_init + len(iter_data_old_exp) numb_new = numb_old + len(iter_data_new_exp) - auto_prob_str = f"prob_sys_size; 0:{numb_old}:{old_ratio}; {numb_old}:{numb_new}:{1.-old_ratio:g}" + auto_prob_str = f"prob_sys_size; 0:{numb_old}:{old_ratio}; {numb_old}:{numb_new}:{1.0 - old_ratio:g}" # update the input dict train_dict = RunDPTrain.write_data_to_input_script( @@ -488,7 +478,7 @@ def decide_init_model( old_data_size_level = int(config["init_model_policy"].split(":")[-1]) if isinstance(init_data, dict): init_data_size = _get_data_size_of_all_systems( - sum(init_data.values(), []) + functools.reduce(operator.iadd, init_data.values(), []) ) else: init_data_size = _get_data_size_of_all_systems(init_data) diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py index 60cd9305..108d53d7 100644 --- a/dpgen2/op/run_lmp.py +++ b/dpgen2/op/run_lmp.py @@ -10,30 +10,23 @@ from typing import ( List, Optional, - Set, - Tuple, ) import numpy as np from dargs import ( Argument, - ArgumentEncoder, - Variant, - dargs, ) from dflow.python import ( OP, OPIO, Artifact, BigParameter, - FatalError, HDF5Datasets, OPIOSign, TransientError, ) from dpgen2.constants import ( - lmp_conf_name, lmp_input_name, lmp_log_name, lmp_model_devi_name, @@ -132,13 +125,13 @@ 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) - model_files = [Path(teacher_model_file).resolve()] + model_files + model_files = [Path(teacher_model_file).resolve(), *model_files] with set_directory(work_dir): # link input files @@ -306,10 +299,10 @@ 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 ") + raise RuntimeError("last matching index should not be -1, terribly wrong ") for ii in range(match_last, len(new_line_split)): if re.fullmatch(pattern, new_line_split[ii]) is not None: raise RuntimeError( diff --git a/dpgen2/op/run_relax.py b/dpgen2/op/run_relax.py index 672275d8..6bae6679 100644 --- a/dpgen2/op/run_relax.py +++ b/dpgen2/op/run_relax.py @@ -1,4 +1,3 @@ -import logging import os from pathlib import ( Path, diff --git a/dpgen2/op/select_confs.py b/dpgen2/op/select_confs.py index e8ba891d..4cf8dcc3 100644 --- a/dpgen2/op/select_confs.py +++ b/dpgen2/op/select_confs.py @@ -1,12 +1,8 @@ -import json -import os from pathlib import ( Path, ) from typing import ( List, - Set, - Tuple, Union, ) @@ -111,11 +107,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/__init__.py b/dpgen2/superop/__init__.py index 0223605f..4082140e 100644 --- a/dpgen2/superop/__init__.py +++ b/dpgen2/superop/__init__.py @@ -16,3 +16,12 @@ from .prep_run_lmp import ( PrepRunLmp, ) + +__all__ = [ + "ConcurrentLearningBlock", + "PrepRunCaly", + "PrepRunDPTrain", + "PrepRunDiffCSP", + "PrepRunFp", + "PrepRunLmp", +] diff --git a/dpgen2/superop/block.py b/dpgen2/superop/block.py index 0e39ab38..50799550 100644 --- a/dpgen2/superop/block.py +++ b/dpgen2/superop/block.py @@ -2,15 +2,11 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( Any, Dict, List, Optional, - Set, Type, Union, ) @@ -25,25 +21,12 @@ Outputs, Step, Steps, - Workflow, - argo_len, - argo_range, - argo_sequence, - download_artifact, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, - Slices, ) -from dpgen2.op import ( - CollectData, -) from dpgen2.utils.step_config import ( init_executor, ) diff --git a/dpgen2/superop/caly_evo_step.py b/dpgen2/superop/caly_evo_step.py index cd2be501..49826bee 100644 --- a/dpgen2/superop/caly_evo_step.py +++ b/dpgen2/superop/caly_evo_step.py @@ -2,13 +2,9 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( List, Optional, - Set, Type, ) @@ -17,30 +13,17 @@ InputParameter, Inputs, OutputArtifact, - OutputParameter, Outputs, Step, Steps, - Workflow, - argo_len, - argo_range, - argo_sequence, - download_artifact, if_expression, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) -from dpgen2.constants import ( - calypso_index_pattern, -) from dpgen2.utils.step_config import ( init_executor, ) @@ -141,7 +124,7 @@ def _caly_evo_step( ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) - prep_template_config = prep_config.pop("template_config") + prep_config.pop("template_config") run_template_config = run_config.pop("template_config") prep_executor_config = prep_config.pop("executor") run_executor_config = run_config.pop("executor") diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py index daa48143..6745b610 100644 --- a/dpgen2/superop/prep_run_calypso.py +++ b/dpgen2/superop/prep_run_calypso.py @@ -2,9 +2,6 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( Any, Dict, @@ -20,29 +17,17 @@ Inputs, OPTemplate, OutputArtifact, - OutputParameter, Outputs, Step, Steps, - Workflow, - argo_len, argo_range, - argo_sequence, - download_artifact, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) -from dpgen2.constants import ( - calypso_index_pattern, -) from dpgen2.utils.step_config import ( init_executor, ) diff --git a/dpgen2/superop/prep_run_diffcsp.py b/dpgen2/superop/prep_run_diffcsp.py index c44851fe..ba1bebbb 100644 --- a/dpgen2/superop/prep_run_diffcsp.py +++ b/dpgen2/superop/prep_run_diffcsp.py @@ -2,40 +2,24 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( - Any, - Dict, List, Optional, Type, - Union, ) from dflow import ( InputArtifact, InputParameter, Inputs, - OPTemplate, OutputArtifact, - OutputParameter, Outputs, Step, Steps, - Workflow, - argo_len, - argo_range, argo_sequence, - download_artifact, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) @@ -138,7 +122,7 @@ def _prep_run_diffcsp( block_id = prep_run_diffcsp_steps.inputs.parameters["block_id"] expl_task_grp = prep_run_diffcsp_steps.inputs.parameters["expl_task_grp"] expl_config = prep_run_diffcsp_steps.inputs.parameters["explore_config"] - type_map = prep_run_diffcsp_steps.inputs.parameters["type_map"] + prep_run_diffcsp_steps.inputs.parameters["type_map"] models = prep_run_diffcsp_steps.inputs.artifacts["models"] diffcsp_gen = Step( diff --git a/dpgen2/superop/prep_run_dp_train.py b/dpgen2/superop/prep_run_dp_train.py index 0fd988e4..6ff10af0 100644 --- a/dpgen2/superop/prep_run_dp_train.py +++ b/dpgen2/superop/prep_run_dp_train.py @@ -1,15 +1,10 @@ -import json import os from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( List, Optional, - Set, Type, ) @@ -23,27 +18,18 @@ S3Artifact, Step, Steps, - Workflow, argo_len, - argo_range, argo_sequence, - download_artifact, upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - BigParameter, - OPIOSign, PythonOPTemplate, Slices, ) from dpgen2.constants import ( train_index_pattern, - train_script_name, - train_task_pattern, ) from dpgen2.op import ( RunDPTrain, diff --git a/dpgen2/superop/prep_run_fp.py b/dpgen2/superop/prep_run_fp.py index ba659c6d..8c5294e8 100644 --- a/dpgen2/superop/prep_run_fp.py +++ b/dpgen2/superop/prep_run_fp.py @@ -2,13 +2,9 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( List, Optional, - Set, Type, ) @@ -21,18 +17,11 @@ Outputs, Step, Steps, - Workflow, argo_len, - argo_range, argo_sequence, - download_artifact, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) diff --git a/dpgen2/superop/prep_run_lmp.py b/dpgen2/superop/prep_run_lmp.py index 3e0a0a0f..c392bd8f 100644 --- a/dpgen2/superop/prep_run_lmp.py +++ b/dpgen2/superop/prep_run_lmp.py @@ -2,13 +2,9 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) from typing import ( List, Optional, - Set, Type, ) @@ -21,18 +17,11 @@ Outputs, Step, Steps, - Workflow, argo_len, - argo_range, argo_sequence, - download_artifact, - upload_artifact, ) from dflow.python import ( OP, - OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) diff --git a/dpgen2/utils/__init__.py b/dpgen2/utils/__init__.py index b7824f57..af0c9299 100644 --- a/dpgen2/utils/__init__.py +++ b/dpgen2/utils/__init__.py @@ -45,3 +45,31 @@ from .step_config import ( step_conf_args, ) + +__all__ = [ + "BinaryFileInput", + "bohrium_config_from_dict", + "chdir", + "dflow_config", + "dflow_s3_config", + "dump_object_to_file", + "find_slice_ranges", + "gen_doc_step_dict", + "get_artifact_from_uri", + "get_iteration", + "get_last_iteration", + "get_last_scheduler", + "get_subkey", + "init_executor", + "load_object_from_file", + "matched_step_key", + "normalize_step_dict", + "print_keys_in_nice_format", + "run_command", + "set_directory", + "setup_ele_temp", + "sort_slice_ops", + "step_conf_args", + "upload_artifact_and_print_uri", + "workflow_config_from_dict", +] diff --git a/dpgen2/utils/binary_file_input.py b/dpgen2/utils/binary_file_input.py index db099a00..644b130d 100644 --- a/dpgen2/utils/binary_file_input.py +++ b/dpgen2/utils/binary_file_input.py @@ -1,38 +1,28 @@ """Binary file inputs""" + import os import warnings from pathlib import ( Path, ) from typing import ( - Any, - List, Optional, - Tuple, Union, ) -from dargs import ( - Argument, - dargs, -) -from dflow.python import ( - TransientError, -) - class BinaryFileInput: def __init__(self, path: Union[str, Path], ext: Optional[str] = None) -> None: path = str(path) - assert os.path.exists(path), f"No such file: {str(path)}" + assert os.path.exists(path), f"No such file: {path!s}" if ext and not ext.startswith("."): ext = "." + ext 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: {path!s}' + ) self.file_name = os.path.basename(path) with open(path, "rb") as f: @@ -42,7 +32,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/bohrium_config.py b/dpgen2/utils/bohrium_config.py index e94dc28a..c0574c09 100644 --- a/dpgen2/utils/bohrium_config.py +++ b/dpgen2/utils/bohrium_config.py @@ -1,7 +1,5 @@ import importlib -import os -import dflow from dflow.config import ( config, s3_config, 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..7ae65341 100644 --- a/dpgen2/utils/dflow_config.py +++ b/dpgen2/utils/dflow_config.py @@ -1,5 +1,3 @@ -import copy - from dflow.config import ( config, s3_config, diff --git a/dpgen2/utils/dflow_query.py b/dpgen2/utils/dflow_query.py index 933c8cc6..8e06ca06 100644 --- a/dpgen2/utils/dflow_query.py +++ b/dpgen2/utils/dflow_query.py @@ -6,8 +6,6 @@ Optional, ) -import numpy as np - def get_subkey( key: str, @@ -74,7 +72,7 @@ def get_last_scheduler( return None else: skey = sorted(scheduler_keys)[-1] - step = [step for step in scheduler_steps if step.key == skey][0] + step = next(step for step in scheduler_steps if step.key == skey) return step.outputs.parameters["exploration_scheduler"].value @@ -165,9 +163,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" - range_fmt = f"%d -> %d" - range_s_fmt = f"%{idx_fmt_len*2+4}s" + normal_fmt = f"%{idx_fmt_len * 2 + 4}d" + range_fmt = "%d -> %d" + range_s_fmt = f"%{idx_fmt_len * 2 + 4}s" idx = 0 ret = [] @@ -187,4 +185,4 @@ def print_keys_in_nice_format( except ValueError: ret.append((normal_fmt + " : " + "%s") % (idx, keys[idx])) idx += 1 - return "\n".join(ret + [""]) + return "\n".join([*ret, ""]) diff --git a/dpgen2/utils/download_dpgen2_artifacts.py b/dpgen2/utils/download_dpgen2_artifacts.py index 67c2aaf0..c81ec86f 100644 --- a/dpgen2/utils/download_dpgen2_artifacts.py +++ b/dpgen2/utils/download_dpgen2_artifacts.py @@ -8,7 +8,6 @@ Optional, ) -import numpy as np from dflow import ( Workflow, download_artifact, diff --git a/dpgen2/utils/run_command.py b/dpgen2/utils/run_command.py index 2d5c5764..69670184 100644 --- a/dpgen2/utils/run_command.py +++ b/dpgen2/utils/run_command.py @@ -1,4 +1,3 @@ -import os from typing import ( List, Tuple, diff --git a/dpgen2/utils/step_config.py b/dpgen2/utils/step_config.py index 591f474e..2fdae5cf 100644 --- a/dpgen2/utils/step_config.py +++ b/dpgen2/utils/step_config.py @@ -1,13 +1,7 @@ -import os - -import dargs from dargs import ( Argument, Variant, ) -from dflow.config import ( - config, -) from dflow.plugins.dispatcher import ( DispatcherExecutor, ) @@ -23,7 +17,7 @@ def dispatcher_args(): def variant_executor(): - doc = f"The type of the executor." + doc = "The type of the executor." return Variant( "type", [ diff --git a/pyproject.toml b/pyproject.toml index 2363c18d..f29e91f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,3 +80,34 @@ exclude = [ [tool.isort] profile = "black" force_grid_wrap = 1 + +[tool.ruff.format] +docstring-code-format = true + +[tool.ruff.lint] +select = [ + "E", # errors + "W", # warning + "F", # pyflakes + "RUF", # ruff + "NPY", # numpy + "DTZ", # datetime +] +ignore = [ + "E501", # line too long + "F841", # local variable is assigned to but never used + "E741", # ambiguous variable name + "E402", # module level import not at top of file + "D100", # TODO: missing docstring in public module + "D101", # TODO: missing docstring in public class + "D102", # TODO: missing docstring in public method + "D103", # TODO: missing docstring in public function + "D104", # TODO: missing docstring in public package + "D105", # TODO: missing docstring in magic method + "D205", # 1 blank line required between summary line and description + "D401", # TODO: first line should be in imperative mood + "D404", # TODO: first word of the docstring should not be This +] + +[tool.ruff.lint.per-file-ignores] +"tests/**/*.py" = ["W291"] # Ignore trailing whitespace in tests to preserve string content diff --git a/tests/conf/context.py b/tests/conf/context.py index d6638725..5504dedc 100644 --- a/tests/conf/context.py +++ b/tests/conf/context.py @@ -2,4 +2,3 @@ import sys sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 diff --git a/tests/conf/test_alloy_conf.py b/tests/conf/test_alloy_conf.py index 0dd4cc5a..cdf484bd 100644 --- a/tests/conf/test_alloy_conf.py +++ b/tests/conf/test_alloy_conf.py @@ -1,7 +1,4 @@ -import json -import os import random -import shutil import tempfile import unittest from pathlib import ( @@ -12,7 +9,6 @@ import numpy as np # isort: off -from .context import dpgen2 from dpgen2.conf.alloy_conf import ( AlloyConf, AlloyConfGenerator, diff --git a/tests/conf/test_file_conf.py b/tests/conf/test_file_conf.py index 381f0d88..5cef4a9b 100644 --- a/tests/conf/test_file_conf.py +++ b/tests/conf/test_file_conf.py @@ -1,8 +1,4 @@ -import json -import os -import random import shutil -import tempfile import textwrap import unittest from pathlib import ( @@ -13,9 +9,6 @@ import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.conf.file_conf import ( FileConfGenerator, ) diff --git a/tests/conf/test_unit_cell.py b/tests/conf/test_unit_cell.py index 19544f0e..72603835 100644 --- a/tests/conf/test_unit_cell.py +++ b/tests/conf/test_unit_cell.py @@ -1,17 +1,8 @@ -import json -import os -import shutil import unittest -from pathlib import ( - Path, -) import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.conf.unit_cells import ( generate_unit_cell, ) diff --git a/tests/context.py b/tests/context.py index 0fce7a85..2cf42f06 100644 --- a/tests/context.py +++ b/tests/context.py @@ -3,7 +3,6 @@ dpgen_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.insert(0, dpgen_path) -import dpgen2 from dpgen2.utils import ( dflow_config, ) diff --git a/tests/entrypoint/context.py b/tests/entrypoint/context.py index ee903610..1445582e 100644 --- a/tests/entrypoint/context.py +++ b/tests/entrypoint/context.py @@ -3,4 +3,3 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 diff --git a/tests/entrypoint/test_argparse.py b/tests/entrypoint/test_argparse.py index 6f494afb..3cab08b3 100644 --- a/tests/entrypoint/test_argparse.py +++ b/tests/entrypoint/test_argparse.py @@ -1,15 +1,8 @@ -import json -import os -import shutil import unittest # isort: off -from .context import ( - dpgen2, -) from dpgen2.entrypoint.main import ( main_parser, - parse_args, workflow_subcommands, ) diff --git a/tests/entrypoint/test_submit.py b/tests/entrypoint/test_submit.py index bc551a2c..8585d081 100644 --- a/tests/entrypoint/test_submit.py +++ b/tests/entrypoint/test_submit.py @@ -1,21 +1,12 @@ import json import os -import random -import shutil -import tempfile import textwrap import unittest from pathlib import ( Path, ) -import dpdata -import numpy as np - # isort: off -from .context import ( - dpgen2, -) from dpgen2.entrypoint.submit import ( copy_scheduler_plans, expand_idx, @@ -28,7 +19,6 @@ TrajRenderLammps, ) from dpgen2.exploration.report import ( - ExplorationReport, ExplorationReportTrustLevelsRandom, ) from dpgen2.exploration.scheduler import ( @@ -38,10 +28,12 @@ from dpgen2.exploration.selector import ( ConfSelectorFrames, ) -from dpgen2.exploration.task import ( - ExplorationStage, - ExplorationTaskGroup, -) + +# isort: off +import sys + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + from mocked_ops import ( MockedExplorationReport, MockedExplorationTaskGroup, @@ -431,7 +423,7 @@ def tearDown(self): ) config["mode"] = None - for ii in self.touched_files + ["POSCAR"]: + for ii in [*self.touched_files, "POSCAR"]: os.remove(ii) def test(self): @@ -481,129 +473,129 @@ def test(self): """ { "default_step_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "step_configs":{ - "run_train_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + "run_train_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_explore_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_explore_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_fp_config" : { - "template_config" : { - "image" : "vasp:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_fp_config" : { + "template_config" : { + "image" : "vasp:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "_comment" : "all" + }, + "_comment" : "all" + }, + "_comment" : "all" }, "inputs": { - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], - "init_data_prefix": null, - "init_data_sys": [ - "init" - ], - "_comment" : "all" + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], + "init_data_prefix": null, + "init_data_sys": [ + "init" + ], + "_comment" : "all" }, "train":{ - "type" : "dp", - "numb_models" : 2, - "config" : {}, - "template_script" : ["foo", "foo1"], + "type" : "dp", + "numb_models" : 2, + "config" : {}, + "template_script" : ["foo", "foo1"], "init_models_paths" : ["bar", "tar"], - "_comment" : "all" + "_comment" : "all" }, "explore" : { - "type" : "lmp", - "config" : { - "command": "lmp -var restart 0" - }, - "convergence": { - "type" : "fixed-levels", - "conv_accuracy" : 0.9, - "level_f_lo": 0.05, - "level_f_hi": 0.50, - "_comment" : "all" - }, - "max_numb_iter" : 5, - "fatal_at_max" : false, - "output_nopbc": false, - "configuration_prefix": null, - "configurations": [ - { - "type": "alloy", - "lattice" : ["fcc", 4.57], - "replicate" : [2, 2, 2], - "numb_confs" : 30, - "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] - } - ], - "_comment" : "Stage is of type List[List[dict]]. ", - "_comment" : "The outer list gives stages, the inner list gives the task groups of the stage, and dict describes the task group.", - "stages": [ - [ - { - "type" : "lmp-md", - "ensemble": "nvt", "nsteps": 50, "press": [1e0], "temps": [50], "trj_freq": 10, - "conf_idx": [0], "n_sample" : 3 - }, - { - "type" : "customized-lmp-template", + "type" : "lmp", + "config" : { + "command": "lmp -var restart 0" + }, + "convergence": { + "type" : "fixed-levels", + "conv_accuracy" : 0.9, + "level_f_lo": 0.05, + "level_f_hi": 0.50, + "_comment" : "all" + }, + "max_numb_iter" : 5, + "fatal_at_max" : false, + "output_nopbc": false, + "configuration_prefix": null, + "configurations": [ + { + "type": "alloy", + "lattice" : ["fcc", 4.57], + "replicate" : [2, 2, 2], + "numb_confs" : 30, + "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] + } + ], + "_comment" : "Stage is of type List[List[dict]]. ", + "_comment" : "The outer list gives stages, the inner list gives the task groups of the stage, and dict describes the task group.", + "stages": [ + [ + { + "type" : "lmp-md", + "ensemble": "nvt", "nsteps": 50, "press": [1e0], "temps": [50], "trj_freq": 10, + "conf_idx": [0], "n_sample" : 3 + }, + { + "type" : "customized-lmp-template", "custom_shell_commands": ["mkdir aaa && cp conf.lmp lmp.template plm.template aaa"], - "input_lmp_tmpl_name": "lmp.template", + "input_lmp_tmpl_name": "lmp.template", "input_plm_tmpl_name": "plm.template", "output_dir_pattern": ["aaa"], "output_lmp_tmpl_name": "lmp.template", "output_plm_tmpl_name": "plm.template", - "conf_idx": [0], "n_sample" : 3 - } - ] - ], - "_comment" : "all" + "conf_idx": [0], "n_sample" : 3 + } + ] + ], + "_comment" : "all" }, "fp" : { - "type" : "vasp", - "task_max": 2, - "inputs_config" : { - "pp_files": {"Al" : "POTCAR.Al", "Mg" : "POTCAR.Mg"}, - "incar": "INCAR", - "kspacing": 0.32, - "kgamma": true - }, - "run_config" : { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" - }, - "_comment" : "all" + "type" : "vasp", + "task_max": 2, + "inputs_config" : { + "pp_files": {"Al" : "POTCAR.Al", "Mg" : "POTCAR.Mg"}, + "incar": "INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "run_config" : { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "_comment" : "all" } } """ @@ -737,8 +729,8 @@ def test(self): "_comment": "all" }, "fp" : { - "type" : "deepmd", - "task_max": 2, + "type" : "deepmd", + "task_max": 2, "run_config" : { "teacher_model_path": "teacher_model.pb", "type_map": ["H", "C"] @@ -753,10 +745,10 @@ def test(self): "command": "lmp -var restart 0" }, "convergence": { - "type" : "fixed-levels", - "conv_accuracy" : 0.9, - "level_f_lo": 0.05, - "level_f_hi": 0.50, + "type" : "fixed-levels", + "conv_accuracy" : 0.9, + "level_f_lo": 0.05, + "level_f_hi": 0.50, "_comment" : "all" }, "max_numb_iter": 2, @@ -791,119 +783,119 @@ def test(self): """ { "default_step_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "step_configs":{ - "run_train_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + "run_train_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_explore_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_explore_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_fp_config" : { - "template_config" : { - "image" : "vasp:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_fp_config" : { + "template_config" : { + "image" : "vasp:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "_comment" : "all" + }, + "_comment" : "all" + }, + "_comment" : "all" }, "inputs": { - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], - "init_data_prefix": null, - "init_data_sys": [ - "init" - ], - "_comment" : "all" + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], + "init_data_prefix": null, + "init_data_sys": [ + "init" + ], + "_comment" : "all" }, "train":{ - "type" : "dp", - "numb_models" : 2, - "config" : {}, - "template_script" : ["foo", "foo1"], + "type" : "dp", + "numb_models" : 2, + "config" : {}, + "template_script" : ["foo", "foo1"], "init_models_paths" : ["bar", "tar"], "do_finetune": true, - "_comment" : "all" + "_comment" : "all" }, "explore" : { - "type" : "lmp", - "config" : { - "command": "lmp -var restart 0" - }, - "convergence": { - "type" : "fixed-levels", - "conv_accuracy" : 0.9, - "level_f_lo": 0.05, - "level_f_hi": 0.50, - "_comment" : "all" - }, - "max_numb_iter" : 5, - "fatal_at_max" : false, - "output_nopbc": false, - "configuration_prefix": null, - "configurations": [ - { - "type": "alloy", - "lattice" : ["fcc", 4.57], - "replicate" : [2, 2, 2], - "numb_confs" : 30, - "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] - } - ], - "_comment" : "Stage is of type List[List[dict]]. ", - "_comment" : "The outer list gives stages, the inner list gives the task groups of the stage, and dict describes the task group.", - "stages": [ - [ - { - "type" : "lmp-md", - "ensemble": "nvt", "nsteps": 50, "press": [1e0], "temps": [50], "trj_freq": 10, - "conf_idx": [0], "n_sample" : 3 - } - ] - ], - "_comment" : "all" + "type" : "lmp", + "config" : { + "command": "lmp -var restart 0" + }, + "convergence": { + "type" : "fixed-levels", + "conv_accuracy" : 0.9, + "level_f_lo": 0.05, + "level_f_hi": 0.50, + "_comment" : "all" + }, + "max_numb_iter" : 5, + "fatal_at_max" : false, + "output_nopbc": false, + "configuration_prefix": null, + "configurations": [ + { + "type": "alloy", + "lattice" : ["fcc", 4.57], + "replicate" : [2, 2, 2], + "numb_confs" : 30, + "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] + } + ], + "_comment" : "Stage is of type List[List[dict]]. ", + "_comment" : "The outer list gives stages, the inner list gives the task groups of the stage, and dict describes the task group.", + "stages": [ + [ + { + "type" : "lmp-md", + "ensemble": "nvt", "nsteps": 50, "press": [1e0], "temps": [50], "trj_freq": 10, + "conf_idx": [0], "n_sample" : 3 + } + ] + ], + "_comment" : "all" }, "fp" : { - "type" : "vasp", - "task_max": 2, - "inputs_config" : { - "pp_files": {"Al" : "POTCAR.Al", "Mg" : "POTCAR.Mg"}, - "incar": "INCAR", - "kspacing": 0.32, - "kgamma": true - }, - "run_config" : { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" - }, - "_comment" : "all" + "type" : "vasp", + "task_max": 2, + "inputs_config" : { + "pp_files": {"Al" : "POTCAR.Al", "Mg" : "POTCAR.Mg"}, + "incar": "INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "run_config" : { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "_comment" : "all" } } """ diff --git a/tests/entrypoint/test_submit_args.py b/tests/entrypoint/test_submit_args.py index ec7e4582..6dd9608e 100644 --- a/tests/entrypoint/test_submit_args.py +++ b/tests/entrypoint/test_submit_args.py @@ -1,21 +1,8 @@ import json -import os -import random -import shutil -import tempfile import textwrap import unittest -from pathlib import ( - Path, -) - -import dpdata -import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.entrypoint.args import ( normalize, ) @@ -164,153 +151,153 @@ def test_bohrium(self): "fp_style" : "vasp", "default_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "run_train_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" + }, + "_comment" : "all" }, "run_explore_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" + }, + "_comment" : "all" }, "run_fp_config" : { - "template_config" : { - "image" : "vasp:wanghan", - "_comment" : "all" - }, - "executor" : { + "template_config" : { + "image" : "vasp:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" + }, + "_comment" : "all" }, "lmp_config": { - "command": "lmp -var restart 0" + "command": "lmp -var restart 0" }, "fp_config": { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std", + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std", "log" : "fp.log", "out" : "data" }, "dflow_config" : { - "host" : "http://60.205.112.9:2746", - "s3_endpoint" : "60.205.112.9:9000", - "_catalog_file_name" : "dflow" + "host" : "http://60.205.112.9:2746", + "s3_endpoint" : "60.205.112.9:9000", + "_catalog_file_name" : "dflow" }, "_comment" : "upload the dpgen2 package if it is not in the images", "upload_python_packages" : "/path/to/dpgen2", - "max_numb_iter" : 5, - "conv_accuracy" : 0.9, - "fatal_at_max" : false, + "max_numb_iter" : 5, + "conv_accuracy" : 0.9, + "fatal_at_max" : false, - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], - "init_data_prefix": "", - "init_data_sys": [ - "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", - "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" + "init_data_prefix": "", + "init_data_sys": [ + "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", + "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" ], "sys_configs_prefix": "", - "sys_configs": [ - { - "lattice" : ["fcc", 4.57], - "replicate" : [2, 2, 2], - "numb_confs" : 30, + "sys_configs": [ + { + "lattice" : ["fcc", 4.57], + "replicate" : [2, 2, 2], + "numb_confs" : 30, "atom_pert_dist" : 0.0, "cell_pert_frac" : 0.0, - "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] - } + "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] + } ], - "_comment": " 00.train ", - "numb_models": 4, + "_comment": " 00.train ", + "numb_models": 4, "default_training_param" : { - "model" : { - "type_map": ["Al", "Mg"], - "descriptor": { - "type": "se_a", - "sel": [90, 90], - "rcut_smth": 1.80, - "rcut": 6.00, - "neuron": [25, 50, 100], - "resnet_dt": false, - "axis_neuron": 4, - "seed": 1 - }, - "fitting_net" : { - "neuron": [128, 128, 128], - "resnet_dt": true, - "seed": 1 - } - }, - - "loss" : { - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0 - }, - - "learning_rate" : { - "start_lr": 0.001, - "stop_lr": 1e-8, - "decay_steps": 100 - }, - - "training" : { - "training_data": { - "systems": [], - "batch_size":"auto" - }, - "numb_steps":1000, - "seed":10, - "disp_file":"lcurve.out", - "disp_freq":100, - "save_freq":1000 - } - }, - - "_comment": " 01.model_devi ", + "model" : { + "type_map": ["Al", "Mg"], + "descriptor": { + "type": "se_a", + "sel": [90, 90], + "rcut_smth": 1.80, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": false, + "axis_neuron": 4, + "seed": 1 + }, + "fitting_net" : { + "neuron": [128, 128, 128], + "resnet_dt": true, + "seed": 1 + } + }, + + "loss" : { + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0 + }, + + "learning_rate" : { + "start_lr": 0.001, + "stop_lr": 1e-8, + "decay_steps": 100 + }, + + "training" : { + "training_data": { + "systems": [], + "batch_size":"auto" + }, + "numb_steps":1000, + "seed":10, + "disp_file":"lcurve.out", + "disp_freq":100, + "save_freq":1000 + } + }, + + "_comment": " 01.model_devi ", "_comment": "model_devi_skip: the first x of the recorded frames", - "model_devi_f_trust_lo": 0.05, - "model_devi_f_trust_hi": 0.50, - "model_devi_jobs": [ - { "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": [1.0,2.0], "sys_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 } + "model_devi_f_trust_lo": 0.05, + "model_devi_f_trust_hi": 0.50, + "model_devi_jobs": [ + { "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": [1.0,2.0], "sys_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 } ], - "_comment": " 02.fp ", - "fp_style": "vasp", - "fp_task_max": 2, - "fp_pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, + "_comment": " 02.fp ", + "fp_style": "vasp", + "fp_task_max": 2, + "fp_pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, "fp_incar": "vasp/INCAR", - "_comment": " that's all " + "_comment": " that's all " } """ ) @@ -320,118 +307,118 @@ def test_bohrium(self): """ { "dflow_config" : { - "host" : "http://address.of.the.host:port" + "host" : "http://address.of.the.host:port" }, "dflow_s3_config" : { - "s3_endpoint" : "address.of.the.s3.sever:port" + "s3_endpoint" : "address.of.the.s3.sever:port" }, "default_step_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "step_configs":{ - "run_train_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + "run_train_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_explore_config" : { - "template_config" : { - "image" : "deepmd-kit:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_explore_config" : { + "template_config" : { + "image" : "deepmd-kit:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "run_fp_config" : { - "template_config" : { - "image" : "vasp:wanghan", - "_comment" : "all" - }, - "executor" : { + }, + "_comment" : "all" + }, + "run_fp_config" : { + "template_config" : { + "image" : "vasp:wanghan", + "_comment" : "all" + }, + "executor" : { "type": "dispatcher", "username": "foo" - }, - "_comment" : "all" - }, - "_comment" : "all" + }, + "_comment" : "all" + }, + "_comment" : "all" }, "upload_python_packages" : "/path/to/dpgen2", "inputs": { - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], - "init_data_prefix": "", - "init_data_sys": [ - "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", - "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" - ], - "_comment" : "all" + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], + "init_data_prefix": "", + "init_data_sys": [ + "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", + "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" + ], + "_comment" : "all" }, "train":{ - "type" : "dp", - "numb_models" : 4, - "config" : {}, - "template_script" : "dp_input_template" + "type" : "dp", + "numb_models" : 4, + "config" : {}, + "template_script" : "dp_input_template" }, "explore" : { - "type" : "lmp", - "config" : { - "command": "lmp -var restart 0" - }, - "max_numb_iter" : 5, - "fatal_at_max" : false, + "type" : "lmp", + "config" : { + "command": "lmp -var restart 0" + }, + "max_numb_iter" : 5, + "fatal_at_max" : false, "convergence":{ "type": "fixed-levels", - "level_f_lo": 0.05, - "level_f_hi": 0.50, - "conv_accuracy" : 0.9 + "level_f_lo": 0.05, + "level_f_hi": 0.50, + "conv_accuracy" : 0.9 }, - "configuration_prefix": null, - "configuration": [ - { + "configuration_prefix": null, + "configuration": [ + { "type" : "alloy", - "lattice" : ["fcc", 4.57], - "replicate" : [2, 2, 2], - "numb_confs" : 30, - "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] - } - ], - "stages": [[ - { "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": [1.0,2.0], "sys_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 - } - ]], - "_comment" : "all" + "lattice" : ["fcc", 4.57], + "replicate" : [2, 2, 2], + "numb_confs" : 30, + "concentration" : [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]] + } + ], + "stages": [[ + { "_idx": 0, "ensemble": "nvt", "nsteps": 20, "press": [1.0,2.0], "sys_idx": [0], "temps": [50,100], "trj_freq": 10, "n_sample" : 3 + } + ]], + "_comment" : "all" }, "fp" : { - "type" : "vasp", - "run_config" : { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" - }, - "task_max": 2, - "inputs_config" : { - "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, - "incar": "vasp/INCAR", - "kspacing": 0.32, - "kgamma": true - }, - "_comment" : "all" + "type" : "vasp", + "run_config" : { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "task_max": 2, + "inputs_config" : { + "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, + "incar": "vasp/INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "_comment" : "all" } } """ @@ -448,70 +435,70 @@ def test_bohrium(self): }, "default_step_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "step_configs":{ - "_comment" : "all" + "_comment" : "all" }, "upload_python_packages" : "/path/to/dpgen2", "inputs": { - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], - "init_data_prefix": "", - "init_data_sys": [ - "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", - "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" - ], - "_comment" : "all" + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], + "init_data_prefix": "", + "init_data_sys": [ + "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", + "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" + ], + "_comment" : "all" }, "train":{ - "type" : "dp", - "numb_models" : 4, - "config" : {}, - "template_script" : "dp_input_template", - "_comment" : "all" + "type" : "dp", + "numb_models" : 4, + "config" : {}, + "template_script" : "dp_input_template", + "_comment" : "all" }, "explore" : { - "type" : "lmp", - "config" : { - "command": "lmp -var restart 0" - }, - "max_numb_iter" : 5, - "fatal_at_max" : false, + "type" : "lmp", + "config" : { + "command": "lmp -var restart 0" + }, + "max_numb_iter" : 5, + "fatal_at_max" : false, "convergence":{ "type": "fixed-levels", - "level_f_lo": 0.05, - "level_f_hi": 0.50, - "conv_accuracy" : 0.9 + "level_f_lo": 0.05, + "level_f_hi": 0.50, + "conv_accuracy" : 0.9 }, - "configuration_prefix": null, - "configuration": [ - ], - "stages": [ - ], - "_comment" : "all" + "configuration_prefix": null, + "configuration": [ + ], + "stages": [ + ], + "_comment" : "all" }, "fp" : { - "type" : "vasp", - "run_config" : { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" - }, - "task_max": 2, - "inputs_config" : { - "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, - "incar": "vasp/INCAR", - "kspacing": 0.32, - "kgamma": true - }, - "_comment" : "all" + "type" : "vasp", + "run_config" : { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "task_max": 2, + "inputs_config" : { + "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, + "incar": "vasp/INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "_comment" : "all" } } """ diff --git a/tests/entrypoint/test_workflow.py b/tests/entrypoint/test_workflow.py index 5b48bcca..84e54bb1 100644 --- a/tests/entrypoint/test_workflow.py +++ b/tests/entrypoint/test_workflow.py @@ -1,19 +1,10 @@ import json -import os -import shutil import textwrap import unittest -import dflow import mock -from dflow import ( - Workflow, -) # isort: off -from .context import ( - dpgen2, -) from dpgen2.entrypoint.workflow import ( execute_workflow_subcommand, ) @@ -61,70 +52,70 @@ def test_resume(self, mocked_f): """ { "default_step_config" : { - "template_config" : { - "image" : "dflow:1.1.4", - "_comment" : "all" - }, - "_comment" : "all" + "template_config" : { + "image" : "dflow:1.1.4", + "_comment" : "all" + }, + "_comment" : "all" }, "step_configs":{ - "_comment" : "all" + "_comment" : "all" }, "upload_python_packages" : "/path/to/dpgen2", "inputs": { - "type_map": ["Al", "Mg"], - "mass_map": [27, 24], - "init_data_prefix": "", - "init_data_sys": [ - "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", - "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" - ], - "_comment" : "all" + "type_map": ["Al", "Mg"], + "mass_map": [27, 24], + "init_data_prefix": "", + "init_data_sys": [ + "init/al.fcc.01x01x01/02.md/sys-0004/deepmd", + "init/mg.fcc.01x01x01/02.md/sys-0004/deepmd" + ], + "_comment" : "all" }, "train":{ - "type" : "dp", - "numb_models" : 4, - "config" : {}, - "template_script" : "dp_input_template", - "_comment" : "all" + "type" : "dp", + "numb_models" : 4, + "config" : {}, + "template_script" : "dp_input_template", + "_comment" : "all" }, "explore" : { - "type" : "lmp", - "config" : { - "command": "lmp -var restart 0" - }, - "max_numb_iter" : 5, - "fatal_at_max" : false, + "type" : "lmp", + "config" : { + "command": "lmp -var restart 0" + }, + "max_numb_iter" : 5, + "fatal_at_max" : false, "convergence":{ "type": "fixed-levels", - "level_f_lo": 0.05, - "level_f_hi": 0.50, - "conv_accuracy" : 0.9 + "level_f_lo": 0.05, + "level_f_hi": 0.50, + "conv_accuracy" : 0.9 }, - "configuration_prefix": null, - "configuration": [ - ], - "stages": [ - ], - "_comment" : "all" + "configuration_prefix": null, + "configuration": [ + ], + "stages": [ + ], + "_comment" : "all" }, "fp" : { - "type" : "vasp", - "run_config" : { - "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" - }, - "task_max": 2, - "inputs_config" : { - "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, - "incar": "vasp/INCAR", - "kspacing": 0.32, - "kgamma": true - }, - "_comment" : "all" + "type" : "vasp", + "run_config" : { + "command": "source /opt/intel/oneapi/setvars.sh && mpirun -n 16 vasp_std" + }, + "task_max": 2, + "inputs_config" : { + "pp_files": {"Al" : "vasp/POTCAR.Al", "Mg" : "vasp/POTCAR.Mg"}, + "incar": "vasp/INCAR", + "kspacing": 0.32, + "kgamma": true + }, + "_comment" : "all" } } """ diff --git a/tests/exploration/context.py b/tests/exploration/context.py index ee903610..1445582e 100644 --- a/tests/exploration/context.py +++ b/tests/exploration/context.py @@ -3,4 +3,3 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 diff --git a/tests/exploration/test_conf_filter.py b/tests/exploration/test_conf_filter.py index a0c36ba1..a82e4c3f 100644 --- a/tests/exploration/test_conf_filter.py +++ b/tests/exploration/test_conf_filter.py @@ -1,19 +1,11 @@ -import os import unittest import dpdata -import numpy as np from fake_data_set import ( fake_system, ) -from mock import ( - patch, -) # isort: off -from .context import ( - dpgen2, -) from dpgen2.exploration.selector import ( ConfFilter, ConfFilters, diff --git a/tests/exploration/test_conf_selector_frame.py b/tests/exploration/test_conf_selector_frame.py index 52b30248..662780fe 100644 --- a/tests/exploration/test_conf_selector_frame.py +++ b/tests/exploration/test_conf_selector_frame.py @@ -7,12 +7,8 @@ ) import dpdata -import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.exploration.render import ( TrajRenderLammps, ) diff --git a/tests/exploration/test_customized_lmp_templ_task_group.py b/tests/exploration/test_customized_lmp_templ_task_group.py index 57462759..9cb930a4 100644 --- a/tests/exploration/test_customized_lmp_templ_task_group.py +++ b/tests/exploration/test_customized_lmp_templ_task_group.py @@ -1,29 +1,10 @@ import itertools import os -import shutil import textwrap import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) - -import numpy as np - -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from unittest.mock import ( - Mock, - patch, -) from dpgen2.constants import ( lmp_conf_name, @@ -32,7 +13,6 @@ ) from dpgen2.exploration.task import ( CustomizedLmpTemplateTaskGroup, - ExplorationStage, ) from .test_lmp_templ_task_group import ( @@ -43,6 +23,9 @@ in_plm_template, ) +# case of upload everything to argo, no context needed + + in_extra_py_file = textwrap.dedent( """from pathlib import Path import os,shutil diff --git a/tests/exploration/test_devi_manager.py b/tests/exploration/test_devi_manager.py index a0a29e68..a2296360 100644 --- a/tests/exploration/test_devi_manager.py +++ b/tests/exploration/test_devi_manager.py @@ -1,15 +1,8 @@ -import os import unittest -from pathlib import ( - Path, -) import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.exploration.deviation import ( DeviManager, DeviManagerStd, @@ -53,7 +46,7 @@ def test_add_invalid_deviation(self): self.assertRaisesRegex( AssertionError, - "Error: deviation\(shape: ", + r"Error: deviation\(shape: ", model_devi.add, DeviManager.MAX_DEVI_F, np.array([[1], [2], [3]]), @@ -61,7 +54,7 @@ def test_add_invalid_deviation(self): self.assertRaisesRegex( AssertionError, - "Error: deviation\(type: ", + r"Error: deviation\(type: ", model_devi.add, DeviManager.MAX_DEVI_F, "foo", @@ -99,13 +92,13 @@ def test_devi_manager_std_check_data(self): model_devi.add(DeviManager.MAX_DEVI_V, np.array([4, 5])) self.assertRaisesRegex( AssertionError, - f"Error: the number of frames in", + "Error: the number of frames in", model_devi.get, DeviManager.MAX_DEVI_F, ) self.assertRaisesRegex( AssertionError, - f"Error: the number of frames in", + "Error: the number of frames in", model_devi.get, DeviManager.MAX_DEVI_V, ) diff --git a/tests/exploration/test_distance_conf_filter.py b/tests/exploration/test_distance_conf_filter.py index 98b7ba4e..3ce49085 100644 --- a/tests/exploration/test_distance_conf_filter.py +++ b/tests/exploration/test_distance_conf_filter.py @@ -2,7 +2,6 @@ import unittest import dpdata -import numpy as np from dpgen2.exploration.selector import ( BoxLengthFilter, @@ -10,10 +9,6 @@ DistanceConfFilter, ) -from .context import ( - dpgen2, -) - POSCAR_valid = """ Er 1.0 7.00390434172054 0.000000000000000E+000 0.000000000000000E+000 diff --git a/tests/exploration/test_exploration_group.py b/tests/exploration/test_exploration_group.py index f71ce8f4..5579cffa 100644 --- a/tests/exploration/test_exploration_group.py +++ b/tests/exploration/test_exploration_group.py @@ -1,31 +1,13 @@ -import os import textwrap import unittest -from pathlib import ( - Path, -) -from typing import ( - List, - Set, -) - -import numpy as np -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from unittest.mock import ( - Mock, patch, ) from dpgen2.constants import ( calypso_input_file, - calypso_run_opt_file, lmp_conf_name, lmp_input_name, ) @@ -532,7 +514,7 @@ def test_nvt_sample(self, mock_random): in_template_nvt % (self.tt[j_idx]), ) - @patch("dpgen2.exploration.task.npt_task_group.random.shuffle") + @patch("dpgen2.exploration.task.conf_sampling_task_group.random.shuffle") @patch("dpgen2.exploration.task.lmp.lmp_input.random.randrange") def test_nvt_sample_random(self, mock_randrange, mock_shuffle): mock_randrange.return_value = 1110 diff --git a/tests/exploration/test_exploration_scheduler.py b/tests/exploration/test_exploration_scheduler.py index 324904b0..2c62cb11 100644 --- a/tests/exploration/test_exploration_scheduler.py +++ b/tests/exploration/test_exploration_scheduler.py @@ -1,23 +1,8 @@ -import os -import textwrap import unittest -from pathlib import ( - Path, -) -from typing import ( - List, - Set, -) import numpy as np -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from dflow.python import ( FatalError, ) @@ -39,13 +24,8 @@ from dpgen2.exploration.selector import ( ConfSelectorFrames, ) -from dpgen2.exploration.task import ( - ExplorationStage, - ExplorationTaskGroup, -) # isort: off -import exploration.context from mocked_ops import ( MockedExplorationReport, MockedExplorationTaskGroup, diff --git a/tests/exploration/test_lmp_templ_task_group.py b/tests/exploration/test_lmp_templ_task_group.py index 6211137d..9886c8f2 100644 --- a/tests/exploration/test_lmp_templ_task_group.py +++ b/tests/exploration/test_lmp_templ_task_group.py @@ -5,24 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) - -import numpy as np - -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from unittest.mock import ( - Mock, - patch, -) from dpgen2.constants import ( lmp_conf_name, @@ -30,10 +12,12 @@ plm_input_name, ) from dpgen2.exploration.task import ( - ExplorationStage, LmpTemplateTaskGroup, ) +# case of upload everything to argo, no context needed + + in_lmp_template = textwrap.dedent( """variable NSTEPS equal V_NSTEPS variable THERMO_FREQ equal 10 diff --git a/tests/exploration/test_make_task_group_from_config.py b/tests/exploration/test_make_task_group_from_config.py index f9fe93be..f14bea45 100644 --- a/tests/exploration/test_make_task_group_from_config.py +++ b/tests/exploration/test_make_task_group_from_config.py @@ -1,24 +1,10 @@ -import itertools import os -import textwrap import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) - -import numpy as np -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from dpgen2.exploration.task import ( CalyTaskGroup, LmpTemplateTaskGroup, diff --git a/tests/exploration/test_report_adaptive_lower.py b/tests/exploration/test_report_adaptive_lower.py index b5f123ca..4c44d594 100644 --- a/tests/exploration/test_report_adaptive_lower.py +++ b/tests/exploration/test_report_adaptive_lower.py @@ -1,9 +1,4 @@ -import os -import textwrap import unittest -from collections import ( - Counter, -) import mock import numpy as np @@ -12,9 +7,6 @@ ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.exploration.deviation import ( DeviManager, DeviManagerStd, @@ -183,8 +175,20 @@ def test_f_inv_pop(self): for jj in ii: expected_fail.add((idx, jj)) expected_cand = set( - [(0, 6), (0, 7), (0, 5)] - + [(0, 1), (0, 3), (0, 4), (1, 0), (1, 1), (1, 5), (1, 6), (1, 7), (1, 8)] + [ + (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([]) diff --git a/tests/exploration/test_report_trust_levels.py b/tests/exploration/test_report_trust_levels.py index 7f6c79a9..e67cc5f2 100644 --- a/tests/exploration/test_report_trust_levels.py +++ b/tests/exploration/test_report_trust_levels.py @@ -1,9 +1,4 @@ -import os -import textwrap import unittest -from collections import ( - Counter, -) import numpy as np from dargs import ( @@ -11,9 +6,6 @@ ) # isort: off -from context import ( - dpgen2, -) from dpgen2.exploration.deviation import ( DeviManager, DeviManagerStd, diff --git a/tests/exploration/test_traj_render_lammps.py b/tests/exploration/test_traj_render_lammps.py index 3659ffd8..b968802f 100644 --- a/tests/exploration/test_traj_render_lammps.py +++ b/tests/exploration/test_traj_render_lammps.py @@ -6,9 +6,6 @@ import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.exploration.render import TrajRenderLammps # isort: on diff --git a/tests/fp/context.py b/tests/fp/context.py index d6638725..5504dedc 100644 --- a/tests/fp/context.py +++ b/tests/fp/context.py @@ -2,4 +2,3 @@ import sys sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 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..b543e295 100644 --- a/tests/fp/data.vasp.kp.gf/make_kp_test.py +++ b/tests/fp/data.vasp.kp.gf/make_kp_test.py @@ -8,10 +8,11 @@ def make_one(out_dir): + rng = np.random.default_rng() # [0.5, 1) - [aa, bb, cc] = np.random.random(3) * 0.5 + 0.5 + [aa, bb, cc] = rng.random(3) * 0.5 + 0.5 # [1, 179) - [alpha, beta, gamma] = np.random.random(3) * (178 / 180) + 1 + [alpha, beta, gamma] = rng.random(3) * (178 / 180) + 1 # make cell cell = ase.geometry.cellpar_to_cell([aa, bb, cc, alpha, beta, gamma]) sys = dpdata.System("POSCAR") diff --git a/tests/fp/test_prep_vasp.py b/tests/fp/test_prep_vasp.py index 08957f79..cd3352e8 100644 --- a/tests/fp/test_prep_vasp.py +++ b/tests/fp/test_prep_vasp.py @@ -1,8 +1,5 @@ -import glob -import json import os import shutil -import sys import textwrap import unittest from pathlib import ( @@ -10,24 +7,14 @@ ) import dpdata -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - FatalError, - OPIOSign, - TransientError, ) from fake_data_set import ( fake_multi_sys, - fake_system, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( fp_task_pattern, ) @@ -39,9 +26,6 @@ vasp_kp_name, vasp_pot_name, ) -from dpgen2.utils import ( - dump_object_to_file, -) # isort: on diff --git a/tests/fp/test_run_vasp.py b/tests/fp/test_run_vasp.py index 7202d6ce..dba507b1 100644 --- a/tests/fp/test_run_vasp.py +++ b/tests/fp/test_run_vasp.py @@ -1,16 +1,11 @@ -import json import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, TransientError, ) from mock import ( diff --git a/tests/fp/test_vasp.py b/tests/fp/test_vasp.py index 8d9bee25..8123a1e2 100644 --- a/tests/fp/test_vasp.py +++ b/tests/fp/test_vasp.py @@ -1,8 +1,6 @@ import glob import json import os -import shutil -import sys import textwrap import unittest from pathlib import ( @@ -13,9 +11,6 @@ import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.fp.vasp import ( PrepVasp, RunVasp, @@ -122,7 +117,7 @@ def test_vasp_input_kp(self): kps = vi.make_kpoints(ss["cells"][0]) self.assertEqual(ref, kps) - def test_vasp_input_kp(self): + def test_vasp_input_kp_monkhorst(self): ref = textwrap.dedent( """K-Points 0 diff --git a/tests/mocked_ops.py b/tests/mocked_ops.py index 6e6cfadf..dc34dfe5 100644 --- a/tests/mocked_ops.py +++ b/tests/mocked_ops.py @@ -1,9 +1,7 @@ from dflow.python import ( OP, OPIO, - Artifact, FatalError, - OPIOSign, upload_packages, ) @@ -11,7 +9,6 @@ import json import os -import pickle import re import shutil from pathlib import ( @@ -23,26 +20,15 @@ Tuple, ) -try: - from flow.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from dpgen2.constants import ( - calypso_check_opt_file, - calypso_run_opt_file, fp_task_pattern, lmp_conf_name, lmp_input_name, lmp_log_name, lmp_model_devi_name, - lmp_task_pattern, lmp_traj_name, model_name_pattern, - train_log_name, - train_script_name, train_task_pattern, ) from dpgen2.exploration.report import ( @@ -64,7 +50,6 @@ RunVasp, ) from dpgen2.fp.vasp import ( - VaspInputs, vasp_conf_name, vasp_input_name, ) @@ -80,9 +65,6 @@ from dpgen2.op.prep_dp_train import ( PrepDPTrain, ) -from dpgen2.op.prep_lmp import ( - PrepExplorationTaskGroup, -) from dpgen2.op.run_caly_dp_optim import ( RunCalyDPOptim, ) @@ -221,15 +203,15 @@ def execute( assert init_model.exists() with log.open("w") as f: - f.write(f"init_model {str(init_model)} OK\n") + f.write(f"init_model {init_model!s} OK\n") for ii in jtmp["data"]: assert Path(ii).exists() assert (ii in init_data_str) or (ii in iter_data_str) with log.open("a") as f: - f.write(f"data {str(ii)} OK\n") + f.write(f"data {ii!s} OK\n") assert script.exists() with log.open("a") as f: - f.write(f"script {str(script)} OK\n") + f.write(f"script {script!s} OK\n") with model.open("w") as f: f.write("read from init model: \n") @@ -328,10 +310,10 @@ def execute( assert Path(ii).exists() assert (ii in init_data_str) or (ii in iter_data_str) with log.open("a") as f: - f.write(f"data {str(ii)} OK\n") + f.write(f"data {ii!s} OK\n") assert script.exists() with log.open("a") as f: - f.write(f"script {str(script)} OK\n") + f.write(f"script {script!s} OK\n") with model.open("w") as f: f.write("read from init model: \n") @@ -366,9 +348,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]) @@ -865,7 +847,7 @@ def select( self, trajs: List[Path], model_devis: List[Path], - type_map: List[str] = None, + type_map: Optional[List[str]] = None, optional_outputs: Optional[List[Path]] = None, ) -> Tuple[List[Path], ExplorationReport]: confs = [] @@ -922,7 +904,7 @@ def __init__( self, stage: ExplorationStage, conv_accuracy: float = 0.9, - max_numb_iter: int = None, + max_numb_iter: Optional[int] = None, ): self.selector = MockedConfSelector( conv_accuracy=conv_accuracy, @@ -982,13 +964,13 @@ def execute( finished = "true" if int(cnt_num) == int(max_step) else "false" if finished == "false": for i in range(5): - Path(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + Path(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") if step is None: Path("step").write_text("2") else: step_num = Path("step").read_text().strip() - Path("step").write_text(f"{int(step_num)+1}") + 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/context.py b/tests/op/context.py index d6638725..5504dedc 100644 --- a/tests/op/context.py +++ b/tests/op/context.py @@ -2,4 +2,3 @@ import sys sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 diff --git a/tests/op/test_collect_data.py b/tests/op/test_collect_data.py index 2e558193..4f2ae88d 100644 --- a/tests/op/test_collect_data.py +++ b/tests/op/test_collect_data.py @@ -1,4 +1,3 @@ -import json import shutil import unittest from pathlib import ( @@ -6,35 +5,15 @@ ) import dpdata -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - TransientError, ) from fake_data_set import ( fake_multi_sys, fake_system, ) -from mock import ( - call, - mock, - patch, -) # isort: off -from .context import ( - dpgen2, -) -from dpgen2.constants import ( - lmp_conf_name, - lmp_input_name, - lmp_log_name, - lmp_model_devi_name, - lmp_traj_name, -) from dpgen2.op.collect_data import ( CollectData, ) diff --git a/tests/op/test_collect_run_caly.py b/tests/op/test_collect_run_caly.py index 22ec37db..e0d32291 100644 --- a/tests/op/test_collect_run_caly.py +++ b/tests/op/test_collect_run_caly.py @@ -1,37 +1,23 @@ -import os import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, TransientError, ) from mock import ( - call, - mock, patch, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( calypso_task_pattern, calypso_input_file, - calypso_log_name, ) from dpgen2.op.collect_run_caly import CollRunCaly, get_value_from_inputdat -from dpgen2.utils import ( - BinaryFileInput, -) # isort: on @@ -71,13 +57,13 @@ def tearDown(self): def test_get_max_step(self): max_step, vsc = get_value_from_inputdat(self.input_file) self.assertTrue(max_step == 3) - self.assertTrue(vsc == True) + self.assertTrue(vsc) temp_input_file = self.input_file_path.joinpath("temp_input_dat") temp_input_file.write_text("input.dat\n") max_step, vsc = get_value_from_inputdat(temp_input_file) self.assertTrue(max_step == 0) - self.assertTrue(vsc == False) + self.assertTrue(not vsc) @patch("dpgen2.op.collect_run_caly.run_command") def test_step_st_maxstep_01(self, mocked_run): @@ -86,7 +72,7 @@ def test_step_st_maxstep_01(self, mocked_run): def side_effect(*args, **kwargs): for i in range(5): - Path().joinpath(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + Path().joinpath(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") Path("step").write_text("3") Path("results").mkdir(parents=True, exist_ok=True) return (0, "foo\n", "") @@ -122,7 +108,7 @@ def test_step_no_eq_maxstep_02(self, mocked_run): def side_effect(*args, **kwargs): for i in range(5): - Path().joinpath(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + Path().joinpath(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") Path("step").write_text("2") Path("results").mkdir(parents=True, exist_ok=True) return (0, "foo\n", "") @@ -153,7 +139,7 @@ def test_step_eq_maxstep_03(self, mocked_run): def side_effect(*args, **kwargs): for i in range(5): - Path().joinpath(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + Path().joinpath(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") Path("step").write_text("4") Path("results").mkdir(parents=True, exist_ok=True) return (0, "foo\n", "") @@ -185,7 +171,7 @@ def test_error_04(self, mocked_run): def side_effect(*args, **kwargs): for i in range(5): - Path().joinpath(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + Path().joinpath(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") Path("step").write_text("3") Path("results").mkdir(parents=True, exist_ok=True) return (1, "foo\n", "") diff --git a/tests/op/test_prep_caly_dp_optim.py b/tests/op/test_prep_caly_dp_optim.py index d1ffd724..4640152a 100644 --- a/tests/op/test_prep_caly_dp_optim.py +++ b/tests/op/test_prep_caly_dp_optim.py @@ -1,38 +1,20 @@ -import os import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - TransientError, -) -from mock import ( - call, - mock, - patch, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( calypso_task_pattern, - model_name_pattern, calypso_run_opt_file, calypso_check_opt_file, ) from dpgen2.op import PrepCalyDPOptim -from dpgen2.utils import ( - BinaryFileInput, -) # isort: on @@ -46,7 +28,7 @@ def setUp(self): self.poscar_dir.mkdir(parents=True, exist_ok=True) nposcar = 10 for i in range(1, nposcar + 1): - self.poscar_dir.joinpath(f"POSCAR_{str(i)}").write_text(f"POSCAR_{str(i)}") + self.poscar_dir.joinpath(f"POSCAR_{i!s}").write_text(f"POSCAR_{i!s}") self.models_dir = Path("models_dir") self.models_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/op/test_prep_caly_input.py b/tests/op/test_prep_caly_input.py index 0fac49c4..f96a8289 100644 --- a/tests/op/test_prep_caly_input.py +++ b/tests/op/test_prep_caly_input.py @@ -1,22 +1,11 @@ -import os import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - TransientError, -) -from mock import ( - call, - mock, - patch, ) # isort: off @@ -27,9 +16,6 @@ calypso_check_opt_file, ) from dpgen2.op.prep_caly_input import PrepCalyInput -from dpgen2.utils import ( - BinaryFileInput, -) from dpgen2.exploration.task import ( BaseExplorationTaskGroup, diff --git a/tests/op/test_prep_caly_model_devi.py b/tests/op/test_prep_caly_model_devi.py index 6cc6337f..875c801c 100644 --- a/tests/op/test_prep_caly_model_devi.py +++ b/tests/op/test_prep_caly_model_devi.py @@ -1,38 +1,15 @@ -import os import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - TransientError, -) -from mock import ( - call, - mock, - patch, ) # isort: off -from .context import ( - dpgen2, -) -from dpgen2.constants import ( - calypso_task_pattern, - model_name_pattern, - calypso_run_opt_file, - calypso_check_opt_file, -) from dpgen2.op import PrepCalyModelDevi -from dpgen2.utils import ( - BinaryFileInput, -) # isort: on diff --git a/tests/op/test_prep_dp_train.py b/tests/op/test_prep_dp_train.py index a380e221..e77594c4 100644 --- a/tests/op/test_prep_dp_train.py +++ b/tests/op/test_prep_dp_train.py @@ -5,21 +5,14 @@ Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, ) from mock import ( mock, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( train_script_name, train_task_pattern, diff --git a/tests/op/test_run_caly_dp_optim.py b/tests/op/test_run_caly_dp_optim.py index ee6afd43..ca8a0e9d 100644 --- a/tests/op/test_run_caly_dp_optim.py +++ b/tests/op/test_run_caly_dp_optim.py @@ -1,38 +1,19 @@ -import os import shutil import unittest from pathlib import ( Path, ) -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, TransientError, ) from mock import ( - call, - mock, patch, ) # isort: off -from .context import ( - dpgen2, -) -from dpgen2.constants import ( - calypso_task_pattern, - model_name_pattern, - calypso_run_opt_file, - calypso_check_opt_file, -) from dpgen2.op import RunCalyDPOptim -from dpgen2.utils import ( - BinaryFileInput, -) # isort: on @@ -57,9 +38,9 @@ def tearDown(self): def test_00_success(self, mocked_run): def side_effect(*args, **kwargs): for i in range(1, 11): - Path().joinpath(f"CONTCAR_{str(i)}").write_text(f"CONTCAR_{str(i)}") - Path().joinpath(f"OUTCAR_{str(i)}").write_text(f"OUTCAR_{str(i)}") - Path().joinpath(f"{str(i)}.traj").write_text(f"{str(i)}.traj") + Path().joinpath(f"CONTCAR_{i!s}").write_text(f"CONTCAR_{i!s}") + Path().joinpath(f"OUTCAR_{i!s}").write_text(f"OUTCAR_{i!s}") + Path().joinpath(f"{i!s}.traj").write_text(f"{i!s}.traj") return (0, "foo\n", "") mocked_run.side_effect = side_effect @@ -106,9 +87,9 @@ def side_effect(*args, **kwargs): def test_01_error(self, mocked_run): def side_effect(*args, **kwargs): for i in range(1, 6): - Path().joinpath(f"CONTCAR_{str(i)}").write_text(f"CONTCAR_{str(i)}") - Path().joinpath(f"OUTCAR_{str(i)}").write_text(f"OUTCAR_{str(i)}") - Path().joinpath(f"{str(i)}.traj").write_text(f"{str(i)}.traj") + Path().joinpath(f"CONTCAR_{i!s}").write_text(f"CONTCAR_{i!s}") + Path().joinpath(f"OUTCAR_{i!s}").write_text(f"OUTCAR_{i!s}") + Path().joinpath(f"{i!s}.traj").write_text(f"{i!s}.traj") return (1, "foo\n", "") mocked_run.side_effect = side_effect diff --git a/tests/op/test_run_caly_model_devi.py b/tests/op/test_run_caly_model_devi.py index 647604a8..fc009763 100644 --- a/tests/op/test_run_caly_model_devi.py +++ b/tests/op/test_run_caly_model_devi.py @@ -1,16 +1,12 @@ -import os import shutil import unittest from pathlib import ( Path, ) from unittest.mock import ( - Mock, - call, patch, ) -import numpy as np from ase import ( Atoms, ) @@ -18,16 +14,10 @@ write, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - TransientError, ) from dpgen2.constants import ( - calypso_input_file, - calypso_log_name, calypso_task_pattern, ) from dpgen2.op.run_caly_model_devi import ( @@ -35,21 +25,14 @@ atoms2lmpdump, parse_traj, ) -from dpgen2.utils import ( - BinaryFileInput, -) # from .context import ( # dpgen2, # ) # isort: on -try: - import deepmd - x = 0 -except: - x = 1 +x = 0 class TestRunCalyModelDevi(unittest.TestCase): diff --git a/tests/op/test_run_dp_train.py b/tests/op/test_run_dp_train.py index 7649b520..dd931139 100644 --- a/tests/op/test_run_dp_train.py +++ b/tests/op/test_run_dp_train.py @@ -1,4 +1,3 @@ -import itertools import json import os import shutil @@ -8,14 +7,9 @@ ) import dpdata -import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, FatalError, - OPIOSign, - TransientError, ) from fake_data_set import ( fake_multi_sys, @@ -27,17 +21,12 @@ ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( train_script_name, - train_task_pattern, ) from dpgen2.op.run_dp_train import ( RunDPTrain, _get_data_size_of_all_mult_sys, - _make_train_command, split_valid, ) @@ -309,7 +298,7 @@ def test_decide_init_model_config_larger_than_no(self): def test_decide_init_model_config_larger_than_yes(self): config = self.config.copy() - config["init_model_policy"] = f"old_data_larger_than:{self.old_data_size-1}" + config["init_model_policy"] = f"old_data_larger_than:{self.old_data_size - 1}" do_init_model = RunDPTrain.decide_init_model( config, self.init_model, self.init_data, self.iter_data ) diff --git a/tests/op/test_run_lmp.py b/tests/op/test_run_lmp.py index 650fd82e..9cd7bded 100644 --- a/tests/op/test_run_lmp.py +++ b/tests/op/test_run_lmp.py @@ -1,4 +1,3 @@ -import json import os import shutil import unittest @@ -9,22 +8,15 @@ import dpdata import numpy as np from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, TransientError, ) from mock import ( call, - mock, patch, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( lmp_conf_name, lmp_input_name, @@ -240,7 +232,7 @@ def test_success(self, mocked_run): # check if the teacher model is linked to model.000.pb ii = 0 self.assertEqual( - (work_dir / (model_name_pattern % ii)).read_text(), f"teacher model" + (work_dir / (model_name_pattern % ii)).read_text(), "teacher model" ) ii = 1 diff --git a/tests/test_block_cl.py b/tests/test_block_cl.py index fa0571a7..9c202b88 100644 --- a/tests/test_block_cl.py +++ b/tests/test_block_cl.py @@ -1,50 +1,12 @@ -import json import os -import pickle import shutil import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) - -import jsonpickle -import numpy as np -from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, - Step, - Steps, - Workflow, - argo_range, - download_artifact, - upload_artifact, -) -from dflow.python import ( - OP, - OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, - upload_packages, -) -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from context import ( default_host, default_image, @@ -52,14 +14,16 @@ skip_ut_with_dflow_reason, upload_python_packages, ) -from mock import ( - patch, +from dflow import ( + Step, + Workflow, + download_artifact, + upload_artifact, ) from mocked_ops import ( MockedCollectData, MockedCollectDataCheckOptParam, MockedConfSelector, - MockedExplorationReport, MockedExplorationTaskGroup, MockedPrepDPTrain, MockedPrepVasp, @@ -78,19 +42,9 @@ from dpgen2.constants import ( fp_task_pattern, - lmp_conf_name, - lmp_input_name, - lmp_log_name, - lmp_traj_name, model_name_pattern, - train_log_name, - train_script_name, train_task_pattern, ) -from dpgen2.exploration.task import ( - ExplorationTask, - ExplorationTaskGroup, -) from dpgen2.fp.vasp import ( VaspInputs, ) @@ -109,10 +63,6 @@ from dpgen2.superop.prep_run_lmp import ( PrepRunLmp, ) -from dpgen2.utils import ( - dump_object_to_file, - load_object_from_file, -) from dpgen2.utils.step_config import normalize as normalize_step_dict default_config = normalize_step_dict( diff --git a/tests/test_caly_evo_step.py b/tests/test_caly_evo_step.py index 2aff3614..2395ca18 100644 --- a/tests/test_caly_evo_step.py +++ b/tests/test_caly_evo_step.py @@ -1,95 +1,42 @@ -import json import os -import pickle import shutil import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import jsonpickle -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, Slices, ) - -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_index_pattern, - calypso_run_opt_file, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, -) from mocked_ops import ( MockedCollRunCaly, - MockedPrepCalyDPOptim, MockedRunCalyDPOptim, mocked_numb_models, ) from dpgen2.constants import ( - lmp_conf_name, - lmp_input_name, - lmp_log_name, - lmp_model_devi_name, - lmp_task_pattern, - lmp_traj_name, - model_name_pattern, - train_log_name, - train_script_name, - train_task_pattern, -) -from dpgen2.exploration.task import ( - ExplorationTask, - ExplorationTaskGroup, + calypso_check_opt_file, + calypso_index_pattern, + calypso_run_opt_file, ) from dpgen2.op import ( PrepCalyDPOptim, - RunCalyDPOptim, -) -from dpgen2.op.collect_run_caly import ( - CollRunCaly, -) -from dpgen2.op.prep_caly_input import ( - PrepCalyInput, ) from dpgen2.superop.caly_evo_step import ( CalyEvoStep, @@ -167,7 +114,7 @@ def setUp(self) -> None: self.file_storage.mkdir(parents=True, exist_ok=True) for i in range(5): self.file_storage.joinpath(f"POSCAR_{i}").write_text(f"POSCAR_{i}") - self.file_storage.joinpath(f"frozen_model.pb").write_text(f"model.{i}.pb") + self.file_storage.joinpath("frozen_model.pb").write_text(f"model.{i}.pb") self.caly_run_opt_file = self.file_storage.joinpath(calypso_run_opt_file) self.caly_run_opt_file.write_text("caly_run_opt_script") self.caly_check_opt_file = self.file_storage.joinpath(calypso_check_opt_file) @@ -233,7 +180,7 @@ def setUp(self): for ii in range(self.nmodels): model_path = self.work_dir.joinpath(f"task.{ii}") model_path.mkdir(exist_ok=True, parents=True) - model = model_path.joinpath(f"model.ckpt.pt") + model = model_path.joinpath("model.ckpt.pt") model.write_text(f"model {ii}") self.model_list.append(model) self.models = upload_artifact(self.model_list) diff --git a/tests/test_collect_data.py b/tests/test_collect_data.py index 95b7e4c9..95235ae2 100644 --- a/tests/test_collect_data.py +++ b/tests/test_collect_data.py @@ -1,55 +1,27 @@ -import json -import os import shutil import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import jsonpickle -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, PythonOPTemplate, - upload_packages, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, ) from mocked_ops import ( MockedCollectData, @@ -65,10 +37,10 @@ def setUp(self): self.labeled_data = [Path(ii) for ii in self.labeled_data] for ii in self.iter_data: ii.mkdir(exist_ok=True, parents=True) - (ii / "data").write_text(f"data of {str(ii)}") + (ii / "data").write_text(f"data of {ii!s}") for ii in self.labeled_data: (ii).mkdir(exist_ok=True, parents=True) - (ii / "data").write_text(f"data of {str(ii)}") + (ii / "data").write_text(f"data of {ii!s}") self.type_map = [] def tearDown(self): @@ -115,10 +87,10 @@ def setUp(self): self.labeled_data = [Path(ii) for ii in self.labeled_data] for ii in self.iter_data: ii.mkdir(exist_ok=True, parents=True) - (ii / "data").write_text(f"data of {str(ii)}") + (ii / "data").write_text(f"data of {ii!s}") for ii in self.labeled_data: (ii).mkdir(exist_ok=True, parents=True) - (ii / "data").write_text(f"data of {str(ii)}") + (ii / "data").write_text(f"data of {ii!s}") self.iter_data = upload_artifact(list(self.iter_data)) self.labeled_data = upload_artifact(self.labeled_data) self.type_map = [] diff --git a/tests/test_dpgen_loop.py b/tests/test_dpgen_loop.py index 49b4c873..d3f625ac 100644 --- a/tests/test_dpgen_loop.py +++ b/tests/test_dpgen_loop.py @@ -1,72 +1,30 @@ -import json import os -import pickle import shutil -import textwrap import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) - -import jsonpickle -import numpy as np -from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, - Step, - Steps, - Workflow, - argo_range, - download_artifact, - upload_artifact, -) -from dflow.python import ( - OP, - OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, - upload_packages, -) -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from context import ( - default_host, default_image, skip_ut_with_dflow, skip_ut_with_dflow_reason, upload_python_packages, ) -from dflow.python import ( - FatalError, +from dflow import ( + Step, + Workflow, + download_artifact, + upload_artifact, ) from mocked_ops import ( MockedCollectData, MockedCollectDataCheckOptParam, MockedCollectDataFailed, MockedCollectDataRestart, - MockedConfSelector, MockedConstTrustLevelStageScheduler, - MockedExplorationReport, - MockedExplorationTaskGroup, - MockedExplorationTaskGroup1, - MockedExplorationTaskGroup2, MockedPrepDPTrain, MockedPrepVasp, MockedRunDPTrain, @@ -79,8 +37,6 @@ MockedStage, MockedStage1, MockedStage2, - make_mocked_init_data, - make_mocked_init_models, mocked_incar_template, mocked_numb_models, mocked_numb_select, @@ -89,34 +45,17 @@ from dpgen2.constants import ( fp_task_pattern, - lmp_conf_name, - lmp_input_name, - lmp_log_name, - lmp_traj_name, model_name_pattern, - train_log_name, - train_script_name, train_task_pattern, ) -from dpgen2.exploration.report import ( - ExplorationReport, -) from dpgen2.exploration.scheduler import ( ExplorationScheduler, ) -from dpgen2.exploration.task import ( - ExplorationStage, - ExplorationTask, - ExplorationTaskGroup, -) from dpgen2.flow.dpgen_loop import ( ConcurrentLearning, ) from dpgen2.fp.vasp import ( VaspInputs, - vasp_conf_name, - vasp_input_name, - vasp_pot_name, ) from dpgen2.op.prep_lmp import ( PrepLmp, diff --git a/tests/test_merge_caly_evo_step.py b/tests/test_merge_caly_evo_step.py index 3f22ba56..8c04cafe 100644 --- a/tests/test_merge_caly_evo_step.py +++ b/tests/test_merge_caly_evo_step.py @@ -1,102 +1,47 @@ -import json import os -import pickle import shutil import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import jsonpickle -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, PythonOPTemplate, Slices, ) - -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_index_pattern, - calypso_run_opt_file, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, -) from mocked_ops import ( MockedCollRunCaly, - MockedPrepCalyDPOptim, MockedRunCalyDPOptim, mocked_numb_models, ) from dpgen2.constants import ( - lmp_conf_name, - lmp_input_name, - lmp_log_name, - lmp_model_devi_name, - lmp_task_pattern, - lmp_traj_name, - model_name_pattern, - train_log_name, - train_script_name, - train_task_pattern, -) -from dpgen2.exploration.task import ( - ExplorationTask, - ExplorationTaskGroup, + calypso_check_opt_file, + calypso_index_pattern, + calypso_run_opt_file, ) from dpgen2.op import ( PrepCalyDPOptim, - RunCalyDPOptim, ) from dpgen2.op.caly_evo_step_merge import ( CalyEvoStepMerge, ) -from dpgen2.op.collect_run_caly import ( - CollRunCaly, -) -from dpgen2.op.prep_caly_input import ( - PrepCalyInput, -) -from dpgen2.superop.caly_evo_step import ( - CalyEvoStep, -) from dpgen2.utils.step_config import normalize as normalize_step_dict default_config = normalize_step_dict( @@ -170,7 +115,7 @@ def setUp(self) -> None: self.file_storage.mkdir(parents=True, exist_ok=True) for i in range(5): self.file_storage.joinpath(f"POSCAR_{i}").write_text(f"POSCAR_{i}") - self.file_storage.joinpath(f"frozen_model.pb").write_text(f"model.{i}.pb") + self.file_storage.joinpath("frozen_model.pb").write_text(f"model.{i}.pb") self.caly_run_opt_file = self.file_storage.joinpath(calypso_run_opt_file) self.caly_run_opt_file.write_text("caly_run_opt_script") self.caly_check_opt_file = self.file_storage.joinpath(calypso_check_opt_file) @@ -236,7 +181,7 @@ def setUp(self): for ii in range(self.nmodels): model_path = self.work_dir.joinpath(f"task.{ii}") model_path.mkdir(exist_ok=True, parents=True) - model = model_path.joinpath(f"model.ckpt.pt") + model = model_path.joinpath("model.ckpt.pt") model.write_text(f"model {ii}") self.model_list.append(model) self.models = upload_artifact(self.model_list) diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py index e949410d..37aa1b3b 100644 --- a/tests/test_prep_run_caly.py +++ b/tests/test_prep_run_caly.py @@ -1,56 +1,11 @@ -import json -import os -import pickle import shutil import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import jsonpickle -import numpy as np -from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, - Step, - Steps, - Workflow, - argo_range, - download_artifact, - upload_artifact, -) -from dflow.python import ( - OP, - OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, -) - -from dpgen2.constants import ( - calypso_check_opt_file, - calypso_index_pattern, - calypso_input_file, - calypso_run_opt_file, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from context import ( default_host, default_image, @@ -58,6 +13,11 @@ skip_ut_with_dflow_reason, upload_python_packages, ) +from dflow import ( + Step, + Workflow, + upload_artifact, +) from mocked_ops import ( MockedCollRunCaly, MockedRunCalyDPOptim, @@ -65,6 +25,11 @@ mocked_numb_models, ) +from dpgen2.constants import ( + calypso_check_opt_file, + calypso_input_file, + calypso_run_opt_file, +) from dpgen2.exploration.task import ( BaseExplorationTaskGroup, ExplorationTask, @@ -81,9 +46,6 @@ from dpgen2.op.prep_caly_model_devi import ( PrepCalyModelDevi, ) -from dpgen2.op.run_caly_model_devi import ( - RunCalyModelDevi, -) from dpgen2.superop.caly_evo_step import ( CalyEvoStep, ) @@ -125,7 +87,7 @@ def setUp(self): for ii in range(self.nmodels): model_path = self.work_dir.joinpath(f"task.{ii}") model_path.mkdir(parents=True, exist_ok=True) - model = model_path.joinpath(f"frozen_model.pb") + model = model_path.joinpath("frozen_model.pb") model.write_text(f"model {ii}") self.model_list.append(model) self.models = upload_artifact(self.model_list) diff --git a/tests/test_prep_run_diffcsp.py b/tests/test_prep_run_diffcsp.py index d5569d17..2f247402 100644 --- a/tests/test_prep_run_diffcsp.py +++ b/tests/test_prep_run_diffcsp.py @@ -119,6 +119,6 @@ def testPrepRunDiffCSP(self): self.assertEqual(len(model_devis), 4) def tearDown(self): - for d in glob.glob("test-prep-run-diffcsp-*") + ["task.000000", "task.000001"]: + for d in [*glob.glob("test-prep-run-diffcsp-*"), "task.000000", "task.000001"]: if os.path.isdir(d): shutil.rmtree(d) diff --git a/tests/test_prep_run_dp_labeling.py b/tests/test_prep_run_dp_labeling.py index 67256f86..bc505e0a 100644 --- a/tests/test_prep_run_dp_labeling.py +++ b/tests/test_prep_run_dp_labeling.py @@ -1,6 +1,4 @@ -import os import shutil -import sys import unittest from pathlib import ( Path, @@ -8,9 +6,6 @@ import dpdata import numpy as np -from dargs import ( - Argument, -) from dflow.python import ( FatalError, ) diff --git a/tests/test_prep_run_dp_train.py b/tests/test_prep_run_dp_train.py index 536ca4b9..f71ddef7 100644 --- a/tests/test_prep_run_dp_train.py +++ b/tests/test_prep_run_dp_train.py @@ -1,4 +1,3 @@ -import json import os import shutil import time @@ -6,48 +5,23 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, ) from mocked_ops import ( MockedPrepDPTrain, @@ -89,7 +63,10 @@ def _check_log( lines.append(" ".join(ww)) else: lines = lines_ - revised_fname = lambda ff: Path(ff).name if only_check_name else Path(ff) + + def revised_fname(ff): + return Path(ff).name if only_check_name else Path(ff) + tcase.assertEqual( lines[0].split(" "), ["init_model", str(revised_fname(Path(path) / init_model)), "OK"], @@ -234,7 +211,7 @@ def setUp(self): Path(self.train_scripts[ii]).write_text("{}") def tearDown(self): - for ii in ["init_data", "iter_data"] + self.task_names: + for ii in ["init_data", "iter_data", *self.task_names]: if Path(ii).exists(): shutil.rmtree(str(ii)) for ii in self.init_models: @@ -302,7 +279,7 @@ def setUp(self): ] def tearDown(self): - for ii in ["init_data", "iter_data"] + self.task_names: + for ii in ["init_data", "iter_data", *self.task_names]: if Path(ii).exists(): shutil.rmtree(str(ii)) for ii in self.str_init_models: diff --git a/tests/test_prep_run_lmp.py b/tests/test_prep_run_lmp.py index 3b350240..416f5ebe 100644 --- a/tests/test_prep_run_lmp.py +++ b/tests/test_prep_run_lmp.py @@ -1,6 +1,4 @@ -import json import os -import pickle import shutil import time import unittest @@ -9,47 +7,24 @@ ) from typing import ( List, - Set, ) -import jsonpickle -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, ) from mocked_ops import ( MockedRunLmp, @@ -63,10 +38,6 @@ lmp_model_devi_name, lmp_task_pattern, lmp_traj_name, - model_name_pattern, - train_log_name, - train_script_name, - train_task_pattern, ) from dpgen2.exploration.task import ( BaseExplorationTaskGroup, diff --git a/tests/test_prep_run_vasp.py b/tests/test_prep_run_vasp.py index 77030bff..4dadab45 100644 --- a/tests/test_prep_run_vasp.py +++ b/tests/test_prep_run_vasp.py @@ -1,4 +1,3 @@ -import json import os import shutil import time @@ -6,49 +5,23 @@ from pathlib import ( Path, ) -from typing import ( - List, - Set, -) -import jsonpickle -import numpy as np +# case of upload everything to argo, no context needed +from context import ( + default_host, + default_image, + skip_ut_with_dflow, + skip_ut_with_dflow_reason, + upload_python_packages, +) from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, Step, - Steps, Workflow, - argo_range, download_artifact, upload_artifact, ) from dflow.python import ( - OP, OPIO, - Artifact, - OPIOSign, - PythonOPTemplate, -) - -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from context import ( - default_host, - default_image, - skip_ut_with_dflow, - skip_ut_with_dflow_reason, - upload_python_packages, ) from mocked_ops import ( MockedPrepVasp, @@ -63,7 +36,6 @@ VaspInputs, vasp_conf_name, vasp_input_name, - vasp_pot_name, ) from dpgen2.superop.prep_run_fp import ( PrepRunFp, @@ -153,7 +125,7 @@ def setUp(self): work_path = Path(fp_task_pattern % ii) work_path.mkdir(exist_ok=True, parents=True) (work_path / vasp_conf_name).write_text(f"conf {ii}") - (work_path / vasp_input_name).write_text(f"incar template") + (work_path / vasp_input_name).write_text("incar template") self.task_list.append(work_path) def check_run_lmp_output( @@ -238,7 +210,7 @@ def check_run_vasp_output( fc = [] ii = int(task_name.split(".")[1]) fc.append(f"conf {ii}") - fc.append(f"incar template") + fc.append("incar template") self.assertEqual(fc, Path("log").read_text().strip().split("\n")) self.assertEqual( f"labeled_data of {task_name}\nconf {ii}", diff --git a/tests/test_select_confs.py b/tests/test_select_confs.py index 491d42f7..dc297924 100644 --- a/tests/test_select_confs.py +++ b/tests/test_select_confs.py @@ -1,53 +1,17 @@ -import json import os -import shutil -import time import unittest from pathlib import ( Path, ) -from typing import ( - List, - Set, - Tuple, -) -import jsonpickle -import numpy as np -from dflow import ( - InputArtifact, - InputParameter, - Inputs, - OutputArtifact, - OutputParameter, - Outputs, - S3Artifact, - Step, - Steps, - Workflow, - argo_range, - download_artifact, - upload_artifact, -) from dflow.python import ( - OP, OPIO, - Artifact, FatalError, - OPIOSign, - PythonOPTemplate, ) -try: - from context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass +# case of upload everything to argo, no context needed from mocked_ops import ( MockedConfSelector, - MockedExplorationReport, MockedSelectConfs, ) diff --git a/tests/utils/context.py b/tests/utils/context.py index d6638725..5504dedc 100644 --- a/tests/utils/context.py +++ b/tests/utils/context.py @@ -2,4 +2,3 @@ import sys sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) -import dpgen2 diff --git a/tests/utils/test_binary_file_input.py b/tests/utils/test_binary_file_input.py index 3e627802..189f22a9 100644 --- a/tests/utils/test_binary_file_input.py +++ b/tests/utils/test_binary_file_input.py @@ -8,9 +8,6 @@ import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils import ( BinaryFileInput, ) @@ -147,7 +144,8 @@ def serialization(obj): ) # check binary file - tensor = np.random.random((3, 2)) + rng = np.random.default_rng() + tensor = rng.random((3, 2)) np.save(self.task_output_path / "tensor.npy", tensor) t = BinaryFileInput(self.task_output_path / "tensor.npy", "npy") t = serialization(t) diff --git a/tests/utils/test_bohrium_config.py b/tests/utils/test_bohrium_config.py index bd084631..311d1482 100644 --- a/tests/utils/test_bohrium_config.py +++ b/tests/utils/test_bohrium_config.py @@ -1,16 +1,4 @@ -import json -import os -import random -import shutil -import tempfile -import unittest -from pathlib import ( - Path, -) - import dflow -import dpdata -import numpy as np import pytest from dflow.config import ( config, @@ -21,9 +9,6 @@ ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils import ( bohrium_config_from_dict, ) diff --git a/tests/utils/test_dflow_config.py b/tests/utils/test_dflow_config.py index e5b2d420..6a4604bd 100644 --- a/tests/utils/test_dflow_config.py +++ b/tests/utils/test_dflow_config.py @@ -1,24 +1,11 @@ -import json -import os -import random -import shutil -import tempfile import unittest -from pathlib import ( - Path, -) -import dpdata -import numpy as np from dflow.config import ( config, s3_config, ) # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils import ( dflow_config, dflow_s3_config, @@ -74,7 +61,7 @@ def test_s3_config(self): dflow_s3_config(config_data) self.assertEqual(s3_config["endpoint"], "bar") - def test_none(self): + def test_none_s3_config(self): config_data = { "endpoint": None, } diff --git a/tests/utils/test_dflow_query.py b/tests/utils/test_dflow_query.py index dad71643..64a59723 100644 --- a/tests/utils/test_dflow_query.py +++ b/tests/utils/test_dflow_query.py @@ -1,41 +1,5 @@ -import os -import textwrap import unittest -from pathlib import ( - Path, -) -from typing import ( - List, - Set, -) - -import numpy as np - -try: - from exploration.context import ( - dpgen2, - ) -except ModuleNotFoundError: - # case of upload everything to argo, no context needed - pass -from dflow.python import ( - FatalError, -) -from dpgen2.exploration.report import ( - ExplorationReport, -) -from dpgen2.exploration.scheduler import ( - ConvergenceCheckStageScheduler, - ExplorationScheduler, -) -from dpgen2.exploration.selector import ( - ConfSelectorFrames, -) -from dpgen2.exploration.task import ( - ExplorationStage, - ExplorationTaskGroup, -) from dpgen2.utils.dflow_query import ( find_slice_ranges, get_all_schedulers, @@ -48,15 +12,10 @@ sort_slice_ops, ) +# case of upload everything to argo, no context needed + + # isort: off -import utils.context -from mocked_ops import ( - MockedExplorationReport, - MockedExplorationTaskGroup, - MockedExplorationTaskGroup1, - MockedStage, - MockedStage1, -) # isort: on @@ -198,7 +157,7 @@ def test_get_last_scheduler(self): ) self.assertEqual(value, 10) - def test_get_last_scheduler(self): + def test_get_last_scheduler_none_global_false(self): value = get_last_scheduler( MockedWF(none_global=False), ["iter-1--scheduler", "foo", "bar", "iter-0--scheduler", "init--scheduler"], @@ -220,7 +179,7 @@ def test_sort_slice_ops(self): idxes = find_slice_ranges(dpgen_keys, "run-lmp") self.assertEqual(idxes, [[8, 14], [30, 36]]) - def test_sort_slice_ops(self): + def test_sort_slice_ops_nice_format(self): expected_output = [ "init--scheduler", "init--id", @@ -266,7 +225,7 @@ def test_print_keys(self): " 3 -> 5 : iter-000000--run-train-0000 -> iter-000000--run-train-0002", " 6 : iter-000000--prep-run-train", ] - expected_output = "\n".join(expected_output + [""]) + expected_output = "\n".join([*expected_output, ""]) ret = print_keys_in_nice_format( dpgen_keys[:7], diff --git a/tests/utils/test_dl_dpgen2_arti.py b/tests/utils/test_dl_dpgen2_arti.py index 037a0d80..37cd8b18 100644 --- a/tests/utils/test_dl_dpgen2_arti.py +++ b/tests/utils/test_dl_dpgen2_arti.py @@ -1,22 +1,15 @@ -import json -import os -import random import shutil -import tempfile import unittest from pathlib import ( Path, ) +from typing import ( + ClassVar, +) -import dflow -import dpdata import mock -import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.entrypoint.watch import ( update_finished_steps, ) @@ -48,7 +41,7 @@ class MockedStep: class Mockedwf: - keys = [ + keys: ClassVar = [ "iter-0--prep-run-train", ] diff --git a/tests/utils/test_dl_dpgen2_arti_by_def.py b/tests/utils/test_dl_dpgen2_arti_by_def.py index e6a30a32..4dfa4a23 100644 --- a/tests/utils/test_dl_dpgen2_arti_by_def.py +++ b/tests/utils/test_dl_dpgen2_arti_by_def.py @@ -1,23 +1,16 @@ -import json -import os -import random import shutil -import tempfile import textwrap import unittest from pathlib import ( Path, ) +from typing import ( + ClassVar, +) -import dflow -import dpdata import mock -import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils.download_dpgen2_artifacts import ( DownloadDefinition, _get_all_iterations, @@ -59,7 +52,7 @@ def __getitem__( class Mockedwf: - keys = [ + keys: ClassVar = [ "iter-000000--prep-run-train", "iter-000001--prep-run-train", "iter-000000--prep-run-explore", diff --git a/tests/utils/test_ele_temp.py b/tests/utils/test_ele_temp.py index 6a7de791..221a3353 100644 --- a/tests/utils/test_ele_temp.py +++ b/tests/utils/test_ele_temp.py @@ -7,9 +7,6 @@ import numpy as np # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils import ( setup_ele_temp, ) diff --git a/tests/utils/test_run_command.py b/tests/utils/test_run_command.py index 4b5b93bd..b65fcd6c 100644 --- a/tests/utils/test_run_command.py +++ b/tests/utils/test_run_command.py @@ -1,4 +1,3 @@ -import json import os import shutil import unittest @@ -6,12 +5,7 @@ Path, ) -import numpy as np - # isort: off -from .context import ( - dpgen2, -) from dpgen2.utils.run_command import ( run_command, ) diff --git a/tests/utils/test_step_config.py b/tests/utils/test_step_config.py index 002902a8..43683fdb 100644 --- a/tests/utils/test_step_config.py +++ b/tests/utils/test_step_config.py @@ -1,6 +1,3 @@ -import json -import os -import shutil import unittest from contextlib import ( contextmanager, @@ -8,26 +5,15 @@ from copy import ( deepcopy, ) -from pathlib import ( - Path, -) import dflow -import numpy as np -from dflow.python import ( - OPIO, -) # isort: off -from .context import ( - dpgen2, -) from dpgen2.constants import ( default_image, ) # from dpgen2.utils.step_config import normalize, gen_doc, init_executor -from dpgen2.utils import gen_doc_step_dict as gen_doc from dpgen2.utils import ( init_executor, )