data_in_memory.py 文件源码

def _handle_slice(self, idx):
        idx = np.index_exp[idx]
        frames, dims = None, None
        if isinstance(idx, (tuple, list)):
            if len(idx) == 1:
                frames, dims = idx[0], slice(None, None, None)
            if len(idx) == 2:
                frames, dims = idx[0], idx[1]
            if len(idx) > 2:
                raise IndexError("Slice was more than two-dimensional, not supported.")

        cumsum = np.cumsum(self._source.trajectory_lengths())
        if not isinstance(frames, (list, np.ndarray)):
            frames = self._get_indices(frames, cumsum[-1])
        dims = self._get_indices(dims, self._source.ndim)

        nframes = len(frames)
        ndims = len(dims)

        data = np.empty((nframes, ndims), dtype=self._source.output_type())

        from chainsaw.clustering import UniformTimeClustering
        for i, x in enumerate(frames):
            traj, idx = UniformTimeClustering._idx_to_traj_idx(x, cumsum)
            data[i, :] = self._source.data[traj][idx, dims]
        return data