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

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

Author: pre-commit-ci[bot]

deepmd/dpmodel/fitting/property_fitting.py

@@ -6,16 +6,16 @@
 
 import numpy as np
 
-from deepmd.dpmodel.output_def import (
-    FittingOutputDef,
-    OutputVariableDef,
-)
 from deepmd.dpmodel.common import (
     DEFAULT_PRECISION,
 )
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )

deepmd/dpmodel/model/dp_model.py

@@ -48,4 +48,4 @@ def update_sel(
 
     def get_fitting_net(self):
         """Get the fitting network."""
-        return self.atomic_model.fitting
+        return self.atomic_model.fitting

deepmd/dpmodel/model/property_model.py

@@ -28,4 +28,4 @@ def __init__(
 
     def get_var_name(self) -> str:
         """Get the name of the property."""
-        return self.get_fitting_net().var_name
+        return self.get_fitting_net().var_name

deepmd/dpmodel/model/transform_output.py

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from typing import (
+    Optional,
+)
+
 import array_api_compat
 import numpy as np
 
@@ -17,15 +21,14 @@
     get_hessian_name,
     get_reduce_name,
 )
-from typing import Optional
 
 
 def fit_output_to_model_output(
     fit_ret: dict[str, np.ndarray],
     fit_output_def: FittingOutputDef,
     coord_ext: np.ndarray,
     do_atomic_virial: bool = False,
-    mask: Optional[np.ndarray] = None
+    mask: Optional[np.ndarray] = None,
 ) -> dict[str, np.ndarray]:
     """Transform the output of the fitting network to
     the model output.
@@ -43,10 +46,15 @@ def fit_output_to_model_output(
             if vdef.intensive:
                 if (mask is not None) and (mask == 0.0).any():
                     mask = mask.astype(bool)
-                    model_ret[kk_redu] = xp.stack([
-                        xp.mean(vv[ii].astype(GLOBAL_ENER_FLOAT_PRECISION)[mask[ii]], axis=atom_axis)
-                        for ii in range(mask.shape[0])
-                    ])
+                    model_ret[kk_redu] = xp.stack(
+                        [
+                            xp.mean(
+                                vv[ii].astype(GLOBAL_ENER_FLOAT_PRECISION)[mask[ii]],
+                                axis=atom_axis,
+                            )
+                            for ii in range(mask.shape[0])
+                        ]
+                    )
                 else:
                     model_ret[kk_redu] = xp.mean(
                         vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis

deepmd/jax/model/base_model.py

@@ -48,12 +48,14 @@ def forward_common_atomic(
             kk_redu = get_reduce_name(kk)
             if vdef.intensive:
                 mask = atomic_ret["mask"] if "mask" in atomic_ret else None
-                if (mask is not None) and (mask==0.0).any():
+                if (mask is not None) and (mask == 0.0).any():
                     mask = mask.astype(jnp.bool_)
-                    model_predict[kk_redu] = jnp.stack([
-                        jnp.mean(vv[ii][mask[ii]], axis=atom_axis)
-                        for ii in range(mask.shape[0])
-                    ])
+                    model_predict[kk_redu] = jnp.stack(
+                        [
+                            jnp.mean(vv[ii][mask[ii]], axis=atom_axis)
+                            for ii in range(mask.shape[0])
+                        ]
+                    )
                 else:
                     model_predict[kk_redu] = jnp.mean(vv, axis=atom_axis)
             else:

source/tests/common/dpmodel/test_padding_atoms.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import sys
 import unittest
+from copy import (
+    deepcopy,
+)
 
 import numpy as np
-from copy import deepcopy
 
 from deepmd.dpmodel.descriptor.se_e2_a import (
     DescrptSeA,
@@ -32,13 +33,13 @@ def setUp(self) -> None:
                     [1, 0, 1],
                     [0, 1, 1],
                     [1, 1, 0],
-                ]
+                ],
             ],
             dtype=np.float64,
         )
-        self.atype = np.array([[0, 0, 1],[1, 1, 0]], dtype=int).reshape([2, self.nloc])
+        self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc])
         self.cell = 2.0 * np.eye(3).reshape([1, 9])
-        self.cell = np.array([self.cell,self.cell]).reshape(2, 9)
+        self.cell = np.array([self.cell, self.cell]).reshape(2, 9)
         self.sel = [16, 8]
         self.rcut = 2.2
         self.rcut_smth = 0.4
@@ -69,7 +70,7 @@ def test_padding_atoms_consistency(self):
         # test intensive
         np.testing.assert_allclose(
             result[f"{var_name}_redu"],
-            np.mean(result[f"{var_name}"],axis=1),
+            np.mean(result[f"{var_name}"], axis=1),
             atol=self.atol,
         )
         # test padding atoms
@@ -80,14 +81,14 @@ def test_padding_atoms_consistency(self):
             atype_padding = np.pad(
                 atype,
                 pad_width=((0, 0), (0, padding_atoms)),
-                mode='constant',
-                constant_values=-1
+                mode="constant",
+                constant_values=-1,
             )
             coord_padding = np.pad(
                 coord,
                 pad_width=((0, 0), (0, padding_atoms), (0, 0)),
-                mode='constant',
-                constant_values=0
+                mode="constant",
+                constant_values=0,
             )
             args = [coord_padding, atype_padding, self.cell]
             result_padding = model.call(*args)

source/tests/jax/test_padding_atoms.py

@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import sys
 import unittest
+from copy import (
+    deepcopy,
+)
 
 import numpy as np
-from copy import deepcopy
 
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -49,13 +51,13 @@ def setUp(self) -> None:
                     [1, 0, 1],
                     [0, 1, 1],
                     [1, 1, 0],
-                ]
+                ],
             ],
             dtype=np.float64,
         )
-        self.atype = np.array([[0, 0, 1],[1, 1, 0]], dtype=int).reshape([2, self.nloc])
+        self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc])
         self.cell = 2.0 * np.eye(3).reshape([1, 9])
-        self.cell = np.array([self.cell,self.cell]).reshape(2, 9)
+        self.cell = np.array([self.cell, self.cell]).reshape(2, 9)
         self.sel = [16, 8]
         self.rcut = 2.2
         self.rcut_smth = 0.4
@@ -90,7 +92,7 @@ def test_padding_atoms_consistency(self):
         # test intensive
         np.testing.assert_allclose(
             to_numpy_array(result[f"{var_name}_redu"]),
-            np.mean(to_numpy_array(result[f"{var_name}"]),axis=1),
+            np.mean(to_numpy_array(result[f"{var_name}"]), axis=1),
             atol=self.atol,
         )
         # test padding atoms
@@ -101,16 +103,18 @@ def test_padding_atoms_consistency(self):
             atype_padding = np.pad(
                 atype,
                 pad_width=((0, 0), (0, padding_atoms)),
-                mode='constant',
-                constant_values=-1
+                mode="constant",
+                constant_values=-1,
             )
             coord_padding = np.pad(
                 coord,
                 pad_width=((0, 0), (0, padding_atoms), (0, 0)),
-                mode='constant',
-                constant_values=0
+                mode="constant",
+                constant_values=0,
             )
-            args = [to_jax_array(ii) for ii in [coord_padding, atype_padding, self.cell]]
+            args = [
+                to_jax_array(ii) for ii in [coord_padding, atype_padding, self.cell]
+            ]
             result_padding = model.call(*args)
             np.testing.assert_allclose(
                 to_numpy_array(result[f"{var_name}_redu"]),

source/tests/pt/test_padding_atoms.py

@@ -1,24 +1,24 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import sys
 import unittest
+from copy import (
+    deepcopy,
+)
 
 import numpy as np
 import torch
-from copy import deepcopy
 
 from deepmd.dpmodel.common import (
     to_numpy_array,
 )
-
 from deepmd.pt.model.descriptor import (
     DescrptSeA,
 )
-from deepmd.pt.model.task import (
-    PropertyFittingNet,
-)
 from deepmd.pt.model.model import (
     PropertyModel,
 )
+from deepmd.pt.model.task import (
+    PropertyFittingNet,
+)
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
@@ -40,13 +40,13 @@ def setUp(self) -> None:
                     [1, 0, 1],
                     [0, 1, 1],
                     [1, 1, 0],
-                ]
+                ],
             ],
             dtype=np.float64,
         )
-        self.atype = np.array([[0, 0, 1],[1, 1, 0]], dtype=int).reshape([2, self.nloc])
+        self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc])
         self.cell = 2.0 * np.eye(3).reshape([1, 9])
-        self.cell = np.array([self.cell,self.cell]).reshape(2, 9)
+        self.cell = np.array([self.cell, self.cell]).reshape(2, 9)
         self.sel = [16, 8]
         self.rcut = 2.2
         self.rcut_smth = 0.4
@@ -78,7 +78,7 @@ def test_padding_atoms_consistency(self):
         # test intensive
         np.testing.assert_allclose(
             to_numpy_array(result[var_name].cpu().detach()),
-            np.mean(to_numpy_array(result[f"atom_{var_name}"].cpu().detach()),axis=1),
+            np.mean(to_numpy_array(result[f"atom_{var_name}"].cpu().detach()), axis=1),
             atol=self.atol,
         )
         # test padding atoms
@@ -89,17 +89,20 @@ def test_padding_atoms_consistency(self):
             atype_padding = np.pad(
                 atype,
                 pad_width=((0, 0), (0, padding_atoms)),
-                mode='constant',
-                constant_values=-1
+                mode="constant",
+                constant_values=-1,
             )
             coord_padding = np.pad(
                 coord,
                 pad_width=((0, 0), (0, padding_atoms), (0, 0)),
-                mode='constant',
-                constant_values=0
+                mode="constant",
+                constant_values=0,
             )
             with torch.device("cpu"):
-                args = [to_torch_tensor(ii) for ii in [coord_padding, atype_padding, self.cell]]
+                args = [
+                    to_torch_tensor(ii)
+                    for ii in [coord_padding, atype_padding, self.cell]
+                ]
                 result_padding = model(*args)
             np.testing.assert_allclose(
                 to_numpy_array(result[var_name].cpu().detach()),