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