Try to handle integer arrays in indexing

Author: Illviljan

xarray/namedarray/_array_api/_utils.py

@@ -20,6 +20,7 @@
     _DimsLike2,
     _dtype,
     _IndexKeys,
+    _IndexKeysDims,
     _IndexKeysNoEllipsis,
     _Shape,
 )
@@ -488,7 +489,7 @@ def _check_indexing_dims(original_dims: _Dims, indexing_dims: _Dims) -> None:
         )
 
 
-def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeys) -> _Dims:
+def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeysDims) -> _Dims:
     """
     Get the expected dims when using tuples in __getitem__.
 
@@ -517,15 +518,18 @@ def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeys) -> _Dims:
     >>> _dims_from_tuple_indexing(("x",), (..., 0))
     ()
 
-    Indexing array
+    Indexing array is converted to dims
 
     >>> import numpy as np
-    >>> key = (0, NamedArray((), np.array(0, dtype=int)))
+    >>> key = (0, NamedArray((), np.array(0, dtype=int)).dims)
     >>> _dims_from_tuple_indexing(("x", "y", "z"), key)
     ('z',)
-    >>> key = (0, NamedArray(("y",), np.array([0], dtype=int)))
+    >>> key = (0, NamedArray(("y",), np.array([0], dtype=int)).dims)
     >>> _dims_from_tuple_indexing(("x", "y", "z"), key)
     ('y', 'z')
+    >>> key = (NamedArray(("x",), np.array([0])).dims, 0)
+    >>> _dims_from_tuple_indexing(("x", "y", "z"), key)
+    ('x', 'z')
     """
     key_no_ellipsis = _replace_ellipsis(key, len(dims))
 
@@ -542,14 +546,20 @@ def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeys) -> _Dims:
         elif isinstance(k, slice):
             # Slice retains the dimension.
             j += 1
-        elif isinstance(k, NamedArray):
-            if len(k.dims) == 0:
+        elif isinstance(k, tuple):
+            _dims = k
+            if len(_dims) == 0:
                 # if 0 dim, removes 1 dimension
                 new_dims.pop(j)
-            else:
+            elif len(_dims) == 1:
                 # same size retains the dimension:
-                _check_indexing_dims(dims[i : i + 1], k.dims)
+                _check_indexing_dims(dims[i : i + 1], _dims)
                 j += 1
+                # new_dims.pop(j)
+            else:
+                raise NotImplementedError(
+                    f"What happens here? {key_no_ellipsis=}, {dims=}, {i=}, {k=}"
+                )
 
     return tuple(new_dims)
 

xarray/namedarray/_typing.py

@@ -143,6 +143,7 @@ def imag(self) -> _T_co: ...
 _IndexKey = Union[_IndexKeyNoEllipsis, EllipsisType]
 _IndexKeysNoEllipsis = tuple[_IndexKeyNoEllipsis, ...]
 _IndexKeys = tuple[_IndexKey, ...]  #  tuple[Union[_IndexKey, None], ...]
+_IndexKeysDims = tuple[Union[_IndexKey, _Dims], ...]
 _IndexKeyLike = Union[_IndexKey, _IndexKeys]
 
 _AttrsLike = Union[Mapping[Any, Any], None]

xarray/namedarray/core.py

@@ -601,29 +601,100 @@ def __getitem__(
         Returns self[key].
 
         Some rules:
-        * Integers removes the dim.
-        * Slices and ellipsis maintains same dim.
-        * None adds a dim.
+        * Integer, removes the dim.
+        * Slice and ellipsis, maintains same dim.
+        * None, adds a dim.
+        * integer array, removes the dim if 1 sized(?)
+        * boolean array,
         * tuple follows above but on that specific axis.
 
         Examples
         --------
 
-        1D
+        Basic indexing:
+
+        >>> x = NamedArray(("x",), np.array([3, 5, 7]))
+        >>> x[0]
+        <xarray.NamedArray ()> Size: 8B
+        np.int64(3)
+        >>> x[-1]
+        <xarray.NamedArray ()> Size: 8B
+        np.int64(7)
+        >>> x[0:2]
+        <xarray.NamedArray (x: 2)> Size: 16B
+        array([3, 5])
+        >>> x[None]
+        <xarray.NamedArray (dim_1: 1, x: 3)> Size: 24B
+        array([[3, 5, 7]])
+        >>> x[...]
+        <xarray.NamedArray (x: 3)> Size: 24B
+        array([3, 5, 7])
 
-        >>> x = NamedArray(("x",), np.array([0, 1, 2]))
-        >>> key = NamedArray(("x",), np.array([1, 0, 0], dtype=bool))
-        >>> xm = x[key]
-        >>> xm.dims, xm.shape
-        (('x',), (1,))
+        Indexing with integer array:
 
-        >>> x = NamedArray(("x",), np.array([0, 1, 2]))
-        >>> key = NamedArray(("x",), np.array([0, 0, 0], dtype=bool))
-        >>> xm = x[key]
-        >>> xm.dims, xm.shape
-        (('x',), (0,))
+        >>> key = NamedArray(("x",), np.array([0, 0, 0, -1], dtype=int))
+        >>> x[key]
+        <xarray.NamedArray (x: 4)> Size: 32B
+        array([3, 3, 3, 7])
+
+        Indexing with boolean array:
 
-        Setup a ND array:
+        >>> key = NamedArray(("x",), np.array([1, 0, 1], dtype=bool))
+        >>> x[key]
+        <xarray.NamedArray (x: 2)> Size: 16B
+        array([3, 7])
+        >>> key = NamedArray(("x",), np.array([0, 0, 0], dtype=bool))
+        >>> x[key]
+        <xarray.NamedArray (x: 0)> Size: 0B
+        array([], dtype=int64)
+
+        Multidmensional basic indexing:
+
+        >>> x = NamedArray(("z", "y", "x"), np.array([[[1], [2], [3]], [[4], [5], [6]]]))
+        >>> x[0, 0, 0]
+        <xarray.NamedArray ()> Size: 8B
+        np.int64(1)
+        >>> x[1:2]
+        <xarray.NamedArray (z: 1, y: 3, x: 1)> Size: 24B
+        array([[[4],
+                [5],
+                [6]]])
+        >>> x[:, None, :, :]
+        <xarray.NamedArray (z: 2, dim_3: 1, y: 3, x: 1)> Size: 48B
+        array([[[[1],
+                 [2],
+                 [3]]],
+        <BLANKLINE>
+        <BLANKLINE>
+               [[[4],
+                 [5],
+                 [6]]]])
+        >>> x[..., 0]
+        <xarray.NamedArray (z: 2, y: 3)> Size: 48B
+        array([[1, 2, 3],
+               [4, 5, 6]])
+
+        Multidimensional indexing with integer array:
+
+        >>> x = NamedArray(("z", "y", "x"), np.array([[[1], [2], [3]], [[4], [5], [6]]]))
+        >>> key = NamedArray(("z",), np.array([1, -1]))
+        >>> x[key]
+        <xarray.NamedArray (z: 2, y: 3, x: 1)> Size: 48B
+        array([[[4],
+                [5],
+                [6]],
+        <BLANKLINE>
+               [[4],
+                [5],
+                [6]]])
+        >>> x[key, 0]
+        <xarray.NamedArray (z: 2, x: 1)> Size: 16B
+        array([[4],
+               [4]])
+
+        OLD
+
+        ND array:
 
         >>> x = NamedArray(("x", "y"), np.arange(3 * 4).reshape((3, 4)))
         >>> xm = x[0]
@@ -635,20 +706,22 @@ def __getitem__(
         >>> xm = x[None]
         >>> xm.dims, xm.shape
         (('dim_2', 'x', 'y'), (1, 3, 4))
-
         >>> key = NamedArray(("x", "y"), np.ones((3, 4), dtype=bool))
         >>> xm = x[key]
         >>> xm.dims, xm.shape
         ((('x', 'y'),), (12,))
 
-        0D
+        0D array:
 
         >>> x = NamedArray((), np.array(False, dtype=np.bool))
         >>> key = NamedArray((), np.array(False, dtype=np.bool))
         >>> xm = x[key]
         >>> xm.dims, xm.shape
         (('dim_0',), (0,))
         """
+        from xarray.namedarray._array_api._data_type_functions import (
+            isdtype,
+        )
         from xarray.namedarray._array_api._manipulation_functions import (
             _broadcast_arrays,
         )
@@ -658,16 +731,30 @@ def __getitem__(
             _flatten_dims,
         )
 
-        if isinstance(key, NamedArray):
+        if isinstance(key, NamedArray) and isdtype(key.dtype, "bool"):
             self_new, key_new = _broadcast_arrays(self, key)
+            # self_new, key_new = self, key
             _data = self_new._data[key_new._data]
             _dims = _flatten_dims(_atleast1d_dims(self_new.dims))
+            # _dims = dims
             return self._new(_dims, _data)
-        elif isinstance(key, int | slice | tuple) or key is None or key is ...:
+        elif (
+            isinstance(key, int | slice | tuple | NamedArray)
+            or key is None
+            or key is ...
+        ):
             # TODO: __getitem__ not always available, use expand_dims
-            _data = self._data[key]
             _key_tuple = key if isinstance(key, tuple) else (key,)
-            _dims = _dims_from_tuple_indexing(self.dims, _key_tuple)
+            # _dims = _dims_from_tuple_indexing(self.dims, _key_tuple)
+            _dims = _dims_from_tuple_indexing(
+                self.dims,
+                tuple(k.dims if isinstance(k, NamedArray) else k for k in _key_tuple),
+            )
+
+            _data = self._data[
+                tuple(k._data if isinstance(k, NamedArray) else k for k in _key_tuple)
+            ]
+
             return self._new(_dims, _data)
         else:
             raise NotImplementedError(f"{key=} is not supported")