Source code for libertem.io.dataset.base.coordinates

import numpy as np

from libertem.common.math import prod
from libertem.common import Shape, Slice
from libertem.io.dataset.base import _roi_to_nd_indices


[docs]def get_coordinates(slice_: Slice, ds_shape: Shape, roi=None) -> np.ndarray: """ Returns `numpy.ndarray` of coordinates that correspond to the frames in the actual navigation space which are part of the current tile or partition. Parameters ---------- slice_: Slice Describes the location within the dataset with navigation dimension flattened and reduced to the ROI. ds_shape: Shape The original shape of the whole dataset, not influenced by the ROI roi: numpy.ndarray, optional Array of type bool, matching the navigation shape of the dataset """ o = slice_.origin s = slice_.shape sig_dims = s.sig.dims start_idx = o[0] end_idx = o[0] + s[0] nav_shape = ds_shape[:-sig_dims] if roi is None: flat_nav_shape = tuple((int(prod(nav_shape)),)) coordinates = np.stack( np.unravel_index( np.ravel_multi_index([np.arange(start_idx, end_idx)], flat_nav_shape), nav_shape ), axis=1 ) else: ds_shape = Shape(ds_shape, sig_dims=sig_dims) ds_slice = Slice(origin=[0] * len(ds_shape), shape=ds_shape) roi = roi.reshape(nav_shape) indices = _roi_to_nd_indices(roi, ds_slice) coordinates = np.array(list(indices)[start_idx:end_idx]) return coordinates