Skip to content

Commit

Permalink
Get 0d slices of ndarrays directly from indexing (#2625)
Browse files Browse the repository at this point in the history
* Add test to ensure that 0d slices are views

* Get 0d slices of ndarrays directly from indexing

* Add 0d slice documentation
  • Loading branch information
danielwe authored and shoyer committed Dec 22, 2018
1 parent a15587d commit ce52341
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 11 deletions.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ Enhancements
- :py:meth:`DataArray.resample` and :py:meth:`Dataset.resample` now supports the
``loffset`` kwarg just like Pandas.
By `Deepak Cherian <https://github.com/dcherian>`_
- 0d slices of ndarrays are now obtained directly through indexing, rather than
extracting and wrapping a scalar, avoiding unnecessary copying. By `Daniel
Wennberg <https://github.com/danielwe>`_.

Bug fixes
~~~~~~~~~
Expand Down
16 changes: 5 additions & 11 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -1142,15 +1142,6 @@ def __init__(self, array):
'Trying to wrap {}'.format(type(array)))
self.array = array

def _ensure_ndarray(self, value):
# We always want the result of indexing to be a NumPy array. If it's
# not, then it really should be a 0d array. Doing the coercion here
# instead of inside variable.as_compatible_data makes it less error
# prone.
if not isinstance(value, np.ndarray):
value = utils.to_0d_array(value)
return value

def _indexing_array_and_key(self, key):
if isinstance(key, OuterIndexer):
array = self.array
Expand All @@ -1160,7 +1151,10 @@ def _indexing_array_and_key(self, key):
key = key.tuple
elif isinstance(key, BasicIndexer):
array = self.array
key = key.tuple
# We want 0d slices rather than scalars. This is achieved by
# appending an ellipsis (see
# https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#detailed-notes). # noqa
key = key.tuple + (Ellipsis,)
else:
raise TypeError('unexpected key type: {}'.format(type(key)))

Expand All @@ -1171,7 +1165,7 @@ def transpose(self, order):

def __getitem__(self, key):
array, key = self._indexing_array_and_key(key)
return self._ensure_ndarray(array[key])
return array[key]

def __setitem__(self, key, value):
array, key = self._indexing_array_and_key(key)
Expand Down
5 changes: 5 additions & 0 deletions xarray/tests/test_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1147,6 +1147,11 @@ def test_getitem_basic(self):
assert v_new.dims == ('x', )
assert_array_equal(v_new, v._data[:, 1])

# test that we obtain a modifiable view when taking a 0d slice
v_new = v[0, 0]
v_new[...] += 99
assert_array_equal(v_new, v._data[0, 0])

def test_getitem_with_mask_2d_input(self):
v = Variable(('x', 'y'), [[0, 1, 2], [3, 4, 5]])
assert_identical(v._getitem_with_mask(([-1, 0], [1, -1])),
Expand Down

0 comments on commit ce52341

Please sign in to comment.