def test_ellipsis(): shape = 3, 5, 7, 11 arr: np.ndarray = np.arange(np.product(shape)).reshape(shape) test_params = [ (...,), (slice(0, 2), ...), (slice(0, 2), slice(1, 4), ...), (slice(0, 2), slice(1, 4), slice(2, 6), ...), (slice(0, 2), slice(1, 4), slice(2, 6), slice(3, 8), ...), (..., slice(0, 2)), (..., slice(0, 2), slice(1, 4)), (..., slice(0, 2), slice(1, 4), slice(2, 6)), (..., slice(0, 2), slice(1, 4), slice(2, 6), slice(3, 8)), (slice(0, 2),), (slice(0, 2), slice(1, 4)), (slice(0, 2), slice(1, 4), slice(2, 6)), (slice(0, 2), slice(1, 4), slice(2, 6), slice(3, 8)), ] for true_sel in test_params: sel = BasicSelection.from_subscript(shape, true_sel) test_sel = sel.selector() assert np.allclose(arr[true_sel], arr[test_sel]) with pytest.raises(ValueError): BasicSelection.from_subscript(shape, (..., ...))
def test_batch_slice_intersection(): shape = (20, 9, 4) block_shape = (4, 3, 2) # shape = (5, 4, 3) # block_shape = (4, 2, 1) size = np.product(shape) arr: np.ndarray = np.random.random_sample(size).reshape(shape) # Test to ensure selections cover basic array. grid: array_utils.OrderedGrid = array_utils.OrderedGrid( shape=shape, block_shape=block_shape, order=(1, 1, 1) ) asgn_test: np.ndarray = np.random.random_sample(size).reshape(shape) for index in grid.index_iterator(): selection = tuple(grid.slices[index]) asgn_test[selection] = arr[selection] assert np.allclose(arr, asgn_test) # Test intersection of a larger selection with variable step sizes with # basic selections of blocks. arr: np.ndarray = np.arange(size).reshape(shape) slices_set = [] for dim in shape: dim_steps = list(filter(lambda i: i != 0, range(-dim * 2, dim * 2))) dim_slices = list( map( lambda step: slice(0, dim, step) if step > 0 else slice(-1, -dim - 1, step), dim_steps, ) ) slices_set.append(dim_slices) slices_set = list(itertools.product(*slices_set)) pbar = tqdm.tqdm(total=len(slices_set) * np.product(grid.grid_shape)) for slices in slices_set: big_sliced_arr: np.ndarray = arr[tuple(slices)] big_bs: BasicSelection = BasicSelection.from_subscript(shape, tuple(slices)) assert np.allclose(big_sliced_arr, arr[big_bs.selector()]) grid: array_utils.OrderedGrid = array_utils.OrderedGrid( shape=shape, block_shape=block_shape, order=big_bs.order() ) small_sliced_arr: np.ndarray = np.empty(grid.grid_shape, dtype=np.ndarray) for index in grid.index_iterator(): small_slices = tuple(grid.slices[index]) small_bs: BasicSelection = BasicSelection.from_subscript( shape, small_slices ) res_bs = big_bs & small_bs assert arr[res_bs.selector()].shape == res_bs.get_output_shape() small_arr = arr[res_bs.selector()] small_sliced_arr[tuple(index)] = small_arr pbar.update(1) stitched_arr = np.block(small_sliced_arr.tolist()) assert stitched_arr.shape == big_sliced_arr.shape, ( stitched_arr.shape, big_sliced_arr.shape, ) assert np.allclose(big_sliced_arr, stitched_arr)
def assign_references(self, dst_sel: BasicSelection, value): # TODO (hme): This seems overly complicated, but correct. Double check it. # Also, revisit some of the variable names. They will likely # be confusing in the future. # The destination has same block shape as value, # but the destination selection may not have the same shape as value. # May need to broadcast value to destination selection output shape. dst_offset = dst_sel.position().value // np.array( self._source.block_shape, dtype=np.int) # Do we need to broadcast? if (isinstance(value, ArrayView) and (dst_sel.get_output_shape() != value.sel.get_output_shape())): value = value.create() if isinstance(value, ArrayView): # This is the best case. # We don't need to create value to perform the reference copy. # No broadcasting required, so this should be okay. src_offset = value.sel.position().value // np.array( value._source.block_shape, dtype=np.int) src_inflated_shape = dst_sel.get_broadcastable_shape() src_inflated_block_shape = dst_sel.get_broadcastable_block_shape( value.block_shape) src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape, src_inflated_block_shape, self.grid.dtype.__name__) for src_grid_entry_inflated in src_inflated_grid.get_entry_iterator( ): # Num axes in value grid may be too small. dst_grid_entry = tuple( (np.array(src_grid_entry_inflated, dtype=np.int) + dst_offset).tolist()) src_grid_entry = tuple( (np.array(src_grid_entry_inflated, dtype=np.int) + src_offset).tolist()) self._source.blocks[dst_grid_entry] = value._source.blocks[ src_grid_entry].copy() elif isinstance(value, BlockArrayBase): # The value has already been created, so just leverage value's existing grid iterator. if value.shape != dst_sel.get_output_shape(): # Need to broadcast. src_ba: BlockArrayBase = value.broadcast_to( dst_sel.get_output_shape()) else: src_ba: BlockArrayBase = value src_inflated_shape = dst_sel.get_broadcastable_shape() src_inflated_block_shape = dst_sel.get_broadcastable_block_shape( src_ba.block_shape) src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape, src_inflated_block_shape, self.grid.dtype.__name__) src_grid_entry_iterator = list(src_ba.grid.get_entry_iterator()) for src_index, src_grid_entry_inflated in \ enumerate(src_inflated_grid.get_entry_iterator()): src_grid_entry = src_grid_entry_iterator[src_index] dst_grid_entry = tuple( (np.array(src_grid_entry_inflated, dtype=np.int) + dst_offset).tolist()) self._source.blocks[dst_grid_entry] = src_ba.blocks[ src_grid_entry].copy()
def test_stepped_slice_selection(): # Ensure that selections from slices of different step sizes are correct. # Thoroughly tested over a subset of the parameter space. arr: np.ndarray = np.arange(3) size = arr.shape[0] slice_params = list(get_slices(3, 3)) pbar = tqdm.tqdm(total=len(slice_params)) for slice_sel in slice_params: pbar.set_description(str(slice_sel)) pbar.update(1) arr_sel: np.ndarray = sel_module.slice_to_range( slice_sel, arr.shape[0]) assert np.all(arr_sel < size) assert len(arr[slice_sel]) == len(arr[arr_sel]), (slice_sel, arr_sel) assert np.allclose(arr[slice_sel], arr[arr_sel]), (slice_sel, arr_sel) sel: BasicSelection = BasicSelection.from_subscript( arr.shape, (slice_sel, )) assert sel.get_output_shape() == arr[slice_sel].shape if isinstance(sel[0], AxisSlice): if sel[0].step is None: assert sel[0].start >= 0 and sel[0].stop >= 0 else: assert (sel[0].step < 0) == (sel[0].step < 0) == (sel[0].stop < 0) ds_sel = sel.selector() assert np.allclose(arr[slice_sel], arr[ds_sel]), (slice_sel, ds_sel) elif isinstance(sel[0], AxisArray): assert slice_sel.step is not None and slice_sel.step != 1 ds_arr = sel[0].array assert np.allclose(arr[slice_sel], arr[ds_arr]), (slice_sel, ds_arr)
def basic_select(self, subscript: tuple): # No support for subscripts of subscripts. # We create new block arrays to deal with nested subscripts. assert self.shape == self._source.shape assert self.block_shape == self._source.block_shape sel: BasicSelection = BasicSelection.from_subscript(self.shape, subscript) result: ArrayView = ArrayView(self._source, sel) return result
def __init__(self, source, sel: BasicSelection = None, block_shape: tuple = None): self._source: BlockArrayBase = source self._system: System = self._source.system if sel is None: sel = BasicSelection.from_shape(self._source.shape) # Currently, this is all we support. assert len(sel.axes) == len(self._source.shape) self.sel = sel self.shape: tuple = self.sel.get_output_shape() if block_shape is None: block_shape: tuple = array_utils.block_shape_from_subscript(self.sel.selector(), self._source.block_shape) self.block_shape = block_shape assert len(self.block_shape) == len(self.shape) self.grid: ArrayGrid = ArrayGrid(self.shape, self.block_shape, dtype=self._source.dtype.__name__)
def from_subscript(cls, bab, subscript): assert isinstance(bab, BlockArrayBase) return cls(source=bab, sel=BasicSelection.from_subscript(bab.shape, subscript))
def basic_assign_single_step(self, dst_sel: BasicSelection, value): assert isinstance(value, (ArrayView, BlockArrayBase)) dst_ba: BlockArrayBase = self._source dst_sel_arr: np.ndarray = selection.BasicSelection.block_selection( dst_ba.shape, dst_ba.block_shape) dst_sel_clipped: np.ndarray = dst_sel_arr & dst_sel assert dst_sel_clipped.shape == self._source.grid.grid_shape # We create value's block array, in case we need to broadcast. # This may not be necessary, but alternative solutions are extremely tedious. # The result is a block array with replicated blocks, # which match the output shape of dst_sel. if isinstance(value, ArrayView): src_ba_bc: BlockArrayBase = value.create().broadcast_to( dst_sel.get_output_shape()) elif isinstance(value, BlockArrayBase): src_ba_bc: BlockArrayBase = value.broadcast_to( dst_sel.get_output_shape()) else: raise Exception("Unexpected value type %s." % type(value)) # Different lengths occur when an index is used to perform # a selection on an axis. Numpy semantics drops such axes. To allow operations # between source and destination selections, dropped axes are restored with dimension 1 # so that selections are of equal length. # We restore the dropped dimensions of the destination selection, because # the source selection must be broadcastable to the destination selection # for the assignment to be valid. src_inflated_shape = dst_sel.get_broadcastable_shape() # The block shapes need not be equal, but the broadcast source block shape must # match the block shape we obtain below, so that there's a 1-to-1 correspondence # between the grid entries. src_inflated_block_shape = dst_sel.get_broadcastable_block_shape( src_ba_bc.block_shape) src_inflated_grid: ArrayGrid = ArrayGrid(src_inflated_shape, src_inflated_block_shape, self.grid.dtype.__name__) src_sel_arr: np.ndarray = selection.BasicSelection.block_selection( src_inflated_shape, src_inflated_block_shape) src_sel_offset: np.ndarray = src_sel_arr + dst_sel.position() # The enumeration of grid entries is identical if the broadcast source grid and # inflated grid have the same number of blocks. src_grid_entry_iterator = list(src_ba_bc.grid.get_entry_iterator()) for dst_grid_entry in dst_ba.grid.get_entry_iterator(): dst_sel_block: BasicSelection = dst_sel_arr[dst_grid_entry] dst_sel_block_clipped: BasicSelection = dst_sel_clipped[ dst_grid_entry] if dst_sel_block_clipped.is_empty(): continue src_intersection_arr = src_sel_offset & dst_sel_block_clipped src_oids = [] src_params = [] dst_params = [] dst_block: Block = dst_ba.blocks[dst_grid_entry] for src_index, src_grid_entry_bc in enumerate( src_inflated_grid.get_entry_iterator()): src_intersection_block: BasicSelection = src_intersection_arr[ src_grid_entry_bc] if src_intersection_block.is_empty(): continue src_grid_entry = src_grid_entry_iterator[src_index] src_block: Block = src_ba_bc.blocks[src_grid_entry] src_oids.append(src_block.oid) src_sel_block_offset: BasicSelection = src_sel_offset[ src_grid_entry_bc] src_dep_sel_loc = src_intersection_block - src_sel_block_offset.position( ) src_params.append((src_dep_sel_loc.selector(), src_sel_block_offset.get_output_shape(), src_block.transposed)) # We're looking at intersection of dst block and src block, so the # location to which we assign must be offset by dst_sel_block. dst_block_sel_loc: BasicSelection = (src_intersection_block - dst_sel_block.position()) dst_params.append( (dst_block_sel_loc.selector(), dst_block.transposed)) if len(src_oids) == 0: continue dst_block.oid = self._cm.update_block(dst_block.oid, *src_oids, src_params=src_params, dst_params=dst_params, syskwargs={ "grid_entry": dst_block.grid_entry, "grid_shape": dst_block.grid_shape })