From 7df79b02fcdf04a58b5cdb3f166b90543ce0e523 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jan 2025 22:23:27 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../atomic_model/linear_atomic_model.py | 6 +-- deepmd/dpmodel/descriptor/dpa1.py | 6 +-- deepmd/dpmodel/descriptor/dpa2.py | 12 +++--- deepmd/dpmodel/descriptor/hybrid.py | 4 +- deepmd/dpmodel/descriptor/se_t_tebd.py | 6 +-- deepmd/dpmodel/fitting/general_fitting.py | 6 +-- .../dpmodel/fitting/polarizability_fitting.py | 12 +++--- deepmd/dpmodel/utils/network.py | 2 +- deepmd/dpmodel/utils/type_embed.py | 18 ++++---- deepmd/infer/model_devi.py | 2 +- deepmd/loggers/training.py | 2 +- deepmd/main.py | 2 +- deepmd/pd/entrypoints/main.py | 24 +++++------ deepmd/pd/infer/deep_eval.py | 12 +++--- deepmd/pd/model/descriptor/descriptor.py | 6 +-- deepmd/pd/model/descriptor/dpa1.py | 12 +++--- deepmd/pd/model/descriptor/dpa2.py | 18 ++++---- deepmd/pd/model/descriptor/se_a.py | 6 +-- deepmd/pd/model/descriptor/se_t_tebd.py | 12 +++--- deepmd/pd/model/network/network.py | 18 ++++---- deepmd/pd/model/task/fitting.py | 12 +++--- deepmd/pd/train/training.py | 26 ++++++------ deepmd/pd/train/wrapper.py | 42 +++++++++---------- deepmd/pd/utils/multi_task.py | 6 +-- deepmd/pd/utils/stat.py | 6 +-- deepmd/pt/entrypoints/main.py | 24 +++++------ deepmd/pt/infer/deep_eval.py | 12 +++--- deepmd/pt/loss/dos.py | 8 ++-- deepmd/pt/loss/tensor.py | 6 +-- .../model/atomic_model/linear_atomic_model.py | 6 +-- deepmd/pt/model/descriptor/descriptor.py | 6 +-- deepmd/pt/model/descriptor/dpa1.py | 18 ++++---- deepmd/pt/model/descriptor/dpa2.py | 24 +++++------ deepmd/pt/model/descriptor/hybrid.py | 6 +-- deepmd/pt/model/descriptor/se_a.py | 6 +-- deepmd/pt/model/descriptor/se_r.py | 6 +-- deepmd/pt/model/descriptor/se_t.py | 6 +-- deepmd/pt/model/descriptor/se_t_tebd.py | 12 +++--- deepmd/pt/model/network/network.py | 18 ++++---- deepmd/pt/model/task/fitting.py | 12 +++--- deepmd/pt/model/task/polarizability.py | 18 ++++---- deepmd/pt/optimizer/LKF.py | 2 +- deepmd/pt/train/training.py | 24 +++++------ deepmd/pt/train/wrapper.py | 42 +++++++++---------- deepmd/pt/utils/multi_task.py | 6 +-- deepmd/pt/utils/stat.py | 6 +-- deepmd/tf/descriptor/se_a.py | 6 +-- deepmd/tf/descriptor/se_atten.py | 18 ++++---- deepmd/tf/descriptor/se_r.py | 6 +-- deepmd/tf/descriptor/se_t.py | 12 +++--- deepmd/tf/fit/dos.py | 12 +++--- deepmd/tf/fit/ener.py | 12 +++--- deepmd/tf/fit/polar.py | 6 +-- deepmd/tf/infer/deep_eval.py | 6 +-- deepmd/tf/infer/deep_tensor.py | 18 ++++---- deepmd/tf/loss/tensor.py | 12 +++--- deepmd/tf/model/frozen.py | 2 +- deepmd/tf/train/trainer.py | 12 +++--- deepmd/tf/utils/finetune.py | 14 +++---- deepmd/tf/utils/graph.py | 6 +-- deepmd/tf/utils/tabulate.py | 6 +-- deepmd/utils/argcheck.py | 3 +- deepmd/utils/data.py | 16 +++---- deepmd/utils/data_system.py | 6 +-- deepmd/utils/pair_tab.py | 6 +-- deepmd/utils/spin.py | 2 +- source/lmp/tests/test_dplr.py | 2 +- source/lmp/tests/test_lammps.py | 12 +++--- source/lmp/tests/test_lammps_3types.py | 4 +- source/lmp/tests/test_lammps_dpa_jax.py | 8 ++-- source/lmp/tests/test_lammps_dpa_pt.py | 8 ++-- source/lmp/tests/test_lammps_dpa_pt_nopbc.py | 8 ++-- source/lmp/tests/test_lammps_dpa_sel_pt.py | 8 ++-- source/lmp/tests/test_lammps_faparam.py | 2 +- source/lmp/tests/test_lammps_jax.py | 8 ++-- source/lmp/tests/test_lammps_pd.py | 8 ++-- source/lmp/tests/test_lammps_pt.py | 8 ++-- source/lmp/tests/write_lmp_data.py | 2 +- source/tests/consistent/io/test_io.py | 2 +- source/tests/pd/test_multitask.py | 12 +++--- source/tests/pt/test_multitask.py | 12 +++--- source/tests/tf/common.py | 6 +-- source/tests/tf/test_descrpt_se_a_mask.py | 6 +-- .../universal/dpmodel/fitting/test_fitting.py | 6 +-- 84 files changed, 426 insertions(+), 423 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 8108292bd2..ce0f1d0cb9 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -457,9 +457,9 @@ def _compute_weight( list[np.ndarray] the atomic ZBL weight for interpolation. (nframes, nloc, 1) """ - assert ( - self.sw_rmax > self.sw_rmin - ), "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`." + assert self.sw_rmax > self.sw_rmin, ( + "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`." + ) xp = array_api_compat.array_namespace(extended_coord, extended_atype) dp_nlist = nlists_[0] diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 20a758b170..37b677cf23 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -432,9 +432,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_atten obj.ntypes = len(type_map) diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 55ae331593..b91347ccb5 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -543,9 +543,9 @@ def init_subclass_params(sub_data, sub_class): ) self.rcsl_list.sort() for ii in range(1, len(self.rcsl_list)): - assert ( - self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1] - ), "rcut and sel are not in the same order" + assert self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1], ( + "rcut and sel are not in the same order" + ) self.rcut_list = [ii[0] for ii in self.rcsl_list] self.nsel_list = [ii[1] for ii in self.rcsl_list] self.use_econf_tebd = use_econf_tebd @@ -678,9 +678,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.type_embedding.change_type_map(type_map=type_map) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 5c1a7d2785..f050bb6222 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -71,7 +71,9 @@ def __init__( for ii in range(1, self.numb_descrpt): assert ( self.descrpt_list[ii].get_ntypes() == self.descrpt_list[0].get_ntypes() - ), f"number of atom types in {ii}th descriptor {self.descrpt_list[0].__class__.__name__} does not match others" + ), ( + f"number of atom types in {ii}th descriptor {self.descrpt_list[0].__class__.__name__} does not match others" + ) # if hybrid sel is larger than sub sel, the nlist needs to be cut for each type hybrid_sel = self.get_sel() nlist_cut_idx: list[np.ndarray] = [] diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index c350e3eb47..570f9a47e8 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -270,9 +270,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_ttebd obj.ntypes = len(type_map) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index c05d84c4a1..388932f297 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -242,9 +242,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 0db6a23377..8acb818a46 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -224,9 +224,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) super().change_type_map(type_map=type_map) @@ -280,9 +280,9 @@ def call( """ xp = array_api_compat.array_namespace(descriptor, atype) nframes, nloc, _ = descriptor.shape - assert ( - gr is not None - ), "Must provide the rotation matrix for polarizability fitting." + assert gr is not None, ( + "Must provide the rotation matrix for polarizability fitting." + ) # (nframes, nloc, _net_out_dim) out = self._call_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ self.var_name diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 18071f8eca..703cbdd339 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -551,7 +551,7 @@ def check_shape_consistency(self) -> None: if self.layers[ii].dim_out() != self.layers[ii + 1].dim_in(): raise ValueError( f"the dim of layer {ii} output {self.layers[ii].dim_out} ", - f"does not match the dim of layer {ii+1} ", + f"does not match the dim of layer {ii + 1} ", f"output {self.layers[ii].dim_out}", ) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index e28b6abb31..17e40f3592 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -166,9 +166,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) if not self.use_econf_tebd: do_resnet = self.neuron[0] in [ @@ -177,9 +177,9 @@ def change_type_map( len(type_map), len(type_map) * 2, ] - assert ( - not do_resnet or self.activation_function == "Linear" - ), "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + assert not do_resnet or self.activation_function == "Linear", ( + "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + ) first_layer_matrix = self.embedding_net.layers[0].w eye_vector = np.eye(self.ntypes, dtype=PRECISION_DICT[self.precision]) # preprocess for resnet connection @@ -227,9 +227,9 @@ def get_econf_tebd(type_map, precision: str = "default"): ) from deepmd.utils.econf_embd import type_map as periodic_table - assert ( - type_map is not None - ), "When using electronic configuration type embedding, type_map must be provided!" + assert type_map is not None, ( + "When using electronic configuration type embedding, type_map must be provided!" + ) missing_types = [t for t in type_map if t not in periodic_table] assert not missing_types, ( diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py index b3e3330835..1828d40c49 100644 --- a/deepmd/infer/model_devi.py +++ b/deepmd/infer/model_devi.py @@ -211,7 +211,7 @@ def write_model_devi_out( f"min_devi_{item}", f"avg_devi_{item}", ) - header += f'{"devi_e":19s}' + header += f"{'devi_e':19s}" if atomic: header += f"{'atm_devi_f(N)':19s}" with open(fname, "ab") as fp: diff --git a/deepmd/loggers/training.py b/deepmd/loggers/training.py index fa643d2257..16d3eb1618 100644 --- a/deepmd/loggers/training.py +++ b/deepmd/loggers/training.py @@ -9,7 +9,7 @@ def format_training_message( wall_time: float, ) -> str: """Format a training message.""" - return f"batch {batch:7d}: " f"total wall time = {wall_time:.2f} s" + return f"batch {batch:7d}: total wall time = {wall_time:.2f} s" def format_training_message_per_task( diff --git a/deepmd/main.py b/deepmd/main.py index 097588ca0a..52e98f6afb 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -72,7 +72,7 @@ class DeprecateAction(argparse.Action): def __init__(self, *args, **kwargs) -> None: self.call_count = 0 if "help" in kwargs: - kwargs["help"] = f'[DEPRECATED] {kwargs["help"]}' + kwargs["help"] = f"[DEPRECATED] {kwargs['help']}" super().__init__(*args, **kwargs) def __call__(self, parser, namespace, values, option_string=None): diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index bcc93432af..e5f06a5487 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -247,9 +247,9 @@ def train( if multi_task: config["model"], shared_links = preprocess_shared_params(config["model"]) # handle the special key - assert ( - "RANDOM" not in config["model"]["model_dict"] - ), "Model name can not be 'RANDOM' in multi-task mode!" + assert "RANDOM" not in config["model"]["model_dict"], ( + "Model name can not be 'RANDOM' in multi-task mode!" + ) # update fine-tuning config finetune_links = None @@ -404,9 +404,9 @@ def change_bias( multi_task = "model_dict" in model_params bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic" if multi_task: - assert ( - model_branch is not None - ), "For multitask model, the model branch must be set!" + assert model_branch is not None, ( + "For multitask model, the model branch must be set!" + ) assert model_branch in model_params["model_dict"], ( f"For multitask model, the model branch must be in the 'model_dict'! " f"Available options are : {list(model_params['model_dict'].keys())}." @@ -427,12 +427,12 @@ def change_bias( if bias_value is not None: # use user-defined bias - assert model_to_change.model_type in [ - "ener" - ], "User-defined bias is only available for energy model!" - assert ( - len(bias_value) == len(type_map) - ), f"The number of elements in the bias should be the same as that in the type_map: {type_map}." + assert model_to_change.model_type in ["ener"], ( + "User-defined bias is only available for energy model!" + ) + assert len(bias_value) == len(type_map), ( + f"The number of elements in the bias should be the same as that in the type_map: {type_map}." + ) old_bias = model_to_change.get_out_bias() bias_to_set = paddle.to_tensor( bias_value, dtype=old_bias.dtype, place=old_bias.place diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index d13e837161..dee5b0176f 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -92,12 +92,12 @@ def __init__( model_keys = list(self.input_param["model_dict"].keys()) if isinstance(head, int): head = model_keys[0] - assert ( - head is not None - ), f"Head must be set for multitask model! Available heads are: {model_keys}" - assert ( - head in model_keys - ), f"No head named {head} in model! Available heads are: {model_keys}" + assert head is not None, ( + f"Head must be set for multitask model! Available heads are: {model_keys}" + ) + assert head in model_keys, ( + f"No head named {head} in model! Available heads are: {model_keys}" + ) self.input_param = self.input_param["model_dict"][head] state_dict_head = {"_extra_state": state_dict["_extra_state"]} for item in state_dict: diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index 36de5b1948..96976e7c2c 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -128,9 +128,9 @@ def share_params(self, base_class, shared_level, resume=False): If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) if shared_level == 0: # link buffers if hasattr(self, "mean"): diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index ebaaf96bc3..42cba44eb4 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -378,9 +378,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA1 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in both type_embedding and se_atten @@ -449,9 +449,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_atten obj.ntypes = len(type_map) diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index 8d4e13edae..9cd3d77f89 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -254,9 +254,9 @@ def init_subclass_params(sub_data, sub_class): ) self.rcsl_list.sort() for ii in range(1, len(self.rcsl_list)): - assert ( - self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1] - ), "rcut and sel are not in the same order" + assert self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1], ( + "rcut and sel are not in the same order" + ) self.rcut_list = [ii[0] for ii in self.rcsl_list] self.nsel_list = [ii[1] for ii in self.rcsl_list] self.use_econf_tebd = use_econf_tebd @@ -386,9 +386,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA2 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in type_embedding, repinit and repformers @@ -419,9 +419,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.type_embedding.change_type_map(type_map=type_map) diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 6e5832670d..c2fd81a02e 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -169,9 +169,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For SeA descriptors, the user-defined share-level # shared_level: 0 # share all parameters in sea diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index b28f79e436..b1f4af788d 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -246,9 +246,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA1 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in both type_embedding and se_ttebd @@ -317,9 +317,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_ttebd obj.ntypes = len(type_map) diff --git a/deepmd/pd/model/network/network.py b/deepmd/pd/model/network/network.py index 1974e526a0..9cdb7b3adc 100644 --- a/deepmd/pd/model/network/network.py +++ b/deepmd/pd/model/network/network.py @@ -103,9 +103,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only TypeEmbedNet of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only TypeEmbedNet of the same type can share params!" + ) if shared_level == 0: # the following will successfully link all the params except buffers, which need manually link. for item in self._sub_layers: @@ -231,9 +231,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) if not self.use_econf_tebd: do_resnet = self.neuron[0] in [ @@ -242,9 +242,9 @@ def change_type_map( len(type_map), len(type_map) * 2, ] - assert ( - not do_resnet or self.activation_function == "Linear" - ), "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + assert not do_resnet or self.activation_function == "Linear", ( + "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + ) first_layer_matrix = self.embedding_net.layers[0].matrix eye_vector = paddle.eye(self.ntypes, dtype=self.prec).to( device=first_layer_matrix.place diff --git a/deepmd/pd/model/task/fitting.py b/deepmd/pd/model/task/fitting.py index 6e96b7b081..a478c12f97 100644 --- a/deepmd/pd/model/task/fitting.py +++ b/deepmd/pd/model/task/fitting.py @@ -60,9 +60,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only fitting nets of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only fitting nets of the same type can share params!" + ) if shared_level == 0: # only not share the bias_atom_e and the case_embd # the following will successfully link all the params except buffers, which need manually link. @@ -259,9 +259,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 0f3c7a9732..4916eda42e 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -86,7 +86,7 @@ def format_training_message( eta: Optional[int] = None, ): """Format a training message.""" - msg = f"batch {batch:7d}: " f"total wall time = {wall_time:.2f} s" + msg = f"batch {batch:7d}: total wall time = {wall_time:.2f} s" if isinstance(eta, int): msg += f", eta = {datetime.timedelta(seconds=int(eta))!s}" return msg @@ -247,9 +247,9 @@ def get_sample(): return get_sample def get_lr(lr_params): - assert ( - lr_params.get("type", "exp") == "exp" - ), "Only learning rate `exp` is supported!" + assert lr_params.get("type", "exp") == "exp", ( + "Only learning rate `exp` is supported!" + ) lr_params["stop_steps"] = self.num_steps - self.warmup_steps lr_exp = LearningRateExp(**lr_params) return lr_exp @@ -260,9 +260,9 @@ def get_lr(lr_params): missing_keys = [ key for key in self.model_keys if key not in self.optim_dict ] - assert ( - not missing_keys - ), f"These keys are not in optim_dict: {missing_keys}!" + assert not missing_keys, ( + f"These keys are not in optim_dict: {missing_keys}!" + ) self.opt_type = {} self.opt_param = {} for model_key in self.model_keys: @@ -386,9 +386,9 @@ def get_lr(lr_params): # Learning rate self.warmup_steps = training_params.get("warmup_steps", 0) self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) - assert ( - self.num_steps - self.warmup_steps > 0 or self.warmup_steps == 0 - ), "Warm up steps must be less than total training steps!" + assert self.num_steps - self.warmup_steps > 0 or self.warmup_steps == 0, ( + "Warm up steps must be less than total training steps!" + ) if self.multi_task and config.get("learning_rate_dict", None) is not None: self.lr_exp = {} for model_key in self.model_keys: @@ -1204,9 +1204,9 @@ def get_model_for_wrapper(_model_params, resuming=False): def get_case_embd_config(_model_params): - assert ( - "model_dict" in _model_params - ), "Only support setting case embedding for multi-task model!" + assert "model_dict" in _model_params, ( + "Only support setting case embedding for multi-task model!" + ) model_keys = list(_model_params["model_dict"]) sorted_model_keys = sorted(model_keys) numb_case_embd_list = [ diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index 2263a6e9b9..bd28b17c88 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -46,9 +46,9 @@ def __init__( elif isinstance(model, dict): self.multi_task = True for task_key in model: - assert isinstance( - model[task_key], paddle.nn.Layer - ), f"{task_key} in model_dict is not a paddle.nn.Layer!" + assert isinstance(model[task_key], paddle.nn.Layer), ( + f"{task_key} in model_dict is not a paddle.nn.Layer!" + ) self.model[task_key] = model[task_key] # Loss self.loss = None @@ -58,9 +58,9 @@ def __init__( self.loss["Default"] = loss elif isinstance(loss, dict): for task_key in loss: - assert isinstance( - loss[task_key], paddle.nn.Layer - ), f"{task_key} in loss_dict is not a paddle.nn.Layer!" + assert isinstance(loss[task_key], paddle.nn.Layer), ( + f"{task_key} in loss_dict is not a paddle.nn.Layer!" + ) self.loss[task_key] = loss[task_key] self.inference_only = self.loss is None @@ -90,12 +90,12 @@ def share_params(self, shared_links, resume=False) -> None: class_type_link = link_item["shared_type"] model_key_link = link_item["model_key"] shared_level_link = int(link_item["shared_level"]) - assert ( - shared_level_link >= shared_level_base - ), "The shared_links must be sorted by shared_level!" - assert ( - "descriptor" in class_type_link - ), f"Class type mismatched: {class_type_base} vs {class_type_link}!" + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert "descriptor" in class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) if class_type_link == "descriptor": link_class = self.model[model_key_link].get_descriptor() elif "hybrid" in class_type_link: @@ -119,12 +119,12 @@ def share_params(self, shared_links, resume=False) -> None: class_type_link = link_item["shared_type"] model_key_link = link_item["model_key"] shared_level_link = int(link_item["shared_level"]) - assert ( - shared_level_link >= shared_level_base - ), "The shared_links must be sorted by shared_level!" - assert ( - class_type_base == class_type_link - ), f"Class type mismatched: {class_type_base} vs {class_type_link}!" + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert class_type_base == class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) link_class = self.model[ model_key_link ].atomic_model.__getattr__(class_type_link) @@ -152,9 +152,9 @@ def forward( if not self.multi_task: task_key = "Default" else: - assert ( - task_key is not None - ), f"Multitask model must specify the inference task! Supported tasks are {list(self.model.keys())}." + assert task_key is not None, ( + f"Multitask model must specify the inference task! Supported tasks are {list(self.model.keys())}." + ) input_dict = { "coord": coord, "atype": atype, diff --git a/deepmd/pd/utils/multi_task.py b/deepmd/pd/utils/multi_task.py index 321883c12e..92d7f380e5 100644 --- a/deepmd/pd/utils/multi_task.py +++ b/deepmd/pd/utils/multi_task.py @@ -105,9 +105,9 @@ def replace_one_item( if ":" in key_in_dict: shared_key = key_in_dict.split(":")[0] shared_level = int(key_in_dict.split(":")[1]) - assert ( - shared_key in shared_dict - ), f"Appointed {shared_type} {shared_key} are not in the shared_dict! Please check the input params." + assert shared_key in shared_dict, ( + f"Appointed {shared_type} {shared_key} are not in the shared_dict! Please check the input params." + ) if index is None: params_dict[shared_type] = deepcopy(shared_dict[shared_key]) else: diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py index a8bdbd6415..9588790abf 100644 --- a/deepmd/pd/utils/stat.py +++ b/deepmd/pd/utils/stat.py @@ -588,9 +588,9 @@ def compute_output_stats_atomic( # correction for missing types missing_types = ntypes - merged_natoms[kk].max() - 1 if missing_types > 0: - assert ( - bias_atom_e[kk].dtype is std_atom_e[kk].dtype - ), "bias and std should be of the same dtypes" + assert bias_atom_e[kk].dtype is std_atom_e[kk].dtype, ( + "bias and std should be of the same dtypes" + ) nan_padding = np.empty( (missing_types, bias_atom_e[kk].shape[1]), dtype=bias_atom_e[kk].dtype, diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index fd4be73e84..576c5a1811 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -269,9 +269,9 @@ def train( if multi_task: config["model"], shared_links = preprocess_shared_params(config["model"]) # handle the special key - assert ( - "RANDOM" not in config["model"]["model_dict"] - ), "Model name can not be 'RANDOM' in multi-task mode!" + assert "RANDOM" not in config["model"]["model_dict"], ( + "Model name can not be 'RANDOM' in multi-task mode!" + ) # update fine-tuning config finetune_links = None @@ -407,9 +407,9 @@ def change_bias( multi_task = "model_dict" in model_params bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic" if multi_task: - assert ( - model_branch is not None - ), "For multitask model, the model branch must be set!" + assert model_branch is not None, ( + "For multitask model, the model branch must be set!" + ) assert model_branch in model_params["model_dict"], ( f"For multitask model, the model branch must be in the 'model_dict'! " f"Available options are : {list(model_params['model_dict'].keys())}." @@ -431,12 +431,12 @@ def change_bias( if bias_value is not None: # use user-defined bias - assert model_to_change.model_type in [ - "ener" - ], "User-defined bias is only available for energy model!" - assert ( - len(bias_value) == len(type_map) - ), f"The number of elements in the bias should be the same as that in the type_map: {type_map}." + assert model_to_change.model_type in ["ener"], ( + "User-defined bias is only available for energy model!" + ) + assert len(bias_value) == len(type_map), ( + f"The number of elements in the bias should be the same as that in the type_map: {type_map}." + ) old_bias = model_to_change.get_out_bias() bias_to_set = torch.tensor( bias_value, dtype=old_bias.dtype, device=old_bias.device diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 06126b5dab..d1cc4fb82f 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -115,12 +115,12 @@ def __init__( model_keys = list(self.input_param["model_dict"].keys()) if isinstance(head, int): head = model_keys[0] - assert ( - head is not None - ), f"Head must be set for multitask model! Available heads are: {model_keys}" - assert ( - head in model_keys - ), f"No head named {head} in model! Available heads are: {model_keys}" + assert head is not None, ( + f"Head must be set for multitask model! Available heads are: {model_keys}" + ) + assert head in model_keys, ( + f"No head named {head} in model! Available heads are: {model_keys}" + ) self.input_param = self.input_param["model_dict"][head] state_dict_head = {"_extra_state": state_dict["_extra_state"]} for item in state_dict: diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 30a3f715ef..493cc85694 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -79,9 +79,11 @@ def __init__( self.has_ados = (start_pref_ados != 0.0 and limit_pref_ados != 0.0) or inference self.has_acdf = (start_pref_acdf != 0.0 and limit_pref_acdf != 0.0) or inference - assert ( - self.has_dos or self.has_cdf or self.has_ados or self.has_acdf - ), AssertionError("Can not assian zero weight both to `pref` and `pref_atomic`") + assert self.has_dos or self.has_cdf or self.has_ados or self.has_acdf, ( + AssertionError( + "Can not assian zero weight both to `pref` and `pref_atomic`" + ) + ) def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): """Return loss on local and global tensors. diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 69b133de58..0acc3989be 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -55,9 +55,9 @@ def __init__( self.inference = inference self.enable_atomic_weight = enable_atomic_weight - assert ( - self.local_weight >= 0.0 and self.global_weight >= 0.0 - ), "Can not assign negative weight to `pref` and `pref_atomic`" + assert self.local_weight >= 0.0 and self.global_weight >= 0.0, ( + "Can not assign negative weight to `pref` and `pref_atomic`" + ) self.has_local_weight = self.local_weight > 0.0 or inference self.has_global_weight = self.global_weight > 0.0 or inference assert self.has_local_weight or self.has_global_weight, AssertionError( diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 36c636ddfb..3d894dc3a0 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -603,9 +603,9 @@ def _compute_weight( list[torch.Tensor] the atomic ZBL weight for interpolation. (nframes, nloc, 1) """ - assert ( - self.sw_rmax > self.sw_rmin - ), "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`." + assert self.sw_rmax > self.sw_rmin, ( + "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`." + ) dp_nlist = nlists_[0] zbl_nlist = nlists_[1] diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index d97f8964c4..46fa4cf933 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -132,9 +132,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) if shared_level == 0: # link buffers if hasattr(self, "mean"): diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 47b2a4d15c..fae50a138c 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -384,9 +384,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA1 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in both type_embedding and se_atten @@ -451,9 +451,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_atten obj.ntypes = len(type_map) @@ -588,9 +588,9 @@ def enable_compression( # do some checks before the mocel compression process if self.compress: raise ValueError("Compression is already enabled.") - assert ( - not self.se_atten.resnet_dt - ), "Model compression error: descriptor resnet_dt must be false!" + assert not self.se_atten.resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) for tt in self.se_atten.exclude_types: if (tt[0] not in range(self.se_atten.ntypes)) or ( tt[1] not in range(self.se_atten.ntypes) diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index f086a346b6..969fdca5fc 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -258,9 +258,9 @@ def init_subclass_params(sub_data, sub_class): ) self.rcsl_list.sort() for ii in range(1, len(self.rcsl_list)): - assert ( - self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1] - ), "rcut and sel are not in the same order" + assert self.rcsl_list[ii - 1][1] <= self.rcsl_list[ii][1], ( + "rcut and sel are not in the same order" + ) self.rcut_list = [ii[0] for ii in self.rcsl_list] self.nsel_list = [ii[1] for ii in self.rcsl_list] self.use_econf_tebd = use_econf_tebd @@ -390,9 +390,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA2 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in type_embedding, repinit and repformers @@ -421,9 +421,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.type_embedding.change_type_map(type_map=type_map) @@ -899,9 +899,9 @@ def enable_compression( # do some checks before the mocel compression process if self.compress: raise ValueError("Compression is already enabled.") - assert ( - not self.repinit.resnet_dt - ), "Model compression error: repinit resnet_dt must be false!" + assert not self.repinit.resnet_dt, ( + "Model compression error: repinit resnet_dt must be false!" + ) for tt in self.repinit.exclude_types: if (tt[0] not in range(self.repinit.ntypes)) or ( tt[1] not in range(self.repinit.ntypes) diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 5caa1243f1..e13b014037 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -170,9 +170,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) if shared_level == 0: for ii, des in enumerate(self.descrpt_list): self.descrpt_list[ii].share_params( diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 550154300b..9d3b850d0f 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -192,9 +192,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For SeA descriptors, the user-defined share-level # shared_level: 0 # share all parameters in sea diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index f25fb93fa7..a91757460c 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -231,9 +231,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For SeR descriptors, the user-defined share-level # shared_level: 0 if shared_level == 0: diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index d41a28d3db..fb444c3d13 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -226,9 +226,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For SeT descriptors, the user-defined share-level # shared_level: 0 # share all parameters in sea diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 0c7b37bb58..7e27805bd5 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -246,9 +246,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only descriptors of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) # For DPA1 descriptors, the user-defined share-level # shared_level: 0 # share all parameters in both type_embedding and se_ttebd @@ -313,9 +313,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) obj = self.se_ttebd obj.ntypes = len(type_map) diff --git a/deepmd/pt/model/network/network.py b/deepmd/pt/model/network/network.py index 5b4d741a3b..8a2bd4f75c 100644 --- a/deepmd/pt/model/network/network.py +++ b/deepmd/pt/model/network/network.py @@ -311,9 +311,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only TypeEmbedNet of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only TypeEmbedNet of the same type can share params!" + ) if shared_level == 0: # the following will successfully link all the params except buffers, which need manually link. for item in self._modules: @@ -434,9 +434,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) if not self.use_econf_tebd: do_resnet = self.neuron[0] in [ @@ -445,9 +445,9 @@ def change_type_map( len(type_map), len(type_map) * 2, ] - assert ( - not do_resnet or self.activation_function == "Linear" - ), "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + assert not do_resnet or self.activation_function == "Linear", ( + "'activation_function' must be 'Linear' when performing type changing on resnet structure!" + ) first_layer_matrix = self.embedding_net.layers[0].matrix.data eye_vector = torch.eye( self.ntypes, dtype=self.prec, device=first_layer_matrix.device diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 2486ab576f..dcbbbe9602 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -60,9 +60,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. """ - assert ( - self.__class__ == base_class.__class__ - ), "Only fitting nets of the same type can share params!" + assert self.__class__ == base_class.__class__, ( + "Only fitting nets of the same type can share params!" + ) if shared_level == 0: # only not share the bias_atom_e and the case_embd # the following will successfully link all the params except buffers, which need manually link. @@ -259,9 +259,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index d9a421d635..a326802918 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -107,9 +107,9 @@ def __init__( self.scale = [1.0 for _ in range(ntypes)] else: if isinstance(self.scale, list): - assert ( - len(self.scale) == ntypes - ), "Scale should be a list of length ntypes." + assert len(self.scale) == ntypes, ( + "Scale should be a list of length ntypes." + ) elif isinstance(self.scale, float): self.scale = [self.scale for _ in range(ntypes)] else: @@ -168,9 +168,9 @@ def change_type_map( """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ - assert ( - self.type_map is not None - ), "'type_map' must be defined when performing type changing!" + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) super().change_type_map(type_map=type_map) @@ -234,9 +234,9 @@ def forward( aparam: Optional[torch.Tensor] = None, ): nframes, nloc, _ = descriptor.shape - assert ( - gr is not None - ), "Must provide the rotation matrix for polarizability fitting." + assert gr is not None, ( + "Must provide the rotation matrix for polarizability fitting." + ) # cast the input to internal precsion gr = gr.to(self.prec) # (nframes, nloc, _net_out_dim) diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index 928fc0943a..c342960e5b 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -40,7 +40,7 @@ def __init__( if len(self.param_groups) != 1 or len(self._params) == 0: raise ValueError( - "LKF doesn't support per-parameter options " "(parameter groups)" + "LKF doesn't support per-parameter options (parameter groups)" ) # NOTE: LKF has only global state, but we register it as state for diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index f3b4548b05..bbf77b4eba 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -240,9 +240,9 @@ def get_sample(): return get_sample def get_lr(lr_params): - assert ( - lr_params.get("type", "exp") == "exp" - ), "Only learning rate `exp` is supported!" + assert lr_params.get("type", "exp") == "exp", ( + "Only learning rate `exp` is supported!" + ) lr_params["stop_steps"] = self.num_steps - self.warmup_steps lr_exp = LearningRateExp(**lr_params) return lr_exp @@ -253,9 +253,9 @@ def get_lr(lr_params): missing_keys = [ key for key in self.model_keys if key not in self.optim_dict ] - assert ( - not missing_keys - ), f"These keys are not in optim_dict: {missing_keys}!" + assert not missing_keys, ( + f"These keys are not in optim_dict: {missing_keys}!" + ) self.opt_type = {} self.opt_param = {} for model_key in self.model_keys: @@ -384,9 +384,9 @@ def get_lr(lr_params): # Learning rate self.warmup_steps = training_params.get("warmup_steps", 0) self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) - assert ( - self.num_steps - self.warmup_steps > 0 or self.warmup_steps == 0 - ), "Warm up steps must be less than total training steps!" + assert self.num_steps - self.warmup_steps > 0 or self.warmup_steps == 0, ( + "Warm up steps must be less than total training steps!" + ) if self.multi_task and config.get("learning_rate_dict", None) is not None: self.lr_exp = {} for model_key in self.model_keys: @@ -1312,9 +1312,9 @@ def get_model_for_wrapper( def get_case_embd_config(_model_params): - assert ( - "model_dict" in _model_params - ), "Only support setting case embedding for multi-task model!" + assert "model_dict" in _model_params, ( + "Only support setting case embedding for multi-task model!" + ) model_keys = list(_model_params["model_dict"]) sorted_model_keys = sorted(model_keys) numb_case_embd_list = [ diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index f0253c283e..9a2cbff295 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -41,9 +41,9 @@ def __init__( elif isinstance(model, dict): self.multi_task = True for task_key in model: - assert isinstance( - model[task_key], torch.nn.Module - ), f"{task_key} in model_dict is not a torch.nn.Module!" + assert isinstance(model[task_key], torch.nn.Module), ( + f"{task_key} in model_dict is not a torch.nn.Module!" + ) self.model[task_key] = model[task_key] # Loss self.loss = None @@ -53,9 +53,9 @@ def __init__( self.loss["Default"] = loss elif isinstance(loss, dict): for task_key in loss: - assert isinstance( - loss[task_key], torch.nn.Module - ), f"{task_key} in loss_dict is not a torch.nn.Module!" + assert isinstance(loss[task_key], torch.nn.Module), ( + f"{task_key} in loss_dict is not a torch.nn.Module!" + ) self.loss[task_key] = loss[task_key] self.inference_only = self.loss is None @@ -88,12 +88,12 @@ def share_params(self, shared_links, resume=False) -> None: class_type_link = link_item["shared_type"] model_key_link = link_item["model_key"] shared_level_link = int(link_item["shared_level"]) - assert ( - shared_level_link >= shared_level_base - ), "The shared_links must be sorted by shared_level!" - assert ( - "descriptor" in class_type_link - ), f"Class type mismatched: {class_type_base} vs {class_type_link}!" + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert "descriptor" in class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) if class_type_link == "descriptor": link_class = self.model[model_key_link].get_descriptor() elif "hybrid" in class_type_link: @@ -120,12 +120,12 @@ def share_params(self, shared_links, resume=False) -> None: class_type_link = link_item["shared_type"] model_key_link = link_item["model_key"] shared_level_link = int(link_item["shared_level"]) - assert ( - shared_level_link >= shared_level_base - ), "The shared_links must be sorted by shared_level!" - assert ( - class_type_base == class_type_link - ), f"Class type mismatched: {class_type_base} vs {class_type_link}!" + assert shared_level_link >= shared_level_base, ( + "The shared_links must be sorted by shared_level!" + ) + assert class_type_base == class_type_link, ( + f"Class type mismatched: {class_type_base} vs {class_type_link}!" + ) link_class = self.model[ model_key_link ].atomic_model.__getattr__(class_type_link) @@ -153,9 +153,9 @@ def forward( if not self.multi_task: task_key = "Default" else: - assert ( - task_key is not None - ), f"Multitask model must specify the inference task! Supported tasks are {list(self.model.keys())}." + assert task_key is not None, ( + f"Multitask model must specify the inference task! Supported tasks are {list(self.model.keys())}." + ) input_dict = { "coord": coord, "atype": atype, diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py index a4de8810af..6c397400bf 100644 --- a/deepmd/pt/utils/multi_task.py +++ b/deepmd/pt/utils/multi_task.py @@ -105,9 +105,9 @@ def replace_one_item( if ":" in key_in_dict: shared_key = key_in_dict.split(":")[0] shared_level = int(key_in_dict.split(":")[1]) - assert ( - shared_key in shared_dict - ), f"Appointed {shared_type} {shared_key} are not in the shared_dict! Please check the input params." + assert shared_key in shared_dict, ( + f"Appointed {shared_type} {shared_key} are not in the shared_dict! Please check the input params." + ) if index is None: params_dict[shared_type] = deepcopy(shared_dict[shared_key]) else: diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 710d392ac3..f3ac3334c2 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -604,9 +604,9 @@ def compute_output_stats_atomic( # correction for missing types missing_types = ntypes - merged_natoms[kk].max() - 1 if missing_types > 0: - assert ( - bias_atom_e[kk].dtype is std_atom_e[kk].dtype - ), "bias and std should be of the same dtypes" + assert bias_atom_e[kk].dtype is std_atom_e[kk].dtype, ( + "bias and std should be of the same dtypes" + ) nan_padding = np.empty( (missing_types, bias_atom_e[kk].shape[1]), dtype=bias_atom_e[kk].dtype, diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index c9a6d6ec57..e3ae7bf99c 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -470,9 +470,9 @@ def enable_compression( The suffix of the scope """ # do some checks before the mocel compression process - assert ( - not self.filter_resnet_dt - ), "Model compression error: descriptor resnet_dt must be false!" + assert not self.filter_resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) for tt in self.exclude_types: if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)): raise RuntimeError( diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index b6d180ee7c..3a9b86a0d6 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -253,9 +253,9 @@ def __init__( Constructor """ if not (nvnmd_cfg.enable and (nvnmd_cfg.version == 1)): - assert Version(TF_VERSION) > Version( - "2" - ), "se_atten only support tensorflow version 2.0 or higher." + assert Version(TF_VERSION) > Version("2"), ( + "se_atten only support tensorflow version 2.0 or higher." + ) if ntypes == 0: raise ValueError("`model/type_map` is not set or empty!") self.stripped_type_embedding = stripped_type_embedding @@ -449,9 +449,9 @@ def enable_compression( The suffix of the type embedding scope, only for DescrptDPA1Compat """ # do some checks before the mocel compression process - assert ( - not self.filter_resnet_dt - ), "Model compression error: descriptor resnet_dt must be false!" + assert not self.filter_resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) for tt in self.exclude_types: if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)): raise RuntimeError( @@ -2363,9 +2363,9 @@ def enable_compression( tebd_suffix : str, optional Same as suffix. """ - assert ( - tebd_suffix == "" - ), "DescrptDPA1Compat must use the same tebd_suffix as suffix!" + assert tebd_suffix == "", ( + "DescrptDPA1Compat must use the same tebd_suffix as suffix!" + ) super().enable_compression( min_nbor_dist, graph, diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 4cc867cf0d..ed66d6ad25 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -346,9 +346,9 @@ def enable_compression( suffix : str, optional The suffix of the scope """ - assert ( - not self.filter_resnet_dt - ), "Model compression error: descriptor resnet_dt must be false!" + assert not self.filter_resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) self.compress = True self.table = DPTabulate( diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 79be372c10..c5d50744af 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -353,9 +353,9 @@ def enable_compression( suffix : str, optional The suffix of the scope """ - assert ( - not self.filter_resnet_dt - ), "Model compression error: descriptor resnet_dt must be false!" + assert not self.filter_resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) self.compress = True self.table = DPTabulate( @@ -830,9 +830,9 @@ def deserialize_network(cls, data: dict, suffix: str = "") -> dict: """ embedding_net_variables = {} embeddings = NetworkCollection.deserialize(data) - assert ( - embeddings.ndim == 2 - ), "Embeddings in descriptor 'se_e3' must have two dimensions." + assert embeddings.ndim == 2, ( + "Embeddings in descriptor 'se_e3' must have two dimensions." + ) for ii in range(embeddings.ntypes**embeddings.ndim): net_idx = [] rest_ii = ii diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 1da0e55a92..8fa3167bfc 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -155,9 +155,9 @@ def __init__( self.trainable = [True for ii in range(len(self.n_neuron) + 1)] if isinstance(self.trainable, bool): self.trainable = [self.trainable] * (len(self.n_neuron) + 1) - assert ( - len(self.trainable) == len(self.n_neuron) + 1 - ), "length of trainable should be that of n_neuron + 1" + assert len(self.trainable) == len(self.n_neuron) + 1, ( + "length of trainable should be that of n_neuron + 1" + ) self.useBN = False self.bias_dos = np.zeros((self.ntypes, self.numb_dos), dtype=np.float64) @@ -173,9 +173,9 @@ def __init__( self.layer_name = layer_name if self.layer_name is not None: assert isinstance(self.layer_name, list), "layer_name should be a list" - assert ( - len(self.layer_name) == len(self.n_neuron) + 1 - ), "length of layer_name should be that of n_neuron + 1" + assert len(self.layer_name) == len(self.n_neuron) + 1, ( + "length of layer_name should be that of n_neuron + 1" + ) self.mixed_types = mixed_types self.type_map = type_map diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index d791a2eba0..e10468df32 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -214,9 +214,9 @@ def __init__( self.trainable = [True for ii in range(len(self.n_neuron) + 1)] if isinstance(self.trainable, bool): self.trainable = [self.trainable] * (len(self.n_neuron) + 1) - assert ( - len(self.trainable) == len(self.n_neuron) + 1 - ), "length of trainable should be that of n_neuron + 1" + assert len(self.trainable) == len(self.n_neuron) + 1, ( + "length of trainable should be that of n_neuron + 1" + ) self.atom_ener = [] self.atom_ener_v = atom_ener for at, ae in enumerate(atom_ener if atom_ener is not None else []): @@ -240,9 +240,9 @@ def __init__( self.layer_name = layer_name if self.layer_name is not None: assert isinstance(self.layer_name, list), "layer_name should be a list" - assert ( - len(self.layer_name) == len(self.n_neuron) + 1 - ), "length of layer_name should be that of n_neuron + 1" + assert len(self.layer_name) == len(self.n_neuron) + 1, ( + "length of layer_name should be that of n_neuron + 1" + ) self.mixed_types = mixed_types self.tebd_dim = 0 diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index 14dd6ee092..3091fd56ca 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -136,9 +136,9 @@ def __init__( self.scale = np.array([1.0 for ii in range(self.ntypes)]) else: if isinstance(self.scale, list): - assert ( - len(self.scale) == ntypes - ), "Scale should be a list of length ntypes." + assert len(self.scale) == ntypes, ( + "Scale should be a list of length ntypes." + ) elif isinstance(self.scale, float): self.scale = [self.scale for _ in range(ntypes)] else: diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index db4faffed9..8e5e3deea5 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -803,9 +803,9 @@ def _prepare_feed_dict( assert aparam is not None aparam = np.array(aparam) if self.has_efield: - assert ( - efield is not None - ), "you are using a model with external field, parameter efield should be provided" + assert efield is not None, ( + "you are using a model with external field, parameter efield should be provided" + ) efield = np.array(efield) # reshape the inputs diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index bdd36b6e09..a4d30c4f2f 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -73,9 +73,9 @@ def __init__( ) # check model type model_type = self.tensors["t_tensor"][2:-2] - assert ( - self.model_type == model_type - ), f"expect {model_type} model but got {self.model_type}" + assert self.model_type == model_type, ( + f"expect {model_type} model but got {self.model_type}" + ) # now load tensors to object attributes for attr_name, tensor_name in self.tensors.items(): @@ -246,14 +246,14 @@ def eval( feed_dict_test[self.t_mesh] = mesh if atomic: - assert ( - "global" not in self.model_type - ), f"cannot do atomic evaluation with model type {self.model_type}" + assert "global" not in self.model_type, ( + f"cannot do atomic evaluation with model type {self.model_type}" + ) t_out = [self.t_tensor] else: - assert ( - self._support_gfv or "global" in self.model_type - ), f"do not support global tensor evaluation with old {self.model_type} model" + assert self._support_gfv or "global" in self.model_type, ( + f"do not support global tensor evaluation with old {self.model_type} model" + ) t_out = [self.t_global_tensor if self._support_gfv else self.t_tensor] v_out = self.sess.run(t_out, feed_dict=feed_dict_test) tensor = v_out[0] diff --git a/deepmd/tf/loss/tensor.py b/deepmd/tf/loss/tensor.py index dbd7da99df..54247a66a4 100644 --- a/deepmd/tf/loss/tensor.py +++ b/deepmd/tf/loss/tensor.py @@ -42,12 +42,12 @@ def __init__(self, jdata, **kwarg) -> None: self.global_weight = jdata.get("pref", None) self.enable_atomic_weight = jdata.get("enable_atomic_weight", False) - assert ( - self.local_weight is not None and self.global_weight is not None - ), "Both `pref` and `pref_atomic` should be provided." - assert ( - self.local_weight >= 0.0 and self.global_weight >= 0.0 - ), "Can not assign negative weight to `pref` and `pref_atomic`" + assert self.local_weight is not None and self.global_weight is not None, ( + "Both `pref` and `pref_atomic` should be provided." + ) + assert self.local_weight >= 0.0 and self.global_weight >= 0.0, ( + "Can not assign negative weight to `pref` and `pref_atomic`" + ) assert (self.local_weight > 0.0) or (self.global_weight > 0.0), AssertionError( "Can not assian zero weight both to `pref` and `pref_atomic`" ) diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 2099a8d947..6ca18ed7bd 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -72,7 +72,7 @@ def __init__(self, model_file: str, **kwargs) -> None: self.model_type = "ener" else: raise NotImplementedError( - "This model type has not been implemented. " "Contribution is welcome!" + "This model type has not been implemented. Contribution is welcome!" ) def build( diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index f8010091d8..3e0495c050 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -175,9 +175,9 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix="") -> Non self.stop_batch = stop_batch if not self.is_compress and data.mixed_type: - assert isinstance( - self.fitting, EnerFitting - ), "Data in mixed_type format must use ener fitting!" + assert isinstance(self.fitting, EnerFitting), ( + "Data in mixed_type format must use ener fitting!" + ) if self.numb_fparam > 0: log.info(f"training with {self.numb_fparam} frame parameter(s)") @@ -872,9 +872,9 @@ def _init_from_pretrained_model( ) from e else: self.model_type = bytes.decode(t_model_type) - assert ( - self.model_type != "compressed_model" - ), "Compressed models are not supported for finetuning!" + assert self.model_type != "compressed_model", ( + "Compressed models are not supported for finetuning!" + ) self.model.init_variables(graph, graph_def, model_type=self.model_type) log.info( f"Changing energy bias in pretrained model for types {origin_type_map!s}... " diff --git a/deepmd/tf/utils/finetune.py b/deepmd/tf/utils/finetune.py index 4c57246ffd..9d331cec9a 100644 --- a/deepmd/tf/utils/finetune.py +++ b/deepmd/tf/utils/finetune.py @@ -39,14 +39,12 @@ def replace_model_params_with_pretrained_model( pretrained_jdata = json.loads(t_jdata) # Check the model type - assert ( - pretrained_jdata["model"]["descriptor"]["type"] - in [ - "se_atten", - "se_atten_v2", - ] - and pretrained_jdata["model"]["fitting_net"]["type"] in ["ener"] - ), "The finetune process only supports models pretrained with 'se_atten' or 'se_atten_v2' descriptor and 'ener' fitting_net!" + assert pretrained_jdata["model"]["descriptor"]["type"] in [ + "se_atten", + "se_atten_v2", + ] and pretrained_jdata["model"]["fitting_net"]["type"] in ["ener"], ( + "The finetune process only supports models pretrained with 'se_atten' or 'se_atten_v2' descriptor and 'ener' fitting_net!" + ) # Check the type map pretrained_type_map = pretrained_jdata["model"]["type_map"] diff --git a/deepmd/tf/utils/graph.py b/deepmd/tf/utils/graph.py index 4fccaac0e8..44020cc20f 100644 --- a/deepmd/tf/utils/graph.py +++ b/deepmd/tf/utils/graph.py @@ -324,9 +324,9 @@ def get_fitting_net_nodes_from_graph_def( fitting_net_pattern = FITTING_NET_PATTERN fitting_net_nodes = get_pattern_nodes_from_graph_def(graph_def, fitting_net_pattern) for key in fitting_net_nodes.keys(): - assert ( - key.find("bias") > 0 or key.find("matrix") > 0 or key.find("idt") > 0 - ), "currently, only support weight matrix, bias and idt at the model compression process!" + assert key.find("bias") > 0 or key.find("matrix") > 0 or key.find("idt") > 0, ( + "currently, only support weight matrix, bias and idt at the model compression process!" + ) return fitting_net_nodes diff --git a/deepmd/tf/utils/tabulate.py b/deepmd/tf/utils/tabulate.py index 2244262786..67ae2f96bd 100644 --- a/deepmd/tf/utils/tabulate.py +++ b/deepmd/tf/utils/tabulate.py @@ -139,9 +139,9 @@ def __init__( self.graph_def, suffix=self.suffix ) for key in self.embedding_net_nodes.keys(): - assert ( - key.find("bias") > 0 or key.find("matrix") > 0 - ), "currently, only support weight matrix and bias matrix at the tabulation op!" + assert key.find("bias") > 0 or key.find("matrix") > 0, ( + "currently, only support weight matrix and bias matrix at the tabulation op!" + ) # move it to the descriptor class # for tt in self.exclude_types: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 66a6308fc5..1d897ceb57 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1147,8 +1147,7 @@ def dpa2_repformer_args(): "and `update_residual_init`." ) doc_update_residual = ( - "When update using residual mode, " - "the initial std of residual vector weights." + "When update using residual mode, the initial std of residual vector weights." ) doc_update_residual_init = ( "When update using residual mode, " diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index d572efd321..c5119cf31f 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -79,9 +79,9 @@ def __init__( self.natoms = len(self.atom_type) # load atom type map self.type_map = self._load_type_map(root) - assert ( - optional_type_map or self.type_map is not None - ), f"System {sys_path} must have type_map.raw in this mode! " + assert optional_type_map or self.type_map is not None, ( + f"System {sys_path} must have type_map.raw in this mode! " + ) if self.type_map is not None: assert len(self.type_map) >= max(self.atom_type) + 1 # check pbc @@ -196,9 +196,9 @@ def reduce(self, key_out: str, key_in: str): assert key_in in self.data_dict, "cannot find input key" assert self.data_dict[key_in]["atomic"], "reduced property should be atomic" assert key_out not in self.data_dict, "output key should not have been added" - assert ( - self.data_dict[key_in]["repeat"] == 1 - ), "reduced properties should not have been repeated" + assert self.data_dict[key_in]["repeat"] == 1, ( + "reduced properties should not have been repeated" + ) self.data_dict[key_out] = { "ndof": self.data_dict[key_in]["ndof"], @@ -569,7 +569,9 @@ def _load_set(self, set_name: DPPath): ).T assert ( atom_type_nums.sum(axis=-1) + ghost_nums.sum(axis=-1) == natoms - ).all(), f"some types in 'real_atom_types.npy' of set {set_name} are not contained in {self.get_ntypes()} types!" + ).all(), ( + f"some types in 'real_atom_types.npy' of set {set_name} are not contained in {self.get_ntypes()} types!" + ) data["real_natoms_vec"] = np.concatenate( ( np.tile(np.array([natoms, natoms], dtype=np.int32), (nframes, 1)), diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index e443fa673a..fbc9c1a684 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -693,9 +693,9 @@ def process_sys_probs(sys_probs, nbatch): type_filter = sys_probs >= 0 assigned_sum_prob = np.sum(type_filter * sys_probs) # 1e-8 is to handle floating point error; See #1917 - assert ( - assigned_sum_prob <= 1.0 + 1e-8 - ), "the sum of assigned probability should be less than 1" + assert assigned_sum_prob <= 1.0 + 1e-8, ( + "the sum of assigned probability should be less than 1" + ) rest_sum_prob = 1.0 - assigned_sum_prob if not np.isclose(rest_sum_prob, 0): rest_nbatch = (1 - type_filter) * nbatch diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 05c68a5d67..93c8b7a1f9 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -64,9 +64,9 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: ncol = self.vdata.shape[1] - 1 n0 = (-1 + np.sqrt(1 + 8 * ncol)) * 0.5 self.ntypes = int(n0 + 0.1) - assert ( - self.ntypes * (self.ntypes + 1) // 2 == ncol - ), f"number of volumes provided in {filename} does not match guessed number of types {self.ntypes}" + assert self.ntypes * (self.ntypes + 1) // 2 == ncol, ( + f"number of volumes provided in {filename} does not match guessed number of types {self.ntypes}" + ) # check table data against rcut and update tab_file if needed, table upper boundary is used as rcut if not provided. self.rcut = rcut if rcut is not None else self.rmax diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index 9ea5fc0713..c6cbb32f52 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -70,7 +70,7 @@ def __init__( else: raise ValueError( f"Invalid length of virtual_scale for spin atoms" - f": Expected {self.ntypes_real} or { self.ntypes_spin} but got {len(virtual_scale)}!" + f": Expected {self.ntypes_real} or {self.ntypes_spin} but got {len(virtual_scale)}!" ) elif isinstance(virtual_scale, float): self.virtual_scale = [virtual_scale for _ in range(self.ntypes_real)] diff --git a/source/lmp/tests/test_dplr.py b/source/lmp/tests/test_dplr.py index 532bafef36..21d1f18658 100644 --- a/source/lmp/tests/test_dplr.py +++ b/source/lmp/tests/test_dplr.py @@ -372,7 +372,7 @@ def test_pair_deepmd_sr_virial(lammps) -> None: lammps.compute("virial real_atom centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 real_atom custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py index 645d826e31..ad8f8cdaac 100644 --- a/source/lmp/tests/test_lammps.py +++ b/source/lmp/tests/test_lammps.py @@ -327,10 +327,10 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.variable("eatom atom c_peatom") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"pressure{jj} equal c_pressure[{ii+1}]") + lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -395,10 +395,10 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.variable("eatom atom c_peatom") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"pressure{jj} equal c_pressure[{ii+1}]") + lammps.variable(f"pressure{jj} equal c_pressure[{ii + 1}]") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -533,7 +533,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -592,7 +592,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_3types.py b/source/lmp/tests/test_lammps_3types.py index c8d3193817..f0cbe19ddf 100644 --- a/source/lmp/tests/test_lammps_3types.py +++ b/source/lmp/tests/test_lammps_3types.py @@ -310,7 +310,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -361,7 +361,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_dpa_jax.py b/source/lmp/tests/test_lammps_dpa_jax.py index 10428b2374..4867b5f84e 100644 --- a/source/lmp/tests/test_lammps_dpa_jax.py +++ b/source/lmp/tests/test_lammps_dpa_jax.py @@ -324,7 +324,7 @@ def test_pair_deepmd_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -375,7 +375,7 @@ def test_pair_deepmd_model_devi_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -499,7 +499,7 @@ def test_pair_deepmd_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -558,7 +558,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_dpa_pt.py b/source/lmp/tests/test_lammps_dpa_pt.py index d92084faa1..e66b93e09e 100644 --- a/source/lmp/tests/test_lammps_dpa_pt.py +++ b/source/lmp/tests/test_lammps_dpa_pt.py @@ -320,7 +320,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -371,7 +371,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -495,7 +495,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -554,7 +554,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py index 2183852770..563650c714 100644 --- a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py +++ b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py @@ -318,7 +318,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -369,7 +369,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -493,7 +493,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -552,7 +552,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_dpa_sel_pt.py b/source/lmp/tests/test_lammps_dpa_sel_pt.py index 57395fda98..9ff2883fc1 100644 --- a/source/lmp/tests/test_lammps_dpa_sel_pt.py +++ b/source/lmp/tests/test_lammps_dpa_sel_pt.py @@ -323,7 +323,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -374,7 +374,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -498,7 +498,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -557,7 +557,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_faparam.py b/source/lmp/tests/test_lammps_faparam.py index b388f0453b..4206aa68fb 100644 --- a/source/lmp/tests/test_lammps_faparam.py +++ b/source/lmp/tests/test_lammps_faparam.py @@ -203,7 +203,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_jax.py b/source/lmp/tests/test_lammps_jax.py index 6d67cd3203..5d88cfca12 100644 --- a/source/lmp/tests/test_lammps_jax.py +++ b/source/lmp/tests/test_lammps_jax.py @@ -322,7 +322,7 @@ def test_pair_deepmd_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -373,7 +373,7 @@ def test_pair_deepmd_model_devi_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -497,7 +497,7 @@ def test_pair_deepmd_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -556,7 +556,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_pd.py b/source/lmp/tests/test_lammps_pd.py index 7824fa5b0f..31ee2e482a 100644 --- a/source/lmp/tests/test_lammps_pd.py +++ b/source/lmp/tests/test_lammps_pd.py @@ -323,7 +323,7 @@ def test_pair_deepmd_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -374,7 +374,7 @@ def test_pair_deepmd_model_devi_virial(lammps): lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -500,7 +500,7 @@ def test_pair_deepmd_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -565,7 +565,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/test_lammps_pt.py b/source/lmp/tests/test_lammps_pt.py index f2a4bac23d..9aed014b62 100644 --- a/source/lmp/tests/test_lammps_pt.py +++ b/source/lmp/tests/test_lammps_pt.py @@ -320,7 +320,7 @@ def test_pair_deepmd_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -371,7 +371,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -495,7 +495,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) @@ -554,7 +554,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.compute("virial all centroid/stress/atom NULL pair") for ii in range(9): jj = [0, 4, 8, 3, 6, 7, 1, 2, 5][ii] - lammps_real.variable(f"virial{jj} atom c_virial[{ii+1}]") + lammps_real.variable(f"virial{jj} atom c_virial[{ii + 1}]") lammps_real.dump( "1 all custom 1 dump id " + " ".join([f"v_virial{ii}" for ii in range(9)]) ) diff --git a/source/lmp/tests/write_lmp_data.py b/source/lmp/tests/write_lmp_data.py index e84fc481f0..1cd0eed2ca 100644 --- a/source/lmp/tests/write_lmp_data.py +++ b/source/lmp/tests/write_lmp_data.py @@ -43,7 +43,7 @@ def write_lmp_data_full( f.write(f"{box[6]:.10e} {box[7]:.10e} {box[8]:.10e} xy xz yz\n") f.write("\nMasses\n\n") for i in range(3): - f.write(f"{i+1:d} {mass_list[i]:.10e}\n") + f.write(f"{i + 1:d} {mass_list[i]:.10e}\n") f.write("\nAtoms\n\n") for i in range(natom): f.write( diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index bc7d3e7af2..d6b7348a91 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -213,7 +213,7 @@ def test_deep_eval(self) -> None: # expect all nan if not supported continue np.testing.assert_allclose( - vv1, vv2, rtol=1e-12, atol=1e-12, err_msg=f"backend {idx+1}" + vv1, vv2, rtol=1e-12, atol=1e-12, err_msg=f"backend {idx + 1}" ) diff --git a/source/tests/pd/test_multitask.py b/source/tests/pd/test_multitask.py index 72ad251068..49c9dbd003 100644 --- a/source/tests/pd/test_multitask.py +++ b/source/tests/pd/test_multitask.py @@ -96,9 +96,9 @@ def test_multitask_train(self) -> None: self.origin_config["training"]["data_dict"]["model_2"] ) self.origin_config["training"]["data_dict"]["model_3"]["stat_file"] = ( - self.origin_config[ - "training" - ]["data_dict"]["model_3"]["stat_file"].replace("model_2", "model_3") + self.origin_config["training"]["data_dict"]["model_3"]["stat_file"].replace( + "model_2", "model_3" + ) ) # add model_4 @@ -115,9 +115,9 @@ def test_multitask_train(self) -> None: self.origin_config["training"]["data_dict"]["model_2"] ) self.origin_config["training"]["data_dict"]["model_4"]["stat_file"] = ( - self.origin_config[ - "training" - ]["data_dict"]["model_4"]["stat_file"].replace("model_2", "model_4") + self.origin_config["training"]["data_dict"]["model_4"]["stat_file"].replace( + "model_2", "model_4" + ) ) # set finetune rules diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py index 62964abad3..38e91067d3 100644 --- a/source/tests/pt/test_multitask.py +++ b/source/tests/pt/test_multitask.py @@ -96,9 +96,9 @@ def test_multitask_train(self) -> None: self.origin_config["training"]["data_dict"]["model_2"] ) self.origin_config["training"]["data_dict"]["model_3"]["stat_file"] = ( - self.origin_config[ - "training" - ]["data_dict"]["model_3"]["stat_file"].replace("model_2", "model_3") + self.origin_config["training"]["data_dict"]["model_3"]["stat_file"].replace( + "model_2", "model_3" + ) ) # add model_4 @@ -115,9 +115,9 @@ def test_multitask_train(self) -> None: self.origin_config["training"]["data_dict"]["model_2"] ) self.origin_config["training"]["data_dict"]["model_4"]["stat_file"] = ( - self.origin_config[ - "training" - ]["data_dict"]["model_4"]["stat_file"].replace("model_2", "model_4") + self.origin_config["training"]["data_dict"]["model_4"]["stat_file"].replace( + "model_2", "model_4" + ) ) # set finetune rules diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index 6095af5ebf..53ee08e0dc 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -1046,9 +1046,9 @@ def process_sys_weights(self, sys_weights): sys_weights = np.array(sys_weights) type_filter = sys_weights >= 0 assigned_sum_prob = np.sum(type_filter * sys_weights) - assert ( - assigned_sum_prob <= 1 - ), "the sum of assigned probability should be less than 1" + assert assigned_sum_prob <= 1, ( + "the sum of assigned probability should be less than 1" + ) rest_sum_prob = 1.0 - assigned_sum_prob rest_nbatch = (1 - type_filter) * self.nbatches rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch) diff --git a/source/tests/tf/test_descrpt_se_a_mask.py b/source/tests/tf/test_descrpt_se_a_mask.py index 6de36adb38..04c02d1c61 100644 --- a/source/tests/tf/test_descrpt_se_a_mask.py +++ b/source/tests/tf/test_descrpt_se_a_mask.py @@ -242,9 +242,9 @@ def test_descriptor_se_a_mask(self) -> None: test_data = data.get_test() numb_test = 1 - assert ( - jdata["model"]["descriptor"]["type"] == "se_a_mask" - ), "Wrong descriptor type" + assert jdata["model"]["descriptor"]["type"] == "se_a_mask", ( + "Wrong descriptor type" + ) descrpt = DescrptSeAMask(**jdata["model"]["descriptor"], uniform_seed=True) t_coord = tf.placeholder( diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index 2fe0060003..90b0668d20 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -119,9 +119,9 @@ def FittingParamDipole( embedding_width=None, numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): - assert ( - embedding_width is not None - ), "embedding_width for dipole fitting is required." + assert embedding_width is not None, ( + "embedding_width for dipole fitting is required." + ) input_dict = { "ntypes": ntypes, "dim_descrpt": dim_descrpt,