Merge branch 'master' into devel_pt_dplr

Signed-off-by: Jia-Xin Zhu <jiaxinzhu@stu.xmu.edu.cn>

Author: ChiahsinChu

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.14.13
+    rev: v0.14.14
     hooks:
       - id: ruff
         args: ["--fix"]

deepmd/dpmodel/modifier/base_modifier.py

@@ -8,6 +8,10 @@
     Any,
 )
 
+from typing_extensions import (
+    Self,
+)
+
 from deepmd.utils.plugin import (
     PluginVariant,
     make_plugin_registry,
@@ -18,7 +22,7 @@ def make_base_modifier() -> type[object]:
     class BaseModifier(ABC, PluginVariant, make_plugin_registry("modifier")):
         """Base class for data modifier."""
 
-        def __new__(cls, *args: Any, **kwargs: Any) -> "BaseModifier":
+        def __new__(cls, *args: Any, **kwargs: Any) -> Self:
             if cls is BaseModifier:
                 cls = cls.get_class_by_type(kwargs["type"])
             return super().__new__(cls)

deepmd/entrypoints/eval_desc.py

@@ -128,6 +128,7 @@ def eval_desc(
             atype,
             fparam=fparam,
             aparam=aparam,
+            mixed_type=mixed_type,
         )
 
         # descriptors are kept in 3D format (nframes, natoms, ndesc)

deepmd/infer/deep_eval.py

@@ -11,6 +11,9 @@
 )
 
 import numpy as np
+from typing_extensions import (
+    Self,
+)
 
 from deepmd.backend.backend import (
     Backend,
@@ -89,9 +92,7 @@ def __init__(
     ) -> None:
         pass
 
-    def __new__(
-        cls, model_file: str, *args: object, **kwargs: object
-    ) -> "DeepEvalBackend":
+    def __new__(cls, model_file: str, *args: object, **kwargs: object) -> Self:
         if cls is DeepEvalBackend:
             backend = Backend.detect_backend_by_model(model_file)
             return super().__new__(backend().deep_eval)
@@ -384,7 +385,7 @@ class DeepEval(ABC):
         Keyword arguments.
     """
 
-    def __new__(cls, model_file: str, *args: object, **kwargs: object) -> "DeepEval":
+    def __new__(cls, model_file: str, *args: object, **kwargs: object) -> Self:
         if cls is DeepEval:
             deep_eval = DeepEvalBackend(
                 model_file,

deepmd/pd/entrypoints/main.py

@@ -574,7 +574,7 @@ def change_bias(
     log.info(f"Saved model to {output_path}")
 
 
-def main(args: list[str] | argparse.Namespace | None = None):
+def main(args: list[str] | argparse.Namespace | None = None) -> None:
     if not isinstance(args, argparse.Namespace):
         FLAGS = parse_args(args=args)
     else:

deepmd/pd/infer/deep_eval.py

@@ -474,7 +474,7 @@ def _eval_model(
         fparam: np.ndarray | None,
         aparam: np.ndarray | None,
         request_defs: list[OutputVariableDef],
-    ):
+    ) -> tuple[np.ndarray, ...]:
         if not self.static_model:
             model = self.dp.to(DEVICE)
         prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PD_FLOAT_PRECISION]]

deepmd/pd/infer/inference.py

@@ -23,9 +23,9 @@
 class Tester:
     def __init__(
         self,
-        model_ckpt,
-        head=None,
-    ):
+        model_ckpt: str,
+        head: str | None = None,
+    ) -> None:
         """Construct a DeePMD tester.
 
         Args:

deepmd/pd/loss/ener.py

@@ -1,4 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
 
 import paddle
 import paddle.nn.functional as F
@@ -20,7 +23,9 @@
 )
 
 
-def custom_huber_loss(predictions, targets, delta=1.0):
+def custom_huber_loss(
+    predictions: paddle.Tensor, targets: paddle.Tensor, delta: float = 1.0
+) -> paddle.Tensor:
     error = targets - predictions
     abs_error = paddle.abs(error)
     quadratic_loss = 0.5 * paddle.pow(error, 2)
