chore: improve type anotations in deepmd.infer

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

deepmd/dpmodel/infer/deep_eval.py

@@ -76,10 +76,10 @@ def __init__(
         self,
         model_file: str,
         output_def: ModelOutputDef,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         self.output_def = output_def
         self.model_path = model_file
@@ -161,12 +161,12 @@ def get_ntypes_spin(self):
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 

deepmd/entrypoints/test.py

@@ -46,14 +46,7 @@
 )
 
 if TYPE_CHECKING:
-    from deepmd.tf.infer import (
-        DeepDipole,
-        DeepDOS,
-        DeepPolar,
-        DeepPot,
-        DeepWFC,
-    )
-    from deepmd.tf.infer.deep_tensor import (
+    from deepmd.infer.deep_tensor import (
         DeepTensor,
     )
 

deepmd/infer/deep_dos.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Any,
-    Dict,
     List,
     Optional,
     Tuple,
@@ -70,7 +69,7 @@ def eval(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Tuple[np.ndarray, ...]:
         """Evaluate energy, force, and virial. If atomic is True,
         also return atomic energy and atomic virial.

deepmd/infer/deep_eval.py

@@ -11,6 +11,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -82,10 +83,10 @@ def __init__(
         self,
         model_file: str,
         output_def: ModelOutputDef,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         pass
 
@@ -99,12 +100,12 @@ def __new__(cls, model_file: str, *args, **kwargs):
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 
@@ -166,13 +167,13 @@ def get_dim_aparam(self) -> int:
     def eval_descriptor(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         efield: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> np.ndarray:
         """Evaluate descriptors by using this DP.
 
@@ -246,11 +247,11 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool:
             # assume mixed_types if there are virtual types, even when
             # the atom types of all frames are the same
             return False
-        return np.all(np.equal(atom_types, atom_types[0]))
+        return np.all(np.equal(atom_types, atom_types[0])).item()
 
     @property
     @abstractmethod
-    def model_type(self) -> "DeepEval":
+    def model_type(self) -> Type["DeepEval"]:
         """The the evaluator of the model type."""
 
     @abstractmethod
@@ -316,10 +317,10 @@ def __new__(cls, model_file: str, *args, **kwargs):
     def __init__(
         self,
         model_file: str,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> None:
         self.deep_eval = DeepEvalBackend(
             model_file,
@@ -387,7 +388,7 @@ def eval_descriptor(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> np.ndarray:
         """Evaluate descriptors by using this DP.
 

deepmd/infer/deep_polar.py

@@ -50,7 +50,7 @@ def eval(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: dict,
+        **kwargs,
     ) -> np.ndarray:
         """Evaluate the model.
 

deepmd/infer/deep_pot.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Any,
-    Dict,
     List,
+    Literal,
     Optional,
     Tuple,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -89,6 +90,48 @@ def output_def_mag(self) -> ModelOutputDef:
             )
         )
 
+    @overload
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: Optional[np.ndarray],
+        atom_types: Union[List[int], np.ndarray],
+        atomic: Literal[True],
+        fparam: Optional[np.ndarray],
+        aparam: Optional[np.ndarray],
+        mixed_type: bool,
+        **kwargs: Any,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        pass
+
+    @overload
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: Optional[np.ndarray],
+        atom_types: Union[List[int], np.ndarray],
+        atomic: Literal[False],
+        fparam: Optional[np.ndarray],
+        aparam: Optional[np.ndarray],
+        mixed_type: bool,
+        **kwargs: Any,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        pass
+
+    @overload
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: Optional[np.ndarray],
+        atom_types: Union[List[int], np.ndarray],
+        atomic: bool,
+        fparam: Optional[np.ndarray],
+        aparam: Optional[np.ndarray],
+        mixed_type: bool,
+        **kwargs: Any,
+    ) -> Tuple[np.ndarray, ...]:
+        pass
+
     def eval(
         self,
         coords: np.ndarray,
@@ -98,7 +141,7 @@ def eval(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         mixed_type: bool = False,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Tuple[np.ndarray, ...]:
         """Evaluate energy, force, and virial. If atomic is True,
         also return atomic energy and atomic virial.

deepmd/infer/model_devi.py

@@ -28,7 +28,7 @@ def calc_model_devi_f(
     fs: np.ndarray,
     real_f: Optional[np.ndarray] = None,
     relative: Optional[float] = None,
-    atomic: Literal[False] = False,
+    atomic: Literal[False] = ...,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: ...
 
 
@@ -37,11 +37,19 @@ def calc_model_devi_f(
     fs: np.ndarray,
     real_f: Optional[np.ndarray] = None,
     relative: Optional[float] = None,
-    *,
-    atomic: Literal[True],
+    atomic: Literal[True] = ...,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ...
 
 
+@overload
+def calc_model_devi_f(
+    fs: np.ndarray,
+    real_f: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
+    atomic: bool = False,
+) -> Tuple[np.ndarray, ...]: ...
+
+
 def calc_model_devi_f(
     fs: np.ndarray,
     real_f: Optional[np.ndarray] = None,

deepmd/pt/infer/deep_eval.py

@@ -7,6 +7,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -87,11 +88,11 @@ def __init__(
         self,
         model_file: str,
         output_def: ModelOutputDef,
-        *args: List[Any],
+        *args: Any,
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
         head: Optional[str] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ):
         self.output_def = output_def
         self.model_path = model_file
@@ -165,7 +166,7 @@ def get_dim_aparam(self) -> int:
         return self.dp.model["Default"].get_dim_aparam()
 
     @property
-    def model_type(self) -> "DeepEvalWrapper":
+    def model_type(self) -> Type["DeepEvalWrapper"]:
         """The the evaluator of the model type."""
         model_output_type = self.dp.model["Default"].model_output_type()
         if "energy" in model_output_type:
@@ -211,12 +212,12 @@ def get_has_spin(self):
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 

deepmd/tf/infer/deep_eval.py

@@ -10,6 +10,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -262,7 +263,7 @@ def _init_attr(self):
 
     @property
     @lru_cache(maxsize=None)
-    def model_type(self) -> "DeepEvalWrapper":
+    def model_type(self) -> Type["DeepEvalWrapper"]:
         """Get type of model.
 
         :type:str
@@ -693,13 +694,13 @@ def _get_natoms_and_nframes(
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         atomic: bool = False,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
         efield: Optional[np.ndarray] = None,
-        **kwargs: Dict[str, Any],
+        **kwargs: Any,
     ) -> Dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
 
@@ -1023,7 +1024,7 @@ def _get_output_shape(self, odef, nframes, natoms):
     def eval_descriptor(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
@@ -1080,7 +1081,7 @@ def eval_descriptor(
     def _eval_descriptor_inner(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: np.ndarray,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,

deepmd/tf/infer/deep_tensor.py

@@ -146,7 +146,7 @@ def get_dim_aparam(self) -> int:
     def eval(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: List[int],
         atomic: bool = True,
         fparam: Optional[np.ndarray] = None,
@@ -276,7 +276,7 @@ def eval(
     def eval_full(
         self,
         coords: np.ndarray,
-        cells: np.ndarray,
+        cells: Optional[np.ndarray],
         atom_types: List[int],
         atomic: bool = False,
         fparam: Optional[np.array] = None,