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

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

deepmd_utils/model_format/__init__.py

@@ -2,6 +2,9 @@
 from .common import (
     PRECISION_DICT,
 )
+from .env_mat import (
+    EnvMat,
+)
 from .network import (
     EmbeddingNet,
     NativeLayer,
@@ -10,10 +13,6 @@
     save_dp_model,
     traverse_model_dict,
 )
-from .env_mat import(
-  EnvMat,
-)
-
 
 __all__ = [
     "EmbeddingNet",

deepmd_utils/model_format/common.py

@@ -1,4 +1,8 @@
-from abc import ABC
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+)
+
 import numpy as np
 
 PRECISION_DICT = {
@@ -8,6 +12,7 @@
     "default": np.float64,
 }
 
+
 class NativeOP(ABC):
     """The unit operation of a native model."""
 

deepmd_utils/model_format/env_mat.py

@@ -1,9 +1,11 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
 
 from .common import (
-  NativeOP,
+    NativeOP,
 )
 
+
 def compute_smooth_weight(
     distance: np.ndarray,
     rmin: float,
@@ -14,9 +16,10 @@ def compute_smooth_weight(
     max_mask = distance >= rmax
     mid_mask = np.logical_not(np.logical_or(min_mask, max_mask))
     uu = (distance - rmin) / (rmax - rmin)
-    vv = uu * uu * uu * (-6. * uu * uu + 15. * uu - 10.) + 1.
+    vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0
     return vv * mid_mask + min_mask
 
+
 def _make_env_mat(
     nlist,
     coord,
@@ -27,59 +30,56 @@ def _make_env_mat(
     nf, nloc, nnei = nlist.shape
     # nf x nall x 3
     coord = coord.reshape(nf, -1, 3)
-    mask = (nlist >= 0)
-    nlist = (nlist * mask)
+    mask = nlist >= 0
+    nlist = nlist * mask
     # nf x (nloc x nnei) x 3
     index = np.tile(nlist.reshape(nf, -1, 1), (1, 1, 3))
     coord_r = np.take_along_axis(coord, index, 1)
     # nf x nloc x nnei x 3
     coord_r = coord_r.reshape(nf, nloc, nnei, 3)
     # nf x nloc x 1 x 3
-    coord_l = coord[:,:nloc].reshape(nf, -1, 1, 3)
+    coord_l = coord[:, :nloc].reshape(nf, -1, 1, 3)
     # nf x nloc x nnei x 3
     diff = coord_r - coord_l
     # nf x nloc x nnei
     length = np.linalg.norm(diff, axis=-1, keepdims=True)
     # for index 0 nloc atom
     length = length + ~np.expand_dims(mask, -1)
     t0 = 1 / length
-    t1 = diff / length ** 2
+    t1 = diff / length**2
     weight = compute_smooth_weight(length, ruct_smth, rcut)
-    env_mat_se_a = (
-      np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1)
-    )
-    return env_mat_se_a, diff * np.expand_dims(mask,-1), weight
-
+    env_mat_se_a = np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1)
+    return env_mat_se_a, diff * np.expand_dims(mask, -1), weight
+
 
 class EnvMat(NativeOP):
-  def __init__(
-      self,
-      rcut,
-      rcut_smth,
-  ):
-    self.rcut = rcut
-    self.rcut_smth = rcut_smth
+    def __init__(
+        self,
+        rcut,
+        rcut_smth,
+    ):
+        self.rcut = rcut
+        self.rcut_smth = rcut_smth
 
-  def call(
-      self,
-      nlist,
-      coord_ext,
-  ):
-    em, diff, ww = _make_env_mat(
-      nlist, coord_ext, self.rcut, self.rcut_smth)
-    return em, ww
+    def call(
+        self,
+        nlist,
+        coord_ext,
+    ):
+        em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth)
+        return em, ww
 
-  def serialize(
-      self,
-  )->dict:
-    return {
-      "rcut": self.rcut,
-      "rcut_smth": self.rcut_smth,
-    }
+    def serialize(
+        self,
+    ) -> dict:
+        return {
+            "rcut": self.rcut,
+            "rcut_smth": self.rcut_smth,
+        }
 
-  @classmethod
-  def deserialize(
-      cls,
-      data: dict,
-  )->"EnvMat":
-    return cls(**data)
+    @classmethod
+    def deserialize(
+        cls,
+        data: dict,
+    ) -> "EnvMat":
+        return cls(**data)

deepmd_utils/model_format/network.py

@@ -18,8 +18,8 @@
     __version__ = "unknown"
 
 from .common import (
-  NativeOP,
-  PRECISION_DICT,
+    PRECISION_DICT,
+    NativeOP,
 )
 
 
@@ -126,8 +126,6 @@ def load_dp_model(filename: str) -> dict:
     return model_dict
 
 
-
-
 class NativeLayer(NativeOP):
     """Native representation of a layer.
 

source/tests/test_model_format_utils.py

@@ -9,8 +9,8 @@
 import numpy as np
 
 from deepmd_utils.model_format import (
-    EnvMat,
     EmbeddingNet,
+    EnvMat,
     NativeLayer,
     NativeNet,
     load_dp_model,
@@ -20,7 +20,10 @@
 
 class TestNativeLayer(unittest.TestCase):
     def test_serialize_deserize(self):
-        for (ni, no), bias, ut, activation_function, resnet, ashp, prec in itertools.product(
+        for (
+            ni,
+            no,
+        ), bias, ut, activation_function, resnet, ashp, prec in itertools.product(
             [(5, 5), (5, 10), (5, 9), (9, 5)],
             [True, False],
             [True, False],
@@ -139,39 +142,38 @@ def tearDown(self) -> None:
 
 
 class TestEnvMat(unittest.TestCase):
-  def setUp(self):
-    # nloc == 3, nall == 4
-    self.nloc = 3
-    self.nall = 4
-    self.coord_ext = np.array(
-      [ [0, 0, 0],
-        [0, 1, 0],
-        [0, 0, 1],
-        [0, -2, 0],
-       ], 
-      dtype = np.float64,
-    ).reshape([1, self.nall*3])
-    self.atype_ext = np.array(
-      [0, 0, 1, 0], dtype=int
-    ).reshape([1, self.nall])
-    # sel = [5, 2]
-    self.nlist = np.array(
-      [
-        [1, 3, -1, -1, -1, 2, -1],
-        [0,-1, -1, -1, -1, 2, -1],
-        [0, 1, -1, -1, -1, 0, -1],
-      ], 
-      dtype=int,
-    ).reshape([1, self.nloc, 7])
-    self.rcut = .4
-    self.rcut_smth = 2.2
-
-  def test_self_consistency(
-      self,
-  ):
-    em0 = EnvMat(self.rcut, self.rcut_smth)
-    em1 = EnvMat.deserialize(em0.serialize())
-    mm0, ww0 = em0.call(self.nlist, self.coord_ext)
-    mm1, ww1 = em1.call(self.nlist, self.coord_ext)
-    np.testing.assert_allclose(mm0, mm1)
-    np.testing.assert_allclose(ww0, ww1)
+    def setUp(self):
+        # nloc == 3, nall == 4
+        self.nloc = 3
+        self.nall = 4
+        self.coord_ext = np.array(
+            [
+                [0, 0, 0],
+                [0, 1, 0],
+                [0, 0, 1],
+                [0, -2, 0],
+            ],
+            dtype=np.float64,
+        ).reshape([1, self.nall * 3])
+        self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall])
+        # sel = [5, 2]
+        self.nlist = np.array(
+            [
+                [1, 3, -1, -1, -1, 2, -1],
+                [0, -1, -1, -1, -1, 2, -1],
+                [0, 1, -1, -1, -1, 0, -1],
+            ],
+            dtype=int,
+        ).reshape([1, self.nloc, 7])
+        self.rcut = 0.4
+        self.rcut_smth = 2.2
+
+    def test_self_consistency(
+        self,
+    ):
+        em0 = EnvMat(self.rcut, self.rcut_smth)
+        em1 = EnvMat.deserialize(em0.serialize())
+        mm0, ww0 = em0.call(self.nlist, self.coord_ext)
+        mm1, ww1 = em1.call(self.nlist, self.coord_ext)
+        np.testing.assert_allclose(mm0, mm1)
+        np.testing.assert_allclose(ww0, ww1)