[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

deepmd/dpmodel/atomic_model/__init__.py

@@ -17,20 +17,17 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
-from .dp_atomic_model import (
-    DPAtomicModel,
-)
-from .energy_atomic_model import (
-    DPEnergyAtomicModel,
+from .dipole_atomic_model import (
+    DPDipoleAtomicModel,
 )
 from .dos_atomic_model import (
     DPDOSAtomicModel,
 )
-from .polar_atomic_model import (
-    DPPolarAtomicModel,
+from .dp_atomic_model import (
+    DPAtomicModel,
 )
-from .dipole_atomic_model import (
-    DPDipoleAtomicModel,
+from .energy_atomic_model import (
+    DPEnergyAtomicModel,
 )
 from .linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel,
@@ -42,6 +39,9 @@
 from .pairtab_atomic_model import (
     PairTabAtomicModel,
 )
+from .polar_atomic_model import (
+    DPPolarAtomicModel,
+)
 
 __all__ = [
     "make_base_atomic_model",

deepmd/dpmodel/atomic_model/dipole_atomic_model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
+
 from deepmd.dpmodel.fitting.dipole_fitting import (
     DipoleFitting,
 )

deepmd/dpmodel/atomic_model/energy_atomic_model.py

@@ -11,8 +11,7 @@
 
 class DPEnergyAtomicModel(DPAtomicModel):
     def __init__(self, descriptor, fitting, type_map, **kwargs):
-        assert (
-            isinstance(fitting, EnergyFittingNet)
-            or isinstance(fitting, InvarFitting)
+        assert isinstance(fitting, EnergyFittingNet) or isinstance(
+            fitting, InvarFitting
         )
         super().__init__(descriptor, fitting, type_map, **kwargs)

deepmd/dpmodel/atomic_model/polar_atomic_model.py

@@ -17,43 +17,42 @@ def __init__(self, descriptor, fitting, type_map, **kwargs):
         super().__init__(descriptor, fitting, type_map, **kwargs)
 
     def apply_out_stat(
-            self,
-            ret: dict[str, np.ndarray],
-            atype: np.ndarray,
-        ):
-            """Apply the stat to each atomic output.
-
-            Parameters
-            ----------
-            ret
-                The returned dict by the forward_atomic method
-            atype
-                The atom types. nf x nloc
-
-            """
-            out_bias, out_std = self._fetch_out_stat(self.bias_keys)
-
-            if self.fitting_net.shift_diag:
-                nframes, nloc = atype.shape
-                dtype = out_bias[self.bias_keys[0]].dtype
-                for kk in self.bias_keys:
-                    ntypes = out_bias[kk].shape[0]
-                    temp = np.zeros(ntypes,dtype=dtype)
-                    for i in range(ntypes):
-                        temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3)))
-                    modified_bias = temp[atype]
-
-                    # (nframes, nloc, 1)
-                    modified_bias = (
-                        modified_bias[..., np.newaxis]
-                        * (self.fitting_net.scale[atype])
-                    )
-
-                    eye = np.eye(3,dtype=dtype)
-                    eye = np.tile(eye, (nframes, nloc, 1, 1))
-                    # (nframes, nloc, 3, 3)
-                    modified_bias = modified_bias[..., np.newaxis] * eye
-
-                    # nf x nloc x odims, out_bias: ntypes x odims
-                    ret[kk] = ret[kk] + modified_bias
-            return ret
+        self,
+        ret: dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        """Apply the stat to each atomic output.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+
+        if self.fitting_net.shift_diag:
+            nframes, nloc = atype.shape
+            dtype = out_bias[self.bias_keys[0]].dtype
+            for kk in self.bias_keys:
+                ntypes = out_bias[kk].shape[0]
+                temp = np.zeros(ntypes, dtype=dtype)
+                for i in range(ntypes):
+                    temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3)))
+                modified_bias = temp[atype]
+
+                # (nframes, nloc, 1)
+                modified_bias = (
+                    modified_bias[..., np.newaxis] * (self.fitting_net.scale[atype])
+                )
+
+                eye = np.eye(3, dtype=dtype)
+                eye = np.tile(eye, (nframes, nloc, 1, 1))
+                # (nframes, nloc, 3, 3)
+                modified_bias = modified_bias[..., np.newaxis] * eye
+
+                # nf x nloc x odims, out_bias: ntypes x odims
+                ret[kk] = ret[kk] + modified_bias
+        return ret

deepmd/dpmodel/model/dos_model.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later