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 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 from_subscript(cls, bab, subscript):
assert isinstance(bab, BlockArrayBase)
return cls(source=bab,
sel=BasicSelection.from_subscript(bab.shape, subscript))
