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

for more information, see https://pre-commit.ci
deepmodeling · Jan 10, 2025 · 9f389ad · 9f389ad
1 parent d5596bf
commit 9f389ad
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py
@@ -62,11 +62,13 @@ def get_frame_index_for_elements(self):
         element_counts = defaultdict(lambda: {"frames": 0, "indices": []})
         set_files = self._data_system.dirs
         base_offset = 0
-        global_type_name = {} 
+        global_type_name = {}
         for set_file in set_files:
             element_data = self._data_system._load_type_mix(set_file)
             unique_elements = np.unique(element_data)
-            type_name = self._data_system.build_reidx_to_name_map(element_data,set_file)
+            type_name = self._data_system.build_reidx_to_name_map(
+                element_data, set_file
+            )
             for new_idx, elem_name in type_name.items():
                 if new_idx not in global_type_name:
                     global_type_name[new_idx] = elem_name

diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
@@ -706,7 +706,7 @@ def _load_type_mix(self, set_name: DPPath):
             real_type = atom_type_mix_
         return real_type
 
-    def build_reidx_to_name_map(self,typemix, set_name: DPPath):
+    def build_reidx_to_name_map(self, typemix, set_name: DPPath):
         type_map = self.type_map
         type_path = set_name / "real_atom_types.npy"
         real_type = type_path.load_numpy().astype(np.int32).reshape([-1, self.natoms])
@@ -715,8 +715,8 @@ def build_reidx_to_name_map(self,typemix, set_name: DPPath):
         N, M = real_type.shape
         for i in range(N):
             for j in range(M):
-                old_val = int(real_type[i, j]) 
-                new_val = int(typemix[i, j])   
+                old_val = int(real_type[i, j])
+                new_val = int(typemix[i, j])
                 reidx_to_name[new_val] = type_map_array[old_val]
         return reidx_to_name