chore: refactor

deepmodeling · Feb 18, 2024 · b5f2582 · b5f2582
1 parent b733c57
commit b5f2582
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 18 deletions.
diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py
@@ -155,11 +155,11 @@ def call(
         self,
         descriptor: np.ndarray,
         atype: np.ndarray,
-        gr: np.ndarray = None,
-        g2: np.ndarray = None,
-        h2: np.ndarray = None,
-        fparam: np.ndarray = None,
-        aparam: np.ndarray = None,
+        gr: Optional[np.ndarray] = None,
+        g2: Optional[np.ndarray] = None,
+        h2: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
     ) -> Dict[str, np.ndarray]:
         """Calculate the fitting.
 

diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py
@@ -246,11 +246,11 @@ def _call_common(
         self,
         descriptor: np.ndarray,
         atype: np.ndarray,
-        gr: np.ndarray = None,
-        g2: np.ndarray = None,
-        h2: np.ndarray = None,
-        fparam: np.ndarray = None,
-        aparam: np.ndarray = None,
+        gr: Optional[np.ndarray] = None,
+        g2: Optional[np.ndarray] = None,
+        h2: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
     ) -> Dict[str, np.ndarray]:
         """Calculate the fitting.
 

diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py
@@ -190,14 +190,14 @@ def output_def(self):
 
     def call(
         self,
-        descriptor: np.array,
-        atype: np.array,
-        gr: Optional[np.array] = None,
-        g2: Optional[np.array] = None,
-        h2: Optional[np.array] = None,
-        fparam: Optional[np.array] = None,
-        aparam: Optional[np.array] = None,
-    ) -> Dict[str, np.array]:
+        descriptor: np.ndarray,
+        atype: np.ndarray,
+        gr: Optional[np.ndarray] = None,
+        g2: Optional[np.ndarray] = None,
+        h2: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+    ) -> Dict[str, np.ndarray]:
         """Calculate the fitting.
 
         Parameters