diff --git a/docs/release.rst b/docs/release.rst index a6c32100ba..f056f621bf 100644 --- a/docs/release.rst +++ b/docs/release.rst @@ -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 ` :issue:`1029`. + .. _release_2.14.2: 2.14.2 diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 43e42faf6b..0f2e1c7345 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/zarr/core.py b/zarr/core.py index b9db6cb2c8..521de80e17 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -28,6 +28,7 @@ err_too_many_indices, is_contiguous_selection, is_pure_fancy_indexing, + is_pure_orthogonal_indexing, is_scalar, pop_fields, ) @@ -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 @@ -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) diff --git a/zarr/indexing.py b/zarr/indexing.py index 268b487105..2f8144fd08 100644 --- a/zarr/indexing.py +++ b/zarr/indexing.py @@ -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 diff --git a/zarr/tests/test_indexing.py b/zarr/tests/test_indexing.py index 5c4c580636..f5f57be010 100644 --- a/zarr/tests/test_indexing.py +++ b/zarr/tests/test_indexing.py @@ -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]) @@ -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(): @@ -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): @@ -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)