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
})
