diff --git a/backend/read_env.py b/backend/read_env.py index 3b217926d6..df358fc83a 100644 --- a/backend/read_env.py +++ b/backend/read_env.py @@ -131,7 +131,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]: ) -def set_scikit_build_env(): +def set_scikit_build_env() -> None: """Set scikit-build environment variables before executing scikit-build.""" cmake_minimum_required_version, cmake_args, _, _, _, _ = get_argument_from_env() os.environ["SKBUILD_CMAKE_MINIMUM_VERSION"] = cmake_minimum_required_version diff --git a/data/json/json2yaml.py b/data/json/json2yaml.py index f4714908bb..8a07b4a6eb 100644 --- a/data/json/json2yaml.py +++ b/data/json/json2yaml.py @@ -13,7 +13,7 @@ import yaml -def _main(): +def _main() -> None: parser = argparse.ArgumentParser( description="convert json config file to yaml", formatter_class=argparse.ArgumentDefaultsHelpFormatter, diff --git a/data/raw/copy_raw.py b/data/raw/copy_raw.py index 2a4078df4d..a0ea45277e 100755 --- a/data/raw/copy_raw.py +++ b/data/raw/copy_raw.py @@ -8,7 +8,7 @@ import numpy as np -def copy(in_dir, out_dir, ncopies=[1, 1, 1]): +def copy(in_dir, out_dir, ncopies=[1, 1, 1]) -> None: has_energy = os.path.isfile(in_dir + "/energy.raw") has_force = os.path.isfile(in_dir + "/force.raw") has_virial = os.path.isfile(in_dir + "/virial.raw") @@ -71,7 +71,7 @@ def copy(in_dir, out_dir, ncopies=[1, 1, 1]): np.savetxt(out_dir + "/ncopies.raw", ncopies, fmt="%d") -def _main(): +def _main() -> None: parser = argparse.ArgumentParser(description="parse copy raw args") parser.add_argument("INPUT", default=".", help="input dir of raw files") parser.add_argument("OUTPUT", default=".", help="output dir of copied raw files") diff --git a/data/raw/shuffle_raw.py b/data/raw/shuffle_raw.py index b4fc1457e5..690307cc2c 100755 --- a/data/raw/shuffle_raw.py +++ b/data/raw/shuffle_raw.py @@ -30,7 +30,7 @@ def detect_raw(path): return raws -def _main(): +def _main() -> None: args = _parse_args() raws = args.raws inpath = args.INPUT diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 0fbc447aaa..c5f742bbec 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -107,7 +107,7 @@ def calculate( atoms: Optional["Atoms"] = None, properties: list[str] = ["energy", "forces", "virial"], system_changes: list[str] = all_changes, - ): + ) -> None: """Run calculation with deepmd model. Parameters diff --git a/deepmd/common.py b/deepmd/common.py index 185722f4a8..1e66113306 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -62,13 +62,11 @@ if TYPE_CHECKING: _DICT_VAL = TypeVar("_DICT_VAL") - __all__.extend( - [ - "_DICT_VAL", - "_PRECISION", - "_ACTIVATION", - ] - ) + __all__ += [ + "_DICT_VAL", + "_PRECISION", + "_ACTIVATION", + ] def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarray: @@ -237,7 +235,7 @@ def get_np_precision(precision: "_PRECISION") -> np.dtype: raise RuntimeError(f"{precision} is not a valid precision") -def symlink_prefix_files(old_prefix: str, new_prefix: str): +def symlink_prefix_files(old_prefix: str, new_prefix: str) -> None: """Create symlinks from old checkpoint prefix to new one. On Windows this function will copy files instead of creating symlinks. diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 4e7620bdda..04cd5b0014 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -43,7 +43,7 @@ def __init__( pair_exclude_types: list[tuple[int, int]] = [], rcond: Optional[float] = None, preset_out_bias: Optional[dict[str, np.ndarray]] = None, - ): + ) -> None: super().__init__() self.type_map = type_map self.reinit_atom_exclude(atom_exclude_types) @@ -51,7 +51,7 @@ def __init__( self.rcond = rcond self.preset_out_bias = preset_out_bias - def init_out_stat(self): + def init_out_stat(self) -> None: """Initialize the output bias.""" ntypes = self.get_ntypes() self.bias_keys: list[str] = list(self.fitting_output_def().keys()) @@ -68,7 +68,7 @@ def init_out_stat(self): self.out_bias = out_bias_data self.out_std = out_std_data - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["out_bias"]: self.out_bias = value elif key in ["out_std"]: @@ -91,7 +91,7 @@ def get_type_map(self) -> list[str]: def reinit_atom_exclude( self, exclude_types: list[int] = [], - ): + ) -> None: self.atom_exclude_types = exclude_types if exclude_types == []: self.atom_excl = None @@ -101,7 +101,7 @@ def reinit_atom_exclude( def reinit_pair_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.pair_exclude_types = exclude_types if exclude_types == []: self.pair_excl = None diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 01b996b80e..749fe6bbf9 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -45,7 +45,7 @@ def __init__( fitting, type_map: list[str], **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) self.type_map = type_map self.descriptor = descriptor diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 243be6f2da..9676b34bfd 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -52,7 +52,7 @@ def __init__( models: list[BaseAtomicModel], type_map: list[str], **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -391,7 +391,7 @@ def __init__( type_map: list[str], smin_alpha: Optional[float] = 0.1, **kwargs, - ): + ) -> None: models = [dp_model, zbl_model] kwargs["models"] = models kwargs["type_map"] = type_map diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index ba7a85a3fe..aefdbf7f1c 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -66,7 +66,7 @@ def __init__( rcond: Optional[float] = None, atom_ener: Optional[list[float]] = None, **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() self.tab_file = tab_file diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index ecc450bcd2..1c5d2d1900 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -9,6 +9,6 @@ class DPPropertyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert isinstance(fitting, PropertyFittingNet) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/descriptor/descriptor.py b/deepmd/dpmodel/descriptor/descriptor.py index 746c02eb68..2c6e8fee57 100644 --- a/deepmd/dpmodel/descriptor/descriptor.py +++ b/deepmd/dpmodel/descriptor/descriptor.py @@ -6,6 +6,7 @@ ) from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -83,7 +84,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -106,7 +107,7 @@ def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -135,7 +136,7 @@ def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" -def extend_descrpt_stat(des, type_map, des_with_stat=None): +def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 259593e731..dd8acba872 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -3,6 +3,7 @@ from typing import ( Any, Callable, + NoReturn, Optional, Union, ) @@ -388,7 +389,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_atten.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -404,7 +405,9 @@ def dim_out(self): def dim_emb(self): return self.get_dim_emb() - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -783,7 +786,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -834,18 +837,18 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.""" raise NotImplementedError - def get_stats(self): + def get_stats(self) -> NoReturn: """Get the statistics of the descriptor.""" raise NotImplementedError def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) @@ -1077,7 +1080,7 @@ def __init__( smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a neighbor-wise attention net.""" super().__init__() self.layer_num = layer_num @@ -1132,7 +1135,7 @@ def __getitem__(self, key): else: raise TypeError(key) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if not isinstance(key, int): raise TypeError(key) if isinstance(value, self.network_type): @@ -1205,7 +1208,7 @@ def __init__( smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a neighbor-wise attention layer.""" super().__init__() self.nnei = nnei @@ -1311,7 +1314,7 @@ def __init__( smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a multi-head neighbor-wise attention net.""" super().__init__() assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 097be2ef09..6aa06c47c3 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + NoReturn, Optional, Union, ) @@ -88,7 +89,7 @@ def __init__( three_body_sel: int = 40, three_body_rcut: float = 4.0, three_body_rcut_smth: float = 0.5, - ): + ) -> None: r"""The constructor for the RepinitArgs class which defines the parameters of the repinit block in DPA2 descriptor. Parameters @@ -212,7 +213,7 @@ def __init__( g1_out_conv: bool = True, g1_out_mlp: bool = True, ln_eps: Optional[float] = 1e-5, - ): + ) -> None: r"""The constructor for the RepformerArgs class which defines the parameters of the repformer block in DPA2 descriptor. Parameters @@ -384,7 +385,7 @@ def __init__( use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: Optional[list[str]] = None, - ): + ) -> None: r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492. Parameters @@ -656,7 +657,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -728,7 +729,9 @@ def dim_emb(self): """Returns the embedding dimension g2.""" return self.get_dim_emb() - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index cde4534853..5c1a7d2785 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -2,6 +2,7 @@ import math from typing import ( Any, + NoReturn, Optional, Union, ) @@ -165,7 +166,7 @@ def get_env_protection(self) -> float: ) return all_protection[0] - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -187,7 +188,9 @@ def change_type_map( else None, ) - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: descrpt.compute_input_stats(merged, path) diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 9f2891d8c0..f45e85e516 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -5,6 +5,7 @@ ) from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -143,7 +144,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index 5658a87a9d..34d9d8f6bc 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -195,7 +196,7 @@ def __init__( g1_out_mlp: bool = True, ln_eps: Optional[float] = 1e-5, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.rcut = rcut self.rcut_smth = rcut_smth @@ -319,7 +320,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.g2_dim - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -366,18 +367,18 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.""" raise NotImplementedError - def get_stats(self): + def get_stats(self) -> NoReturn: """Get the statistics of the descriptor.""" raise NotImplementedError def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) @@ -796,7 +797,7 @@ def __init__( attnw_shift: float = 20.0, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() self.input_dim = input_dim @@ -919,7 +920,7 @@ def __init__( head_num: int, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num @@ -1007,7 +1008,7 @@ def __init__( head_num: int, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num @@ -1082,7 +1083,7 @@ def __init__( attnw_shift: float = 20.0, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim @@ -1244,7 +1245,7 @@ def __init__( g1_out_mlp: bool = True, ln_eps: Optional[float] = 1e-5, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei self.rcut = rcut diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 63402b6f84..a43b92082c 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -2,6 +2,7 @@ import itertools from typing import ( Any, + NoReturn, Optional, Union, ) @@ -217,7 +218,7 @@ def __init__( self.orig_sel = self.sel self.sel_cumsum = [0, *np.cumsum(self.sel).tolist()] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -258,7 +259,7 @@ def get_sel(self): """Returns cutoff radius.""" return self.sel - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -276,7 +277,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -304,7 +305,9 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -336,7 +339,7 @@ def cal_g( def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index d652eb1420..f1260bbab6 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, + NoReturn, Optional, Union, ) @@ -173,7 +174,7 @@ def __init__( self.orig_sel = self.sel self.sel_cumsum = [0, *np.cumsum(self.sel).tolist()] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -198,7 +199,7 @@ def get_dim_out(self): """Returns the output dimension of this descriptor.""" return self.neuron[-1] - def get_dim_emb(self): + def get_dim_emb(self) -> NoReturn: """Returns the embedding (g2) dimension of this descriptor.""" raise NotImplementedError @@ -214,7 +215,7 @@ def get_sel(self): """Returns cutoff radius.""" return self.sel - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -232,7 +233,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -260,7 +261,9 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index be587c77da..a0f60baebb 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools from typing import ( + NoReturn, Optional, Union, ) @@ -153,7 +154,7 @@ def __init__( ) self.orig_sel = self.sel - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -206,7 +207,7 @@ def get_sel(self): """Returns cutoff radius.""" return self.sel - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -224,7 +225,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -240,7 +241,9 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -260,7 +263,7 @@ def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index 298f823690..0079c2f6aa 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -227,7 +228,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_ttebd.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -243,7 +244,9 @@ def dim_out(self): def dim_emb(self): return self.get_dim_emb() - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -570,7 +573,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -621,18 +624,18 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.""" raise NotImplementedError - def get_stats(self): + def get_stats(self) -> NoReturn: """Get the statistics of the descriptor.""" raise NotImplementedError def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 5988951445..2a03934f3b 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -106,7 +106,7 @@ def __init__( c_differentiable: bool = True, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") if spin is not None: diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index cd70a35d06..b4b1ee3cb2 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -45,7 +45,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: if bias_dos is not None: self.bias_dos = bias_dos else: diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 1b6e1d28b2..53bedb4cec 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -45,7 +45,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__( var_name="energy", ntypes=ntypes, diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 344dab7ff1..f71f322dd1 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -119,7 +119,7 @@ def __init__( remove_vaccum_contribution: Optional[list[bool]] = None, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: self.var_name = var_name self.ntypes = ntypes self.dim_descrpt = dim_descrpt @@ -244,7 +244,7 @@ def change_type_map( ) self.bias_atom_e = self.bias_atom_e[remap_index] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["bias_atom_e"]: self.bias_atom_e = value elif key in ["fparam_avg"]: @@ -279,7 +279,7 @@ def __getitem__(self, key): def reinit_exclude( self, exclude_types: list[int] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index afae455441..d840d76149 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, + NoReturn, Optional, Union, ) @@ -133,7 +134,7 @@ def __init__( exclude_types: list[int] = [], type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") if spin is not None: @@ -186,7 +187,7 @@ def _net_out_dim(self): """Set the FittingNet output dim.""" return self.dim_out - def compute_output_stats(self, merged): + def compute_output_stats(self, merged) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index a67273356d..201b5e27d1 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -4,6 +4,7 @@ abstractmethod, ) from typing import ( + NoReturn, Optional, ) @@ -62,7 +63,7 @@ def fwd( """Calculate fitting.""" pass - def compute_output_stats(self, merged): + def compute_output_stats(self, merged) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 8cf3c4faaa..616ac20437 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -115,7 +115,7 @@ def __init__( shift_diag: bool = True, type_map: Optional[list[str]] = None, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") if spin is not None: @@ -171,7 +171,7 @@ def _net_out_dim(self): else self.embedding_width * self.embedding_width ) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["constant_matrix"]: self.constant_matrix = value else: diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 2cb49c3236..18a56e3bf9 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -85,7 +85,7 @@ def __init__( type_map: Optional[list[str]] = None, # not used seed: Optional[int] = None, - ): + ) -> None: self.task_dim = task_dim self.intensive = intensive self.bias_method = bias_method diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 5463743ada..ce176f5f45 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -80,7 +80,7 @@ def __init__( auto_batch_size: Union[bool, int, AutoBatchSize] = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, - ): + ) -> None: self.output_def = output_def self.model_path = model_file @@ -150,11 +150,11 @@ def get_numb_dos(self) -> int: """Get the number of DOS.""" return 0 - def get_has_efield(self): + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" return 0 diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index ba19785235..7bf22dfc6b 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -31,7 +31,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) @classmethod diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 5f21681830..643f260bff 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -22,6 +22,6 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPEnergyModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 95d97262df..a539adb292 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -159,7 +159,7 @@ def __init__( # underscore to prevent conflict with normal inputs atomic_model_: Optional[T_AtomicModel] = None, **kwargs, - ): + ) -> None: BaseModel.__init__(self) if atomic_model_ is not None: self.atomic_model: T_AtomicModel = atomic_model_ diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index dc64bb27fb..16fdedd36e 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -22,6 +22,6 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPPropertyModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index b0801fe59e..d149d427e0 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -29,7 +29,7 @@ def __init__( self, backbone_model, spin: Spin, - ): + ) -> None: super().__init__() self.backbone_model = backbone_model self.spin = spin diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py index bfee338d64..c2a1147786 100644 --- a/deepmd/dpmodel/output_def.py +++ b/deepmd/dpmodel/output_def.py @@ -8,7 +8,7 @@ def check_shape( shape: list[int], def_shape: list[int], -): +) -> None: """Check if the shape satisfies the defined shape.""" assert len(shape) == len(def_shape) if def_shape[-1] == -1: @@ -19,7 +19,7 @@ def check_shape( raise ValueError(f"{shape} shape not matching def {def_shape}") -def check_var(var, var_def): +def check_var(var, var_def) -> None: if var_def.atomic: # var.shape == [nf, nloc, *var_def.shape] if len(var.shape) != len(var_def.shape) + 2: @@ -47,7 +47,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) self.md = self.output_def() @@ -89,7 +89,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) self.md = self.output_def() @@ -197,7 +197,7 @@ def __init__( r_hessian: bool = False, magnetic: bool = False, intensive: bool = False, - ): + ) -> None: self.name = name self.shape = list(shape) # jit doesn't support math.prod(self.shape) @@ -230,7 +230,7 @@ def __init__( def size(self): return self.output_size - def squeeze(self, dim): + def squeeze(self, dim) -> None: # squeeze the shape on given dimension if -len(self.shape) <= dim < len(self.shape) and self.shape[dim] == 1: self.shape.pop(dim) @@ -252,7 +252,7 @@ class FittingOutputDef: def __init__( self, var_defs: list[OutputVariableDef], - ): + ) -> None: self.var_defs = {vv.name: vv for vv in var_defs} def __getitem__( @@ -286,7 +286,7 @@ class ModelOutputDef: def __init__( self, fit_defs: FittingOutputDef, - ): + ) -> None: self.def_outp = fit_defs self.def_redu = do_reduce(self.def_outp.get_data()) self.def_derv_r, self.def_derv_c = do_derivative(self.def_outp.get_data()) diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index aa8520202e..abbd68945b 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Optional, - Union, ) import array_api_compat @@ -83,7 +82,7 @@ def __init__( rcut, rcut_smth, protection: float = 0.0, - ): + ) -> None: self.rcut = rcut self.rcut_smth = rcut_smth self.protection = protection @@ -96,7 +95,7 @@ def call( davg: Optional[np.ndarray] = None, dstd: Optional[np.ndarray] = None, radial_only: bool = False, - ) -> Union[np.ndarray, np.ndarray]: + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Compute the environment matrix. Parameters diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py index b09a9b3e47..372eba133e 100644 --- a/deepmd/dpmodel/utils/exclude_mask.py +++ b/deepmd/dpmodel/utils/exclude_mask.py @@ -15,7 +15,7 @@ def __init__( self, ntypes: int, exclude_types: list[int] = [], - ): + ) -> None: self.ntypes = ntypes self.exclude_types = exclude_types type_mask = np.array( @@ -64,7 +64,7 @@ def __init__( self, ntypes: int, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.ntypes = ntypes self.exclude_types = set() for tt in exclude_types: @@ -135,5 +135,5 @@ def build_type_exclude_mask( ) return mask - def __contains__(self, item): + def __contains__(self, item) -> bool: return item in self.exclude_types diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index 5997b7d63a..90c18fca22 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -11,7 +11,7 @@ def __init__( stop_steps, decay_rate=None, **kwargs, - ): + ) -> None: """ Construct an exponential-decayed learning rate. diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 9ae5fd2b40..971ab00894 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -35,7 +35,7 @@ class Identity(NativeOP): - def __init__(self): + def __init__(self) -> None: super().__init__() def call(self, x: np.ndarray) -> np.ndarray: @@ -160,7 +160,7 @@ def deserialize(cls, data: dict) -> "NativeLayer": obj.check_shape_consistency() return obj - def check_shape_consistency(self): + def check_shape_consistency(self) -> None: if self.b is not None and self.w.shape[1] != self.b.shape[0]: raise ValueError( f"dim 1 of w {self.w.shape[1]} is not equal to shape " @@ -172,10 +172,10 @@ def check_shape_consistency(self): f"of idt {self.idt.shape[0]}", ) - def check_type_consistency(self): + def check_type_consistency(self) -> None: precision = self.precision - def check_var(var): + def check_var(var) -> None: if var is not None: # array api standard doesn't provide a API to get the dtype name # this is really hacked @@ -187,7 +187,7 @@ def check_var(var): check_var(self.b) check_var(self.idt) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("w", "matrix"): self.w = value elif key in ("b", "bias"): @@ -421,14 +421,14 @@ def deserialize(cls, data: dict) -> "LayerNorm": obj._check_shape_consistency() return obj - def _check_shape_consistency(self): + def _check_shape_consistency(self) -> None: if self.b is not None and self.w.shape[0] != self.b.shape[0]: raise ValueError( f"dim 1 of w {self.w.shape[0]} is not equal to shape " f"of b {self.b.shape[0]}", ) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("w", "matrix"): self.w = value elif key in ("b", "bias"): @@ -538,11 +538,11 @@ def __getitem__(self, key): assert isinstance(key, int) return self.layers[key] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: assert isinstance(key, int) self.layers[key] = value - def check_shape_consistency(self): + def check_shape_consistency(self) -> None: for ii in range(len(self.layers) - 1): if self.layers[ii].dim_out() != self.layers[ii + 1].dim_in(): raise ValueError( @@ -568,7 +568,7 @@ def call(self, x): x = layer(x) return x - def clear(self): + def clear(self) -> None: """Clear the network parameters to zero.""" for layer in self.layers: xp = array_api_compat.array_namespace(layer.w) @@ -616,7 +616,7 @@ def __init__( precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, bias: bool = True, - ): + ) -> None: layers = [] i_in = in_dim for idx, ii in enumerate(neuron): @@ -721,7 +721,7 @@ def __init__( precision: str = DEFAULT_PRECISION, bias_out: bool = True, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__( in_dim, neuron=neuron, @@ -822,7 +822,7 @@ def __init__( ntypes: int, network_type: str = "network", networks: list[Union[NativeNet, dict]] = [], - ): + ) -> None: self.ndim = ndim self.ntypes = ntypes self.network_type = self.NETWORK_TYPE_MAP[network_type] @@ -832,7 +832,7 @@ def __init__( if len(networks): self.check_completeness() - def check_completeness(self): + def check_completeness(self) -> None: """Check whether the collection is complete. Raises @@ -862,7 +862,7 @@ def _convert_key(self, key): def __getitem__(self, key): return self._networks[self._convert_key(key)] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if isinstance(value, self.network_type): pass elif isinstance(value, dict): diff --git a/deepmd/dpmodel/utils/serialization.py b/deepmd/dpmodel/utils/serialization.py index 37702cc9f0..1c07038a23 100644 --- a/deepmd/dpmodel/utils/serialization.py +++ b/deepmd/dpmodel/utils/serialization.py @@ -64,7 +64,7 @@ class Counter: 1 """ - def __init__(self): + def __init__(self) -> None: self.count = -1 def __call__(self): diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py index 8ef4490d1c..5679d838ac 100644 --- a/deepmd/entrypoints/doc.py +++ b/deepmd/entrypoints/doc.py @@ -10,7 +10,9 @@ __all__ = ["doc_train_input"] -def doc_train_input(*, out_type: str = "rst", multi_task: bool = False, **kwargs): +def doc_train_input( + *, out_type: str = "rst", multi_task: bool = False, **kwargs +) -> None: """Print out trining input arguments to console.""" if out_type == "rst": doc_str = gen_doc(make_anchor=True, multi_task=multi_task) diff --git a/deepmd/entrypoints/gui.py b/deepmd/entrypoints/gui.py index 8b6b9e0a09..7a61eb07ee 100644 --- a/deepmd/entrypoints/gui.py +++ b/deepmd/entrypoints/gui.py @@ -2,7 +2,7 @@ """DP-GUI entrypoint.""" -def start_dpgui(*, port: int, bind_all: bool, **kwargs): +def start_dpgui(*, port: int, bind_all: bool, **kwargs) -> None: """Host DP-GUI server. Parameters diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 05f660cb9a..2c91ca5f29 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -38,7 +38,7 @@ ) -def main(args: argparse.Namespace): +def main(args: argparse.Namespace) -> None: """DeePMD-Kit entry point. Parameters diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index 4cad5f312c..6ea92a6b7b 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -13,7 +13,7 @@ def show( INPUT: str, ATTRIBUTES: list[str], **kwargs, -): +) -> None: model = DeepEval(INPUT, head=0) model_params = model.get_model_def_script() model_is_multi_task = "model_dict" in model_params diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index fd0393c914..0f2c116f80 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -34,9 +34,6 @@ from deepmd.infer.deep_property import ( DeepProperty, ) -from deepmd.infer.deep_wfc import ( - DeepWFC, -) from deepmd.utils import random as dp_random from deepmd.utils.data import ( DeepmdData, @@ -49,6 +46,9 @@ from deepmd.infer.deep_tensor import ( DeepTensor, ) + from deepmd.infer.deep_wfc import ( + DeepWFC, + ) __all__ = ["test"] @@ -67,7 +67,7 @@ def test( atomic: bool, head: Optional[str] = None, **kwargs, -): +) -> None: """Test model predictions. Parameters @@ -236,7 +236,7 @@ def rmse(diff: np.ndarray) -> float: def save_txt_file( fname: Path, data: np.ndarray, header: str = "", append: bool = False -): +) -> None: """Save numpy array to test file. Parameters @@ -561,7 +561,7 @@ def test_ener( } -def print_ener_sys_avg(avg: dict[str, float]): +def print_ener_sys_avg(avg: dict[str, float]) -> None: """Print errors summary for energy type potential. Parameters @@ -733,7 +733,7 @@ def test_dos( } -def print_dos_sys_avg(avg: dict[str, float]): +def print_dos_sys_avg(avg: dict[str, float]) -> None: """Print errors summary for DOS type potential. Parameters @@ -887,7 +887,7 @@ def test_property( } -def print_property_sys_avg(avg: dict[str, float]): +def print_property_sys_avg(avg: dict[str, float]) -> None: """Print errors summary for Property type potential. Parameters @@ -983,7 +983,7 @@ def test_wfc( return {"rmse": (rmse_f, wfc.size)} -def print_wfc_sys_avg(avg): +def print_wfc_sys_avg(avg) -> None: """Print errors summary for wfc type potential. Parameters @@ -1125,7 +1125,7 @@ def test_polar( return {"rmse": (rmse_f, polar.size)} -def print_polar_sys_avg(avg): +def print_polar_sys_avg(avg) -> None: """Print errors summary for polar type potential. Parameters @@ -1239,7 +1239,7 @@ def test_dipole( return {"rmse": (rmse_f, dipole.size)} -def print_dipole_sys_avg(avg): +def print_dipole_sys_avg(avg) -> None: """Print errors summary for dipole type potential. Parameters diff --git a/deepmd/env.py b/deepmd/env.py index 50e52fd719..c1dd9d6f25 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -48,7 +48,7 @@ ) -def set_env_if_empty(key: str, value: str, verbose: bool = True): +def set_env_if_empty(key: str, value: str, verbose: bool = True) -> None: """Set environment variable only if it is empty. Parameters @@ -68,7 +68,7 @@ def set_env_if_empty(key: str, value: str, verbose: bool = True): ) -def set_default_nthreads(): +def set_default_nthreads() -> None: """Set internal number of threads to default=automatic selection. Notes diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index e08dc88674..6e2b2942ed 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -268,11 +268,11 @@ def get_numb_dos(self) -> int: """Get the number of DOS.""" raise NotImplementedError - def get_has_efield(self): + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False - def get_has_spin(self): + def get_has_spin(self) -> bool: """Check if the model has spin atom types.""" return False diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 4a3283cf32..389a0e8512 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -53,7 +53,7 @@ def output_def(self) -> ModelOutputDef: ) ) - def change_output_def(self): + def change_output_def(self) -> None: self.output_def["property"].shape = self.task_dim self.output_def["property"].intensive = self.get_intensive() diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index b9d1974c27..a83964329e 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -86,7 +86,7 @@ def __init__( auto_batch_size: Union[bool, int, AutoBatchSize] = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, - ): + ) -> None: self.output_def = output_def self.model_path = model_file @@ -178,11 +178,11 @@ def get_numb_dos(self) -> int: """Get the number of DOS.""" return 0 - def get_has_efield(self): + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model.""" return 0 diff --git a/deepmd/jax/utils/auto_batch_size.py b/deepmd/jax/utils/auto_batch_size.py index eec6766ae2..ef93d85d96 100644 --- a/deepmd/jax/utils/auto_batch_size.py +++ b/deepmd/jax/utils/auto_batch_size.py @@ -25,7 +25,7 @@ def __init__( self, initial_batch_size: int = 1024, factor: float = 2.0, - ): + ) -> None: super().__init__( initial_batch_size=initial_batch_size, factor=factor, diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 1ed26f2d40..5d4da49e08 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -167,7 +167,7 @@ def serialize_from_file(model_file: str) -> dict: state = data.state # convert str "1" to int 1 key - def convert_str_to_int_key(item: dict): + def convert_str_to_int_key(item: dict) -> None: for key, value in item.copy().items(): if isinstance(value, dict): convert_str_to_int_key(value) diff --git a/deepmd/loggers/loggers.py b/deepmd/loggers/loggers.py index f42c032cfa..060ac1a4b0 100644 --- a/deepmd/loggers/loggers.py +++ b/deepmd/loggers/loggers.py @@ -5,6 +5,7 @@ import os from typing import ( TYPE_CHECKING, + NoReturn, Optional, ) @@ -43,7 +44,7 @@ class _AppFilter(logging.Filter): """Add field `app_name` to log messages.""" - def filter(self, record): + def filter(self, record) -> bool: record.app_name = "DEEPMD" return True @@ -55,7 +56,7 @@ def __init__(self, rank: int) -> None: super().__init__(name="MPI_rank_id") self.mpi_rank = str(rank) - def filter(self, record): + def filter(self, record) -> bool: record.rank = self.mpi_rank return True @@ -67,7 +68,7 @@ def __init__(self, rank: int) -> None: super().__init__(name="MPI_master_log") self.mpi_rank = rank - def filter(self, record): + def filter(self, record) -> bool: if self.mpi_rank == 0: return True else: @@ -94,7 +95,7 @@ def __init__( self.stream.Set_atomicity(True) self.name = "MPIfilestream" - def write(self, msg: str): + def write(self, msg: str) -> None: """Write to MPI shared file stream. Parameters @@ -106,7 +107,7 @@ def write(self, msg: str): b.extend(map(ord, msg)) self.stream.Write_shared(b) - def close(self): + def close(self) -> None: """Synchronize and close MPI file stream.""" self.stream.Sync() self.stream.Close() @@ -137,14 +138,14 @@ def __init__( def _open(self): return _MPIFileStream(self.baseFilename, self.MPI, self.mode) - def setStream(self, stream): + def setStream(self, stream) -> NoReturn: """Stream cannot be reasigned in MPI mode.""" raise NotImplementedError("Unable to do for MPI file handler!") def set_log_handles( level: int, log_path: Optional["Path"] = None, mpi_log: Optional[str] = None -): +) -> None: """Set desired level for package loggers and add file handlers. Parameters diff --git a/deepmd/loggers/training.py b/deepmd/loggers/training.py index b2fff4788b..fa643d2257 100644 --- a/deepmd/loggers/training.py +++ b/deepmd/loggers/training.py @@ -7,7 +7,7 @@ def format_training_message( batch: int, wall_time: float, -): +) -> str: """Format a training message.""" return f"batch {batch:7d}: " f"total wall time = {wall_time:.2f} s" @@ -17,7 +17,7 @@ def format_training_message_per_task( task_name: str, rmse: dict[str, float], learning_rate: Optional[float], -): +) -> str: if task_name: task_name += ": " if learning_rate is None: diff --git a/deepmd/main.py b/deepmd/main.py index d2a52568b4..b3daf75963 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -69,7 +69,7 @@ def __call__(self, parser, namespace, values, option_string=None): class DeprecateAction(argparse.Action): # See https://stackoverflow.com/a/69052677/9567349 by Ibolit under CC BY-SA 4.0 - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: self.call_count = 0 if "help" in kwargs: kwargs["help"] = f'[DEPRECATED] {kwargs["help"]}' @@ -878,7 +878,7 @@ def parse_args(args: Optional[list[str]] = None) -> argparse.Namespace: return parsed_args -def main(args: Optional[list[str]] = None): +def main(args: Optional[list[str]] = None) -> None: """DeePMD-kit new entry point. Parameters diff --git a/deepmd/pt/entrypoints/compress.py b/deepmd/pt/entrypoints/compress.py index d94a34215c..3a4efaeea5 100644 --- a/deepmd/pt/entrypoints/compress.py +++ b/deepmd/pt/entrypoints/compress.py @@ -36,7 +36,7 @@ def enable_compression( extrapolate: int = 5, check_frequency: int = -1, training_script: Optional[str] = None, -): +) -> None: saved_model = torch.jit.load(input_file, map_location="cpu") model_def_script = json.loads(saved_model.model_def_script) model = get_model(model_def_script) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index fe85a3301c..1e5314a821 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -252,7 +252,7 @@ def train( use_pretrain_script: bool = False, force_load: bool = False, output: str = "out.json", -): +) -> None: log.info("Configuration path: %s", input_file) SummaryPrinter()() with open(input_file) as fin: @@ -364,7 +364,7 @@ def freeze( model: str, output: str = "frozen_model.pth", head: Optional[str] = None, -): +) -> None: model = inference.Tester(model, head=head).model model.eval() model = torch.jit.script(model) @@ -386,7 +386,7 @@ def change_bias( numb_batch: int = 0, model_branch: Optional[str] = None, output: Optional[str] = None, -): +) -> None: if input_file.endswith(".pt"): old_state_dict = torch.load( input_file, map_location=env.DEVICE, weights_only=True @@ -509,7 +509,7 @@ def change_bias( @record -def main(args: Optional[Union[list[str], argparse.Namespace]] = None): +def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: if not isinstance(args, argparse.Namespace): FLAGS = parse_args(args=args) else: diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 934cafdb47..89bb8a0286 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -99,7 +99,7 @@ def __init__( neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, head: Optional[Union[str, int]] = None, **kwargs: Any, - ): + ) -> None: self.output_def = output_def self.model_path = model_file if str(self.model_path).endswith(".pt"): @@ -221,11 +221,11 @@ def get_task_dim(self) -> int: """Get the output dimension.""" return self.dp.model["Default"].get_task_dim() - def get_has_efield(self): + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model. Only used in old implement.""" return 0 diff --git a/deepmd/pt/infer/inference.py b/deepmd/pt/infer/inference.py index b3d120cbc4..0e3bc31057 100644 --- a/deepmd/pt/infer/inference.py +++ b/deepmd/pt/infer/inference.py @@ -27,7 +27,7 @@ def __init__( self, model_ckpt, head=None, - ): + ) -> None: """Construct a DeePMD tester. Args: diff --git a/deepmd/pt/loss/denoise.py b/deepmd/pt/loss/denoise.py index 57691558cb..574210adb6 100644 --- a/deepmd/pt/loss/denoise.py +++ b/deepmd/pt/loss/denoise.py @@ -22,7 +22,7 @@ def __init__( mask_loss_coord=True, mask_loss_token=True, **kwargs, - ): + ) -> None: """Construct a layer to compute loss on coord, and type reconstruction.""" super().__init__() self.ntypes = ntypes diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 84513b6bf9..7a64c6fbd3 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -28,7 +28,7 @@ def __init__( limit_pref_acdf: float = 0.0, inference=False, **kwargs, - ): + ) -> None: r"""Construct a loss for local and global tensors. Parameters diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index f40110a749..327d75c2cd 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -42,7 +42,7 @@ def __init__( use_l1_all: bool = False, inference=False, **kwargs, - ): + ) -> None: r"""Construct a layer to compute loss on energy, force and virial. Parameters diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 09a053451f..6a926f4051 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -35,7 +35,7 @@ def __init__( use_l1_all: bool = False, inference=False, **kwargs, - ): + ) -> None: r"""Construct a layer to compute loss on energy, real force, magnetic force and virial. Parameters diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index 5447c8735b..dfe62f4da0 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -3,6 +3,9 @@ ABC, abstractmethod, ) +from typing import ( + NoReturn, +) import torch @@ -15,11 +18,11 @@ class TaskLoss(torch.nn.Module, ABC, make_plugin_registry("loss")): - def __init__(self, **kwargs): + def __init__(self, **kwargs) -> None: """Construct loss.""" super().__init__() - def forward(self, input_dict, model, label, natoms, learning_rate): + def forward(self, input_dict, model, label, natoms, learning_rate) -> NoReturn: """Return loss .""" raise NotImplementedError diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index ba120e3d6c..07e394650a 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -25,7 +25,7 @@ def __init__( metric: list = ["mae"], beta: float = 1.00, **kwargs, - ): + ) -> None: r"""Construct a layer to compute loss on property. Parameters diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 32d25cc9f1..8f2f937a07 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -23,7 +23,7 @@ def __init__( pref: float = 0.0, inference=False, **kwargs, - ): + ) -> None: r"""Construct a loss for local and global tensors. Parameters diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 4730de53fa..a64eca0fe9 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -3,6 +3,7 @@ import logging from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -78,7 +79,7 @@ def __init__( pair_exclude_types: list[tuple[int, int]] = [], rcond: Optional[float] = None, preset_out_bias: Optional[dict[str, np.ndarray]] = None, - ): + ) -> None: torch.nn.Module.__init__(self) BaseAtomicModel_.__init__(self) self.type_map = type_map @@ -87,7 +88,7 @@ def __init__( self.rcond = rcond self.preset_out_bias = preset_out_bias - def init_out_stat(self): + def init_out_stat(self) -> None: """Initialize the output bias.""" ntypes = self.get_ntypes() self.bias_keys: list[str] = list(self.fitting_output_def().keys()) @@ -103,7 +104,7 @@ def init_out_stat(self): def set_out_bias(self, out_bias: torch.Tensor) -> None: self.out_bias = out_bias - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["out_bias"]: self.out_bias = value elif key in ["out_std"]: @@ -127,7 +128,7 @@ def get_type_map(self) -> list[str]: def reinit_atom_exclude( self, exclude_types: list[int] = [], - ): + ) -> None: self.atom_exclude_types = exclude_types if exclude_types == []: self.atom_excl = None @@ -137,7 +138,7 @@ def reinit_atom_exclude( def reinit_pair_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.pair_exclude_types = exclude_types if exclude_types == []: self.pair_excl = None @@ -352,7 +353,7 @@ def compute_or_load_stat( self, merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[DPPath] = None, - ): + ) -> NoReturn: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -375,7 +376,7 @@ def compute_or_load_out_stat( self, merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -542,7 +543,7 @@ def _store_out_stat( out_bias: dict[str, torch.Tensor], out_std: dict[str, torch.Tensor], add: bool = False, - ): + ) -> None: ntypes = self.get_ntypes() out_bias_data = torch.clone(self.out_bias) out_std_data = torch.clone(self.out_std) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index aa28294cc5..fd0879e707 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -12,7 +12,7 @@ class DPDipoleAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert isinstance(fitting, DipoleFittingNet) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index 5e399f2aff..1f7d7a9917 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -9,6 +9,6 @@ class DPDOSAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert isinstance(fitting, DOSFittingNet) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 1efac7fd1f..2cdc97f934 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -51,7 +51,7 @@ def __init__( fitting, type_map: list[str], **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) ntypes = len(type_map) self.type_map = type_map @@ -256,7 +256,7 @@ def compute_or_load_stat( self, sampled_func, stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index 7cedaa1ab3..855c1213ec 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -11,7 +11,7 @@ class DPEnergyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert ( isinstance(fitting, EnergyFittingNet) or isinstance(fitting, EnergyFittingNetDirect) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 318079ef68..80c864ad53 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -57,7 +57,7 @@ def __init__( type_map: list[str], weights: Optional[Union[str, list[float]]] = "mean", **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -459,7 +459,7 @@ def compute_or_load_out_stat( self, merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -483,7 +483,7 @@ def compute_or_load_stat( self, sampled_func, stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. @@ -533,7 +533,7 @@ def __init__( type_map: list[str], smin_alpha: Optional[float] = 0.1, **kwargs, - ): + ) -> None: models = [dp_model, zbl_model] kwargs["models"] = models kwargs["type_map"] = type_map diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 8cf76351b7..c535808d83 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -69,7 +69,7 @@ def __init__( sel: Union[int, list[int]], type_map: list[str], **kwargs, - ): + ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() self.tab_file = tab_file @@ -217,7 +217,7 @@ def compute_or_load_stat( self, merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 39cda2650d..8ec27b5762 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -12,7 +12,7 @@ class DPPolarAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert isinstance(fitting, PolarFittingNet) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index 2fac90100f..6bacaf5d72 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -12,7 +12,7 @@ class DPPropertyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs) -> None: assert isinstance(fitting, PropertyFittingNet) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 5d36606760..3c738d59e3 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -6,6 +6,7 @@ ) from typing import ( Callable, + NoReturn, Optional, Union, ) @@ -102,7 +103,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -125,7 +126,7 @@ def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -194,7 +195,7 @@ def make_default_type_embedding( return TypeEmbedNet(ntypes, aux["tebd_dim"]), aux -def extend_descrpt_stat(des, type_map, des_with_stat=None): +def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 76115b2810..ce2a8a220e 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -251,7 +251,7 @@ def __init__( # not implemented spin=None, type: Optional[str] = None, - ): + ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. if stripped_type_embedding is not None: @@ -376,7 +376,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_atten.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 77e9f1d936..0bf97c5f23 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -96,7 +96,7 @@ def __init__( use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: Optional[list[str]] = None, - ): + ) -> None: r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492. Parameters @@ -374,7 +374,7 @@ def get_env_protection(self) -> float: # the env_protection of repinit is the same as that of the repformer return self.repinit.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -490,7 +490,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index 00984df238..5caa1243f1 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -164,7 +164,7 @@ def get_env_protection(self) -> float: ) return all_protection[0] - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -195,7 +195,9 @@ def change_type_map( else None, ) - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: Optional[DPPath] = None + ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: descrpt.compute_input_stats(merged, path) diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py index b4fac5fce7..86b09e9b40 100644 --- a/deepmd/pt/model/descriptor/repformer_layer.py +++ b/deepmd/pt/model/descriptor/repformer_layer.py @@ -160,7 +160,7 @@ def __init__( attnw_shift: float = 20.0, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() self.input_dim = input_dim @@ -285,7 +285,7 @@ def __init__( head_num: int, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num @@ -370,7 +370,7 @@ def __init__( head_num: int, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num @@ -443,7 +443,7 @@ def __init__( attnw_shift: float = 20.0, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim @@ -602,7 +602,7 @@ def __init__( g1_out_conv: bool = True, g1_out_mlp: bool = True, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei self.rcut = float(rcut) diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index bd1109e8b7..62f693edef 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -108,7 +108,7 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ): + ) -> None: r""" The repformer descriptor block. @@ -328,7 +328,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.g2_dim - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -378,7 +378,7 @@ def dim_emb(self): def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) @@ -532,7 +532,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 9b5ee6d2c4..6c0bf57447 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -111,7 +111,7 @@ def __init__( type_map: Optional[list[str]] = None, # not implemented spin=None, - ): + ) -> None: del ntypes if spin is not None: raise NotImplementedError("old implementation of spin is not supported.") @@ -185,7 +185,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.sea.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -246,7 +246,7 @@ def compute_input_stats( def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: """Update the type exclusions.""" self.sea.reinit_exclude(exclude_types) @@ -456,7 +456,7 @@ def __init__( trainable: bool = True, seed: Optional[Union[int, list[int]]] = None, **kwargs, - ): + ) -> None: """Construct an embedding net of type `se_a`. Args: @@ -592,11 +592,11 @@ def dim_out(self): return self.filter_neuron[-1] * self.axis_neuron @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return 0 - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -616,7 +616,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -666,7 +666,7 @@ def get_stats(self) -> dict[str, StatItem]: def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index a4e3d44bf0..ca86423a7f 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -100,7 +100,7 @@ def __init__( ln_eps: Optional[float] = 1e-5, seed: Optional[Union[int, list[int]]] = None, type: Optional[str] = None, - ): + ) -> None: r"""Construct an embedding net of type `se_atten`. Parameters @@ -315,7 +315,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -366,7 +366,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -416,7 +416,7 @@ def get_stats(self) -> dict[str, StatItem]: def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.is_sorted = len(self.exclude_types) == 0 self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) @@ -659,7 +659,7 @@ def __init__( smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a neighbor-wise attention net.""" super().__init__() self.layer_num = layer_num @@ -730,7 +730,7 @@ def __getitem__(self, key): else: raise TypeError(key) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if not isinstance(key, int): raise TypeError(key) if isinstance(value, self.network_type): @@ -802,7 +802,7 @@ def __init__( ln_eps: float = 1e-5, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a neighbor-wise attention layer.""" super().__init__() self.nnei = nnei @@ -909,7 +909,7 @@ def __init__( smooth: bool = True, precision: str = DEFAULT_PRECISION, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Construct a multi-head neighbor-wise attention net.""" super().__init__() assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads" diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index beb8acd5d5..190f4c1007 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -95,7 +95,7 @@ def __init__( seed: Optional[Union[int, list[int]]] = None, type_map: Optional[list[str]] = None, **kwargs, - ): + ) -> None: super().__init__() self.rcut = float(rcut) self.rcut_smth = float(rcut_smth) @@ -225,7 +225,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -282,7 +282,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -329,7 +329,7 @@ def get_stats(self) -> dict[str, StatItem]: ) return self.stats - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -348,7 +348,7 @@ def __getitem__(self, key): def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index ae9b3a9c1a..b1c0922b47 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -147,7 +147,7 @@ def __init__( ntypes: Optional[int] = None, # to be compat with input # not implemented spin=None, - ): + ) -> None: del ntypes if spin is not None: raise NotImplementedError("old implementation of spin is not supported.") @@ -219,7 +219,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.seat.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -328,7 +328,7 @@ def enable_compression( def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: """Update the type exclusions.""" self.seat.reinit_exclude(exclude_types) @@ -484,7 +484,7 @@ def __init__( precision: str = "float64", trainable: bool = True, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: r"""Construct an embedding net of type `se_e3`. The embedding takes angles between two neighboring atoms as input. @@ -638,11 +638,11 @@ def dim_out(self): return self.filter_neuron[-1] @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return 0 - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -662,7 +662,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -712,7 +712,7 @@ def get_stats(self) -> dict[str, StatItem]: def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 82ccb06f32..810e4b7705 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -142,7 +142,7 @@ def __init__( use_econf_tebd: bool = False, use_tebd_bias=False, smooth: bool = True, - ): + ) -> None: super().__init__() self.se_ttebd = DescrptBlockSeTTebd( rcut, @@ -238,7 +238,7 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_ttebd.get_env_protection() - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -510,7 +510,7 @@ def __init__( env_protection: float = 0.0, smooth: bool = True, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.rcut = float(rcut) self.rcut_smth = float(rcut_smth) @@ -615,7 +615,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -666,7 +666,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -716,7 +716,7 @@ def get_stats(self) -> dict[str, StatItem]: def reinit_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 03896f3e59..580b83cb3d 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -13,6 +13,9 @@ import copy import json +from typing import ( + Optional, +) import numpy as np @@ -200,7 +203,7 @@ def get_zbl_model(model_params): ) -def _can_be_converted_to_float(value): +def _can_be_converted_to_float(value) -> Optional[bool]: try: float(value) return True diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 537a0fd6a4..79490bc20c 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -30,7 +30,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 1be99624da..41d85a559e 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -30,7 +30,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index 9d11bca622..fe4487a495 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -33,7 +33,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) def translated_output_def(self): diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index eacc3cbb3b..3c019b9376 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -33,7 +33,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) def translated_output_def(self): diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 8f7ee2094a..9487bcc5bb 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -30,7 +30,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPEnergyModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index 37149303d4..27284ec276 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -2,6 +2,7 @@ import json import tempfile from typing import ( + NoReturn, Optional, ) @@ -31,7 +32,7 @@ class FrozenModel(BaseModel): The path to the frozen model """ - def __init__(self, model_file: str, **kwargs): + def __init__(self, model_file: str, **kwargs) -> None: super().__init__(**kwargs) self.model_file = model_file if model_file.endswith(".pth"): @@ -157,7 +158,7 @@ def serialize(self) -> dict: return model.serialize() @classmethod - def deserialize(cls, data: dict): + def deserialize(cls, data: dict) -> NoReturn: raise RuntimeError("Should not touch here.") @torch.jit.export diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py index d2541a815e..4104314225 100644 --- a/deepmd/pt/model/model/make_hessian_model.py +++ b/deepmd/pt/model/model/make_hessian_model.py @@ -36,7 +36,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: super().__init__( *args, **kwargs, @@ -46,7 +46,7 @@ def __init__( def requires_hessian( self, keys: Union[str, list[str]], - ): + ) -> None: """Set which output variable(s) requires hessian.""" if isinstance(keys, str): keys = [keys] @@ -189,7 +189,7 @@ def __init__( box: Optional[torch.Tensor], fparam: Optional[torch.Tensor], aparam: Optional[torch.Tensor], - ): + ) -> None: self.atype, self.box, self.fparam, self.aparam = atype, box, fparam, aparam self.ci = ci self.obj = obj diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 709fce8150..6bb5f6b8e9 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -69,7 +69,7 @@ def __init__( # underscore to prevent conflict with normal inputs atomic_model_: Optional[T_AtomicModel] = None, **kwargs, - ): + ) -> None: super().__init__(*args, **kwargs) if atomic_model_ is not None: self.atomic_model: T_AtomicModel = atomic_model_ diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index c7b5986f0c..44600ab804 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + NoReturn, Optional, ) @@ -17,7 +18,7 @@ class BaseModel(torch.nn.Module, make_base_model()): - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: """Construct a basic model for different tasks.""" torch.nn.Module.__init__(self) self.model_def_script = "" @@ -29,7 +30,7 @@ def compute_or_load_stat( self, sampled_func, stat_file_path: Optional[DPPath] = None, - ): + ) -> NoReturn: """ Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index c37f0a32fc..ea6316dc91 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -30,7 +30,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 2ee13c3aa3..4581a2bc3e 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -30,7 +30,7 @@ def __init__( self, *args, **kwargs, - ): + ) -> None: DPModelCommon.__init__(self) DPPropertyModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 455c240493..ac94668039 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -40,7 +40,7 @@ def __init__( self, backbone_model, spin: Spin, - ): + ) -> None: super().__init__() self.backbone_model = backbone_model self.spin = spin @@ -345,7 +345,7 @@ def compute_or_load_stat( self, sampled_func, stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. @@ -526,7 +526,7 @@ def __init__( self, backbone_model, spin: Spin, - ): + ) -> None: super().__init__(backbone_model, spin) def translated_output_def(self): diff --git a/deepmd/pt/model/network/layernorm.py b/deepmd/pt/model/network/layernorm.py index 76ce90b627..89bd16d569 100644 --- a/deepmd/pt/model/network/layernorm.py +++ b/deepmd/pt/model/network/layernorm.py @@ -45,7 +45,7 @@ def __init__( precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() self.eps = eps self.uni_init = uni_init diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index f2137bd004..43989a3cab 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -47,7 +47,7 @@ def empty_t(shape, precision): class Identity(nn.Module): - def __init__(self): + def __init__(self) -> None: super().__init__() def forward( @@ -82,7 +82,7 @@ def __init__( precision: str = DEFAULT_PRECISION, init: str = "default", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: super().__init__() # only use_timestep when skip connection is established. self.use_timestep = use_timestep and ( @@ -126,10 +126,10 @@ def __init__( else: raise ValueError(f"Unknown initialization method: {init}") - def check_type_consistency(self): + def check_type_consistency(self) -> None: precision = self.precision - def check_var(var): + def check_var(var) -> None: if var is not None: # assertion "float64" == "double" would fail assert PRECISION_DICT[var.dtype.name] is PRECISION_DICT[precision] @@ -149,7 +149,7 @@ def _default_normal_init( bavg: float = 0.0, stddev: float = 1.0, generator: Optional[torch.Generator] = None, - ): + ) -> None: normal_( self.matrix.data, std=stddev / np.sqrt(self.num_out + self.num_in), @@ -162,7 +162,7 @@ def _default_normal_init( def _trunc_normal_init( self, scale=1.0, generator: Optional[torch.Generator] = None - ): + ) -> None: # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.) TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 _, fan_in = self.matrix.shape @@ -170,17 +170,17 @@ def _trunc_normal_init( std = (scale**0.5) / TRUNCATED_NORMAL_STDDEV_FACTOR trunc_normal_(self.matrix, mean=0.0, std=std, generator=generator) - def _glorot_uniform_init(self, generator: Optional[torch.Generator] = None): + def _glorot_uniform_init(self, generator: Optional[torch.Generator] = None) -> None: xavier_uniform_(self.matrix, gain=1, generator=generator) - def _zero_init(self, use_bias=True): + def _zero_init(self, use_bias=True) -> None: with torch.no_grad(): self.matrix.fill_(0.0) if use_bias and self.bias is not None: with torch.no_grad(): self.bias.fill_(1.0) - def _normal_init(self, generator: Optional[torch.Generator] = None): + def _normal_init(self, generator: Optional[torch.Generator] = None) -> None: kaiming_normal_(self.matrix, nonlinearity="linear", generator=generator) def forward( @@ -280,7 +280,7 @@ def check_load_param(ss): class MLP(MLP_): - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.layers = torch.nn.ModuleList(self.layers) @@ -301,7 +301,7 @@ class NetworkCollection(DPNetworkCollection, nn.Module): "fitting_network": FittingNet, } - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs) -> None: # init both two base classes DPNetworkCollection.__init__(self, *args, **kwargs) nn.Module.__init__(self) diff --git a/deepmd/pt/model/network/network.py b/deepmd/pt/model/network/network.py index 88ea108ce7..1998cc0dce 100644 --- a/deepmd/pt/model/network/network.py +++ b/deepmd/pt/model/network/network.py @@ -56,7 +56,7 @@ def __init__( use_timestep=False, activate=None, bias: bool = True, - ): + ) -> None: """Construct a linear layer. Args: @@ -99,7 +99,7 @@ def __init__( d_out: int, bias: bool = True, init: str = "default", - ): + ) -> None: super().__init__( d_in, d_out, @@ -129,7 +129,7 @@ def __init__( else: raise ValueError("Invalid init method.") - def _trunc_normal_init(self, scale=1.0): + def _trunc_normal_init(self, scale=1.0) -> None: # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.) TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 _, fan_in = self.weight.shape @@ -137,22 +137,22 @@ def _trunc_normal_init(self, scale=1.0): std = (scale**0.5) / TRUNCATED_NORMAL_STDDEV_FACTOR nn.init.trunc_normal_(self.weight, mean=0.0, std=std) - def _glorot_uniform_init(self): + def _glorot_uniform_init(self) -> None: nn.init.xavier_uniform_(self.weight, gain=1) - def _zero_init(self, use_bias=True): + def _zero_init(self, use_bias=True) -> None: with torch.no_grad(): self.weight.fill_(0.0) if use_bias: with torch.no_grad(): self.bias.fill_(1.0) - def _normal_init(self): + def _normal_init(self) -> None: nn.init.kaiming_normal_(self.weight, nonlinearity="linear") class NonLinearHead(nn.Module): - def __init__(self, input_dim, out_dim, activation_fn, hidden=None): + def __init__(self, input_dim, out_dim, activation_fn, hidden=None) -> None: super().__init__() hidden = input_dim if not hidden else hidden self.linear1 = SimpleLinear(input_dim, hidden, activate=activation_fn) @@ -167,7 +167,7 @@ def forward(self, x): class MaskLMHead(nn.Module): """Head for masked language modeling.""" - def __init__(self, embed_dim, output_dim, activation_fn, weight=None): + def __init__(self, embed_dim, output_dim, activation_fn, weight=None) -> None: super().__init__() self.dense = SimpleLinear(embed_dim, embed_dim) self.activation_fn = ActivationFn(activation_fn) @@ -199,7 +199,7 @@ def forward(self, features, masked_tokens: Optional[torch.Tensor] = None, **kwar class ResidualDeep(nn.Module): def __init__( self, type_id, embedding_width, neuron, bias_atom_e, out_dim=1, resnet_dt=False - ): + ) -> None: """Construct a filter on the given element as neighbor. Args: @@ -261,7 +261,7 @@ def __init__( use_econf_tebd=False, use_tebd_bias: bool = False, type_map=None, - ): + ) -> None: """Construct a type embedding net.""" super().__init__() self.type_nums = type_nums @@ -296,7 +296,7 @@ def forward(self, atype): """ return self.embedding(atype.device)[atype] - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -364,7 +364,7 @@ def __init__( use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: Optional[list[str]] = None, - ): + ) -> None: """Construct a type embedding net.""" super().__init__() self.ntypes = ntypes diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index df65f1cd18..fc9e8943e9 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -32,7 +32,7 @@ def __init__( prefactor=[0.5, 0.5], activation_function="gelu", **kwargs, - ): + ) -> None: """Construct a denoise net. Args: diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index bc09fa4d0f..bf445b39a1 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -91,7 +91,7 @@ def __init__( c_differentiable: bool = True, type_map: Optional[list[str]] = None, **kwargs, - ): + ) -> None: self.embedding_width = embedding_width self.r_differentiable = r_differentiable self.c_differentiable = c_differentiable @@ -149,7 +149,7 @@ def compute_output_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[DPPath] = None, - ): + ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index 531ad0a9f3..a71117e587 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -56,7 +56,7 @@ def __init__( exclude_types: list[int] = [], mixed_types: bool = True, type_map: Optional[list[str]] = None, - ): + ) -> None: if bias_dos is not None: self.bias_dos = bias_dos else: diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 0d80006b94..543d987e31 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -56,7 +56,7 @@ def __init__( seed: Optional[Union[int, list[int]]] = None, type_map: Optional[list[str]] = None, **kwargs, - ): + ) -> None: super().__init__( "energy", ntypes, @@ -109,7 +109,7 @@ def __init__( use_tebd=True, return_energy=False, **kwargs, - ): + ) -> None: """Construct a fitting net for energy. Args: diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 470b420c89..a9bb8119a0 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -54,7 +54,7 @@ def __new__(cls, *args, **kwargs): return BaseFitting.__new__(BaseFitting, *args, **kwargs) return super().__new__(cls) - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -150,7 +150,7 @@ def __init__( type_map: Optional[list[str]] = None, use_aparam_as_mask: bool = False, **kwargs, - ): + ) -> None: super().__init__() self.var_name = var_name self.ntypes = ntypes @@ -242,7 +242,7 @@ def __init__( def reinit_exclude( self, exclude_types: list[int] = [], - ): + ) -> None: self.exclude_types = exclude_types self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) @@ -349,7 +349,7 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["bias_atom_e"]: value = value.view([self.ntypes, self._net_out_dim()]) self.bias_atom_e = value diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index acdd5b0fda..c3fc1b014a 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -101,7 +101,7 @@ def __init__( type_map: Optional[list[str]] = None, use_aparam_as_mask: bool = False, **kwargs, - ): + ) -> None: self.dim_out = dim_out self.atom_ener = atom_ener super().__init__( diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 7b7f92c3af..d9e9e0d43b 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -96,7 +96,7 @@ def __init__( shift_diag: bool = True, type_map: Optional[list[str]] = None, **kwargs, - ): + ) -> None: self.embedding_width = embedding_width self.fit_diag = fit_diag self.scale = scale @@ -146,7 +146,7 @@ def _net_out_dim(self): else self.embedding_width * self.embedding_width ) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: if key in ["constant_matrix"]: self.constant_matrix = value else: diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 0108b83dff..1c2b9e7c9c 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -88,7 +88,7 @@ def __init__( mixed_types: bool = True, seed: Optional[int] = None, **kwargs, - ): + ) -> None: self.task_dim = task_dim self.intensive = intensive self.bias_method = bias_method diff --git a/deepmd/pt/model/task/type_predict.py b/deepmd/pt/model/task/type_predict.py index e8a5db62b5..e4a980c3ea 100644 --- a/deepmd/pt/model/task/type_predict.py +++ b/deepmd/pt/model/task/type_predict.py @@ -14,7 +14,9 @@ class TypePredictNet(Fitting): - def __init__(self, feature_dim, ntypes, activation_function="gelu", **kwargs): + def __init__( + self, feature_dim, ntypes, activation_function="gelu", **kwargs + ) -> None: """Construct a type predict net. Args: diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index 35d5f4d659..928fc0943a 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -31,7 +31,7 @@ def __init__( kalman_lambda=0.98, kalman_nue=0.9987, block_size=5120, - ): + ) -> None: defaults = {"lr": 0.1, "kalman_nue": kalman_nue, "block_size": block_size} super().__init__(params, defaults) @@ -53,7 +53,7 @@ def __init__( self.remainder = 0 self.__init_P() - def __init_P(self): + def __init_P(self) -> None: param_nums = [] param_sum = 0 block_size = self.__get_blocksize() @@ -179,7 +179,7 @@ def __split_weights(self, weight): res.append(weight[i * block_size :]) return res - def __update(self, H, error, weights): + def __update(self, H, error, weights) -> None: P = self._state.get("P") kalman_lambda = self._state.get("kalman_lambda") weights_num = self._state.get("weights_num") @@ -253,10 +253,10 @@ def __update(self, H, error, weights): i += 1 param.data = tmp_weight.reshape(param.data.T.shape).T.contiguous() - def set_grad_prefactor(self, grad_prefactor): + def set_grad_prefactor(self, grad_prefactor) -> None: self.grad_prefactor = grad_prefactor - def step(self, error): + def step(self, error) -> None: params_packed_index = self._state.get("params_packed_index") weights = [] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 466080d34c..94f93ee180 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -99,7 +99,7 @@ def __init__( shared_links=None, finetune_links=None, init_frz_model=None, - ): + ) -> None: """Construct a DeePMD trainer. Args: @@ -477,7 +477,7 @@ def collect_single_finetune_params( _new_state_dict, _origin_state_dict, _random_state_dict, - ): + ) -> None: _new_fitting = _finetune_rule_single.get_random_fitting() _model_key_from = _finetune_rule_single.get_model_branch() target_keys = [ @@ -627,7 +627,7 @@ def warm_up_linear(step, warmup_steps): self.profiling = training_params.get("profiling", False) self.profiling_file = training_params.get("profiling_file", "timeline.json") - def run(self): + def run(self) -> None: fout = ( open( self.disp_file, @@ -662,7 +662,7 @@ def run(self): ) prof.start() - def step(_step_id, task_key="Default"): + def step(_step_id, task_key="Default") -> None: # PyTorch Profiler if self.enable_profiler or self.profiling: prof.step() @@ -1037,7 +1037,7 @@ def log_loss_valid(_task_key="Default"): f"The profiling trace have been saved to: {self.profiling_file}" ) - def save_model(self, save_path, lr=0.0, step=0): + def save_model(self, save_path, lr=0.0, step=0) -> None: module = ( self.wrapper.module if dist.is_available() and dist.is_initialized() @@ -1139,7 +1139,7 @@ def get_data(self, is_train=True, task_key="Default"): log_dict["sid"] = batch_data["sid"] return input_dict, label_dict, log_dict - def print_header(self, fout, train_results, valid_results): + def print_header(self, fout, train_results, valid_results) -> None: train_keys = sorted(train_results.keys()) print_str = "" print_str += "# %5s" % "step" @@ -1170,7 +1170,9 @@ def print_header(self, fout, train_results, valid_results): fout.write(print_str) fout.flush() - def print_on_training(self, fout, step_id, cur_lr, train_results, valid_results): + def print_on_training( + self, fout, step_id, cur_lr, train_results, valid_results + ) -> None: train_keys = sorted(train_results.keys()) print_str = "" print_str += "%7d" % step_id diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index 17fb8477a5..48119caf19 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -21,7 +21,7 @@ def __init__( loss: Union[torch.nn.Module, dict] = None, model_params=None, shared_links=None, - ): + ) -> None: """Construct a DeePMD model wrapper. Args: @@ -59,7 +59,7 @@ def __init__( self.loss[task_key] = loss[task_key] self.inference_only = self.loss is None - def share_params(self, shared_links, resume=False): + def share_params(self, shared_links, resume=False) -> None: """ Share the parameters of classes following rules defined in shared_links during multitask training. If not start from checkpoint (resume is False), @@ -182,7 +182,7 @@ def forward( ) return model_pred, loss, more_loss - def set_extra_state(self, state: dict): + def set_extra_state(self, state: dict) -> None: self.model_params = state["model_params"] self.train_infos = state["train_infos"] return None diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py index 3d6054296f..53942a176a 100644 --- a/deepmd/pt/utils/auto_batch_size.py +++ b/deepmd/pt/utils/auto_batch_size.py @@ -22,7 +22,7 @@ def __init__( self, initial_batch_size: int = 1024, factor: float = 2.0, - ): + ) -> None: super().__init__( initial_batch_size=initial_batch_size, factor=factor, diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 581f67196c..8b8c218d04 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -46,7 +46,7 @@ torch.multiprocessing.set_sharing_strategy("file_system") -def setup_seed(seed): +def setup_seed(seed) -> None: torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True @@ -76,7 +76,7 @@ def __init__( type_map, seed=None, shuffle=True, - ): + ) -> None: if seed is not None: setup_seed(seed) if isinstance(systems, str): @@ -152,7 +152,7 @@ def construct_dataset(system): for item in self.dataloaders: self.iters.append(iter(item)) - def set_noise(self, noise_settings): + def set_noise(self, noise_settings) -> None: # noise_settings['noise_type'] # "trunc_normal", "normal", "uniform" # noise_settings['noise'] # float, default 1.0 # noise_settings['noise_mode'] # "prob", "fix_num" @@ -162,7 +162,7 @@ def set_noise(self, noise_settings): for system in self.systems: system.set_noise(noise_settings) - def __len__(self): + def __len__(self) -> int: return len(self.dataloaders) def __getitem__(self, idx): @@ -175,7 +175,7 @@ def __getitem__(self, idx): batch["sid"] = idx return batch - def add_data_requirement(self, data_requirement: list[DataRequirementItem]): + def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: """Add data requirement for each system in multiple systems.""" for system in self.systems: system.add_data_requirement(data_requirement) @@ -184,7 +184,7 @@ def print_summary( self, name: str, prob: list[float], - ): + ) -> None: print_summary( name, len(self.systems), @@ -205,13 +205,13 @@ def print_summary( class BackgroundConsumer(Thread): - def __init__(self, queue, source, max_len): + def __init__(self, queue, source, max_len) -> None: Thread.__init__(self) self._queue = queue self._source = source # Main DL iterator self._max_len = max_len # - def run(self): + def run(self) -> None: for item in self._source: self._queue.put(item) # Blocking if the queue is full @@ -220,7 +220,7 @@ def run(self): class BufferedIterator: - def __init__(self, iterable): + def __init__(self, iterable) -> None: self._queue = queue.Queue(QUEUESIZE) self._iterable = iterable self._consumer = None @@ -229,7 +229,7 @@ def __init__(self, iterable): self.warning_time = None self.total = len(iterable) - def _create_consumer(self): + def _create_consumer(self) -> None: self._consumer = BackgroundConsumer(self._queue, self._iterable, self.total) self._consumer.daemon = True self._consumer.start() @@ -237,7 +237,7 @@ def _create_consumer(self): def __iter__(self): return self - def __len__(self): + def __len__(self) -> int: return self.total def __next__(self): diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 4a29f3f045..3043839308 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -16,7 +16,7 @@ class DeepmdDataSetForLoader(Dataset): - def __init__(self, system: str, type_map: Optional[list[str]] = None): + def __init__(self, system: str, type_map: Optional[list[str]] = None) -> None: """Construct DeePMD-style dataset containing frames cross different systems. Args: @@ -31,7 +31,7 @@ def __init__(self, system: str, type_map: Optional[list[str]] = None): self._natoms = self._data_system.get_natoms() self._natoms_vec = self._data_system.get_natoms_vec(self._ntypes) - def __len__(self): + def __len__(self) -> int: return self._data_system.nframes def __getitem__(self, index): @@ -40,7 +40,7 @@ def __getitem__(self, index): b_data["natoms"] = self._natoms_vec return b_data - def add_data_requirement(self, data_requirement: list[DataRequirementItem]): + def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: """Add data requirement for this data system.""" for data_item in data_requirement: self._data_system.add( diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index b253a1b55e..23e8627bcd 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -69,7 +69,7 @@ class EnvMatStatSe(EnvMatStat): The descriptor of the model. """ - def __init__(self, descriptor: "DescriptorBlock"): + def __init__(self, descriptor: "DescriptorBlock") -> None: super().__init__() self.descriptor = descriptor self.last_dim = ( diff --git a/deepmd/pt/utils/exclude_mask.py b/deepmd/pt/utils/exclude_mask.py index a5de969c07..0a99c0777f 100644 --- a/deepmd/pt/utils/exclude_mask.py +++ b/deepmd/pt/utils/exclude_mask.py @@ -15,7 +15,7 @@ def __init__( self, ntypes: int, exclude_types: list[int] = [], - ): + ) -> None: super().__init__() self.reinit(ntypes, exclude_types) @@ -23,7 +23,7 @@ def reinit( self, ntypes: int, exclude_types: list[int] = [], - ): + ) -> None: self.ntypes = ntypes self.exclude_types = exclude_types self.type_mask = np.array( @@ -68,7 +68,7 @@ def __init__( self, ntypes: int, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: super().__init__() self.reinit(ntypes, exclude_types) @@ -76,7 +76,7 @@ def reinit( self, ntypes: int, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.ntypes = ntypes self._exclude_types: set[tuple[int, int]] = set() for tt in exclude_types: diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py index e2076b3b2b..a4de8810af 100644 --- a/deepmd/pt/utils/multi_task.py +++ b/deepmd/pt/utils/multi_task.py @@ -96,7 +96,9 @@ def preprocess_shared_params(model_config): shared_links = {} type_map_keys = [] - def replace_one_item(params_dict, key_type, key_in_dict, suffix="", index=None): + def replace_one_item( + params_dict, key_type, key_in_dict, suffix="", index=None + ) -> None: shared_type = key_type shared_key = key_in_dict shared_level = 0 diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 4028d89fc9..2bcc87ce42 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -117,7 +117,7 @@ def _save_to_file( stat_file_path: DPPath, bias_out: dict, std_out: dict, -): +) -> None: assert stat_file_path is not None stat_file_path.mkdir(exist_ok=True, parents=True) for kk, vv in bias_out.items(): diff --git a/deepmd/pt/utils/tabulate.py b/deepmd/pt/utils/tabulate.py index e21d2ec9a6..db743ff98c 100644 --- a/deepmd/pt/utils/tabulate.py +++ b/deepmd/pt/utils/tabulate.py @@ -293,7 +293,7 @@ def _layer_1(self, x, w, b): t = torch.cat([x, x], dim=1) return t, self.activation_fn(torch.matmul(x, w) + b) + t - def _get_descrpt_type(self): + def _get_descrpt_type(self) -> str: if isinstance( self.descrpt, ( @@ -423,7 +423,7 @@ def _get_bias(self): def _get_matrix(self): return self._get_network_variable("w") - def _convert_numpy_to_tensor(self): + def _convert_numpy_to_tensor(self) -> None: """Convert self.data from np.ndarray to torch.Tensor.""" for ii in self.data: self.data[ii] = torch.tensor(self.data[ii], device=env.DEVICE) # pylint: disable=no-explicit-dtype diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index 43b82efcc1..6ce4f5d6fc 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -19,7 +19,7 @@ class ActivationFn(torch.nn.Module): - def __init__(self, activation: Optional[str]): + def __init__(self, activation: Optional[str]) -> None: super().__init__() self.activation: str = activation if activation is not None else "linear" @@ -101,7 +101,7 @@ def to_torch_tensor( return torch.tensor(xx, dtype=prec, device=DEVICE) -def dict_to_device(sample_dict): +def dict_to_device(sample_dict) -> None: for key in sample_dict: if isinstance(sample_dict[key], list): sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]] diff --git a/deepmd/tf/common.py b/deepmd/tf/common.py index 9244b4f7ee..df7842545a 100644 --- a/deepmd/tf/common.py +++ b/deepmd/tf/common.py @@ -282,7 +282,7 @@ def wrapper(self, *args, **kwargs): return wrapper -def clear_session(): +def clear_session() -> None: """Reset all state generated by DeePMD-kit.""" tf.reset_default_graph() _TF_VERSION = Version(TF_VERSION) diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py index e055efb31f..2ee35d9ebb 100644 --- a/deepmd/tf/descriptor/hybrid.py +++ b/deepmd/tf/descriptor/hybrid.py @@ -184,7 +184,7 @@ def compute_input_stats( **kwargs, ) - def merge_input_stats(self, stat_dict): + def merge_input_stats(self, stat_dict) -> None: """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd. Parameters diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index 746ea8c628..2863704143 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -103,7 +103,7 @@ def pass_tensors_from_frz_model( descrpt_deriv: tf.Tensor, rij: tf.Tensor, nlist: tf.Tensor, - ): + ) -> None: """Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def. Parameters diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index a0b6b810e4..c9a6d6ec57 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -396,7 +396,7 @@ def compute_input_stats( } self.merge_input_stats(stat_dict) - def merge_input_stats(self, stat_dict): + def merge_input_stats(self, stat_dict) -> None: """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd. Parameters diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index bf891e6032..d5476cbe89 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -411,7 +411,7 @@ def compute_input_stats( mesh, input_dict, **kwargs, - ): + ) -> None: data_efield = input_dict["efield"] all_davg = [] all_dstd = [] diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index b5896c9510..b6d180ee7c 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -1400,7 +1400,7 @@ def init_variables( graph_def, suffix=suffix ) - def compat_ln_pattern(old_key): + def compat_ln_pattern(old_key) -> None: pattern = r"attention_layer_(\d+)/(layer_normalization)_\d+" replacement = r"attention_layer_\1/\2" if bool(re.search(pattern, old_key)): diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 752642c1d5..d74b6df9ee 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -250,7 +250,7 @@ def compute_input_stats( mesh, input_dict, **kwargs, - ): + ) -> None: """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. Parameters @@ -283,7 +283,7 @@ def compute_input_stats( stat_dict = {"sumr": sumr, "sumn": sumn, "sumr2": sumr2} self.merge_input_stats(stat_dict) - def merge_input_stats(self, stat_dict): + def merge_input_stats(self, stat_dict) -> None: """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd. Parameters diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 464839aeac..79be372c10 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -280,7 +280,7 @@ def compute_input_stats( } self.merge_input_stats(stat_dict) - def merge_input_stats(self, stat_dict): + def merge_input_stats(self, stat_dict) -> None: """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd. Parameters @@ -763,7 +763,7 @@ def serialize_network( network_type="embedding_network", ) - def clear_ij(type_i, type_j): + def clear_ij(type_i, type_j) -> None: # initialize an empty network embeddings[(type_i, type_j)] = EmbeddingNet( in_dim=in_dim, diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py index 24982d3120..e8ceec7e9c 100644 --- a/deepmd/tf/entrypoints/compress.py +++ b/deepmd/tf/entrypoints/compress.py @@ -61,7 +61,7 @@ def compress( log_path: Optional[str], log_level: int, **kwargs, -): +) -> None: """Compress model. The table is composed of fifth-order polynomial coefficients and is assembled from @@ -199,7 +199,7 @@ def compress( ) from e -def _check_compress_type(graph: tf.Graph): +def _check_compress_type(graph: tf.Graph) -> None: try: t_model_type = bytes.decode(get_tensor_by_name_from_graph(graph, "model_type")) except GraphWithoutTensorError as e: diff --git a/deepmd/tf/entrypoints/convert.py b/deepmd/tf/entrypoints/convert.py index 17c8667362..84d9eb9f79 100644 --- a/deepmd/tf/entrypoints/convert.py +++ b/deepmd/tf/entrypoints/convert.py @@ -17,7 +17,7 @@ def convert( input_model: str, output_model: str, **kwargs, -): +) -> None: if output_model[-6:] == ".pbtxt": if input_model[-6:] != ".pbtxt": convert_pb_to_pbtxt(input_model, output_model) diff --git a/deepmd/tf/entrypoints/freeze.py b/deepmd/tf/entrypoints/freeze.py index 484faa4d19..83d2ff33f4 100755 --- a/deepmd/tf/entrypoints/freeze.py +++ b/deepmd/tf/entrypoints/freeze.py @@ -230,7 +230,7 @@ def freeze_graph( out_graph_name, node_names=None, out_suffix="", -): +) -> None: """Freeze the single graph with chosen out_suffix. Parameters @@ -300,7 +300,7 @@ def freeze( node_names: Optional[str] = None, nvnmd_weight: Optional[str] = None, **kwargs, -): +) -> None: """Freeze the graph in supplied folder. Parameters diff --git a/deepmd/tf/entrypoints/ipi.py b/deepmd/tf/entrypoints/ipi.py index a08a2293a9..189212a4e6 100644 --- a/deepmd/tf/entrypoints/ipi.py +++ b/deepmd/tf/entrypoints/ipi.py @@ -4,6 +4,9 @@ import os import subprocess import sys +from typing import ( + NoReturn, +) from deepmd.tf.lmp import ( get_op_dir, @@ -25,7 +28,7 @@ def _program(name: str, args: list[str]): return subprocess.call([os.path.join(ROOT_DIR, name), *args], close_fds=False) -def dp_ipi(): +def dp_ipi() -> NoReturn: """dp_ipi.""" suffix = ".exe" if os.name == "nt" else "" raise SystemExit(_program("dp_ipi" + suffix, sys.argv[1:])) diff --git a/deepmd/tf/entrypoints/main.py b/deepmd/tf/entrypoints/main.py index 9c3759288a..799a6280f6 100644 --- a/deepmd/tf/entrypoints/main.py +++ b/deepmd/tf/entrypoints/main.py @@ -38,7 +38,7 @@ __all__ = ["main", "parse_args", "get_ll", "main_parser"] -def main(args: Optional[Union[list[str], argparse.Namespace]] = None): +def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: """DeePMD-Kit entry point. Parameters diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 3d965ea71c..590fa04635 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -66,7 +66,7 @@ def train( finetune: Optional[str] = None, use_pretrain_script: bool = False, **kwargs, -): +) -> None: """Run DeePMD model training. Parameters @@ -185,7 +185,9 @@ def train( _do_work(jdata, run_opt, is_compress) -def _do_work(jdata: dict[str, Any], run_opt: RunOptions, is_compress: bool = False): +def _do_work( + jdata: dict[str, Any], run_opt: RunOptions, is_compress: bool = False +) -> None: """Run serial model training. Parameters diff --git a/deepmd/tf/entrypoints/transfer.py b/deepmd/tf/entrypoints/transfer.py index 52bf56c4fd..0ad022776a 100644 --- a/deepmd/tf/entrypoints/transfer.py +++ b/deepmd/tf/entrypoints/transfer.py @@ -57,7 +57,7 @@ def convert_matrix( return conv -def transfer(*, old_model: str, raw_model: str, output: str, **kwargs): +def transfer(*, old_model: str, raw_model: str, output: str, **kwargs) -> None: """Transfer operation from old from graph to new prepared raw graph. Parameters @@ -189,14 +189,14 @@ def __init__(self, node) -> None: def from_array( self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None - ): + ) -> None: if shape is None: shape = tensor.shape self.node.attr["value"].CopyFrom( tf.AttrValue(tensor=tf.make_tensor_proto(tensor, dtype, shape)) ) - def from_str(self, tensor: np.ndarray): + def from_str(self, tensor: np.ndarray) -> None: self.node.attr["value"].tensor.tensor_content = tensor.tobytes() @@ -209,7 +209,9 @@ def load_tensor(node: tf.Tensor, dtype_old: type, dtype_new: type) -> np.ndarray return tensor -def check_dim(raw_graph_node: tf.Tensor, old_graph_node: tf.Tensor, node_name: str): +def check_dim( + raw_graph_node: tf.Tensor, old_graph_node: tf.Tensor, node_name: str +) -> None: """Check if dimensions of tensor in old and new graph is equal. Parameters diff --git a/deepmd/tf/env.py b/deepmd/tf/env.py index 16ad4735fd..9a50f1281f 100644 --- a/deepmd/tf/env.py +++ b/deepmd/tf/env.py @@ -44,7 +44,7 @@ ) -def dlopen_library(module: str, filename: str): +def dlopen_library(module: str, filename: str) -> None: """Dlopen a library from a module. Parameters @@ -84,7 +84,7 @@ def dlopen_library(module: str, filename: str): class TFWarningFilter(logging.Filter): - def filter(self, record): + def filter(self, record) -> bool: return not any(msg in record.getMessage().strip() for msg in FILTER_MSGS) @@ -221,7 +221,7 @@ def filter(self, record): ) -def set_mkl(): +def set_mkl() -> None: """Tuning MKL for the best performance. References @@ -293,7 +293,7 @@ def get_tf_session_config() -> Any: default_tf_session_config = get_tf_session_config() -def reset_default_tf_session_config(cpu_only: bool): +def reset_default_tf_session_config(cpu_only: bool) -> None: """Limit tensorflow session to CPU or not. Parameters diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 040ec47cf7..7a3ee8eade 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - TYPE_CHECKING, Optional, ) @@ -67,9 +66,6 @@ check_version_compatibility, ) -if TYPE_CHECKING: - pass - log = logging.getLogger(__name__) diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index b5a21012bd..dd6a2d2963 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -160,7 +160,7 @@ def get_out_size(self) -> int: """Get the output size. Should be 9.""" return 9 - def compute_output_stats(self, all_stat): + def compute_output_stats(self, all_stat) -> None: """Compute the output statistics. Parameters diff --git a/deepmd/tf/infer/__init__.py b/deepmd/tf/infer/__init__.py index 9ef9c0d348..3596ff033e 100644 --- a/deepmd/tf/infer/__init__.py +++ b/deepmd/tf/infer/__init__.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Submodule containing all the implemented potentials.""" -from typing import ( - TYPE_CHECKING, -) - from deepmd.infer import ( DeepPotential, ) +from deepmd.infer.deep_eval import ( + DeepEval, +) from .data_modifier import ( DipoleChargeModifier, @@ -18,9 +17,6 @@ from .deep_dos import ( DeepDOS, ) -from .deep_eval import ( - DeepEval, -) from .deep_polar import ( DeepGlobalPolar, DeepPolar, @@ -38,11 +34,6 @@ calc_model_devi, ) -if TYPE_CHECKING: - from deepmd.infer.deep_eval import ( - DeepEval, - ) - __all__ = [ "DeepPotential", "DeepDipole", diff --git a/deepmd/tf/infer/deep_dipole.py b/deepmd/tf/infer/deep_dipole.py index b493af8552..09252d1d6c 100644 --- a/deepmd/tf/infer/deep_dipole.py +++ b/deepmd/tf/infer/deep_dipole.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from pathlib import ( - Path, -) from typing import ( + TYPE_CHECKING, Optional, ) @@ -13,6 +11,11 @@ DeepTensor, ) +if TYPE_CHECKING: + from pathlib import ( + Path, + ) + __all__ = [ "DeepDipole", ] diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 9527cb2ae8..9348b973b9 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -97,7 +97,7 @@ def __init__( input_map: Optional[dict] = None, neighbor_list=None, **kwargs: dict, - ): + ) -> None: self.graph = self._load_graph( model_file, prefix=load_prefix, @@ -169,7 +169,7 @@ def __init__( ewald_beta=ewald_beta, ) - def _init_tensors(self): + def _init_tensors(self) -> None: tensor_names = { # descrpt attrs "ntypes": "descrpt_attr/ntypes:0", @@ -224,7 +224,7 @@ def _init_tensors(self): for ii in sorted(removed_defs, reverse=True): del self.output_def.var_defs[list(self.output_def.var_defs.keys())[ii]] - def _init_attr(self): + def _init_attr(self) -> None: [ self.ntypes, self.rcut, @@ -1157,7 +1157,7 @@ def __init__( auto_batch_size: Union[bool, int, AutoBatchSize] = False, input_map: Optional[dict] = None, neighbor_list=None, - ): + ) -> None: self.graph = self._load_graph( model_file, prefix=load_prefix, diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index a1edaaa409..bdd36b6e09 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -104,7 +104,7 @@ def __init__( self._run_default_sess() self.tmap = self.tmap.decode("UTF-8").split() - def _run_default_sess(self): + def _run_default_sess(self) -> None: [self.ntypes, self.rcut, self.tmap, self.tselt, self.output_dim] = run_sess( self.sess, [ diff --git a/deepmd/tf/infer/ewald_recp.py b/deepmd/tf/infer/ewald_recp.py index f4b7d86588..ecfb243fe6 100644 --- a/deepmd/tf/infer/ewald_recp.py +++ b/deepmd/tf/infer/ewald_recp.py @@ -16,7 +16,7 @@ class EwaldRecp: """Evaluate the reciprocal part of the Ewald sum.""" - def __init__(self, hh, beta): + def __init__(self, hh, beta) -> None: """Constructor. Parameters diff --git a/deepmd/tf/loss/tensor.py b/deepmd/tf/loss/tensor.py index a5bcbbe025..aca9182ff6 100644 --- a/deepmd/tf/loss/tensor.py +++ b/deepmd/tf/loss/tensor.py @@ -21,7 +21,7 @@ class TensorLoss(Loss): """Loss function for tensorial properties.""" - def __init__(self, jdata, **kwarg): + def __init__(self, jdata, **kwarg) -> None: model = kwarg.get("model", None) if model is not None: self.type_sel = model.get_sel_type() diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 7ab068da63..1bebb4b971 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -90,7 +90,7 @@ def get_numb_aparam(self) -> int: """Get the number of atomic parameters.""" return self.numb_aparam - def data_stat(self, data): + def data_stat(self, data) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) self._compute_input_stat( @@ -99,7 +99,7 @@ def data_stat(self, data): # self._compute_output_stat(all_stat, mixed_type=data.mixed_type) # self.bias_atom_e = data.compute_energy_shift(self.rcond) - def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False): + def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False) -> None: if mixed_type: self.descrpt.compute_input_stats( all_stat["coord"], @@ -122,7 +122,7 @@ def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False): ) self.fitting.compute_input_stats(all_stat, protection=protection) - def _compute_output_stat(self, all_stat, mixed_type=False): + def _compute_output_stat(self, all_stat, mixed_type=False) -> None: if mixed_type: self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type) else: diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 57aaa2acf4..6d2ff4615f 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -135,7 +135,7 @@ def get_numb_aparam(self) -> int: """Get the number of atomic parameters.""" return self.numb_aparam - def data_stat(self, data): + def data_stat(self, data) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) self._compute_input_stat( @@ -144,7 +144,7 @@ def data_stat(self, data): self._compute_output_stat(all_stat, mixed_type=data.mixed_type) # self.bias_atom_e = data.compute_energy_shift(self.rcond) - def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False): + def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False) -> None: if mixed_type: self.descrpt.compute_input_stats( all_stat["coord"], @@ -167,7 +167,7 @@ def _compute_input_stat(self, all_stat, protection=1e-2, mixed_type=False): ) self.fitting.compute_input_stats(all_stat, protection=protection) - def _compute_output_stat(self, all_stat, mixed_type=False): + def _compute_output_stat(self, all_stat, mixed_type=False) -> None: if mixed_type: self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type) else: diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 7501a5cbd1..2099a8d947 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -6,6 +6,7 @@ Enum, ) from typing import ( + NoReturn, Optional, Union, ) @@ -56,7 +57,7 @@ class FrozenModel(Model): The path to the frozen model """ - def __init__(self, model_file: str, **kwargs): + def __init__(self, model_file: str, **kwargs) -> None: super().__init__(**kwargs) self.model_file = model_file if not model_file.endswith(".pb"): @@ -199,7 +200,7 @@ def get_rcut(self): def get_ntypes(self) -> int: return self.model.get_ntypes() - def data_stat(self, data): + def data_stat(self, data) -> None: pass def init_variables( @@ -279,7 +280,7 @@ def serialize(self, suffix: str = "") -> dict: return model.serialize() @classmethod - def deserialize(cls, data: dict, suffix: str = ""): + def deserialize(cls, data: dict, suffix: str = "") -> NoReturn: raise RuntimeError("Should not touch here.") @property diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 7cf3c5194d..63f55eae9e 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -48,7 +48,7 @@ class LinearModel(Model): If "sum", the weights are set to be 1. """ - def __init__(self, models: list[dict], weights: list[float], **kwargs): + def __init__(self, models: list[dict], weights: list[float], **kwargs) -> None: super().__init__(**kwargs) self.models = [Model(**model) for model in models] if isinstance(weights, list): @@ -90,7 +90,7 @@ def get_ntypes(self) -> int: raise ValueError("Models have different ntypes") return self.models[0].get_ntypes() - def data_stat(self, data): + def data_stat(self, data) -> None: for model in self.models: model.data_stat(data) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 5e3f99bc2d..5dcbad37ca 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -7,6 +7,7 @@ Enum, ) from typing import ( + NoReturn, Optional, Union, ) @@ -377,7 +378,7 @@ def _import_graph_def_from_ckpt_meta( sub_graph_def, input_map=feed_dict, return_elements=return_elements, name="" ) - def enable_mixed_precision(self, mixed_prec: dict): + def enable_mixed_precision(self, mixed_prec: dict) -> NoReturn: """Enable mixed precision for the model. Parameters @@ -415,7 +416,7 @@ def change_energy_bias( """ raise RuntimeError("Not supported") - def enable_compression(self, suffix: str = ""): + def enable_compression(self, suffix: str = "") -> NoReturn: """Enable compression. Parameters @@ -699,7 +700,7 @@ def __init__( else: self.typeebd = None - def enable_mixed_precision(self, mixed_prec: dict): + def enable_mixed_precision(self, mixed_prec: dict) -> None: """Enable mixed precision for the model. Parameters @@ -710,7 +711,7 @@ def enable_mixed_precision(self, mixed_prec: dict): self.descrpt.enable_mixed_precision(mixed_prec) self.fitting.enable_mixed_precision(mixed_prec) - def enable_compression(self, suffix: str = ""): + def enable_compression(self, suffix: str = "") -> None: """Enable compression. Parameters diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index 80e68d7825..d91554b98d 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -68,7 +68,7 @@ class PairTabModel(Model): def __init__( self, tab_file: str, rcut: float, sel: Union[int, list[int]], **kwargs - ): + ) -> None: super().__init__() self.tab_file = tab_file self.tab = PairTab(self.tab_file, rcut=rcut) @@ -255,7 +255,7 @@ def get_ntypes(self) -> int: """Get the number of types.""" return self.ntypes - def data_stat(self, data: dict): + def data_stat(self, data: dict) -> None: """Data staticis.""" # nothing needs to do diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index a0eaa1385f..5ed98f0c49 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -317,7 +317,7 @@ def get_rcut(self): def get_ntypes(self) -> int: return self.ntypes - def data_stat(self, data): + def data_stat(self, data) -> None: self.qm_model.data_stat(data) self.qmmm_model.data_stat(data) diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 8514844e03..8f361ce281 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -81,13 +81,13 @@ def get_sel_type(self): def get_out_size(self): return self.fitting.get_out_size() - def data_stat(self, data): + def data_stat(self, data) -> None: all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) self._compute_input_stat(m_all_stat, protection=self.data_stat_protect) self._compute_output_stat(m_all_stat) - def _compute_input_stat(self, all_stat, protection=1e-2): + def _compute_input_stat(self, all_stat, protection=1e-2) -> None: self.descrpt.compute_input_stats( all_stat["coord"], all_stat["box"], @@ -99,7 +99,7 @@ def _compute_input_stat(self, all_stat, protection=1e-2): if hasattr(self.fitting, "compute_input_stats"): self.fitting.compute_input_stats(all_stat, protection=protection) - def _compute_output_stat(self, all_stat): + def _compute_output_stat(self, all_stat) -> None: if hasattr(self.fitting, "compute_output_stats"): self.fitting.compute_output_stats(all_stat) diff --git a/deepmd/tf/nvnmd/descriptor/se_a.py b/deepmd/tf/nvnmd/descriptor/se_a.py index cc90df7a5c..0dbc0e3ac7 100644 --- a/deepmd/tf/nvnmd/descriptor/se_a.py +++ b/deepmd/tf/nvnmd/descriptor/se_a.py @@ -35,7 +35,7 @@ def build_davg_dstd(): return davg, dstd -def check_switch_range(davg, dstd): +def check_switch_range(davg, dstd) -> None: r"""Check the range of switch, let it in range [-2, 14].""" rmin = nvnmd_cfg.dscp["rcut_smth"] # diff --git a/deepmd/tf/nvnmd/descriptor/se_atten.py b/deepmd/tf/nvnmd/descriptor/se_atten.py index d9e93cf07f..1ab2369148 100644 --- a/deepmd/tf/nvnmd/descriptor/se_atten.py +++ b/deepmd/tf/nvnmd/descriptor/se_atten.py @@ -31,7 +31,7 @@ def build_davg_dstd(): return davg, dstd -def check_switch_range(davg, dstd): +def check_switch_range(davg, dstd) -> None: r"""Check the range of switch, let it in range [-2, 14].""" rmin = nvnmd_cfg.dscp["rcut_smth"] ntype = nvnmd_cfg.dscp["ntype"] diff --git a/deepmd/tf/nvnmd/entrypoints/freeze.py b/deepmd/tf/nvnmd/entrypoints/freeze.py index 2a2b8d9179..fe3d62c54e 100644 --- a/deepmd/tf/nvnmd/entrypoints/freeze.py +++ b/deepmd/tf/nvnmd/entrypoints/freeze.py @@ -60,7 +60,7 @@ def filter_tensorVariableList(tensorVariableList) -> dict: return dic_name_tv -def save_weight(sess, file_name: str = "nvnmd/weight.npy"): +def save_weight(sess, file_name: str = "nvnmd/weight.npy") -> None: r"""Save the dictionary of weight to a npy file.""" tvs = tf.global_variables() dic_key_tv = filter_tensorVariableList(tvs) diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 7a50ceae30..2e6e56bf51 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -77,7 +77,7 @@ class MapTable: DOI: 10.1038/s41524-022-00773-z """ - def __init__(self, config_file: str, weight_file: str, map_file: str): + def __init__(self, config_file: str, weight_file: str, map_file: str) -> None: self.config_file = config_file self.weight_file = weight_file self.map_file = map_file @@ -250,7 +250,7 @@ def mapping2(self, x, dic_map, cfgs): dic_val[key] = dats return dic_val - def plot_lines(self, x, dic1, dic2=None): + def plot_lines(self, x, dic1, dic2=None) -> None: r"""Plot lines to see accuracy.""" pass @@ -653,7 +653,7 @@ def mapt( nvnmd_weight: Optional[str] = "nvnmd/weight.npy", nvnmd_map: Optional[str] = "nvnmd/map.npy", **kwargs, -): +) -> None: # build mapping table mapObj = MapTable(nvnmd_config, nvnmd_weight, nvnmd_map) mapObj.build_map() diff --git a/deepmd/tf/nvnmd/entrypoints/train.py b/deepmd/tf/nvnmd/entrypoints/train.py index e6ca374977..c690190c0d 100644 --- a/deepmd/tf/nvnmd/entrypoints/train.py +++ b/deepmd/tf/nvnmd/entrypoints/train.py @@ -123,7 +123,7 @@ def train_nvnmd( step: str, skip_neighbor_stat: bool = False, **kwargs, -): +) -> None: # test input if not os.path.exists(INPUT): log.warning(f"The input script {INPUT} does not exist") diff --git a/deepmd/tf/nvnmd/entrypoints/wrap.py b/deepmd/tf/nvnmd/entrypoints/wrap.py index 7811f90f75..c4bee0acb5 100755 --- a/deepmd/tf/nvnmd/entrypoints/wrap.py +++ b/deepmd/tf/nvnmd/entrypoints/wrap.py @@ -71,7 +71,7 @@ class Wrap: def __init__( self, config_file: str, weight_file: str, map_file: str, model_file: str - ): + ) -> None: self.config_file = config_file self.weight_file = weight_file self.map_file = map_file @@ -85,7 +85,7 @@ def __init__( nvnmd_cfg.init_from_jdata(jdata) - def wrap(self): + def wrap(self) -> None: e = Encode() # cfg bcfg = self.wrap_dscp() @@ -613,6 +613,6 @@ def wrap( nvnmd_map: Optional[str] = "nvnmd/map.npy", nvnmd_model: Optional[str] = "nvnmd/model.pb", **kwargs, -): +) -> None: wrapObj = Wrap(nvnmd_config, nvnmd_weight, nvnmd_map, nvnmd_model) wrapObj.wrap() diff --git a/deepmd/tf/nvnmd/utils/config.py b/deepmd/tf/nvnmd/utils/config.py index 15998069b3..8b35a1a2aa 100644 --- a/deepmd/tf/nvnmd/utils/config.py +++ b/deepmd/tf/nvnmd/utils/config.py @@ -41,7 +41,7 @@ class NvnmdConfig: DOI: 10.1038/s41524-022-00773-z """ - def __init__(self, jdata: dict): + def __init__(self, jdata: dict) -> None: self.version = 0 self.enable = False self.map = {} @@ -50,7 +50,7 @@ def __init__(self, jdata: dict): self.weight = {} self.init_from_jdata(jdata) - def init_from_jdata(self, jdata: dict = {}): + def init_from_jdata(self, jdata: dict = {}) -> None: r"""Initialize this class with `jdata` loaded from input script.""" if jdata == {}: return None @@ -78,7 +78,7 @@ def init_from_jdata(self, jdata: dict = {}): # if load the file, set net_size self.init_net_size() - def init_value(self): + def init_value(self) -> None: r"""Initialize member with dict.""" self.dscp = self.config["dscp"] self.fitn = self.config["fitn"] @@ -87,7 +87,7 @@ def init_value(self): self.ctrl = self.config["ctrl"] self.nbit = self.config["nbit"] - def update_config(self): + def update_config(self) -> None: r"""Update config from dict.""" self.config["dscp"] = self.dscp self.config["fitn"] = self.fitn @@ -96,7 +96,7 @@ def update_config(self): self.config["ctrl"] = self.ctrl self.config["nbit"] = self.nbit - def init_train_mode(self, mod="cnn"): + def init_train_mode(self, mod="cnn") -> None: r"""Configure for taining cnn or qnn.""" if mod == "cnn": self.restore_descriptor = False @@ -109,7 +109,7 @@ def init_train_mode(self, mod="cnn"): self.quantize_descriptor = True self.quantize_fitting_net = True - def init_from_config(self, jdata): + def init_from_config(self, jdata) -> None: r"""Initialize member element one by one.""" if "ctrl" in jdata.keys(): if "VERSION" in jdata["ctrl"].keys(): @@ -128,7 +128,7 @@ def init_from_config(self, jdata): self.config["nbit"] = self.init_nbit(self.config["nbit"], self.config) self.init_value() - def init_config_by_version(self, version, max_nnei): + def init_config_by_version(self, version, max_nnei) -> None: r"""Initialize version-dependent parameters.""" self.version = version self.max_nnei = max_nnei @@ -152,13 +152,13 @@ def init_config_by_version(self, version, max_nnei): else: log.error("The max_nnei only can be set as 128|256 for version 1") - def init_net_size(self): + def init_net_size(self) -> None: r"""Initialize net_size.""" self.net_size = self.config["fitn"]["neuron"][0] if self.enable: self.config["fitn"]["neuron"] = [self.net_size] * 3 - def init_from_deepmd_input(self, jdata): + def init_from_deepmd_input(self, jdata) -> None: r"""Initialize members with input script of deepmd.""" fioObj = FioDic() self.config["dscp"] = fioObj.update(jdata["descriptor"], self.config["dscp"]) @@ -269,7 +269,7 @@ def init_nbit(self, jdata: dict, jdata_parent: dict = {}) -> dict: jdata["NBIT_SEL"] = int(np.ceil(np.log2(NSEL))) return jdata - def save(self, file_name=None): + def save(self, file_name=None) -> None: r"""Save all configuration to file.""" if file_name is None: file_name = self.save_path @@ -278,13 +278,13 @@ def save(self, file_name=None): self.update_config() FioDic().save(file_name, self.config) - def set_ntype(self, ntype): + def set_ntype(self, ntype) -> None: r"""Set the number of type.""" self.dscp["ntype"] = ntype self.config["dscp"]["ntype"] = ntype nvnmd_cfg.save() - def get_s_range(self, davg, dstd): + def get_s_range(self, davg, dstd) -> None: r"""Get the range of switch function.""" rmin = nvnmd_cfg.dscp["rcut_smth"] rmax = nvnmd_cfg.dscp["rcut"] @@ -380,7 +380,7 @@ def get_dp_init_weights(self): dic[key2] = self.weight[key] return dic - def disp_message(self): + def disp_message(self) -> None: r"""Display the log of NVNMD.""" NVNMD_CONFIG = ( f"enable: {self.enable}", diff --git a/deepmd/tf/nvnmd/utils/encode.py b/deepmd/tf/nvnmd/utils/encode.py index 46209e5230..4b82658262 100644 --- a/deepmd/tf/nvnmd/utils/encode.py +++ b/deepmd/tf/nvnmd/utils/encode.py @@ -13,7 +13,7 @@ class Encode: r"""Encoding value as hex, bin, and dec format.""" - def __init__(self): + def __init__(self) -> None: pass def qr(self, v, nbit: int = 14): @@ -95,7 +95,7 @@ def byte2hex(self, bs, nbyte): hs.append(h) return hs - def check_dec(self, idec, nbit, signed=False, name=""): + def check_dec(self, idec, nbit, signed=False, name="") -> None: r"""Check whether the data (idec) is in the range range is :math:`[0, 2^nbit-1]` for unsigned range is :math:`[-2^{nbit-1}, 2^{nbit-1}-1]` for signed. diff --git a/deepmd/tf/nvnmd/utils/fio.py b/deepmd/tf/nvnmd/utils/fio.py index 9daff62183..ab2f5baa57 100644 --- a/deepmd/tf/nvnmd/utils/fio.py +++ b/deepmd/tf/nvnmd/utils/fio.py @@ -12,7 +12,7 @@ class Fio: r"""Basic class for FIO.""" - def __init__(self): + def __init__(self) -> None: pass def exits(self, file_name=""): @@ -20,11 +20,11 @@ def exits(self, file_name=""): return True return os.path.exists(file_name) - def mkdir(self, path_name=""): + def mkdir(self, path_name="") -> None: if not self.exits(path_name): os.makedirs(path_name) - def create_file_path(self, file_name=""): + def create_file_path(self, file_name="") -> None: pars = file_name.split("/") if len(pars) > 0: path_name = "/".join(pars[:-1]) @@ -68,7 +68,7 @@ def load(self, file_name="", default_value={}): else: return FioNpyDic().load(file_name, default_value) - def save(self, file_name="", dic={}): + def save(self, file_name="", dic={}) -> None: if file_name.endswith(".json"): FioJsonDic().save(file_name, dic) elif file_name.endswith(".npy"): @@ -104,7 +104,7 @@ def update(self, jdata, jdata_o): class FioNpyDic: r"""Input and output for .npy file containing dictionary.""" - def __init__(self): + def __init__(self) -> None: pass def load(self, file_name="", default_value={}): @@ -116,7 +116,7 @@ def load(self, file_name="", default_value={}): log.warning(f"can not find {file_name}") return default_value - def save(self, file_name="", dic={}): + def save(self, file_name="", dic={}) -> None: Fio().create_file_path(file_name) np.save(file_name, [dic]) @@ -124,7 +124,7 @@ def save(self, file_name="", dic={}): class FioJsonDic: r"""Input and output for .json file containing dictionary.""" - def __init__(self): + def __init__(self) -> None: pass def load(self, file_name="", default_value={}): @@ -139,7 +139,7 @@ def load(self, file_name="", default_value={}): log.warning(f"can not find {file_name}") return default_value - def save(self, file_name="", dic={}): + def save(self, file_name="", dic={}) -> None: r"""Save dict into .json file.""" log.info(f"write jdata to {file_name}") Fio().create_file_path(file_name) @@ -150,7 +150,7 @@ def save(self, file_name="", dic={}): class FioBin: r"""Input and output for binary file.""" - def __init__(self): + def __init__(self) -> None: pass def load(self, file_name="", default_value=""): @@ -165,7 +165,7 @@ def load(self, file_name="", default_value=""): log.warning(f"can not find {file_name}") return default_value - def save(self, file_name: str, data: list[str]): + def save(self, file_name: str, data: list[str]) -> None: r"""Save hex string into binary file.""" log.info(f"write binary to {file_name}") Fio().create_file_path(file_name) @@ -183,7 +183,7 @@ def save(self, file_name: str, data: list[str]): class FioTxt: r"""Input and output for .txt file with string.""" - def __init__(self): + def __init__(self) -> None: pass def load(self, file_name="", default_value=[]): @@ -198,7 +198,7 @@ def load(self, file_name="", default_value=[]): log.info(f"can not find {file_name}") return default_value - def save(self, file_name: str = "", data: list = []): + def save(self, file_name: str = "", data: list = []) -> None: r"""Save string list into .txt file.""" log.info(f"write string to txt file {file_name}") Fio().create_file_path(file_name) diff --git a/deepmd/tf/op/__init__.py b/deepmd/tf/op/__init__.py index 805dc148a7..8c9b7e740d 100644 --- a/deepmd/tf/op/__init__.py +++ b/deepmd/tf/op/__init__.py @@ -13,7 +13,7 @@ log = logging.getLogger(__name__) -def import_ops(): +def import_ops() -> None: """Import all custom TF ops that are present in this submodule. Notes diff --git a/deepmd/tf/train/run_options.py b/deepmd/tf/train/run_options.py index c7f7b92674..fb6e09e8a9 100644 --- a/deepmd/tf/train/run_options.py +++ b/deepmd/tf/train/run_options.py @@ -115,7 +115,7 @@ def __init__( log_path: Optional[str] = None, log_level: int = 0, mpi_log: str = "master", - ): + ) -> None: self._try_init_distrib() # model init options @@ -145,7 +145,7 @@ def is_chief(self): """Whether my rank is 0.""" return self.my_rank == 0 - def print_resource_summary(self): + def print_resource_summary(self) -> None: """Print build and current running cluster configuration summary.""" SummaryPrinter(self.my_device, len(self.gpus or []))() @@ -154,7 +154,7 @@ def _setup_logger( log_path: Optional[Path], log_level: int, mpi_log: Optional[str], - ): + ) -> None: """Set up package loggers. Parameters @@ -183,7 +183,7 @@ def _setup_logger( f"reset them{', especially when running with MPI!' if self._HVD else ''}" ) - def _try_init_distrib(self): + def _try_init_distrib(self) -> None: try: import horovod.tensorflow as HVD @@ -201,7 +201,7 @@ def _try_init_distrib(self): self._init_serial() self._HVD = None - def _init_distributed(self, HVD: "HVD"): + def _init_distributed(self, HVD: "HVD") -> None: """Initialize settings for distributed training. Parameters @@ -231,7 +231,7 @@ def _init_distributed(self, HVD: "HVD"): else: self.my_device = "cpu:0" - def _init_serial(self): + def _init_serial(self) -> None: """Initialize setting for serial training.""" nodename, _, gpus = get_resource() diff --git a/deepmd/tf/train/trainer.py b/deepmd/tf/train/trainer.py index 58be9e8176..c6affdef7b 100644 --- a/deepmd/tf/train/trainer.py +++ b/deepmd/tf/train/trainer.py @@ -69,7 +69,7 @@ ) -def _is_subdir(path, directory): +def _is_subdir(path, directory) -> bool: path = os.path.realpath(path) directory = os.path.realpath(directory) if path == directory: @@ -79,12 +79,12 @@ def _is_subdir(path, directory): class DPTrainer: - def __init__(self, jdata, run_opt, is_compress=False): + def __init__(self, jdata, run_opt, is_compress=False) -> None: self.run_opt = run_opt self._init_param(jdata) self.is_compress = is_compress - def _init_param(self, jdata): + def _init_param(self, jdata) -> None: # model config model_param = jdata["model"] @@ -170,7 +170,7 @@ def get_lr_and_coef(lr_param): self.ckpt_meta = None self.model_type = None - def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): + def build(self, data=None, stop_batch=0, origin_type_map=None, suffix="") -> None: self.ntypes = self.model.get_ntypes() self.stop_batch = stop_batch @@ -240,7 +240,7 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): self._build_network(data, suffix) self._build_training() - def _build_lr(self): + def _build_lr(self) -> None: self._extra_train_ops = [] self.global_step = tf.train.get_or_create_global_step() self.learning_rate = self.lr.build(self.global_step, self.stop_batch) @@ -263,7 +263,7 @@ def _build_loss(self): return l2_l, l2_more - def _build_network(self, data, suffix=""): + def _build_network(self, data, suffix="") -> None: self.place_holders = {} if self.is_compress: for kk in ["coord", "box"]: @@ -328,7 +328,7 @@ def _build_optimizer(self): ) return optimizer - def _build_training(self): + def _build_training(self) -> None: if self.stop_batch == 0: # self.train_op is not used if stop_batch is zero self.train_op = None @@ -347,7 +347,7 @@ def _build_training(self): self.train_op = tf.group(*train_ops) log.info("built training") - def _init_session(self): + def _init_session(self) -> None: config = get_tf_session_config() device, idx = self.run_opt.my_device.split(":", 1) if device == "gpu": @@ -402,7 +402,7 @@ def _init_session(self): log.info("receive global variables from task#0") run_sess(self.sess, bcast_op) - def train(self, train_data=None, valid_data=None): + def train(self, train_data=None, valid_data=None) -> None: # if valid_data is None: # no validation set specified. # valid_data = train_data # using training set as validation set. @@ -626,7 +626,7 @@ def train(self, train_data=None, valid_data=None): if self.enable_profiler and self.run_opt.is_chief: tfv2.profiler.experimental.stop() - def save_checkpoint(self, cur_batch: int): + def save_checkpoint(self, cur_batch: int) -> None: try: ckpt_prefix = self.saver.save( self.sess, @@ -674,7 +674,7 @@ def get_global_step(self): def valid_on_the_fly( self, fp, train_batches, valid_batches, print_header=False, fitting_key=None - ): + ) -> None: train_results = self.get_evaluation_results(train_batches) valid_results = self.get_evaluation_results(valid_batches) @@ -691,7 +691,7 @@ def valid_on_the_fly( ) @staticmethod - def print_header(fp, train_results, valid_results): + def print_header(fp, train_results, valid_results) -> None: print_str = "" print_str += "# %5s" % "step" if valid_results is not None: @@ -714,7 +714,7 @@ def print_on_training( valid_results, cur_batch, cur_lr, - ): + ) -> None: print_str = "" print_str += "%7d" % cur_batch if valid_results is not None: @@ -778,13 +778,13 @@ def get_evaluation_results(self, batch_list): ) return avg_results - def save_compressed(self): + def save_compressed(self) -> None: """Save the compressed graph.""" self._init_session() if self.is_compress: self.saver.save(self.sess, os.path.join(os.getcwd(), self.save_ckpt)) - def _get_place_holders(self, data_dict): + def _get_place_holders(self, data_dict) -> None: for kk in data_dict.keys(): if kk == "type": continue @@ -796,7 +796,7 @@ def _get_place_holders(self, data_dict): tf.float32, name="t_find_" + kk ) - def _init_from_frz_model(self): + def _init_from_frz_model(self) -> None: try: graph, graph_def = load_graph_def(self.run_opt.init_frz_model) except FileNotFoundError as e: @@ -819,7 +819,7 @@ def _init_from_frz_model(self): self.frz_model = self.run_opt.init_frz_model self.model.init_variables(graph, graph_def, model_type=self.model_type) - def _init_from_ckpt(self, ckpt_meta: str): + def _init_from_ckpt(self, ckpt_meta: str) -> None: with tf.Graph().as_default() as graph: tf.train.import_meta_graph(f"{ckpt_meta}.meta", clear_devices=True) # get the model type from the model @@ -834,7 +834,7 @@ def _init_from_ckpt(self, ckpt_meta: str): def _init_from_pretrained_model( self, data, origin_type_map=None, bias_adjust_mode="change-by-statistic" - ): + ) -> None: """Init the embedding net variables with the given frozen model. Parameters @@ -886,7 +886,7 @@ def _change_energy_bias( frozen_model, origin_type_map, bias_adjust_mode="change-by-statistic", - ): + ) -> None: full_type_map = data.get_type_map() self.model.change_energy_bias( data, @@ -921,7 +921,7 @@ class DatasetLoader: >>> data_dict = loader.get_data_dict(data_list) """ - def __init__(self, train_data: DeepmdDataSystem): + def __init__(self, train_data: DeepmdDataSystem) -> None: self.train_data = train_data # get the keys of the data batch_data = self.train_data.get_batch() diff --git a/deepmd/tf/utils/convert.py b/deepmd/tf/utils/convert.py index 9005fb7d8e..461d870f80 100644 --- a/deepmd/tf/utils/convert.py +++ b/deepmd/tf/utils/convert.py @@ -59,7 +59,9 @@ def detect_model_version(input_model: str): return version -def convert_to_21(input_model: str, output_model: str, version: Optional[str] = None): +def convert_to_21( + input_model: str, output_model: str, version: Optional[str] = None +) -> None: """Convert DP graph to 2.1 graph. Parameters @@ -99,7 +101,7 @@ def convert_to_21(input_model: str, output_model: str, version: Optional[str] = ) -def convert_13_to_21(input_model: str, output_model: str): +def convert_13_to_21(input_model: str, output_model: str) -> None: """Convert DP 1.3 graph to 2.1 graph. Parameters @@ -112,7 +114,7 @@ def convert_13_to_21(input_model: str, output_model: str): convert_to_21(input_model, output_model, version="1.3") -def convert_12_to_21(input_model: str, output_model: str): +def convert_12_to_21(input_model: str, output_model: str) -> None: """Convert DP 1.2 graph to 2.1 graph. Parameters @@ -125,7 +127,7 @@ def convert_12_to_21(input_model: str, output_model: str): convert_to_21(input_model, output_model, version="1.2") -def convert_10_to_21(input_model: str, output_model: str): +def convert_10_to_21(input_model: str, output_model: str) -> None: """Convert DP 1.0 graph to 2.1 graph. Parameters @@ -138,7 +140,7 @@ def convert_10_to_21(input_model: str, output_model: str): convert_to_21(input_model, output_model, version="1.0") -def convert_012_to_21(input_model: str, output_model: str): +def convert_012_to_21(input_model: str, output_model: str) -> None: """Convert DP 0.12 graph to 2.1 graph. Parameters @@ -151,7 +153,7 @@ def convert_012_to_21(input_model: str, output_model: str): convert_to_21(input_model, output_model, version="0.12") -def convert_20_to_21(input_model: str, output_model: str): +def convert_20_to_21(input_model: str, output_model: str) -> None: """Convert DP 2.0 graph to 2.1 graph. Parameters @@ -164,7 +166,7 @@ def convert_20_to_21(input_model: str, output_model: str): convert_to_21(input_model, output_model, version="2.0") -def convert_pb_to_pbtxt(pbfile: str, pbtxtfile: str): +def convert_pb_to_pbtxt(pbfile: str, pbtxtfile: str) -> None: """Convert DP graph to graph text. Parameters @@ -181,7 +183,7 @@ def convert_pb_to_pbtxt(pbfile: str, pbtxtfile: str): tf.train.write_graph(graph_def, "./", pbtxtfile, as_text=True) -def convert_pbtxt_to_pb(pbtxtfile: str, pbfile: str): +def convert_pbtxt_to_pb(pbtxtfile: str, pbfile: str) -> None: """Convert DP graph text to graph. Parameters @@ -199,7 +201,7 @@ def convert_pbtxt_to_pb(pbtxtfile: str, pbfile: str): tf.train.write_graph(graph_def, "./", pbfile, as_text=False) -def convert_dp012_to_dp10(file: str): +def convert_dp012_to_dp10(file: str) -> None: """Convert DP 0.12 graph text to 1.0 graph text. Parameters @@ -302,7 +304,7 @@ def convert_dp012_to_dp10(file: str): fp.write(file_content) -def convert_dp10_to_dp11(file: str): +def convert_dp10_to_dp11(file: str) -> None: """Convert DP 1.0 graph text to 1.1 graph text. Parameters @@ -339,7 +341,7 @@ def convert_dp10_to_dp11(file: str): ) -def convert_dp12_to_dp13(file: str): +def convert_dp12_to_dp13(file: str) -> None: """Convert DP 1.2 graph text to 1.3 graph text. Parameters @@ -373,7 +375,7 @@ def convert_dp12_to_dp13(file: str): f.write(file_data) -def convert_dp13_to_dp20(fname: str): +def convert_dp13_to_dp20(fname: str) -> None: """Convert DP 1.3 graph text to 2.0 graph text. Parameters @@ -415,7 +417,7 @@ def convert_dp13_to_dp20(fname: str): fp.write(file_content) -def convert_dp20_to_dp21(fname: str): +def convert_dp20_to_dp21(fname: str) -> None: with open(fname) as fp: file_content = fp.read() old_model_version_node = textwrap.dedent( diff --git a/deepmd/tf/utils/network.py b/deepmd/tf/utils/network.py index c4a0646705..bd1b03e511 100644 --- a/deepmd/tf/utils/network.py +++ b/deepmd/tf/utils/network.py @@ -11,7 +11,7 @@ ) -def one_layer_rand_seed_shift(): +def one_layer_rand_seed_shift() -> int: return 3 @@ -392,7 +392,7 @@ def embedding_net( return xx -def variable_summaries(var: tf.Variable, name: str): +def variable_summaries(var: tf.Variable, name: str) -> None: """Attach a lot of summaries to a Tensor (for TensorBoard visualization). Parameters diff --git a/deepmd/tf/utils/spin.py b/deepmd/tf/utils/spin.py index ab70bdf319..8919bbd16a 100644 --- a/deepmd/tf/utils/spin.py +++ b/deepmd/tf/utils/spin.py @@ -38,7 +38,7 @@ def build( self, reuse=None, suffix="", - ): + ) -> None: """Build the computational graph for the spin. Parameters diff --git a/deepmd/tf/utils/tabulate.py b/deepmd/tf/utils/tabulate.py index 30171b12db..2244262786 100644 --- a/deepmd/tf/utils/tabulate.py +++ b/deepmd/tf/utils/tabulate.py @@ -169,7 +169,7 @@ def _load_sub_graph(self): tf.import_graph_def(sub_graph_def, name="") return sub_graph, sub_graph_def - def _get_descrpt_type(self): + def _get_descrpt_type(self) -> str: if isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeAtten): return "Atten" elif isinstance(self.descrpt, deepmd.tf.descriptor.DescrptSeAEbdV2): @@ -507,7 +507,7 @@ def _n_all_excluded(self) -> int: """Then number of types excluding all types.""" return sum(int(self._all_excluded(ii)) for ii in range(0, self.ntypes)) - def _convert_numpy_to_tensor(self): + def _convert_numpy_to_tensor(self) -> None: """Convert self.data from np.ndarray to tf.Tensor.""" for ii in self.data: self.data[ii] = tf.constant(self.data[ii]) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 3c4a4e6204..af9e2984dd 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -53,7 +53,7 @@ def list_to_doc(xx): return "".join(items) -def make_link(content, ref_key): +def make_link(content, ref_key) -> str: return ( f"`{content} <{ref_key}_>`_" if not dargs.RAW_ANCHOR @@ -70,7 +70,7 @@ def deprecate_argument_extra_check(key: str) -> Callable[[dict], bool]: The name of the deprecated argument. """ - def deprecate_something(data: Optional[dict]): + def deprecate_something(data: Optional[dict]) -> bool: if data is not None and key in data: warnings.warn(f"{key} has been removed and takes no effect.", FutureWarning) data.pop(key) @@ -2084,7 +2084,7 @@ def learning_rate_args(fold_subdoc: bool = False) -> Argument: # --- Loss configurations: --- # -def start_pref(item, label=None, abbr=None): +def start_pref(item, label=None, abbr=None) -> str: if label is None: label = item if abbr is None: @@ -2092,7 +2092,7 @@ def start_pref(item, label=None, abbr=None): return f"The prefactor of {item} loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the {label} label should be provided by file {label}.npy in each data system. If both start_pref_{abbr} and limit_pref_{abbr} are set to 0, then the {item} will be ignored." -def limit_pref(item): +def limit_pref(item) -> str: return f"The prefactor of {item} loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity." diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 5ab06e55e2..7cd125c97b 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -137,7 +137,7 @@ def execute( self._adjust_batch_size(self.factor) return n_batch, result - def _adjust_batch_size(self, factor: float): + def _adjust_batch_size(self, factor: float) -> None: old_batch_size = self.current_batch_size self.current_batch_size = int(self.current_batch_size * factor) log.info( diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index 83cbe46fad..72948c96f4 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -54,7 +54,7 @@ def convert_input_v0_v1( return output -def _warning_input_v0_v1(fname: Optional[Union[str, Path]]): +def _warning_input_v0_v1(fname: Optional[Union[str, Path]]) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 0.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" @@ -251,7 +251,7 @@ def _training(jdata: dict[str, Any]) -> dict[str, Any]: return training -def _jcopy(src: dict[str, Any], dst: dict[str, Any], keys: Sequence[str]): +def _jcopy(src: dict[str, Any], dst: dict[str, Any], keys: Sequence[str]) -> None: """Copy specified keys from one dict to another. Parameters @@ -268,7 +268,7 @@ def _jcopy(src: dict[str, Any], dst: dict[str, Any], keys: Sequence[str]): dst[k] = src[k] -def remove_decay_rate(jdata: dict[str, Any]): +def remove_decay_rate(jdata: dict[str, Any]) -> None: """Convert decay_rate to stop_lr. Parameters @@ -324,7 +324,7 @@ def convert_input_v1_v2( return jdata -def _warning_input_v1_v2(fname: Optional[Union[str, Path]]): +def _warning_input_v1_v2(fname: Optional[Union[str, Path]]) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 1.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 7d58d65578..493a9d8d54 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -57,7 +57,7 @@ def __init__( modifier=None, trn_all_set: bool = False, sort_atoms: bool = True, - ): + ) -> None: """Constructor.""" root = DPPath(sys_path) self.dirs = root.glob(set_prefix + ".*") @@ -403,7 +403,7 @@ def _get_subdata(self, data, idx=None): new_data[ii] = dd return new_data - def _load_batch_set(self, set_name: DPPath): + def _load_batch_set(self, set_name: DPPath) -> None: if not hasattr(self, "batch_set") or self.get_numb_set() > 1: self.batch_set = self._load_set(set_name) if self.modifier is not None: @@ -411,10 +411,10 @@ def _load_batch_set(self, set_name: DPPath): self.batch_set, _ = self._shuffle_data(self.batch_set) self.reset_get_batch() - def reset_get_batch(self): + def reset_get_batch(self) -> None: self.iterator = 0 - def _load_test_set(self, shuffle_test: bool): + def _load_test_set(self, shuffle_test: bool) -> None: test_sets = [] for ii in self.dirs: test_set = self._load_set(ii) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 03a399106f..0e960d0ba1 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -56,7 +56,7 @@ def __init__( sys_probs=None, auto_prob_style="prob_sys_size", sort_atoms: bool = True, - ): + ) -> None: """Constructor. Parameters @@ -231,7 +231,7 @@ def __init__( % (self.system_dirs[ii], chk_ret[0], self.test_size[ii], chk_ret[1]) ) - def _load_test(self, ntests=-1): + def _load_test(self, ntests=-1) -> None: self.test_data = collections.defaultdict(list) for ii in range(self.nsystems): test_system_data = self.data_systems[ii].get_test(ntests=ntests) @@ -313,7 +313,7 @@ def add( default: float = 0.0, dtype: Optional[np.dtype] = None, output_natoms_for_type_sel: bool = False, - ): + ) -> None: """Add a data item that to be loaded. Parameters @@ -356,7 +356,7 @@ def add( output_natoms_for_type_sel=output_natoms_for_type_sel, ) - def reduce(self, key_out, key_in): + def reduce(self, key_out, key_in) -> None: """Generate a new item from the reduction of another atom. Parameters @@ -372,7 +372,9 @@ def reduce(self, key_out, key_in): def get_data_dict(self, ii: int = 0) -> dict: return self.data_systems[ii].get_data_dict() - def set_sys_probs(self, sys_probs=None, auto_prob_style: str = "prob_sys_size"): + def set_sys_probs( + self, sys_probs=None, auto_prob_style: str = "prob_sys_size" + ) -> None: if sys_probs is None: if auto_prob_style == "prob_uniform": prob_v = 1.0 / float(self.nsystems) @@ -575,7 +577,7 @@ def get_batch_size(self) -> int: """Get the batch size.""" return self.batch_size - def print_summary(self, name: str): + def print_summary(self, name: str) -> None: print_summary( name, self.nsystems, @@ -642,7 +644,7 @@ def print_summary( nbatches: list[int], sys_probs: list[float], pbc: list[bool], -): +) -> None: """Print summary of systems. Parameters diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py index e33e07cee7..ce0da9e8fe 100644 --- a/deepmd/utils/econf_embd.py +++ b/deepmd/utils/econf_embd.py @@ -257,7 +257,7 @@ def transform(ii): return ret -def print_econf_embedding(res: dict[str, np.ndarray]): +def print_econf_embedding(res: dict[str, np.ndarray]) -> None: """Print electron configuration embedding.""" for kk, vv in res.items(): vvstr = ",".join([str(ii) for ii in vv]) diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index d8d035a853..644da3649d 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -12,7 +12,7 @@ def __init__( model_branch: str = "Default", random_fitting: bool = False, resuming: bool = False, - ): + ) -> None: """ The rules for fine-tuning the model from pretrained model. diff --git a/deepmd/utils/plugin.py b/deepmd/utils/plugin.py index ce8b015ddf..37ff784d61 100644 --- a/deepmd/utils/plugin.py +++ b/deepmd/utils/plugin.py @@ -29,7 +29,7 @@ def xxx(): >>> print(plugin.plugins["xx"]) """ - def __init__(self): + def __init__(self) -> None: self.plugins = {} def __add__(self, other) -> "Plugin": diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 440faca177..b0c75600fb 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -56,7 +56,7 @@ def random(size=None): return _RANDOM_GENERATOR.random_sample(size) -def seed(val: Optional[int] = None): +def seed(val: Optional[int] = None) -> None: """Seed the generator. Parameters @@ -67,7 +67,7 @@ def seed(val: Optional[int] = None): _RANDOM_GENERATOR.seed(val) -def shuffle(x: np.ndarray): +def shuffle(x: np.ndarray) -> None: """Modify a sequence in-place by shuffling its contents. Parameters diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 545b265b88..6e5e988b4a 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -240,7 +240,7 @@ def build( def _build_lower( self, net, xx, idx, upper, lower, stride0, stride1, extrapolate, nspline - ): + ) -> None: vv, dd, d2 = self._make_data(xx, idx) self.data[net] = np.zeros( [nspline, 6 * self.last_layer_size], dtype=self.data_type @@ -422,7 +422,7 @@ def _convert_numpy_to_tensor(self): """Convert self.data from np.ndarray to torch.Tensor.""" pass - def _convert_numpy_float_to_int(self): + def _convert_numpy_float_to_int(self) -> None: """Convert self.lower and self.upper from np.float32 or np.float64 to int.""" self.lower = {k: int(v) for k, v in self.lower.items()} self.upper = {k: int(v) for k, v in self.upper.items()} diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index ba1457b19c..c47a1d2fc9 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -55,7 +55,7 @@ def update_one_sel( ) return min_nbor_dist, sel - def parse_auto_sel(self, sel): + def parse_auto_sel(self, sel) -> bool: if not isinstance(sel, str): return False words = sel.split(":") diff --git a/deepmd/utils/version.py b/deepmd/utils/version.py index a0b479778d..5ba772b479 100644 --- a/deepmd/utils/version.py +++ b/deepmd/utils/version.py @@ -3,7 +3,7 @@ def check_version_compatibility( current_version: int, maximum_supported_version: int, minimal_supported_version: int = 1, -): +) -> None: """Check if the current version is compatible with the supported versions. Parameters diff --git a/doc/sphinx_contrib_exhale_multiproject.py b/doc/sphinx_contrib_exhale_multiproject.py index e26cc158a4..23b91f9f8d 100644 --- a/doc/sphinx_contrib_exhale_multiproject.py +++ b/doc/sphinx_contrib_exhale_multiproject.py @@ -85,7 +85,7 @@ import exhale.utils -def exhale_environment_ready(app): +def exhale_environment_ready(app) -> None: default_project = app.config.breathe_default_project default_exhale_args = dict(app.config.exhale_args) diff --git a/pyproject.toml b/pyproject.toml index f2312aba85..8f72d7457f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -375,6 +375,8 @@ select = [ "N804", # invalid-first-argument-name-for-class-method "N805", # invalid-first-argument-name-for-method "DTZ", # datetime + "TCH", # flake8-type-checking + "PYI", # flake8-pyi ] ignore = [ @@ -414,6 +416,9 @@ banned-module-level-imports = [ [tool.ruff.lint.flake8-tidy-imports.banned-api] "torch.testing.assert_allclose".msg = "Use `torch.testing.assert_close()` instead, see https://github.com/pytorch/pytorch/issues/61844." +[tool.ruff.lint.flake8-type-checking] +runtime-evaluated-base-classes = ["torch.nn.Module"] + [tool.ruff.lint.extend-per-file-ignores] # Also ignore `E402` in all `__init__.py` files. "deepmd/tf/**" = ["TID253"] diff --git a/source/checker/deepmd_checker.py b/source/checker/deepmd_checker.py index 0e11ed71c7..c066fe0c10 100644 --- a/source/checker/deepmd_checker.py +++ b/source/checker/deepmd_checker.py @@ -33,7 +33,7 @@ class DPChecker(BaseChecker): ), } - def visit_call(self, node): + def visit_call(self, node) -> None: if ( isinstance(node.func, Attribute) and isinstance(node.func.expr, Name) @@ -73,5 +73,5 @@ def register(linter: "PyLinter") -> None: pass -def load_configuration(linter): +def load_configuration(linter) -> None: linter.register_checker(DPChecker(linter)) diff --git a/source/install/build_tf.py b/source/install/build_tf.py index 0239ebfa46..d49458c032 100755 --- a/source/install/build_tf.py +++ b/source/install/build_tf.py @@ -56,6 +56,7 @@ ignore_patterns, ) from typing import ( + NoReturn, Optional, ) @@ -91,7 +92,7 @@ # Common utils -def download_file(url: str, filename: str): +def download_file(url: str, filename: str) -> None: """Download files from remote URL. Parameters @@ -154,7 +155,7 @@ def __call__(self): ) self.post_process() - def post_process(self): + def post_process(self) -> None: if self.executable: self.path.chmod(self.path.stat().st_mode | stat.S_IEXEC) if self.gzip is not None: @@ -168,7 +169,9 @@ def is_within_directory(directory, target): return prefix == abs_directory - def safe_extract(tar, path=".", members=None, *, numeric_owner=False): + def safe_extract( + tar, path=".", members=None, *, numeric_owner=False + ) -> None: for member in tar.getmembers(): member_path = os.path.join(path, member.name) if not is_within_directory(path, member_path): @@ -178,7 +181,7 @@ def safe_extract(tar, path=".", members=None, *, numeric_owner=False): safe_extract(tar, path=self.gzip_path) - def download(self): + def download(self) -> None: """Download the target file.""" download_file(self.url, self.path) @@ -230,7 +233,7 @@ def resources(self) -> dict[str, OnlineResource]: def dependencies(self) -> dict[str, "Build"]: """Required dependencies.""" - def download_all_resources(self): + def download_all_resources(self) -> None: """All resources, including dependencies' resources.""" for res in self.resources.values(): res() @@ -270,7 +273,7 @@ def prefix(self): """Tmp prefix.""" return self._prefix - def copy_from_tmp_to_prefix(self): + def copy_from_tmp_to_prefix(self) -> None: """Copy from tmp prefix to real prefix.""" copytree2(str(self.prefix), str(PREFIX)) @@ -306,7 +309,7 @@ def list2env(l: list) -> str: return ":".join(map(str, l)) -def get_shlib_ext(): +def get_shlib_ext() -> str: """Return the shared library extension.""" plat = sys.platform if plat.startswith("win"): @@ -324,7 +327,7 @@ def copy3(src: Path, dst: Path, *args, **kwargs): return copy2(str(src), str(dst), *args, **kwargs) -def copytree2(src: Path, dst: Path, *args, **kwargs): +def copytree2(src: Path, dst: Path, *args, **kwargs) -> None: """Wrapper to copytree and cp to support Pathlib, pattern, and override.""" with tempfile.TemporaryDirectory() as td: # hack to support override @@ -362,7 +365,7 @@ def _ignore_patterns(path, names): return _ignore_patterns -def call(commands: list[str], env={}, **kwargs): +def call(commands: list[str], env={}, **kwargs) -> None: """Call commands and print to screen for debug. Raises @@ -431,7 +434,7 @@ def resources(self) -> dict[str, OnlineResource]: def dependencies(self) -> dict[str, Build]: return {} - def build(self): + def build(self) -> None: bazel_res = self.resources["bazelisk"] bin_dst = self.prefix / "bin" bin_dst.mkdir(exist_ok=True) @@ -459,7 +462,7 @@ def dependencies(self) -> dict[str, Build]: def built(self) -> bool: return importlib.util.find_spec("numpy") is not None - def build(self): + def build(self) -> None: try: call( [ @@ -487,7 +490,7 @@ def resources(self) -> dict[str, OnlineResource]: def dependencies(self) -> dict[str, Build]: return {} - def build(self): + def build(self) -> NoReturn: raise RuntimeError( "NVCC is not found. Please manually install CUDA" "Toolkit and cuDNN!\n" @@ -534,7 +537,7 @@ def cudnn_version(self): @property @lru_cache - def cuda_compute_capabilities(self): + def cuda_compute_capabilities(self) -> str: """Get cuda compute capabilities.""" cuda_version = tuple(map(int, self.cuda_version.split("."))) if (10, 0, 0) <= cuda_version < (11, 0, 0): @@ -560,7 +563,7 @@ def resources(self) -> dict[str, OnlineResource]: def dependencies(self) -> dict[str, Build]: return {} - def build(self): + def build(self) -> NoReturn: raise RuntimeError("ROCm is not found!") @property @@ -616,7 +619,7 @@ def dependencies(self) -> dict[str, Build]: **optional_dep, } - def build(self): + def build(self) -> None: tf_res = self.resources["tensorflow"] src = tf_res.gzip_path / (f"tensorflow-{self.version}") with set_directory(src): @@ -718,7 +721,7 @@ def build(self): self.copy_lib("libtensorflow_framework" + ext, lib_src, lib_dst) self.copy_lib("libtensorflow_cc" + ext, lib_src, lib_dst) - def copy_lib(self, libname, src, dst): + def copy_lib(self, libname, src, dst) -> None: """Copy library and make symlink.""" copy3(src / (libname + "." + self.version), dst) libname_v = libname + "." + self.version @@ -806,7 +809,7 @@ def built(self): ).exists() -def clean_package(): +def clean_package() -> None: """Clean the unused files.""" clean_files = [ PACKAGE_DIR, diff --git a/source/ipi/tests/test_driver.py b/source/ipi/tests/test_driver.py index b0fbf53b01..54c632a2b6 100644 --- a/source/ipi/tests/test_driver.py +++ b/source/ipi/tests/test_driver.py @@ -27,7 +27,7 @@ class DPiPICalculator(FileIOCalculator): - def __init__(self, model: str, use_unix: bool = True, **kwargs): + def __init__(self, model: str, use_unix: bool = True, **kwargs) -> None: self.xyz_file = "test_ipi.xyz" self.config_file = "config.json" config = { @@ -49,20 +49,20 @@ def __init__(self, model: str, use_unix: bool = True, **kwargs): self, command=command, label=self.config_file, **kwargs ) - def write_input(self, atoms, **kwargs): + def write_input(self, atoms, **kwargs) -> None: atoms.write(self.xyz_file, format="xyz") class TestDPIPI(unittest.TestCase): # copy from test_deeppot_a.py @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.model_file = "deeppot.pb" convert_pbtxt_to_pb( str(tests_path / os.path.join("infer", "deeppot.pbtxt")), "deeppot.pb" ) - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -179,11 +179,11 @@ def setUp(self): ) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deeppot.pb") cls.dp = None - def test_ase_unix(self): + def test_ase_unix(self) -> None: with SocketIOCalculator( DPiPICalculator(self.model_file), log=sys.stdout, unixsocket="localhost" ) as calc: @@ -202,7 +202,7 @@ def test_ase_unix(self): expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1) np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) - def test_ase_nounix(self): + def test_ase_nounix(self) -> None: with SocketIOCalculator( DPiPICalculator(self.model_file, use_unix=False), log=sys.stdout, @@ -222,7 +222,7 @@ def test_ase_nounix(self): expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1) np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) - def test_normalize_coords(self): + def test_normalize_coords(self) -> None: # coordinate nomarlization should happen inside the interface cell = self.box.reshape((3, 3)) coord = self.coords.reshape((-1, 3)) @@ -255,10 +255,10 @@ def test_normalize_coords(self): class TestDPIPIPt(TestDPIPI): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.model_file = str(tests_path / "infer" / "deeppot_sea.pth") - def setUp(self): + def setUp(self) -> None: super().setUp() self.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @@ -354,5 +354,5 @@ def setUp(self): ) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: cls.dp = None diff --git a/source/jax2tf_tests/test_format_nlist.py b/source/jax2tf_tests/test_format_nlist.py index 9b95db2cf1..cd157ac974 100644 --- a/source/jax2tf_tests/test_format_nlist.py +++ b/source/jax2tf_tests/test_format_nlist.py @@ -14,7 +14,7 @@ class TestFormatNlist(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.nf = 3 self.nloc = 3 self.ns = 5 * 5 * 3 @@ -42,11 +42,11 @@ def setUp(self): distinguish_types=False, ) - def test_format_nlist_equal(self): + def test_format_nlist_equal(self) -> None: nlist = format_nlist(self.ecoord, self.nlist, sum(self.nsel), self.rcut) self.assertAllEqual(nlist, self.nlist) - def test_format_nlist_less(self): + def test_format_nlist_less(self) -> None: nlist = build_neighbor_list( self.ecoord, self.eatype, @@ -58,7 +58,7 @@ def test_format_nlist_less(self): nlist = format_nlist(self.ecoord, nlist, sum(self.nsel), self.rcut) self.assertAllEqual(nlist, self.nlist) - def test_format_nlist_large(self): + def test_format_nlist_large(self) -> None: nlist = build_neighbor_list( self.ecoord, self.eatype, @@ -74,7 +74,7 @@ def test_format_nlist_large(self): # we only need to ensure the result is correct, no need to check the order self.assertAllEqual(tnp.sort(nlist, axis=-1), tnp.sort(self.nlist, axis=-1)) - def test_format_nlist_larger_rcut(self): + def test_format_nlist_larger_rcut(self) -> None: nlist = build_neighbor_list( self.ecoord, self.eatype, diff --git a/source/jax2tf_tests/test_nlist.py b/source/jax2tf_tests/test_nlist.py index 5b13e4231c..c5906fa2fc 100644 --- a/source/jax2tf_tests/test_nlist.py +++ b/source/jax2tf_tests/test_nlist.py @@ -16,7 +16,7 @@ class TestNeighList(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.nf = 3 self.nloc = 3 self.ns = 5 * 5 * 3 @@ -43,7 +43,7 @@ def setUp(self): ] ) - def test_build_notype(self): + def test_build_notype(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -64,7 +64,7 @@ def test_build_notype(self): tnp.sort(self.ref_nlist, axis=-1), ) - def test_build_type(self): + def test_build_type(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -86,7 +86,7 @@ def test_build_type(self): tnp.sort(tnp.split(self.ref_nlist, self.nsel, axis=-1)[ii], axis=-1), ) - def test_extend_coord(self): + def test_extend_coord(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) diff --git a/source/jax2tf_tests/test_region.py b/source/jax2tf_tests/test_region.py index 2becf08c94..c7bd182889 100644 --- a/source/jax2tf_tests/test_region.py +++ b/source/jax2tf_tests/test_region.py @@ -13,7 +13,7 @@ class TestRegion(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.cell = tnp.array( [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], ) @@ -21,7 +21,7 @@ def setUp(self): self.cell = tnp.tile(self.cell, [4, 5, 1, 1]) self.prec = 1e-8 - def test_inter_to_phys(self): + def test_inter_to_phys(self) -> None: rng = tf.random.Generator.from_seed(GLOBAL_SEED) inter = rng.normal(shape=[4, 5, 3, 3]) phys = inter2phys(inter, self.cell) @@ -32,7 +32,7 @@ def test_inter_to_phys(self): phys[ii, jj], expected_phys, rtol=self.prec, atol=self.prec ) - def test_to_face_dist(self): + def test_to_face_dist(self) -> None: cell0 = self.cell[0][0] vol = tf.linalg.det(cell0) # area of surfaces xy, xz, yz diff --git a/source/lmp/tests/test_deeptensor.py b/source/lmp/tests/test_deeptensor.py index 6df0a8617a..6fb7cde746 100644 --- a/source/lmp/tests/test_deeptensor.py +++ b/source/lmp/tests/test_deeptensor.py @@ -65,7 +65,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) # TODO # write_lmp_data(box, coord, type_HO, data_type_map_file) @@ -77,7 +77,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) # os.remove(data_type_map_file) os.remove(data_file_si) @@ -135,7 +135,7 @@ def lammps_si(): lmp.close() -def test_compute_deeptensor_atom(lammps): +def test_compute_deeptensor_atom(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute(f"tensor all deeptensor/atom {pb_file2.resolve()}") @@ -148,7 +148,7 @@ def test_compute_deeptensor_atom(lammps): ) -def test_compute_deeptensor_atom_si(lammps_si): +def test_compute_deeptensor_atom_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.compute(f"tensor all deeptensor/atom {pb_file2.resolve()}") diff --git a/source/lmp/tests/test_dplr.py b/source/lmp/tests/test_dplr.py index 4750849860..532bafef36 100644 --- a/source/lmp/tests/test_dplr.py +++ b/source/lmp/tests/test_dplr.py @@ -270,7 +270,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data_full( box, coord, mol_list, type_OH, charge, data_file, bond_list, mass_list ) @@ -292,7 +292,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_file2) os.remove(data_type_map_file) @@ -352,7 +352,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd_sr(lammps): +def test_pair_deepmd_sr(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -365,7 +365,7 @@ def test_pair_deepmd_sr(lammps): lammps.run(1) -def test_pair_deepmd_sr_virial(lammps): +def test_pair_deepmd_sr_virial(lammps) -> None: lammps.group("real_atom type 1 2") lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") @@ -392,7 +392,7 @@ def test_pair_deepmd_sr_virial(lammps): os.remove("dump") -def test_pair_deepmd_lr(lammps): +def test_pair_deepmd_lr(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.bond_style("zero") @@ -420,7 +420,7 @@ def test_pair_deepmd_lr(lammps): lammps.run(1) -def test_pair_deepmd_lr_run0(lammps2): +def test_pair_deepmd_lr_run0(lammps2) -> None: lammps2.pair_style(f"deepmd {pb_file.resolve()}") lammps2.pair_coeff("* *") lammps2.bond_style("zero") @@ -434,7 +434,7 @@ def test_pair_deepmd_lr_run0(lammps2): lammps2.run(0) -def test_pair_deepmd_lr_efield_constant(lammps): +def test_pair_deepmd_lr_efield_constant(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.bond_style("zero") @@ -469,7 +469,7 @@ def test_pair_deepmd_lr_efield_constant(lammps): ) -def test_pair_deepmd_lr_efield_variable(lammps): +def test_pair_deepmd_lr_efield_variable(lammps) -> None: lammps.variable("EFIELD_Z equal 2*sin(2*PI*time/0.006)") lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") @@ -505,7 +505,7 @@ def test_pair_deepmd_lr_efield_variable(lammps): ) -def test_min_dplr(lammps): +def test_min_dplr(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.bond_style("zero") @@ -529,7 +529,7 @@ def test_min_dplr(lammps): ) -def test_pair_deepmd_lr_type_map(lammps_type_map): +def test_pair_deepmd_lr_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.bond_style("zero") @@ -559,7 +559,7 @@ def test_pair_deepmd_lr_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_lr_si(lammps_si): +def test_pair_deepmd_lr_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.bond_style("zero") diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py index bc554b1c6a..645d826e31 100644 --- a/source/lmp/tests/test_lammps.py +++ b/source/lmp/tests/test_lammps.py @@ -229,7 +229,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -240,7 +240,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -306,7 +306,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -318,7 +318,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("peatom all pe/atom pair") @@ -358,7 +358,7 @@ def test_pair_deepmd_virial(lammps): ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -384,7 +384,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -439,7 +439,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -468,7 +468,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -501,7 +501,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) @@ -513,7 +513,7 @@ def test_pair_deepmd_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_real(lammps_real): +def test_pair_deepmd_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.run(0) @@ -527,7 +527,7 @@ def test_pair_deepmd_real(lammps_real): lammps_real.run(1) -def test_pair_deepmd_virial_real(lammps_real): +def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.compute("virial all centroid/stress/atom NULL pair") @@ -554,7 +554,7 @@ def test_pair_deepmd_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_real(lammps_real): +def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -584,7 +584,7 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) -def test_pair_deepmd_model_devi_virial_real(lammps_real): +def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -628,7 +628,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative * constants.force_metal2real}" @@ -661,7 +661,7 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative * constants.ener_metal2real}" @@ -698,7 +698,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): ) -def test_pair_deepmd_si(lammps_si): +def test_pair_deepmd_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.run(0) @@ -720,7 +720,7 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_3types.py b/source/lmp/tests/test_lammps_3types.py index e4e64d9ecf..c8d3193817 100644 --- a/source/lmp/tests/test_lammps_3types.py +++ b/source/lmp/tests/test_lammps_3types.py @@ -252,12 +252,12 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -292,7 +292,7 @@ def lammps_type_map(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -304,7 +304,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") @@ -327,7 +327,7 @@ def test_pair_deepmd_virial(lammps): ) / nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -353,7 +353,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -391,7 +391,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -420,7 +420,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -453,7 +453,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) diff --git a/source/lmp/tests/test_lammps_dpa_pt.py b/source/lmp/tests/test_lammps_dpa_pt.py index 8747884ef6..d92084faa1 100644 --- a/source/lmp/tests/test_lammps_dpa_pt.py +++ b/source/lmp/tests/test_lammps_dpa_pt.py @@ -225,7 +225,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -236,7 +236,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -302,7 +302,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -314,7 +314,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") @@ -337,7 +337,7 @@ def test_pair_deepmd_virial(lammps): ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -363,7 +363,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -401,7 +401,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -430,7 +430,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -463,7 +463,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) @@ -475,7 +475,7 @@ def test_pair_deepmd_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_real(lammps_real): +def test_pair_deepmd_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.run(0) @@ -489,7 +489,7 @@ def test_pair_deepmd_real(lammps_real): lammps_real.run(1) -def test_pair_deepmd_virial_real(lammps_real): +def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.compute("virial all centroid/stress/atom NULL pair") @@ -516,7 +516,7 @@ def test_pair_deepmd_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_real(lammps_real): +def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -546,7 +546,7 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) -def test_pair_deepmd_model_devi_virial_real(lammps_real): +def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -590,7 +590,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative * constants.force_metal2real}" @@ -623,7 +623,7 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative * constants.ener_metal2real}" @@ -660,7 +660,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): ) -def test_pair_deepmd_si(lammps_si): +def test_pair_deepmd_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.run(0) @@ -682,7 +682,7 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py index b0909bfc03..2183852770 100644 --- a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py +++ b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py @@ -223,7 +223,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -234,7 +234,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -300,7 +300,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -312,7 +312,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") @@ -335,7 +335,7 @@ def test_pair_deepmd_virial(lammps): ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -361,7 +361,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -399,7 +399,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -428,7 +428,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -461,7 +461,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) @@ -473,7 +473,7 @@ def test_pair_deepmd_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_real(lammps_real): +def test_pair_deepmd_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.run(0) @@ -487,7 +487,7 @@ def test_pair_deepmd_real(lammps_real): lammps_real.run(1) -def test_pair_deepmd_virial_real(lammps_real): +def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.compute("virial all centroid/stress/atom NULL pair") @@ -514,7 +514,7 @@ def test_pair_deepmd_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_real(lammps_real): +def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -544,7 +544,7 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) -def test_pair_deepmd_model_devi_virial_real(lammps_real): +def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -588,7 +588,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative * constants.force_metal2real}" @@ -621,7 +621,7 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative * constants.ener_metal2real}" @@ -658,7 +658,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): ) -def test_pair_deepmd_si(lammps_si): +def test_pair_deepmd_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.run(0) @@ -680,7 +680,7 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_dpa_sel_pt.py b/source/lmp/tests/test_lammps_dpa_sel_pt.py index 45e148fcaa..57395fda98 100644 --- a/source/lmp/tests/test_lammps_dpa_sel_pt.py +++ b/source/lmp/tests/test_lammps_dpa_sel_pt.py @@ -228,7 +228,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -239,7 +239,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -305,7 +305,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -317,7 +317,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") @@ -340,7 +340,7 @@ def test_pair_deepmd_virial(lammps): ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -366,7 +366,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -404,7 +404,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -433,7 +433,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -466,7 +466,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) @@ -478,7 +478,7 @@ def test_pair_deepmd_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_real(lammps_real): +def test_pair_deepmd_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.run(0) @@ -492,7 +492,7 @@ def test_pair_deepmd_real(lammps_real): lammps_real.run(1) -def test_pair_deepmd_virial_real(lammps_real): +def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.compute("virial all centroid/stress/atom NULL pair") @@ -519,7 +519,7 @@ def test_pair_deepmd_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_real(lammps_real): +def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -549,7 +549,7 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) -def test_pair_deepmd_model_devi_virial_real(lammps_real): +def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -593,7 +593,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative * constants.force_metal2real}" @@ -626,7 +626,7 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative * constants.ener_metal2real}" @@ -663,7 +663,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): ) -def test_pair_deepmd_si(lammps_si): +def test_pair_deepmd_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.run(0) @@ -685,7 +685,7 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_faparam.py b/source/lmp/tests/test_lammps_faparam.py index f78639a96b..b388f0453b 100644 --- a/source/lmp/tests/test_lammps_faparam.py +++ b/source/lmp/tests/test_lammps_faparam.py @@ -139,11 +139,11 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) @@ -185,7 +185,7 @@ def lammps(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()} fparam 0.25852028 aparam 0.25852028") lammps.pair_coeff("* *") lammps.run(0) @@ -197,7 +197,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()} fparam 0.25852028 aparam 0.25852028") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") diff --git a/source/lmp/tests/test_lammps_pt.py b/source/lmp/tests/test_lammps_pt.py index 245f3eaf6d..f2a4bac23d 100644 --- a/source/lmp/tests/test_lammps_pt.py +++ b/source/lmp/tests/test_lammps_pt.py @@ -225,7 +225,7 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -236,7 +236,7 @@ def setup_module(): ) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) os.remove(data_type_map_file) @@ -302,7 +302,7 @@ def lammps_si(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -314,7 +314,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepmd {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("virial all centroid/stress/atom NULL pair") @@ -337,7 +337,7 @@ def test_pair_deepmd_virial(lammps): ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -363,7 +363,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_virial(lammps): +def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -401,7 +401,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -430,7 +430,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_model_devi_atomic_relative_v(lammps): +def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative}" @@ -463,7 +463,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) -def test_pair_deepmd_type_map(lammps_type_map): +def test_pair_deepmd_type_map(lammps_type_map) -> None: lammps_type_map.pair_style(f"deepmd {pb_file.resolve()}") lammps_type_map.pair_coeff("* * H O") lammps_type_map.run(0) @@ -475,7 +475,7 @@ def test_pair_deepmd_type_map(lammps_type_map): lammps_type_map.run(1) -def test_pair_deepmd_real(lammps_real): +def test_pair_deepmd_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.run(0) @@ -489,7 +489,7 @@ def test_pair_deepmd_real(lammps_real): lammps_real.run(1) -def test_pair_deepmd_virial_real(lammps_real): +def test_pair_deepmd_virial_real(lammps_real) -> None: lammps_real.pair_style(f"deepmd {pb_file.resolve()}") lammps_real.pair_coeff("* *") lammps_real.compute("virial all centroid/stress/atom NULL pair") @@ -516,7 +516,7 @@ def test_pair_deepmd_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_real(lammps_real): +def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -546,7 +546,7 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) -def test_pair_deepmd_model_devi_virial_real(lammps_real): +def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" ) @@ -590,7 +590,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative * constants.force_metal2real}" @@ -623,7 +623,7 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) -def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): +def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative_v {relative * constants.ener_metal2real}" @@ -660,7 +660,7 @@ def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): ) -def test_pair_deepmd_si(lammps_si): +def test_pair_deepmd_si(lammps_si) -> None: lammps_si.pair_style(f"deepmd {pb_file.resolve()}") lammps_si.pair_coeff("* *") lammps_si.run(0) @@ -682,7 +682,7 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_spin.py b/source/lmp/tests/test_lammps_spin.py index cd65f6d3ce..39e12b03fc 100644 --- a/source/lmp/tests/test_lammps_spin.py +++ b/source/lmp/tests/test_lammps_spin.py @@ -99,11 +99,11 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data_spin(box, coord, spin, type_NiO, data_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) @@ -138,7 +138,7 @@ def lammps(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -150,7 +150,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("peatom all pe/atom pair") @@ -190,7 +190,7 @@ def test_pair_deepmd_virial(lammps): # ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" ) @@ -213,7 +213,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[9] == pytest.approx(np.mean(expected_md_fm)) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -251,7 +251,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_spin_nopbc.py b/source/lmp/tests/test_lammps_spin_nopbc.py index e507a61add..3ef9a7a355 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc.py +++ b/source/lmp/tests/test_lammps_spin_nopbc.py @@ -98,11 +98,11 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data_spin(box, coord, spin, type_NiO, data_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) @@ -137,7 +137,7 @@ def lammps(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -149,7 +149,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" ) @@ -172,7 +172,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[9] == pytest.approx(np.mean(expected_md_fm)) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -210,7 +210,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_spin_nopbc_pt.py b/source/lmp/tests/test_lammps_spin_nopbc_pt.py index b34d2281a5..8a4dc7a7f1 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc_pt.py +++ b/source/lmp/tests/test_lammps_spin_nopbc_pt.py @@ -95,11 +95,11 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data_spin(box, coord, spin, type_NiO, data_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) @@ -134,7 +134,7 @@ def lammps(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -146,7 +146,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" ) @@ -169,7 +169,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[9] == pytest.approx(np.mean(expected_md_fm)) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -207,7 +207,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/test_lammps_spin_pt.py b/source/lmp/tests/test_lammps_spin_pt.py index c3bd27be38..e15a10ee72 100644 --- a/source/lmp/tests/test_lammps_spin_pt.py +++ b/source/lmp/tests/test_lammps_spin_pt.py @@ -95,11 +95,11 @@ ) -def setup_module(): +def setup_module() -> None: write_lmp_data_spin(box, coord, spin, type_NiO, data_file) -def teardown_module(): +def teardown_module() -> None: os.remove(data_file) @@ -134,7 +134,7 @@ def lammps(): lmp.close() -def test_pair_deepmd(lammps): +def test_pair_deepmd(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.run(0) @@ -146,7 +146,7 @@ def test_pair_deepmd(lammps): lammps.run(1) -def test_pair_deepmd_virial(lammps): +def test_pair_deepmd_virial(lammps) -> None: lammps.pair_style(f"deepspin {pb_file.resolve()}") lammps.pair_coeff("* *") lammps.compute("peatom all pe/atom pair") @@ -186,7 +186,7 @@ def test_pair_deepmd_virial(lammps): # ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) -def test_pair_deepmd_model_devi(lammps): +def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" ) @@ -209,7 +209,7 @@ def test_pair_deepmd_model_devi(lammps): assert md[9] == pytest.approx(np.mean(expected_md_fm)) -def test_pair_deepmd_model_devi_atomic_relative(lammps): +def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic relative {relative}" @@ -247,7 +247,7 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): ("balance_args",), [(["--balance"],), ([],)], ) -def test_pair_deepmd_mpi(balance_args: list): +def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( [ diff --git a/source/lmp/tests/write_lmp_data.py b/source/lmp/tests/write_lmp_data.py index fd86c8d61a..747b42160d 100644 --- a/source/lmp/tests/write_lmp_data.py +++ b/source/lmp/tests/write_lmp_data.py @@ -4,7 +4,7 @@ comment_lmp_data = "# the first line must be comment" -def write_lmp_data(box, coord, type_list, file_name): +def write_lmp_data(box, coord, type_list, file_name) -> None: natom = coord.shape[0] ntype = np.unique(type_list).shape[0] with open(file_name, "w") as f: @@ -25,7 +25,7 @@ def write_lmp_data(box, coord, type_list, file_name): def write_lmp_data_full( box, coord, mol_list, type_list, charge, file_name, bond_list, mass_list -): +) -> None: natom = coord.shape[0] ntype = np.unique(type_list).shape[0] nbond_type = len(bond_list) @@ -71,7 +71,7 @@ def write_lmp_data_full( f.write("\n") -def write_lmp_data_spin(box, coord, spin, type_list, file_name): +def write_lmp_data_spin(box, coord, spin, type_list, file_name) -> None: natom = coord.shape[0] ntype = np.unique(type_list).shape[0] sp_norm = np.linalg.norm(spin, axis=1, keepdims=True) diff --git a/source/tests/common/dpmodel/array_api/test_env_mat.py b/source/tests/common/dpmodel/array_api/test_env_mat.py index 0c0a69fc2e..607b0515bf 100644 --- a/source/tests/common/dpmodel/array_api/test_env_mat.py +++ b/source/tests/common/dpmodel/array_api/test_env_mat.py @@ -13,7 +13,7 @@ class TestEnvMat(unittest.TestCase, ArrayAPITest): - def test_compute_smooth_weight(self): + def test_compute_smooth_weight(self) -> None: d = xp.arange(10, dtype=xp.float64) w = compute_smooth_weight( d, diff --git a/source/tests/common/dpmodel/array_api/utils.py b/source/tests/common/dpmodel/array_api/utils.py index 7e422c2ead..1e03c5aae0 100644 --- a/source/tests/common/dpmodel/array_api/utils.py +++ b/source/tests/common/dpmodel/array_api/utils.py @@ -8,20 +8,20 @@ class ArrayAPITest: """Utils for array API tests.""" - def set_array_api_version(self, func): + def set_array_api_version(self, func) -> None: """Set the array API version for a function.""" set_array_api_strict_flags(api_version=func.array_api_version) - def assert_namespace_equal(self, a, b): + def assert_namespace_equal(self, a, b) -> None: """Assert two array has the same namespace.""" self.assertEqual( array_api_compat.array_namespace(a), array_api_compat.array_namespace(b) ) - def assert_dtype_equal(self, a, b): + def assert_dtype_equal(self, a, b) -> None: """Assert two array has the same dtype.""" self.assertEqual(a.dtype, b.dtype) - def assert_device_equal(self, a, b): + def assert_device_equal(self, a, b) -> None: """Assert two array has the same device.""" self.assertEqual(array_api_compat.device(a), array_api_compat.device(b)) diff --git a/source/tests/common/dpmodel/case_single_frame_with_nlist.py b/source/tests/common/dpmodel/case_single_frame_with_nlist.py index f674c7a3a9..7ec92a1de1 100644 --- a/source/tests/common/dpmodel/case_single_frame_with_nlist.py +++ b/source/tests/common/dpmodel/case_single_frame_with_nlist.py @@ -3,7 +3,7 @@ class TestCaseSingleFrameWithoutNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nf, self.nt = 1, 2 @@ -25,7 +25,7 @@ def setUp(self): class TestCaseSingleFrameWithNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nall = 4 @@ -80,7 +80,7 @@ def setUp(self): class TestCaseSingleFrameWithNlistWithVirtual: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 4 self.nall = 5 diff --git a/source/tests/common/dpmodel/test_descriptor_dpa1.py b/source/tests/common/dpmodel/test_descriptor_dpa1.py index f441895f15..fbbf6e3994 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa1.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa1.py @@ -16,12 +16,12 @@ class TestDescrptDPA1(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -37,7 +37,7 @@ def test_self_consistency( for ii in [0, 1, 4]: np.testing.assert_allclose(mm0[ii], mm1[ii]) - def test_multiple_frames(self): + def test_multiple_frames(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_descriptor_dpa2.py b/source/tests/common/dpmodel/test_descriptor_dpa2.py index 4f43b81566..7867fee874 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa2.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa2.py @@ -20,12 +20,12 @@ class TestDescrptDPA2(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_descriptor_hybrid.py b/source/tests/common/dpmodel/test_descriptor_hybrid.py index 14f4dd66be..b241f71d04 100644 --- a/source/tests/common/dpmodel/test_descriptor_hybrid.py +++ b/source/tests/common/dpmodel/test_descriptor_hybrid.py @@ -25,13 +25,13 @@ class TestDescrptHybrid(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: unittest.TestCase.setUp(self) TestCaseSingleFrameWithNlist.setUp(self) def test_get_parameters( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -74,7 +74,7 @@ def test_get_parameters( def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_descriptor_se_e2_a.py b/source/tests/common/dpmodel/test_descriptor_se_e2_a.py index 9028dac714..30ba47cb9d 100644 --- a/source/tests/common/dpmodel/test_descriptor_se_e2_a.py +++ b/source/tests/common/dpmodel/test_descriptor_se_e2_a.py @@ -16,12 +16,12 @@ class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_descriptor_se_r.py b/source/tests/common/dpmodel/test_descriptor_se_r.py index df0e65cde0..2e7b33b014 100644 --- a/source/tests/common/dpmodel/test_descriptor_se_r.py +++ b/source/tests/common/dpmodel/test_descriptor_se_r.py @@ -16,12 +16,12 @@ class TestDescrptSeR(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 1)) diff --git a/source/tests/common/dpmodel/test_descriptor_se_t.py b/source/tests/common/dpmodel/test_descriptor_se_t.py index dc78dfb3f5..fa240d1048 100644 --- a/source/tests/common/dpmodel/test_descriptor_se_t.py +++ b/source/tests/common/dpmodel/test_descriptor_se_t.py @@ -16,12 +16,12 @@ class TestDescrptSeT(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index 2a24cfdada..8ebd214865 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -21,10 +21,10 @@ class TestDPAtomicModel(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_methods(self): + def test_methods(self) -> None: ds = DescrptSeA( self.rcut, self.rcut_smth, @@ -55,7 +55,7 @@ def test_methods(self): def test_self_consistency( self, - ): + ) -> None: ds = DescrptSeA( self.rcut, self.rcut_smth, @@ -81,7 +81,7 @@ def test_self_consistency( np.testing.assert_allclose(ret0["energy"], ret1["energy"]) - def test_excl_consistency(self): + def test_excl_consistency(self) -> None: type_map = ["foo", "bar"] # test the case of exclusion @@ -146,13 +146,13 @@ def test_excl_consistency(self): class TestDPAtomicModelVirtualConsistency(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.case0 = TestCaseSingleFrameWithNlist() self.case1 = TestCaseSingleFrameWithNlistWithVirtual() self.case0.setUp() self.case1.setUp() - def test_virtual_consistency(self): + def test_virtual_consistency(self) -> None: nf, _, _ = self.case0.nlist.shape ds = DescrptSeA( self.case0.rcut, diff --git a/source/tests/common/dpmodel/test_dp_model.py b/source/tests/common/dpmodel/test_dp_model.py index b1310a835b..af4eea624d 100644 --- a/source/tests/common/dpmodel/test_dp_model.py +++ b/source/tests/common/dpmodel/test_dp_model.py @@ -23,12 +23,12 @@ class TestDPModelLower(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -52,7 +52,7 @@ def test_self_consistency( np.testing.assert_allclose(ret0["energy"], ret1["energy"]) np.testing.assert_allclose(ret0["energy_redu"], ret1["energy_redu"]) - def test_prec_consistency(self): + def test_prec_consistency(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( @@ -101,10 +101,10 @@ def test_prec_consistency(self): class TestDPModel(unittest.TestCase, TestCaseSingleFrameWithoutNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithoutNlist.setUp(self) - def test_prec_consistency(self): + def test_prec_consistency(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc = self.atype.shape ds = DescrptSeA( diff --git a/source/tests/common/dpmodel/test_env_mat.py b/source/tests/common/dpmodel/test_env_mat.py index a24380a8b3..39f5003d98 100644 --- a/source/tests/common/dpmodel/test_env_mat.py +++ b/source/tests/common/dpmodel/test_env_mat.py @@ -16,12 +16,12 @@ class TestEnvMat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/common/dpmodel/test_exclusion_mask.py b/source/tests/common/dpmodel/test_exclusion_mask.py index a6fdce317a..29d9f9b56b 100644 --- a/source/tests/common/dpmodel/test_exclusion_mask.py +++ b/source/tests/common/dpmodel/test_exclusion_mask.py @@ -14,7 +14,7 @@ class TestAtomExcludeMask(unittest.TestCase): - def test_build_type_exclude_mask(self): + def test_build_type_exclude_mask(self) -> None: nf = 2 nt = 3 exclude_types = [0, 2] @@ -38,10 +38,10 @@ def test_build_type_exclude_mask(self): # to be merged with the tf test case class TestPairExcludeMask(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_build_type_exclude_mask(self): + def test_build_type_exclude_mask(self) -> None: exclude_types = [[0, 1]] expected_mask = np.array( [ diff --git a/source/tests/common/dpmodel/test_fitting_invar_fitting.py b/source/tests/common/dpmodel/test_fitting_invar_fitting.py index 355bab4c1d..d82d3f43d8 100644 --- a/source/tests/common/dpmodel/test_fitting_invar_fitting.py +++ b/source/tests/common/dpmodel/test_fitting_invar_fitting.py @@ -20,12 +20,12 @@ class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_self_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel) @@ -72,7 +72,7 @@ def test_self_consistency( self.assertEqual(sel_set | exclude_set, set(range(self.nt))) self.assertEqual(sel_set & exclude_set, set()) - def test_mask(self): + def test_mask(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel) dd = ds.call(self.coord_ext, self.atype_ext, self.nlist) @@ -104,7 +104,7 @@ def test_mask(self): def test_self_exception( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape ds = DescrptSeA(self.rcut, self.rcut_smth, self.sel) @@ -156,7 +156,7 @@ def test_self_exception( ifn0(dd[0], atype, fparam=ifp, aparam=iap) self.assertIn("input aparam", context.exception) - def test_get_set(self): + def test_get_set(self) -> None: ifn0 = InvarFitting( "energy", self.nt, diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index b7bf310676..bc7238bc46 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -25,7 +25,7 @@ class TestWeightCalculation(unittest.TestCase): @patch("numpy.loadtxt") - def test_pairwise(self, mock_loadtxt): + def test_pairwise(self, mock_loadtxt) -> None: file_path = "dummy_path" mock_loadtxt.return_value = np.array( [ @@ -100,7 +100,7 @@ def test_pairwise(self, mock_loadtxt): class TestIntegration(unittest.TestCase): @patch("numpy.loadtxt") - def setUp(self, mock_loadtxt): + def setUp(self, mock_loadtxt) -> None: self.nloc = 3 self.nall = 4 self.nf, self.nt = 1, 2 @@ -162,7 +162,7 @@ def setUp(self, mock_loadtxt): ) self.md1 = DPZBLLinearEnergyAtomicModel.deserialize(self.md0.serialize()) - def test_self_consistency(self): + def test_self_consistency(self) -> None: ret0 = self.md0.forward_atomic(self.coord_ext, self.atype_ext, self.nlist) ret1 = self.md1.forward_atomic(self.coord_ext, self.atype_ext, self.nlist) np.testing.assert_allclose( diff --git a/source/tests/common/dpmodel/test_network.py b/source/tests/common/dpmodel/test_network.py index 381c542272..dde5992746 100644 --- a/source/tests/common/dpmodel/test_network.py +++ b/source/tests/common/dpmodel/test_network.py @@ -20,7 +20,7 @@ class TestNativeLayer(unittest.TestCase): - def test_serialize_deserize(self): + def test_serialize_deserize(self) -> None: for ( ni, no, @@ -49,7 +49,7 @@ def test_serialize_deserize(self): inp = np.arange(np.prod(inp_shap)).reshape(inp_shap) np.testing.assert_allclose(nl0.call(inp), nl1.call(inp)) - def test_shape_error(self): + def test_shape_error(self) -> None: self.w0 = np.full((2, 3), 3.0) self.b0 = np.full((2,), 4.0) self.b1 = np.full((3,), 4.0) @@ -81,7 +81,7 @@ def setUp(self) -> None: self.w1 = np.full((3, 4), 3.0) self.b1 = np.full((4,), 4.0) - def test_serialize(self): + def test_serialize(self) -> None: network = NativeNet( [ NativeLayer(2, 3).serialize(), @@ -106,7 +106,7 @@ def test_serialize(self): np.testing.assert_array_equal(jdata["layers"][0]["resnet"], True) np.testing.assert_array_equal(jdata["layers"][1]["resnet"], True) - def test_deserialize(self): + def test_deserialize(self) -> None: network = NativeNet.deserialize( { "layers": [ @@ -132,7 +132,7 @@ def test_deserialize(self): np.testing.assert_array_equal(network[0]["resnet"], True) np.testing.assert_array_equal(network[1]["resnet"], True) - def test_shape_error(self): + def test_shape_error(self) -> None: with self.assertRaises(ValueError) as context: NativeNet.deserialize( { @@ -154,7 +154,7 @@ def test_shape_error(self): class TestEmbeddingNet(unittest.TestCase): - def test_embedding_net(self): + def test_embedding_net(self) -> None: for ni, act, idt, prec in itertools.product( [1, 10], ["tanh", "none"], @@ -173,7 +173,7 @@ def test_embedding_net(self): class TestFittingNet(unittest.TestCase): - def test_fitting_net(self): + def test_fitting_net(self) -> None: for ni, no, act, idt, prec, bo in itertools.product( [1, 10], [1, 7], @@ -218,7 +218,7 @@ def setUp(self) -> None: ], } - def test_two_dim(self): + def test_two_dim(self) -> None: networks = NetworkCollection(ndim=2, ntypes=2) networks[(0, 0)] = self.network networks[(1, 1)] = self.network @@ -235,7 +235,7 @@ def test_two_dim(self): networks[(0, 0)].serialize(), networks.serialize()["networks"][0] ) - def test_one_dim(self): + def test_one_dim(self) -> None: networks = NetworkCollection(ndim=1, ntypes=2) networks[(0,)] = self.network with self.assertRaises(RuntimeError): @@ -250,7 +250,7 @@ def test_one_dim(self): networks[(0,)].serialize(), networks.serialize()["networks"][0] ) - def test_zero_dim(self): + def test_zero_dim(self) -> None: networks = NetworkCollection(ndim=0, ntypes=2) networks[()] = self.network networks.check_completeness() @@ -285,14 +285,14 @@ def setUp(self) -> None: self.filename = "test_dp_dpmodel.dp" self.filename_yaml = "test_dp_dpmodel.yaml" - def test_save_load_model(self): + def test_save_load_model(self) -> None: save_dp_model(self.filename, {"model": deepcopy(self.model_dict)}) model = load_dp_model(self.filename) np.testing.assert_equal(model["model"], self.model_dict) assert "software" in model assert "version" in model - def test_save_load_model_yaml(self): + def test_save_load_model_yaml(self) -> None: save_dp_model(self.filename_yaml, {"model": deepcopy(self.model_dict)}) model = load_dp_model(self.filename_yaml) np.testing.assert_equal(model["model"], self.model_dict) diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index 404232013b..baed11a961 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -22,7 +22,7 @@ class TestDPModelFormatNlist(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nall = 5 @@ -67,7 +67,7 @@ def setUp(self): type_map = ["foo", "bar"] self.md = EnergyModel(ds, ft, type_map=type_map) - def test_nlist_eq(self): + def test_nlist_eq(self) -> None: # n_nnei == nnei nlist = np.array( [ @@ -84,7 +84,7 @@ def test_nlist_eq(self): ) np.testing.assert_allclose(self.expected_nlist, nlist1) - def test_nlist_st(self): + def test_nlist_st(self) -> None: # n_nnei < nnei nlist = np.array( [ @@ -101,7 +101,7 @@ def test_nlist_st(self): ) np.testing.assert_allclose(self.expected_nlist, nlist1) - def test_nlist_lt(self): + def test_nlist_lt(self) -> None: # n_nnei > nnei nlist = np.array( [ @@ -123,7 +123,7 @@ def test_nlist_lt(self): class TestNeighList(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.nf = 3 self.nloc = 3 self.ns = 5 * 5 * 3 @@ -150,7 +150,7 @@ def setUp(self): ] ) - def test_build_notype(self): + def test_build_notype(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -171,7 +171,7 @@ def test_build_notype(self): np.sort(self.ref_nlist, axis=-1), ) - def test_build_type(self): + def test_build_type(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -193,7 +193,7 @@ def test_build_type(self): np.sort(np.split(self.ref_nlist, self.nsel, axis=-1)[ii], axis=-1), ) - def test_build_multiple_nlist(self): + def test_build_multiple_nlist(self) -> None: rcuts = [1.01, 2.01] nsels = [20, 80] ecoord, eatype, mapping = extend_coord_with_ghosts( @@ -232,7 +232,7 @@ def test_build_multiple_nlist(self): nlist2, ) - def test_extend_coord(self): + def test_extend_coord(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) diff --git a/source/tests/common/dpmodel/test_output_def.py b/source/tests/common/dpmodel/test_output_def.py index 03ceb67d01..ad7d42d5a1 100644 --- a/source/tests/common/dpmodel/test_output_def.py +++ b/source/tests/common/dpmodel/test_output_def.py @@ -25,14 +25,14 @@ def __init__( name: str, shape: list[int], atomic: bool = True, - ): + ) -> None: self.name = name self.shape = list(shape) self.atomic = atomic class TestDef(unittest.TestCase): - def test_model_output_def(self): + def test_model_output_def(self) -> None: defs = [ OutputVariableDef( "energy", @@ -495,7 +495,7 @@ def test_model_output_def(self): with self.assertRaises(ValueError): apply_operation(hession_vardef, OVO.DERV_R) - def test_no_raise_no_redu_deriv(self): + def test_no_raise_no_redu_deriv(self) -> None: OutputVariableDef( "energy", [1], @@ -504,7 +504,7 @@ def test_no_raise_no_redu_deriv(self): c_differentiable=False, ) - def test_raise_requires_r_deriv(self): + def test_raise_requires_r_deriv(self) -> None: with self.assertRaises(ValueError) as context: OutputVariableDef( "energy", @@ -514,11 +514,11 @@ def test_raise_requires_r_deriv(self): c_differentiable=True, ) - def test_raise_redu_not_atomic(self): + def test_raise_redu_not_atomic(self) -> None: with self.assertRaises(ValueError) as context: (OutputVariableDef("energy", [1], reducible=True, atomic=False),) - def test_hessian_not_reducible(self): + def test_hessian_not_reducible(self) -> None: with self.assertRaises(ValueError) as context: ( OutputVariableDef( @@ -531,7 +531,7 @@ def test_hessian_not_reducible(self): ), ) - def test_hessian_not_r_differentiable(self): + def test_hessian_not_r_differentiable(self) -> None: with self.assertRaises(ValueError) as context: ( OutputVariableDef( @@ -544,7 +544,7 @@ def test_hessian_not_r_differentiable(self): ), ) - def test_energy_magnetic(self): + def test_energy_magnetic(self) -> None: with self.assertRaises(ValueError) as context: ( OutputVariableDef( @@ -558,7 +558,7 @@ def test_energy_magnetic(self): ), ) - def test_inten_requires_redu(self): + def test_inten_requires_redu(self) -> None: with self.assertRaises(ValueError) as context: ( OutputVariableDef( @@ -573,7 +573,7 @@ def test_inten_requires_redu(self): ), ) - def test_model_decorator(self): + def test_model_decorator(self) -> None: nf = 2 nloc = 3 nall = 4 @@ -603,14 +603,14 @@ def call(self): ff = Foo() ff() - def test_model_decorator_keyerror(self): + def test_model_decorator_keyerror(self) -> None: nf = 2 nloc = 3 nall = 4 @model_check_output class Foo(NativeOP): - def __init__(self): + def __init__(self) -> None: super().__init__() def output_def(self): @@ -637,7 +637,7 @@ def call(self): ff() self.assertIn("energy_derv_r", context.exception) - def test_model_decorator_shapeerror(self): + def test_model_decorator_shapeerror(self) -> None: nf = 2 nloc = 3 nall = 4 @@ -648,7 +648,7 @@ def __init__( self, shape_rd=[nf, 1], shape_dr=[nf, nall, 1, 3], - ): + ) -> None: self.shape_rd, self.shape_dr = shape_rd, shape_dr def output_def(self): @@ -696,7 +696,7 @@ def call(self): ff() self.assertIn("not matching", context.exception) - def test_fitting_decorator(self): + def test_fitting_decorator(self) -> None: nf = 2 nloc = 3 nall = 4 @@ -723,7 +723,7 @@ def call(self): ff = Foo() ff() - def test_fitting_decorator_shapeerror(self): + def test_fitting_decorator_shapeerror(self) -> None: nf = 2 nloc = 3 @@ -732,7 +732,7 @@ class Foo(NativeOP): def __init__( self, shape=[nf, nloc, 1], - ): + ) -> None: self.shape = shape def output_def(self): @@ -764,7 +764,7 @@ def call(self): ff() self.assertIn("not matching", context.exception) - def test_check_var(self): + def test_check_var(self) -> None: var_def = VariableDef("foo", [2, 3], atomic=True) with self.assertRaises(ValueError) as context: check_var(np.zeros([2, 3, 4, 5, 6]), var_def) @@ -801,7 +801,7 @@ def test_check_var(self): self.assertIn("shape not matching", context.exception) check_var(np.zeros([2, 2, 8]), var_def) - def test_squeeze(self): + def test_squeeze(self) -> None: out_var = OutputVariableDef("foo", []) out_var.squeeze(0) self.assertEqual(out_var.shape, []) diff --git a/source/tests/common/dpmodel/test_pairtab_atomic_model.py b/source/tests/common/dpmodel/test_pairtab_atomic_model.py index e2866d3766..f8145c258a 100644 --- a/source/tests/common/dpmodel/test_pairtab_atomic_model.py +++ b/source/tests/common/dpmodel/test_pairtab_atomic_model.py @@ -51,7 +51,7 @@ def setUp(self, mock_loadtxt) -> None: # nframes=2, nloc=2, nnei=2 self.nlist = np.array([[[1, 2], [0, 2]], [[1, 2], [0, 3]]]) - def test_without_mask(self): + def test_without_mask(self) -> None: result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) @@ -59,7 +59,7 @@ def test_without_mask(self): np.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) - def test_with_mask(self): + def test_with_mask(self) -> None: self.nlist = np.array([[[1, -1], [0, 2]], [[1, 2], [0, 3]]]) result = self.model.forward_atomic( @@ -69,7 +69,7 @@ def test_with_mask(self): np.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) - def test_deserialize(self): + def test_deserialize(self) -> None: model1 = PairTabAtomicModel.deserialize(self.model.serialize()) np.testing.assert_allclose(self.model.tab_data, model1.tab_data) np.testing.assert_allclose(self.model.tab_info, model1.tab_info) diff --git a/source/tests/common/dpmodel/test_pairtab_preprocess.py b/source/tests/common/dpmodel/test_pairtab_preprocess.py index 7f4058dedd..93a61bc1f6 100644 --- a/source/tests/common/dpmodel/test_pairtab_preprocess.py +++ b/source/tests/common/dpmodel/test_pairtab_preprocess.py @@ -30,7 +30,7 @@ def setUp(self, mock_loadtxt) -> None: self.tab4 = PairTab(filename=file_path, rcut=0.03) self.tab5 = PairTab(filename=file_path, rcut=0.032) - def test_deserialize(self): + def test_deserialize(self) -> None: deserialized_tab = PairTab.deserialize(self.tab1.serialize()) np.testing.assert_allclose(self.tab1.vdata, deserialized_tab.vdata) np.testing.assert_allclose(self.tab1.rmin, deserialized_tab.rmin) @@ -42,7 +42,7 @@ def test_deserialize(self): np.testing.assert_allclose(self.tab1.tab_info, deserialized_tab.tab_info) np.testing.assert_allclose(self.tab1.tab_data, deserialized_tab.tab_data) - def test_preprocess(self): + def test_preprocess(self) -> None: np.testing.assert_allclose( self.tab1.vdata, np.array( @@ -139,7 +139,7 @@ def setUp(self, mock_loadtxt) -> None: self.tab3 = PairTab(filename=file_path, rcut=0.028) self.tab4 = PairTab(filename=file_path, rcut=0.033) - def test_preprocess(self): + def test_preprocess(self) -> None: np.testing.assert_allclose( self.tab1.vdata, np.array( @@ -210,7 +210,7 @@ def setUp(self, mock_loadtxt) -> None: self.tab3 = PairTab(filename=file_path, rcut=0.03) self.tab4 = PairTab(filename=file_path, rcut=0.037) - def test_preprocess(self): + def test_preprocess(self) -> None: np.testing.assert_allclose( self.tab1.vdata, np.array( diff --git a/source/tests/common/dpmodel/test_region.py b/source/tests/common/dpmodel/test_region.py index 1c0b6ff6f4..c21405e781 100644 --- a/source/tests/common/dpmodel/test_region.py +++ b/source/tests/common/dpmodel/test_region.py @@ -14,7 +14,7 @@ class TestRegion(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.cell = np.array( [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], ) @@ -22,7 +22,7 @@ def setUp(self): self.cell = np.tile(self.cell, [4, 5, 1, 1]) self.prec = 1e-8 - def test_inter_to_phys(self): + def test_inter_to_phys(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) inter = rng.normal(size=[4, 5, 3, 3]) phys = inter2phys(inter, self.cell) @@ -33,7 +33,7 @@ def test_inter_to_phys(self): phys[ii, jj], expected_phys, rtol=self.prec, atol=self.prec ) - def test_to_face_dist(self): + def test_to_face_dist(self) -> None: cell0 = self.cell[0][0] vol = np.linalg.det(cell0) # area of surfaces xy, xz, yz diff --git a/source/tests/common/dpmodel/test_update_sel.py b/source/tests/common/dpmodel/test_update_sel.py index 200b1e50f5..ff8313b204 100644 --- a/source/tests/common/dpmodel/test_update_sel.py +++ b/source/tests/common/dpmodel/test_update_sel.py @@ -31,7 +31,7 @@ def setUp(self) -> None: return super().setUp() @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_one_sel(self, sel_mock): + def test_update_one_sel(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") @@ -44,7 +44,7 @@ def test_update_one_sel(self, sel_mock): self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_hybrid(self, sel_mock): + def test_update_sel_hybrid(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] jdata = { @@ -75,7 +75,7 @@ def test_update_sel_hybrid(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel(self, sel_mock): + def test_update_sel(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] jdata = { @@ -90,7 +90,7 @@ def test_update_sel(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_auto(self, sel_mock): + def test_update_sel_atten_auto(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -117,7 +117,7 @@ def test_update_sel_atten_auto(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_int(self, sel_mock): + def test_update_sel_atten_int(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -144,7 +144,7 @@ def test_update_sel_atten_int(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_list(self, sel_mock): + def test_update_sel_atten_list(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -170,7 +170,7 @@ def test_update_sel_atten_list(self, sel_mock): jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_wrap_up_4(self): + def test_wrap_up_4(self) -> None: self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) diff --git a/source/tests/common/test_argument_parser.py b/source/tests/common/test_argument_parser.py index 2c67c1f6cb..4e39df8659 100644 --- a/source/tests/common/test_argument_parser.py +++ b/source/tests/common/test_argument_parser.py @@ -71,7 +71,7 @@ class TestParserOutput(unittest.TestCase): def attr_and_type_check( self, namespace: Namespace, mapping: "TEST_DICT", command: str, test_value: bool - ): + ) -> None: """Check attributes of `argparse.Manespace` types and values are as expected. First check for attribute existence, if it exists check its type and if type is @@ -131,7 +131,7 @@ def attr_and_type_check( f"{getattr(namespace, attribute)}", ) - def run_test(self, *, command: str, mapping: "TEST_DICT"): + def run_test(self, *, command: str, mapping: "TEST_DICT") -> None: """Run test first for specified arguments and then for default. Parameters @@ -195,16 +195,16 @@ def run_test(self, *, command: str, mapping: "TEST_DICT"): ) from e self.attr_and_type_check(namespace, mapping, command, test_value=False) - def test_no_command(self): + def test_no_command(self) -> None: """Test that parser outputs nothing when no command is input and does not fail.""" self.assertIsNone(parse_args([]).command) - def test_wrong_command(self): + def test_wrong_command(self) -> None: """Test that parser fails if no command is passed in.""" with self.assertRaises(SystemExit): parse_args(["RANDOM_WRONG_COMMAND"]) - def test_parser_log(self): + def test_parser_log(self) -> None: """Check if logging associated attributes are present in specified parsers.""" ARGS = { "--log-level": {"type": int, "value": "INFO", "expected": 20}, @@ -227,7 +227,7 @@ def test_parser_log(self): self.run_test(command=parser, mapping=args) - def test_parser_mpi(self): + def test_parser_mpi(self) -> None: """Check if mpi-log attribute is present in specified parsers.""" ARGS = {"--mpi-log": {"type": str, "value": "master"}} @@ -238,7 +238,7 @@ def test_parser_mpi(self): args = ARGS self.run_test(command=parser, mapping=args) - def test_parser_transfer(self): + def test_parser_transfer(self) -> None: """Test transfer subparser.""" ARGS = { "--raw-model": {"type": str, "value": "INFILE.PB"}, @@ -248,7 +248,7 @@ def test_parser_transfer(self): self.run_test(command="transfer", mapping=ARGS) - def test_parser_train_init_model(self): + def test_parser_train_init_model(self) -> None: """Test train init-model subparser.""" ARGS = { "INPUT": {"type": str, "value": "INFILE"}, @@ -258,7 +258,7 @@ def test_parser_train_init_model(self): self.run_test(command="train", mapping=ARGS) - def test_parser_train_restart(self): + def test_parser_train_restart(self) -> None: """Test train restart subparser.""" ARGS = { "INPUT": {"type": str, "value": "INFILE"}, @@ -268,7 +268,7 @@ def test_parser_train_restart(self): self.run_test(command="train", mapping=ARGS) - def test_parser_train_init_frz_model(self): + def test_parser_train_init_frz_model(self) -> None: """Test train init-frz-model subparser.""" ARGS = { "INPUT": {"type": str, "value": "INFILE"}, @@ -278,7 +278,7 @@ def test_parser_train_init_frz_model(self): self.run_test(command="train", mapping=ARGS) - def test_parser_train_finetune(self): + def test_parser_train_finetune(self) -> None: """Test train finetune subparser.""" ARGS = { "INPUT": {"type": str, "value": "INFILE"}, @@ -288,7 +288,7 @@ def test_parser_train_finetune(self): self.run_test(command="train", mapping=ARGS) - def test_parser_train_wrong_subcommand(self): + def test_parser_train_wrong_subcommand(self) -> None: """Test train with multiple subparsers.""" ARGS = { "INPUT": {"type": str, "value": "INFILE"}, @@ -299,7 +299,7 @@ def test_parser_train_wrong_subcommand(self): with self.assertRaises(SystemExit): self.run_test(command="train", mapping=ARGS) - def test_parser_freeze(self): + def test_parser_freeze(self) -> None: """Test freeze subparser.""" ARGS = { "--checkpoint-folder": {"type": str, "value": "FOLDER"}, @@ -309,7 +309,7 @@ def test_parser_freeze(self): self.run_test(command="freeze", mapping=ARGS) - def test_parser_test(self): + def test_parser_test(self) -> None: """Test test subparser.""" ARGS = { "--model": {"type": str, "value": "MODEL.PB"}, @@ -322,7 +322,7 @@ def test_parser_test(self): self.run_test(command="test", mapping=ARGS) - def test_parser_compress(self): + def test_parser_compress(self) -> None: """Test compress subparser.""" ARGS = { "--output": {"type": str, "value": "OUTFILE"}, @@ -334,7 +334,7 @@ def test_parser_compress(self): self.run_test(command="compress", mapping=ARGS) - def test_parser_doc(self): + def test_parser_doc(self) -> None: """Test doc subparser.""" ARGS = { "--out-type": {"type": str, "value": "rst"}, @@ -342,7 +342,7 @@ def test_parser_doc(self): self.run_test(command="doc-train-input", mapping=ARGS) - def test_parser_model_devi(self): + def test_parser_model_devi(self) -> None: """Test model-devi subparser.""" ARGS = { "--models": { @@ -357,7 +357,7 @@ def test_parser_model_devi(self): self.run_test(command="model-devi", mapping=ARGS) - def test_get_log_level(self): + def test_get_log_level(self) -> None: MAPPING = { "DEBUG": 10, "INFO": 20, diff --git a/source/tests/common/test_auto_batch_size.py b/source/tests/common/test_auto_batch_size.py index 9cb86d4923..e65e8d6f90 100644 --- a/source/tests/common/test_auto_batch_size.py +++ b/source/tests/common/test_auto_batch_size.py @@ -16,7 +16,7 @@ class CustomizedAutoBatchSizeCPU(AutoBatchSize): - def is_gpu_available(self): + def is_gpu_available(self) -> bool: return False def is_oom_error(self, e): @@ -24,7 +24,7 @@ def is_oom_error(self, e): class CustomizedAutoBatchSizeGPU(AutoBatchSize): - def is_gpu_available(self): + def is_gpu_available(self) -> bool: return True def is_oom_error(self, e): @@ -37,7 +37,7 @@ def oom(self, batch_size, start_index): raise OutOfMemoryError return batch_size, xp.zeros((batch_size, 2)) - def test_execute_oom_gpu(self): + def test_execute_oom_gpu(self) -> None: # initial batch size 256 = 128 * 2 auto_batch_size = CustomizedAutoBatchSizeGPU(256, 2.0) # no error - 128 @@ -61,7 +61,7 @@ def test_execute_oom_gpu(self): self.assertEqual(nb, 256) self.assertEqual(result.shape, (256, 2)) - def test_execute_oom_cpu(self): + def test_execute_oom_cpu(self) -> None: # initial batch size 256 = 128 * 2, nb is always 128 auto_batch_size = CustomizedAutoBatchSizeCPU(256, 2.0) nb, result = auto_batch_size.execute(self.oom, 1, 2) @@ -81,7 +81,7 @@ def test_execute_oom_cpu(self): self.assertEqual(result.shape, (128, 2)) @patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "256"}, clear=True) - def test_execute_oom_environment_variables(self): + def test_execute_oom_environment_variables(self) -> None: # DP_INFER_BATCH_SIZE = 256 = 128 * 2, nb is always 128 auto_batch_size = CustomizedAutoBatchSizeGPU(999, 2.0) nb, result = auto_batch_size.execute(self.oom, 1, 2) @@ -100,13 +100,13 @@ def test_execute_oom_environment_variables(self): self.assertEqual(nb, 128) self.assertEqual(result.shape, (128, 2)) - def test_execute_all(self): + def test_execute_all(self) -> None: dd1 = xp.zeros((10000, 2, 1)) auto_batch_size = CustomizedAutoBatchSizeGPU(256, 2.0) dd2 = auto_batch_size.execute_all(xp.asarray, 10000, 2, dd1) assert xp.all(dd1 == dd2) - def test_execute_all_dict(self): + def test_execute_all_dict(self) -> None: dd0 = xp.zeros((10000, 2, 1, 3, 4)) dd1 = xp.ones((10000, 2, 1, 3, 4)) auto_batch_size = CustomizedAutoBatchSizeGPU(256, 2.0) @@ -121,7 +121,7 @@ def func(dd1): assert xp.all(dd0 == dd2["foo"]) assert xp.all(dd1 == dd2["bar"]) - def test_execute_all_dict_oom(self): + def test_execute_all_dict_oom(self) -> None: # to reproduce #4036 when commenting "if n_batch == 0: continue" dd0 = xp.zeros((10, 2, 1, 3, 4)) dd1 = xp.ones((10, 2, 1, 3, 4)) diff --git a/source/tests/common/test_common.py b/source/tests/common/test_common.py index fe9054d6ad..dad224b2d8 100644 --- a/source/tests/common/test_common.py +++ b/source/tests/common/test_common.py @@ -14,7 +14,7 @@ class TestGetXPPrecision(unittest.TestCase): - def test(self): + def test(self) -> None: aa = np.zeros(3) xp = array_api_compat.array_namespace(aa) self.assertEqual(get_xp_precision(xp, "float16"), xp.float16) diff --git a/source/tests/common/test_doc_train_input.py b/source/tests/common/test_doc_train_input.py index fe6c8ab1ac..3405c3b774 100644 --- a/source/tests/common/test_doc_train_input.py +++ b/source/tests/common/test_doc_train_input.py @@ -23,20 +23,20 @@ class TestDocTrainInput(unittest.TestCase): def multi_task(self): return self.param[0] - def test_rst(self): + def test_rst(self) -> None: f = io.StringIO() with redirect_stdout(f): doc_train_input(out_type="rst", multi_task=self.multi_task) self.assertNotEqual(f.getvalue(), "") - def test_json(self): + def test_json(self) -> None: f = io.StringIO() with redirect_stdout(f): doc_train_input(out_type="json", multi_task=self.multi_task) # validate json json.loads(f.getvalue()) - def test_json_schema(self): + def test_json_schema(self) -> None: f = io.StringIO() with redirect_stdout(f): doc_train_input(out_type="json_schema", multi_task=self.multi_task) diff --git a/source/tests/common/test_econf_embd.py b/source/tests/common/test_econf_embd.py index 242ea9ca65..6c58420e7a 100644 --- a/source/tests/common/test_econf_embd.py +++ b/source/tests/common/test_econf_embd.py @@ -10,7 +10,7 @@ class TestEConfEmbd(unittest.TestCase): - def test_fe(self): + def test_fe(self) -> None: res = make_econf_embedding(["Fe"], flatten=False)["Fe"] expected_res = { (1, "s"): [2], @@ -35,14 +35,14 @@ def test_fe(self): } self.assertDictEqual({kk: list(vv) for kk, vv in res.items()}, expected_res) - def test_fe_flatten(self): + def test_fe_flatten(self) -> None: res = make_econf_embedding(["Fe"], flatten=True)["Fe"] # fmt: off expected_res = [2,2,2,2,2,2,2,2,2,2,1,1,1,1,2,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0] # fmt: on self.assertEqual(list(res), expected_res) - def test_fe_spin(self): + def test_fe_spin(self) -> None: res = make_econf_embedding(["Fe"], flatten=True) res = transform_to_spin_rep(res)["Fe"] # fmt: off @@ -50,7 +50,7 @@ def test_fe_spin(self): # fmt: on self.assertEqual(list(res), expected_res) - def test_dict(self): + def test_dict(self) -> None: res = electronic_configuration_embedding["Fe"] # fmt: off expected_res = [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,-1,1,-1,1,-1,1,-1,1,1,1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1] diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 246e767f01..8236040f27 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -67,7 +67,7 @@ class TestExamples(unittest.TestCase): - def test_arguments(self): + def test_arguments(self) -> None: for fn in input_files + input_files_multi: multi_task = fn in input_files_multi fn = str(fn) diff --git a/source/tests/common/test_gui.py b/source/tests/common/test_gui.py index 25fd7e6651..2b2fb85191 100644 --- a/source/tests/common/test_gui.py +++ b/source/tests/common/test_gui.py @@ -7,5 +7,5 @@ class TestDPGUI(unittest.TestCase): - def test_dpgui_entrypoints(self): + def test_dpgui_entrypoints(self) -> None: self.assertTrue(len(generate_dpgui_templates()) > 0) diff --git a/source/tests/common/test_out_stat.py b/source/tests/common/test_out_stat.py index c0cfc25071..c175d7c643 100644 --- a/source/tests/common/test_out_stat.py +++ b/source/tests/common/test_out_stat.py @@ -34,7 +34,7 @@ def setUp(self) -> None: return super().setUp() - def test_compute_stats_from_redu(self): + def test_compute_stats_from_redu(self) -> None: bias, std = compute_stats_from_redu(self.output_redu, self.natoms) np.testing.assert_allclose(bias, self.mean, rtol=1e-7) reference_std = np.array( @@ -58,7 +58,7 @@ def test_compute_stats_from_redu(self): rtol=1e-7, ) - def test_compute_stats_from_redu_with_assigned_bias(self): + def test_compute_stats_from_redu_with_assigned_bias(self) -> None: assigned_bias = np.full_like(self.mean, np.nan) assigned_bias[0] = self.mean[0] bias, std = compute_stats_from_redu( @@ -89,7 +89,7 @@ def test_compute_stats_from_redu_with_assigned_bias(self): rtol=1e-7, ) - def test_compute_stats_from_atomic(self): + def test_compute_stats_from_atomic(self) -> None: bias, std = compute_stats_from_atomic(self.output, self.atype) np.testing.assert_allclose(bias, self.mean) reference_std = np.array( diff --git a/source/tests/common/test_path.py b/source/tests/common/test_path.py index 7dcb3a031c..4eb74edc53 100644 --- a/source/tests/common/test_path.py +++ b/source/tests/common/test_path.py @@ -16,7 +16,7 @@ class PathTest: path: DPPath - def test_numpy(self): + def test_numpy(self) -> None: numpy_path = self.path / "testcase" arr1 = np.ones(3) self.assertFalse(numpy_path.is_file()) @@ -25,7 +25,7 @@ def test_numpy(self): arr2 = numpy_path.load_numpy() np.testing.assert_array_equal(arr1, arr2) - def test_dir(self): + def test_dir(self) -> None: dir_path = self.path / "testcase" self.assertFalse(dir_path.is_dir()) dir_path.mkdir() @@ -33,21 +33,21 @@ def test_dir(self): class TestOSPath(PathTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.tempdir = tempfile.TemporaryDirectory() self.path = DPPath(self.tempdir.name, "a") - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() class TestH5Path(PathTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.tempdir = tempfile.TemporaryDirectory() h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) with h5py.File(h5file, "w") as f: pass self.path = DPPath(h5file, "a") - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() diff --git a/source/tests/common/test_sel_idx.py b/source/tests/common/test_sel_idx.py index d6630e3e83..a9a97046f9 100644 --- a/source/tests/common/test_sel_idx.py +++ b/source/tests/common/test_sel_idx.py @@ -9,7 +9,7 @@ class TestSelIdx(unittest.TestCase): - def test_add(self): + def test_add(self) -> None: atom_type = np.array([0, 1, 2, 2, 1, 0], dtype=int) type_sel = np.array([1, 0], dtype=int) idx_map = select_idx_map(atom_type, type_sel) diff --git a/source/tests/common/test_spin.py b/source/tests/common/test_spin.py index c3bca50b09..5249c4f3d1 100644 --- a/source/tests/common/test_spin.py +++ b/source/tests/common/test_spin.py @@ -12,7 +12,7 @@ class SpinTest(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: type_map_1 = ["H", "O"] self.use_spin_1 = [False, False] self.virtual_scale_1 = [0.1, 0.1] @@ -105,7 +105,7 @@ def setUp(self): [9, 9], ] - def test_ntypes(self): + def test_ntypes(self) -> None: self.assertEqual(self.spin_1.get_ntypes_real(), 2) self.assertEqual(self.spin_1.get_ntypes_spin(), 0) self.assertEqual(self.spin_1.get_ntypes_real_and_spin(), 2) @@ -121,12 +121,12 @@ def test_ntypes(self): self.assertEqual(self.spin_3.get_ntypes_real_and_spin(), 6) self.assertEqual(self.spin_3.get_ntypes_input(), 10) - def test_use_spin(self): + def test_use_spin(self) -> None: np.testing.assert_allclose(self.spin_1.get_use_spin(), self.use_spin_1) np.testing.assert_allclose(self.spin_2.get_use_spin(), self.use_spin_2) np.testing.assert_allclose(self.spin_3.get_use_spin(), self.use_spin_3) - def test_mask(self): + def test_mask(self) -> None: np.testing.assert_allclose( self.spin_1.get_virtual_scale_mask(), self.expect_virtual_scale_mask_1 ) @@ -137,7 +137,7 @@ def test_mask(self): self.spin_3.get_virtual_scale_mask(), self.expect_virtual_scale_mask_3 ) - def test_exclude_types(self): + def test_exclude_types(self) -> None: self.assertEqual( sorted(self.spin_1.get_pair_exclude_types()), sorted(self.expect_pair_exclude_types_1), @@ -151,7 +151,7 @@ def test_exclude_types(self): sorted(self.expect_pair_exclude_types_3), ) - def test_virtual_scale_consistence(self): + def test_virtual_scale_consistence(self) -> None: np.testing.assert_allclose( self.spin_3.get_virtual_scale(), self.spin_3_float.get_virtual_scale() ) diff --git a/source/tests/common/test_type_index_map.py b/source/tests/common/test_type_index_map.py index cd7e761ac2..bfe8bc20f5 100644 --- a/source/tests/common/test_type_index_map.py +++ b/source/tests/common/test_type_index_map.py @@ -9,7 +9,7 @@ class TestTypeIndexMap(unittest.TestCase): - def test_get_index_between_two_maps(self): + def test_get_index_between_two_maps(self) -> None: tm_1 = [ "Al", "F", @@ -88,7 +88,7 @@ def test_get_index_between_two_maps(self): self.assertEqual(expected_map, result_map) self.assertEqual(expected_has_new, result_has_new) - def test_map_exclude_types(self): + def test_map_exclude_types(self) -> None: old_tm = [ "Al", "F", diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 734486becb..358ac8d542 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -104,10 +104,10 @@ class CommonTest(ABC): atol = 1e-10 """Absolute tolerance for comparing the return value. Override for float32.""" - def setUp(self): + def setUp(self) -> None: self.unique_id = uuid4().hex - def reset_unique_id(self): + def reset_unique_id(self) -> None: self.unique_id = uuid4().hex def init_backend_cls(self, cls) -> Any: @@ -303,7 +303,7 @@ def get_reference_ret_serialization(self, ref: RefBackend): return self.get_array_api_strict_ret_serialization_from_cls(obj) raise ValueError("No available reference") - def test_tf_consistent_with_ref(self): + def test_tf_consistent_with_ref(self) -> None: """Test whether TF and reference are consistent.""" if self.skip_tf: self.skipTest("Unsupported backend") @@ -333,7 +333,7 @@ def test_tf_consistent_with_ref(self): ) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" - def test_tf_self_consistent(self): + def test_tf_self_consistent(self) -> None: """Test whether TF is self consistent.""" if self.skip_tf: self.skipTest("Unsupported backend") @@ -349,7 +349,7 @@ def test_tf_self_consistent(self): assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_dp_consistent_with_ref(self): + def test_dp_consistent_with_ref(self) -> None: """Test whether DP and reference are consistent.""" if self.skip_dp: self.skipTest("Unsupported backend") @@ -370,7 +370,7 @@ def test_dp_consistent_with_ref(self): assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_dp_self_consistent(self): + def test_dp_self_consistent(self) -> None: """Test whether DP is self consistent.""" if self.skip_dp: self.skipTest("Unsupported backend") @@ -386,7 +386,7 @@ def test_dp_self_consistent(self): else: self.assertEqual(rr1, rr2) - def test_pt_consistent_with_ref(self): + def test_pt_consistent_with_ref(self) -> None: """Test whether PT and reference are consistent.""" if self.skip_pt: self.skipTest("Unsupported backend") @@ -409,7 +409,7 @@ def test_pt_consistent_with_ref(self): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" - def test_pt_self_consistent(self): + def test_pt_self_consistent(self) -> None: """Test whether PT is self consistent.""" if self.skip_pt: self.skipTest("Unsupported backend") @@ -425,7 +425,7 @@ def test_pt_self_consistent(self): else: self.assertEqual(rr1, rr2) - def test_jax_consistent_with_ref(self): + def test_jax_consistent_with_ref(self) -> None: """Test whether JAX and reference are consistent.""" if self.skip_jax: self.skipTest("Unsupported backend") @@ -443,7 +443,7 @@ def test_jax_consistent_with_ref(self): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" - def test_jax_self_consistent(self): + def test_jax_self_consistent(self) -> None: """Test whether JAX is self consistent.""" if self.skip_jax: self.skipTest("Unsupported backend") @@ -460,7 +460,7 @@ def test_jax_self_consistent(self): self.assertEqual(rr1, rr2) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_array_api_strict_consistent_with_ref(self): + def test_array_api_strict_consistent_with_ref(self) -> None: """Test whether array_api_strict and reference are consistent.""" if self.skip_array_api_strict: self.skipTest("Unsupported backend") @@ -479,7 +479,7 @@ def test_array_api_strict_consistent_with_ref(self): assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_array_api_strict_self_consistent(self): + def test_array_api_strict_self_consistent(self) -> None: """Test whether array_api_strict is self consistent.""" if self.skip_array_api_strict: self.skipTest("Unsupported backend") diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py index 3d80e310d0..8be219f5ea 100644 --- a/source/tests/consistent/descriptor/test_dpa1.py +++ b/source/tests/consistent/descriptor/test_dpa1.py @@ -292,7 +292,7 @@ def skip_tf(self) -> bool: args = descrpt_se_atten_args().append(Argument("ntypes", int, optional=False)) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_dpa2.py b/source/tests/consistent/descriptor/test_dpa2.py index 17c55db368..72c0967a78 100644 --- a/source/tests/consistent/descriptor/test_dpa2.py +++ b/source/tests/consistent/descriptor/test_dpa2.py @@ -290,7 +290,7 @@ def skip_tf(self) -> bool: array_api_strict_class = DescrptDPA2Strict args = descrpt_dpa2_args().append(Argument("ntypes", int, optional=False)) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_hybrid.py b/source/tests/consistent/descriptor/test_hybrid.py index c43652b498..2cbbd92d84 100644 --- a/source/tests/consistent/descriptor/test_hybrid.py +++ b/source/tests/consistent/descriptor/test_hybrid.py @@ -87,7 +87,7 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_se_atten_v2.py b/source/tests/consistent/descriptor/test_se_atten_v2.py index f4a8119ca3..af041c6cce 100644 --- a/source/tests/consistent/descriptor/test_se_atten_v2.py +++ b/source/tests/consistent/descriptor/test_se_atten_v2.py @@ -255,7 +255,7 @@ def skip_array_api_strict(self) -> bool: array_api_strict_class = DescrptSeAttenV2Strict args = descrpt_se_atten_args().append(Argument("ntypes", int, optional=False)) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_se_e2_a.py b/source/tests/consistent/descriptor/test_se_e2_a.py index 286703e21d..a3ed19e8f3 100644 --- a/source/tests/consistent/descriptor/test_se_e2_a.py +++ b/source/tests/consistent/descriptor/test_se_e2_a.py @@ -140,7 +140,7 @@ def skip_array_api_strict(self) -> bool: array_api_strict_class = DescrptSeAArrayAPIStrict args = descrpt_se_a_args() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_se_r.py b/source/tests/consistent/descriptor/test_se_r.py index e851106c44..3420c5592f 100644 --- a/source/tests/consistent/descriptor/test_se_r.py +++ b/source/tests/consistent/descriptor/test_se_r.py @@ -121,7 +121,7 @@ def skip_array_api_strict(self) -> bool: array_api_strict_class = DescrptSeRArrayAPIStrict args = descrpt_se_r_args() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_se_t.py b/source/tests/consistent/descriptor/test_se_t.py index 1e6110705a..56655cabe1 100644 --- a/source/tests/consistent/descriptor/test_se_t.py +++ b/source/tests/consistent/descriptor/test_se_t.py @@ -111,7 +111,7 @@ def skip_tf(self) -> bool: array_api_strict_class = DescrptSeTStrict args = descrpt_se_t_args() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/descriptor/test_se_t_tebd.py b/source/tests/consistent/descriptor/test_se_t_tebd.py index 4712c28e53..bb4a5db6e7 100644 --- a/source/tests/consistent/descriptor/test_se_t_tebd.py +++ b/source/tests/consistent/descriptor/test_se_t_tebd.py @@ -156,7 +156,7 @@ def skip_tf(self) -> bool: array_api_strict_class = DescrptSeTTebdStrict args = descrpt_se_e3_tebd_args().append(Argument("ntypes", int, optional=False)) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 088cb30238..396ee2d492 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -95,7 +95,7 @@ def skip_pt(self) -> bool: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/fitting/test_dos.py b/source/tests/consistent/fitting/test_dos.py index 0649681ccb..a098f04d27 100644 --- a/source/tests/consistent/fitting/test_dos.py +++ b/source/tests/consistent/fitting/test_dos.py @@ -112,7 +112,7 @@ def skip_array_api_strict(self) -> bool: array_api_strict_class = DOSFittingStrict args = fitting_dos() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py index 7be0382b16..1ef846dbcc 100644 --- a/source/tests/consistent/fitting/test_ener.py +++ b/source/tests/consistent/fitting/test_ener.py @@ -122,7 +122,7 @@ def skip_array_api_strict(self) -> bool: array_api_strict_class = EnerFittingStrict args = fitting_ener() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/fitting/test_polar.py b/source/tests/consistent/fitting/test_polar.py index 12f13d1e08..3ad3e00ceb 100644 --- a/source/tests/consistent/fitting/test_polar.py +++ b/source/tests/consistent/fitting/test_polar.py @@ -95,7 +95,7 @@ def skip_pt(self) -> bool: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/fitting/test_property.py b/source/tests/consistent/fitting/test_property.py index d8a56447a4..3abd672c88 100644 --- a/source/tests/consistent/fitting/test_property.py +++ b/source/tests/consistent/fitting/test_property.py @@ -115,7 +115,7 @@ def skip_tf(self) -> bool: array_api_strict_class = PropertyFittingStrict args = fitting_property() - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index e8e751cee6..bc7d3e7af2 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -64,7 +64,7 @@ def save_data_to_model(self, model_file: str, data: dict) -> None: out_hook = out_backend.deserialize_hook out_hook(model_file, data) - def tearDown(self): + def tearDown(self) -> None: prefix = "test_consistent_io_" + self.__class__.__name__.lower() for ii in Path(".").glob(prefix + ".*"): if Path(ii).is_file(): @@ -73,7 +73,7 @@ def tearDown(self): shutil.rmtree(ii) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_data_equal(self): + def test_data_equal(self) -> None: prefix = "test_consistent_io_" + self.__class__.__name__.lower() for backend_name, suffix_idx in ( ("tensorflow", 0) if not DP_TEST_TF2_ONLY else ("jax", 0), @@ -107,7 +107,7 @@ def test_data_equal(self): reference_data.pop(kk, None) np.testing.assert_equal(data, reference_data) - def test_deep_eval(self): + def test_deep_eval(self) -> None: self.coords = np.array( [ 12.83, @@ -218,7 +218,7 @@ def test_deep_eval(self): class TestDeepPot(unittest.TestCase, IOTest): - def setUp(self): + def setUp(self) -> None: model_def_script = { "type_map": ["O", "H"], "descriptor": { @@ -255,12 +255,12 @@ def setUp(self): "model_def_script": model_def_script, } - def tearDown(self): + def tearDown(self) -> None: IOTest.tearDown(self) class TestDeepPotFparamAparam(unittest.TestCase, IOTest): - def setUp(self): + def setUp(self) -> None: model_def_script = { "type_map": ["O", "H"], "descriptor": { @@ -299,5 +299,5 @@ def setUp(self): "model_def_script": model_def_script, } - def tearDown(self): + def tearDown(self) -> None: IOTest.tearDown(self) diff --git a/source/tests/consistent/model/test_dpa1.py b/source/tests/consistent/model/test_dpa1.py index 32b523e1ba..774c624ac7 100644 --- a/source/tests/consistent/model/test_dpa1.py +++ b/source/tests/consistent/model/test_dpa1.py @@ -109,7 +109,7 @@ def get_reference_backend(self): raise ValueError("No available reference") @property - def skip_jax(self): + def skip_jax(self) -> bool: return not INSTALLED_JAX def pass_data_to_cls(self, cls, data) -> Any: @@ -123,7 +123,7 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_jax(data) return cls(**data, **self.additional_data) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 5d0253c5e8..45d85861fd 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -132,7 +132,7 @@ def skip_tf(self): ) @property - def skip_jax(self): + def skip_jax(self) -> bool: return not INSTALLED_JAX def pass_data_to_cls(self, cls, data) -> Any: @@ -146,7 +146,7 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_jax(data) return cls(**data, **self.additional_data) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 @@ -323,12 +323,12 @@ def get_reference_backend(self): raise ValueError("No available reference") @property - def skip_tf(self): + def skip_tf(self) -> bool: # TF does not have lower interface return True @property - def skip_jax(self): + def skip_jax(self) -> bool: return not INSTALLED_JAX def pass_data_to_cls(self, cls, data) -> Any: @@ -342,7 +342,7 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_jax(data) return cls(**data, **self.additional_data) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/model/test_frozen.py b/source/tests/consistent/model/test_frozen.py index f11a11914b..ff7d651e7e 100644 --- a/source/tests/consistent/model/test_frozen.py +++ b/source/tests/consistent/model/test_frozen.py @@ -44,14 +44,14 @@ dp_model = "deeppot_for_consistent_frozen.dp" -def setUpModule(): +def setUpModule() -> None: case = get_cases()["se_e2_a"] case.get_model(".dp", dp_model) case.get_model(".pb", tf_model) case.get_model(".pth", pt_model) -def tearDownModule(): +def tearDownModule() -> None: for model_file in (dp_model, pt_model, tf_model): try: os.remove(model_file) @@ -78,10 +78,10 @@ def data(self) -> dict: pt_class = FrozenModelPT args = model_args() - def skip_dp(self): + def skip_dp(self) -> bool: return True - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/model/test_zbl_ener.py b/source/tests/consistent/model/test_zbl_ener.py index a63543ab74..6fb44a59ed 100644 --- a/source/tests/consistent/model/test_zbl_ener.py +++ b/source/tests/consistent/model/test_zbl_ener.py @@ -111,11 +111,11 @@ def get_reference_backend(self): raise ValueError("No available reference") @property - def skip_tf(self): + def skip_tf(self) -> bool: return True @property - def skip_jax(self): + def skip_jax(self) -> bool: return not INSTALLED_JAX def pass_data_to_cls(self, cls, data) -> Any: @@ -129,7 +129,7 @@ def pass_data_to_cls(self, cls, data) -> Any: return get_model_jax(data) return cls(**data, **self.additional_data) - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/consistent/test_activation.py b/source/tests/consistent/test_activation.py index 2af545fc35..2368b6c473 100644 --- a/source/tests/consistent/test_activation.py +++ b/source/tests/consistent/test_activation.py @@ -43,7 +43,7 @@ tuple([x.capitalize() for x in VALID_ACTIVATION]), ) class TestActivationFunctionConsistent(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.activation,) = self.param self.random_input = np.random.default_rng(GLOBAL_SEED).normal( scale=10, size=(10, 10) @@ -51,7 +51,7 @@ def setUp(self): self.ref = get_activation_fn_dp(self.activation)(self.random_input) @unittest.skipUnless(INSTALLED_TF, "TensorFlow is not installed") - def test_tf_consistent_with_ref(self): + def test_tf_consistent_with_ref(self) -> None: if INSTALLED_TF: place_holder = tf.placeholder(tf.float64, self.random_input.shape) t_test = get_activation_fn_tf(self.activation)(place_holder) @@ -60,7 +60,7 @@ def test_tf_consistent_with_ref(self): np.testing.assert_allclose(self.ref, test, atol=1e-10) @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") - def test_pt_consistent_with_ref(self): + def test_pt_consistent_with_ref(self) -> None: if INSTALLED_PT: test = torch_to_numpy( ActivationFn_pt(self.activation)(to_torch_tensor(self.random_input)) @@ -70,7 +70,7 @@ def test_pt_consistent_with_ref(self): @unittest.skipUnless( sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8" ) - def test_arary_api_strict(self): + def test_arary_api_strict(self) -> None: import array_api_strict as xp input = xp.asarray(self.random_input) @@ -78,7 +78,7 @@ def test_arary_api_strict(self): np.testing.assert_allclose(self.ref, to_numpy_array(test), atol=1e-10) @unittest.skipUnless(INSTALLED_JAX, "JAX is not installed") - def test_jax_consistent_with_ref(self): + def test_jax_consistent_with_ref(self) -> None: input = jnp.from_dlpack(self.random_input) test = get_activation_fn_dp(self.activation)(input) self.assertTrue(isinstance(test, jnp.ndarray)) diff --git a/source/tests/consistent/test_neighbor_stat.py b/source/tests/consistent/test_neighbor_stat.py index 55181a6903..573e367267 100644 --- a/source/tests/consistent/test_neighbor_stat.py +++ b/source/tests/consistent/test_neighbor_stat.py @@ -38,16 +38,16 @@ def gen_sys(nframes): class TestNeighborStat(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: data0 = gen_sys(1) sys0 = dpdata.LabeledSystem() sys0.data = data0 sys0.to_deepmd_npy("system_0", set_size=1) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree("system_0") - def run_neighbor_stat(self, backend): + def run_neighbor_stat(self, backend) -> None: for rcut in (0.0, 1.0, 2.0, 4.0): for mixed_type in (True, False): with self.subTest(rcut=rcut, mixed_type=mixed_type): @@ -74,16 +74,16 @@ def run_neighbor_stat(self, backend): np.testing.assert_array_equal(max_nbor_size, ret) @unittest.skipUnless(INSTALLED_TF, "tensorflow is not installed") - def test_neighbor_stat_tf(self): + def test_neighbor_stat_tf(self) -> None: self.run_neighbor_stat("tensorflow") @unittest.skipUnless(INSTALLED_PT, "pytorch is not installed") - def test_neighbor_stat_pt(self): + def test_neighbor_stat_pt(self) -> None: self.run_neighbor_stat("pytorch") - def test_neighbor_stat_dp(self): + def test_neighbor_stat_dp(self) -> None: self.run_neighbor_stat("numpy") @unittest.skipUnless(INSTALLED_JAX, "jax is not installed") - def test_neighbor_stat_jax(self): + def test_neighbor_stat_jax(self) -> None: self.run_neighbor_stat("jax") diff --git a/source/tests/consistent/test_type_embedding.py b/source/tests/consistent/test_type_embedding.py index 10cbd1837d..1c56abea0c 100644 --- a/source/tests/consistent/test_type_embedding.py +++ b/source/tests/consistent/test_type_embedding.py @@ -100,7 +100,7 @@ def additional_data(self) -> dict: "type_map": ["O", "H"] if use_econf_tebd else None, } - def setUp(self): + def setUp(self) -> None: CommonTest.setUp(self) self.ntypes = 2 diff --git a/source/tests/infer/case.py b/source/tests/infer/case.py index 4a5ce638d6..8b8481a194 100644 --- a/source/tests/infer/case.py +++ b/source/tests/infer/case.py @@ -135,7 +135,7 @@ class Case: The path to the test case file. """ - def __init__(self, filename: str): + def __init__(self, filename: str) -> None: with open(filename) as file: config = yaml.safe_load(file) self.key = config["key"] diff --git a/source/tests/infer/test_models.py b/source/tests/infer/test_models.py index 2b0f292046..a6cde3206c 100644 --- a/source/tests/infer/test_models.py +++ b/source/tests/infer/test_models.py @@ -34,24 +34,24 @@ class TestDeepPot(unittest.TestCase): # moved from tests/tf/test_deeppot_a.py @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: key, extension = cls.param cls.case = get_cases()[key] cls.model_name = cls.case.get_model(extension) cls.dp = DeepEval(cls.model_name) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: cls.dp = None - def setUp(self): + def setUp(self) -> None: key, extension = self.param if key == "se_e2_r" and extension == ".pth": self.skipTest( reason="se_e2_r type_one_side is not supported for PyTorch models" ) - def test_attrs(self): + def test_attrs(self) -> None: assert isinstance(self.dp, DeepPot) self.assertEqual(self.dp.get_ntypes(), self.case.ntypes) self.assertAlmostEqual( @@ -61,7 +61,7 @@ def test_attrs(self): self.assertEqual(self.dp.get_dim_fparam(), self.case.dim_fparam) self.assertEqual(self.dp.get_dim_aparam(), self.case.dim_aparam) - def test_1frame(self): + def test_1frame(self) -> None: for ii, result in enumerate(self.case.results): ee, ff, vv = self.dp.eval( result.coord, @@ -99,7 +99,7 @@ def test_1frame(self): err_msg=f"Result {ii} virial", ) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: for ii, result in enumerate(self.case.results): ee, ff, vv, ae, av = self.dp.eval( result.coord, @@ -151,7 +151,7 @@ def test_1frame_atm(self): err_msg=f"Result {ii} virial", ) - def test_descriptor(self): + def test_descriptor(self) -> None: _, extension = self.param for ii, result in enumerate(self.case.results): if result.descriptor is None: @@ -160,7 +160,7 @@ def test_descriptor(self): expected_descpt = result.descriptor np.testing.assert_almost_equal(descpt.ravel(), expected_descpt.ravel()) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: for ii, result in enumerate(self.case.results): coords2 = np.concatenate((result.coord, result.coord)) if result.box is not None: @@ -209,7 +209,7 @@ def test_2frame_atm(self): vv.ravel(), expected_sv.ravel(), default_places ) - def test_zero_input(self): + def test_zero_input(self) -> None: _, extension = self.param if extension == ".pb": from deepmd.tf.env import ( @@ -243,7 +243,7 @@ def test_zero_input(self): np.testing.assert_almost_equal(ee.ravel(), 0, default_places) np.testing.assert_almost_equal(vv.ravel(), 0, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -280,7 +280,7 @@ def test_ase(self): err_msg=f"Result {ii} energy", ) - def test_dpdata_driver(self): + def test_dpdata_driver(self) -> None: if self.case.dim_fparam or self.case.dim_aparam: self.skipTest("fparam and aparam not supported") @@ -323,7 +323,7 @@ def test_dpdata_driver(self): err_msg=f"Result {ii} virial", ) - def test_model_script_def(self): + def test_model_script_def(self) -> None: if self.case.model_def_script is not None: self.assertDictEqual( self.case.model_def_script, self.dp.get_model_def_script() @@ -336,7 +336,7 @@ def test_model_script_def(self): ) class TestDeepPotNeighborList(TestDeepPot): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: key, extension = cls.param cls.case = get_cases()[key] model_name = cls.case.get_model(extension) @@ -348,9 +348,9 @@ def setUpClass(cls): ) @unittest.skip("multiple frames not supported") - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: pass @unittest.skip("Zero atoms not supported") - def test_zero_input(self): + def test_zero_input(self) -> None: pass diff --git a/source/tests/pt/model/test_atomic_model_atomic_stat.py b/source/tests/pt/model/test_atomic_model_atomic_stat.py index 6a21fc6e5a..e6e873e6c1 100644 --- a/source/tests/pt/model/test_atomic_model_atomic_stat.py +++ b/source/tests/pt/model/test_atomic_model_atomic_stat.py @@ -5,6 +5,7 @@ Path, ) from typing import ( + NoReturn, Optional, ) @@ -114,10 +115,10 @@ def forward( class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.merged_output_stat = [ { @@ -171,7 +172,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -240,7 +241,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) @@ -284,10 +285,10 @@ def raise_error(): class TestAtomicModelStatMergeGlobalAtomic( unittest.TestCase, TestCaseSingleFrameWithNlist ): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.merged_output_stat = [ { @@ -341,7 +342,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -401,7 +402,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py index 9ce5784bfa..fbb0279c09 100644 --- a/source/tests/pt/model/test_atomic_model_global_stat.py +++ b/source/tests/pt/model/test_atomic_model_global_stat.py @@ -5,6 +5,7 @@ Path, ) from typing import ( + NoReturn, Optional, ) @@ -137,10 +138,10 @@ def forward( class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) nf, nloc, nnei = self.nlist.shape self.merged_output_stat = [ @@ -173,7 +174,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -243,7 +244,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) @@ -280,7 +281,7 @@ def raise_error(): # bar is too complicated to be manually computed. np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std) - def test_preset_bias(self): + def test_preset_bias(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -355,7 +356,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) @@ -392,7 +393,7 @@ def raise_error(): np.testing.assert_almost_equal(ret3[kk], expected_ret3[kk]) # bar is too complicated to be manually computed. - def test_preset_bias_all_none(self): + def test_preset_bias_all_none(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -462,7 +463,7 @@ def cvt_ret(x): for kk in ["foo", "pix", "bar"]: np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) - def test_serialize(self): + def test_serialize(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, diff --git a/source/tests/pt/model/test_autodiff.py b/source/tests/pt/model/test_autodiff.py index 1adcff55fc..31e06af751 100644 --- a/source/tests/pt/model/test_autodiff.py +++ b/source/tests/pt/model/test_autodiff.py @@ -60,7 +60,7 @@ def stretch_box(old_coord, old_box, new_box): class ForceTest: def test( self, - ): + ) -> None: places = 5 delta = 1e-5 natoms = 5 @@ -132,7 +132,7 @@ def ff_spin(_spin): class VirialTest: def test( self, - ): + ) -> None: places = 5 delta = 1e-4 natoms = 5 @@ -176,77 +176,77 @@ def ff(bb): class TestEnergyModelSeAForce(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSeAVirial(unittest.TestCase, VirialTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1Force(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1Virial(unittest.TestCase, VirialTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2Force(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPAUniVirial(unittest.TestCase, VirialTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelHybridForce(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelHybridVirial(unittest.TestCase, VirialTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelZBLForce(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelZBLVirial(unittest.TestCase, VirialTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSpinSeAForce(unittest.TestCase, ForceTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_compressed_descriptor_dpa2.py b/source/tests/pt/model/test_compressed_descriptor_dpa2.py index 05b1143eb1..5e84572b6d 100644 --- a/source/tests/pt/model/test_compressed_descriptor_dpa2.py +++ b/source/tests/pt/model/test_compressed_descriptor_dpa2.py @@ -51,7 +51,7 @@ def eval_pt_descriptor( @parameterized(("float32", "float64"), (True, False)) class TestDescriptorDPA2(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.dtype, self.type_one_side) = self.param if self.dtype == "float32": self.skipTest("FP32 has bugs:") @@ -119,7 +119,7 @@ def setUp(self): precision=self.dtype, ) - def test_compressed_forward(self): + def test_compressed_forward(self) -> None: result_pt = eval_pt_descriptor( self.descriptor, self.natoms, diff --git a/source/tests/pt/model/test_compressed_descriptor_se_a.py b/source/tests/pt/model/test_compressed_descriptor_se_a.py index 14d82a452c..1165878759 100644 --- a/source/tests/pt/model/test_compressed_descriptor_se_a.py +++ b/source/tests/pt/model/test_compressed_descriptor_se_a.py @@ -47,7 +47,7 @@ def eval_pt_descriptor( @parameterized(("float32", "float64"), (True, False)) class TestDescriptorSeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.dtype, self.type_one_side) = self.param if self.dtype == "float32": self.atol = 1e-5 @@ -101,7 +101,7 @@ def setUp(self): precision=self.dtype, ) - def test_compressed_forward(self): + def test_compressed_forward(self) -> None: result_pt = eval_pt_descriptor( self.se_a, self.natoms, diff --git a/source/tests/pt/model/test_compressed_descriptor_se_atten.py b/source/tests/pt/model/test_compressed_descriptor_se_atten.py index a439255396..d1c246a6b4 100644 --- a/source/tests/pt/model/test_compressed_descriptor_se_atten.py +++ b/source/tests/pt/model/test_compressed_descriptor_se_atten.py @@ -47,7 +47,7 @@ def eval_pt_descriptor( @parameterized(("float32", "float64"), (True, False)) class TestDescriptorSeAtten(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.dtype, self.type_one_side) = self.param if self.dtype == "float32": self.atol = 1e-5 @@ -106,7 +106,7 @@ def setUp(self): tebd_input_mode="strip", ) - def test_compressed_forward(self): + def test_compressed_forward(self) -> None: result_pt = eval_pt_descriptor( self.se_atten, self.natoms, diff --git a/source/tests/pt/model/test_compressed_descriptor_se_r.py b/source/tests/pt/model/test_compressed_descriptor_se_r.py index 156cb9a06d..ffe0646d02 100644 --- a/source/tests/pt/model/test_compressed_descriptor_se_r.py +++ b/source/tests/pt/model/test_compressed_descriptor_se_r.py @@ -47,7 +47,7 @@ def eval_pt_descriptor( @parameterized(("float32", "float64")) class TestDescriptorSeR(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.dtype,) = self.param if self.dtype == "float32": self.atol = 1e-5 @@ -98,7 +98,7 @@ def setUp(self): precision=self.dtype, ) - def test_compressed_forward(self): + def test_compressed_forward(self) -> None: result_pt = eval_pt_descriptor( self.se_r, self.natoms, diff --git a/source/tests/pt/model/test_compressed_descriptor_se_t.py b/source/tests/pt/model/test_compressed_descriptor_se_t.py index aa3054bc0d..5507c31749 100644 --- a/source/tests/pt/model/test_compressed_descriptor_se_t.py +++ b/source/tests/pt/model/test_compressed_descriptor_se_t.py @@ -47,7 +47,7 @@ def eval_pt_descriptor( @parameterized(("float32", "float64")) class TestDescriptorSeT(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (self.dtype,) = self.param if self.dtype == "float32": self.atol = 1e-5 @@ -98,7 +98,7 @@ def setUp(self): precision=self.dtype, ) - def test_compressed_forward(self): + def test_compressed_forward(self) -> None: result_pt = eval_pt_descriptor( self.se_t, self.natoms, diff --git a/source/tests/pt/model/test_deeppot.py b/source/tests/pt/model/test_deeppot.py index 7f530b0a5e..0ae5b5f3a3 100644 --- a/source/tests/pt/model/test_deeppot.py +++ b/source/tests/pt/model/test_deeppot.py @@ -23,7 +23,7 @@ class TestDeepPot(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -47,12 +47,12 @@ def setUp(self): trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0) self.model = "model.pt" - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f in ["lcurve.out", self.input_json]: os.remove(f) - def test_dp_test(self): + def test_dp_test(self) -> None: dp = DeepPot(str(self.model)) cell = np.array( [ @@ -102,12 +102,12 @@ def test_dp_test(self): self.assertEqual(dp.get_dim_aparam(), 0) self.assertEqual(dp.deep_eval.model_type, DeepPot) - def test_uni(self): + def test_uni(self) -> None: dp = DeepPotUni("model.pt") self.assertIsInstance(dp, DeepPot) # its methods has been tested in test_dp_test - def test_eval_typeebd(self): + def test_eval_typeebd(self) -> None: dp = DeepPot(str(self.model)) eval_typeebd = dp.eval_typeebd() self.assertEqual( @@ -117,7 +117,7 @@ def test_eval_typeebd(self): class TestDeepPotFrozen(TestDeepPot): - def setUp(self): + def setUp(self) -> None: super().setUp() frozen_model = "frozen_model.pth" freeze( @@ -132,7 +132,7 @@ def setUp(self): @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") @unittest.mock.patch("deepmd.pt.utils.env.DEVICE", torch.device("cpu")) @unittest.mock.patch("deepmd.pt.infer.deep_eval.DEVICE", torch.device("cpu")) - def test_dp_test_cpu(self): + def test_dp_test_cpu(self) -> None: self.test_dp_test() diff --git a/source/tests/pt/model/test_descriptor.py b/source/tests/pt/model/test_descriptor.py index ff1fd0c959..daec358003 100644 --- a/source/tests/pt/model/test_descriptor.py +++ b/source/tests/pt/model/test_descriptor.py @@ -97,7 +97,7 @@ def base_se_a(rcut, rcut_smth, sel, batch, mean, stddev): class TestSeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: dp_random.seed(20) with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: content = fin.read() @@ -123,7 +123,7 @@ def setUp(self): self.ntypes = len(self.sel) self.nnei = sum(self.sel) - def test_consistency(self): + def test_consistency(self) -> None: avg_zero = torch.zeros( [self.ntypes, self.nnei * 4], dtype=GLOBAL_PT_FLOAT_PRECISION, diff --git a/source/tests/pt/model/test_descriptor_dpa1.py b/source/tests/pt/model/test_descriptor_dpa1.py index 9652a63944..abf5d1af01 100644 --- a/source/tests/pt/model/test_descriptor_dpa1.py +++ b/source/tests/pt/model/test_descriptor_dpa1.py @@ -26,7 +26,7 @@ class TestDPA1(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: cell = [ 5.122106549439247480e00, 4.016537340154059388e-01, @@ -235,7 +235,7 @@ def setUp(self): self.file_model_param = Path(CUR_DIR) / "models" / "dpa1.pth" self.file_type_embed = Path(CUR_DIR) / "models" / "dpa2_tebd.pth" - def test_descriptor_block(self): + def test_descriptor_block(self) -> None: # torch.manual_seed(0) model_dpa1 = self.model_json dparams = model_dpa1["descriptor"] @@ -291,7 +291,7 @@ def test_descriptor_block(self): descriptor.view(-1), self.ref_d, atol=1e-10, rtol=1e-10 ) - def test_descriptor(self): + def test_descriptor(self) -> None: with open(Path(CUR_DIR) / "models" / "dpa1.json") as fp: self.model_json = json.load(fp) model_dpa2 = self.model_json diff --git a/source/tests/pt/model/test_descriptor_dpa2.py b/source/tests/pt/model/test_descriptor_dpa2.py index 7efbe0a921..6a859a497a 100644 --- a/source/tests/pt/model/test_descriptor_dpa2.py +++ b/source/tests/pt/model/test_descriptor_dpa2.py @@ -22,7 +22,7 @@ class TestDPA2(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: cell = [ 5.122106549439247480e00, 4.016537340154059388e-01, @@ -109,7 +109,7 @@ def setUp(self): self.file_model_param = Path(CUR_DIR) / "models" / "dpa2.pth" self.file_type_embed = Path(CUR_DIR) / "models" / "dpa2_tebd.pth" - def test_descriptor(self): + def test_descriptor(self) -> None: with open(Path(CUR_DIR) / "models" / "dpa2.json") as fp: self.model_json = json.load(fp) model_dpa2 = self.model_json diff --git a/source/tests/pt/model/test_descriptor_hybrid.py b/source/tests/pt/model/test_descriptor_hybrid.py index 074af4da4e..b05a843b4a 100644 --- a/source/tests/pt/model/test_descriptor_hybrid.py +++ b/source/tests/pt/model/test_descriptor_hybrid.py @@ -31,12 +31,12 @@ class TestDescrptHybrid(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_jit( self, - ): + ) -> None: ddsub0 = DescrptSeA( self.rcut, self.rcut_smth, @@ -54,7 +54,7 @@ def test_jit( def test_get_parameters( self, - ): + ) -> None: nf, nloc, nnei = self.nlist.shape ddsub0 = DescrptSeA( rcut=self.rcut, @@ -85,7 +85,7 @@ def test_get_parameters( with self.assertRaises(ValueError): self.assertAlmostEqual(em0.get_env_protection(), 0.0) - def test_hybrid_mixed_and_no_mixed(self): + def test_hybrid_mixed_and_no_mixed(self) -> None: coord_ext = to_torch_tensor(self.coord_ext) atype_ext = to_torch_tensor(self.atype_ext) nlist1 = to_torch_tensor(self.nlist) diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py index e4aa405dd8..8c6e2c6776 100644 --- a/source/tests/pt/model/test_descriptor_se_r.py +++ b/source/tests/pt/model/test_descriptor_se_r.py @@ -34,12 +34,12 @@ # to be merged with the tf test case class TestDescrptSeR(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 1)) @@ -109,7 +109,7 @@ def test_consistency( err_msg=err_msg, ) - def test_load_stat(self): + def test_load_stat(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 1)) @@ -159,7 +159,7 @@ def test_load_stat(self): def test_jit( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 1)) diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py index 71da2781ac..ecefe0dc28 100644 --- a/source/tests/pt/model/test_dipole_fitting.py +++ b/source/tests/pt/model/test_dipole_fitting.py @@ -58,7 +58,7 @@ def finite_difference(f, x, a, delta=1e-6): class TestDipoleFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.rng = np.random.default_rng(GLOBAL_SEED) self.nf, self.nloc, _ = self.nlist.shape @@ -66,7 +66,7 @@ def setUp(self): def test_consistency( self, - ): + ) -> None: rd0, gr, _, _, _ = self.dd0( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), @@ -127,7 +127,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: for mixed_types, nfp, nap in itertools.product( [True, False], [0, 3], @@ -164,7 +164,7 @@ def setUp(self) -> None: ) self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) - def test_rot(self): + def test_rot(self) -> None: atype = self.atype.reshape(1, 5) rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) coord_rot = torch.matmul(self.coord, rmat) @@ -228,7 +228,7 @@ def test_rot(self): to_numpy_array(res[1]), to_numpy_array(torch.matmul(res[0], rmat)) ) - def test_permu(self): + def test_permu(self) -> None: coord = torch.matmul(self.coord, self.cell) ft0 = DipoleFittingNet( 3, # ntype @@ -269,7 +269,7 @@ def test_permu(self): to_numpy_array(res[0][:, idx_perm]), to_numpy_array(res[1]) ) - def test_trans(self): + def test_trans(self) -> None: atype = self.atype.reshape(1, 5) coord_s = torch.matmul( torch.remainder( @@ -310,7 +310,7 @@ def test_trans(self): class TestDipoleModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.natoms = 5 self.rcut = 4.0 self.nt = 3 @@ -338,7 +338,7 @@ def setUp(self): self.model = DipoleModel(self.dd0, self.ft0, self.type_mapping) self.file_path = "model_output.pth" - def test_auto_diff(self): + def test_auto_diff(self) -> None: places = 5 delta = 1e-5 atype = self.atype.view(self.nf, self.natoms) @@ -360,7 +360,7 @@ def ff(coord, atype): np.testing.assert_almost_equal(fdf, rff.transpose(0, 2, 1, 3), decimal=places) - def test_deepdipole_infer(self): + def test_deepdipole_infer(self) -> None: atype = to_numpy_array(self.atype.view(self.nf, self.natoms)) coord = to_numpy_array(self.coord.reshape(1, 5, 3)) cell = to_numpy_array(self.cell.reshape(1, 9)) diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py index aeb2b591ef..6d6a22f357 100644 --- a/source/tests/pt/model/test_dp_atomic_model.py +++ b/source/tests/pt/model/test_dp_atomic_model.py @@ -34,10 +34,10 @@ class TestDPAtomicModel(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_self_consistency(self): + def test_self_consistency(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -74,7 +74,7 @@ def test_self_consistency(self): to_numpy_array(ret1["energy"]), ) - def test_dp_consistency(self): + def test_dp_consistency(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DPDescrptSeA( self.rcut, @@ -102,7 +102,7 @@ def test_dp_consistency(self): to_numpy_array(ret1["energy"]), ) - def test_jit(self): + def test_jit(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -122,7 +122,7 @@ def test_jit(self): self.assertEqual(md0.get_rcut(), self.rcut) self.assertEqual(md0.get_type_map(), type_map) - def test_excl_consistency(self): + def test_excl_consistency(self) -> None: type_map = ["foo", "bar"] # test the case of exclusion @@ -190,13 +190,13 @@ def test_excl_consistency(self): class TestDPAtomicModelVirtualConsistency(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.case0 = TestCaseSingleFrameWithNlist() self.case1 = TestCaseSingleFrameWithNlistWithVirtual() self.case0.setUp() self.case1.setUp() - def test_virtual_consistency(self): + def test_virtual_consistency(self) -> None: nf, _, _ = self.case0.nlist.shape ds = DescrptSeA( self.case0.rcut, diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index db4fab1c63..93153ce6d5 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -41,10 +41,10 @@ class TestDPModel(unittest.TestCase, TestCaseSingleFrameWithoutNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithoutNlist.setUp(self) - def test_self_consistency(self): + def test_self_consistency(self) -> None: nf, nloc = self.atype.shape ds = DescrptSeA( self.rcut, @@ -114,7 +114,7 @@ def test_self_consistency(self): atol=self.atol, ) - def test_dp_consistency(self): + def test_dp_consistency(self) -> None: nf, nloc = self.atype.shape nfp, nap = 2, 3 ds = DPDescrptSeA( @@ -153,7 +153,7 @@ def test_dp_consistency(self): atol=self.atol, ) - def test_dp_consistency_nopbc(self): + def test_dp_consistency_nopbc(self) -> None: nf, nloc = self.atype.shape nfp, nap = 2, 3 ds = DPDescrptSeA( @@ -192,7 +192,7 @@ def test_dp_consistency_nopbc(self): atol=self.atol, ) - def test_prec_consistency(self): + def test_prec_consistency(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc = self.atype.shape ds = DPDescrptSeA( @@ -243,10 +243,10 @@ def test_prec_consistency(self): class TestDPModelLower(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_self_consistency(self): + def test_self_consistency(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -294,7 +294,7 @@ def test_self_consistency(self): atol=self.atol, ) - def test_dp_consistency(self): + def test_dp_consistency(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DPDescrptSeA( self.rcut, @@ -326,7 +326,7 @@ def test_dp_consistency(self): atol=self.atol, ) - def test_prec_consistency(self): + def test_prec_consistency(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape ds = DPDescrptSeA( @@ -377,7 +377,7 @@ def test_prec_consistency(self): atol=self.atol, ) - def test_jit(self): + def test_jit(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -397,7 +397,7 @@ def test_jit(self): class TestDPModelFormatNlist(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nall = 5 @@ -440,7 +440,7 @@ def setUp(self): type_map = ["foo", "bar"] self.md = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE) - def test_nlist_eq(self): + def test_nlist_eq(self) -> None: # n_nnei == nnei nlist = np.array( [ @@ -457,7 +457,7 @@ def test_nlist_eq(self): ) np.testing.assert_equal(self.expected_nlist, to_numpy_array(nlist1)) - def test_nlist_st(self): + def test_nlist_st(self) -> None: # n_nnei < nnei nlist = np.array( [ @@ -474,7 +474,7 @@ def test_nlist_st(self): ) np.testing.assert_equal(self.expected_nlist, to_numpy_array(nlist1)) - def test_nlist_lt(self): + def test_nlist_lt(self) -> None: # n_nnei > nnei nlist = np.array( [ @@ -493,10 +493,10 @@ def test_nlist_lt(self): class TestEnergyModel(unittest.TestCase, TestCaseSingleFrameWithoutNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithoutNlist.setUp(self) - def test_self_consistency(self): + def test_self_consistency(self) -> None: nf, nloc = self.atype.shape ds = DescrptSeA( self.rcut, @@ -561,10 +561,10 @@ def test_self_consistency(self): class TestEnergyModelLower(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_self_consistency(self): + def test_self_consistency(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, @@ -612,7 +612,7 @@ def test_self_consistency(self): atol=self.atol, ) - def test_jit(self): + def test_jit(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptSeA( self.rcut, diff --git a/source/tests/pt/model/test_dpa1.py b/source/tests/pt/model/test_dpa1.py index d168ceb2ae..9186774110 100644 --- a/source/tests/pt/model/test_dpa1.py +++ b/source/tests/pt/model/test_dpa1.py @@ -30,12 +30,12 @@ class TestDescrptSeAtten(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(100) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -110,7 +110,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_dpa2.py b/source/tests/pt/model/test_dpa2.py index 2eac49d573..38f32164e6 100644 --- a/source/tests/pt/model/test_dpa2.py +++ b/source/tests/pt/model/test_dpa2.py @@ -34,12 +34,12 @@ class TestDescrptDPA2(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(100) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -195,7 +195,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: rng = np.random.default_rng(100) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_embedding_net.py b/source/tests/pt/model/test_embedding_net.py index 2cfcaa820e..0dd2c192a8 100644 --- a/source/tests/pt/model/test_embedding_net.py +++ b/source/tests/pt/model/test_embedding_net.py @@ -124,7 +124,7 @@ def base_se_a(descriptor, coord, atype, natoms, box): class TestSeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: dp_random.seed(0) with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: content = fin.read() @@ -149,7 +149,7 @@ def setUp(self): self.axis_neuron = model_config["descriptor"]["axis_neuron"] self.np_batch, self.torch_batch = get_single_batch(ds) - def test_consistency(self): + def test_consistency(self) -> None: dp_d = DescrptSeA_tf( rcut=self.rcut, rcut_smth=self.rcut_smth, diff --git a/source/tests/pt/model/test_ener_fitting.py b/source/tests/pt/model/test_ener_fitting.py index acf0a47769..1d0b140c7d 100644 --- a/source/tests/pt/model/test_ener_fitting.py +++ b/source/tests/pt/model/test_ener_fitting.py @@ -30,12 +30,12 @@ class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: # ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 1600 is different from 1604) rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape @@ -107,7 +107,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: for od, mixed_types, nfp, nap, et, use_aparam_as_mask in itertools.product( [1, 3], [True, False], @@ -130,7 +130,7 @@ def test_jit( ).to(env.DEVICE) torch.jit.script(ft0) - def test_get_set(self): + def test_get_set(self) -> None: ifn0 = InvarFitting( "energy", self.nt, @@ -152,7 +152,7 @@ def test_get_set(self): foo, np.reshape(ifn0[ii].detach().cpu().numpy(), foo.shape) ) - def test_use_aparam_as_mask(self): + def test_use_aparam_as_mask(self) -> None: nap = 4 dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index a0cc621c9b..ddea392f33 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -57,7 +57,7 @@ def reduce_tensor(extended_tensor, mapping, nloc: int): class SpinTest: - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 natoms = 5 self.ntypes = 3 # ["O", "H", "B"] for test @@ -95,7 +95,7 @@ def setUp(self): def test_output_shape( self, - ): + ) -> None: result = self.model( self.coord, self.atype, @@ -111,7 +111,7 @@ def test_output_shape( torch.testing.assert_close(result["force"].shape, [nframes, nloc, 3]) torch.testing.assert_close(result["force_mag"].shape, [nframes, nloc, 3]) - def test_input_output_process(self): + def test_input_output_process(self) -> None: nframes, nloc = self.coord.shape[:2] self.real_ntypes = self.model.spin.get_ntypes_real() # 1. test forward input process @@ -278,13 +278,13 @@ def test_input_output_process(self): force_mag, force_all_switched[:, nall:] * virtual_scale.unsqueeze(-1) ) - def test_jit(self): + def test_jit(self) -> None: model = torch.jit.script(self.model) self.assertEqual(model.get_rcut(), self.rcut) self.assertEqual(model.get_nsel(), self.nsel) self.assertEqual(model.get_type_map(), self.type_map) - def test_self_consistency(self): + def test_self_consistency(self) -> None: if hasattr(self, "serial_test") and not self.serial_test: # not implement serialize and deserialize return @@ -307,7 +307,7 @@ def test_self_consistency(self): ) model1 = torch.jit.script(model1) - def test_dp_consistency(self): + def test_dp_consistency(self) -> None: if hasattr(self, "serial_test") and not self.serial_test: # not implement serialize and deserialize return @@ -384,7 +384,7 @@ def test_dp_consistency(self): class TestEnergyModelSpinSeA(unittest.TestCase, SpinTest): - def setUp(self): + def setUp(self) -> None: SpinTest.setUp(self) model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_se_e2_a["descriptor"]) @@ -395,7 +395,7 @@ def setUp(self): class TestEnergyModelSpinDPA1(unittest.TestCase, SpinTest): - def setUp(self): + def setUp(self) -> None: SpinTest.setUp(self) model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_dpa1["descriptor"]) @@ -408,7 +408,7 @@ def setUp(self): class TestEnergyModelSpinDPA2(unittest.TestCase, SpinTest): - def setUp(self): + def setUp(self) -> None: SpinTest.setUp(self) model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_dpa2["descriptor"]) diff --git a/source/tests/pt/model/test_env_mat.py b/source/tests/pt/model/test_env_mat.py index f8cdd62310..391f50da03 100644 --- a/source/tests/pt/model/test_env_mat.py +++ b/source/tests/pt/model/test_env_mat.py @@ -22,7 +22,7 @@ class TestCaseSingleFrameWithNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nall = 4 @@ -76,7 +76,7 @@ def setUp(self): class TestCaseSingleFrameWithNlistWithVirtual: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 4 self.nall = 5 @@ -128,7 +128,7 @@ def setUp(self): class TestCaseSingleFrameWithoutNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nf, self.nt = 1, 2 @@ -153,12 +153,12 @@ def setUp(self): # to be merged with the tf test case class TestEnvMat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_exclusion_mask.py b/source/tests/pt/model/test_exclusion_mask.py index b50f163eb6..31683888d7 100644 --- a/source/tests/pt/model/test_exclusion_mask.py +++ b/source/tests/pt/model/test_exclusion_mask.py @@ -23,7 +23,7 @@ class TestAtomExcludeMask(unittest.TestCase): - def test_build_type_exclude_mask(self): + def test_build_type_exclude_mask(self) -> None: nf = 2 nt = 3 exclude_types = [0, 2] @@ -47,10 +47,10 @@ def test_build_type_exclude_mask(self): # to be merged with the tf test case class TestPairExcludeMask(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_build_type_exclude_mask(self): + def test_build_type_exclude_mask(self) -> None: exclude_types = [[0, 1]] expected_mask = np.array( [ diff --git a/source/tests/pt/model/test_fitting_net.py b/source/tests/pt/model/test_fitting_net.py index e08eed4f8b..7cb0d13a2e 100644 --- a/source/tests/pt/model/test_fitting_net.py +++ b/source/tests/pt/model/test_fitting_net.py @@ -27,7 +27,7 @@ class FakeDescriptor: - def __init__(self, ntypes, embedding_width): + def __init__(self, ntypes, embedding_width) -> None: self._ntypes = ntypes self._dim_out = embedding_width @@ -84,7 +84,7 @@ def base_fitting_net(dp_fn, embedding, natoms, atype): class TestFittingNet(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: nloc = 7 self.embedding_width = 30 self.natoms = np.array([nloc, nloc, 2, 5], dtype=np.int32) @@ -104,7 +104,7 @@ def setUp(self): ) self.dp_fn.bias_atom_e = rng.uniform(size=[self.ntypes]) - def test_consistency(self): + def test_consistency(self) -> None: dp_energy, values = base_fitting_net( self.dp_fn, self.embedding, self.natoms, self.atype ) diff --git a/source/tests/pt/model/test_force_grad.py b/source/tests/pt/model/test_force_grad.py index d3cd11f71d..44dbffa536 100644 --- a/source/tests/pt/model/test_force_grad.py +++ b/source/tests/pt/model/test_force_grad.py @@ -32,7 +32,7 @@ def __init__( self, sys_path: str, type_map: Optional[list[str]] = None, - ): + ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) self.add("force", 3, atomic=True, must=False, high_prec=False) @@ -61,7 +61,7 @@ def get_data(batch): class TestForceGrad(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: self.config = json.load(fin) data_file = [str(Path(__file__).parent / "water/data/data_0")] @@ -72,10 +72,10 @@ def setUp(self): self.get_dataset(self.system_index, self.batch_index) self.get_model() - def get_model(self): + def get_model(self) -> None: self.model = get_model(self.config["model"]).to(env.DEVICE) - def get_dataset(self, system_index=0, batch_index=0): + def get_dataset(self, system_index=0, batch_index=0) -> None: systems = self.config["training"]["training_data"]["systems"] rcut = self.config["model"]["descriptor"]["rcut"] sel = self.config["model"]["descriptor"]["sel"] @@ -87,7 +87,7 @@ def get_dataset(self, system_index=0, batch_index=0): self.origin_batch = self.dpdatasystem.get_item_torch(batch_index) @unittest.skip("it can be replaced by autodiff") - def test_force_grad(self, threshold=1e-2, delta0=1e-6, seed=20): + def test_force_grad(self, threshold=1e-2, delta0=1e-6, seed=20) -> None: rng = np.random.default_rng(GLOBAL_SEED) result0 = self.model(**get_data(self.origin_batch)) np.random.default_rng(seed) diff --git a/source/tests/pt/model/test_forward_lower.py b/source/tests/pt/model/test_forward_lower.py index 87a3f5b06e..234d611a6e 100644 --- a/source/tests/pt/model/test_forward_lower.py +++ b/source/tests/pt/model/test_forward_lower.py @@ -57,7 +57,7 @@ def reduce_tensor(extended_tensor, mapping, nloc: int): class ForwardLowerTest: def test( self, - ): + ) -> None: prec = self.prec natoms = 5 cell = 4.0 * torch.eye(3, dtype=dtype, device=env.DEVICE) @@ -135,35 +135,35 @@ def test( class TestEnergyModelSeA(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_se_e2_a) self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_dpa1) self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_dpa2) self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelZBL(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_zbl) self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSpinSeA(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_spin) self.test_spin = True @@ -171,7 +171,7 @@ def setUp(self): class TestEnergyModelSpinDPA1(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_dpa1)["descriptor"] @@ -182,7 +182,7 @@ def setUp(self): class TestEnergyModelSpinDPA2(unittest.TestCase, ForwardLowerTest): - def setUp(self): + def setUp(self) -> None: self.prec = 1e-10 model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_dpa2)["descriptor"] diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py index 4774582f57..e323c95ce0 100644 --- a/source/tests/pt/model/test_get_model.py +++ b/source/tests/pt/model/test_get_model.py @@ -45,7 +45,7 @@ class TestGetModel(unittest.TestCase): - def test_model_attr(self): + def test_model_attr(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.model = get_model(model_params).to(env.DEVICE) atomic_model = self.model.atomic_model @@ -63,7 +63,7 @@ def test_model_attr(self): self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) - def test_model_attr_energy_float(self): + def test_model_attr_energy_float(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": ["1.", 3, None]} self.model = get_model(model_params).to(env.DEVICE) @@ -82,19 +82,19 @@ def test_model_attr_energy_float(self): self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) - def test_model_attr_energy_unsupported_type(self): + def test_model_attr_energy_unsupported_type(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": [1.0 + 2.0j, 3, None]} with self.assertRaises(ValueError): self.model = get_model(model_params).to(env.DEVICE) - def test_model_attr_energy_unsupported_value(self): + def test_model_attr_energy_unsupported_value(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": ["1.0 + 2.0j", 3, None]} with self.assertRaises(ValueError): self.model = get_model(model_params).to(env.DEVICE) - def test_notset_model_attr(self): + def test_notset_model_attr(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params.pop("atom_exclude_types") model_params.pop("pair_exclude_types") @@ -106,7 +106,7 @@ def test_notset_model_attr(self): self.assertEqual(atomic_model.atom_exclude_types, []) self.assertEqual(atomic_model.pair_exclude_types, []) - def test_preset_wrong_len(self): + def test_preset_wrong_len(self) -> None: model_params = copy.deepcopy(model_se_e2_a) model_params["preset_out_bias"] = {"energy": [None]} with self.assertRaises(ValueError): diff --git a/source/tests/pt/model/test_jit.py b/source/tests/pt/model/test_jit.py index 1f1034c330..0fc672db86 100644 --- a/source/tests/pt/model/test_jit.py +++ b/source/tests/pt/model/test_jit.py @@ -29,13 +29,13 @@ class JITTest: - def test_jit(self): + def test_jit(self) -> None: trainer = get_trainer(deepcopy(self.config)) trainer.run() model = torch.jit.script(inference.Tester("./model.pt").model) torch.jit.save(model, "./frozen_model.pth", {}) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith("pt"): os.remove(f) @@ -48,7 +48,7 @@ def tearDown(self): class TestEnergyModelSeA(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -59,12 +59,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestDOSModelSeA(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent.parent / "dos/input.json") with open(input_json) as f: self.config = json.load(f) @@ -75,12 +75,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestEnergyModelDPA1(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -91,12 +91,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestEnergyModelDPA2(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -107,12 +107,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestEnergyModelHybrid(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -123,12 +123,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestEnergyModelHybrid2(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -140,12 +140,12 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) class TestEnergyModelDPA2IntRcut(unittest.TestCase, JITTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -162,7 +162,7 @@ def setUp(self): self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 - def tearDown(self): + def tearDown(self) -> None: JITTest.tearDown(self) diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py index c4b2648028..038bb04a5a 100644 --- a/source/tests/pt/model/test_linear_atomic_model.py +++ b/source/tests/pt/model/test_linear_atomic_model.py @@ -44,7 +44,7 @@ class TestWeightCalculation(unittest.TestCase): @patch("numpy.loadtxt") - def test_pairwise(self, mock_loadtxt): + def test_pairwise(self, mock_loadtxt) -> None: file_path = "dummy_path" mock_loadtxt.return_value = np.array( [ @@ -121,7 +121,7 @@ def test_pairwise(self, mock_loadtxt): class TestIntegration(unittest.TestCase, TestCaseSingleFrameWithNlist): @patch("numpy.loadtxt") - def setUp(self, mock_loadtxt): + def setUp(self, mock_loadtxt) -> None: TestCaseSingleFrameWithNlist.setUp(self) file_path = "dummy_path" mock_loadtxt.return_value = np.array( @@ -165,7 +165,7 @@ def setUp(self, mock_loadtxt): dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25, type_map=type_map ) - def test_self_consistency(self): + def test_self_consistency(self) -> None: args = [ to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] ] @@ -181,7 +181,7 @@ def test_self_consistency(self): to_numpy_array(ret0["energy"]), ret2["energy"], atol=0.001, rtol=0.001 ) - def test_jit(self): + def test_jit(self) -> None: md1 = torch.jit.script(self.md1) # atomic model no more export methods # self.assertEqual(md1.get_rcut(), self.rcut) @@ -193,7 +193,7 @@ def test_jit(self): class TestRemmapMethod(unittest.TestCase): - def test_valid(self): + def test_valid(self) -> None: generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) atype = torch.randint(0, 3, (4, 20), device=env.DEVICE, generator=generator) commonl = ["H", "O", "S"] diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py index 49b7a3821f..90758526b9 100644 --- a/source/tests/pt/model/test_linear_atomic_model_stat.py +++ b/source/tests/pt/model/test_linear_atomic_model_stat.py @@ -149,10 +149,10 @@ def forward( class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) nf, nloc, nnei = self.nlist.shape self.merged_output_stat = [ @@ -179,7 +179,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_linear_atomic_model_stat_with_bias(self): + def test_linear_atomic_model_stat_with_bias(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index df58d115a2..f2387ca8b0 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -63,7 +63,7 @@ def finite_hessian(f, x, delta=1e-6): class HessianTest: def test( self, - ): + ) -> None: # setup test case places = 6 delta = 1e-3 @@ -142,7 +142,7 @@ def ff(xx): class TestDPModel(unittest.TestCase, HessianTest): - def setUp(self): + def setUp(self) -> None: torch.manual_seed(2) self.nf = 2 self.nloc = 3 @@ -174,7 +174,7 @@ def setUp(self): self.model_valu = EnergyModel.deserialize(self.model_hess.serialize()) self.model_hess.requires_hessian("energy") - def test_output_def(self): + def test_output_def(self) -> None: self.assertTrue(self.model_hess.atomic_output_def()["energy"].r_hessian) self.assertFalse(self.model_valu.atomic_output_def()["energy"].r_hessian) self.assertTrue(self.model_hess.model_output_def()["energy"].r_hessian) diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py index ca2bb6d2cf..e6dde660cb 100644 --- a/source/tests/pt/model/test_mlp.py +++ b/source/tests/pt/model/test_mlp.py @@ -38,7 +38,7 @@ def get_tols(prec): class TestMLPLayer(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.test_cases = itertools.product( [(5, 5), (5, 10), (5, 8), (8, 5)], # inp, out [True, False], # bias @@ -51,7 +51,7 @@ def setUp(self): def test_match_native_layer( self, - ): + ) -> None: for (ninp, nout), bias, ut, ac, resnet, ashp, prec in self.test_cases: # input inp_shap = [ninp] @@ -85,7 +85,7 @@ def test_match_native_layer( err_msg=f"(i={ninp}, o={nout}) bias={bias} use_dt={ut} act={ac} resnet={resnet} prec={prec}", ) - def test_jit(self): + def test_jit(self) -> None: for (ninp, nout), bias, ut, ac, resnet, _, prec in self.test_cases: ml = MLPLayer(ninp, nout, bias, ut, ac, resnet, precision=prec) model = torch.jit.script(ml) @@ -94,7 +94,7 @@ def test_jit(self): class TestMLP(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.test_cases = itertools.product( [[2, 2, 4, 8], [1, 3, 3]], # inp and hiddens [True, False], # bias @@ -107,7 +107,7 @@ def setUp(self): def test_match_native_net( self, - ): + ) -> None: for ndims, bias, ut, ac, resnet, ashp, prec in self.test_cases: # input inp_shap = [ndims[0]] @@ -146,7 +146,7 @@ def test_match_native_net( err_msg=f"net={ndims} bias={bias} use_dt={ut} act={ac} resnet={resnet} prec={prec}", ) - def test_jit(self): + def test_jit(self) -> None: for ndims, bias, ut, ac, resnet, _, prec in self.test_cases: layers = [] for ii in range(1, len(ndims)): @@ -162,7 +162,7 @@ def test_jit(self): class TestEmbeddingNet(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.test_cases = itertools.product( [1, 3], # inp [[24, 48, 96], [24, 36]], # and hiddens @@ -173,7 +173,7 @@ def setUp(self): def test_match_embedding_net( self, - ): + ) -> None: for idim, nn, act, idt, prec in self.test_cases: # input rtol, atol = get_tols(prec) @@ -202,7 +202,7 @@ def test_match_embedding_net( def test_jit( self, - ): + ) -> None: for idim, nn, act, idt, prec in self.test_cases: # def MLP ml = EmbeddingNet(idim, nn, act, idt, prec).to(env.DEVICE) @@ -212,7 +212,7 @@ def test_jit( class TestFittingNet(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.test_cases = itertools.product( [1, 3], # inp [1, 5], # out @@ -225,7 +225,7 @@ def setUp(self): def test_match_fitting_net( self, - ): + ) -> None: for idim, odim, nn, act, idt, prec, ob in self.test_cases: # input rtol, atol = get_tols(prec) @@ -262,7 +262,7 @@ def test_match_fitting_net( def test_jit( self, - ): + ) -> None: for idim, odim, nn, act, idt, prec, ob in self.test_cases: # def MLP ml = FittingNet( diff --git a/source/tests/pt/model/test_model.py b/source/tests/pt/model/test_model.py index 0fa6baec68..e0fb850ca8 100644 --- a/source/tests/pt/model/test_model.py +++ b/source/tests/pt/model/test_model.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import collections import json import unittest +from typing import ( + NamedTuple, +) import numpy as np import tensorflow.compat.v1 as tf @@ -54,7 +56,10 @@ energy_data_requirement, ) -VariableState = collections.namedtuple("VariableState", ["value", "gradient"]) + +class VariableState(NamedTuple): + value: np.ndarray + gradient: np.ndarray def torch2tf(torch_name, last_layer_id=None): @@ -80,7 +85,7 @@ def torch2tf(torch_name, last_layer_id=None): class DpTrainer: - def __init__(self): + def __init__(self) -> None: with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: content = fin.read() config = json.loads(content) @@ -262,7 +267,7 @@ def _get_feed_dict(self, batch, place_holders): class TestEnergy(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.dp_trainer = DpTrainer() self.wanted_step = 0 for key in dir(self.dp_trainer): @@ -270,7 +275,7 @@ def setUp(self): value = getattr(self.dp_trainer, key) setattr(self, key, value) - def test_consistency(self): + def test_consistency(self) -> None: batch, head_dict, stat_dict, vs_dict = self.dp_trainer.get_intermediate_state( self.wanted_step ) @@ -406,7 +411,7 @@ def test_consistency(self): optimizer = torch.optim.Adam(my_model.parameters(), lr=cur_lr) optimizer.zero_grad() - def step(step_id): + def step(step_id) -> None: bdata = self.training_data.get_trainning_batch() optimizer.zero_grad() diff --git a/source/tests/pt/model/test_nlist.py b/source/tests/pt/model/test_nlist.py index 7558a2a7d9..c445ec337f 100644 --- a/source/tests/pt/model/test_nlist.py +++ b/source/tests/pt/model/test_nlist.py @@ -21,7 +21,7 @@ class TestNeighList(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.nf = 3 self.nloc = 3 self.ns = 5 * 5 * 3 @@ -59,7 +59,7 @@ def setUp(self): device=env.DEVICE, ) - def test_build_notype(self): + def test_build_notype(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -96,7 +96,7 @@ def test_build_notype(self): torch.sort(self.ref_nlist, descending=True, dim=-1)[0], ) - def test_build_type(self): + def test_build_type(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -120,7 +120,7 @@ def test_build_type(self): ], ) - def test_build_multiple_nlist(self): + def test_build_multiple_nlist(self) -> None: rcuts = [1.01, 2.01] nsels = [20, 80] ecoord, eatype, mapping = extend_coord_with_ghosts( @@ -161,7 +161,7 @@ def test_build_multiple_nlist(self): nlist2, ) - def test_extend_coord(self): + def test_extend_coord(self) -> None: ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut ) @@ -236,7 +236,7 @@ def test_extend_coord(self): atol=self.prec, ) - def test_build_directional_nlist(self): + def test_build_directional_nlist(self) -> None: """Directional nlist is tested against the standard nlist implementation.""" ecoord, eatype, mapping = extend_coord_with_ghosts( self.coord, self.atype, self.cell, self.rcut diff --git a/source/tests/pt/model/test_null_input.py b/source/tests/pt/model/test_null_input.py index a2e0fa66db..063b551abf 100644 --- a/source/tests/pt/model/test_null_input.py +++ b/source/tests/pt/model/test_null_input.py @@ -36,7 +36,7 @@ class NullTest: def test_nloc_1( self, - ): + ) -> None: natoms = 1 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) # torch.manual_seed(1000) @@ -62,7 +62,7 @@ def test_nloc_1( def test_nloc_2_far( self, - ): + ) -> None: natoms = 2 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE, generator=generator) @@ -87,28 +87,28 @@ def test_nloc_2_far( class TestEnergyModelSeA(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelDPA2(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -117,14 +117,14 @@ def setUp(self): class TestEnergyModelHybrid(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelHybrid(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -133,7 +133,7 @@ def setUp(self): class TestEnergyModelZBL(unittest.TestCase, NullTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_zbl_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index 165e3dead7..0a6f6ea933 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -64,7 +64,7 @@ def setUp(self, mock_loadtxt) -> None: [[[1, 2], [0, 2]], [[1, 2], [0, 3]]], device=env.DEVICE ) - def test_without_mask(self): + def test_without_mask(self) -> None: result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) @@ -78,7 +78,7 @@ def test_without_mask(self): result["energy"], expected_result, rtol=0.0001, atol=0.0001 ) - def test_with_mask(self): + def test_with_mask(self) -> None: self.nlist = torch.tensor( [[[1, -1], [0, 2]], [[1, 2], [0, 3]]], device=env.DEVICE ) @@ -96,13 +96,13 @@ def test_with_mask(self): result["energy"], expected_result, rtol=0.0001, atol=0.0001 ) - def test_jit(self): + def test_jit(self) -> None: model = torch.jit.script(self.model) # atomic model no more export methods # self.assertEqual(model.get_rcut(), 0.02) # self.assertEqual(model.get_type_map(), ["H", "O"]) - def test_deserialize(self): + def test_deserialize(self) -> None: model1 = PairTabAtomicModel.deserialize(self.model.serialize()) torch.testing.assert_close(self.model.tab_data, model1.tab_data) torch.testing.assert_close(self.model.tab_info, model1.tab_info) @@ -126,7 +126,7 @@ def test_deserialize(self): # self.assertEqual(model1.get_rcut(), 0.02) # self.assertEqual(model1.get_type_map(), ["H", "O"]) - def test_cross_deserialize(self): + def test_cross_deserialize(self) -> None: model_dict = self.model.serialize() # pytorch model to dict model1 = DPPairTabAtomicModel.deserialize(model_dict) # dict to numpy model np.testing.assert_allclose(self.model.tab_data, model1.tab_data) diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index 2d391c7115..5c7b8db9a4 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -343,7 +343,7 @@ class PermutationTest: def test( self, - ): + ) -> None: natoms = 5 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE, generator=generator) @@ -396,35 +396,35 @@ def test( class TestEnergyModelSeA(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestDOSModelSeA(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dos) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelDPA2(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -433,14 +433,14 @@ def setUp(self): class TestEnergyModelHybrid(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelHybrid(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -449,14 +449,14 @@ def setUp(self): class TestEnergyModelZBL(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSpinSeA(unittest.TestCase, PermutationTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_permutation_denoise.py b/source/tests/pt/model/test_permutation_denoise.py index 53bf55fb0f..389520daa3 100644 --- a/source/tests/pt/model/test_permutation_denoise.py +++ b/source/tests/pt/model/test_permutation_denoise.py @@ -39,7 +39,7 @@ class PermutationDenoiseTest: def test( self, - ): + ) -> None: generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) natoms = 5 cell = torch.rand([3, 3], dtype=dtype, generator=generator).to(env.DEVICE) @@ -71,7 +71,7 @@ def test( @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, PermutationDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -79,7 +79,7 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, PermutationDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model( diff --git a/source/tests/pt/model/test_polar_atomic_model_stat.py b/source/tests/pt/model/test_polar_atomic_model_stat.py index d9ddfcd3e6..19fb2a4d3f 100644 --- a/source/tests/pt/model/test_polar_atomic_model_stat.py +++ b/source/tests/pt/model/test_polar_atomic_model_stat.py @@ -5,6 +5,7 @@ Path, ) from typing import ( + NoReturn, Optional, ) @@ -77,10 +78,10 @@ def forward( class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.merged_output_stat = [ { @@ -144,7 +145,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -217,7 +218,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(to_numpy_array(md0.out_std), expected_std) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py index 1ca563a8c2..4dcde309b1 100644 --- a/source/tests/pt/model/test_polarizability_fitting.py +++ b/source/tests/pt/model/test_polarizability_fitting.py @@ -43,7 +43,7 @@ class TestPolarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.rng = np.random.default_rng(GLOBAL_SEED) self.nf, self.nloc, _ = self.nlist.shape @@ -52,7 +52,7 @@ def setUp(self): def test_consistency( self, - ): + ) -> None: rd0, gr, _, _, _ = self.dd0( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), @@ -129,7 +129,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: for mixed_types, nfp, nap, fit_diag in itertools.product( [True, False], [0, 3], @@ -166,7 +166,7 @@ def setUp(self) -> None: self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) self.scale = self.rng.uniform(0, 1, self.nt).tolist() - def test_rot(self): + def test_rot(self) -> None: atype = self.atype.reshape(1, 5) rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) coord_rot = torch.matmul(self.coord, rmat) @@ -239,7 +239,7 @@ def test_rot(self): ), ) - def test_permu(self): + def test_permu(self) -> None: coord = torch.matmul(self.coord, self.cell) for fit_diag, scale in itertools.product([True, False], [None, self.scale]): ft0 = PolarFittingNet( @@ -284,7 +284,7 @@ def test_permu(self): to_numpy_array(res[1]), ) - def test_trans(self): + def test_trans(self) -> None: atype = self.atype.reshape(1, 5) coord_s = torch.matmul( torch.remainder( @@ -333,7 +333,7 @@ def test_trans(self): class TestPolarModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.natoms = 5 self.rcut = 4.0 self.nt = 3 @@ -358,7 +358,7 @@ def setUp(self): self.model = PolarModel(self.dd0, self.ft0, self.type_mapping) self.file_path = "model_output.pth" - def test_deepdipole_infer(self): + def test_deepdipole_infer(self) -> None: atype = self.atype.view(self.nf, self.natoms) coord = self.coord.reshape(1, 5, 3) cell = self.cell.reshape(1, 9) diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index dfe2725f3b..320560aa9f 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -43,7 +43,7 @@ class TestPropertyFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.rng = np.random.default_rng() self.nf, self.nloc, _ = self.nlist.shape @@ -51,7 +51,7 @@ def setUp(self): def test_consistency( self, - ): + ) -> None: rd0, gr, _, _, _ = self.dd0( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), @@ -134,7 +134,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: for nfp, nap, intensive, bias_method in itertools.product( [0, 3], [0, 4], @@ -189,7 +189,7 @@ def setUp(self) -> None: ) self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) - def test_trans(self): + def test_trans(self) -> None: atype = self.atype.reshape(1, 5) coord_s = torch.matmul( torch.remainder( @@ -250,7 +250,7 @@ def setUp(self) -> None: self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) - def test_rot(self): + def test_rot(self) -> None: atype = self.atype.reshape(1, 5) rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) coord_rot = torch.matmul(self.coord, rmat) @@ -318,7 +318,7 @@ def test_rot(self): to_numpy_array(res[0]), ) - def test_permu(self): + def test_permu(self) -> None: coord = torch.matmul(self.coord, self.cell) ft0 = PropertyFittingNet( self.nt, @@ -360,7 +360,7 @@ def test_permu(self): to_numpy_array(res[1]), ) - def test_trans(self): + def test_trans(self) -> None: atype = self.atype.reshape(1, 5) coord_s = torch.matmul( torch.remainder( @@ -406,7 +406,7 @@ def test_trans(self): class TestPropertyModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.natoms = 5 self.rcut = 4.0 self.nt = 3 @@ -432,7 +432,7 @@ def setUp(self): self.model = PropertyModel(self.dd0, self.ft0, self.type_mapping) self.file_path = "model_output.pth" - def test_deepproperty_infer(self): + def test_deepproperty_infer(self) -> None: atype = self.atype.view(self.nf, self.natoms) coord = self.coord.reshape(1, 5, 3) cell = self.cell.reshape(1, 9) diff --git a/source/tests/pt/model/test_region.py b/source/tests/pt/model/test_region.py index f92fb24f52..444613b65e 100644 --- a/source/tests/pt/model/test_region.py +++ b/source/tests/pt/model/test_region.py @@ -17,7 +17,7 @@ class TestRegion(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.cell = torch.tensor( [[1, 0, 0], [0.4, 0.8, 0], [0.1, 0.3, 2.1]], dtype=dtype, device="cpu" ) @@ -25,7 +25,7 @@ def setUp(self): self.cell = torch.tile(self.cell, [4, 5, 1, 1]) self.prec = 1e-8 - def test_inter_to_phys(self): + def test_inter_to_phys(self) -> None: generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED) inter = torch.rand([4, 5, 3, 3], dtype=dtype, device="cpu", generator=generator) phys = inter2phys(inter, self.cell) @@ -36,7 +36,7 @@ def test_inter_to_phys(self): phys[ii, jj], expected_phys, rtol=self.prec, atol=self.prec ) - def test_to_face_dist(self): + def test_to_face_dist(self) -> None: cell0 = self.cell[0][0].numpy() vol = np.linalg.det(cell0) # area of surfaces xy, xz, yz diff --git a/source/tests/pt/model/test_rot.py b/source/tests/pt/model/test_rot.py index ca6a6375c8..283dbb31d7 100644 --- a/source/tests/pt/model/test_rot.py +++ b/source/tests/pt/model/test_rot.py @@ -33,7 +33,7 @@ class RotTest: def test( self, - ): + ) -> None: generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) prec = 1e-10 natoms = 5 @@ -146,35 +146,35 @@ def test( class TestEnergyModelSeA(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestDOSModelSeA(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dos) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelDPA2(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -183,14 +183,14 @@ def setUp(self): class TestEnergyModelHybrid(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelHybrid(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -199,14 +199,14 @@ def setUp(self): class TestEnergyModelZBL(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSpinSeA(unittest.TestCase, RotTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_rot_denoise.py b/source/tests/pt/model/test_rot_denoise.py index 9828ba5225..fcae4b23d7 100644 --- a/source/tests/pt/model/test_rot_denoise.py +++ b/source/tests/pt/model/test_rot_denoise.py @@ -28,7 +28,7 @@ class RotDenoiseTest: def test( self, - ): + ) -> None: generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) prec = 1e-10 natoms = 5 @@ -104,7 +104,7 @@ def test( @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, RotDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -112,7 +112,7 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, RotDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_rotation.py b/source/tests/pt/model/test_rotation.py index cf947c30b2..713a957e99 100644 --- a/source/tests/pt/model/test_rotation.py +++ b/source/tests/pt/model/test_rotation.py @@ -30,7 +30,7 @@ def __init__( self, sys_path: str, type_map: Optional[list[str]] = None, - ): + ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) self.add("force", 3, atomic=True, must=False, high_prec=False) @@ -66,7 +66,7 @@ def get_data(batch): class TestRotation(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: self.config = json.load(fin) data_file = [str(Path(__file__).parent / "water/data/data_0")] @@ -77,17 +77,17 @@ def setUp(self): self.get_dataset(0) self.get_model() - def get_model(self): + def get_model(self) -> None: self.model = get_model(self.config["model"]).to(env.DEVICE) - def get_dataset(self, system_index=0, batch_index=0): + def get_dataset(self, system_index=0, batch_index=0) -> None: systems = self.config["training"]["training_data"]["systems"] type_map = self.config["model"]["type_map"] dpdatasystem = CheckSymmetry(sys_path=systems[system_index], type_map=type_map) self.origin_batch = dpdatasystem.get_item_torch(batch_index) self.rotated_batch = dpdatasystem.get_rotation(batch_index, self.rotation) - def test_rotation(self): + def test_rotation(self) -> None: result1 = self.model(**get_data(self.origin_batch)) result2 = self.model(**get_data(self.rotated_batch)) rotation = torch.from_numpy(self.rotation).to(env.DEVICE) diff --git a/source/tests/pt/model/test_saveload_dpa1.py b/source/tests/pt/model/test_saveload_dpa1.py index 5b2b6cd583..c97313231c 100644 --- a/source/tests/pt/model/test_saveload_dpa1.py +++ b/source/tests/pt/model/test_saveload_dpa1.py @@ -53,7 +53,7 @@ def get_dataset(config): class TestSaveLoadDPA1(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as fin: self.config = json.load(fin) @@ -129,7 +129,7 @@ def get_data(self): label_dict[item] = batch_data[item].to(env.DEVICE) return input_dict, label_dict - def test_saveload(self): + def test_saveload(self) -> None: result1 = self.get_model_result() result2 = self.get_model_result(read=True) final_result = all( diff --git a/source/tests/pt/model/test_saveload_se_e2_a.py b/source/tests/pt/model/test_saveload_se_e2_a.py index d226f628bc..ba3697d47a 100644 --- a/source/tests/pt/model/test_saveload_se_e2_a.py +++ b/source/tests/pt/model/test_saveload_se_e2_a.py @@ -53,7 +53,7 @@ def get_dataset(config): class TestSaveLoadSeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_e2_a.json") with open(input_json) as fin: self.config = json.load(fin) @@ -123,7 +123,7 @@ def get_data(self): label_dict[item] = batch_data[item].to(env.DEVICE) return input_dict, label_dict - def test_saveload(self): + def test_saveload(self) -> None: result1 = self.get_model_result() result2 = self.get_model_result(read=True) final_result = all( diff --git a/source/tests/pt/model/test_se_atten_v2.py b/source/tests/pt/model/test_se_atten_v2.py index 462b2aca34..e2d34ead5e 100644 --- a/source/tests/pt/model/test_se_atten_v2.py +++ b/source/tests/pt/model/test_se_atten_v2.py @@ -30,12 +30,12 @@ class TestDescrptSeAttenV2(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(100) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -106,7 +106,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: rng = np.random.default_rng() _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_se_e2_a.py b/source/tests/pt/model/test_se_e2_a.py index da9e69243c..bdedd440cb 100644 --- a/source/tests/pt/model/test_se_e2_a.py +++ b/source/tests/pt/model/test_se_e2_a.py @@ -31,12 +31,12 @@ # to be merged with the tf test case class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -107,7 +107,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_se_t.py b/source/tests/pt/model/test_se_t.py index d3968d7f03..fb60ca011e 100644 --- a/source/tests/pt/model/test_se_t.py +++ b/source/tests/pt/model/test_se_t.py @@ -31,12 +31,12 @@ # to be merged with the tf test case class TestDescrptSeT(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -111,7 +111,7 @@ def test_consistency( def test_jit( self, - ): + ) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) diff --git a/source/tests/pt/model/test_smooth.py b/source/tests/pt/model/test_smooth.py index 9a7040f9cc..1c6303d14c 100644 --- a/source/tests/pt/model/test_smooth.py +++ b/source/tests/pt/model/test_smooth.py @@ -33,7 +33,7 @@ class SmoothTest: def test( self, - ): + ) -> None: generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) # displacement of atoms epsilon = 1e-5 if self.epsilon is None else self.epsilon @@ -121,7 +121,7 @@ def test( ) ret3 = {key: result_3[key].squeeze(0) for key in test_keys} - def compare(ret0, ret1): + def compare(ret0, ret1) -> None: for key in test_keys: if key in ["energy"]: torch.testing.assert_close( @@ -146,7 +146,7 @@ def compare(ret0, ret1): class TestEnergyModelSeA(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) @@ -154,7 +154,7 @@ def setUp(self): class TestDOSModelSeA(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dos) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) @@ -162,7 +162,7 @@ def setUp(self): class TestEnergyModelDPA1(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -173,7 +173,7 @@ def setUp(self): class TestEnergyModelDPA1Excl1(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) model_params["pair_exclude_types"] = [[0, 1]] self.type_split = True @@ -185,7 +185,7 @@ def setUp(self): class TestEnergyModelDPA1Excl12(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) model_params["pair_exclude_types"] = [[0, 1], [0, 2]] self.type_split = True @@ -197,7 +197,7 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["descriptor"]["repinit"]["rcut"] = 8 model_params["descriptor"]["repinit"]["rcut_smth"] = 3.5 @@ -207,7 +207,7 @@ def setUp(self): class TestEnergyModelDPA2_1(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "ener" self.type_split = True @@ -217,7 +217,7 @@ def setUp(self): class TestEnergyModelDPA2_2(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "ener" self.type_split = True @@ -227,7 +227,7 @@ def setUp(self): class TestEnergyModelHybrid(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -235,7 +235,7 @@ def setUp(self): class TestEnergyModelZBL(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) @@ -243,7 +243,7 @@ def setUp(self): class TestEnergyModelSpinSeA(unittest.TestCase, SmoothTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_smooth_denoise.py b/source/tests/pt/model/test_smooth_denoise.py index faa892c5d0..199d6664a1 100644 --- a/source/tests/pt/model/test_smooth_denoise.py +++ b/source/tests/pt/model/test_smooth_denoise.py @@ -27,7 +27,7 @@ class SmoothDenoiseTest: def test( self, - ): + ) -> None: # displacement of atoms epsilon = 1e-5 if self.epsilon is None else self.epsilon # required prec. relative prec is not checked. @@ -88,7 +88,7 @@ def test( ) ret3 = {"updated_coord": update_c3.squeeze(0), "logits": logits3.squeeze(0)} - def compare(ret0, ret1): + def compare(ret0, ret1) -> None: torch.testing.assert_close( ret0["updated_coord"], ret1["updated_coord"], rtol=rprec, atol=aprec ) @@ -103,7 +103,7 @@ def compare(ret0, ret1): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, SmoothDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["descriptor"]["sel"] = 8 model_params["descriptor"]["rcut_smth"] = 3.5 @@ -116,7 +116,7 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2_1(unittest.TestCase, SmoothDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) # model_params["descriptor"]["combine_grrg"] = True self.type_split = True diff --git a/source/tests/pt/model/test_trans.py b/source/tests/pt/model/test_trans.py index b62fac1312..2e39cc4bd5 100644 --- a/source/tests/pt/model/test_trans.py +++ b/source/tests/pt/model/test_trans.py @@ -33,7 +33,7 @@ class TransTest: def test( self, - ): + ) -> None: natoms = 5 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE, generator=generator) @@ -88,35 +88,35 @@ def test( class TestEnergyModelSeA(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_se_e2_a) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestDOSModelSeA(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dos) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA1(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelDPA2(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelDPA2(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -125,14 +125,14 @@ def setUp(self): class TestEnergyModelHybrid(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) class TestForceModelHybrid(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True @@ -141,14 +141,14 @@ def setUp(self): class TestEnergyModelZBL(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_zbl) self.type_split = False self.model = get_model(model_params).to(env.DEVICE) class TestEnergyModelSpinSeA(unittest.TestCase, TransTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_spin) self.type_split = False self.test_spin = True diff --git a/source/tests/pt/model/test_trans_denoise.py b/source/tests/pt/model/test_trans_denoise.py index 84ec21929c..77bff7980a 100644 --- a/source/tests/pt/model/test_trans_denoise.py +++ b/source/tests/pt/model/test_trans_denoise.py @@ -29,7 +29,7 @@ class TransDenoiseTest: def test( self, - ): + ) -> None: natoms = 5 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) cell = torch.rand([3, 3], dtype=dtype, generator=generator).to(env.DEVICE) @@ -63,7 +63,7 @@ def test( @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA1(unittest.TestCase, TransDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa1) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -71,7 +71,7 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, TransDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -79,7 +79,7 @@ def setUp(self): @unittest.skip("hybrid not supported at the moment") class TestDenoiseModelHybrid(unittest.TestCase, TransDenoiseTest): - def setUp(self): + def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_unused_params.py b/source/tests/pt/model/test_unused_params.py index 8c223d7590..3df1c82173 100644 --- a/source/tests/pt/model/test_unused_params.py +++ b/source/tests/pt/model/test_unused_params.py @@ -25,7 +25,7 @@ class TestUnusedParamsDPA2(unittest.TestCase): - def test_unused(self): + def test_unused(self) -> None: import itertools for conv, drrd, grrg, attn1, g1g1, attn2, h2 in itertools.product( @@ -56,7 +56,7 @@ def test_unused(self): model["fitting_net"]["neuron"] = [12, 12, 12] self._test_unused(model) - def _test_unused(self, model_params): + def _test_unused(self, model_params) -> None: self.model = get_model(model_params).to(env.DEVICE) natoms = 5 generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) diff --git a/source/tests/pt/test_LKF.py b/source/tests/pt/test_LKF.py index 33aeac7f4f..9a524f1c85 100644 --- a/source/tests/pt/test_LKF.py +++ b/source/tests/pt/test_LKF.py @@ -12,7 +12,7 @@ class TestLKF(unittest.TestCase): - def test_lkf(self): + def test_lkf(self) -> None: with open(str(Path(__file__).parent / "water/lkf.json")) as fin: content = fin.read() self.config = json.loads(content) @@ -27,7 +27,7 @@ def test_lkf(self): json.dump(self.config, fp, indent=4) main(["train", self.input_json]) - def tearDown(self): + def tearDown(self) -> None: os.remove(self.input_json) diff --git a/source/tests/pt/test_auto_batch_size.py b/source/tests/pt/test_auto_batch_size.py index 71194e001e..c67a23df52 100644 --- a/source/tests/pt/test_auto_batch_size.py +++ b/source/tests/pt/test_auto_batch_size.py @@ -9,7 +9,7 @@ class TestAutoBatchSize(unittest.TestCase): - def test_execute_all(self): + def test_execute_all(self) -> None: dd0 = np.zeros((10000, 2, 1, 3, 4)) dd1 = np.ones((10000, 2, 1, 3, 4)) auto_batch_size = AutoBatchSize(256, 2.0) @@ -21,7 +21,7 @@ def func(dd1): np.testing.assert_equal(dd0, dd2[0]) np.testing.assert_equal(dd1, dd2[1]) - def test_execute_all_dict(self): + def test_execute_all_dict(self) -> None: dd0 = np.zeros((10000, 2, 1, 3, 4)) dd1 = np.ones((10000, 2, 1, 3, 4)) auto_batch_size = AutoBatchSize(256, 2.0) diff --git a/source/tests/pt/test_calculator.py b/source/tests/pt/test_calculator.py index 5f5470c611..c73bbad415 100644 --- a/source/tests/pt/test_calculator.py +++ b/source/tests/pt/test_calculator.py @@ -26,7 +26,7 @@ class TestCalculator(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -50,7 +50,7 @@ def setUp(self): self.calculator = DPCalculator("model.pt") - def test_calculator(self): + def test_calculator(self) -> None: from ase import ( Atoms, ) diff --git a/source/tests/pt/test_change_bias.py b/source/tests/pt/test_change_bias.py index febc439f50..a3cf3edbbc 100644 --- a/source/tests/pt/test_change_bias.py +++ b/source/tests/pt/test_change_bias.py @@ -51,7 +51,7 @@ class TestChangeBias(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -88,7 +88,7 @@ def setUp(self): model_name + "user_bias" + ".pt" ) - def test_change_bias_with_data(self): + def test_change_bias_with_data(self) -> None: run_dp( f"dp --pt change-bias {self.model_path!s} -s {self.data_file[0]} -o {self.model_path_data_bias!s}" ) @@ -108,7 +108,7 @@ def test_change_bias_with_data(self): expected_bias = expected_model.get_out_bias() torch.testing.assert_close(updated_bias, expected_bias) - def test_change_bias_with_data_sys_file(self): + def test_change_bias_with_data_sys_file(self) -> None: tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt") with open(tmp_file.name, "w") as f: f.writelines([sys + "\n" for sys in self.data_file]) @@ -131,7 +131,7 @@ def test_change_bias_with_data_sys_file(self): expected_bias = expected_model.get_out_bias() torch.testing.assert_close(updated_bias, expected_bias) - def test_change_bias_with_user_defined(self): + def test_change_bias_with_user_defined(self) -> None: user_bias = [0.1, 3.2, -0.5] run_dp( f"dp --pt change-bias {self.model_path!s} -b {' '.join([str(_) for _ in user_bias])} -o {self.model_path_user_bias!s}" @@ -147,7 +147,7 @@ def test_change_bias_with_user_defined(self): expected_bias = to_torch_tensor(np.array(user_bias)).view(updated_bias.shape) torch.testing.assert_close(updated_bias, expected_bias) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("change-bias-model") and f.endswith(".pt"): os.remove(f) diff --git a/source/tests/pt/test_dp_show.py b/source/tests/pt/test_dp_show.py index 5d6cb9bd36..3ccc3d5d15 100644 --- a/source/tests/pt/test_dp_show.py +++ b/source/tests/pt/test_dp_show.py @@ -30,7 +30,7 @@ class TestSingleTaskModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -45,7 +45,7 @@ def setUp(self): trainer.run() run_dp("dp --pt freeze") - def test_checkpoint(self): + def test_checkpoint(self) -> None: INPUT = "model.pt" ATTRIBUTES = "type-map descriptor fitting-net" with redirect_stderr(io.StringIO()) as f: @@ -61,7 +61,7 @@ def test_checkpoint(self): in results[-1] ) - def test_frozen_model(self): + def test_frozen_model(self) -> None: INPUT = "frozen_model.pth" ATTRIBUTES = "type-map descriptor fitting-net" with redirect_stderr(io.StringIO()) as f: @@ -77,7 +77,7 @@ def test_frozen_model(self): in results[-1] ) - def test_checkpoint_error(self): + def test_checkpoint_error(self) -> None: INPUT = "model.pt" ATTRIBUTES = "model-branch type-map descriptor fitting-net" with self.assertRaisesRegex( @@ -85,7 +85,7 @@ def test_checkpoint_error(self): ): run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}") - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith("pt"): os.remove(f) @@ -96,7 +96,7 @@ def tearDown(self): class TestMultiTaskModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/multitask.json") with open(input_json) as f: self.config = json.load(f) @@ -142,7 +142,7 @@ def setUp(self): trainer.run() run_dp("dp --pt freeze --head model_1") - def test_checkpoint(self): + def test_checkpoint(self) -> None: INPUT = "model.ckpt.pt" ATTRIBUTES = "model-branch type-map descriptor fitting-net" with redirect_stderr(io.StringIO()) as f: @@ -177,7 +177,7 @@ def test_checkpoint(self): in results[-1] ) - def test_frozen_model(self): + def test_frozen_model(self) -> None: INPUT = "frozen_model.pth" ATTRIBUTES = "type-map descriptor fitting-net" with redirect_stderr(io.StringIO()) as f: @@ -193,7 +193,7 @@ def test_frozen_model(self): in results[-1] ) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith("pt"): os.remove(f) diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index 0427f2b14a..dbec472cc0 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -30,7 +30,7 @@ class DPTest: - def test_dp_test_1_frame(self): + def test_dp_test_1_frame(self) -> None: trainer = get_trainer(deepcopy(self.config)) with torch.device("cpu"): input_dict, label_dict, _ = trainer.get_data(is_train=False) @@ -93,7 +93,7 @@ def test_dp_test_1_frame(self): ).reshape(-1, 3), ) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) @@ -106,7 +106,7 @@ def tearDown(self): class TestDPTestSeA(DPTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.detail_file = "test_dp_test_ener_detail" input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: @@ -123,7 +123,7 @@ def setUp(self): class TestDPTestSeASpin(DPTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.detail_file = "test_dp_test_ener_spin_detail" input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: @@ -141,7 +141,7 @@ def setUp(self): class TestDPTestPropertySeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.detail_file = "test_dp_test_property_detail" input_json = str(Path(__file__).parent / "property/input.json") with open(input_json) as f: @@ -159,7 +159,7 @@ def setUp(self): with open(self.input_json, "w") as fp: json.dump(self.config, fp, indent=4) - def test_dp_test_1_frame(self): + def test_dp_test_1_frame(self) -> None: trainer = get_trainer(deepcopy(self.config)) with torch.device("cpu"): input_dict, label_dict, _ = trainer.get_data(is_train=False) @@ -186,7 +186,7 @@ def test_dp_test_1_frame(self): to_numpy_array(result["property"])[0], ) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py index dff3a4c8a4..8998392830 100644 --- a/source/tests/pt/test_finetune.py +++ b/source/tests/pt/test_finetune.py @@ -99,7 +99,7 @@ class FinetuneTest: - def test_finetune_change_out_bias(self): + def test_finetune_change_out_bias(self) -> None: self.testkey = "energy" if self.testkey is None else self.testkey # get data data = DpLoaderSet( @@ -193,7 +193,7 @@ def test_finetune_change_out_bias(self): self.tearDown() - def test_finetune_change_type(self): + def test_finetune_change_type(self) -> None: if not self.mixed_types: # skip when not mixed_types return @@ -280,7 +280,7 @@ def test_finetune_change_type(self): self.tearDown() - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) @@ -291,7 +291,7 @@ def tearDown(self): class TestEnergyModelSeA(FinetuneTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -306,7 +306,7 @@ def setUp(self): class TestEnergyZBLModelSeA(FinetuneTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -321,7 +321,7 @@ def setUp(self): class TestEnergyDOSModelSeA(FinetuneTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "dos/input.json") with open(input_json) as f: self.config = json.load(f) @@ -336,7 +336,7 @@ def setUp(self): class TestEnergyModelDPA1(FinetuneTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -351,7 +351,7 @@ def setUp(self): class TestEnergyModelDPA2(FinetuneTest, unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) diff --git a/source/tests/pt/test_init_frz_model.py b/source/tests/pt/test_init_frz_model.py index 69c738d6bd..62fbeede3c 100644 --- a/source/tests/pt/test_init_frz_model.py +++ b/source/tests/pt/test_init_frz_model.py @@ -27,7 +27,7 @@ class TestInitFrzModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: config = json.load(f) @@ -74,7 +74,7 @@ def setUp(self): ) self.models.append(frozen_model) - def test_dp_test(self): + def test_dp_test(self) -> None: dp1 = DeepPot(str(self.models[0])) dp2 = DeepPot(str(self.models[1])) dp3 = DeepPot(str(self.models[2])) @@ -129,7 +129,7 @@ def test_dp_test(self): np.testing.assert_allclose(av1, av2, rtol=1e-10, atol=1e-10) np.testing.assert_allclose(av1, av3, rtol=1e-10, atol=1e-10) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("frozen_model") and f.endswith(".pth"): os.remove(f) diff --git a/source/tests/pt/test_init_model.py b/source/tests/pt/test_init_model.py index dd264fbe89..07994f083d 100644 --- a/source/tests/pt/test_init_model.py +++ b/source/tests/pt/test_init_model.py @@ -26,7 +26,7 @@ class TestInitModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: config = json.load(f) @@ -71,7 +71,7 @@ def setUp(self): trainer.run() self.models.append(ckpt_model + ".pt") - def test_dp_test(self): + def test_dp_test(self) -> None: dp1 = DeepPot(str(self.models[0])) dp2 = DeepPot(str(self.models[1])) dp3 = DeepPot(str(self.models[2])) @@ -126,7 +126,7 @@ def test_dp_test(self): np.testing.assert_allclose(av1, av2, rtol=1e-10, atol=1e-10) np.testing.assert_allclose(av1, av3, rtol=1e-10, atol=1e-10) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py index ddce7b6322..d0746c1368 100644 --- a/source/tests/pt/test_loss.py +++ b/source/tests/pt/test_loss.py @@ -47,7 +47,7 @@ def get_batch(system, type_map, data_requirement): class LossCommonTest(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.cur_lr = 1.2 if not self.spin: self.system = str(Path(__file__).parent / "water/data/data_0") @@ -281,7 +281,7 @@ def tearDown(self) -> None: class TestEnerStdLoss(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -312,7 +312,7 @@ def setUp(self): self.spin = False super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -349,7 +349,7 @@ def fake_model(): class TestEnerStdLossAePfGf(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -401,7 +401,7 @@ def setUp(self): self.spin = False super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -438,7 +438,7 @@ def fake_model(): class TestEnerStdLossAecoeff(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -471,7 +471,7 @@ def setUp(self): self.spin = False super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -508,7 +508,7 @@ def fake_model(): class TestEnerStdLossRelativeF(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -541,7 +541,7 @@ def setUp(self): self.spin = False super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -578,7 +578,7 @@ def fake_model(): class TestEnerSpinLoss(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -612,7 +612,7 @@ def setUp(self): self.spin = True super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -649,7 +649,7 @@ def fake_model(): class TestEnerSpinLossAe(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -689,7 +689,7 @@ def setUp(self): self.spin = True super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict @@ -726,7 +726,7 @@ def fake_model(): class TestEnerSpinLossAecoeff(LossCommonTest): - def setUp(self): + def setUp(self) -> None: self.start_lr = 1.1 self.start_pref_e = 0.02 self.limit_pref_e = 1.0 @@ -762,7 +762,7 @@ def setUp(self): self.spin = True super().setUp() - def test_consistency(self): + def test_consistency(self) -> None: with tf.Session(graph=self.g) as sess: tf_loss, tf_more_loss = sess.run( self.tf_loss_sess, feed_dict=self.feed_dict diff --git a/source/tests/pt/test_lr.py b/source/tests/pt/test_lr.py index 9fbde599bb..2d6bf156e1 100644 --- a/source/tests/pt/test_lr.py +++ b/source/tests/pt/test_lr.py @@ -15,13 +15,13 @@ class TestLearningRate(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.start_lr = 0.001 self.stop_lr = 3.51e-8 self.decay_steps = np.arange(400, 601, 100) self.stop_steps = np.arange(500, 1600, 500) - def test_consistency(self): + def test_consistency(self) -> None: for decay_step in self.decay_steps: for stop_step in self.stop_steps: self.decay_step = decay_step @@ -29,7 +29,7 @@ def test_consistency(self): self.judge_it() self.decay_rate_pt() - def judge_it(self): + def judge_it(self) -> None: base_lr = learning_rate.LearningRateExp( self.start_lr, self.stop_lr, self.decay_step ) @@ -55,7 +55,7 @@ def judge_it(self): self.assertTrue(np.allclose(base_vals, my_vals)) tf.reset_default_graph() - def decay_rate_pt(self): + def decay_rate_pt(self) -> None: my_lr = LearningRateExp( self.start_lr, self.stop_lr, self.decay_step, self.stop_step ) diff --git a/source/tests/pt/test_model_compression_se_a.py b/source/tests/pt/test_model_compression_se_a.py index 3738b61c13..7688743523 100644 --- a/source/tests/pt/test_model_compression_se_a.py +++ b/source/tests/pt/test_model_compression_se_a.py @@ -24,7 +24,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -102,7 +102,7 @@ def _init_models_skip_neighbor_stat(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global \ INPUT, \ FROZEN_MODEL, \ @@ -121,7 +121,7 @@ def setUpModule(): class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepEval(FROZEN_MODEL) cls.dp_compressed = DeepEval(COMPRESSED_MODEL) cls.coords = np.array( @@ -149,7 +149,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -164,7 +164,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -185,7 +185,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -212,7 +212,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -245,7 +245,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepEval(FROZEN_MODEL) cls.dp_compressed = DeepEval(COMPRESSED_MODEL) cls.coords = np.array( @@ -273,7 +273,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -294,7 +294,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -321,7 +321,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -353,7 +353,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepEval(FROZEN_MODEL) cls.dp_compressed = DeepEval(COMPRESSED_MODEL) cls.coords = np.array( @@ -381,7 +381,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -402,7 +402,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -429,7 +429,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -461,7 +461,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepEval(FROZEN_MODEL_ET) cls.dp_compressed = DeepEval(COMPRESSED_MODEL_ET) cls.coords = np.array( @@ -490,7 +490,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(INPUT_ET) _file_delete(FROZEN_MODEL_ET) _file_delete(COMPRESSED_MODEL_ET) @@ -511,7 +511,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -526,7 +526,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -547,7 +547,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -574,7 +574,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -607,7 +607,7 @@ def test_2frame_atm(self): class TestSkipNeighborStat(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepEval(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) cls.dp_compressed = DeepEval(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) cls.coords = np.array( @@ -635,7 +635,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -650,7 +650,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -671,7 +671,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -698,7 +698,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py index 0f0224006d..a59d6f8e54 100644 --- a/source/tests/pt/test_multitask.py +++ b/source/tests/pt/test_multitask.py @@ -36,7 +36,7 @@ ) -def setUpModule(): +def setUpModule() -> None: global multitask_template multitask_template_json = str(Path(__file__).parent / "water/multitask.json") with open(multitask_template_json) as f: @@ -44,7 +44,7 @@ def setUpModule(): class MultiTaskTrainTest: - def test_multitask_train(self): + def test_multitask_train(self) -> None: # test multitask training self.config = update_deepmd_input(self.config, warning=True) self.config = normalize(self.config, multi_task=True) @@ -174,7 +174,7 @@ def test_multitask_train(self): trainer_finetune.run() self.tearDown() - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) @@ -185,7 +185,7 @@ def tearDown(self): class TestMultiTaskSeA(unittest.TestCase, MultiTaskTrainTest): - def setUp(self): + def setUp(self) -> None: multitask_se_e2_a = deepcopy(multitask_template) multitask_se_e2_a["model"]["shared_dict"]["my_descriptor"] = model_se_e2_a[ "descriptor" @@ -224,7 +224,7 @@ def tearDown(self) -> None: class TestMultiTaskDPA1(unittest.TestCase, MultiTaskTrainTest): - def setUp(self): + def setUp(self) -> None: multitask_DPA1 = deepcopy(multitask_template) multitask_DPA1["model"]["shared_dict"]["my_descriptor"] = model_dpa1[ "descriptor" @@ -263,7 +263,7 @@ def tearDown(self) -> None: class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest): - def setUp(self): + def setUp(self) -> None: multitask_DPA2 = deepcopy(multitask_template) multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2[ "descriptor" @@ -302,7 +302,7 @@ def tearDown(self) -> None: class TestMultiTaskDPA2Tebd(unittest.TestCase, MultiTaskTrainTest): - def setUp(self): + def setUp(self) -> None: multitask_DPA2 = deepcopy(multitask_template) multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2tebd[ "descriptor" diff --git a/source/tests/pt/test_sampler.py b/source/tests/pt/test_sampler.py index 4f1091c936..420aacd8be 100644 --- a/source/tests/pt/test_sampler.py +++ b/source/tests/pt/test_sampler.py @@ -28,7 +28,7 @@ class TestSampler(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: with open(str(Path(__file__).parent / "water/se_e2_a.json")) as fin: content = fin.read() config = json.loads(content) @@ -54,7 +54,7 @@ def setUp(self): tf_random.seed(10) self.dp_dataset = DeepmdDataSystem(self.systems, self.batch_size, 1, self.rcut) - def test_sampler_debug_info(self): + def test_sampler_debug_info(self) -> None: dataloader = DataLoader( self.my_dataset, sampler=get_weighted_sampler(self.my_dataset, prob_style="prob_sys_size"), @@ -71,7 +71,7 @@ def test_sampler_debug_info(self): frame = self.my_dataset.systems[sid].__getitem__(fid) self.assertTrue(np.allclose(coord, frame["coord"])) - def test_auto_prob_uniform(self): + def test_auto_prob_uniform(self) -> None: auto_prob_style = "prob_uniform" sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style) my_probs = np.array(sampler.weights) @@ -79,7 +79,7 @@ def test_auto_prob_uniform(self): dp_probs = np.array(self.dp_dataset.sys_probs) self.assertTrue(np.allclose(my_probs, dp_probs)) - def test_auto_prob_sys_size(self): + def test_auto_prob_sys_size(self) -> None: auto_prob_style = "prob_sys_size" sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style) my_probs = np.array(sampler.weights) @@ -87,7 +87,7 @@ def test_auto_prob_sys_size(self): dp_probs = np.array(self.dp_dataset.sys_probs) self.assertTrue(np.allclose(my_probs, dp_probs)) - def test_auto_prob_sys_size_ext(self): + def test_auto_prob_sys_size_ext(self) -> None: auto_prob_style = "prob_sys_size;0:1:0.2;1:3:0.8" sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style) my_probs = np.array(sampler.weights) @@ -95,7 +95,7 @@ def test_auto_prob_sys_size_ext(self): dp_probs = np.array(self.dp_dataset.sys_probs) self.assertTrue(np.allclose(my_probs, dp_probs)) - def test_sys_probs(self): + def test_sys_probs(self) -> None: sys_probs = [0.1, 0.4, 0.5] sampler = get_weighted_sampler( self.my_dataset, prob_style=sys_probs, sys_prob=True diff --git a/source/tests/pt/test_tabulate.py b/source/tests/pt/test_tabulate.py index 5a812f420c..164819408f 100644 --- a/source/tests/pt/test_tabulate.py +++ b/source/tests/pt/test_tabulate.py @@ -19,16 +19,16 @@ ) -def setUpModule(): +def setUpModule() -> None: tf.compat.v1.enable_eager_execution() -def tearDownModule(): +def tearDownModule() -> None: tf.compat.v1.disable_eager_execution() class TestDPTabulate(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.w = np.array( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1.0, 1.1, 1.2]], dtype=np.float64, @@ -45,7 +45,7 @@ def setUp(self): self.y = np.tanh(self.xbar) - def test_ops(self): + def test_ops(self) -> None: dy_tf = op_module.unaggregated_dy_dx_s( tf.constant(self.y, dtype="double"), tf.constant(self.w, dtype="double"), diff --git a/source/tests/pt/test_tabulate_fusion_se_a.py b/source/tests/pt/test_tabulate_fusion_se_a.py index f811463ff0..8861f0564e 100644 --- a/source/tests/pt/test_tabulate_fusion_se_a.py +++ b/source/tests/pt/test_tabulate_fusion_se_a.py @@ -18,7 +18,7 @@ @parameterized((torch.float64, torch.float32)) @unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") class TestTabulateFusionSeAOp(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (dtype,) = self.param if dtype == torch.float64: self.prec = 1e-10 @@ -1418,7 +1418,7 @@ def setUp(self): device=env.DEVICE, ).reshape(4, 4, 4) - def test_forward(self): + def test_forward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_a( self.table_tensor, @@ -1441,7 +1441,7 @@ def test_forward(self): rtol=self.prec, ) - def test_backward(self): + def test_backward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_a( self.table_tensor, diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py index 99c6e5c7d0..09b06dcc89 100644 --- a/source/tests/pt/test_tabulate_fusion_se_atten.py +++ b/source/tests/pt/test_tabulate_fusion_se_atten.py @@ -18,7 +18,7 @@ @parameterized((torch.float64, torch.float32)) @unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") class TestTabulateFusionSeAttenOp(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (dtype,) = self.param if dtype == torch.float64: self.prec = 1e-10 @@ -1553,7 +1553,7 @@ def setUp(self): device=env.DEVICE, ).reshape(4, 4, 4) - def test_forward(self): + def test_forward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_atten( self.table_tensor, @@ -1578,7 +1578,7 @@ def test_forward(self): rtol=self.prec, ) - def test_backward(self): + def test_backward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_atten( self.table_tensor, diff --git a/source/tests/pt/test_tabulate_fusion_se_r.py b/source/tests/pt/test_tabulate_fusion_se_r.py index 9abfc72dc3..34d8dbf182 100644 --- a/source/tests/pt/test_tabulate_fusion_se_r.py +++ b/source/tests/pt/test_tabulate_fusion_se_r.py @@ -18,7 +18,7 @@ @parameterized((torch.float64, torch.float32)) @unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") class TestTabulateFusionSeAOp(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (dtype,) = self.param if dtype == torch.float64: self.prec = 1e-10 @@ -1276,7 +1276,7 @@ def setUp(self): device=env.DEVICE, ).reshape(4, 4) - def test_forward(self): + def test_forward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_r( self.table_tensor, @@ -1298,7 +1298,7 @@ def test_forward(self): rtol=self.prec, ) - def test_backward(self): + def test_backward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_r( self.table_tensor, diff --git a/source/tests/pt/test_tabulate_fusion_se_t.py b/source/tests/pt/test_tabulate_fusion_se_t.py index 3fe34c702e..50654d557d 100644 --- a/source/tests/pt/test_tabulate_fusion_se_t.py +++ b/source/tests/pt/test_tabulate_fusion_se_t.py @@ -18,7 +18,7 @@ @parameterized((torch.float64, torch.float32)) @unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") class TestTabulateFusionSeTOp(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: (dtype,) = self.param if dtype == torch.float64: self.prec = 1e-10 @@ -1678,7 +1678,7 @@ def setUp(self): device=env.DEVICE, ).reshape(4, 4, 4) - def test_forward(self): + def test_forward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_t( self.table_tensor, @@ -1701,7 +1701,7 @@ def test_forward(self): rtol=self.prec, ) - def test_backward(self): + def test_backward(self) -> None: # Call the forward function forward_result = torch.ops.deepmd.tabulate_fusion_se_t( self.table_tensor, diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index a7fcedcede..1fbd01c39f 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -30,7 +30,7 @@ class DPTrainTest: - def test_dp_train(self): + def test_dp_train(self) -> None: # test training from scratch trainer = get_trainer(deepcopy(self.config)) trainer.run() @@ -100,7 +100,7 @@ def test_dp_train(self): trainer_finetune_empty.run() trainer_finetune_random.run() - def test_trainable(self): + def test_trainable(self) -> None: fix_params = deepcopy(self.config) fix_params["model"]["descriptor"]["trainable"] = False fix_params["model"]["fitting_net"]["trainable"] = False @@ -128,7 +128,7 @@ def test_trainable(self): model_dict_before_training[key], model_dict_after_training[key] ) - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) @@ -139,7 +139,7 @@ def tearDown(self): class TestEnergyModelSeA(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -155,7 +155,7 @@ def tearDown(self) -> None: class TestDOSModelSeA(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "dos/input.json") with open(input_json) as f: self.config = json.load(f) @@ -173,7 +173,7 @@ def tearDown(self) -> None: class TestEnergyZBLModelSeA(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/zbl.json") with open(input_json) as f: self.config = json.load(f) @@ -191,7 +191,7 @@ def tearDown(self) -> None: class TestFparam(unittest.TestCase, DPTrainTest): """Test if `fparam` can be loaded correctly.""" - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -211,7 +211,7 @@ def tearDown(self) -> None: class TestEnergyModelDPA1(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -227,7 +227,7 @@ def tearDown(self) -> None: class TestEnergyModelDPA2(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -244,7 +244,7 @@ def tearDown(self) -> None: @unittest.skip("hybrid not supported at the moment") class TestEnergyModelHybrid(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -260,7 +260,7 @@ def tearDown(self) -> None: class TestDipoleModelSeA(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -289,7 +289,7 @@ def tearDown(self) -> None: class TestDipoleModelDPA1(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -318,7 +318,7 @@ def tearDown(self) -> None: class TestDipoleModelDPA2(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -347,7 +347,7 @@ def tearDown(self) -> None: class TestPolarModelSeA(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -381,7 +381,7 @@ def tearDown(self) -> None: class TestPolarModelDPA1(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -415,7 +415,7 @@ def tearDown(self) -> None: class TestPolarModelDPA2(unittest.TestCase, DPTrainTest): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json") with open(input_json) as f: self.config = json.load(f) @@ -449,7 +449,7 @@ def tearDown(self) -> None: class TestPropFintuFromEnerModel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") with open(input_json) as f: self.config = json.load(f) @@ -471,7 +471,7 @@ def setUp(self): self.config_property["training"]["numb_steps"] = 1 self.config_property["training"]["save_freq"] = 1 - def test_dp_train(self): + def test_dp_train(self) -> None: # test training from scratch trainer = get_trainer(deepcopy(self.config)) trainer.run() @@ -506,7 +506,7 @@ def test_dp_train(self): # check running trainer_finetune.run() - def tearDown(self): + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): os.remove(f) diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py index 810351b74d..c87c3330e9 100644 --- a/source/tests/pt/test_update_sel.py +++ b/source/tests/pt/test_update_sel.py @@ -31,7 +31,7 @@ def setUp(self) -> None: return super().setUp() @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_one_sel(self, sel_mock): + def test_update_one_sel(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") @@ -44,7 +44,7 @@ def test_update_one_sel(self, sel_mock): self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_hybrid(self, sel_mock): + def test_update_sel_hybrid(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] jdata = { @@ -75,7 +75,7 @@ def test_update_sel_hybrid(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel(self, sel_mock): + def test_update_sel(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] jdata = { @@ -90,7 +90,7 @@ def test_update_sel(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_auto(self, sel_mock): + def test_update_sel_atten_auto(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -117,7 +117,7 @@ def test_update_sel_atten_auto(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_int(self, sel_mock): + def test_update_sel_atten_int(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -144,7 +144,7 @@ def test_update_sel_atten_int(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_list(self, sel_mock): + def test_update_sel_atten_list(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -171,7 +171,7 @@ def test_update_sel_atten_list(self, sel_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_dpa2_auto(self, sel_mock): + def test_update_sel_dpa2_auto(self, sel_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] jdata = { @@ -213,7 +213,7 @@ def test_update_sel_dpa2_auto(self, sel_mock): jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_frozen(self): + def test_skip_frozen(self) -> None: jdata = { "model": { "type": "frozen", @@ -224,7 +224,7 @@ def test_skip_frozen(self): jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_wrap_up_4(self): + def test_wrap_up_4(self) -> None: self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) diff --git a/source/tests/pt/test_utils.py b/source/tests/pt/test_utils.py index 21fad0ecea..42bf2679a8 100644 --- a/source/tests/pt/test_utils.py +++ b/source/tests/pt/test_utils.py @@ -15,7 +15,7 @@ class TestCvt(unittest.TestCase): - def test_to_numpy(self): + def test_to_numpy(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) foo = rng.normal([3, 4]) for ptp, npp in zip( diff --git a/source/tests/tf/common.py b/source/tests/tf/common.py index 2b912c7a10..7e93af94c8 100644 --- a/source/tests/tf/common.py +++ b/source/tests/tf/common.py @@ -39,14 +39,14 @@ def j_loader(filename): return dp_j_loader(tests_path / filename) -def del_data(): +def del_data() -> None: if os.path.isdir("system"): shutil.rmtree("system") if os.path.isdir("system_mixed_type"): shutil.rmtree("system_mixed_type") -def gen_data_type_specific(nframes=1, dim_fparam=2): +def gen_data_type_specific(nframes=1, dim_fparam=2) -> None: tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam) sys = dpdata.LabeledSystem() sys.data["atom_names"] = ["foo", "bar"] @@ -67,7 +67,7 @@ def gen_data_type_specific(nframes=1, dim_fparam=2): ) -def gen_data_mixed_type(nframes=1, dim_fparam=2): +def gen_data_mixed_type(nframes=1, dim_fparam=2) -> None: tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam) sys = dpdata.LabeledSystem() real_type_map = ["foo", "bar"] @@ -94,7 +94,7 @@ def gen_data_mixed_type(nframes=1, dim_fparam=2): ) -def gen_data_virtual_type(nframes=1, nghost=4, dim_fparam=2): +def gen_data_virtual_type(nframes=1, nghost=4, dim_fparam=2) -> None: tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam) sys = dpdata.LabeledSystem() real_type_map = ["foo", "bar"] @@ -145,7 +145,7 @@ def gen_data_virtual_type(nframes=1, nghost=4, dim_fparam=2): ) -def gen_data(nframes=1, mixed_type=False, virtual_type=False, dim_fparam=2): +def gen_data(nframes=1, mixed_type=False, virtual_type=False, dim_fparam=2) -> None: if not mixed_type: gen_data_type_specific(nframes, dim_fparam=dim_fparam) elif virtual_type: @@ -155,7 +155,9 @@ def gen_data(nframes=1, mixed_type=False, virtual_type=False, dim_fparam=2): class Data: - def __init__(self, rand_pert=0.1, seed=1, box_scale=20, nframes=1, dim_fparam=2): + def __init__( + self, rand_pert=0.1, seed=1, box_scale=20, nframes=1, dim_fparam=2 + ) -> None: coord = [ [0.0, 0.0, 0.1], [1.1, 0.0, 0.1], @@ -282,7 +284,7 @@ def get_test_box_data(self, hh, rand_pert=0.1): def force_test( inter, testCase, places=global_default_places, hh=global_default_fv_hh, suffix="" -): +) -> None: # set weights w0 = np.ones(inter.ndescrpt) inter.net_w_i = np.copy(w0) @@ -348,7 +350,7 @@ def comp_vol(box): def virial_test( inter, testCase, places=global_default_places, hh=global_default_fv_hh, suffix="" -): +) -> None: # set weights w0 = np.ones(inter.ndescrpt) inter.net_w_i = np.copy(w0) @@ -385,7 +387,7 @@ def virial_test( def force_dw_test( inter, testCase, places=global_default_places, hh=global_default_dw_hh, suffix="" -): +) -> None: dcoord, dbox, dtype = inter.data.get_data() defield = inter.data.efield feed_dict_test0 = { @@ -451,7 +453,7 @@ def force_dw_test( def virial_dw_test( inter, testCase, places=global_default_places, hh=global_default_dw_hh, suffix="" -): +) -> None: dcoord, dbox, dtype = inter.data.get_data() defield = inter.data.efield feed_dict_test0 = { @@ -646,7 +648,7 @@ class DataSets: This class is not maintained any more. """ - def __init__(self, sys_path, set_prefix, seed=None, shuffle_test=True): + def __init__(self, sys_path, set_prefix, seed=None, shuffle_test=True) -> None: self.dirs = glob.glob(os.path.join(sys_path, set_prefix + ".*")) self.dirs.sort() # load atom type @@ -824,14 +826,14 @@ def load_set(self, set_name, shuffle=True): data[ii] = data[ii][:, self.idx3_map] return data - def load_batch_set(self, set_name): + def load_batch_set(self, set_name) -> None: self.batch_set = self.load_set(set_name, True) self.reset_iter() - def load_test_set(self, set_name, shuffle_test): + def load_test_set(self, set_name, shuffle_test) -> None: self.test_set = self.load_set(set_name, shuffle_test) - def reset_iter(self): + def reset_iter(self) -> None: self.iterator = 0 self.set_count += 1 @@ -914,7 +916,9 @@ class DataSystem: This class is not maintained any more. """ - def __init__(self, systems, set_prefix, batch_size, test_size, rcut, run_opt=None): + def __init__( + self, systems, set_prefix, batch_size, test_size, rcut, run_opt=None + ) -> None: self.system_dirs = systems self.nsystems = len(self.system_dirs) self.batch_size = batch_size @@ -1017,7 +1021,7 @@ def format_name_length(self, name, width): name = "-- " + name return name - def print_summary(self): + def print_summary(self) -> None: tmp_msg = "" # width 65 sys_width = 42 diff --git a/source/tests/tf/test_activation_fn_gelu.py b/source/tests/tf/test_activation_fn_gelu.py index 9be0885b74..7511357e86 100644 --- a/source/tests/tf/test_activation_fn_gelu.py +++ b/source/tests/tf/test_activation_fn_gelu.py @@ -15,7 +15,7 @@ class TestGelu(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.places = 6 self.sess = self.cached_session().__enter__() self.inputs = tf.reshape( @@ -64,7 +64,7 @@ def setUp(self): ], ] - def test_activation_function_gelu_custom(self): + def test_activation_function_gelu_custom(self) -> None: network_size = [2, 4, 8] out = embedding_net( self.inputs, @@ -79,7 +79,7 @@ def test_activation_function_gelu_custom(self): myout = self.sess.run(out) np.testing.assert_almost_equal(self.refout, myout, self.places) - def test_activation_function_gelu_tensorflow(self): + def test_activation_function_gelu_tensorflow(self) -> None: network_size = [2, 4, 8] out = embedding_net( self.inputs, diff --git a/source/tests/tf/test_adjust_sel.py b/source/tests/tf/test_adjust_sel.py index dc94f81458..c0df54eefa 100644 --- a/source/tests/tf/test_adjust_sel.py +++ b/source/tests/tf/test_adjust_sel.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -72,7 +72,7 @@ def _init_models(): class TestDeepPotAAdjustSel(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: INPUT, FROZEN_MODEL, DECREASED_MODEL, INCREASED_MODEL = _init_models() cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_decreased = DeepPot(DECREASED_MODEL) @@ -102,7 +102,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -121,7 +121,7 @@ def test_attrs(self): self.assertEqual(self.dp_increased.get_dim_fparam(), 0) self.assertEqual(self.dp_increased.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -151,7 +151,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee2, default_places) np.testing.assert_almost_equal(vv0, vv2, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -191,7 +191,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee2, default_places) np.testing.assert_almost_equal(vv0, vv2, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -234,7 +234,7 @@ def test_2frame_atm(self): np.testing.assert_almost_equal(ee0, ee2, default_places) np.testing.assert_almost_equal(vv0, vv2, default_places) - def test_descriptor(self): + def test_descriptor(self) -> None: dd0 = self.dp_original.eval_descriptor(self.coords, self.box, self.atype) dd1 = self.dp_decreased.eval_descriptor(self.coords, self.box, self.atype) dd2 = self.dp_increased.eval_descriptor(self.coords, self.box, self.atype) diff --git a/source/tests/tf/test_auto_batch_size.py b/source/tests/tf/test_auto_batch_size.py index 3316e186b6..84f15e58e4 100644 --- a/source/tests/tf/test_auto_batch_size.py +++ b/source/tests/tf/test_auto_batch_size.py @@ -19,7 +19,7 @@ def oom(self, batch_size, start_index): return batch_size, np.zeros((batch_size, 2)) @unittest.mock.patch("tensorflow.compat.v1.test.is_gpu_available") - def test_execute_oom_gpu(self, mock_is_gpu_available): + def test_execute_oom_gpu(self, mock_is_gpu_available) -> None: mock_is_gpu_available.return_value = True # initial batch size 256 = 128 * 2 auto_batch_size = AutoBatchSize(256, 2.0) @@ -46,7 +46,9 @@ def test_execute_oom_gpu(self, mock_is_gpu_available): @unittest.mock.patch("tensorflow.compat.v1.test.is_gpu_available") @unittest.mock.patch("tensorflow.compat.v1.config.experimental.get_visible_devices") - def test_execute_oom_cpu(self, mock_is_gpu_available, mock_get_visible_devices): + def test_execute_oom_cpu( + self, mock_is_gpu_available, mock_get_visible_devices + ) -> None: mock_is_gpu_available.return_value = False mock_get_visible_devices.return_value = [] # initial batch size 256 = 128 * 2, nb is always 128 @@ -68,7 +70,7 @@ def test_execute_oom_cpu(self, mock_is_gpu_available, mock_get_visible_devices): self.assertEqual(result.shape, (128, 2)) @unittest.mock.patch.dict(os.environ, {"DP_INFER_BATCH_SIZE": "256"}, clear=True) - def test_execute_oom_environment_variables(self): + def test_execute_oom_environment_variables(self) -> None: # DP_INFER_BATCH_SIZE = 256 = 128 * 2, nb is always 128 auto_batch_size = AutoBatchSize(999, 2.0) nb, result = auto_batch_size.execute(self.oom, 1, 2) @@ -87,7 +89,7 @@ def test_execute_oom_environment_variables(self): self.assertEqual(nb, 128) self.assertEqual(result.shape, (128, 2)) - def test_execute_all(self): + def test_execute_all(self) -> None: dd1 = np.zeros((10000, 2, 1)) auto_batch_size = AutoBatchSize(256, 2.0) dd2 = auto_batch_size.execute_all(np.array, 10000, 2, dd1) diff --git a/source/tests/tf/test_cluster.py b/source/tests/tf/test_cluster.py index ea90e1ea6d..1b2eceebfc 100644 --- a/source/tests/tf/test_cluster.py +++ b/source/tests/tf/test_cluster.py @@ -12,7 +12,7 @@ class FakePopen: - def __init__(self, stdout=b"", stderr=b"", returncode=0): + def __init__(self, stdout=b"", stderr=b"", returncode=0) -> None: self._stdout = stdout self._stderr = stderr self._returncode = returncode @@ -28,7 +28,7 @@ def returncode(self): class TestGPU(unittest.TestCase): @mock.patch("tensorflow.compat.v1.test.is_built_with_cuda") @mock.patch("subprocess.Popen") - def test_none(self, mock_Popen, mock_is_built_with_cuda): + def test_none(self, mock_Popen, mock_is_built_with_cuda) -> None: mock_Popen.return_value.__enter__.return_value = FakePopen(b"0", b"") mock_is_built_with_cuda.return_value = True gpus = local.get_gpus() @@ -36,7 +36,7 @@ def test_none(self, mock_Popen, mock_is_built_with_cuda): @mock.patch("tensorflow.compat.v1.test.is_built_with_cuda") @mock.patch("subprocess.Popen") - def test_valid(self, mock_Popen, mock_is_built_with_cuda): + def test_valid(self, mock_Popen, mock_is_built_with_cuda) -> None: mock_Popen.return_value.__enter__.return_value = FakePopen(b"2", b"") mock_is_built_with_cuda.return_value = True gpus = local.get_gpus() @@ -44,7 +44,7 @@ def test_valid(self, mock_Popen, mock_is_built_with_cuda): @mock.patch("tensorflow.compat.v1.test.is_built_with_cuda") @mock.patch("subprocess.Popen") - def test_error(self, mock_Popen, mock_is_built_with_cuda): + def test_error(self, mock_Popen, mock_is_built_with_cuda) -> None: mock_Popen.return_value.__enter__.return_value = FakePopen( stderr=b"!", returncode=1 ) @@ -55,7 +55,7 @@ def test_error(self, mock_Popen, mock_is_built_with_cuda): @mock.patch("tensorflow.compat.v1.test.is_built_with_rocm", create=True) @mock.patch("tensorflow.compat.v1.test.is_built_with_cuda") - def test_cpu(self, mock_is_built_with_cuda, mock_is_built_with_rocm): + def test_cpu(self, mock_is_built_with_cuda, mock_is_built_with_rocm) -> None: mock_is_built_with_cuda.return_value = False mock_is_built_with_rocm.return_value = False gpus = local.get_gpus() @@ -64,7 +64,7 @@ def test_cpu(self, mock_is_built_with_cuda, mock_is_built_with_rocm): class TestLocal(unittest.TestCase): @mock.patch("socket.gethostname") - def test_resource(self, mock_gethostname): + def test_resource(self, mock_gethostname) -> None: mock_gethostname.return_value = kHostName nodename, nodelist, _ = local.get_resource() self.assertEqual(nodename, kHostName) diff --git a/source/tests/tf/test_common.py b/source/tests/tf/test_common.py index 95948f29bb..d4a1539f85 100644 --- a/source/tests/tf/test_common.py +++ b/source/tests/tf/test_common.py @@ -29,7 +29,7 @@ def relpath(path_to, path_from): class TestCommonExpandSysDir(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.match_file = Path("type.raw") Path("test_sys").mkdir() self.dir = Path("test_sys") @@ -56,10 +56,10 @@ def setUp(self): ] self.expected_out.sort() - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree("test_sys") - def test_expand(self): + def test_expand(self) -> None: ret = expand_sys_str("test_sys") ret.sort() self.assertEqual(ret, self.expected_out) @@ -74,7 +74,7 @@ def precision(self): return tf.float64 return tf.float32 - def test_cast_precision(self): + def test_cast_precision(self) -> None: x = tf.zeros(1, dtype=GLOBAL_TF_FLOAT_PRECISION) y = tf.zeros(1, dtype=tf.int64) self.assertEqual(x.dtype, GLOBAL_TF_FLOAT_PRECISION) diff --git a/source/tests/tf/test_compat_input.py b/source/tests/tf/test_compat_input.py index f7c605380c..4e60b6bf3e 100644 --- a/source/tests/tf/test_compat_input.py +++ b/source/tests/tf/test_compat_input.py @@ -13,25 +13,25 @@ class TestConvertInput(unittest.TestCase): - def test_convert_smth(self): + def test_convert_smth(self) -> None: jdata0 = j_loader(os.path.join("compat_inputs", "water_se_a_v0.json")) jdata1 = j_loader(os.path.join("compat_inputs", "water_se_a_v1.json")) jdata = convert_input_v0_v1(jdata0, warning=False, dump=None) self.assertEqual(jdata, jdata1) - def test_convert_nonsmth(self): + def test_convert_nonsmth(self) -> None: jdata0 = j_loader(os.path.join("compat_inputs", "water_v0.json")) jdata1 = j_loader(os.path.join("compat_inputs", "water_v1.json")) jdata = convert_input_v0_v1(jdata0, warning=False, dump=None) self.assertEqual(jdata, jdata1) - def test_convert_v1_v2(self): + def test_convert_v1_v2(self) -> None: jdata0 = j_loader(os.path.join("compat_inputs", "water_v1.json")) jdata1 = j_loader(os.path.join("compat_inputs", "water_v2.json")) jdata = convert_input_v1_v2(jdata0, warning=False, dump=None) self.assertDictAlmostEqual(jdata, jdata1) - def assertDictAlmostEqual(self, d1, d2, msg=None, places=7): + def assertDictAlmostEqual(self, d1, d2, msg=None, places=7) -> None: self.assertEqual(d1.keys(), d2.keys()) for kk, vv in d1.items(): if isinstance(vv, dict): @@ -39,7 +39,7 @@ def assertDictAlmostEqual(self, d1, d2, msg=None, places=7): else: self.assertAlmostEqual(d1[kk], d2[kk], places=places, msg=msg) - def test_json_yaml_equal(self): + def test_json_yaml_equal(self) -> None: inputs = ("water_v1", "water_se_a_v1") for i in inputs: diff --git a/source/tests/tf/test_compressed_training.py b/source/tests/tf/test_compressed_training.py index 998ef8cb59..3cfe339d3e 100644 --- a/source/tests/tf/test_compressed_training.py +++ b/source/tests/tf/test_compressed_training.py @@ -45,7 +45,7 @@ def setUp(self) -> None: with open(self.input_file, "w") as fp: json.dump(jdata, fp, indent=4) - def test_compressed_training(self): + def test_compressed_training(self) -> None: run_dp(f"dp train {self.input_file}") run_dp(f"dp freeze -c {self.checkpoint_dir} -o {self.frozen_model}") run_dp(f"dp compress -i {self.frozen_model} -o {self.compressed_model}") diff --git a/source/tests/tf/test_data_large_batch.py b/source/tests/tf/test_data_large_batch.py index d9bb00de40..858a6e69b3 100644 --- a/source/tests/tf/test_data_large_batch.py +++ b/source/tests/tf/test_data_large_batch.py @@ -39,11 +39,11 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestDataLargeBatch(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data(mixed_type=True) self.filename = __file__ - def test_data_mixed_type(self): + def test_data_mixed_type(self) -> None: jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) @@ -242,7 +242,7 @@ def test_data_mixed_type(self): np.testing.assert_almost_equal(v, refv, places) sess.close() - def test_stripped_data_mixed_type(self): + def test_stripped_data_mixed_type(self) -> None: jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) @@ -441,7 +441,7 @@ def test_stripped_data_mixed_type(self): np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) - def test_compressible_data_mixed_type(self): + def test_compressible_data_mixed_type(self) -> None: jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index 890de86034..12471f3817 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -40,15 +40,15 @@ class TestDataModifier(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: # with tf.variable_scope('load', reuse = False) : tf.reset_default_graph() self._setUp() - def tearDown(self): + def tearDown(self) -> None: tf.reset_default_graph() - def _setUp(self): + def _setUp(self) -> None: run_opt = RunOptions( restart=None, init_model=None, log_path=None, log_level=30, mpi_log="master" ) @@ -91,11 +91,11 @@ def _setUp(self): with tf.gfile.GFile(output_graph, "wb") as f: f.write(output_graph_def.SerializeToString()) - def test_fv(self): + def test_fv(self) -> None: # with tf.variable_scope('load', reuse = False) : self._test_fv() - def _test_fv(self): + def _test_fv(self) -> None: dcm = DipoleChargeModifier( str(tests_path / os.path.join(modifier_datapath, "dipole.pb")), [-8], diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index f1822a9346..002b4f5746 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -41,19 +41,19 @@ class TestDataModifier(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: # with tf.variable_scope('load', reuse = False) : tf.reset_default_graph() self._setUp() - def tearDown(self): + def tearDown(self) -> None: tf.reset_default_graph() if os.path.isdir(os.path.join(modifier_datapath, "sys_test_0")): shutil.rmtree(os.path.join(modifier_datapath, "sys_test_0")) if os.path.isfile(os.path.join(modifier_datapath, "dipole.pb")): os.remove(os.path.join(modifier_datapath, "dipole.pb")) - def _setUp(self): + def _setUp(self) -> None: run_opt = RunOptions( restart=None, init_model=None, log_path=None, log_level=30, mpi_log="master" ) @@ -94,7 +94,7 @@ def _setUp(self): with tf.gfile.GFile(output_graph, "wb") as f: f.write(output_graph_def.SerializeToString()) - def _setUp_data(self): + def _setUp_data(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) jdata = self._setUp_jdata() # sys0 @@ -129,7 +129,7 @@ def _setUp_data(self): self.sel_mask0 = np.isin(self.atom_types0, self.sel_type) self.sel_mask1 = np.isin(self.atom_types1, self.sel_type) - def _write_sys_data(self, dirname, atom_types, coords, dipoles, box): + def _write_sys_data(self, dirname, atom_types, coords, dipoles, box) -> None: os.makedirs(dirname, exist_ok=True) os.makedirs(dirname + "/set.0", exist_ok=True) np.savetxt(os.path.join(dirname, "type.raw"), atom_types, fmt="%d") @@ -182,7 +182,7 @@ def _setUp_jdata(self): } return jdata - def test_z_dipole(self): + def test_z_dipole(self) -> None: dd = DeepDipole(os.path.join(modifier_datapath, "dipole.pb")) dv0 = dd.eval(self.coords0, self.box0, self.atom_types0)[:, self.sel_mask0] @@ -195,7 +195,7 @@ def test_z_dipole(self): np.testing.assert_almost_equal(dv01, dv1, err_msg="dipole dose not match") - def test_modify(self): + def test_modify(self) -> None: dcm = DipoleChargeModifier( os.path.join(modifier_datapath, "dipole.pb"), [-1, -3], diff --git a/source/tests/tf/test_deepdipole.py b/source/tests/tf/test_deepdipole.py index 1e2f6dd45a..ca8f44fc7d 100644 --- a/source/tests/tf/test_deepdipole.py +++ b/source/tests/tf/test_deepdipole.py @@ -31,13 +31,13 @@ class TestDeepDipolePBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole.pbtxt")), "deepdipole.pb" ) cls.dp = DeepDipole("deepdipole.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -75,17 +75,17 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepdipole.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 4.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_sel_type(), [0]) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -95,7 +95,7 @@ def test_1frame_atm(self): # check values np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) dd = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] @@ -111,13 +111,13 @@ def test_2frame_atm(self): class TestDeepDipoleNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole.pbtxt")), "deepdipole.pb" ) cls.dp = DeepDipole("deepdipole.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -155,11 +155,11 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepdipole.pb") cls.dp = None - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: dd = self.dp.eval(self.coords, None, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -169,7 +169,7 @@ def test_1frame_atm(self): # check values np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) - def test_1frame_atm_large_box(self): + def test_1frame_atm_large_box(self) -> None: dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -186,14 +186,14 @@ def test_1frame_atm_large_box(self): ) class TestDeepDipoleNewPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole_new.pbtxt")), "deepdipole_new.pb", ) cls.dp = DeepDipole("deepdipole_new.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -460,17 +460,17 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepdipole_new.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 4.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_sel_type(), [0]) - def test_1frame_old(self): + def test_1frame_old(self) -> None: gt = self.dp.eval(self.coords, self.box, self.atype, atomic=False) # check shape of the returns nframes = 1 @@ -478,7 +478,7 @@ def test_1frame_old(self): # check values np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) - def test_1frame_old_atm(self): + def test_1frame_old_atm(self) -> None: at = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -488,7 +488,7 @@ def test_1frame_old_atm(self): # check values np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places) - def test_2frame_old_atm(self): + def test_2frame_old_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) at = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] @@ -501,7 +501,7 @@ def test_2frame_old_atm(self): expected_d = np.concatenate((self.expected_t, self.expected_t)) np.testing.assert_almost_equal(at.ravel(), expected_d, default_places) - def test_1frame_full(self): + def test_1frame_full(self) -> None: gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic=False) # check shape of the returns nframes = 1 @@ -514,7 +514,7 @@ def test_1frame_full(self): np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places) - def test_1frame_full_atm(self): + def test_1frame_full_atm(self) -> None: gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) @@ -545,7 +545,7 @@ def test_1frame_full_atm(self): vv.reshape([-1]), self.expected_gv.reshape([-1]), decimal=default_places ) - def test_1frame_full_atm_shuffle(self): + def test_1frame_full_atm_shuffle(self) -> None: i_sf = [2, 1, 3, 0, 5, 4] isel_sf = [1, 0] gt, ff, vv, at, av = self.dp.eval_full( @@ -588,7 +588,7 @@ def test_1frame_full_atm_shuffle(self): vv.reshape([-1]), self.expected_gv.reshape([-1]), decimal=default_places ) - def test_1frame_num_deriv(self): + def test_1frame_num_deriv(self) -> None: # numerical force num_f = -finite_difference( lambda coord: self.dp.eval( @@ -618,7 +618,7 @@ def test_1frame_num_deriv(self): num_v.reshape([-1]), self.expected_gv.reshape([-1]), atol=1e-5 ) - def test_2frame_full_atm(self): + def test_2frame_full_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) gt, ff, vv, at, av = self.dp.eval_full(coords2, box2, self.atype, atomic=True) @@ -661,14 +661,14 @@ def test_2frame_full_atm(self): ) class TestDeepDipoleFakePBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole_fake.pbtxt")), "deepdipole_fake.pb", ) cls.dp = DeepDipole("deepdipole_fake.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -958,18 +958,18 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepdipole_fake.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 2.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_sel_type().tolist(), [0, 1]) np.testing.assert_allclose(self.target_t, self.expected_t, atol=3e-2) - def test_1frame_full_atm(self): + def test_1frame_full_atm(self) -> None: gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) @@ -1000,7 +1000,7 @@ def test_1frame_full_atm(self): vv.reshape([-1]), self.expected_gv.reshape([-1]), decimal=default_places ) - def test_1frame_full_atm_shuffle(self): + def test_1frame_full_atm_shuffle(self) -> None: i_sf = [2, 1, 3, 0, 5, 4] isel_sf = i_sf gt, ff, vv, at, av = self.dp.eval_full( @@ -1050,7 +1050,7 @@ def test_1frame_full_atm_shuffle(self): ) class TestDeepDipoleNewPBCNeighborList(TestDeepDipoleNewPBC): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole_new.pbtxt")), "deepdipole_new.pb", @@ -1063,9 +1063,9 @@ def setUpClass(cls): ) @unittest.skip("multiple frames not supported") - def test_2frame_full_atm(self): + def test_2frame_full_atm(self) -> None: pass @unittest.skip("multiple frames not supported") - def test_2frame_old_atm(self): + def test_2frame_old_atm(self) -> None: pass diff --git a/source/tests/tf/test_deepdos.py b/source/tests/tf/test_deepdos.py index d94c2c3f2d..f65c6717fa 100644 --- a/source/tests/tf/test_deepdos.py +++ b/source/tests/tf/test_deepdos.py @@ -26,13 +26,13 @@ class TestDeepDOS(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdos.pbtxt")), "deepdos.pb" ) cls.dp = DeepDOS("deepdos.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 2.288635, @@ -832,17 +832,17 @@ def setUp(self): ) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepdos.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 1) self.assertAlmostEqual(self.dp.get_rcut(), 5.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["Si"]) self.assertEqual(self.dp.get_numb_dos(), 250) - def test_1frame_atomic(self): + def test_1frame_atomic(self) -> None: dd = self.dp.eval(self.coords, self.box, self.atype, atomic=True) # check shape of the returns nframes = 1 @@ -857,7 +857,7 @@ def test_1frame_atomic(self): np.testing.assert_almost_equal(ados_list[1], self.expected_ados_2, 4) np.testing.assert_almost_equal(dd[0].ravel(), self.expected_dos, 4) - def test_2frame_atomic(self): + def test_2frame_atomic(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) dd = self.dp.eval(coords2, box2, self.atype, atomic=True) diff --git a/source/tests/tf/test_deepmd_data.py b/source/tests/tf/test_deepmd_data.py index 0969a9baf1..87fdd6862d 100644 --- a/source/tests/tf/test_deepmd_data.py +++ b/source/tests/tf/test_deepmd_data.py @@ -27,7 +27,7 @@ class TestDataTypeSel(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) self.data_name = "test_data" os.makedirs(self.data_name, exist_ok=True) @@ -60,10 +60,10 @@ def setUp(self): self.value_2 = np.reshape(self.value_2, [self.nframes, 4]) np.save(path, self.value_2) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree(self.data_name) - def test_load_set_1(self): + def test_load_set_1(self) -> None: dd = DeepmdData(self.data_name).add( "value_1", 1, atomic=True, must=True, type_sel=[0] ) @@ -71,7 +71,7 @@ def test_load_set_1(self): self.assertEqual(data["value_1"].shape, (self.nframes, 2)) np.testing.assert_almost_equal(data["value_1"], self.value_1) - def test_load_set_2(self): + def test_load_set_2(self) -> None: dd = DeepmdData(self.data_name).add( "value_2", 1, atomic=True, must=True, type_sel=[1] ) @@ -81,7 +81,7 @@ def test_load_set_2(self): class TestData(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) self.data_name = "test_data" os.makedirs(self.data_name, exist_ok=True) @@ -157,10 +157,10 @@ def setUp(self): self.tensor_nsel = self.tensor_natoms[:, 1, :] np.save(path, self.tensor_nsel) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree(self.data_name) - def test_init(self): + def test_init(self) -> None: dd = DeepmdData(self.data_name) self.assertEqual(dd.idx_map[0], 1) self.assertEqual(dd.idx_map[1], 0) @@ -169,7 +169,7 @@ def test_init(self): dd.dirs, ["test_data/set.bar", "test_data/set.foo", "test_data/set.tar"] ) - def test_init_type_map(self): + def test_init_type_map(self) -> None: dd = DeepmdData(self.data_name, type_map=["bar", "foo", "tar"]) self.assertEqual(dd.idx_map[0], 0) self.assertEqual(dd.idx_map[1], 1) @@ -177,7 +177,7 @@ def test_init_type_map(self): self.assertEqual(dd.atom_type[1], 1) self.assertEqual(dd.type_map, ["bar", "foo", "tar"]) - def test_load_set(self): + def test_load_set(self) -> None: dd = ( DeepmdData(self.data_name) .add("test_atomic", 7, atomic=True, must=True) @@ -202,7 +202,7 @@ def test_load_set(self): self.assertEqual(data["find_test_null"], 0) self._comp_np_mat2(data["test_null"], self.test_null) - def test_shuffle(self): + def test_shuffle(self) -> None: dd = ( DeepmdData(self.data_name) .add("test_atomic", 7, atomic=True, must=True) @@ -215,7 +215,7 @@ def test_shuffle(self): self._comp_np_mat2(data_bk["test_atom"][idx, :], data["test_atom"]) self._comp_np_mat2(data_bk["test_frame"][idx, :], data["test_frame"]) - def test_shuffle_with_numb_copy(self): + def test_shuffle_with_numb_copy(self) -> None: path = os.path.join(self.data_name, "set.foo", "numb_copy.npy") prob = np.arange(self.nframes) np.save(path, prob) @@ -232,7 +232,7 @@ def test_shuffle_with_numb_copy(self): self._comp_np_mat2(data_bk["test_atom"][idx, :], data["test_atom"]) self._comp_np_mat2(data_bk["test_frame"][idx, :], data["test_frame"]) - def test_reduce(self): + def test_reduce(self) -> None: dd = DeepmdData(self.data_name).add("test_atomic", 7, atomic=True, must=True) dd.reduce("redu", "test_atomic") data = dd._load_set(os.path.join(self.data_name, "set.foo")) @@ -241,7 +241,7 @@ def test_reduce(self): self.assertEqual(data["find_redu"], 1) self._comp_np_mat2(data["redu"], self.redu_atomic) - def test_reduce_null(self): + def test_reduce_null(self) -> None: dd = DeepmdData(self.data_name).add("test_atomic_1", 7, atomic=True, must=False) dd.reduce("redu", "test_atomic_1") data = dd._load_set(os.path.join(self.data_name, "set.foo")) @@ -252,12 +252,12 @@ def test_reduce_null(self): self.assertEqual(data["find_redu"], 0) self._comp_np_mat2(data["redu"], np.zeros([self.nframes, 7])) - def test_load_null_must(self): + def test_load_null_must(self) -> None: dd = DeepmdData(self.data_name).add("test_atomic_1", 7, atomic=True, must=True) with self.assertRaises(RuntimeError): data = dd._load_set(os.path.join(self.data_name, "set.foo")) - def test_avg(self): + def test_avg(self) -> None: dd = DeepmdData(self.data_name).add("test_frame", 5, atomic=False, must=True) favg = dd.avg("test_frame") fcmp = np.average( @@ -268,7 +268,7 @@ def test_avg(self): ) np.testing.assert_almost_equal(favg, fcmp, places) - def test_check_batch_size(self): + def test_check_batch_size(self) -> None: dd = DeepmdData(self.data_name) ret = dd.check_batch_size(10) self.assertEqual(ret, (os.path.join(self.data_name, "set.bar"), 5)) @@ -277,7 +277,7 @@ def test_check_batch_size(self): ret = dd.check_batch_size(1) self.assertEqual(ret, None) - def test_check_test_size(self): + def test_check_test_size(self) -> None: dd = DeepmdData(self.data_name) ret = dd.check_test_size(10) self.assertEqual(ret, (os.path.join(self.data_name, "set.bar"), 5)) @@ -286,7 +286,7 @@ def test_check_test_size(self): ret = dd.check_test_size(1) self.assertEqual(ret, None) - def test_get_batch(self): + def test_get_batch(self) -> None: dd = DeepmdData(self.data_name) data = dd.get_batch(5) self._comp_np_mat2( @@ -305,7 +305,7 @@ def test_get_batch(self): data = dd.get_batch(5) self._comp_np_mat2(np.sort(data["coord"], axis=0), np.sort(self.coord, axis=0)) - def test_get_test(self): + def test_get_test(self) -> None: dd = DeepmdData(self.data_name) data = dd.get_test() expected_coord = np.concatenate( @@ -315,14 +315,14 @@ def test_get_test(self): np.sort(data["coord"], axis=0), np.sort(expected_coord, axis=0) ) - def test_get_nbatch(self): + def test_get_nbatch(self) -> None: dd = DeepmdData(self.data_name) nb = dd.get_numb_batch(1, 0) self.assertEqual(nb, 5) nb = dd.get_numb_batch(2, 0) self.assertEqual(nb, 2) - def test_get_tensor(self): + def test_get_tensor(self) -> None: dd_natoms = ( DeepmdData(self.data_name) .add( @@ -374,20 +374,20 @@ def test_get_tensor(self): data_nsel["tensor_natoms"], ) - def _comp_np_mat2(self, first, second): + def _comp_np_mat2(self, first, second) -> None: np.testing.assert_almost_equal(first, second, places) class TestH5Data(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.data_name = str(tests_path / "test.hdf5") - def test_init(self): + def test_init(self) -> None: dd = DeepmdData(self.data_name) self.assertEqual(dd.idx_map[0], 0) self.assertEqual(dd.type_map, ["X"]) self.assertEqual(dd.dirs[0], self.data_name + "#/set.000") - def test_get_batch(self): + def test_get_batch(self) -> None: dd = DeepmdData(self.data_name) data = dd.get_batch(5) diff --git a/source/tests/tf/test_deepmd_data_sys.py b/source/tests/tf/test_deepmd_data_sys.py index 893177ac4f..19abab1d21 100644 --- a/source/tests/tf/test_deepmd_data_sys.py +++ b/source/tests/tf/test_deepmd_data_sys.py @@ -27,7 +27,7 @@ class TestDataSystem(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) self.nsys = 4 self.nframes = [3, 6, 5, 4] @@ -56,12 +56,12 @@ def setUp(self): ) np.save(path, val) - def tearDown(self): + def tearDown(self) -> None: for ii in range(self.nsys): sys_name = "sys_%d" % ii shutil.rmtree(sys_name) - def test_ntypes(self): + def test_ntypes(self) -> None: batch_size = 3 test_size = 2 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) @@ -72,25 +72,25 @@ def test_ntypes(self): self.assertEqual(ds.get_nsystems(), self.nsys) self.assertEqual(list(ds.get_batch_size()), [batch_size] * 4) - def test_batch_size_5(self): + def test_batch_size_5(self) -> None: batch_size = "auto:5" test_size = 2 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) self.assertEqual(ds.batch_size, [2, 2, 1, 1]) - def test_batch_size_null(self): + def test_batch_size_null(self) -> None: batch_size = "auto:3" test_size = 2 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) self.assertEqual(ds.batch_size, [1, 1, 1, 1]) - def test_batch_size_raise(self): + def test_batch_size_raise(self) -> None: batch_size = "foo" test_size = 2 with self.assertRaises(RuntimeError): ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) - def test_get_test(self): + def test_get_test(self) -> None: batch_size = 3 test_size = 2 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) @@ -193,7 +193,7 @@ def test_get_test(self): 0.0, ) - def test_get_batch(self): + def test_get_batch(self) -> None: batch_size = 3 test_size = 2 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) @@ -412,7 +412,7 @@ def test_get_batch(self): 0.0, ) - def test_prob_sys_size_1(self): + def test_prob_sys_size_1(self) -> None: batch_size = 1 test_size = 1 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) @@ -436,7 +436,7 @@ def test_prob_sys_size_1(self): / float(self.nframes[2] * (self.nset) + shift), ) - def test_prob_sys_size_2(self): + def test_prob_sys_size_2(self) -> None: batch_size = 1 test_size = 1 ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0) @@ -464,7 +464,7 @@ def _idx_map(self, target, idx_map, ndof): target = target.reshape([-1, natoms * ndof]) return target - def _in_array(self, target, idx_map, ndof, array): + def _in_array(self, target, idx_map, ndof, array) -> None: target = self._idx_map(target, idx_map, ndof) all_find = [] for ii in array: @@ -476,7 +476,7 @@ def _in_array(self, target, idx_map, ndof, array): for idx, ii in enumerate(all_find): self.assertTrue(ii, msg="does not find frame %d in array" % idx) - def test_sys_prob_floating_point_error(self): + def test_sys_prob_floating_point_error(self) -> None: # test floating point error; See #1917 sys_probs = [ 0.010, @@ -513,7 +513,7 @@ def test_sys_prob_floating_point_error(self): ds = DeepmdDataSystem(self.sys_name, 3, 2, 2.0, sys_probs=sys_probs) self.assertEqual(ds.sys_probs.size, len(sys_probs)) - def test_get_mixed_batch(self): + def test_get_mixed_batch(self) -> None: """Test get_batch with mixed system.""" batch_size = "mixed:3" test_size = 2 diff --git a/source/tests/tf/test_deeppolar.py b/source/tests/tf/test_deeppolar.py index b4f3fe7d0d..090603fd4f 100644 --- a/source/tests/tf/test_deeppolar.py +++ b/source/tests/tf/test_deeppolar.py @@ -29,13 +29,13 @@ class TestDeepPolarPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deeppolar.pbtxt")), "deeppolar.pb" ) cls.dp = DeepPolar("deeppolar.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -85,17 +85,17 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deeppolar.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_sel_type(), [0]) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -105,7 +105,7 @@ def test_1frame_atm(self): # check values np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) dd = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] @@ -121,13 +121,13 @@ def test_2frame_atm(self): class TestDeepPolarNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deeppolar.pbtxt")), "deeppolar.pb" ) cls.dp = DeepPolar("deeppolar.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -177,11 +177,11 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deeppolar.pb") cls.dp = None - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: dd = self.dp.eval(self.coords, None, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -191,7 +191,7 @@ def test_1frame_atm(self): # check values np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) - def test_1frame_atm_large_box(self): + def test_1frame_atm_large_box(self) -> None: dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -208,14 +208,14 @@ def test_1frame_atm_large_box(self): ) class TestDeepPolarNewPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deeppolar_new.pbtxt")), "deeppolar_new.pb", ) cls.dp = DeepPolar("deeppolar_new.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -926,17 +926,17 @@ def setUp(self): self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deeppolar_new.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_sel_type(), [0]) - def test_1frame_old(self): + def test_1frame_old(self) -> None: gt = self.dp.eval(self.coords, self.box, self.atype, atomic=False) # check shape of the returns nframes = 1 @@ -944,7 +944,7 @@ def test_1frame_old(self): # check values np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) - def test_1frame_old_atm(self): + def test_1frame_old_atm(self) -> None: at = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -954,7 +954,7 @@ def test_1frame_old_atm(self): # check values np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places) - def test_2frame_old_atm(self): + def test_2frame_old_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) at = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] @@ -967,7 +967,7 @@ def test_2frame_old_atm(self): expected_d = np.concatenate((self.expected_t, self.expected_t)) np.testing.assert_almost_equal(at.ravel(), expected_d, default_places) - def test_1frame_full(self): + def test_1frame_full(self) -> None: gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic=False) # check shape of the returns nframes = 1 @@ -980,7 +980,7 @@ def test_1frame_full(self): np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places) - def test_1frame_full_atm(self): + def test_1frame_full_atm(self) -> None: gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) @@ -1012,7 +1012,7 @@ def test_1frame_full_atm(self): vv.reshape([-1]), self.expected_gv.reshape([-1]), decimal=default_places ) - def test_1frame_full_atm_shuffle(self): + def test_1frame_full_atm_shuffle(self) -> None: i_sf = [2, 1, 3, 0, 5, 4] isel_sf = [1, 0] gt, ff, vv, at, av = self.dp.eval_full( @@ -1055,7 +1055,7 @@ def test_1frame_full_atm_shuffle(self): vv.reshape([-1]), self.expected_gv.reshape([-1]), decimal=default_places ) - def test_2frame_full_atm(self): + def test_2frame_full_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) gt, ff, vv, at, av = self.dp.eval_full(coords2, box2, self.atype, atomic=True) @@ -1098,7 +1098,7 @@ def test_2frame_full_atm(self): ) class TestDeepPolarNewPBCNeighborList(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deeppolar_new.pbtxt")), "deeppolar_new.pb", @@ -1111,9 +1111,9 @@ def setUpClass(cls): ) @unittest.skip("multiple frames not supported") - def test_2frame_full_atm(self): + def test_2frame_full_atm(self) -> None: pass @unittest.skip("multiple frames not supported") - def test_2frame_old_atm(self): + def test_2frame_old_atm(self) -> None: pass diff --git a/source/tests/tf/test_deeppot_a.py b/source/tests/tf/test_deeppot_a.py index 25477cafea..4f92cdee27 100644 --- a/source/tests/tf/test_deeppot_a.py +++ b/source/tests/tf/test_deeppot_a.py @@ -36,7 +36,7 @@ class TestModelMajorCompatability(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: model_file = str(infer_path / os.path.join("deeppot.pbtxt")) with open(model_file) as fp: # data = fp.read().replace('\n', '') @@ -53,11 +53,11 @@ def setUp(self): fp.write("\n".join(data)) convert_pbtxt_to_pb(self.version_pbtxt, self.version_pb) - def tearDown(self): + def tearDown(self) -> None: os.remove(self.version_pbtxt) os.remove(self.version_pb) - def test(self): + def test(self) -> None: with self.assertRaises(RuntimeError) as context: DeepPot(str(self.version_pb)) self.assertTrue("incompatible" in str(context.exception)) @@ -66,7 +66,7 @@ def test(self): class TestModelMinorCompatability(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: model_file = str(infer_path / os.path.join("deeppot.pbtxt")) with open(model_file) as fp: # data = fp.read().replace('\n', '') @@ -83,11 +83,11 @@ def setUp(self): fp.write("\n".join(data)) convert_pbtxt_to_pb(self.version_pbtxt, self.version_pb) - def tearDown(self): + def tearDown(self) -> None: os.remove(self.version_pbtxt) os.remove(self.version_pb) - def test(self): + def test(self) -> None: with self.assertRaises(RuntimeError) as context: DeepPot(self.version_pb) self.assertTrue("incompatible" in str(context.exception)) @@ -99,7 +99,7 @@ def test(self): class TestModelConvert(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -125,7 +125,7 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_convert_012(self): + def test_convert_012(self) -> None: old_model = "deeppot.pb" new_model = "deeppot-new.pb" convert_pbtxt_to_pb(str(infer_path / "sea_012.pbtxt"), old_model) @@ -135,7 +135,7 @@ def test_convert_012(self): os.remove(old_model) os.remove(new_model) - def test_convert(self): + def test_convert(self) -> None: old_model = "deeppot.pb" new_model = "deeppot-new.pb" convert_pbtxt_to_pb(str(infer_path / "sea_012.pbtxt"), old_model) @@ -145,7 +145,7 @@ def test_convert(self): os.remove(old_model) os.remove(new_model) - def test_detect(self): + def test_detect(self) -> None: old_model = "deeppot.pb" new_model_txt = "deeppot_new.pbtxt" new_model_pb = "deeppot_new.pb" @@ -184,14 +184,14 @@ def test_detect(self): class TestTypeEmbed(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("se_e2_a_tebd.pbtxt")), "se_e2_a_tebd.pb", ) cls.dp = DeepPot("se_e2_a_tebd.pb") - def test_eval_typeebd(self): + def test_eval_typeebd(self) -> None: expected_typeebd = np.array( [ [ diff --git a/source/tests/tf/test_deeppot_spin.py b/source/tests/tf/test_deeppot_spin.py index d64cdf2dd6..7ca0e27202 100644 --- a/source/tests/tf/test_deeppot_spin.py +++ b/source/tests/tf/test_deeppot_spin.py @@ -26,13 +26,13 @@ class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("deepspin.pbtxt")), "deepspin.pb" ) cls.dp = DeepPot("deepspin.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 0.0, @@ -163,11 +163,11 @@ def setUp(self): ) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("deepspin.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 3) self.assertEqual(self.dp.get_ntypes_spin(), 1) self.assertAlmostEqual(self.dp.get_rcut(), 5.6, places=default_places) @@ -175,7 +175,7 @@ def test_attrs(self): self.assertEqual(self.dp.get_dim_fparam(), 0) self.assertEqual(self.dp.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee, ff, vv = self.dp.eval(self.coords, self.box, self.atype, atomic=False) # check shape of the returns @@ -199,7 +199,7 @@ def test_1frame(self): expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee, ff, vv, ae, av = self.dp.eval( self.coords, self.box, self.atype, atomic=True ) @@ -232,7 +232,7 @@ def test_1frame_atm(self): expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee, ff, vv, ae, av = self.dp.eval(coords2, box2, self.atype, atomic=True) diff --git a/source/tests/tf/test_descrpt_hybrid.py b/source/tests/tf/test_descrpt_hybrid.py index cadf2f83ae..87f772e99b 100644 --- a/source/tests/tf/test_descrpt_hybrid.py +++ b/source/tests/tf/test_descrpt_hybrid.py @@ -30,10 +30,10 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestHybrid(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data(nframes=2) - def test_descriptor_hybrid(self): + def test_descriptor_hybrid(self) -> None: jfile = "water_hybrid.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_descrpt_nonsmth.py b/source/tests/tf/test_descrpt_nonsmth.py index 63c5f15c85..c971a6fa59 100644 --- a/source/tests/tf/test_descrpt_nonsmth.py +++ b/source/tests/tf/test_descrpt_nonsmth.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, comp=0, pbc=True, sess=None): + def setUp(self, data, comp=0, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -158,26 +158,26 @@ class TestNonSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_se") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_se") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_se") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_se") class TestLFPbc(tf.test.TestCase): - def test_pbc(self): + def test_pbc(self) -> None: data = Data() inter0 = Inter() inter1 = Inter() @@ -229,7 +229,7 @@ def test_pbc(self): np.testing.assert_almost_equal(f0[0], f1[0]) np.testing.assert_almost_equal(v0[0], v1[0]) - def test_pbc_small_box(self): + def test_pbc_small_box(self) -> None: data0 = Data() data1 = Data(box_scale=2) inter0 = Inter() diff --git a/source/tests/tf/test_descrpt_se_a_mask.py b/source/tests/tf/test_descrpt_se_a_mask.py index b495072c7d..6de36adb38 100644 --- a/source/tests/tf/test_descrpt_se_a_mask.py +++ b/source/tests/tf/test_descrpt_se_a_mask.py @@ -30,14 +30,14 @@ class TestModel(tf.test.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("dp4mask.pbtxt")), str(infer_path / os.path.join("dp4mask.pb")), ) cls.dp = DeepPot(str(infer_path / os.path.join("dp4mask.pb"))) - def test_dp_mask_model(self): + def test_dp_mask_model(self) -> None: dcoord = np.array( [ 3.345, @@ -218,7 +218,7 @@ def test_dp_mask_model(self): assert normalized_diff_ff < 100 - def test_descriptor_se_a_mask(self): + def test_descriptor_se_a_mask(self) -> None: jfile = "zinc_se_a_mask.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_descrpt_se_a_type.py b/source/tests/tf/test_descrpt_se_a_type.py index ccb2ef96a6..b9d9a66e86 100644 --- a/source/tests/tf/test_descrpt_se_a_type.py +++ b/source/tests/tf/test_descrpt_se_a_type.py @@ -23,10 +23,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data(nframes=2) - def test_descriptor_two_sides(self): + def test_descriptor_two_sides(self) -> None: jfile = "water_se_a_type.json" jdata = j_loader(jfile) @@ -188,7 +188,7 @@ def test_descriptor_two_sides(self): places = 10 np.testing.assert_almost_equal(model_dout, ref_dout, places) - def test_descriptor_one_side(self): + def test_descriptor_one_side(self) -> None: jfile = "water_se_a_type.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_descrpt_se_atten.py b/source/tests/tf/test_descrpt_se_atten.py index 6393cc222b..db47657e54 100644 --- a/source/tests/tf/test_descrpt_se_atten.py +++ b/source/tests/tf/test_descrpt_se_atten.py @@ -31,11 +31,11 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data(nframes=2) self.filename = __file__ - def test_descriptor_two_sides(self): + def test_descriptor_two_sides(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -212,7 +212,7 @@ def test_descriptor_two_sides(self): places = 10 np.testing.assert_almost_equal(model_dout, ref_dout, places) - def test_descriptor_one_side(self): + def test_descriptor_one_side(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -390,7 +390,7 @@ def test_descriptor_one_side(self): places = 10 np.testing.assert_almost_equal(model_dout, ref_dout, places) - def test_stripped_type_embedding_descriptor_two_sides(self): + def test_stripped_type_embedding_descriptor_two_sides(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -561,7 +561,7 @@ def test_stripped_type_embedding_descriptor_two_sides(self): places = 10 np.testing.assert_almost_equal(model_dout, ref_dout, places) - def test_compressible_descriptor_two_sides(self): + def test_compressible_descriptor_two_sides(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_descrpt_se_r.py b/source/tests/tf/test_descrpt_se_r.py index d95c8fbb21..acd9983186 100644 --- a/source/tests/tf/test_descrpt_se_r.py +++ b/source/tests/tf/test_descrpt_se_r.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, pbc=True, sess=None): + def setUp(self, data, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -133,26 +133,26 @@ class TestSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_se_r") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_se_r") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_se_r") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_se_r") class TestSeRPbc(tf.test.TestCase): - def test_pbc(self): + def test_pbc(self) -> None: data = Data() inter0 = Inter() inter1 = Inter() @@ -204,7 +204,7 @@ def test_pbc(self): np.testing.assert_almost_equal(f0[0], f1[0]) np.testing.assert_almost_equal(v0[0], v1[0]) - def test_pbc_small_box(self): + def test_pbc_small_box(self) -> None: data0 = Data() data1 = Data(box_scale=2) inter0 = Inter() diff --git a/source/tests/tf/test_descrpt_sea_ef.py b/source/tests/tf/test_descrpt_sea_ef.py index e9e990a659..24f8bdfe9b 100644 --- a/source/tests/tf/test_descrpt_sea_ef.py +++ b/source/tests/tf/test_descrpt_sea_ef.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, pbc=True, sess=None): + def setUp(self, data, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -152,21 +152,21 @@ class TestSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_sea_ef") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_sea_ef") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_sea_ef") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_sea_ef") diff --git a/source/tests/tf/test_descrpt_sea_ef_para.py b/source/tests/tf/test_descrpt_sea_ef_para.py index 1af6ea648a..05e87d2ba5 100644 --- a/source/tests/tf/test_descrpt_sea_ef_para.py +++ b/source/tests/tf/test_descrpt_sea_ef_para.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, pbc=True, sess=None): + def setUp(self, data, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -152,21 +152,21 @@ class TestSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_sea_ef_para") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_sea_ef_para") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_sea_ef_para") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_sea_ef_para") diff --git a/source/tests/tf/test_descrpt_sea_ef_rot.py b/source/tests/tf/test_descrpt_sea_ef_rot.py index ac0064130e..356e9dd5bc 100644 --- a/source/tests/tf/test_descrpt_sea_ef_rot.py +++ b/source/tests/tf/test_descrpt_sea_ef_rot.py @@ -20,7 +20,7 @@ class TestEfRot(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.natoms = [5, 5, 2, 3] self.ntypes = 2 @@ -137,7 +137,7 @@ def rotate_mat(self, axis_, theta): ] return np.array(r).reshape(3, 3) - def test_rot_axis(self, suffix=""): + def test_rot_axis(self, suffix="") -> None: suffix = "_axis" t_p_e, t_p_f, t_p_f, t_p_ae, t_p_av = self.build_efv( self.coord, @@ -283,7 +283,7 @@ def test_rot_axis(self, suffix=""): self.assertNotAlmostEqual(p_ae0[ii], p_ae1[ii]) self.assertNotAlmostEqual(v_ae0[ii], v_ae1[ii]) - def test_rot_diff_axis(self, suffix=""): + def test_rot_diff_axis(self, suffix="") -> None: suffix = "_diff_axis" t_p_e, t_p_f, t_p_f, t_p_ae, t_p_av = self.build_efv( self.coord, @@ -398,7 +398,7 @@ def test_rot_diff_axis(self, suffix=""): self.assertNotAlmostEqual(p_ae0[ii], p_ae1[ii]) self.assertNotAlmostEqual(v_ae0[ii], v_ae1[ii]) - def test_rot_field_corot(self, suffix=""): + def test_rot_field_corot(self, suffix="") -> None: suffix = "_field_corot" t_p_e, t_p_f, t_p_f, t_p_ae, t_p_av = self.build_efv( self.coord, diff --git a/source/tests/tf/test_descrpt_sea_ef_vert.py b/source/tests/tf/test_descrpt_sea_ef_vert.py index 09bca9b754..bee9b2e19c 100644 --- a/source/tests/tf/test_descrpt_sea_ef_vert.py +++ b/source/tests/tf/test_descrpt_sea_ef_vert.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, pbc=True, sess=None): + def setUp(self, data, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -152,21 +152,21 @@ class TestSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_sea_ef_vert") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_sea_ef_vert") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_sea_ef_vert") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_sea_ef_vert") diff --git a/source/tests/tf/test_descrpt_smooth.py b/source/tests/tf/test_descrpt_smooth.py index 91a3f7dbf0..be7f44e6c5 100644 --- a/source/tests/tf/test_descrpt_smooth.py +++ b/source/tests/tf/test_descrpt_smooth.py @@ -21,7 +21,7 @@ class Inter: - def setUp(self, data, pbc=True, sess=None): + def setUp(self, data, pbc=True, sess=None) -> None: self.sess = sess self.data = data self.natoms = self.data.get_natoms() @@ -151,26 +151,26 @@ class TestSmooth(Inter, tf.test.TestCase): # tf.test.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() Inter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, suffix="_smth") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, suffix="_smth") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, suffix="_smth") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, suffix="_smth") class TestSeAPbc(tf.test.TestCase): - def test_pbc(self): + def test_pbc(self) -> None: data = Data() inter0 = Inter() inter1 = Inter() @@ -222,7 +222,7 @@ def test_pbc(self): np.testing.assert_almost_equal(f0[0], f1[0]) np.testing.assert_almost_equal(v0[0], v1[0]) - def test_pbc_small_box(self): + def test_pbc_small_box(self) -> None: data0 = Data() data1 = Data(box_scale=2) inter0 = Inter() diff --git a/source/tests/tf/test_dipole_se_a.py b/source/tests/tf/test_dipole_se_a.py index 77e8fc4f0a..bab61157ba 100644 --- a/source/tests/tf/test_dipole_se_a.py +++ b/source/tests/tf/test_dipole_se_a.py @@ -28,10 +28,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "polar_se_a.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_dipole_se_a_tebd.py b/source/tests/tf/test_dipole_se_a_tebd.py index ada0643dbf..1d2fee1fa3 100644 --- a/source/tests/tf/test_dipole_se_a_tebd.py +++ b/source/tests/tf/test_dipole_se_a_tebd.py @@ -38,10 +38,10 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "polar_se_a_tebd.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_dipolecharge.py b/source/tests/tf/test_dipolecharge.py index 408b1bbdf2..d4ad3254bc 100644 --- a/source/tests/tf/test_dipolecharge.py +++ b/source/tests/tf/test_dipolecharge.py @@ -26,7 +26,7 @@ class TestDipoleCharge(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("dipolecharge_d.pbtxt")), "dipolecharge_d.pb", @@ -35,7 +35,7 @@ def setUpClass(cls): "dipolecharge_d.pb", [-1.0, -3.0], [1.0, 1.0, 1.0, 1.0, 1.0], 4.0, 0.2 ) - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 4.6067455554, @@ -117,16 +117,16 @@ def setUp(self): self.coords = self.coords.reshape([-1, self.natoms, 3]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove("dipolecharge_d.pb") cls.dp = None - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 5) self.assertAlmostEqual(self.dp.get_rcut(), 4.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["A", "B", "C", "D", "E"]) - def test_1frame(self): + def test_1frame(self) -> None: ee, ff, vv = self.dp.eval(self.coords, self.box, self.atype, eval_fv=True) # check shape of the returns nframes = 1 @@ -147,7 +147,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ff, self.expected_f) np.testing.assert_almost_equal(vv, self.expected_v) - def test_2frame(self): + def test_2frame(self) -> None: nframes = 2 self.coords = np.tile(self.coords, [nframes, 1, 1]) self.box = np.tile(self.box, [nframes, 1]) diff --git a/source/tests/tf/test_dp_test.py b/source/tests/tf/test_dp_test.py index 2d9af234dd..dfbe4fe878 100644 --- a/source/tests/tf/test_dp_test.py +++ b/source/tests/tf/test_dp_test.py @@ -25,7 +25,7 @@ class TestDPTest: - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -62,28 +62,28 @@ def setUp(self): } ).to_deepmd_npy(self.test_data) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree(self.test_data, ignore_errors=True) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: os.remove(cls.model_name) class TestDPTestEner(unittest.TestCase, TestDPTest): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.case = get_cases()["se_e2_a"] cls.model_name = cls.case.get_model(".pb") - def setUp(self): + def setUp(self) -> None: self.result = self.case.results[0] TestDPTest.setUp(self) self.expected_e = self.result.atomic_energy self.expected_f = self.result.force self.expected_v = self.result.atomic_virial - def test_1frame(self): + def test_1frame(self) -> None: detail_file = "test_dp_test_ener_detail" dp_test( model=self.model_name, @@ -120,13 +120,13 @@ def test_1frame(self): class TestDPTestDipole(unittest.TestCase, TestDPTest): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.model_name = "deepdipole.pb" convert_pbtxt_to_pb( str(infer_path / os.path.join("deepdipole.pbtxt")), cls.model_name ) - def setUp(self): + def setUp(self) -> None: TestDPTest.setUp(self) self.expected_d = np.array( [ @@ -142,7 +142,7 @@ def setUp(self): np.save(Path(self.test_data) / "set.000" / "atomic_dipole.npy", self.expected_d) np.save(Path(self.test_data) / "set.000" / "dipole.npy", self.expected_global_d) - def test_1frame(self): + def test_1frame(self) -> None: detail_file = "test_dp_test_dipole_detail" dp_test( model=self.model_name, @@ -158,7 +158,7 @@ def test_1frame(self): dipole = np.loadtxt(detail_file + ".out", ndmin=2)[0, 6:12] np.testing.assert_almost_equal(dipole, self.expected_d, decimal=default_places) - def test_1frame_global(self): + def test_1frame_global(self) -> None: detail_file = "test_dp_test_global_dipole_detail" dp_test( model=self.model_name, @@ -179,13 +179,13 @@ def test_1frame_global(self): class TestDPTestPolar(unittest.TestCase, TestDPTest): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.model_name = "deeppolar.pb" convert_pbtxt_to_pb( str(infer_path / os.path.join("deeppolar.pbtxt")), cls.model_name ) - def setUp(self): + def setUp(self) -> None: TestDPTest.setUp(self) self.expected_d = np.array( [ @@ -219,7 +219,7 @@ def setUp(self): self.expected_global_d, ) - def test_1frame(self): + def test_1frame(self) -> None: detail_file = "test_dp_test_polar_detail" dp_test( model=self.model_name, @@ -235,7 +235,7 @@ def test_1frame(self): polar = np.loadtxt(detail_file + ".out", ndmin=2)[0, 18:36] np.testing.assert_almost_equal(polar, self.expected_d, decimal=default_places) - def test_1frame_global(self): + def test_1frame_global(self) -> None: detail_file = "test_dp_test_global_polar_detail" dp_test( model=self.model_name, diff --git a/source/tests/tf/test_dplr.py b/source/tests/tf/test_dplr.py index 6788a22547..195b1d0820 100644 --- a/source/tests/tf/test_dplr.py +++ b/source/tests/tf/test_dplr.py @@ -15,7 +15,7 @@ class TestDPLR(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: # a bit strange path, need to move to the correct directory pbtxt_file = ( Path(__file__).parent.parent.parent / "lmp" / "tests" / "lrmodel.pbtxt" @@ -49,7 +49,7 @@ def setUp(self): ).reshape(1, 8, 3) self.atype = np.array([0, 0, 1, 1, 1, 1, 2, 2]) - def test_eval(self): + def test_eval(self) -> None: dp = DeepPot("lrmodel.pb") e, f, v, ae, av = dp.eval( self.coord[:, :6], self.box, self.atype[:6], atomic=True diff --git a/source/tests/tf/test_embedding_net.py b/source/tests/tf/test_embedding_net.py index f766fff8b3..17f752cb3a 100644 --- a/source/tests/tf/test_embedding_net.py +++ b/source/tests/tf/test_embedding_net.py @@ -12,14 +12,14 @@ class Inter(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.inputs = tf.constant([0.0, 1.0, 2.0], dtype=tf.float64) self.ndata = 3 self.inputs = tf.reshape(self.inputs, [-1, 1]) self.places = 6 - def test_enlarger_net(self): + def test_enlarger_net(self) -> None: network_size = [3, 4] out = embedding_net( self.inputs, @@ -38,7 +38,7 @@ def test_enlarger_net(self): ] np.testing.assert_almost_equal(refout, myout, self.places) - def test_enlarger_net_1(self): + def test_enlarger_net_1(self) -> None: network_size = [4, 4] out = embedding_net( self.inputs, @@ -57,7 +57,7 @@ def test_enlarger_net_1(self): ] np.testing.assert_almost_equal(refout, myout, self.places) - def test_enlarger_net_1_idt(self): + def test_enlarger_net_1_idt(self) -> None: network_size = [4, 4] out = embedding_net( self.inputs, @@ -77,7 +77,7 @@ def test_enlarger_net_1_idt(self): ] np.testing.assert_almost_equal(refout, myout, self.places) - def test_enlarger_net_2(self): + def test_enlarger_net_2(self) -> None: network_size = [2, 4] out = embedding_net( self.inputs, @@ -96,7 +96,7 @@ def test_enlarger_net_2(self): ] np.testing.assert_almost_equal(refout, myout, self.places) - def test_enlarger_net_3(self): + def test_enlarger_net_3(self) -> None: network_size = [2, 4] out = embedding_net( self.inputs, diff --git a/source/tests/tf/test_env.py b/source/tests/tf/test_env.py index cd066b06a5..dc2e61913f 100644 --- a/source/tests/tf/test_env.py +++ b/source/tests/tf/test_env.py @@ -11,7 +11,7 @@ class TestTFThreadCount(unittest.TestCase): @mock.patch.dict("os.environ", values={}) - def test_empty(self): + def test_empty(self) -> None: intra, inter = env.get_tf_default_nthreads() self.assertEqual(intra, 0) self.assertEqual(inter, 0) @@ -23,26 +23,26 @@ def test_empty(self): "DP_INTER_OP_PARALLELISM_THREADS": "3", }, ) - def test_given(self): + def test_given(self) -> None: intra, inter = env.get_tf_default_nthreads() self.assertEqual(intra, 5) self.assertEqual(inter, 3) class TestTFSessionConfig(unittest.TestCase): - def test_default(self): + def test_default(self) -> None: shared = env.default_tf_session_config new = env.get_tf_session_config() self.assertNotEqual(id(shared), id(new)) @mock.patch("deepmd.tf.env.get_tf_default_nthreads") - def test_get(self, mock_method): + def test_get(self, mock_method) -> None: mock_method.return_value = (5, 3) config = env.get_tf_session_config() self.assertEqual(config.intra_op_parallelism_threads, 5) self.assertEqual(config.inter_op_parallelism_threads, 3) - def test_reset(self): + def test_reset(self) -> None: shared = env.default_tf_session_config env.reset_default_tf_session_config(True) self.assertEqual(shared.device_count["GPU"], 0) diff --git a/source/tests/tf/test_ewald.py b/source/tests/tf/test_ewald.py index a465f5bc7a..c108c44fe7 100644 --- a/source/tests/tf/test_ewald.py +++ b/source/tests/tf/test_ewald.py @@ -26,7 +26,7 @@ class TestEwaldRecp(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) boxl = 4.5 # NOTICE grid should not change before and after box pert... box_pert = 0.2 @@ -66,7 +66,7 @@ def setUp(self): self.box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_box") self.nloc = tf.placeholder(tf.int32, [1], name="t_nloc") - def test_py_interface(self): + def test_py_interface(self) -> None: hh = 1e-4 places = 4 sess = self.cached_session().__enter__() @@ -93,7 +93,7 @@ def test_py_interface(self): np.testing.assert_almost_equal(f, f1, places, err_msg="force component failed") np.testing.assert_almost_equal(v, v, places, err_msg="virial component failed") - def test_force(self): + def test_force(self) -> None: hh = 1e-4 places = 6 sess = self.cached_session().__enter__() @@ -146,7 +146,7 @@ def test_force(self): err_msg="force component [%d,%d] failed" % (idx, dd), ) - def test_virial(self): + def test_virial(self) -> None: hh = 1e-4 places = 6 sess = self.cached_session().__enter__() diff --git a/source/tests/tf/test_finetune_se_atten.py b/source/tests/tf/test_finetune_se_atten.py index 441eb88003..a7500917ec 100644 --- a/source/tests/tf/test_finetune_se_atten.py +++ b/source/tests/tf/test_finetune_se_atten.py @@ -38,7 +38,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -145,15 +145,15 @@ class TestFinetuneSeAtten(unittest.TestCase): def setUpClass(cls) -> None: if not parse_version(tf.__version__) < parse_version("1.15"): - def previous_se_atten(jdata): + def previous_se_atten(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "concat" jdata["model"]["descriptor"]["attn_layer"] = 2 - def stripped_model(jdata): + def stripped_model(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "strip" jdata["model"]["descriptor"]["attn_layer"] = 2 - def compressible_model(jdata): + def compressible_model(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "strip" jdata["model"]["descriptor"]["attn_layer"] = 0 @@ -199,7 +199,7 @@ def compressible_model(jdata): cls.VALID_DATAS = VALID_DATAS @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: for i in range(len(cls.INPUT_PRES)): _file_delete(cls.INPUT_PRES[i]) _file_delete(cls.INPUT_FINETUNES[i]) @@ -221,7 +221,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_finetune_standard(self): + def test_finetune_standard(self) -> None: for i in range(len(self.INPUT_PRES)): self.valid_data = self.VALID_DATAS[i] pretrained_bias = get_tensor_by_name( @@ -264,7 +264,7 @@ def test_finetune_standard(self): ) np.testing.assert_almost_equal(finetune_results, 0.0, default_places) - def test_finetune_mixed_type(self): + def test_finetune_mixed_type(self) -> None: for i in range(len(self.INPUT_PRES)): self.valid_data = self.VALID_DATAS[i] pretrained_bias = get_tensor_by_name( diff --git a/source/tests/tf/test_fitting_dos.py b/source/tests/tf/test_fitting_dos.py index e66f0078c6..81adb2761e 100644 --- a/source/tests/tf/test_fitting_dos.py +++ b/source/tests/tf/test_fitting_dos.py @@ -23,10 +23,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_fitting(self): + def test_fitting(self) -> None: jfile = "train_dos.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_fitting_ener_type.py b/source/tests/tf/test_fitting_ener_type.py index 94ee99c386..dca7e1c4cd 100644 --- a/source/tests/tf/test_fitting_ener_type.py +++ b/source/tests/tf/test_fitting_ener_type.py @@ -23,10 +23,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_fitting(self): + def test_fitting(self) -> None: jfile = "water_se_a_type.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_fitting_stat.py b/source/tests/tf/test_fitting_stat.py index 100868fd18..eb28ef2bfa 100644 --- a/source/tests/tf/test_fitting_stat.py +++ b/source/tests/tf/test_fitting_stat.py @@ -74,7 +74,7 @@ def _brute_aparam(data, ndim): class TestEnerFittingStat(unittest.TestCase): - def test(self): + def test(self) -> None: jdata = j_loader(input_json) jdata = jdata["model"] # descrpt = DescrptSeA(jdata['descriptor']) diff --git a/source/tests/tf/test_gen_stat_data.py b/source/tests/tf/test_gen_stat_data.py index ebede15fbb..d0fc3b6b3d 100644 --- a/source/tests/tf/test_gen_stat_data.py +++ b/source/tests/tf/test_gen_stat_data.py @@ -46,7 +46,7 @@ def gen_sys(nframes, atom_types): class TestGenStatData(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: data0 = gen_sys(20, [0, 1, 0, 2, 1]) data1 = gen_sys(30, [0, 1, 0, 0]) sys0 = dpdata.LabeledSystem() @@ -56,14 +56,14 @@ def setUp(self): sys0.to_deepmd_npy("system_0", set_size=10) sys1.to_deepmd_npy("system_1", set_size=10) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree("system_0") shutil.rmtree("system_1") - def _comp_data(self, d0, d1): + def _comp_data(self, d0, d1) -> None: np.testing.assert_almost_equal(d0, d1) - def test_merge_all_stat(self): + def test_merge_all_stat(self) -> None: dp_random.seed(0) data0 = DeepmdDataSystem(["system_0", "system_1"], 5, 10, 1.0) data0.add("energy", 1, must=True) @@ -103,7 +103,7 @@ def test_merge_all_stat(self): class TestEnerShift(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: data0 = gen_sys(30, [0, 1, 0, 2, 1]) data1 = gen_sys(30, [0, 1, 0, 0]) sys0 = dpdata.LabeledSystem() @@ -113,11 +113,11 @@ def setUp(self): sys0.to_deepmd_npy("system_0", set_size=10) sys1.to_deepmd_npy("system_1", set_size=10) - def tearDown(self): + def tearDown(self) -> None: shutil.rmtree("system_0") shutil.rmtree("system_1") - def test_ener_shift(self): + def test_ener_shift(self) -> None: dp_random.seed(0) data = DeepmdDataSystem(["system_0", "system_1"], 5, 10, 1.0) data.add("energy", 1, must=True) @@ -133,7 +133,7 @@ def test_ener_shift(self): ener_shift1 = fitting._compute_output_stats(all_stat, rcond=1) np.testing.assert_almost_equal(ener_shift0, ener_shift1) - def test_ener_shift_assigned(self): + def test_ener_shift_assigned(self) -> None: dp_random.seed(0) ae0 = dp_random.random() data = DeepmdDataSystem(["system_0"], 5, 10, 1.0) diff --git a/source/tests/tf/test_get_potential.py b/source/tests/tf/test_get_potential.py index fb39d41d2e..8723b6bfa8 100644 --- a/source/tests/tf/test_get_potential.py +++ b/source/tests/tf/test_get_potential.py @@ -26,7 +26,7 @@ class TestGetPotential(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.work_dir = infer_path convert_pbtxt_to_pb( @@ -56,11 +56,11 @@ def setUp(self): f.flush() convert_pbtxt_to_pb(f.name, str(self.work_dir / "deep_wfc.pb")) - def tearDown(self): + def tearDown(self) -> None: for f in self.work_dir.glob("*.pb"): f.unlink() - def test_factory(self): + def test_factory(self) -> None: msg = "Returned wrong type of potential. Expected: {}, got: {}" dp = DeepPotential(self.work_dir / "deep_dipole.pb") diff --git a/source/tests/tf/test_init_frz_model_se_a.py b/source/tests/tf/test_init_frz_model_se_a.py index 5d4ed1063c..5fa722fdb2 100644 --- a/source/tests/tf/test_init_frz_model_se_a.py +++ b/source/tests/tf/test_init_frz_model_se_a.py @@ -37,7 +37,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -131,7 +131,7 @@ def _init_models(): class TestInitFrzModelA(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT, cls.CKPT, @@ -147,7 +147,7 @@ def setUpClass(cls): cls.stop_batch = STOP_BATCH @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT) _file_delete(cls.FROZEN_MODEL) _file_delete("out.json") @@ -164,7 +164,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: valid_batch = self.valid_data.get_batch() natoms = valid_batch["natoms_vec"] tf.reset_default_graph() diff --git a/source/tests/tf/test_init_frz_model_se_a_tebd.py b/source/tests/tf/test_init_frz_model_se_a_tebd.py index afc1e46ed8..a515e52126 100644 --- a/source/tests/tf/test_init_frz_model_se_a_tebd.py +++ b/source/tests/tf/test_init_frz_model_se_a_tebd.py @@ -37,7 +37,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -132,7 +132,7 @@ def _init_models(): class TestInitFrzModelA(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT, cls.CKPT, @@ -149,7 +149,7 @@ def setUpClass(cls): cls.stop_batch = STOP_BATCH @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT) _file_delete(cls.FROZEN_MODEL) _file_delete("out.json") @@ -166,7 +166,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: valid_batch = self.valid_data.get_batch() natoms = valid_batch["natoms_vec"] tf.reset_default_graph() diff --git a/source/tests/tf/test_init_frz_model_se_a_type.py b/source/tests/tf/test_init_frz_model_se_a_type.py index 48ff4eb294..3719e4e9db 100644 --- a/source/tests/tf/test_init_frz_model_se_a_type.py +++ b/source/tests/tf/test_init_frz_model_se_a_type.py @@ -37,7 +37,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -135,7 +135,7 @@ def _init_models(): class TestInitFrzModelAType(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT, cls.CKPT, @@ -151,7 +151,7 @@ def setUpClass(cls): cls.stop_batch = STOP_BATCH @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT) _file_delete(cls.FROZEN_MODEL) _file_delete("out.json") @@ -168,7 +168,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: valid_batch = self.valid_data.get_batch() natoms = valid_batch["natoms_vec"] tf.reset_default_graph() diff --git a/source/tests/tf/test_init_frz_model_se_atten.py b/source/tests/tf/test_init_frz_model_se_atten.py index 25f629511c..fa2711ca4e 100644 --- a/source/tests/tf/test_init_frz_model_se_atten.py +++ b/source/tests/tf/test_init_frz_model_se_atten.py @@ -38,7 +38,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -135,15 +135,15 @@ def _init_models(model_setup, i): if not parse_version(tf.__version__) < parse_version("1.15"): - def previous_se_atten(jdata): + def previous_se_atten(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "concat" jdata["model"]["descriptor"]["attn_layer"] = 2 - def stripped_model(jdata): + def stripped_model(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "strip" jdata["model"]["descriptor"]["attn_layer"] = 2 - def compressible_model(jdata): + def compressible_model(jdata) -> None: jdata["model"]["descriptor"]["tebd_input_mode"] = "strip" jdata["model"]["descriptor"]["attn_layer"] = 0 @@ -154,7 +154,7 @@ def compressible_model(jdata): ) class TestInitFrzModelAtten(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: models = [previous_se_atten, stripped_model, compressible_model] INPUTS = [] CKPTS = [] @@ -193,7 +193,7 @@ def setUpClass(cls): cls.stop_batchs = STOP_BATCHS @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: for i in range(len(cls.dp_ckpts)): _file_delete(cls.INPUTS[i]) _file_delete(cls.FROZEN_MODELS[i]) @@ -211,7 +211,7 @@ def tearDownClass(cls): _file_delete(f"input_v2_compat{i}.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: for i in range(len(self.dp_ckpts)): self.dp_ckpt = self.CKPT_TRAINERS[i] self.dp_frz = self.FRZ_TRAINERS[i] diff --git a/source/tests/tf/test_init_frz_model_se_r.py b/source/tests/tf/test_init_frz_model_se_r.py index 100c09196e..8c6febbe49 100644 --- a/source/tests/tf/test_init_frz_model_se_r.py +++ b/source/tests/tf/test_init_frz_model_se_r.py @@ -37,7 +37,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -139,7 +139,7 @@ def _init_models(): class TestInitFrzModelR(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT, cls.CKPT, @@ -156,7 +156,7 @@ def setUpClass(cls): cls.stop_batch = STOP_BATCH @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT) _file_delete(cls.FROZEN_MODEL) _file_delete("out.json") @@ -173,7 +173,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: valid_batch = self.valid_data.get_batch() natoms = valid_batch["natoms_vec"] tf.reset_default_graph() diff --git a/source/tests/tf/test_init_frz_model_spin.py b/source/tests/tf/test_init_frz_model_spin.py index a99a052391..30a9ae6513 100644 --- a/source/tests/tf/test_init_frz_model_spin.py +++ b/source/tests/tf/test_init_frz_model_spin.py @@ -37,7 +37,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -141,7 +141,7 @@ def _init_models(): class TestInitFrzModelR(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT, cls.CKPT, @@ -158,7 +158,7 @@ def setUpClass(cls): cls.stop_batch = STOP_BATCH @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT) _file_delete(cls.FROZEN_MODEL) _file_delete("out.json") @@ -175,7 +175,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_single_frame(self): + def test_single_frame(self) -> None: valid_batch = self.valid_data.get_batch() natoms = valid_batch["natoms_vec"] tf.reset_default_graph() diff --git a/source/tests/tf/test_lammps.py b/source/tests/tf/test_lammps.py index b295d5212a..b8f42e93e5 100644 --- a/source/tests/tf/test_lammps.py +++ b/source/tests/tf/test_lammps.py @@ -20,13 +20,13 @@ class TestLAMMPS(unittest.TestCase): """Test LAMMPS in cibuildwheel environment.""" @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.work_dir = infer_path convert_pbtxt_to_pb( str(cls.work_dir / "deeppot.pbtxt"), str(cls.work_dir / "deep_pot.pb") ) - def test_lmp(self): + def test_lmp(self) -> None: in_file = (self.work_dir / "in.test").absolute() subprocess.check_call(["lmp", "-in", str(in_file)], cwd=str(self.work_dir)) diff --git a/source/tests/tf/test_linear_model.py b/source/tests/tf/test_linear_model.py index 95ece9c19f..e99c1042dd 100644 --- a/source/tests/tf/test_linear_model.py +++ b/source/tests/tf/test_linear_model.py @@ -27,7 +27,7 @@ class TestLinearModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() self.data_dir = "system" with open(os.path.join(self.data_dir, "type_map.raw"), "w") as f: @@ -41,7 +41,7 @@ def setUp(self): convert_pbtxt_to_pb(pbtxt, pb) self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] - def test_linear_ener_model(self): + def test_linear_ener_model(self) -> None: numb_test = 1 data = DataSystem([self.data_dir], "set", 1, 1, 6, run_opt=None) test_data = data.get_test() @@ -117,7 +117,7 @@ def test_linear_ener_model(self): np.testing.assert_allclose(f, np.mean(fs, axis=0), rtol=1e-5, atol=1e-5) np.testing.assert_allclose(v, np.mean(vs, axis=0), rtol=1e-5, atol=1e-5) - def tearDown(self): + def tearDown(self) -> None: for pb in self.graph_dirs: os.remove(pb) del_data() diff --git a/source/tests/tf/test_loss_gf.py b/source/tests/tf/test_loss_gf.py index 116b98b649..fd293b1ea0 100644 --- a/source/tests/tf/test_loss_gf.py +++ b/source/tests/tf/test_loss_gf.py @@ -11,7 +11,7 @@ class TestLossGf(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.loss = EnerStdLoss( 1e-3, start_pref_e=0, @@ -29,7 +29,7 @@ def setUp(self): numb_generalized_coord=2, ) - def test_label_requirements(self): + def test_label_requirements(self) -> None: """Test label_requirements are expected.""" self.assertCountEqual( self.loss.label_requirement, @@ -85,7 +85,7 @@ def test_label_requirements(self): ], ) - def test_build_loss(self): + def test_build_loss(self) -> None: natoms = tf.constant([6, 6]) model_dict = { "energy": tf.zeros((1, 1), dtype=tf.float64), diff --git a/source/tests/tf/test_mixed_prec_training.py b/source/tests/tf/test_mixed_prec_training.py index b43ad7ce2f..5436015d76 100644 --- a/source/tests/tf/test_mixed_prec_training.py +++ b/source/tests/tf/test_mixed_prec_training.py @@ -23,7 +23,7 @@ ) -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -31,7 +31,7 @@ def _file_delete(file): class TestMixedPrecTraining(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: data_file = str(tests_path / os.path.join("model_compression", "data")) self.INPUT = str(tests_path / "input.json") jdata = j_loader( @@ -45,14 +45,14 @@ def setUp(self): with open(self.INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - def test_training(self): + def test_training(self) -> None: _TF_VERSION = Version(TF_VERSION) # check the TF_VERSION, when TF < 1.12, mixed precision is not allowed if _TF_VERSION >= Version("1.14.0"): ret = run_dp("dp train " + self.INPUT) np.testing.assert_equal(ret, 0, "DP train failed!") - def tearDown(self): + def tearDown(self) -> None: _file_delete(self.INPUT) _file_delete("out.json") _file_delete("checkpoint") diff --git a/source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py b/source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py index 5557305a7a..0dac14acfa 100644 --- a/source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py +++ b/source/tests/tf/test_model_compression_dpa1_compat_suffix_only.py @@ -59,7 +59,7 @@ def build_eval_tf(sess, obj, natoms, coords, atype, box, suffix): class TestDescriptorSeA(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.device = "cpu" self.seed = 21 self.sel = [9, 10] @@ -117,7 +117,7 @@ def setUp(self): def test_tf_pt_consistent( self, - ): + ) -> None: with tf.Session(config=default_tf_session_config) as sess: graph = tf.get_default_graph() ret = build_eval_tf( diff --git a/source/tests/tf/test_model_compression_se_a.py b/source/tests/tf/test_model_compression_se_a.py index 0740ada288..adcd781d10 100644 --- a/source/tests/tf/test_model_compression_se_a.py +++ b/source/tests/tf/test_model_compression_se_a.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -73,7 +73,7 @@ def _init_models_exclude_types(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global \ INPUT, \ FROZEN_MODEL, \ @@ -87,7 +87,7 @@ def setUpModule(): class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -115,7 +115,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -130,7 +130,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -151,7 +151,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -178,7 +178,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -211,7 +211,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -239,7 +239,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -260,7 +260,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -287,7 +287,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -319,7 +319,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -347,7 +347,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -368,7 +368,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -395,7 +395,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -427,7 +427,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL_ET) cls.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) cls.coords = np.array( @@ -456,7 +456,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(INPUT_ET) _file_delete(FROZEN_MODEL_ET) _file_delete(COMPRESSED_MODEL_ET) @@ -477,7 +477,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -492,7 +492,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -513,7 +513,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -540,7 +540,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_model_compression_se_a_ebd.py b/source/tests/tf/test_model_compression_se_a_ebd.py index 38fdaf7570..ddeb617fc8 100644 --- a/source/tests/tf/test_model_compression_se_a_ebd.py +++ b/source/tests/tf/test_model_compression_se_a_ebd.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -85,7 +85,7 @@ def _init_models_exclude_types(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global \ INPUT, \ FROZEN_MODEL, \ @@ -99,7 +99,7 @@ def setUpModule(): class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -127,7 +127,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -142,7 +142,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -163,7 +163,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -190,7 +190,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -223,7 +223,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -251,7 +251,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -272,7 +272,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -299,7 +299,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -331,7 +331,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -359,7 +359,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -380,7 +380,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -407,7 +407,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -439,7 +439,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL_ET) cls.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) cls.coords = np.array( @@ -468,7 +468,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(INPUT_ET) _file_delete(FROZEN_MODEL_ET) _file_delete(COMPRESSED_MODEL_ET) @@ -489,7 +489,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -504,7 +504,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -525,7 +525,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -552,7 +552,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py b/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py index c873675d7c..933b950a60 100644 --- a/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py +++ b/source/tests/tf/test_model_compression_se_a_ebd_type_one_side.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -85,14 +85,14 @@ def _init_models_exclude_types(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global INPUT, FROZEN_MODEL, COMPRESSED_MODEL INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -120,7 +120,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -135,7 +135,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -156,7 +156,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -183,7 +183,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -216,7 +216,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -244,7 +244,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -265,7 +265,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -292,7 +292,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -324,7 +324,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -352,7 +352,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -373,7 +373,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -400,7 +400,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -432,7 +432,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ( cls.INPUT_ET, cls.FROZEN_MODEL_ET, @@ -466,7 +466,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.INPUT_ET) _file_delete(cls.FROZEN_MODEL_ET) _file_delete(cls.COMPRESSED_MODEL_ET) @@ -487,7 +487,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -502,7 +502,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -523,7 +523,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -550,7 +550,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_model_compression_se_a_type_one_side_exclude_types.py b/source/tests/tf/test_model_compression_se_a_type_one_side_exclude_types.py index 31c17ebd68..ef90e7ba1a 100644 --- a/source/tests/tf/test_model_compression_se_a_type_one_side_exclude_types.py +++ b/source/tests/tf/test_model_compression_se_a_type_one_side_exclude_types.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -56,7 +56,7 @@ def _init_models(): class TestDeepPotAPBCTypeOneSideExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() cls.dp_original = DeepPot(FROZEN_MODEL) @@ -86,7 +86,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -101,7 +101,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -122,7 +122,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -149,7 +149,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_model_compression_se_atten.py b/source/tests/tf/test_model_compression_se_atten.py index 0a231fde2b..a49b526c86 100644 --- a/source/tests/tf/test_model_compression_se_atten.py +++ b/source/tests/tf/test_model_compression_se_atten.py @@ -21,7 +21,7 @@ ) -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -154,7 +154,7 @@ def _init_models_exclude_types(): return inputs, frozen_models, compressed_models -def setUpModule(): +def setUpModule() -> None: global \ INPUTS, \ FROZEN_MODELS, \ @@ -172,7 +172,7 @@ def setUpModule(): ) class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(len(tests))] cls.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(len(tests))] cls.coords = np.array( @@ -200,7 +200,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -218,7 +218,7 @@ def test_attrs(self): self.assertEqual(dp_compressed.get_dim_fparam(), 0) self.assertEqual(dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -244,7 +244,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -276,7 +276,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -318,7 +318,7 @@ def test_2frame_atm(self): ) class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(len(tests))] cls.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(len(tests))] cls.coords = np.array( @@ -346,7 +346,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -372,7 +372,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -404,7 +404,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -445,7 +445,7 @@ def test_2frame_atm(self): ) class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_originals = [DeepPot(FROZEN_MODELS[i]) for i in range(len(tests))] cls.dp_compresseds = [DeepPot(COMPRESSED_MODELS[i]) for i in range(len(tests))] cls.coords = np.array( @@ -473,7 +473,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -505,7 +505,7 @@ def test_1frame(self): vv0, vv1, default_places, err_msg=str(tests[i]) ) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -537,7 +537,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: for i in range(len(tests)): default_places = tests[i]["precision_digit"] from ase import ( @@ -575,7 +575,7 @@ def test_ase(self): ) class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_originals = [DeepPot(FROZEN_MODELS_ET[i]) for i in range(len(tests))] cls.dp_compresseds = [ DeepPot(COMPRESSED_MODELS_ET[i]) for i in range(len(tests)) @@ -606,7 +606,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: for i in range(len(tests)): _file_delete(INPUTS_ET[i]) _file_delete(FROZEN_MODELS_ET[i]) @@ -628,7 +628,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -646,7 +646,7 @@ def test_attrs(self): self.assertEqual(dp_compressed.get_dim_fparam(), 0) self.assertEqual(dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -672,7 +672,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] @@ -704,7 +704,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: for i in range(len(tests)): dp_original = self.dp_originals[i] dp_compressed = self.dp_compresseds[i] diff --git a/source/tests/tf/test_model_compression_se_r.py b/source/tests/tf/test_model_compression_se_r.py index d4f8549abb..c42538f05d 100644 --- a/source/tests/tf/test_model_compression_se_r.py +++ b/source/tests/tf/test_model_compression_se_r.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -60,14 +60,14 @@ def _init_models(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global INPUT, FROZEN_MODEL, COMPRESSED_MODEL INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -95,7 +95,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -110,7 +110,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -131,7 +131,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -158,7 +158,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -191,7 +191,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -219,7 +219,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -240,7 +240,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -267,7 +267,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -299,7 +299,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -327,7 +327,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -348,7 +348,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -375,7 +375,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -407,7 +407,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -436,7 +436,7 @@ def setUpClass(cls): cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(INPUT) _file_delete(FROZEN_MODEL) _file_delete(COMPRESSED_MODEL) @@ -457,7 +457,7 @@ def tearDownClass(cls): _file_delete("input_v2_compat.json") _file_delete("lcurve.out") - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -472,7 +472,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -493,7 +493,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -520,7 +520,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_model_compression_se_t.py b/source/tests/tf/test_model_compression_se_t.py index 6228455890..36373d13d5 100644 --- a/source/tests/tf/test_model_compression_se_t.py +++ b/source/tests/tf/test_model_compression_se_t.py @@ -25,7 +25,7 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.isdir(file): os.rmdir(file) elif os.path.isfile(file): @@ -60,12 +60,12 @@ def _init_models(): return INPUT, frozen_model, compressed_model -def setUpModule(): +def setUpModule() -> None: global INPUT, FROZEN_MODEL, COMPRESSED_MODEL INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() -def tearDownModule(): +def tearDownModule() -> None: _file_delete(INPUT) _file_delete(FROZEN_MODEL) _file_delete(COMPRESSED_MODEL) @@ -89,7 +89,7 @@ def tearDownModule(): class TestDeepPotAPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -117,7 +117,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -132,7 +132,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -153,7 +153,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -180,7 +180,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -213,7 +213,7 @@ def test_2frame_atm(self): class TestDeepPotANoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -241,7 +241,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = None - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -262,7 +262,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -289,7 +289,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( coords2, self.box, self.atype, atomic=True @@ -321,7 +321,7 @@ def test_2frame_atm(self): class TestDeepPotALargeBoxNoPBC(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -349,7 +349,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -370,7 +370,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -397,7 +397,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_ase(self): + def test_ase(self) -> None: from ase import ( Atoms, ) @@ -429,7 +429,7 @@ def test_ase(self): class TestDeepPotAPBCExcludeTypes(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( @@ -457,7 +457,7 @@ def setUpClass(cls): cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) @@ -472,7 +472,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=False ) @@ -493,7 +493,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, self.atype, atomic=True ) @@ -520,7 +520,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( @@ -553,7 +553,7 @@ def test_2frame_atm(self): class TestDeepPotAPBC2(TestDeepPotAPBC): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.dp_original = DeepPot(FROZEN_MODEL) cls.dp_compressed = DeepPot(COMPRESSED_MODEL) cls.coords = np.array( diff --git a/source/tests/tf/test_model_devi.py b/source/tests/tf/test_model_devi.py index 58a6266ca9..be9486cf02 100644 --- a/source/tests/tf/test_model_devi.py +++ b/source/tests/tf/test_model_devi.py @@ -24,7 +24,7 @@ class TestMakeModelDevi(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() self.data_dir = "system" with open(os.path.join(self.data_dir, "type_map.raw"), "w") as f: @@ -58,7 +58,7 @@ def setUp(self): ] ) - def test_calc_model_devi(self): + def test_calc_model_devi(self) -> None: model_devi = calc_model_devi( self.coord, None, @@ -73,7 +73,7 @@ def test_calc_model_devi(self): np.testing.assert_almost_equal(model_devi[0][1:8], model_devi[1][1:8], 6) self.assertTrue(os.path.isfile(self.output)) - def test_make_model_devi(self): + def test_make_model_devi(self) -> None: make_model_devi( models=self.graph_dirs, system=self.data_dir, @@ -84,7 +84,7 @@ def test_make_model_devi(self): x = np.loadtxt(self.output) np.testing.assert_allclose(x, self.expect, 6) - def test_make_model_devi_real_erorr(self): + def test_make_model_devi_real_erorr(self) -> None: make_model_devi( models=self.graph_dirs, system=self.data_dir, @@ -111,7 +111,7 @@ def test_make_model_devi_real_erorr(self): 6, ) - def test_make_model_devi_atomic_relative(self): + def test_make_model_devi_atomic_relative(self) -> None: _, expected_f, expected_v = self.graphs[0].eval( self.coord[0], self.box[0], self.atype ) @@ -152,7 +152,7 @@ def test_make_model_devi_atomic_relative(self): np.testing.assert_allclose(md[2], np.min(expected_md_v), 6) np.testing.assert_allclose(md[3], np.sqrt(np.mean(np.square(expected_md_v))), 6) - def test_make_model_devi_atomic_relative_v(self): + def test_make_model_devi_atomic_relative_v(self) -> None: _, expected_f, expected_v = self.graphs[0].eval( self.coord[0], self.box[0], self.atype ) @@ -200,7 +200,7 @@ def test_make_model_devi_atomic_relative_v(self): np.testing.assert_allclose(md[2], np.min(expected_md_v), 6) np.testing.assert_allclose(md[3], np.sqrt(np.mean(np.square(expected_md_v))), 6) - def tearDown(self): + def tearDown(self) -> None: for pb in self.graph_dirs: os.remove(pb) os.remove(self.output) @@ -211,7 +211,7 @@ class TestMakeModelDeviFparamAparam(unittest.TestCase): """Ensure dp model_devi accepts fparam and aparam.""" @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.pbtxts = [ os.path.join(infer_path, "fparam_aparam.pbtxt"), ] @@ -221,12 +221,12 @@ def setUpClass(cls): cls.graphs = [DeepPotential(pb) for pb in cls.graph_dirs] @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: for pb in cls.graph_dirs: os.remove(pb) cls.graphs = None - def setUp(self): + def setUp(self) -> None: gen_data(dim_fparam=1) self.data_dir = "system" coord = np.load(os.path.join(self.data_dir, "set.000/coord.npy")) @@ -246,7 +246,7 @@ def setUp(self): (nframes, self.atype.size, 1) ) - def test_calc_model_devi(self): + def test_calc_model_devi(self) -> None: model_devi = calc_model_devi( self.coord, None, @@ -263,7 +263,7 @@ def test_calc_model_devi(self): np.testing.assert_almost_equal(model_devi[0][1:8], model_devi[1][1:8], 6) self.assertTrue(os.path.isfile(self.output)) - def test_make_model_devi(self): + def test_make_model_devi(self) -> None: make_model_devi( models=self.graph_dirs, system=self.data_dir, @@ -274,6 +274,6 @@ def test_make_model_devi(self): x = np.loadtxt(self.output) np.testing.assert_allclose(x, self.expect, 6) - def tearDown(self): + def tearDown(self) -> None: os.remove(self.output) del_data() diff --git a/source/tests/tf/test_model_devi_mix.py b/source/tests/tf/test_model_devi_mix.py index d9062e939a..4ef9f39cab 100644 --- a/source/tests/tf/test_model_devi_mix.py +++ b/source/tests/tf/test_model_devi_mix.py @@ -32,7 +32,7 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestMakeModelDeviMix(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() self.data_dir = "system" with open(os.path.join(self.data_dir, "type_map.raw"), "w") as f: @@ -87,7 +87,7 @@ def setUp(self): ] ) - def test_calc_model_devi(self): + def test_calc_model_devi(self) -> None: model_devi = calc_model_devi( self.coord, None, @@ -102,7 +102,7 @@ def test_calc_model_devi(self): np.testing.assert_almost_equal(model_devi[0][1:8], model_devi[1][1:8], 6) self.assertTrue(os.path.isfile(self.output)) - def test_calc_model_devi_mixed(self): + def test_calc_model_devi_mixed(self) -> None: model_devi = calc_model_devi( self.coord, None, @@ -118,7 +118,7 @@ def test_calc_model_devi_mixed(self): np.testing.assert_almost_equal(model_devi[1][1:8], self.expect[1][1:8], 6) self.assertTrue(os.path.isfile(self.output)) - def test_make_model_devi_mixed(self): + def test_make_model_devi_mixed(self) -> None: make_model_devi( models=self.graph_dirs, system=self.data_dir, @@ -129,7 +129,7 @@ def test_make_model_devi_mixed(self): x = np.loadtxt(self.output) np.testing.assert_allclose(x, self.expect, 6) - def tearDown(self): + def tearDown(self) -> None: for pb in self.graph_dirs: os.remove(pb) os.remove(self.output) diff --git a/source/tests/tf/test_model_dos.py b/source/tests/tf/test_model_dos.py index fc6e4fdb2d..a37eae9d75 100644 --- a/source/tests/tf/test_model_dos.py +++ b/source/tests/tf/test_model_dos.py @@ -27,13 +27,13 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def tearDown(self): + def tearDown(self) -> None: del_data() - def test_model(self): + def test_model(self) -> None: jfile = "train_dos.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_loc_frame.py b/source/tests/tf/test_model_loc_frame.py index a043d16677..bf97916842 100644 --- a/source/tests/tf/test_model_loc_frame.py +++ b/source/tests/tf/test_model_loc_frame.py @@ -26,10 +26,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water.json" jdata = j_loader(jfile) systems = jdata["systems"] diff --git a/source/tests/tf/test_model_pairtab.py b/source/tests/tf/test_model_pairtab.py index 0a09e70430..203cc3a311 100644 --- a/source/tests/tf/test_model_pairtab.py +++ b/source/tests/tf/test_model_pairtab.py @@ -21,10 +21,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water.json" jdata = j_loader(jfile) systems = jdata["systems"] diff --git a/source/tests/tf/test_model_se_a.py b/source/tests/tf/test_model_se_a.py index 4d691834ae..da20291087 100644 --- a/source/tests/tf/test_model_se_a.py +++ b/source/tests/tf/test_model_se_a.py @@ -31,13 +31,13 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def tearDown(self): + def tearDown(self) -> None: del_data() - def test_model_atom_ener(self): + def test_model_atom_ener(self) -> None: jfile = "water_se_a.json" jdata = j_loader(jfile) set_atom_ener = [0.02, 0.01] @@ -132,7 +132,7 @@ def test_model_atom_ener(self): [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) self.assertAlmostEqual(e[0], set_atom_ener[1], places=10) - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a.json" jdata = j_loader(jfile) @@ -263,7 +263,7 @@ def test_model(self): # test input requirement for the model self.assertCountEqual(model.input_requirement, []) - def test_model_atom_ener_type_embedding(self): + def test_model_atom_ener_type_embedding(self) -> None: """Test atom ener with type embedding.""" jfile = "water_se_a.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_a_aparam.py b/source/tests/tf/test_model_se_a_aparam.py index fcc7dd52d1..46d7538d4b 100644 --- a/source/tests/tf/test_model_se_a_aparam.py +++ b/source/tests/tf/test_model_se_a_aparam.py @@ -29,10 +29,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_aparam.json" jdata = j_loader(jfile) systems = jdata["systems"] diff --git a/source/tests/tf/test_model_se_a_ebd.py b/source/tests/tf/test_model_se_a_ebd.py index 9f48076de8..5c3343ebaf 100644 --- a/source/tests/tf/test_model_se_a_ebd.py +++ b/source/tests/tf/test_model_se_a_ebd.py @@ -26,10 +26,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_ebd.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_a_ebd_v2.py b/source/tests/tf/test_model_se_a_ebd_v2.py index 69d44b7c56..cb3cdd946b 100644 --- a/source/tests/tf/test_model_se_a_ebd_v2.py +++ b/source/tests/tf/test_model_se_a_ebd_v2.py @@ -29,10 +29,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_ebd.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_a_fparam.py b/source/tests/tf/test_model_se_a_fparam.py index 7529873e83..4397ff0b26 100644 --- a/source/tests/tf/test_model_se_a_fparam.py +++ b/source/tests/tf/test_model_se_a_fparam.py @@ -29,10 +29,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_fparam.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_a_srtab.py b/source/tests/tf/test_model_se_a_srtab.py index dddd007006..dc21f259e1 100644 --- a/source/tests/tf/test_model_se_a_srtab.py +++ b/source/tests/tf/test_model_se_a_srtab.py @@ -27,7 +27,7 @@ GLOBAL_NP_FLOAT_PRECISION = np.float64 -def _make_tab(ntype): +def _make_tab(ntype) -> None: xx = np.arange(0, 9, 0.001) yy = 1000 / (xx + 0.5) ** 6 prt = xx @@ -39,14 +39,14 @@ def _make_tab(ntype): class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() _make_tab(2) - def tearDown(self): + def tearDown(self) -> None: os.remove("tab.xvg") - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_srtab.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_a_type.py b/source/tests/tf/test_model_se_a_type.py index e9fb8db466..a6ad67549f 100644 --- a/source/tests/tf/test_model_se_a_type.py +++ b/source/tests/tf/test_model_se_a_type.py @@ -29,10 +29,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_a_type.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_atten.py b/source/tests/tf/test_model_se_atten.py index 4f661c2cc3..b76bf0f8b6 100644 --- a/source/tests/tf/test_model_se_atten.py +++ b/source/tests/tf/test_model_se_atten.py @@ -39,11 +39,11 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() self.filename = __file__ - def test_model(self): + def test_model(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -193,7 +193,7 @@ def test_model(self): np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) - def test_exclude_types(self): + def test_exclude_types(self) -> None: """In this test, we make type 0 has no interaction with type 0 and type 1, so the descriptor should be zero for type 0 atoms. """ @@ -268,7 +268,7 @@ def test_exclude_types(self): with self.assertRaises(AssertionError): np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) - def test_compressible_model(self): + def test_compressible_model(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -419,7 +419,7 @@ def test_compressible_model(self): np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) - def test_compressible_exclude_types(self): + def test_compressible_exclude_types(self) -> None: """In this test, we make type 0 has no interaction with type 0 and type 1, so the descriptor should be zero for type 0 atoms. """ @@ -497,7 +497,7 @@ def test_compressible_exclude_types(self): with self.assertRaises(AssertionError): np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) - def test_stripped_type_embedding_model(self): + def test_stripped_type_embedding_model(self) -> None: jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -652,7 +652,7 @@ def test_stripped_type_embedding_model(self): np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) - def test_stripped_type_embedding_exclude_types(self): + def test_stripped_type_embedding_exclude_types(self) -> None: """In this test, we make type 0 has no interaction with type 0 and type 1, so the descriptor should be zero for type 0 atoms. """ @@ -735,7 +735,7 @@ def test_stripped_type_embedding_exclude_types(self): with self.assertRaises(AssertionError): np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) - def test_smoothness_of_stripped_type_embedding_smooth_model(self): + def test_smoothness_of_stripped_type_embedding_smooth_model(self) -> None: """test: auto-diff, continuity of e,f,v.""" jfile = "water_se_atten.json" jdata = j_loader(jfile) @@ -882,7 +882,9 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self): np.testing.assert_allclose(df[0], df[1], rtol=0, atol=deltad) np.testing.assert_allclose(dv[0], dv[1], rtol=0, atol=deltad) - def test_smoothness_of_stripped_type_embedding_smooth_model_excluded_types(self): + def test_smoothness_of_stripped_type_embedding_smooth_model_excluded_types( + self, + ) -> None: """test: auto-diff, continuity of e,f,v.""" jfile = "water_se_atten.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_r.py b/source/tests/tf/test_model_se_r.py index 628a59d320..fd41979a4f 100644 --- a/source/tests/tf/test_model_se_r.py +++ b/source/tests/tf/test_model_se_r.py @@ -26,10 +26,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_r.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_se_t.py b/source/tests/tf/test_model_se_t.py index 6167e2e224..75b2981f6a 100644 --- a/source/tests/tf/test_model_se_t.py +++ b/source/tests/tf/test_model_se_t.py @@ -26,10 +26,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "water_se_t.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_model_spin.py b/source/tests/tf/test_model_spin.py index 0e9c0678c2..a450dff178 100644 --- a/source/tests/tf/test_model_spin.py +++ b/source/tests/tf/test_model_spin.py @@ -33,13 +33,13 @@ class TestModelSpin(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def tearDown(self): + def tearDown(self) -> None: del_data() - def test_model_spin(self): + def test_model_spin(self) -> None: jfile = "test_model_spin.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_nlist.py b/source/tests/tf/test_nlist.py index 9e66b185ee..808dde1bfa 100644 --- a/source/tests/tf/test_nlist.py +++ b/source/tests/tf/test_nlist.py @@ -20,7 +20,7 @@ class TestNeighList(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: self.nf = 3 self.nloc = 2 self.ns = 5 * 5 * 3 @@ -50,7 +50,7 @@ def setUp(self): ] ) - def test_extend_coord(self): + def test_extend_coord(self) -> None: t_coord = tf.placeholder( GLOBAL_TF_FLOAT_PRECISION, [None, None], name="i_coord" ) diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py index 32e90d9d9b..eaf8bfafd5 100644 --- a/source/tests/tf/test_nvnmd_entrypoints.py +++ b/source/tests/tf/test_nvnmd_entrypoints.py @@ -48,7 +48,7 @@ class TestNvnmdEntrypointsV0(tf.test.TestCase): @pytest.mark.run(order=0) - def test_mapt_cnn_v0(self): + def test_mapt_cnn_v0(self) -> None: config_file = str(tests_path / "nvnmd" / "ref" / "config_v0_cnn.npy") weight_file = str(tests_path / "nvnmd" / "ref" / "weight_v0_cnn.npy") map_file = str(tests_path / "nvnmd" / "out" / "map_v0_cnn.npy") @@ -393,7 +393,7 @@ def test_mapt_cnn_v0(self): np.testing.assert_almost_equal(pred, ref_dout, 8) @pytest.mark.run(order=1) - def test_model_qnn_v0(self): + def test_model_qnn_v0(self) -> None: # without calling test_mapt_cnn_v0, this test will fail when running individually self.test_mapt_cnn_v0() @@ -512,14 +512,14 @@ def test_model_qnn_v0(self): ref_dout = -62.60181403 np.testing.assert_almost_equal(pred, ref_dout, 8) - def tearDown(self): + def tearDown(self) -> None: # close nvnmd_cfg.enable = False class TestNvnmdEntrypointsV1(tf.test.TestCase): @pytest.mark.run(order=0) - def test_mapt_cnn_v1(self): + def test_mapt_cnn_v1(self) -> None: config_file = str(tests_path / "nvnmd" / "ref" / "config_v1_cnn.npy") weight_file = str(tests_path / "nvnmd" / "ref" / "weight_v1_cnn.npy") map_file = str(tests_path / "nvnmd" / "out" / "map_v1_cnn.npy") @@ -705,7 +705,7 @@ def test_mapt_cnn_v1(self): np.testing.assert_almost_equal(pred, ref_dout, 8) @pytest.mark.run(order=1) - def test_model_qnn_v1(self): + def test_model_qnn_v1(self) -> None: # without calling test_mapt_cnn_v1, this test will fail when running individually self.test_mapt_cnn_v1() @@ -856,7 +856,7 @@ def test_model_qnn_v1(self): np.testing.assert_almost_equal(pred, ref_dout, 8) @pytest.mark.run(order=2) - def test_wrap_qnn_v1(self): + def test_wrap_qnn_v1(self) -> None: # without calling test_mapt_cnn_v1, this test will fail when running individually self.test_mapt_cnn_v1() @@ -875,7 +875,7 @@ def test_wrap_qnn_v1(self): red_dout = [249, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 254, 95, 24, 176] np.testing.assert_equal(pred, red_dout) - def tearDown(self): + def tearDown(self) -> None: # close nvnmd_cfg.enable = False diff --git a/source/tests/tf/test_nvnmd_op.py b/source/tests/tf/test_nvnmd_op.py index beff8375b8..acf345b61d 100644 --- a/source/tests/tf/test_nvnmd_op.py +++ b/source/tests/tf/test_nvnmd_op.py @@ -11,7 +11,7 @@ class TestOpAddFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -19,7 +19,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 1], "t_x") t_y = op_module.add_flt_nvnmd(t_x, t_x) @@ -104,7 +104,7 @@ def test_op(self): class TestOpCopyFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -112,7 +112,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 1], "t_x") t_y, t_y2 = op_module.copy_flt_nvnmd(t_x) @@ -134,7 +134,7 @@ def test_op(self): class TestOpDotmulFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -142,7 +142,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_y = op_module.dotmul_flt_nvnmd(t_x, t_x) @@ -160,7 +160,7 @@ def test_op(self): class TestOpFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -168,7 +168,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_y = op_module.flt_nvnmd(t_x) @@ -186,7 +186,7 @@ def test_op(self): class TestOpMatmulFitnetNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -194,7 +194,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_w = tf.placeholder(tf.float64, [4, 4], "t_w") @@ -232,7 +232,7 @@ def test_op(self): class TestOpMatmulFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -240,7 +240,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_w = tf.placeholder(tf.float64, [4, 4], "t_w") @@ -278,7 +278,7 @@ def test_op(self): class TestOpMatmulFlt2fixNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -286,7 +286,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_w = tf.placeholder(tf.float64, [4, 4], "t_w") @@ -324,7 +324,7 @@ def test_op(self): class TestOpMulFltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -332,7 +332,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_w = tf.placeholder(tf.float64, [4, 4], "t_w") @@ -370,7 +370,7 @@ def test_op(self): class TestOpQuantizeNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -378,7 +378,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_y = op_module.quantize_nvnmd(t_x, 0, 23, 23, -1) @@ -396,7 +396,7 @@ def test_op(self): class TestOpTanh4FltNvnmd(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -404,7 +404,7 @@ def setUp(self): ) self.sess = self.cached_session(config=config).__enter__() - def test_op(self): + def test_op(self) -> None: # graph t_x = tf.placeholder(tf.float64, [None, 4], "t_x") t_y = op_module.tanh4_flt_nvnmd(t_x) diff --git a/source/tests/tf/test_pairwise_dprc.py b/source/tests/tf/test_pairwise_dprc.py index 22c73beb3b..450134a36c 100644 --- a/source/tests/tf/test_pairwise_dprc.py +++ b/source/tests/tf/test_pairwise_dprc.py @@ -51,7 +51,7 @@ class TestPairwiseOP(tf.test.TestCase): """Test dprc_pairwise_idx OP.""" - def test_op_single_frame(self): + def test_op_single_frame(self) -> None: """Test dprc_pairwise_idx OP with a single frame.""" # same as C++ tests idxs = np.array([[1, 1, 1, 0, 0, 2, 2, 2, 3, 3, 0, 1]], dtype=int) @@ -104,7 +104,7 @@ def test_op_single_frame(self): class TestConvertForwardMapOP(tf.test.TestCase): """Test convert_forward_map OP.""" - def test_convert_forward_map(self): + def test_convert_forward_map(self) -> None: forward_qmmm_map = np.array( [ [3, 4, 0, 1, 2, 10, 11], @@ -269,7 +269,7 @@ def test_convert_forward_map(self): f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestPairwiseModel(tf.test.TestCase): - def test_gather_placeholder(self): + def test_gather_placeholder(self) -> None: coord = np.arange(12 * 3, dtype=np.float64).reshape(1, 12, 3) idxs = np.array([[1, 1, 1, 0, 0, 2, 2, 2, 3, 3, 0, 1]], dtype=int) natoms = np.array([10, 12, 10], dtype=int) @@ -343,7 +343,7 @@ def test_gather_placeholder(self): ), ) - def test_model_ener(self): + def test_model_ener(self) -> None: jfile = tests_path / "pairwise_dprc.json" jdata = j_loader(jfile) model = Model(**jdata["model"]) @@ -531,7 +531,7 @@ def test_model_ener(self): [DataRequirementItem("aparam", 1, atomic=True, must=True, high_prec=False)], ) - def test_nloc(self): + def test_nloc(self) -> None: jfile = tests_path / "pairwise_dprc.json" jdata = j_loader(jfile) model = Model(**jdata["model"]) @@ -826,7 +826,7 @@ def _init_models(): ) class TestPairwiseCompress(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() cls.dp_original = DeepPotential(FROZEN_MODEL) cls.dp_compressed = DeepPotential(COMPRESSED_MODEL) @@ -907,7 +907,7 @@ def setUp(self) -> None: # self.idxs = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0,0,0,0,0,0,0,0]).astype(np.float64) self.type_map = ["C", "N", "O", "H", "OW", "HW"] - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), len(self.type_map)) self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp_original.get_type_map(), self.type_map) @@ -922,7 +922,7 @@ def test_attrs(self): self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) self.assertEqual(self.dp_compressed.get_dim_aparam(), 1) - def test_1frame(self): + def test_1frame(self) -> None: ee0, ff0, vv0 = self.dp_original.eval( self.coords, self.box, @@ -951,7 +951,7 @@ def test_1frame(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( self.coords, self.box, @@ -986,7 +986,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(ee0, ee1, default_places) np.testing.assert_almost_equal(vv0, vv1, default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = None ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( diff --git a/source/tests/tf/test_parallel_training.py b/source/tests/tf/test_parallel_training.py index 0d9d7c1b94..fd17052c9f 100644 --- a/source/tests/tf/test_parallel_training.py +++ b/source/tests/tf/test_parallel_training.py @@ -13,7 +13,7 @@ class TestSingleMachine(unittest.TestCase): - def setUp(self): + def setUp(self) -> None: try: import horovod # noqa: F401 except ImportError as e: @@ -22,7 +22,7 @@ def setUp(self): ) from e self.input_file = str(tests_path / "model_compression" / "input.json") - def test_two_workers(self): + def test_two_workers(self) -> None: command = "horovodrun -np 2 dp train -m workers " + self.input_file penv = os.environ.copy() num_gpus = len(get_gpus() or []) diff --git a/source/tests/tf/test_polar_se_a.py b/source/tests/tf/test_polar_se_a.py index f71ead0d0f..a97e0b5945 100644 --- a/source/tests/tf/test_polar_se_a.py +++ b/source/tests/tf/test_polar_se_a.py @@ -35,10 +35,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "polar_se_a.json" jdata = j_loader(jfile) @@ -200,7 +200,7 @@ def test_model(self): places = 10 np.testing.assert_almost_equal(pv, spv, places) - def test_data_stat(self): + def test_data_stat(self) -> None: jfile = "polar_se_a.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_polar_se_a_tebd.py b/source/tests/tf/test_polar_se_a_tebd.py index 7f8d304af9..72b2aa373d 100644 --- a/source/tests/tf/test_polar_se_a_tebd.py +++ b/source/tests/tf/test_polar_se_a_tebd.py @@ -38,10 +38,10 @@ f"The current tf version {tf.__version__} is too low to run the new testing model.", ) class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data() - def test_model(self): + def test_model(self) -> None: jfile = "polar_se_a_tebd.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_prod_env_mat.py b/source/tests/tf/test_prod_env_mat.py index ac1c16bf97..a1ea2d7ce7 100644 --- a/source/tests/tf/test_prod_env_mat.py +++ b/source/tests/tf/test_prod_env_mat.py @@ -10,7 +10,7 @@ class TestProdEnvMat(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.nframes = 2 self.dcoord = [ @@ -1027,7 +1027,7 @@ def setUp(self): self.t_avg = tf.constant(davg.astype(GLOBAL_NP_FLOAT_PRECISION)) self.t_std = tf.constant(dstd.astype(GLOBAL_NP_FLOAT_PRECISION)) - def test_pbc_self_built_nlist(self): + def test_pbc_self_built_nlist(self) -> None: tem, tem_deriv, trij, tnlist = op_module.prod_env_mat_a( self.tcoord, self.ttype, @@ -1059,7 +1059,7 @@ def test_pbc_self_built_nlist(self): for ff in range(self.nframes): np.testing.assert_almost_equal(dem[ff], self.pbc_expected_output, 5) - def test_pbc_self_built_nlist_deriv(self): + def test_pbc_self_built_nlist_deriv(self) -> None: hh = 1e-4 tem, tem_deriv, trij, tnlist = op_module.prod_env_mat_a( self.tcoord, @@ -1078,7 +1078,7 @@ def test_pbc_self_built_nlist_deriv(self): self.sess.run(tf.global_variables_initializer()) self.check_deriv_numerical_deriv(hh, tem, tem_deriv, trij, tnlist) - def test_nopbc_self_built_nlist(self): + def test_nopbc_self_built_nlist(self) -> None: tem, tem_deriv, trij, tnlist = op_module.prod_env_mat_a( self.tcoord, self.ttype, @@ -1110,7 +1110,7 @@ def test_nopbc_self_built_nlist(self): for ff in range(self.nframes): np.testing.assert_almost_equal(dem[ff], self.nopbc_expected_output, 5) - def test_nopbc_self_built_nlist_deriv(self): + def test_nopbc_self_built_nlist_deriv(self) -> None: hh = 1e-4 tem, tem_deriv, trij, tnlist = op_module.prod_env_mat_a( self.tcoord, @@ -1129,7 +1129,7 @@ def test_nopbc_self_built_nlist_deriv(self): self.sess.run(tf.global_variables_initializer()) self.check_deriv_numerical_deriv(hh, tem, tem_deriv, trij, tnlist) - def check_deriv_numerical_deriv(self, hh, tem, tem_deriv, trij, tnlist): + def check_deriv_numerical_deriv(self, hh, tem, tem_deriv, trij, tnlist) -> None: dem_, dem_deriv_, drij_, dnlist_ = self.sess.run( [tem, tem_deriv, trij, tnlist], feed_dict={ diff --git a/source/tests/tf/test_prod_force.py b/source/tests/tf/test_prod_force.py index 7d3bcee6ce..33c2ad45fc 100644 --- a/source/tests/tf/test_prod_force.py +++ b/source/tests/tf/test_prod_force.py @@ -12,7 +12,7 @@ class TestProdForce(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: config = tf.ConfigProto() if int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)): config.graph_options.rewrite_options.custom_optimizers.add().name = ( @@ -1016,7 +1016,7 @@ def setUp(self): ) self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms") - def test_prod_force(self): + def test_prod_force(self) -> None: tforce = op_module.prod_force_se_a( self.tnet_deriv, self.tem_deriv, @@ -1040,7 +1040,7 @@ def test_prod_force(self): np.testing.assert_almost_equal(dforce[ff], self.expected_force, 5) @unittest.skipIf(tf.test.is_gpu_available(), reason="Not supported in GPUs") - def test_prod_force_parallel(self): + def test_prod_force_parallel(self) -> None: forces = [] for ii in range(4): tforce = op_module.parallel_prod_force_se_a( diff --git a/source/tests/tf/test_prod_force_grad.py b/source/tests/tf/test_prod_force_grad.py index 49e63d161c..6fea4a631d 100644 --- a/source/tests/tf/test_prod_force_grad.py +++ b/source/tests/tf/test_prod_force_grad.py @@ -9,7 +9,7 @@ class TestProdForceGrad(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.nframes = 2 self.dcoord = [ @@ -1112,7 +1112,7 @@ def setUp(self): ) self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms") - def test_prod_force_grad(self): + def test_prod_force_grad(self) -> None: tgrad_net = op_grads_module.prod_force_se_a_grad( self.tgrad, self.tnet_deriv, diff --git a/source/tests/tf/test_prod_virial.py b/source/tests/tf/test_prod_virial.py index fa6347382e..e99e06dcc8 100644 --- a/source/tests/tf/test_prod_virial.py +++ b/source/tests/tf/test_prod_virial.py @@ -9,7 +9,7 @@ class TestProdVirial(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.nframes = 2 self.dcoord = [ @@ -1495,7 +1495,7 @@ def setUp(self): ) self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms") - def test_prod_virial(self): + def test_prod_virial(self) -> None: tvirial, tatom_virial = op_module.prod_virial_se_a( self.tnet_deriv, self.tem_deriv, diff --git a/source/tests/tf/test_prod_virial_grad.py b/source/tests/tf/test_prod_virial_grad.py index 470441a939..67d9e3050f 100644 --- a/source/tests/tf/test_prod_virial_grad.py +++ b/source/tests/tf/test_prod_virial_grad.py @@ -9,7 +9,7 @@ class TestProdVirialGrad(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: self.sess = self.cached_session().__enter__() self.nframes = 2 self.dcoord = [ @@ -1299,7 +1299,7 @@ def setUp(self): ) self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms") - def test_prod_virial_grad(self): + def test_prod_virial_grad(self) -> None: tgrad_net = op_grads_module.prod_virial_se_a_grad( self.tgrad, self.tnet_deriv, diff --git a/source/tests/tf/test_tab_nonsmth.py b/source/tests/tf/test_tab_nonsmth.py index 7132d0c206..c13289cfb8 100644 --- a/source/tests/tf/test_tab_nonsmth.py +++ b/source/tests/tf/test_tab_nonsmth.py @@ -26,7 +26,7 @@ ) -def _make_tab(ntype): +def _make_tab(ntype) -> None: xx = np.arange(0, 9, 0.001) yy = 1000 / (xx + 0.5) ** 6 prt = xx @@ -38,7 +38,7 @@ def _make_tab(ntype): class IntplInter(Inter): - def setUp(self, data, sess=None): + def setUp(self, data, sess=None) -> None: # tabulated Inter.setUp(self, data, sess=sess) _make_tab(data.get_ntypes()) @@ -63,7 +63,7 @@ def setUp(self, data, sess=None): initializer=tf.constant_initializer(tab_data), ) - def tearDown(self): + def tearDown(self) -> None: os.remove("tab.xvg") def comp_interpl_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): @@ -176,21 +176,21 @@ class TestTabNonSmooth(IntplInter, tf.test.TestCase): # unittest.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() IntplInter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, places=5, suffix="_tab") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, places=5, suffix="_tab") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, places=5, suffix="_tab") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, places=5, suffix="_tab") diff --git a/source/tests/tf/test_tab_smooth.py b/source/tests/tf/test_tab_smooth.py index e0cf564cd6..df5fc6a147 100644 --- a/source/tests/tf/test_tab_smooth.py +++ b/source/tests/tf/test_tab_smooth.py @@ -25,7 +25,7 @@ ) -def _make_tab(ntype): +def _make_tab(ntype) -> None: xx = np.arange(0, 9, 0.001) yy = 1000 / (xx + 0.5) ** 6 prt = xx @@ -37,7 +37,7 @@ def _make_tab(ntype): class IntplInter(Inter): - def setUp(self, data, sess=None): + def setUp(self, data, sess=None) -> None: # tabulated Inter.setUp(self, data, sess=sess) _make_tab(data.get_ntypes()) @@ -62,7 +62,7 @@ def setUp(self, data, sess=None): initializer=tf.constant_initializer(tab_data), ) - def tearDown(self): + def tearDown(self) -> None: os.remove("tab.xvg") def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): @@ -173,21 +173,21 @@ class TestTabSmooth(IntplInter, tf.test.TestCase): # unittest.TestCase.__init__(self, *args, **kwargs) # self.controller = object() - def setUp(self): + def setUp(self) -> None: self.places = 5 data = Data() IntplInter.setUp(self, data, sess=self.cached_session().__enter__()) - def test_force(self): + def test_force(self) -> None: force_test(self, self, places=5, suffix="_tab_smth") - def test_virial(self): + def test_virial(self) -> None: virial_test(self, self, places=5, suffix="_tab_smth") - def test_force_dw(self): + def test_force_dw(self) -> None: force_dw_test(self, self, places=8, suffix="_tab_smth") - def test_virial_dw(self): + def test_virial_dw(self) -> None: virial_dw_test(self, self, places=8, suffix="_tab_smth") diff --git a/source/tests/tf/test_tabulate.py b/source/tests/tf/test_tabulate.py index 0d46293b62..63a04afa9b 100644 --- a/source/tests/tf/test_tabulate.py +++ b/source/tests/tf/test_tabulate.py @@ -15,7 +15,7 @@ class TestDPTabulate(unittest.TestCase): - def test_op_tanh(self): + def test_op_tanh(self) -> None: w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype="double", @@ -61,7 +61,7 @@ def test_op_tanh(self): places = 15 np.testing.assert_almost_equal(dy_array, answer, places) - def test_op_gelu(self): + def test_op_gelu(self) -> None: w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype="double", diff --git a/source/tests/tf/test_train.py b/source/tests/tf/test_train.py index 2d296558fe..4a0fb55494 100644 --- a/source/tests/tf/test_train.py +++ b/source/tests/tf/test_train.py @@ -23,14 +23,14 @@ def setUp(self) -> None: self.mock_min_nbor_dist = random.Random(GLOBAL_SEED).random() return super().setUp() - def test_train_parse_auto_sel(self): + def test_train_parse_auto_sel(self) -> None: self.assertTrue(self.update_sel.parse_auto_sel("auto")) self.assertTrue(self.update_sel.parse_auto_sel("auto:12")) self.assertTrue(self.update_sel.parse_auto_sel("auto:12:13")) self.assertFalse(self.update_sel.parse_auto_sel([1, 2])) self.assertFalse(self.update_sel.parse_auto_sel("abc:12:13")) - def test_train_parse_auto_sel_ratio(self): + def test_train_parse_auto_sel_ratio(self) -> None: self.assertEqual(self.update_sel.parse_auto_sel_ratio("auto"), 1.1) self.assertEqual(self.update_sel.parse_auto_sel_ratio("auto:1.2"), 1.2) with self.assertRaises(RuntimeError): @@ -42,7 +42,7 @@ def test_train_parse_auto_sel_ratio(self): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_one_sel(self, sel_mock, get_data_mock): + def test_update_one_sel(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] get_data_mock.return_value = None min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") @@ -56,7 +56,7 @@ def test_update_one_sel(self, sel_mock, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_hybrid(self, sel_mock, get_data_mock): + def test_update_sel_hybrid(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] get_data_mock.return_value = None jdata = { @@ -88,7 +88,7 @@ def test_update_sel_hybrid(self, sel_mock, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel(self, sel_mock, get_data_mock): + def test_update_sel(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] get_data_mock.return_value = None jdata = { @@ -104,7 +104,7 @@ def test_update_sel(self, sel_mock, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_auto(self, sel_mock, get_data_mock): + def test_update_sel_atten_auto(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] get_data_mock.return_value = None jdata = { @@ -132,7 +132,7 @@ def test_update_sel_atten_auto(self, sel_mock, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_int(self, sel_mock, get_data_mock): + def test_update_sel_atten_int(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] get_data_mock.return_value = None jdata = { @@ -160,7 +160,7 @@ def test_update_sel_atten_int(self, sel_mock, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") - def test_update_sel_atten_list(self, sel_mock, get_data_mock): + def test_update_sel_atten_list(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist, [25] get_data_mock.return_value = None jdata = { @@ -187,7 +187,7 @@ def test_update_sel_atten_list(self, sel_mock, get_data_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.tf.entrypoints.train.get_data") - def test_skip_loc_frame(self, get_data_mock): + def test_skip_loc_frame(self, get_data_mock) -> None: get_data_mock.return_value = None jdata = { "model": { @@ -211,7 +211,7 @@ def test_skip_loc_frame(self, get_data_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.tf.entrypoints.train.get_data") - def test_skip_frozen(self, get_data_mock): + def test_skip_frozen(self, get_data_mock) -> None: get_data_mock.return_value = None jdata = { "model": { @@ -224,7 +224,7 @@ def test_skip_frozen(self, get_data_mock): self.assertEqual(jdata, expected_out) @patch("deepmd.tf.entrypoints.train.get_data") - def test_skip_linear_frozen(self, get_data_mock): + def test_skip_linear_frozen(self, get_data_mock) -> None: get_data_mock.return_value = None jdata = { "model": { @@ -244,7 +244,7 @@ def test_skip_linear_frozen(self, get_data_mock): @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_min_nbor_dist") - def test_pairwise_dprc(self, sel_mock, get_data_mock): + def test_pairwise_dprc(self, sel_mock, get_data_mock) -> None: sel_mock.return_value = self.mock_min_nbor_dist get_data_mock.return_value = None jdata = { @@ -263,7 +263,7 @@ def test_pairwise_dprc(self, sel_mock, get_data_mock): jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_wrap_up_4(self): + def test_wrap_up_4(self) -> None: self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) diff --git a/source/tests/tf/test_transfer.py b/source/tests/tf/test_transfer.py index 04a239724a..8cc2ff92c5 100644 --- a/source/tests/tf/test_transfer.py +++ b/source/tests/tf/test_transfer.py @@ -26,14 +26,14 @@ default_places = 10 -def _file_delete(file): +def _file_delete(file) -> None: if os.path.exists(file): os.remove(file) class TestTransform(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: cls.old_model = str(tests_path / "dp-old.pb") cls.raw_model = str(tests_path / "dp-raw.pb") cls.new_model = str(tests_path / "dp-new.pb") @@ -170,19 +170,19 @@ def setUpClass(cls): ) @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: _file_delete(cls.old_model) _file_delete(cls.raw_model) _file_delete(cls.new_model) - def test_attrs(self): + def test_attrs(self) -> None: self.assertEqual(self.dp.get_ntypes(), 2) self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places=default_places) self.assertEqual(self.dp.get_type_map(), ["O", "H"]) self.assertEqual(self.dp.get_dim_fparam(), 0) self.assertEqual(self.dp.get_dim_aparam(), 0) - def test_1frame(self): + def test_1frame(self) -> None: ee, ff, vv = self.dp.eval(self.coords, self.box, self.atype, atomic=False) # check shape of the returns nframes = 1 @@ -199,7 +199,7 @@ def test_1frame(self): expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) - def test_1frame_atm(self): + def test_1frame_atm(self) -> None: ee, ff, vv, ae, av = self.dp.eval( self.coords, self.box, self.atype, atomic=True ) @@ -226,7 +226,7 @@ def test_1frame_atm(self): expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1) np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) - def test_2frame_atm(self): + def test_2frame_atm(self) -> None: coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) ee, ff, vv, ae, av = self.dp.eval(coords2, box2, self.atype, atomic=True) diff --git a/source/tests/tf/test_type_embed.py b/source/tests/tf/test_type_embed.py index fd89f8fd10..da3311f2e2 100644 --- a/source/tests/tf/test_type_embed.py +++ b/source/tests/tf/test_type_embed.py @@ -11,7 +11,7 @@ class TestTypeEbd(tf.test.TestCase): - def test_embed_atom_type(self): + def test_embed_atom_type(self) -> None: ntypes = 3 natoms = tf.constant([5, 5, 3, 0, 2]) type_embedding = tf.constant( @@ -27,7 +27,7 @@ def test_embed_atom_type(self): atom_embed = sess.run(atom_embed) np.testing.assert_almost_equal(atom_embed, expected_out, 10) - def test_type_embed_net(self): + def test_type_embed_net(self) -> None: ten = TypeEmbedNet( ntypes=2, neuron=[2, 4, 8], seed=1, uniform_seed=True, use_tebd_bias=True ) diff --git a/source/tests/tf/test_type_one_side.py b/source/tests/tf/test_type_one_side.py index 2ab7cc03a7..e26c4a9710 100644 --- a/source/tests/tf/test_type_one_side.py +++ b/source/tests/tf/test_type_one_side.py @@ -20,10 +20,10 @@ class TestModel(tf.test.TestCase): - def setUp(self): + def setUp(self) -> None: gen_data(nframes=2) - def test_descriptor_one_side_exclude_types(self): + def test_descriptor_one_side_exclude_types(self) -> None: """When we enable type_one_side, the descriptor should be the same for different types, when its environments are the same. @@ -140,7 +140,7 @@ def test_descriptor_one_side_exclude_types(self): model_dout1_failed[0], model_dout2_failed[0], 10 ) - def test_se_r_one_side_exclude_types(self): + def test_se_r_one_side_exclude_types(self) -> None: """se_r.""" jfile = "water_se_r.json" jdata = j_loader(jfile) diff --git a/source/tests/tf/test_virtual_type.py b/source/tests/tf/test_virtual_type.py index 40ec2674dd..697467c565 100644 --- a/source/tests/tf/test_virtual_type.py +++ b/source/tests/tf/test_virtual_type.py @@ -28,7 +28,7 @@ class TestVirtualType(unittest.TestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: convert_pbtxt_to_pb( str(infer_path / os.path.join("virtual_type.pbtxt")), "virtual_type.pb", @@ -36,7 +36,7 @@ def setUpClass(cls): cls.dp = DeepPot("virtual_type.pb") os.remove("virtual_type.pb") - def setUp(self): + def setUp(self) -> None: self.coords = np.array( [ 12.83, @@ -62,7 +62,7 @@ def setUp(self): self.atype = [0, 1, 1, 0, 1, 1] self.box = None - def test_virtual_type(self): + def test_virtual_type(self) -> None: nloc = len(self.atype) nghost = 10 e1, f1, v1, ae1, av1 = self.dp.eval( @@ -82,7 +82,7 @@ def test_virtual_type(self): np.testing.assert_almost_equal(ae1, ae2[:, :nloc]) np.testing.assert_almost_equal(av1, av2[:, :nloc]) - def test_infer_mixed_type(self): + def test_infer_mixed_type(self) -> None: nloc = len(self.atype) nghost = 10 e, f, v, ae, av = self.dp.eval( @@ -123,7 +123,7 @@ class TestTrainVirtualType(unittest.TestCase): def setUp(self) -> None: gen_data(mixed_type=True, virtual_type=True) - def test_data_mixed_type(self): + def test_data_mixed_type(self) -> None: jfile = "water_se_atten_mixed_type.json" jdata = j_loader(jfile) diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py index 97a6cf707b..054272b8c6 100644 --- a/source/tests/universal/common/cases/atomic_model/utils.py +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -48,55 +48,55 @@ class AtomicModelTestCase: epsilon_dict: dict[str, Optional[float]] """Dictionary of epsilons in each test.""" - def test_get_type_map(self): + def test_get_type_map(self) -> None: """Test get_type_map.""" for module in self.modules_to_test: self.assertEqual(module.get_type_map(), self.expected_type_map) - def test_get_rcut(self): + def test_get_rcut(self) -> None: """Test get_rcut.""" for module in self.modules_to_test: self.assertAlmostEqual(module.get_rcut(), self.expected_rcut) - def test_get_dim_fparam(self): + def test_get_dim_fparam(self) -> None: """Test get_dim_fparam.""" for module in self.modules_to_test: self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam) - def test_get_dim_aparam(self): + def test_get_dim_aparam(self) -> None: """Test get_dim_aparam.""" for module in self.modules_to_test: self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam) - def test_get_sel_type(self): + def test_get_sel_type(self) -> None: """Test get_sel_type.""" for module in self.modules_to_test: self.assertEqual(module.get_sel_type(), self.expected_sel_type) - def test_is_aparam_nall(self): + def test_is_aparam_nall(self) -> None: """Test is_aparam_nall.""" for module in self.modules_to_test: self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall) - def test_get_nnei(self): + def test_get_nnei(self) -> None: """Test get_nnei.""" expected_nnei = sum(self.expected_sel) for module in self.modules_to_test: self.assertEqual(module.get_nnei(), expected_nnei) - def test_get_ntypes(self): + def test_get_ntypes(self) -> None: """Test get_ntypes.""" for module in self.modules_to_test: self.assertEqual(module.get_ntypes(), len(self.expected_type_map)) - def test_has_message_passing(self): + def test_has_message_passing(self) -> None: """Test has_message_passing.""" for module in self.modules_to_test: self.assertEqual( module.has_message_passing(), self.expected_has_message_passing ) - def test_forward(self): + def test_forward(self) -> None: """Test forward.""" nf = 1 rng = np.random.default_rng(GLOBAL_SEED) diff --git a/source/tests/universal/common/cases/cases.py b/source/tests/universal/common/cases/cases.py index 7830a20989..d625ef0d35 100644 --- a/source/tests/universal/common/cases/cases.py +++ b/source/tests/universal/common/cases/cases.py @@ -4,7 +4,7 @@ # originally copied from source/tests/pt/model/test_env_mat.py class TestCaseSingleFrameWithNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nall = 4 @@ -71,7 +71,7 @@ def setUp(self): class TestCaseSingleFrameWithNlistWithVirtual: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 4 self.nall = 5 @@ -123,7 +123,7 @@ def setUp(self): class TestCaseSingleFrameWithoutNlist: - def setUp(self): + def setUp(self) -> None: # nloc == 3, nall == 4 self.nloc = 3 self.nf, self.nt = 1, 2 diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index 4c8219d347..96eee7b757 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -24,10 +24,10 @@ class DescriptorTestCase(TestCaseSingleFrameWithNlist): """Common test case for descriptor.""" - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) - def test_forward_consistency(self): + def test_forward_consistency(self) -> None: ret = [] for module in self.modules_to_test: module = self.forward_wrapper(module) @@ -57,7 +57,7 @@ def test_forward_consistency(self): def test_exclude_types( self, - ): + ) -> None: coord_ext_device = self.coord_ext atype_ext_device = self.atype_ext nlist_device = self.nlist @@ -95,7 +95,7 @@ def test_exclude_types( ) np.testing.assert_allclose(rd0, rd_ex) - def test_change_type_map(self): + def test_change_type_map(self) -> None: if ( not self.module.mixed_types() or getattr(self.module, "sel_no_mixed_types", None) is not None @@ -180,7 +180,7 @@ def test_change_type_map(self): ) np.testing.assert_allclose(rd_old_tm, rd_new_tm) - def test_change_type_map_extend_stat(self): + def test_change_type_map_extend_stat(self) -> None: if ( not self.module.mixed_types() or getattr(self.module, "sel_no_mixed_types", None) is not None @@ -317,7 +317,9 @@ def get_rand_stat_item(self, rng, mean, std): std_rand = self.convert_from_numpy(std_rand) return mean_rand, std_rand - def check_expect_stat(self, type_index_map, stat_small, stat_large, stat_result): + def check_expect_stat( + self, type_index_map, stat_small, stat_large, stat_result + ) -> None: if not isinstance(stat_small, list): self.check_expect_stat_item( type_index_map, stat_small, stat_large, stat_result @@ -330,7 +332,7 @@ def check_expect_stat(self, type_index_map, stat_small, stat_large, stat_result) def check_expect_stat_item( self, type_index_map, stat_small, stat_large, stat_result - ): + ) -> None: stat_small = self.convert_to_numpy(stat_small) stat_large = self.convert_to_numpy(stat_large) stat_result = self.convert_to_numpy(stat_result) diff --git a/source/tests/universal/common/cases/fitting/utils.py b/source/tests/universal/common/cases/fitting/utils.py index 2ab5fd911b..de6b12c3a2 100644 --- a/source/tests/universal/common/cases/fitting/utils.py +++ b/source/tests/universal/common/cases/fitting/utils.py @@ -21,7 +21,7 @@ class FittingTestCase(TestCaseSingleFrameWithNlist): """Common test case for descriptor.""" - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.input_dict = { "ntypes": self.nt, @@ -30,7 +30,7 @@ def setUp(self): "type_map": ["O", "H"], } - def test_forward_consistency(self): + def test_forward_consistency(self) -> None: serialize_dict = self.module.serialize() # set random bias rng = np.random.default_rng() @@ -66,7 +66,7 @@ def test_forward_consistency(self): def test_exclude_types( self, - ): + ) -> None: atype_device = self.atype_ext[:, : self.nloc] serialize_dict = self.module.serialize() # set random bias @@ -104,7 +104,7 @@ def test_exclude_types( )[var_name] np.testing.assert_allclose(rd, rd_ex) - def test_change_type_map(self): + def test_change_type_map(self) -> None: if not self.module.mixed_types: # skip if not mixed_types return diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 628c415eb2..87c0209169 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -61,62 +61,62 @@ class ModelTestCase: epsilon_dict: dict[str, Optional[float]] """Dictionary of epsilons in each test.""" - def test_get_type_map(self): + def test_get_type_map(self) -> None: """Test get_type_map.""" for module in self.modules_to_test: self.assertEqual(module.get_type_map(), self.expected_type_map) - def test_get_rcut(self): + def test_get_rcut(self) -> None: """Test get_rcut.""" for module in self.modules_to_test: self.assertAlmostEqual(module.get_rcut(), self.expected_rcut) - def test_get_dim_fparam(self): + def test_get_dim_fparam(self) -> None: """Test get_dim_fparam.""" for module in self.modules_to_test: self.assertEqual(module.get_dim_fparam(), self.expected_dim_fparam) - def test_get_dim_aparam(self): + def test_get_dim_aparam(self) -> None: """Test get_dim_aparam.""" for module in self.modules_to_test: self.assertEqual(module.get_dim_aparam(), self.expected_dim_aparam) - def test_get_sel_type(self): + def test_get_sel_type(self) -> None: """Test get_sel_type.""" for module in self.modules_to_test: self.assertEqual(module.get_sel_type(), self.expected_sel_type) - def test_is_aparam_nall(self): + def test_is_aparam_nall(self) -> None: """Test is_aparam_nall.""" for module in self.modules_to_test: self.assertEqual(module.is_aparam_nall(), self.expected_aparam_nall) - def test_model_output_type(self): + def test_model_output_type(self) -> None: """Test model_output_type.""" for module in self.modules_to_test: self.assertEqual( module.model_output_type(), self.expected_model_output_type ) - def test_get_nnei(self): + def test_get_nnei(self) -> None: """Test get_nnei.""" expected_nnei = sum(self.expected_sel) for module in self.modules_to_test: self.assertEqual(module.get_nnei(), expected_nnei) - def test_get_ntypes(self): + def test_get_ntypes(self) -> None: """Test get_ntypes.""" for module in self.modules_to_test: self.assertEqual(module.get_ntypes(), len(self.expected_type_map)) - def test_has_message_passing(self): + def test_has_message_passing(self) -> None: """Test has_message_passing.""" for module in self.modules_to_test: self.assertEqual( module.has_message_passing(), self.expected_has_message_passing ) - def test_forward(self): + def test_forward(self) -> None: """Test forward and forward_lower.""" test_spin = getattr(self, "test_spin", False) nf = 1 @@ -225,7 +225,7 @@ def test_forward(self): continue np.testing.assert_allclose(rr1, rr2, atol=aprec) - def test_zero_forward(self): + def test_zero_forward(self) -> None: test_spin = getattr(self, "test_spin", False) nf = 1 natoms = 0 @@ -329,7 +329,7 @@ def test_zero_forward(self): np.testing.assert_allclose(rr1, rr2, atol=aprec) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_permutation(self): + def test_permutation(self) -> None: """Test permutation.""" if getattr(self, "skip_test_permutation", False): self.skipTest("Skip test permutation.") @@ -415,7 +415,7 @@ def test_permutation(self): raise RuntimeError(f"Unknown output key: {kk}") @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_trans(self): + def test_trans(self) -> None: """Test translation.""" if getattr(self, "skip_test_trans", False): self.skipTest("Skip test translation.") @@ -484,7 +484,7 @@ def test_trans(self): raise RuntimeError(f"Unknown output key: {kk}") @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_rot(self): + def test_rot(self) -> None: """Test rotation.""" if getattr(self, "skip_test_rot", False): self.skipTest("Skip test rotation.") @@ -674,7 +674,7 @@ def test_rot(self): raise RuntimeError(f"Unknown output key: {kk}") @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_smooth(self): + def test_smooth(self) -> None: """Test smooth.""" if getattr(self, "skip_test_smooth", False): self.skipTest("Skip test smooth.") @@ -781,7 +781,7 @@ def test_smooth(self): raise RuntimeError(f"Unknown output key: {kk}") @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") - def test_autodiff(self): + def test_autodiff(self) -> None: """Test autodiff.""" if getattr(self, "skip_test_autodiff", False): self.skipTest("Skip test autodiff.") @@ -921,7 +921,7 @@ def ff_cell(bb): pass @unittest.skipIf(TEST_DEVICE == "cpu" and CI, "Skip test on CPU.") - def test_device_consistence(self): + def test_device_consistence(self) -> None: """Test forward consistency between devices.""" test_spin = getattr(self, "test_spin", False) nf = 1 diff --git a/source/tests/universal/common/cases/utils/utils.py b/source/tests/universal/common/cases/utils/utils.py index 9f86ca1feb..f9fc697cf2 100644 --- a/source/tests/universal/common/cases/utils/utils.py +++ b/source/tests/universal/common/cases/utils/utils.py @@ -17,7 +17,7 @@ class TypeEmbdTestCase(TestCaseSingleFrameWithNlist): """Common test case for type embedding network.""" - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.input_dict = { "ntypes": self.nt, @@ -27,7 +27,7 @@ def setUp(self): } self.module_input = {} - def test_change_type_map(self): + def test_change_type_map(self) -> None: atype_ext_device = self.atype_ext # type_map for data and exclude_types original_type_map = ["O", "H"] diff --git a/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py b/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py index 8e7324e2bc..7b579ae82c 100644 --- a/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py +++ b/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py @@ -102,7 +102,7 @@ @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestEnergyAtomicModelDP(unittest.TestCase, EnerAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: EnerAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -169,7 +169,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestDosAtomicModelDP(unittest.TestCase, DosAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DosAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -231,7 +231,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestDipoleAtomicModelDP(unittest.TestCase, DipoleAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DipoleAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -294,7 +294,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestPolarAtomicModelDP(unittest.TestCase, PolarAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PolarAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -355,7 +355,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestZBLAtomicModelDP(unittest.TestCase, ZBLAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ZBLAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -433,7 +433,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestPropertyAtomicModelDP(unittest.TestCase, PropertyAtomicModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PropertyAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] diff --git a/source/tests/universal/dpmodel/backend.py b/source/tests/universal/dpmodel/backend.py index 4f624ae501..daf0213311 100644 --- a/source/tests/universal/dpmodel/backend.py +++ b/source/tests/universal/dpmodel/backend.py @@ -54,7 +54,7 @@ def modules_to_test(self): return modules @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: super().tearDownClass() if hasattr(cls, "module"): del cls.module diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index fc7ee8b075..7911cb9395 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -522,7 +522,7 @@ def DescriptorParamHybridMixedTTebd(ntypes, rcut, rcut_smth, sel, type_map, **kw ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestDescriptorDP(unittest.TestCase, DescriptorTest, DPTestCase): - def setUp(self): + def setUp(self) -> None: DescriptorTest.setUp(self) (DescriptorParam, Descrpt) = self.param[0] self.module_class = Descrpt diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index f64faee76f..fe6ffd2e09 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -239,7 +239,7 @@ def FittingParamProperty( ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestFittingDP(unittest.TestCase, FittingTest, DPTestCase): - def setUp(self): + def setUp(self) -> None: ((FittingParam, Fitting), self.mixed_types) = self.param FittingTest.setUp(self) self.module_class = Fitting diff --git a/source/tests/universal/dpmodel/model/test_model.py b/source/tests/universal/dpmodel/model/test_model.py index 265dc43c6c..a07a29dcf5 100644 --- a/source/tests/universal/dpmodel/model/test_model.py +++ b/source/tests/universal/dpmodel/model/test_model.py @@ -116,7 +116,7 @@ def skip_model_tests(test_obj): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestEnergyModelDP(unittest.TestCase, EnerModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: EnerModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -204,7 +204,7 @@ def setUpClass(cls): @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestSpinEnergyModelDP(unittest.TestCase, SpinEnerModelTest, DPTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: SpinEnerModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] diff --git a/source/tests/universal/dpmodel/utils/test_type_embed.py b/source/tests/universal/dpmodel/utils/test_type_embed.py index ee3063af7d..32f4d1c99c 100644 --- a/source/tests/universal/dpmodel/utils/test_type_embed.py +++ b/source/tests/universal/dpmodel/utils/test_type_embed.py @@ -19,7 +19,7 @@ @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") class TestTypeEmbd(unittest.TestCase, TypeEmbdTest, DPTestCase): - def setUp(self): + def setUp(self) -> None: TypeEmbdTest.setUp(self) self.module_class = TypeEmbedNet self.module = TypeEmbedNet(**self.input_dict) diff --git a/source/tests/universal/pt/atomc_model/test_atomic_model.py b/source/tests/universal/pt/atomc_model/test_atomic_model.py index b8dd9e684d..f41d384b6b 100644 --- a/source/tests/universal/pt/atomc_model/test_atomic_model.py +++ b/source/tests/universal/pt/atomc_model/test_atomic_model.py @@ -97,7 +97,7 @@ ) class TestEnergyAtomicModelPT(unittest.TestCase, EnerAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: EnerAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -163,7 +163,7 @@ def setUpClass(cls): ) class TestDosAtomicModelPT(unittest.TestCase, DosAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DosAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -224,7 +224,7 @@ def setUpClass(cls): ) class TestDipoleAtomicModelPT(unittest.TestCase, DipoleAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DipoleAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -286,7 +286,7 @@ def setUpClass(cls): ) class TestPolarAtomicModelPT(unittest.TestCase, PolarAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PolarAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -346,7 +346,7 @@ def setUpClass(cls): ) class TestZBLAtomicModelPT(unittest.TestCase, ZBLAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ZBLAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -423,7 +423,7 @@ def setUpClass(cls): ) class TestPropertyAtomicModelPT(unittest.TestCase, PropertyAtomicModelTest, PTTestCase): @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PropertyAtomicModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] diff --git a/source/tests/universal/pt/backend.py b/source/tests/universal/pt/backend.py index 5146fdc79b..ae857d6105 100644 --- a/source/tests/universal/pt/backend.py +++ b/source/tests/universal/pt/backend.py @@ -55,14 +55,14 @@ def modules_to_test(self): return modules @classmethod - def tearDownClass(cls): + def tearDownClass(cls) -> None: super().tearDownClass() if hasattr(cls, "module"): del cls.module cls._get_deserialized_module.cache_clear() cls._get_script_module.cache_clear() - def test_jit(self): + def test_jit(self) -> None: if getattr(self, "skip_test_jit", False): self.skipTest("Skip test jit.") self.script_module diff --git a/source/tests/universal/pt/descriptor/test_descriptor.py b/source/tests/universal/pt/descriptor/test_descriptor.py index 4a5d511deb..349eb65588 100644 --- a/source/tests/universal/pt/descriptor/test_descriptor.py +++ b/source/tests/universal/pt/descriptor/test_descriptor.py @@ -45,7 +45,7 @@ ) # class_param & class ) class TestDescriptorPT(unittest.TestCase, DescriptorTest, PTTestCase): - def setUp(self): + def setUp(self) -> None: DescriptorTest.setUp(self) (DescriptorParam, Descrpt) = self.param[0] self.module_class = Descrpt diff --git a/source/tests/universal/pt/fitting/test_fitting.py b/source/tests/universal/pt/fitting/test_fitting.py index 84934a23ad..efda3b9619 100644 --- a/source/tests/universal/pt/fitting/test_fitting.py +++ b/source/tests/universal/pt/fitting/test_fitting.py @@ -38,7 +38,7 @@ (True, False), # mixed_types ) class TestFittingPT(unittest.TestCase, FittingTest, PTTestCase): - def setUp(self): + def setUp(self) -> None: ((FittingParam, Fitting), self.mixed_types) = self.param FittingTest.setUp(self) self.module_class = Fitting diff --git a/source/tests/universal/pt/model/test_model.py b/source/tests/universal/pt/model/test_model.py index 81c32eb94c..3eb1484c45 100644 --- a/source/tests/universal/pt/model/test_model.py +++ b/source/tests/universal/pt/model/test_model.py @@ -152,7 +152,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: EnerModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -254,7 +254,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DosModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -347,7 +347,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: DipoleModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -440,7 +440,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PolarModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -527,7 +527,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: ZBLModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -639,7 +639,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: SpinEnerModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -755,7 +755,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: PropertyModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] @@ -842,7 +842,7 @@ def modules_to_test(self): return modules @classmethod - def setUpClass(cls): + def setUpClass(cls) -> None: LinearEnerModelTest.setUpClass() (DescriptorParam, Descrpt) = cls.param[0] (FittingParam, Fitting) = cls.param[1] diff --git a/source/tests/universal/pt/utils/test_type_embed.py b/source/tests/universal/pt/utils/test_type_embed.py index 0a53eeeccb..3d28965c29 100644 --- a/source/tests/universal/pt/utils/test_type_embed.py +++ b/source/tests/universal/pt/utils/test_type_embed.py @@ -17,7 +17,7 @@ class TestTypeEmbd(unittest.TestCase, TypeEmbdTest, PTTestCase): - def setUp(self): + def setUp(self) -> None: TypeEmbdTest.setUp(self) self.module_class = TypeEmbedNetConsistent self.module = TypeEmbedNetConsistent(**self.input_dict)