diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py index 5028bc43a3..682859fb0e 100644 --- a/deepmd/dpmodel/output_def.py +++ b/deepmd/dpmodel/output_def.py @@ -267,7 +267,7 @@ def __getitem__( def get_data(self) -> dict[str, OutputVariableDef]: return self.var_defs - def keys(self): # noqa: ANN201 + def keys(self): return self.var_defs.keys() @@ -319,25 +319,25 @@ def get_data( ) -> dict[str, OutputVariableDef]: return self.var_defs - def keys(self): # noqa: ANN201 + def keys(self): return self.var_defs.keys() - def keys_outp(self): # noqa: ANN201 + def keys_outp(self): return self.def_outp.keys() - def keys_redu(self): # noqa: ANN201 + def keys_redu(self): return self.def_redu.keys() - def keys_derv_r(self): # noqa: ANN201 + def keys_derv_r(self): return self.def_derv_r.keys() - def keys_hess_r(self): # noqa: ANN201 + def keys_hess_r(self): return self.def_hess_r.keys() - def keys_derv_c(self): # noqa: ANN201 + def keys_derv_c(self): return self.def_derv_c.keys() - def keys_derv_c_redu(self): # noqa: ANN201 + def keys_derv_c_redu(self): return self.def_derv_c_redu.keys() diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index d48c42ad08..0230b678c6 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -38,7 +38,7 @@ ) -def sigmoid_t(x): # noqa: ANN001, ANN201 +def sigmoid_t(x): """Sigmoid.""" if array_api_compat.is_jax_array(x): from deepmd.jax.env import ( @@ -55,7 +55,7 @@ class Identity(NativeOP): def __init__(self) -> None: super().__init__() - def call(self, x): # noqa: ANN001, ANN201 + def call(self, x): """The Identity operation layer.""" return x @@ -260,7 +260,7 @@ def dim_out(self) -> int: return self.w.shape[1] @support_array_api(version="2022.12") - def call(self, x): # noqa: ANN001, ANN201 + def call(self, x): """Forward pass. Parameters @@ -301,14 +301,14 @@ def get_activation_fn(activation_function: str) -> Callable[[np.ndarray], np.nda activation_function = activation_function.lower() if activation_function == "tanh": - def fn(x): # noqa: ANN001, ANN202 # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) return xp.tanh(x) return fn elif activation_function == "relu": - def fn(x): # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) # https://stackoverflow.com/a/47936476/9567349 return x * xp.astype(x > 0, x.dtype) @@ -316,7 +316,7 @@ def fn(x): # noqa: ANN001, ANN202 return fn elif activation_function in ("gelu", "gelu_tf"): - def fn(x): # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return ( @@ -328,7 +328,7 @@ def fn(x): # noqa: ANN001, ANN202 return fn elif activation_function == "relu6": - def fn(x): # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return xp.where( @@ -338,7 +338,7 @@ def fn(x): # noqa: ANN001, ANN202 return fn elif activation_function == "softplus": - def fn(x): # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return xp.log(1 + xp.exp(x)) @@ -346,14 +346,14 @@ def fn(x): # noqa: ANN001, ANN202 return fn elif activation_function == "sigmoid": - def fn(x): # noqa: ANN001, ANN202 + def fn(x): # generated by GitHub Copilot return sigmoid_t(x) return fn elif activation_function == "silu": - def fn(x): # noqa: ANN001, ANN202 + def fn(x): # generated by GitHub Copilot return x * sigmoid_t(x) @@ -362,13 +362,13 @@ def fn(x): # noqa: ANN001, ANN202 "custom_silu" ): - def sigmoid(x): # noqa: ANN001, ANN202 + def sigmoid(x): return 1 / (1 + np.exp(-x)) - def silu(x): # noqa: ANN001, ANN202 + def silu(x): return x * sigmoid(x) - def silu_grad(x): # noqa: ANN001, ANN202 + def silu_grad(x): sig = sigmoid(x) return sig + x * sig * (1 - sig) @@ -380,7 +380,7 @@ def silu_grad(x): # noqa: ANN001, ANN202 slope = float(silu_grad(threshold)) const = float(silu(threshold)) - def fn(x): # noqa: ANN001, ANN202 + def fn(x): xp = array_api_compat.array_namespace(x) return xp.where( x < threshold, @@ -391,7 +391,7 @@ def fn(x): # noqa: ANN001, ANN202 return fn elif activation_function.lower() in ("none", "linear"): - def fn(x): # noqa: ANN001, ANN202 + def fn(x): return x return fn @@ -535,7 +535,7 @@ def __getitem__(self, key: str) -> Any: def dim_out(self) -> int: return self.w.shape[0] - def call(self, x): # noqa: ANN001, ANN201 + def call(self, x): """Forward pass. Parameters @@ -552,11 +552,11 @@ def call(self, x): # noqa: ANN001, ANN201 return y @staticmethod - def layer_norm_numpy( # noqa: ANN205 - x, # noqa: ANN001 + def layer_norm_numpy( + x, shape: tuple[int, ...], - weight=None, # noqa: ANN001 - bias=None, # noqa: ANN001 + weight=None, + bias=None, eps: float = 1e-5, ): xp = array_api_compat.array_namespace(x) @@ -633,7 +633,7 @@ def check_shape_consistency(self) -> None: f"output {self.layers[ii].dim_out}", ) - def call(self, x): # noqa: ANN001, ANN202 + def call(self, x): """Forward pass. Parameters @@ -650,7 +650,7 @@ def call(self, x): # noqa: ANN001, ANN202 x = layer(x) return x - def call_until_last(self, x): # noqa: ANN001, ANN202 + def call_until_last(self, x): """Return the output before last layer. Parameters @@ -1025,9 +1025,9 @@ def deserialize(cls, data: dict) -> "NetworkCollection": return cls(**data) -def aggregate( # noqa: ANN201 - data, # noqa: ANN001 - owners, # noqa: ANN001 +def aggregate( + data, + owners, average: bool = True, num_owner: Optional[int] = None, ): @@ -1065,10 +1065,10 @@ def aggregate( # noqa: ANN201 return output -def get_graph_index( # noqa: ANN201 - nlist, # noqa: ANN001 - nlist_mask, # noqa: ANN001 - a_nlist_mask, # noqa: ANN001 +def get_graph_index( + nlist, + nlist_mask, + a_nlist_mask, nall: int, use_loc_mapping: bool = True, ): diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index 4e47dbfe77..006afec901 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -7,6 +7,7 @@ Path, ) from typing import ( + Any, Optional, Union, ) @@ -80,15 +81,15 @@ def get_trainer( - config, - init_model=None, - restart_model=None, - finetune_model=None, - force_load=False, - init_frz_model=None, - shared_links=None, - finetune_links=None, -): + config: dict[str, Any], + init_model: Optional[str] = None, + restart_model: Optional[str] = None, + finetune_model: Optional[str] = None, + force_load: bool = False, + init_frz_model: Optional[str] = None, + shared_links: Optional[dict[str, Any]] = None, + finetune_links: Optional[dict[str, Any]] = None, +) -> training.Trainer: multi_task = "model_dict" in config.get("model", {}) # Initialize DDP @@ -98,8 +99,11 @@ def get_trainer( fleet.init(is_collective=True) def prepare_trainer_input_single( - model_params_single, data_dict_single, rank=0, seed=None - ): + model_params_single: dict[str, Any], + data_dict_single: dict[str, Any], + rank: int = 0, + seed: Optional[int] = None, + ) -> tuple[Any, Any, Any, Optional[Any]]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data", None) validation_systems = ( @@ -535,7 +539,7 @@ def change_bias( log.info(f"Saved model to {output_path}") -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/pd/infer/inference.py b/deepmd/pd/infer/inference.py index ae1b8e8516..1f717eed35 100644 --- a/deepmd/pd/infer/inference.py +++ b/deepmd/pd/infer/inference.py @@ -3,6 +3,9 @@ from copy import ( deepcopy, ) +from typing import ( + Optional, +) import paddle @@ -23,9 +26,9 @@ class Tester: def __init__( self, - model_ckpt, - head=None, - ): + model_ckpt: str, + head: Optional[str] = None, + ) -> None: """Construct a DeePMD tester. Args: diff --git a/deepmd/pd/loss/loss.py b/deepmd/pd/loss/loss.py index f825f9ff61..053ac98120 100644 --- a/deepmd/pd/loss/loss.py +++ b/deepmd/pd/loss/loss.py @@ -3,6 +3,9 @@ ABC, abstractmethod, ) +from typing import ( + NoReturn, +) import paddle @@ -15,11 +18,11 @@ class TaskLoss(paddle.nn.Layer, ABC, make_plugin_registry("loss")): - def __init__(self, **kwargs): + def __init__(self) -> None: """Construct loss.""" super().__init__() - def forward(self, input_dict, model, label, natoms, learning_rate): + def forward(self, input_dict: dict[str, paddle.Tensor], model: paddle.nn.Layer, label: dict[str, paddle.Tensor], natoms: int, learning_rate: float) -> NoReturn: """Return loss .""" raise NotImplementedError diff --git a/deepmd/pd/model/atomic_model/energy_atomic_model.py b/deepmd/pd/model/atomic_model/energy_atomic_model.py index 708ec9db7f..406f1f7b96 100644 --- a/deepmd/pd/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pd/model/atomic_model/energy_atomic_model.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Union, +) + from deepmd.pd.model.task.ener import ( EnergyFittingNet, InvarFitting, @@ -10,7 +14,13 @@ class DPEnergyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, + descriptor: object, + fitting: Union[EnergyFittingNet, InvarFitting], + type_map: list[str], + **kwargs: object, + ) -> None: assert isinstance(fitting, EnergyFittingNet) or isinstance( fitting, InvarFitting ) diff --git a/deepmd/pd/model/model/dp_model.py b/deepmd/pd/model/model/dp_model.py index e014be5b68..0c7cfb6c43 100644 --- a/deepmd/pd/model/model/dp_model.py +++ b/deepmd/pd/model/model/dp_model.py @@ -47,11 +47,11 @@ def update_sel( ) return local_jdata_cpy, min_nbor_dist - def get_fitting_net(self): + def get_fitting_net(self) -> object: """Get the fitting network.""" return self.atomic_model.fitting_net - def get_descriptor(self): + def get_descriptor(self) -> object: """Get the descriptor.""" return self.atomic_model.descriptor diff --git a/deepmd/pd/model/task/ener.py b/deepmd/pd/model/task/ener.py index 738990b2d8..f4b8655454 100644 --- a/deepmd/pd/model/task/ener.py +++ b/deepmd/pd/model/task/ener.py @@ -48,8 +48,8 @@ def __init__( mixed_types: bool = True, seed: Optional[Union[int, list[int]]] = None, type_map: Optional[list[str]] = None, - **kwargs, - ): + **kwargs: object, + ) -> None: super().__init__( "energy", ntypes, diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 4e5fea081f..86f4dc535a 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -11,6 +11,8 @@ ) from typing import ( Any, + Optional, + Union, ) import numpy as np @@ -86,16 +88,16 @@ class Trainer: def __init__( self, config: dict[str, Any], - training_data, - stat_file_path=None, - validation_data=None, - init_model=None, - restart_model=None, - finetune_model=None, - force_load=False, - shared_links=None, - finetune_links=None, - init_frz_model=None, + training_data: Any, + stat_file_path: Optional[Union[str, Path]] = None, + validation_data: Optional[Any] = None, + init_model: Optional[str] = None, + restart_model: Optional[str] = None, + finetune_model: Optional[str] = None, + force_load: bool = False, + shared_links: Optional[dict[str, Any]] = None, + finetune_links: Optional[dict[str, Any]] = None, + init_frz_model: Optional[str] = None, ) -> None: """Construct a DeePMD trainer. @@ -148,7 +150,7 @@ def __init__( ) self.lcurve_should_print_header = True - def get_opt_param(params): + def get_opt_param(params: dict[str, Any]) -> tuple[str, dict[str, Any]]: opt_type = params.get("opt_type", "Adam") opt_param = { "kf_blocksize": params.get("kf_blocksize", 5120), @@ -159,8 +161,12 @@ def get_opt_param(params): } return opt_type, opt_param - def get_data_loader(_training_data, _validation_data, _training_params): - def get_dataloader_and_buffer(_data, _params): + def get_data_loader( + _training_data: Any, _validation_data: Any, _training_params: dict[str, Any] + ) -> tuple[Any, Any, Any, Any]: + def get_dataloader_and_buffer( + _data: Any, _params: dict[str, Any] + ) -> tuple[Any, Any]: _sampler = get_sampler_from_params(_data, _params) if _sampler is None: log.warning( @@ -207,14 +213,14 @@ def get_dataloader_and_buffer(_data, _params): ) def single_model_stat( - _model, - _data_stat_nbatch, - _training_data, - _validation_data, - _stat_file_path, - _data_requirement, - finetune_has_new_type=False, - ): + _model: Any, + _data_stat_nbatch: int, + _training_data: Any, + _validation_data: Optional[Any], + _stat_file_path: Optional[Union[str, Path]], + _data_requirement: list[DataRequirementItem], + finetune_has_new_type: bool = False, + ) -> Any: _data_requirement += get_additional_data_requirement(_model) _training_data.add_data_requirement(_data_requirement) if _validation_data is not None: @@ -1057,7 +1063,7 @@ def log_loss_valid(_task_key="Default"): "files, which can be viewd in NVIDIA Nsight Systems software" ) - def save_model(self, save_path, lr=0.0, step=0) -> None: + def save_model(self, save_path: str, lr: float = 0.0, step: int = 0) -> None: module = ( self.wrapper._layers if dist.is_available() and dist.is_initialized() @@ -1079,7 +1085,9 @@ def save_model(self, save_path, lr=0.0, step=0) -> None: checkpoint_files.sort(key=lambda x: x.stat().st_mtime) checkpoint_files[0].unlink() - def get_data(self, is_train=True, task_key="Default"): + def get_data( + self, is_train: bool = True, task_key: str = "Default" + ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: if not self.multi_task: if is_train: try: @@ -1155,7 +1163,9 @@ 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) -> None: + def print_header( + self, fout: Any, train_results: dict[str, Any], valid_results: dict[str, Any] + ) -> None: train_keys = sorted(train_results.keys()) print_str = "" print_str += "# {:5s}".format("step") @@ -1187,7 +1197,12 @@ def print_header(self, fout, train_results, valid_results) -> None: fout.flush() def print_on_training( - self, fout, step_id, cur_lr, train_results, valid_results + self, + fout: Any, + step_id: int, + cur_lr: float, + train_results: dict[str, Any], + valid_results: dict[str, Any], ) -> None: train_keys = sorted(train_results.keys()) print_str = "" @@ -1219,7 +1234,7 @@ def print_on_training( fout.flush() -def get_additional_data_requirement(_model): +def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: additional_data_requirement = [] if _model.get_dim_fparam() > 0: fparam_requirement_items = [ @@ -1246,12 +1261,14 @@ def get_additional_data_requirement(_model): return additional_data_requirement -def whether_hessian(loss_params): +def whether_hessian(loss_params: dict[str, Any]) -> bool: loss_type = loss_params.get("type", "ener") return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0 -def get_loss(loss_params, start_lr, _ntypes, _model): +def get_loss( + loss_params: dict[str, Any], start_lr: float, _ntypes: int, _model: Any +) -> TaskLoss: loss_type = loss_params.get("type", "ener") if whether_hessian(loss_params): loss_params["starter_learning_rate"] = start_lr @@ -1265,17 +1282,17 @@ def get_loss(loss_params, start_lr, _ntypes, _model): def get_single_model( - _model_params, -): + _model_params: dict[str, Any], +) -> Any: model = get_model(deepcopy(_model_params)).to(DEVICE) return model def get_model_for_wrapper( - _model_params, - resuming=False, - _loss_params=None, -): + _model_params: dict[str, Any], + resuming: bool = False, + _loss_params: Optional[dict[str, Any]] = None, +) -> Any: if "model_dict" not in _model_params: if _loss_params is not None and whether_hessian(_loss_params): _model_params["hessian_mode"] = True @@ -1298,7 +1315,7 @@ def get_model_for_wrapper( return _model -def get_case_embd_config(_model_params): +def get_case_embd_config(_model_params: dict[str, Any]) -> tuple[bool, dict[str, Any]]: assert "model_dict" in _model_params, ( "Only support setting case embedding for multi-task model!" ) @@ -1323,10 +1340,10 @@ def get_case_embd_config(_model_params): def model_change_out_bias( - _model, - _sample_func, - _bias_adjust_mode="change-by-statistic", -): + _model: Any, + _sample_func: Any, + _bias_adjust_mode: str = "change-by-statistic", +) -> None: old_bias = deepcopy(_model.get_out_bias()) _model.change_out_bias( _sample_func, diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index bd28b17c88..c3d5bd0495 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -24,8 +24,8 @@ def __init__( self, model: paddle.nn.Layer | dict, loss: paddle.nn.Layer | dict = None, - model_params=None, - shared_links=None, + model_params: dict[str, Any] | None = None, + shared_links: dict[str, Any] | None = None, ) -> None: """Construct a DeePMD model wrapper. @@ -64,7 +64,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) -> None: + def share_params(self, shared_links: dict[str, Any], resume: bool = False) -> None: """ Share the parameters of classes following rules defined in shared_links during multitask training. If not start from checkpoint (resume is False), @@ -137,18 +137,18 @@ def share_params(self, shared_links, resume=False) -> None: def forward( self, - coord, - atype, + coord: paddle.Tensor, + atype: paddle.Tensor, spin: paddle.Tensor | None = None, box: paddle.Tensor | None = None, cur_lr: paddle.Tensor | None = None, label: paddle.Tensor | None = None, task_key: paddle.Tensor | None = None, - inference_only=False, - do_atomic_virial=False, + inference_only: bool = False, + do_atomic_virial: bool = False, fparam: paddle.Tensor | None = None, aparam: paddle.Tensor | None = None, - ): + ) -> dict[str, paddle.Tensor]: if not self.multi_task: task_key = "Default" else: @@ -196,13 +196,13 @@ def set_state_dict( ) -> tuple[list[str], list[str]]: return self.load_state_dict(state_dict) - def state_dict(self): + def state_dict(self) -> dict[str, Any]: state_dict = super().state_dict() extra_state = self.get_extra_state() state_dict.update({"_extra_state": extra_state}) return state_dict - def set_extra_state(self, extra_state: dict): + def set_extra_state(self, extra_state: dict[str, Any]) -> None: self.model_params = extra_state["model_params"] self.train_infos = extra_state["train_infos"] return None diff --git a/deepmd/pd/utils/dataloader.py b/deepmd/pd/utils/dataloader.py index 0cb8adbc63..f96ded17d1 100644 --- a/deepmd/pd/utils/dataloader.py +++ b/deepmd/pd/utils/dataloader.py @@ -12,6 +12,10 @@ from threading import ( Thread, ) +from typing import ( + Optional, + Union, +) import h5py import numpy as np @@ -53,7 +57,7 @@ # paddle.multiprocessing.set_sharing_strategy("file_system") -def setup_seed(seed): +def setup_seed(seed: Union[int, list, tuple]) -> None: if isinstance(seed, (list, tuple)): mixed_seed = mix_entropy(seed) else: @@ -82,12 +86,12 @@ class DpLoaderSet(Dataset): def __init__( self, - systems, - batch_size, - type_map, - seed=None, - shuffle=True, - ): + systems: Union[str, list[str]], + batch_size: int, + type_map: list[str], + seed: Optional[int] = None, + shuffle: bool = True, + ) -> None: if seed is not None: setup_seed(seed) if isinstance(systems, str): diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index 1f0533d8fc..685dc1b23e 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -2,6 +2,7 @@ from typing import ( + Any, Optional, ) @@ -16,7 +17,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,16 +32,16 @@ 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): + def __getitem__(self, index: int) -> dict[str, Any]: """Get a frame from the selected system.""" b_data = self._data_system.get_item_paddle(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/pd/utils/env.py b/deepmd/pd/utils/env.py index 28606d0945..94a5a3ccb7 100644 --- a/deepmd/pd/utils/env.py +++ b/deepmd/pd/utils/env.py @@ -121,7 +121,7 @@ def to_bool(flag: int | bool | str) -> bool: # os.environ['CPU_NUM'] = str(intra_nthreads) -def enable_prim(enable: bool = True): +def enable_prim(enable: bool = True) -> None: # NOTE: operators in list below will not use composite # operator but kernel instead for better performance EAGER_COMP_OP_BLACK_LIST = [ diff --git a/deepmd/pd/utils/exclude_mask.py b/deepmd/pd/utils/exclude_mask.py index cde8730c9a..e22379ea3e 100644 --- a/deepmd/pd/utils/exclude_mask.py +++ b/deepmd/pd/utils/exclude_mask.py @@ -32,10 +32,10 @@ def reinit( ) self.type_mask = to_paddle_tensor(self.type_mask).reshape([-1]) - def get_exclude_types(self): + def get_exclude_types(self) -> list[int]: return self.exclude_types - def get_type_mask(self): + def get_type_mask(self) -> paddle.Tensor: return self.type_mask def forward( @@ -98,7 +98,7 @@ def reinit( self.type_mask = to_paddle_tensor(self.type_mask).reshape([-1]) self.no_exclusion = len(self._exclude_types) == 0 - def get_exclude_types(self): + def get_exclude_types(self) -> set[tuple[int, int]]: return self._exclude_types # may have a better place for this method... diff --git a/deepmd/pd/utils/preprocess.py b/deepmd/pd/utils/preprocess.py index 3be42b522e..ba10c6848c 100644 --- a/deepmd/pd/utils/preprocess.py +++ b/deepmd/pd/utils/preprocess.py @@ -6,7 +6,7 @@ log = logging.getLogger(__name__) -def compute_smooth_weight(distance, rmin: float, rmax: float): +def compute_smooth_weight(distance: paddle.Tensor, rmin: float, rmax: float) -> paddle.Tensor: """Compute smooth weight for descriptor elements.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") @@ -17,7 +17,7 @@ def compute_smooth_weight(distance, rmin: float, rmax: float): return vv -def compute_exp_sw(distance, rmin: float, rmax: float): +def compute_exp_sw(distance: paddle.Tensor, rmin: float, rmax: float) -> paddle.Tensor: """Compute the exponential switch function for neighbor update.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") diff --git a/deepmd/pd/utils/region.py b/deepmd/pd/utils/region.py index d2600ef16e..237fa84b26 100644 --- a/deepmd/pd/utils/region.py +++ b/deepmd/pd/utils/region.py @@ -75,7 +75,7 @@ def to_face_distance( return dist.reshape(list(cshape[:-2]) + [3]) # noqa:RUF005 -def b_to_face_distance(cell): +def b_to_face_distance(cell: paddle.Tensor) -> paddle.Tensor: volume = paddle.linalg.det(cell) c_yz = paddle.cross(cell[:, 1], cell[:, 2], axis=-1) _h2yz = volume / paddle.linalg.norm(c_yz, axis=-1) diff --git a/deepmd/pd/utils/spin.py b/deepmd/pd/utils/spin.py index 27bc355877..83fa01a8d0 100644 --- a/deepmd/pd/utils/spin.py +++ b/deepmd/pd/utils/spin.py @@ -4,10 +4,10 @@ def concat_switch_virtual( - extended_tensor, - extended_tensor_virtual, + extended_tensor: paddle.Tensor, + extended_tensor_virtual: paddle.Tensor, nloc: int, -): +) -> paddle.Tensor: """ Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms. - [:, :nloc]: original nloc real atoms. diff --git a/deepmd/pd/utils/utils.py b/deepmd/pd/utils/utils.py index 175ac5019b..166cd500f8 100644 --- a/deepmd/pd/utils/utils.py +++ b/deepmd/pd/utils/utils.py @@ -83,7 +83,7 @@ def silut_double_backward( class SiLUTScript(paddle.nn.Layer): - def __init__(self, threshold: float = 3.0): + def __init__(self, threshold: float = 3.0) -> None: super().__init__() self.threshold = threshold @@ -95,7 +95,7 @@ def __init__(self, threshold: float = 3.0): self.const_val = float(threshold * sigmoid_threshold) self.get_script_code() - def get_script_code(self): + def get_script_code(self) -> None: silut_forward_script = paddle.jit.to_static(silut_forward, full_graph=True) silut_backward_script = paddle.jit.to_static(silut_backward, full_graph=True) silut_double_backward_script = paddle.jit.to_static( @@ -104,7 +104,13 @@ def get_script_code(self): class SiLUTFunction(paddle.autograd.PyLayer): @staticmethod - def forward(ctx, x, threshold, slope, const_val): + def forward( + ctx: paddle.autograd.PyLayerContext, + x: paddle.Tensor, + threshold: float, + slope: float, + const_val: float, + ) -> paddle.Tensor: ctx.save_for_backward(x) ctx.threshold = threshold ctx.slope = slope @@ -112,7 +118,9 @@ def forward(ctx, x, threshold, slope, const_val): return silut_forward_script(x, threshold, slope, const_val) @staticmethod - def backward(ctx, grad_output): + def backward( + ctx: paddle.autograd.PyLayerContext, grad_output: paddle.Tensor + ) -> paddle.Tensor: (x,) = ctx.saved_tensor() threshold = ctx.threshold slope = ctx.slope @@ -122,7 +130,13 @@ def backward(ctx, grad_output): class SiLUTGradFunction(paddle.autograd.PyLayer): @staticmethod - def forward(ctx, x, grad_output, threshold, slope): + def forward( + ctx: paddle.autograd.PyLayerContext, + x: paddle.Tensor, + grad_output: paddle.Tensor, + threshold: float, + slope: float, + ) -> paddle.Tensor: ctx.threshold = threshold ctx.slope = slope grad_input = silut_backward_script(x, grad_output, threshold, slope) @@ -142,21 +156,21 @@ def backward(ctx, grad_grad_output): self.SiLUTFunction = SiLUTFunction - def forward(self, x): + def forward(self, x: paddle.Tensor) -> paddle.Tensor: return self.SiLUTFunction.apply(x, self.threshold, self.slope, self.const_val) class SiLUT(paddle.nn.Layer): - def __init__(self, threshold=3.0): + def __init__(self, threshold: float = 3.0) -> None: super().__init__() - def sigmoid(x): + def sigmoid(x: paddle.Tensor) -> paddle.Tensor: return F.sigmoid(x) - def silu(x): + def silu(x: paddle.Tensor) -> paddle.Tensor: return F.silu(x) - def silu_grad(x): + def silu_grad(x: paddle.Tensor) -> paddle.Tensor: sig = sigmoid(x) return sig + x * sig * (1 - sig) @@ -281,7 +295,9 @@ def to_paddle_tensor( return paddle.to_tensor(xx, dtype=prec, place=DEVICE) -def dict_to_device(sample_dict): +def dict_to_device( + sample_dict: dict[str, paddle.Tensor | list[paddle.Tensor] | None], +) -> 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/progress.tmp b/progress.tmp new file mode 100644 index 0000000000..c634e3046e --- /dev/null +++ b/progress.tmp @@ -0,0 +1 @@ +# Progress update diff --git a/pyproject.toml b/pyproject.toml index cb11d0258d..81ffc60aec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -427,7 +427,48 @@ runtime-evaluated-base-classes = ["torch.nn.Module"] "deepmd/pt/**" = ["TID253"] "deepmd/jax/**" = ["TID253"] "deepmd/pd/**" = ["TID253", "ANN"] - +# Paddle backend: Gradually enabling ANN rule +# Completed files with full type annotations: +"deepmd/pd/entrypoints/main.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/entrypoints/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/train/wrapper.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/train/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/learning_rate.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/dp_random.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/update_sel.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/preprocess.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/spin.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/dataset.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/decomp.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/env.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/auto_batch_size.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/loss/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/loss/loss.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/network/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/task/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/task/task.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/task/base_fitting.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/descriptor/base_descriptor.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/infer/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/infer/inference.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/atomic_model/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/atomic_model/energy_atomic_model.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/cxx_op.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/descriptor/__init__.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/task/ener.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/model/model/dp_model.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/serialization.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/region.py" = ["TID253"] # ✅ Fully typed +"deepmd/pd/utils/exclude_mask.py" = ["TID253"] # ✅ Fully typed +# TODO: Complete type hints and remove ANN exclusion for remaining files: +"deepmd/pd/train/**" = ["TID253", "ANN"] # 🚧 Partial progress - training.py still needs work +"deepmd/pd/utils/**" = ["TID253", "ANN"] # 🚧 Partial progress - utils.py partially done +"deepmd/pd/loss/**" = ["TID253", "ANN"] # 🚧 Partial progress - other loss files need work +"deepmd/pd/model/**" = ["TID253", "ANN"] # 🚧 Partial progress - some files completed +"deepmd/pd/infer/**" = ["TID253", "ANN"] # 🚧 Partial progress - inference.py completed "source/**" = ["ANN"] "source/tests/tf/**" = ["TID253", "ANN"] "source/tests/pt/**" = ["TID253", "ANN"] diff --git a/pyproject.toml.backup b/pyproject.toml.backup new file mode 100644 index 0000000000..6d43fd9211 --- /dev/null +++ b/pyproject.toml.backup @@ -0,0 +1,556 @@ +[build-system] +requires = [ + # TODO: unpin the upper bound when scikit-build dynamic metadata API is stable + # dynamic metadata API is still unstable + "scikit-build-core>=0.5,<0.11,!=0.6.0", + "packaging", + 'tomli >= 1.1.0 ; python_version < "3.11"', +] +build-backend = "backend.dp_backend" +backend-path = ["."] + +[project] +name = "deepmd-kit" +dynamic = ["version", "optional-dependencies", "scripts", "readme"] +description = "A deep learning package for many-body potential energy representation and molecular dynamics" +authors = [ + {name = "DeepModeling"}, + {name = "Han Wang", email = "wang_han@iapcm.ac.cn"}, +] +license = {file = "LICENSE"} +classifiers = [ + "Natural Language :: English", + "Operating System :: POSIX :: Linux", + "Operating System :: Microsoft :: Windows", + "Development Status :: 5 - Production/Stable", + "Programming Language :: C", + "Programming Language :: C++", + "Programming Language :: Python :: 3 :: Only", + "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.2", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Scientific/Engineering :: Physics", + "Topic :: Scientific/Engineering :: Chemistry", + "Environment :: Console", +] +dependencies = [ + # array-api-compat requires numpy>=1.21 + 'numpy>=1.21', + 'scipy', + 'pyyaml', + 'dargs >= 0.4.7', + 'typing_extensions; python_version < "3.8"', + 'importlib_metadata>=1.4; python_version < "3.8"', + 'h5py', + "h5py>=3.6.0,!=3.11.0; platform_system=='Linux' and platform_machine=='aarch64'", + 'wcmatch', + 'packaging', + 'ml_dtypes', + 'mendeleev', + 'array-api-compat', +] +requires-python = ">=3.9" +keywords = ["deepmd"] + +[project.entry-points."lammps.plugins"] +deepmd = "deepmd.lmp_check_build:get_op_dir" + +[project.entry-points."dpgui"] +"DeePMD-kit" = "deepmd.utils.argcheck:gen_args" +"DeePMD-kit Multi-task" = "deepmd.utils.argcheck:gen_args_multi_task" + +[project.entry-points."dpdata.plugins"] +deepmd_driver = "deepmd.driver:DPDriver" + +[project.urls] +Homepage = "https://github.com/deepmodeling/deepmd-kit" +documentation = "https://docs.deepmodeling.com/projects/deepmd" +repository = "https://github.com/deepmodeling/deepmd-kit" + +# Metadata below is dynamic. However, it still has static parts, +# which can be read by the build backend. +[tool.deepmd_build_backend.optional-dependencies] +test = [ + "dpdata>=0.2.7", + # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843 + # fixed in 3.23.0 + "ase>=3.23.0", + "pytest", + "pytest-cov", + "pytest-sugar", + "pytest-split", + "dpgui", + 'array-api-strict>=2,!=2.1.1;python_version>="3.9"', +] +docs = [ + "sphinx>=3.1.1", + "sphinx-book-theme", + "myst-nb>=1.0.0", + "myst-parser>=0.19.2", + "sphinx-design", + "breathe", + "exhale>=0.3.7", + "numpydoc", + "ase", + "deepmodeling-sphinx>=0.3.0", + "dargs>=0.3.4", + "sphinx-argparse<0.5.0", + "pygments-lammps", + "sphinxcontrib-bibtex", + "sphinx-autoapi>=3.0.0", + "sphinxcontrib-programoutput", + "sphinxcontrib-moderncmakedomain", + "sphinx-remove-toctrees", +] +lmp = [ + "lammps[mpi]~=2025.7.22.1.0", +] +ipi = [ + "ipi", +] +gui = [ + "dpgui", +] +cu11 = [ + "nvidia-cuda-runtime-cu11", + "nvidia-cublas-cu11", + "nvidia-cufft-cu11", + "nvidia-curand-cu11", + "nvidia-cusolver-cu11", + "nvidia-cusparse-cu11", + "nvidia-cudnn-cu11<9", + "nvidia-cuda-nvcc-cu11", +] +cu12 = [ + "nvidia-cuda-runtime-cu12", + "nvidia-cublas-cu12", + "nvidia-cufft-cu12", + "nvidia-curand-cu12", + "nvidia-cusolver-cu12", + "nvidia-cusparse-cu12", + "nvidia-cudnn-cu12", + "nvidia-cuda-nvcc-cu12", +] +jax = [ + # below is a funny workaround for + # https://github.com/astral-sh/uv/issues/8601 + 'jax>=0.4.33;python_version>="3.10"', + 'jax>=0.4.33;python_version>="3.10"', + 'flax>=0.10.0;python_version>="3.10"', + 'flax>=0.10.0;python_version>="3.10"', + 'orbax-checkpoint;python_version>="3.10"', + 'orbax-checkpoint;python_version>="3.10"', + # The pinning of ml_dtypes may conflict with TF + # 'jax-ai-stack;python_version>="3.10"', +] + +[tool.deepmd_build_backend.scripts] +dp = "deepmd.main:main" + +[dependency-groups] +dev = [ + "pre-commit", + "cmake", + "mpich", +] + +[tool.setuptools_scm] + +[tool.scikit-build] +experimental = true +minimum-version = "0.5" +cmake.source-dir = "source" +sdist.include = [ + "/deepmd/_version.py", +] +sdist.exclude = [ + "/source/tests", + "/source/api_c/tests", + "/source/api_cc/tests", + "/source/lib/tests", + "/source/lmp/tests", + "/doc", + "/examples", + "/data", + "/.github", +] +wheel.packages = [ + "deepmd", +] +wheel.py-api = "py37" +build-dir = "build/{wheel_tag}" + +[tool.scikit-build.metadata.version] +provider = "scikit_build_core.metadata.setuptools_scm" + +[tool.scikit-build.metadata.optional-dependencies] +provider = "backend.dynamic_metadata" +provider-path = "backend" + +[tool.scikit-build.metadata.scripts] +provider = "backend.dynamic_metadata" +provider-path = "backend" + +[tool.scikit-build.metadata.readme] +provider = "scikit_build_core.metadata.fancy_pypi_readme" + +[[tool.scikit-build.generate]] +path = "deepmd/_version.py" +template = ''' +version = "${version}" +''' + +[tool.hatch.metadata.hooks.fancy-pypi-readme] +content-type = "text/markdown" + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.fragments]] +path = "README.md" + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +# links +pattern = '\[(.+?)\]\(((?!https?://)\S+?)\)' +replacement = '[\1](https://github.com/deepmodeling/deepmd-kit/tree/master/\g<2>)' + +[[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]] +# image +pattern = '(srcset|src)="((?!https?://)\S+?)"' +replacement = '\1="https://github.com/deepmodeling/deepmd-kit/raw/master/\g<2>"' + +[tool.cibuildwheel] +test-command = [ + "python -m deepmd -h", + """python -c "import deepmd.tf;import deepmd.pt;import deepmd.pd" """, + "dp -h", + "dp_ipi", + "pytest {project}/source/tests/common/test_lammps.py" +] +test-extras = ["cpu", "test", "lmp", "ipi", "torch", "paddle"] +build = ["cp311-*"] +skip = ["*-win32", "*-manylinux_i686", "*-musllinux*"] +# TODO: uncomment to use the latest image when CUDA 11 is deprecated +# manylinux-x86_64-image = "manylinux_2_28" +manylinux-x86_64-image = "quay.io/pypa/manylinux_2_28_x86_64:2022-11-19-1b19e81" +manylinux-aarch64-image = "manylinux_2_28" + +[tool.cibuildwheel.macos] +repair-wheel-command = """delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies""" + +[tool.cibuildwheel.macos.environment] +PIP_PREFER_BINARY = "1" +DP_LAMMPS_VERSION = "stable_22Jul2025_update1" +DP_ENABLE_IPI = "1" +DP_ENABLE_PYTORCH = "1" +DP_ENABLE_PADDLE = "1" +# for unclear reason, when enabling PyTorch, OpenMP is found accidentally +CMAKE_ARGS = "-DCMAKE_DISABLE_FIND_PACKAGE_OpenMP=1" + +[[tool.cibuildwheel.overrides]] +# error: 'value' is unavailable: introduced in macOS 10.13 +select = "*-macosx_x86_64" +inherit.environment = "append" +environment.MACOSX_DEPLOYMENT_TARGET = "11.0" + +[tool.cibuildwheel.linux] +repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so --exclude _pywrap_tensorflow_internal.so --exclude libtensorflow_cc.so.2 --exclude libc10.so --exclude libtorch.so --exclude libtorch_cpu.so --exclude libmpi.so.12 -w {dest_dir} {wheel}" +environment-pass = [ + "CIBW_BUILD", + "DP_VARIANT", + "CUDA_VERSION", + "DP_PKG_NAME", + "SETUPTOOLS_SCM_PRETEND_VERSION", +] +before-all = [ + """if [ ! -z "${DP_PKG_NAME}" ]; then sed -i "s/name = \\"deepmd-kit\\"/name = \\"${DP_PKG_NAME}\\"/g" pyproject.toml; fi""", + # https://almalinux.org/blog/2023-12-20-almalinux-8-key-update/ + """rpm --import https://repo.almalinux.org/almalinux/RPM-GPG-KEY-AlmaLinux""", + """{ if [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""", + # uv is not available in the old manylinux image + """{ if [ "$(uname -m)" = "x86_64" ] ; then pipx install uv; fi }""", +] +before-build = [ + # old build doesn't support uv + """{ if [ "$(uname -m)" = "x86_64" ] ; then uv pip install --system -U build; fi }""", +] +[tool.cibuildwheel.linux.environment] +PIP_PREFER_BINARY = "1" +DP_LAMMPS_VERSION = "stable_22Jul2025_update1" +DP_ENABLE_IPI = "1" +DP_ENABLE_PYTORCH = "1" +DP_ENABLE_PADDLE = "1" +# use CPU version of torch for building, which should also work for GPU +# note: uv has different behavior from pip on extra index url +# https://github.com/astral-sh/uv/blob/main/PIP_COMPATIBILITY.md#packages-that-exist-on-multiple-indexes +UV_EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cpu" + +[tool.cibuildwheel.windows] +test-extras = ["cpu", "torch", "paddle"] +test-command = [ + "python -m deepmd -h", + "dp -h", +] +[tool.cibuildwheel.windows.environment] +PIP_PREFER_BINARY = "1" +DP_ENABLE_PYTORCH = "1" +DP_ENABLE_PADDLE = "1" + +# One can run `tox` or `tox -e gpu` +# to run pytest in an isolated environment +# Use with pipx: +# $ pip install -U pipx +# $ pipx tox +[tool.tox] +legacy_tox_ini = """ + [tox] + min_version = 4.0 + + [testenv] + extras = + test + cpu + commands = pytest source/tests + + [testenv:gpu] + extras = + test + gpu + commands = pytest source/tests + setenv = + DP_VARIANT = cuda +""" + +# selectively turn of lintner warnings, always include reasoning why any warning should +# be silenced + +# W504 - line break after binary operator - there is conflict between W503 and W504 in +# some lintners. One recommends line bread after and one before binary operator so we +# switch W504 off and recommend this coding style: +# a = (b + -> instead of -> a = (b +# c) + c) +[tool.autopep8] +ignore = "W504" + +# D413 - Missing blank line after last section - makes no sense only adds empty lines in +# docstrings +# D416 - Section name should end with a colon - only applicable to RST type docstrings, +# we are using numpy style +# D203 - 1 blank line required before class docstring - only adds unnecessary empty space +# D107 - Missing docstring in __init__ - Nupmy style documents __init__ parameters in +# class docstring +# D213 - Multi-line docstring summary should start at the second line - unnecessary waste +# of space, start on the first line +[tool.pydocstyle] +ignore = "D413, D416, D203, D107, D213" + +[tool.isort] +profile = "black" +force_grid_wrap = 1 + +[tool.ruff.format] +docstring-code-format = true + +[tool.ruff.lint] +select = [ + "E", # errors + "W", # warning + "F", # pyflakes + "D", # pydocstyle + "UP", # pyupgrade + "C4", # flake8-comprehensions + "RUF", # ruff + "NPY", # numpy + "TID251", # banned-api + "TID253", # banned-module-level-imports + "T20", # ban print + "B904", # raise-without-from-inside-except + "N804", # invalid-first-argument-name-for-class-method + "N805", # invalid-first-argument-name-for-method + "DTZ", # datetime + "TCH", # flake8-type-checking + "PYI", # flake8-pyi + "ANN", # type annotations +] + +ignore = [ + "ANN401", # Allow Any due to too many violations + "E501", # line too long + "F841", # local variable is assigned to but never used + "RUF059", # unused-unpacked-variable + "E741", # ambiguous variable name + "E402", # module level import not at top of file + "D100", # TODO: missing docstring in public module + "D101", # TODO: missing docstring in public class + "D102", # TODO: missing docstring in public method + "D103", # TODO: missing docstring in public function + "D104", # TODO: missing docstring in public package + "D105", # TODO: missing docstring in magic method + "D205", # 1 blank line required between summary line and description + "D401", # TODO: first line should be in imperative mood + "D404", # TODO: first word of the docstring should not be This +] + +exclude = [ + "source/3rdparty/**", +] + +[tool.ruff.lint.pydocstyle] +convention = "numpy" + +[tool.ruff.lint.flake8-tidy-imports] +banned-module-level-imports = [ + "deepmd.tf", + "deepmd.pt", + "deepmd.pd", + "deepmd.jax", + "tensorflow", + "torch", + "jax", + "paddle", +] + +[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. +"source/3rdparty/**" = ["ALL"] +"backend/**" = ["ANN"] +"data/**" = ["ANN"] +"deepmd/tf/**" = ["TID253", "ANN"] +"deepmd/pt/**" = ["TID253"] +"deepmd/jax/**" = ["TID253"] +# Temporarily disabled to check violations +# "deepmd/pd/**" = ["TID253", "ANN"] +# Paddle backend: Gradually enabling ANN rule +# Completed files with full type annotations: +"deepmd/pd/entrypoints/main.py" = ["TID253"] # ✅ Fully typed +# TODO: Complete type hints and remove ANN exclusion for remaining files: +"deepmd/pd/train/**" = ["TID253", "ANN"] # 🚧 Partial progress +"deepmd/pd/utils/**" = ["TID253", "ANN"] # 🚧 Partial progress +"deepmd/pd/loss/**" = ["TID253", "ANN"] # ❌ Not started +"deepmd/pd/model/**" = ["TID253", "ANN"] # ❌ Not started +"deepmd/pd/infer/**" = ["TID253", "ANN"] # ❌ Not started +"deepmd/pd/cxx_op.py" = ["ANN"] # ❌ Not started +"deepmd/dpmodel/**" = ["ANN"] +"source/**" = ["ANN"] +"source/tests/tf/**" = ["TID253", "ANN"] +"source/tests/pt/**" = ["TID253", "ANN"] +"source/tests/jax/**" = ["TID253", "ANN"] +"source/tests/pd/**" = ["TID253", "ANN"] +"source/tests/universal/pt/**" = ["TID253", "ANN"] +"source/tests/universal/pd/**" = ["TID253", "ANN"] +"source/tests/**" = ["ANN"] +"source/jax2tf_tests/**" = ["TID253", "ANN"] +"source/ipi/tests/**" = ["TID253", "ANN"] +"source/lmp/tests/**" = ["TID253", "ANN"] +"**/tests/**/test_*.py" = ["ANN"] +"**/tests/**/*_test.py" = ["ANN"] +"**/*.ipynb" = ["T20"] # printing in a nb file is expected + +[tool.pytest.ini_options] +markers = "run" + +[tool.coverage.run] +plugins = ["source.3rdparty.coverage_plugins.jit_plugin"] + +[tool.pylint.'MESSAGES CONTROL'] +load-plugins = "deepmd_checker" +disable = "all" +enable = "E8001,E8002" + +[tool.flake8] +select = [ + "TOR0", + "TOR1", + "TOR2", +] + +[[tool.uv.dependency-metadata]] +# Fix https://github.com/deepmodeling/deepmd-kit/issues/4679 +name = "tensorflow" +version = "2.19.0" +requires-dist = [ + 'absl-py >=1.0.0', + 'astunparse >=1.6.0', + 'flatbuffers >=24.3.25', + 'gast !=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1', + 'google-pasta >=0.1.1', + 'libclang >=13.0.0', + 'opt-einsum >=2.3.2', + 'packaging', + 'protobuf !=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<6.0.0dev,>=3.20.3', + 'requests <3,>=2.21.0', + 'setuptools', + 'six >=1.12.0', + 'termcolor >=1.1.0', + 'typing-extensions >=3.6.6', + 'wrapt >=1.11.0', + 'grpcio <2.0,>=1.24.3', + 'tensorboard ~=2.19.0', + 'keras >=3.5.0', + 'numpy <2.2.0,>=1.26.0', + 'h5py >=3.11.0', + 'ml-dtypes <1.0.0,>=0.5.1', + # 'tensorflow-intel ==2.19.0 ; platform_system == "Windows"', + 'tensorflow-io-gcs-filesystem >=0.23.1 ; python_version < "3.12"', + 'nvidia-cublas-cu12 ==12.5.3.2 ; extra == "and-cuda"', + 'nvidia-cuda-cupti-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-nvcc-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-nvrtc-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-runtime-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cudnn-cu12 ==9.3.0.75 ; extra == "and-cuda"', + 'nvidia-cufft-cu12 ==11.2.3.61 ; extra == "and-cuda"', + 'nvidia-curand-cu12 ==10.3.6.82 ; extra == "and-cuda"', + 'nvidia-cusolver-cu12 ==11.6.3.83 ; extra == "and-cuda"', + 'nvidia-cusparse-cu12 ==12.5.1.3 ; extra == "and-cuda"', + 'nvidia-nccl-cu12 ==2.23.4 ; extra == "and-cuda"', + 'nvidia-nvjitlink-cu12 ==12.5.82 ; extra == "and-cuda"', +] + +[[tool.uv.dependency-metadata]] +name = "tensorflow-cpu" +version = "2.19.0" +requires-dist = [ + 'absl-py >=1.0.0', + 'astunparse >=1.6.0', + 'flatbuffers >=24.3.25', + 'gast !=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1', + 'google-pasta >=0.1.1', + 'libclang >=13.0.0', + 'opt-einsum >=2.3.2', + 'packaging', + 'protobuf !=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<6.0.0dev,>=3.20.3', + 'requests <3,>=2.21.0', + 'setuptools', + 'six >=1.12.0', + 'termcolor >=1.1.0', + 'typing-extensions >=3.6.6', + 'wrapt >=1.11.0', + 'grpcio <2.0,>=1.24.3', + 'tensorboard ~=2.19.0', + 'keras >=3.5.0', + 'numpy <2.2.0,>=1.26.0', + 'h5py >=3.11.0', + 'ml-dtypes <1.0.0,>=0.5.1', + # 'tensorflow-intel ==2.19.0 ; platform_system == "Windows"', + 'tensorflow-io-gcs-filesystem >=0.23.1 ; python_version < "3.12"', + 'nvidia-cublas-cu12 ==12.5.3.2 ; extra == "and-cuda"', + 'nvidia-cuda-cupti-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-nvcc-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-nvrtc-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cuda-runtime-cu12 ==12.5.82 ; extra == "and-cuda"', + 'nvidia-cudnn-cu12 ==9.3.0.75 ; extra == "and-cuda"', + 'nvidia-cufft-cu12 ==11.2.3.61 ; extra == "and-cuda"', + 'nvidia-curand-cu12 ==10.3.6.82 ; extra == "and-cuda"', + 'nvidia-cusolver-cu12 ==11.6.3.83 ; extra == "and-cuda"', + 'nvidia-cusparse-cu12 ==12.5.1.3 ; extra == "and-cuda"', + 'nvidia-nccl-cu12 ==2.23.4 ; extra == "and-cuda"', + 'nvidia-nvjitlink-cu12 ==12.5.82 ; extra == "and-cuda"', +]