def test_get_slice_stack_nav_only():
slice_ = Slice(
origin=(0, 0, 0, 0),
shape=Shape((1, 1, 1, 1), sig_dims=2)
)
data = np.arange(4 * 4 * 4 * 4).reshape(4, 4, 4, 4)
assert slice_.get(data, nav_only=True).shape[0:2] == tuple(slice_.shape.nav)
assert np.all(slice_.get(data, nav_only=True) == data[0:1, 0:1])
def test_get_slice_stack_signal_only():
slice_ = Slice(
origin=(0, 0, 0, 0),
shape=Shape((1, 1, 1, 1), sig_dims=2)
)
data = np.arange(4 * 4 * 4 * 4).reshape(4, 4, 4, 4)
assert slice_.get(data, sig_only=True).shape[2:4] == tuple(slice_.shape.sig)
assert np.all(slice_.get(data, sig_only=True) == data[..., 0:1, 0:1])
def test_get_slice_2():
slice_ = Slice(
origin=(1, 1, 1, 1),
shape=Shape((1, 1, 2, 2), sig_dims=2),
)
data = np.arange(4 * 4 * 4 * 4).reshape(4, 4, 4, 4)
assert slice_.get(data).shape == tuple(slice_.shape)
assert np.all(slice_.get(data) == np.array([[[
[85, 86],
[89, 90],
]]]))
def apply(self, data: np.ndarray, tile_slice: Slice):
"""
Apply corrections in-place to `data`, cropping the
correction data to the `tile_slice`.
"""
dark_frame = self.get_dark_frame()
gain_map = self.get_gain_map()
excluded_pixels = self.get_excluded_pixels()
if not self.have_corrections():
return
sig_slice = tile_slice.get(sig_only=True)
if dark_frame is not None:
dark_frame = dark_frame[sig_slice]
if gain_map is not None:
gain_map = gain_map[sig_slice]
if excluded_pixels is not None:
excluded_pixels = excluded_pixels[sig_slice]
excluded_pixels = excluded_pixels.coords
correct(
buffer=data,
dark_image=dark_frame,
gain_map=gain_map,
excluded_pixels=excluded_pixels,
inplace=True,
sig_shape=tuple(tile_slice.shape.sig),
)
def _roi_to_nd_indices(roi, part_slice: Slice):
"""
Helper function to calculate indices from roi mask
Parameters
----------
roi : numpy.ndarray of type bool, matching the navigation shape of the dataset
part_slice
Slice indicating what part of the roi to operate on, for example,
corresponding to a partition.
"""
roi_slice = roi[part_slice.get(nav_only=True)]
nav_dims = part_slice.shape.nav.dims
total = 0
frames_in_roi = np.count_nonzero(roi)
for idx, value in np.ndenumerate(roi_slice):
if not value:
continue
yield tuple(a + b for a, b in zip(idx, part_slice.origin[:nav_dims]))
# early exit: we know we don't have more frames in the roi
total += 1
if total == frames_in_roi:
break
def apply(self, data: np.ndarray, tile_slice: Slice):
"""
Apply corrections in-place to `data`, cropping the
correction data to the `tile_slice`.
"""
dark_frame = self.get_dark_frame()
gain_map = self.get_gain_map()
if not self.have_corrections():
return
sig_slice = tile_slice.get(sig_only=True)
if dark_frame is not None:
dark_frame = dark_frame[sig_slice]
if gain_map is not None:
gain_map = gain_map[sig_slice]
correct(buffer=data,
dark_image=dark_frame,
gain_map=gain_map,
repair_descriptor=self.repair_descriptor(
tile_slice.discard_nav()),
inplace=True,
sig_shape=tuple(tile_slice.shape.sig),
allow_empty=self._allow_empty)
def test_get_signal_only():
s = Slice(origin=(0, 0, 0, 0), shape=Shape((1, 1, 1, 1), sig_dims=2))
assert s.get(sig_only=True) == (
slice(0, 1),
slice(0, 1),
)
def test_get():
s = Slice(origin=(0, 0, 0, 0), shape=Shape((1, 1, 1, 1), sig_dims=2))
assert s.get() == (
slice(0, 1),
slice(0, 1),
slice(0, 1),
slice(0, 1),
)
def test_get_slice_1():
slice_ = Slice(
origin=(0, 0, 0, 0),
shape=Shape((4, 4, 4, 4), sig_dims=2),
)
assert slice_.get() == (
slice(0, 4),
slice(0, 4),
slice(0, 4),
slice(0, 4),
)
def _get_tiles_roi(self, tiling_scheme, fileset, read_ranges, roi, sync_offset):
ds_sig_shape = tiling_scheme.dataset_shape.sig
sig_dims = tiling_scheme.shape.sig.dims
slices, ranges, scheme_indices = read_ranges
fh = fileset[0]
memmap = fh.mmap().reshape((fh.num_frames,) + tuple(ds_sig_shape))
if sync_offset == 0:
flat_roi = roi.reshape((-1,))
elif sync_offset > 0:
flat_roi = np.full(roi.reshape((-1,)).shape, False)
flat_roi[:sync_offset] = False
flat_roi[sync_offset:] = roi.reshape((-1,))[:-sync_offset]
else:
flat_roi = np.full(roi.reshape((-1,)).shape, False)
flat_roi[sync_offset:] = False
flat_roi[:sync_offset] = roi.reshape((-1,))[-sync_offset:]
data_w_roi = memmap[flat_roi]
for idx in range(slices.shape[0]):
origin, shape = slices[idx]
scheme_idx = scheme_indices[idx]
tile_slice = Slice(
origin=origin,
shape=Shape(shape, sig_dims=sig_dims)
)
if sync_offset >= 0:
data_slice = tile_slice.get()
else:
frames_to_skip = np.count_nonzero(roi.reshape((-1,))[:abs(sync_offset)])
data_slice = Slice(
origin=(origin[0] - frames_to_skip,) + tuple(origin[-sig_dims:]),
shape=Shape(shape, sig_dims=sig_dims)
)
data_slice = data_slice.get()
data = data_w_roi[data_slice]
yield DataTile(
data,
tile_slice=tile_slice,
scheme_idx=scheme_idx,
)
def _get_tiles_mmap(self, crop_to, full_frames, dest_dtype):
start_at_frame = self._start_frame
num_frames = self._num_frames
sig_shape = self.meta.shape.sig
sig_origin = tuple([0] * len(sig_shape))
if dest_dtype != self.meta.raw_dtype:
raise ValueError(
"using mmap with dtype conversion is not efficient")
if crop_to is not None:
sig_origin = tuple(crop_to.origin[-sig_shape.dims:])
sig_shape = crop_to.shape.sig
if full_frames:
sig_shape = self.meta.shape.sig
fileset = self._fileset.get_for_range(
start=start_at_frame,
stop=start_at_frame + num_frames,
)
with self._fileset as fileset:
for f in fileset:
# global start/stop indices:
start = max(f.start_idx, self._start_frame)
stop = min(f.end_idx, self._start_frame + num_frames)
tile_slice = Slice(origin=(start, ) + sig_origin,
shape=Shape(
(stop - start, ) + tuple(sig_shape),
sig_dims=sig_shape.dims))
arr = f.mmap()
# limit to this partition (translate to file-local coords)
arr = arr[start - f.start_idx:stop - f.start_idx]
if crop_to is not None:
intersection = tile_slice.intersection_with(crop_to)
if intersection.is_null():
continue
# crop to, signal part:
arr = arr[(..., ) + tile_slice.get(sig_only=True)]
yield DataTile(data=arr, tile_slice=tile_slice)
def _get_tiles_roi(self, tiling_scheme, fileset, read_ranges, roi):
ds_sig_shape = tiling_scheme.dataset_shape.sig
sig_dims = tiling_scheme.shape.sig.dims
slices, ranges, scheme_indices = read_ranges
fh = fileset[0]
memmap = fh.mmap().reshape((fh.num_frames, ) + tuple(ds_sig_shape))
data_w_roi = memmap[roi.reshape((-1, ))]
for idx in range(slices.shape[0]):
origin, shape = slices[idx]
scheme_idx = scheme_indices[idx]
tile_slice = Slice(origin=origin,
shape=Shape(shape, sig_dims=sig_dims))
data_slice = tile_slice.get()
data = data_w_roi[data_slice]
yield DataTile(
data,
tile_slice=tile_slice,
scheme_idx=scheme_idx,
)
def _get_tiles_straight(self, tiling_scheme, fileset, read_ranges):
"""
Parameters
----------
fileset : FileSet
To ensure best performance, should be limited to the files
that are part of the current partition (otherwise we will
spend more time finding the right file for a given frame
index)
read_ranges : Tuple[np.ndarray, np.ndarray]
As returned by `get_read_ranges`
"""
ds_sig_shape = tiling_scheme.dataset_shape.sig
sig_dims = tiling_scheme.shape.sig.dims
slices, ranges, scheme_indices = read_ranges
for idx in range(slices.shape[0]):
origin, shape = slices[idx]
tile_ranges = ranges[idx]
scheme_idx = scheme_indices[idx]
# FIXME: for straight mmap, read_ranges must not contain tiles
# that span multiple files!
file_idx = tile_ranges[0][0]
fh = fileset[file_idx]
memmap = fh.mmap().reshape((fh.num_frames, ) + tuple(ds_sig_shape))
tile_slice = Slice(origin=origin,
shape=Shape(shape, sig_dims=sig_dims))
data_slice = tile_slice.get()
data = memmap[data_slice]
yield DataTile(
data,
tile_slice=tile_slice,
scheme_idx=scheme_idx,
)
def open_file(self):
f = np.memmap(self._path, dtype=self._meta.dtype, mode='r',
shape=self._scan_size + self._detector_size_raw)
ds_slice = Slice(origin=(0, 0, 0, 0), shape=self._meta.shape)
return f[ds_slice.get()] # crop off the two extra rows
