bug fix for corner case: single frame dataset

Author: OutisLi

deepmd/utils/data.py

@@ -854,26 +854,26 @@ def _load_single_data(
         mmap_obj = self._get_memmap(path)
         # Slice the single frame and make an in-memory copy for modification
         if mmap_obj.ndim == 0:
-            # Scalar array
-            data = mmap_obj.copy().astype(dtype, copy=False)
+            # case: single frame data && non-atomic
+            data = mmap_obj.copy().astype(dtype, copy=False).reshape(1, -1)
         elif mmap_obj.ndim == 1:
-            # Single-frame file (shape: [ndof]); only frame_idx==0 is valid
+            # case: single frame data && atomic
             if frame_idx != 0:
                 raise IndexError(
                     f"frame index {frame_idx} out of range for single-frame file: {path}"
                 )
-            data = mmap_obj.copy().astype(dtype, copy=False)
+            data = mmap_obj.copy().astype(dtype, copy=False).reshape(1, -1)
         else:
-            # Regular [nframes, ...]
-            data = mmap_obj[frame_idx].copy().astype(dtype, copy=False)
+            # case: multi-frame data
+            data = mmap_obj[frame_idx].copy().astype(dtype, copy=False).reshape(1, -1)
 
         try:
             if vv["atomic"]:
                 # Handle type_sel logic
                 if vv["type_sel"] is not None:
                     sel_mask = np.isin(self.atom_type, vv["type_sel"])
 
-                    if mmap_obj.shape[1] == natoms_sel * ndof:
+                    if data.shape[1] == natoms_sel * ndof:
                         if vv["output_natoms_for_type_sel"]:
                             tmp = np.zeros([natoms, ndof], dtype=data.dtype)
                             # sel_mask needs to be applied to the original atom layout
@@ -882,7 +882,7 @@ def _load_single_data(
                         else:  # output is natoms_sel
                             natoms = natoms_sel
                             idx_map = idx_map_sel
-                    elif mmap_obj.shape[1] == natoms * ndof:
+                    elif data.shape[1] == natoms * ndof:
                         data = data.reshape([natoms, ndof])
                         if vv["output_natoms_for_type_sel"]:
                             pass
@@ -892,7 +892,7 @@ def _load_single_data(
                             natoms = natoms_sel
                     else:  # Shape mismatch error
                         raise ValueError(
-                            f"The shape of the data {key} in {set_dir} has width {mmap_obj.shape[1]}, which doesn't match either ({natoms_sel * ndof}) or ({natoms * ndof})"
+                            f"The shape of the data {key} in {set_dir} has width {data.shape[1]}, which doesn't match either ({natoms_sel * ndof}) or ({natoms * ndof})"
                         )
 
                 # Handle special case for Hessian