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

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

deepmd/pt/utils/dataloader.py

@@ -12,7 +12,7 @@
 from typing import (
     List,
 )
-import numpy as np
+
 import h5py
 import torch
 import torch.distributed as dist
@@ -22,12 +22,13 @@
     Dataset,
     WeightedRandomSampler,
 )
+from torch.utils.data._utils.collate import (
+    collate_tensor_fn,
+)
 from torch.utils.data.distributed import (
     DistributedSampler,
 )
-from torch.utils.data._utils.collate import (
-    collate_tensor_fn
-)
+
 from deepmd.pt.model.descriptor import (
     Descriptor,
 )
@@ -231,6 +232,7 @@ def __next__(self):
             raise StopIteration
         return item
 
+
 def collate_batch(batch):
     example = batch[0]
     result = {}
@@ -242,10 +244,14 @@ def collate_batch(batch):
                 result[key] = None
             elif key == "fid":
                 result[key] = [d[key] for d in batch]
-            elif key == 'type':
-                result['atype'] = collate_tensor_fn([torch.as_tensor(d[key]) for d in batch])
+            elif key == "type":
+                result["atype"] = collate_tensor_fn(
+                    [torch.as_tensor(d[key]) for d in batch]
+                )
             else:
-                result[key] = collate_tensor_fn([torch.as_tensor(d[key]) for d in batch])
+                result[key] = collate_tensor_fn(
+                    [torch.as_tensor(d[key]) for d in batch]
+                )
     return result
 
 

deepmd/pt/utils/dataset.py

@@ -29,6 +29,7 @@
     DeepmdData,
 )
 
+
 class DeepmdDataSystem:
     def __init__(
         self,
@@ -574,9 +575,7 @@ def __init__(
         """
         self._type_map = type_map
         self._data_system = DeepmdData(
-            sys_path=system,
-            shuffle_test=shuffle,
-            type_map=self._type_map
+            sys_path=system, shuffle_test=shuffle, type_map=self._type_map
         )
         self._data_system.add("energy", 1, atomic=False, must=False, high_prec=True)
         self._data_system.add("force", 3, atomic=True, must=False, high_prec=False)

deepmd/utils/data.py

@@ -257,7 +257,8 @@ def check_test_size(self, test_size):
             return self.test_dir, tmpe.shape[0]
         else:
             return None
-    def get_item(self,index:int) -> dict:
+
+    def get_item(self, index: int) -> dict:
         """Get a single frame data . The frame is picked from the data system by index.
 
         Parameters
@@ -275,7 +276,7 @@ def get_item(self,index:int) -> dict:
         frame = self.preprocess(frame)
         frame["fid"] = index
         return frame
-    
+
     def get_batch(self, batch_size: int) -> dict:
         """Get a batch of data with `batch_size` frames. The frames are randomly picked from the data system.
 
@@ -466,8 +467,8 @@ def _shuffle_data(self, data):
             else:
                 ret[kk] = data[kk]
         return ret, idx
-    
-    def _get_nframes(self,set_name:DPPath):
+
+    def _get_nframes(self, set_name: DPPath):
         # get nframes
         if not isinstance(set_name, DPPath):
             set_name = DPPath(set_name)
@@ -480,7 +481,7 @@ def _get_nframes(self,set_name:DPPath):
             coord = coord.reshape([1, -1])
         nframes = coord.shape[0]
         return nframes
-    
+
     def preprocess(self, data):
         for kk in self.data_dict.keys():
             if "find_" in kk: