Skip to content

More extensive orthogonal indexing in get/setitem #1333

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 10, 2023
Merged

More extensive orthogonal indexing in get/setitem #1333

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1667199
More extensive orthogonal indexing in get/setitem
AndreasAlbertQC Jan 25, 2023
caf433b
Add test cases for indexing with single integer iterable
AndreasAlbertQC Jan 30, 2023
64489cd
Merge branch 'main' into 2023-01-24_orthogonal_indexing
joshmoore Mar 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions docs/release.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,18 @@ Release notes
# to document your changes. On releases it will be
# re-indented so that it does not show up in the notes.

.. _unreleased:
.. _unreleased:

Unreleased
----------
Unreleased
----------

..
# .. warning::
# Pre-release! Use :command:`pip install --pre zarr` to evaluate this release.

* Implement more extensive fallback of getitem/setitem for orthogonal indexing.
By :user:`Andreas Albert <AndreasAlbertQC>` :issue:`1029`.

.. _release_2.14.2:

2.14.2
Expand Down
7 changes: 7 additions & 0 deletions docs/tutorial.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -634,6 +634,13 @@ For convenience, the orthogonal indexing functionality is also available via the
Any combination of integer, slice, 1D integer array and/or 1D Boolean array can
be used for orthogonal indexing.

If the index contains at most one iterable, and otherwise contains only slices and integers,
orthogonal indexing is also available directly on the array:

>>> z = zarr.array(np.arange(15).reshape(3, 5))
>>> all(z.oindex[[0, 2], :] == z[[0, 2], :])
True

Indexing fields in structured arrays
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Expand Down
5 changes: 5 additions & 0 deletions zarr/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
err_too_many_indices,
is_contiguous_selection,
is_pure_fancy_indexing,
is_pure_orthogonal_indexing,
is_scalar,
pop_fields,
)
Expand Down Expand Up @@ -817,6 +818,8 @@ def __getitem__(self, selection):
fields, pure_selection = pop_fields(selection)
if is_pure_fancy_indexing(pure_selection, self.ndim):
result = self.vindex[selection]
elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
result = self.get_orthogonal_selection(pure_selection, fields=fields)
else:
result = self.get_basic_selection(pure_selection, fields=fields)
return result
Expand Down Expand Up @@ -1387,6 +1390,8 @@ def __setitem__(self, selection, value):
fields, pure_selection = pop_fields(selection)
if is_pure_fancy_indexing(pure_selection, self.ndim):
self.vindex[selection] = value
elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
self.set_orthogonal_selection(pure_selection, value, fields=fields)
else:
self.set_basic_selection(pure_selection, value, fields=fields)

Expand Down
20 changes: 20 additions & 0 deletions zarr/indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,26 @@ def is_pure_fancy_indexing(selection, ndim):
)


def is_pure_orthogonal_indexing(selection, ndim):
if not ndim:
return False

# Case 1: Selection is a single iterable of integers
if is_integer_list(selection) or is_integer_array(selection, ndim=1):
return True

# Case two: selection contains either zero or one integer iterables.
# All other selection elements are slices or integers
return (
isinstance(selection, tuple) and len(selection) == ndim and
sum(is_integer_list(elem) or is_integer_array(elem) for elem in selection) <= 1 and
all(
is_integer_list(elem) or is_integer_array(elem)
or isinstance(elem, slice) or isinstance(elem, int) for
elem in selection)
)


def normalize_integer_selection(dim_sel, dim_len):

# normalize type to int
Expand Down
214 changes: 199 additions & 15 deletions zarr/tests/test_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,8 +283,6 @@ def test_get_basic_selection_2d():
for selection in bad_selections:
with pytest.raises(IndexError):
z.get_basic_selection(selection)
with pytest.raises(IndexError):
z[selection]
# check fallback on fancy indexing
fancy_selection = ([0, 1], [0, 1])
np.testing.assert_array_equal(z[fancy_selection], [0, 11])
Expand Down Expand Up @@ -317,14 +315,179 @@ def test_fancy_indexing_fallback_on_get_setitem():
)


