def test_orthogonal_indexing_edge_cases():
a = np.arange(6).reshape(1, 2, 3)
z = zarr.create(shape=a.shape, chunks=(1, 2, 3), dtype=a.dtype)
z[:] = a
expect = oindex(a, (0, slice(None), [0, 1, 2]))
actual = z.oindex[0, :, [0, 1, 2]]
assert_array_equal(expect, actual)
expect = oindex(a, (0, slice(None), [True, True, True]))
actual = z.oindex[0, :, [True, True, True]]
assert_array_equal(expect, actual)
def _test_set_orthogonal_selection(v, a, z, selection):
for value in 42, oindex(v, selection), oindex(v, selection).tolist():
if isinstance(value, list) and value == []:
# skip these cases as cannot preserve all dimensions
continue
# setup expectation
a[:] = 0
oindex_set(a, selection, value)
# long-form API
z[:] = 0
z.set_orthogonal_selection(selection, value)
assert_array_equal(a, z[:])
# short-form API
z[:] = 0
z.oindex[selection] = value
assert_array_equal(a, z[:])
def _test_get_orthogonal_selection(a, z, selection):
expect = oindex(a, selection)
actual = z.get_orthogonal_selection(selection)
assert_array_equal(expect, actual)
actual = z.oindex[selection]
assert_array_equal(expect, actual)
def test_get_selection_out():
# basic selections
a = np.arange(1050)
z = zarr.create(shape=1050, chunks=100, dtype=a.dtype)
z[:] = a
selections = [
slice(50, 150),
slice(0, 1050),
slice(1, 2),
]
for selection in selections:
expect = a[selection]
out = zarr.create(shape=expect.shape, chunks=10, dtype=expect.dtype, fill_value=0)
z.get_basic_selection(selection, out=out)
assert_array_equal(expect, out[:])
with pytest.raises(TypeError):
z.get_basic_selection(Ellipsis, out=[])
# orthogonal selections
a = np.arange(10000, dtype=int).reshape(1000, 10)
z = zarr.create(shape=a.shape, chunks=(300, 3), dtype=a.dtype)
z[:] = a
np.random.seed(42)
# test with different degrees of sparseness
for p in 0.5, 0.1, 0.01:
ix0 = np.random.binomial(1, p, size=a.shape[0]).astype(bool)
ix1 = np.random.binomial(1, .5, size=a.shape[1]).astype(bool)
selections = [
# index both axes with array
(ix0, ix1),
# mixed indexing with array / slice
(ix0, slice(1, 5)),
(slice(250, 350), ix1),
# mixed indexing with array / int
(ix0, 4),
(42, ix1),
# mixed int array / bool array
(ix0, np.nonzero(ix1)[0]),
(np.nonzero(ix0)[0], ix1),
]
for selection in selections:
expect = oindex(a, selection)
# out = zarr.create(shape=expect.shape, chunks=10, dtype=expect.dtype,
# fill_value=0)
out = np.zeros(expect.shape, dtype=expect.dtype)
z.get_orthogonal_selection(selection, out=out)
assert_array_equal(expect, out[:])
# coordinate selections
a = np.arange(10000, dtype=int).reshape(1000, 10)
z = zarr.create(shape=a.shape, chunks=(300, 3), dtype=a.dtype)
z[:] = a
np.random.seed(42)
# test with different degrees of sparseness
for p in 0.5, 0.1, 0.01:
n = int(a.size * p)
ix0 = np.random.choice(a.shape[0], size=n, replace=True)
ix1 = np.random.choice(a.shape[1], size=n, replace=True)
selections = [
# index both axes with array
(ix0, ix1),
# mixed indexing with array / int
(ix0, 4),
(42, ix1),
]
for selection in selections:
expect = a[selection]
out = np.zeros(expect.shape, dtype=expect.dtype)
z.get_coordinate_selection(selection, out=out)
assert_array_equal(expect, out[:])
