Skip to content

Commit

Permalink
indexing are always 1d?
Browse files Browse the repository at this point in the history
  • Loading branch information
Illviljan committed Sep 24, 2024
1 parent 6cf9fae commit 18400d5
Showing 1 changed file with 18 additions and 9 deletions.
27 changes: 18 additions & 9 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -505,9 +505,9 @@ def __ge__(self, other: int | float | NamedArray, /):

def __getitem__(self, key: _IndexKeyLike | NamedArray) -> NamedArray:
"""
sdfd
Returns self[key].
Some rules
Some rules:
* Integers removes the dim.
* Slices and ellipsis maintains same dim.
* None adds a dim.
Expand All @@ -519,14 +519,14 @@ def __getitem__(self, key: _IndexKeyLike | NamedArray) -> NamedArray:
1D
>>> x = NamedArray(("x",), np.array([0, 1, 2]))
>>> mask = NamedArray(("x",), np.array([1, 0, 0], dtype=bool))
>>> xm = x[mask]
>>> key = NamedArray(("x",), np.array([1, 0, 0], dtype=bool))
>>> xm = x[key]
>>> xm.dims, xm.shape
(('x',), (1,))
>>> x = NamedArray(("x",), np.array([0, 1, 2]))
>>> mask = NamedArray(("x",), np.array([0, 0, 0], dtype=bool))
>>> xm = x[mask]
>>> key = NamedArray(("x",), np.array([0, 0, 0], dtype=bool))
>>> xm = x[key]
>>> xm.dims, xm.shape
(('x',), (0,))
Expand All @@ -543,23 +543,32 @@ def __getitem__(self, key: _IndexKeyLike | NamedArray) -> NamedArray:
>>> xm.dims, xm.shape
(('dim_2', 'x', 'y'), (1, 3, 4))
>>> mask = NamedArray(("x", "y"), np.ones((3, 4), dtype=bool))
>>> xm = x[mask]
>>> key = NamedArray(("x", "y"), np.ones((3, 4), dtype=bool))
>>> xm = x[key]
>>> xm.dims, xm.shape
((('x', 'y'),), (12,))
0D
>>> 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._manipulation_functions import (
_broadcast_arrays,
)
from xarray.namedarray._array_api._utils import (
_atleast1d_dims,
_dims_from_tuple_indexing,
_flatten_dims,
)

if isinstance(key, NamedArray):
self_new, key_new = _broadcast_arrays(self, key)
_data = self_new._data[key_new._data]
_dims = _flatten_dims(self_new.dims)
_dims = _flatten_dims(_atleast1d_dims(self_new.dims))
return self._new(_dims, _data)
# elif isinstance(key, int):
# return self._new(self.dims[1:], self._data[key])
Expand Down

0 comments on commit 18400d5

Please sign in to comment.