def test_fancy_indexing_doesnt_mix_with_slicing():
z = zarr.zeros((20, 20))
with pytest.raises(IndexError):
z[[1, 2, 3], :] = 2
with pytest.raises(IndexError):
np.testing.assert_array_equal(
z[[1, 2, 3], :], 0
@pytest.mark.parametrize("index,expected_result",
[
# Single iterable of integers
(
[0, 1],
[[0, 1, 2],
[3, 4, 5]]
),
# List first, then slice
(
([0, 1], slice(None)),
[[0, 1, 2],
[3, 4, 5]]
),
# List first, then slice
(
([0, 1], slice(1, None)),
[[1, 2],
[4, 5]]
),
# Slice first, then list
(
(slice(0, 2), [0, 2]),
[[0, 2],
[3, 5]]
),
# Slices only
(
(slice(0, 2), slice(0, 2)),
[[0, 1],
[3, 4]]
),
# List with repeated index
(
([1, 0, 1], slice(1, None)),
[[4, 5],
[1, 2],
[4, 5]]
),
# 1D indexing
(
([1, 0, 1]),
[
[3, 4, 5],
[0, 1, 2],
[3, 4, 5]
]
)

])
def test_orthogonal_indexing_fallback_on_getitem_2d(index, expected_result):
"""
Tests the orthogonal indexing fallback on __getitem__ for a 2D matrix.

In addition to checking expected behavior, all indexing
is also checked against numpy.
"""
# [0, 1, 2],
# [3, 4, 5],
# [6, 7, 8]
a = np.arange(9).reshape(3, 3)
z = zarr.array(a)

np.testing.assert_array_equal(z[index], a[index], err_msg="Indexing disagrees with numpy")
np.testing.assert_array_equal(z[index], expected_result)


@pytest.mark.parametrize("index,expected_result",
[
# Single iterable of integers
(
[0, 1],
[[[0, 1, 2],
[3, 4, 5],
[6, 7, 8]],
[[9, 10, 11],
[12, 13, 14],
[15, 16, 17]]]
),
# One slice, two integers
(
(slice(0, 2), 1, 1),
[4, 13]
),
# One integer, two slices
(
(slice(0, 2), 1, slice(0, 2)),
[[3, 4], [12, 13]]
),
# Two slices and a list
(
(slice(0, 2), [1, 2], slice(0, 2)),
[[[3, 4], [6, 7]], [[12, 13], [15, 16]]]
),
])
def test_orthogonal_indexing_fallback_on_getitem_3d(index, expected_result):
"""
Tests the orthogonal indexing fallback on __getitem__ for a 3D matrix.

In addition to checking expected behavior, all indexing
is also checked against numpy.
"""
# [[[ 0, 1, 2],
# [ 3, 4, 5],
# [ 6, 7, 8]],

# [[ 9, 10, 11],
# [12, 13, 14],
# [15, 16, 17]],

# [[18, 19, 20],
# [21, 22, 23],
# [24, 25, 26]]]
a = np.arange(27).reshape(3, 3, 3)
z = zarr.array(a)

np.testing.assert_array_equal(z[index], a[index], err_msg="Indexing disagrees with numpy")
np.testing.assert_array_equal(z[index], expected_result)


@pytest.mark.parametrize(
"index,expected_result",
[
# Single iterable of integers
(
[0, 1],
[
[1, 1, 1],
[1, 1, 1],
[0, 0, 0]
]
),
# List and slice combined
(
([0, 1], slice(1, 3)),
[[0, 1, 1],
[0, 1, 1],
[0, 0, 0]]
),
# Index repetition is ignored on setitem
(
([0, 1, 1, 1, 1, 1, 1], slice(1, 3)),
[[0, 1, 1],
[0, 1, 1],
[0, 0, 0]]
),
# Slice with step
(
([0, 2], slice(None, None, 2)),
[[1, 0, 1],
[0, 0, 0],
[1, 0, 1]]
)
]
)
def test_orthogonal_indexing_fallback_on_setitem_2d(index, expected_result):
"""
Tests the orthogonal indexing fallback on __setitem__ for a 3D matrix.

In addition to checking expected behavior, all indexing
is also checked against numpy.
"""
# Slice + fancy index
a = np.zeros((3, 3))
z = zarr.array(a)
z[index] = 1
a[index] = 1
np.testing.assert_array_equal(
z, expected_result
)
np.testing.assert_array_equal(
z, a, err_msg="Indexing disagrees with numpy"
)


def test_fancy_indexing_doesnt_mix_with_implicit_slicing():
Expand All @@ -335,12 +498,6 @@ def test_fancy_indexing_doesnt_mix_with_implicit_slicing():
np.testing.assert_array_equal(
z2[[1, 2, 3], [1, 2, 3]], 0
)
with pytest.raises(IndexError):
z2[[1, 2, 3]] = 2
with pytest.raises(IndexError):
np.testing.assert_array_equal(
z2[[1, 2, 3]], 0
)
with pytest.raises(IndexError):
z2[..., [1, 2, 3]] = 2
with pytest.raises(IndexError):
Expand Down Expand Up @@ -770,6 +927,33 @@ def test_set_orthogonal_selection_3d():
_test_set_orthogonal_selection_3d(v, a, z, ix0, ix1, ix2)


def test_orthogonal_indexing_fallback_on_get_setitem():
z = zarr.zeros((20, 20))
z[[1, 2, 3], [1, 2, 3]] = 1
np.testing.assert_array_equal(
z[:4, :4],
[
[0, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1],
],
)
np.testing.assert_array_equal(
z[[1, 2, 3], [1, 2, 3]], 1
)
# test broadcasting
np.testing.assert_array_equal(
z[1, [1, 2, 3]], [1, 0, 0]
)
# test 1D fancy indexing
z2 = zarr.zeros(5)
z2[[1, 2, 3]] = 1
np.testing.assert_array_equal(
z2, [0, 1, 1, 1, 0]
)


def _test_get_coordinate_selection(a, z, selection):
expect = a[selection]
actual = z.get_coordinate_selection(selection)
Expand Down