Method __getitem__

cdslib/core/data/chunked_memmap.py:144–179

Return the ith's sample (in global index) as if calling samples[idx]. Note: 1. Each sample is returned in their native shape (no padding) 2. The sample is returned as ndarray (so int's shape will be (1,))

(self, idx: int)

            return len(self.chunk_idxs)

    def __getitem__(self, idx: int):
        """Return the ith's sample (in global index) as if calling samples[idx].

        Note:
            1. Each sample is returned in their native shape (no padding)
            2. The sample is returned as ndarray (so int's shape will be (1,))
        """
        if isinstance(idx, slice):
            # do not support slicing
            raise NotImplementedError

        if self.chunk_filename_dict is None or self.chunk_shape_dict is None or self.chunk_idxs is None:
            return None

        # get chunk idx of ith sample
        target_chunk_idx = self.chunk_idxs[idx]

        # open memmap
        memmap_filename = self.chunk_filename_dict[target_chunk_idx]
        memmap_shape = self.chunk_shape_dict[target_chunk_idx]
        sample_shape = self.chunk_sample_shapes_dict[target_chunk_idx][self.chunk_local_idxs[idx]]

        # check if memmap_shape or sample_shape contains any 0
        if np.prod(memmap_shape) == 0 or np.prod(sample_shape) == 0:
            out = np.zeros(sample_shape, dtype=self.dtype)
            return out

        offset = self.chunk_offsets_dict[target_chunk_idx][self.chunk_local_idxs[idx]] * self.itemsize  # offset in byte
        arr = np.memmap(
            memmap_filename,
            dtype=self.dtype,
            mode="r",
            offset=offset,
            shape=tuple(sample_shape),
        )
        return np.array(arr)

    def get_shape(self, idx: int) -> T.Sequence[int]:
        """Returns the shape of idx-th sample."""