@@ -32,13 +37,13 @@ def custom_huber_loss(predictions, targets, delta=1.0):
 class EnergyStdLoss(TaskLoss):
     def __init__(
         self,
-        starter_learning_rate=1.0,
-        start_pref_e=0.0,
-        limit_pref_e=0.0,
-        start_pref_f=0.0,
-        limit_pref_f=0.0,
-        start_pref_v=0.0,
-        limit_pref_v=0.0,
+        starter_learning_rate: float = 1.0,
+        start_pref_e: float = 0.0,
+        limit_pref_e: float = 0.0,
+        start_pref_f: float = 0.0,
+        limit_pref_f: float = 0.0,
+        start_pref_v: float = 0.0,
+        limit_pref_v: float = 0.0,
         start_pref_ae: float = 0.0,
         limit_pref_ae: float = 0.0,
         start_pref_pf: float = 0.0,
@@ -49,10 +54,10 @@ def __init__(
         limit_pref_gf: float = 0.0,
         numb_generalized_coord: int = 0,
         use_l1_all: bool = False,
-        inference=False,
-        use_huber=False,
-        huber_delta=0.01,
-        **kwargs,
+        inference: bool = False,
+        use_huber: bool = False,
+        huber_delta: float = 0.01,
+        **kwargs: Any,
     ) -> None:
         r"""Construct a layer to compute loss on energy, force and virial.
 
@@ -146,7 +151,15 @@ def __init__(
                 "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
             )
 
-    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+    def forward(
+        self,
+        input_dict: dict[str, paddle.Tensor],
+        model: paddle.nn.Layer,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool = False,
+    ) -> tuple[dict[str, paddle.Tensor], paddle.Tensor, dict[str, paddle.Tensor]]:
         """Return loss on energy and force.
 
         Parameters
@@ -535,10 +548,10 @@ def deserialize(cls, data: dict) -> "TaskLoss":
 class EnergyHessianStdLoss(EnergyStdLoss):
     def __init__(
         self,
-        start_pref_h=0.0,
-        limit_pref_h=0.0,
-        **kwargs,
-    ):
+        start_pref_h: float = 0.0,
+        limit_pref_h: float = 0.0,
+        **kwargs: Any,
+    ) -> None:
         r"""Enable the layer to compute loss on hessian.
 
         Parameters
@@ -556,7 +569,15 @@ def __init__(
         self.start_pref_h = start_pref_h
         self.limit_pref_h = limit_pref_h
 
-    def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
+    def forward(
+        self,
+        input_dict: dict[str, paddle.Tensor],
+        model: paddle.nn.Module,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool = False,
+    ) -> tuple[dict[str, paddle.Tensor], paddle.Tensor, dict[str, paddle.Tensor]]:
         model_pred, loss, more_loss = super().forward(
             input_dict, model, label, natoms, learning_rate, mae=mae
         )

deepmd/pd/loss/loss.py

@@ -3,6 +3,9 @@
     ABC,
     abstractmethod,
 )
+from typing import (
+    Any,
+)
 
 import paddle
 
@@ -15,11 +18,19 @@
 
 
 class TaskLoss(paddle.nn.Layer, ABC, make_plugin_registry("loss")):
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any) -> 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.Module,
+        label: dict[str, paddle.Tensor],
+        natoms: int,
+        learning_rate: float,
+        mae: bool | None = None,
+    ) -> paddle.Tensor:
         """Return loss ."""
         raise NotImplementedError
 

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -578,7 +578,7 @@ def _default_bias(self) -> paddle.Tensor:
             device=device
         )
 
-    def _default_std(self):
+    def _default_std(self) -> paddle.Tensor:
         ntypes = self.get_ntypes()
         return paddle.ones([self.n_out, ntypes, self.max_out_size], dtype=dtype).to(
             device=device
@@ -626,7 +626,7 @@ def _store_out_stat(
         paddle.assign(out_bias_data, self.out_bias)
         paddle.assign(out_std_data, self.out_std)
 
-    def get_ntypes(self):
+    def get_ntypes(self) -> int:
         return len(self.type_map)
 
     def get_buffer_ntypes(self) -> paddle.Tensor: