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

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

deepmd/pt/utils/dataset.py

@@ -489,9 +489,7 @@ def preprocess(self, batch):
 
         for kk in ["type", "real_natoms_vec"]:
             if kk in batch.keys():
-                batch[kk] = torch.tensor(
-                    batch[kk], dtype=torch.long, device=env.DEVICE
-                )
+                batch[kk] = torch.tensor(batch[kk], dtype=torch.long, device=env.DEVICE)
         batch["atype"] = batch.pop("type")
 
         keys = ["nlist", "nlist_loc", "nlist_type", "shift", "mapping"]
@@ -878,9 +876,7 @@ def __getitem__(self, index=None):
         if index is None:
             index = dp_random.choice(np.arange(self.nsystems), p=self.probs)
         b_data = self._data_systems[index].get_batch(self._batch_size)
-        b_data["natoms"] = torch.tensor(
-            self._natoms_vec[index], device=env.DEVICE
-        )
+        b_data["natoms"] = torch.tensor(self._natoms_vec[index], device=env.DEVICE)
         batch_size = b_data["coord"].shape[0]
         b_data["natoms"] = b_data["natoms"].unsqueeze(0).expand(batch_size, -1)
         return b_data
@@ -891,9 +887,7 @@ def get_training_batch(self, index=None):
         if index is None:
             index = dp_random.choice(np.arange(self.nsystems), p=self.probs)
         b_data = self._data_systems[index].get_batch_for_train(self._batch_size)
-        b_data["natoms"] = torch.tensor(
-            self._natoms_vec[index], device=env.DEVICE
-        )
+        b_data["natoms"] = torch.tensor(self._natoms_vec[index], device=env.DEVICE)
         batch_size = b_data["coord"].shape[0]
         b_data["natoms"] = b_data["natoms"].unsqueeze(0).expand(batch_size, -1)
         return b_data

deepmd/pt/utils/preprocess.py

@@ -165,9 +165,7 @@ def append_neighbors(coord, region: Region3D, atype, rcut: float):
     merged_coord = torch.cat([coord, tmp_coord])
     merged_coord_shift = torch.cat([torch.zeros_like(coord), coord_shift[tmp]])
     merged_atype = torch.cat([atype, tmp_atype])
-    merged_mapping = torch.cat(
-        [torch.arange(atype.numel(), device=env.DEVICE), aid]
-    )
+    merged_mapping = torch.cat([torch.arange(atype.numel(), device=env.DEVICE), aid])
     return merged_coord_shift, merged_atype, merged_mapping
 
 
@@ -198,12 +196,8 @@ def build_neighbor_list(
         sec = sec[-1:]
     lst = []
     nlist = torch.zeros((nloc, sec[-1].item()), device=env.DEVICE).long() - 1
-    nlist_loc = (
-        torch.zeros((nloc, sec[-1].item()), device=env.DEVICE).long() - 1
-    )
-    nlist_type = (
-        torch.zeros((nloc, sec[-1].item()), device=env.DEVICE).long() - 1
-    )
+    nlist_loc = torch.zeros((nloc, sec[-1].item()), device=env.DEVICE).long() - 1
+    nlist_type = torch.zeros((nloc, sec[-1].item()), device=env.DEVICE).long() - 1
     for i, nnei in enumerate(sec):
         if i > 0:
             nnei = nnei - sec[i - 1]
@@ -217,10 +211,7 @@ def build_neighbor_list(
         else:
             # when nnei > nall
             indices = torch.zeros((nloc, nnei), device=env.DEVICE).long() - 1
-            _sorted = (
-                torch.ones((nloc, nnei), device=env.DEVICE).long()
-                * DISTANCE_INF
-            )
+            _sorted = torch.ones((nloc, nnei), device=env.DEVICE).long() * DISTANCE_INF
             _sorted_nnei, indices_nnei = torch.topk(
                 tmp, tmp.shape[1], dim=1, largest=False
             )

source/tests/pt/test_descriptor.py

@@ -12,15 +12,13 @@
 from pathlib import (
     Path,
 )
-from deepmd.pt.utils import (
-    env,
-)
 
 from deepmd.pt.model.descriptor import (
     prod_env_mat_se_a,
 )
 from deepmd.pt.utils import (
     dp_random,
+    env,
 )
 from deepmd.pt.utils.dataset import (
     DeepmdDataSet,
@@ -115,10 +113,14 @@ def setUp(self):
 
     def test_consistency(self):
         avg_zero = torch.zeros(
-            [self.ntypes, self.nnei * 4], dtype=GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+            [self.ntypes, self.nnei * 4],
+            dtype=GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
         )
         std_ones = torch.ones(
-            [self.ntypes, self.nnei * 4], dtype=GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
+            [self.ntypes, self.nnei * 4],
+            dtype=GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
         )
         base_d, base_force, nlist = base_se_a(
             rcut=self.rcut,

source/tests/pt/test_mlp.py

@@ -5,12 +5,12 @@
 import numpy as np
 import torch
 
-from deepmd.pt.utils.env import (
-    PRECISION_DICT,
-)
 from deepmd.pt.utils import (
     env,
 )
+from deepmd.pt.utils.env import (
+    PRECISION_DICT,
+)
 
 try:
     from deepmd.pt.model.network.mlp import (
@@ -107,9 +107,13 @@ def test_match_native_layer(
                 inp_shap = ashp + inp_shap
             rtol, atol = get_tols(prec)
             dtype = PRECISION_DICT[prec]
-            xx = torch.arange(np.prod(inp_shap), dtype=dtype, device=env.DEVICE).view(inp_shap)
+            xx = torch.arange(np.prod(inp_shap), dtype=dtype, device=env.DEVICE).view(
+                inp_shap
+            )
             # def mlp layer
-            ml = MLPLayer(ninp, nout, bias, ut, ac, resnet, precision=prec).to(env.DEVICE)
+            ml = MLPLayer(ninp, nout, bias, ut, ac, resnet, precision=prec).to(
+                env.DEVICE
+            )
             # check consistency
             nl = NativeLayer.deserialize(ml.serialize())
             np.testing.assert_allclose(
@@ -160,7 +164,9 @@ def test_match_native_net(
                 inp_shap = ashp + inp_shap
             rtol, atol = get_tols(prec)
             dtype = PRECISION_DICT[prec]
-            xx = torch.arange(np.prod(inp_shap), dtype=dtype, device=env.DEVICE).view(inp_shap)
+            xx = torch.arange(np.prod(inp_shap), dtype=dtype, device=env.DEVICE).view(
+                inp_shap
+            )
             # def MLP
             layers = []
             for ii in range(1, len(ndims)):