Skip to content

Commit ce52341

Browse files
danielweshoyer
authored andcommitted
Get 0d slices of ndarrays directly from indexing (#2625)
* Add test to ensure that 0d slices are views * Get 0d slices of ndarrays directly from indexing * Add 0d slice documentation
1 parent a15587d commit ce52341

File tree

3 files changed

+13
-11
lines changed

3 files changed

+13
-11
lines changed

doc/whats-new.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ Enhancements
5656
- :py:meth:`DataArray.resample` and :py:meth:`Dataset.resample` now supports the
5757
``loffset`` kwarg just like Pandas.
5858
By `Deepak Cherian <https://github.com/dcherian>`_
59+
- 0d slices of ndarrays are now obtained directly through indexing, rather than
60+
extracting and wrapping a scalar, avoiding unnecessary copying. By `Daniel
61+
Wennberg <https://github.com/danielwe>`_.
5962

6063
Bug fixes
6164
~~~~~~~~~

xarray/core/indexing.py

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1142,15 +1142,6 @@ def __init__(self, array):
11421142
'Trying to wrap {}'.format(type(array)))
11431143
self.array = array
11441144

1145-
def _ensure_ndarray(self, value):
1146-
# We always want the result of indexing to be a NumPy array. If it's
1147-
# not, then it really should be a 0d array. Doing the coercion here
1148-
# instead of inside variable.as_compatible_data makes it less error
1149-
# prone.
1150-
if not isinstance(value, np.ndarray):
1151-
value = utils.to_0d_array(value)
1152-
return value
1153-
11541145
def _indexing_array_and_key(self, key):
11551146
if isinstance(key, OuterIndexer):
11561147
array = self.array
@@ -1160,7 +1151,10 @@ def _indexing_array_and_key(self, key):
11601151
key = key.tuple
11611152
elif isinstance(key, BasicIndexer):
11621153
array = self.array
1163-
key = key.tuple
1154+
# We want 0d slices rather than scalars. This is achieved by
1155+
# appending an ellipsis (see
1156+
# https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#detailed-notes). # noqa
1157+
key = key.tuple + (Ellipsis,)
11641158
else:
11651159
raise TypeError('unexpected key type: {}'.format(type(key)))
11661160

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

11721166
def __getitem__(self, key):
11731167
array, key = self._indexing_array_and_key(key)
1174-
return self._ensure_ndarray(array[key])
1168+
return array[key]
11751169

11761170
def __setitem__(self, key, value):
11771171
array, key = self._indexing_array_and_key(key)

xarray/tests/test_variable.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1147,6 +1147,11 @@ def test_getitem_basic(self):
11471147
assert v_new.dims == ('x', )
11481148
assert_array_equal(v_new, v._data[:, 1])
11491149

1150+
# test that we obtain a modifiable view when taking a 0d slice
1151+
v_new = v[0, 0]
1152+
v_new[...] += 99
1153+
assert_array_equal(v_new, v._data[0, 0])
1154+
11501155
def test_getitem_with_mask_2d_input(self):
11511156
v = Variable(('x', 'y'), [[0, 1, 2], [3, 4, 5]])
11521157
assert_identical(v._getitem_with_mask(([-1, 0], [1, -1])),

0 commit comments

Comments
 (0)