diff options
| -rw-r--r-- | doc/release/1.8.0-notes.rst | 1 | ||||
| -rw-r--r-- | numpy/core/fromnumeric.py | 9 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 4 | ||||
| -rw-r--r-- | numpy/core/tests/test_item_selection.py | 43 |
4 files changed, 54 insertions, 3 deletions
diff --git a/doc/release/1.8.0-notes.rst b/doc/release/1.8.0-notes.rst index c2cd90b5c..52abd7df4 100644 --- a/doc/release/1.8.0-notes.rst +++ b/doc/release/1.8.0-notes.rst @@ -42,6 +42,7 @@ Changes General ------- +The function np.take now allows 0-d arrays as indices. Deprecations ============ diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index d73f1313c..7596e3707 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -57,6 +57,10 @@ def take(a, indices, axis=None, out=None, mode='raise'): The source array. indices : array_like The indices of the values to extract. + + .. versionadded:: 1.8.0 + + Also allow scalars for indices. axis : int, optional The axis over which to select values. By default, the flattened input array is used. @@ -96,6 +100,11 @@ def take(a, indices, axis=None, out=None, mode='raise'): >>> a[indices] array([4, 3, 6]) + If `indices` is not one dimensional, the output also has these dimensions. + + >>> np.take(a, [[0, 1], [2, 3]]) + array([[4, 3], + [5, 7]]) """ try: take = a.take diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index ab8480857..0da921b62 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -44,13 +44,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis, } indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0, NPY_INTP, - 1, 0); + 0, 0); if (indices == NULL) { goto fail; } - - n = m = chunk = 1; nd = PyArray_NDIM(self) + PyArray_NDIM(indices) - 1; for (i = 0; i < nd; i++) { diff --git a/numpy/core/tests/test_item_selection.py b/numpy/core/tests/test_item_selection.py new file mode 100644 index 000000000..f35e04c4f --- /dev/null +++ b/numpy/core/tests/test_item_selection.py @@ -0,0 +1,43 @@ +import numpy as np +from numpy.testing import * +import sys, warnings + +def test_take(): + a = [[1, 2], [3, 4]] + a_str = [['1','2'],['3','4']] + modes = ['raise', 'wrap', 'clip'] + indices = [-1, 4] + index_arrays = [np.empty(0, dtype=np.intp), + np.empty(tuple(), dtype=np.intp), + np.empty((1,1), dtype=np.intp)] + real_indices = {} + real_indices['raise'] = {-1:1, 4:IndexError} + real_indices['wrap'] = {-1:1, 4:0} + real_indices['clip'] = {-1:0, 4:1} + # Currently all types but object, use the same function generation. + # So it should not be necessary to test all, but the code does support it. + types = np.int, np.object + for t in types: + ta = np.array(a if issubclass(t, np.number) else a_str, dtype=t) + tresult = list(ta.T.copy()) + for index_array in index_arrays: + if index_array.size != 0: + tresult[0].shape = (2,) + index_array.shape + tresult[1].shape = (2,) + index_array.shape + for mode in modes: + for index in indices: + real_index = real_indices[mode][index] + if real_index is IndexError and index_array.size != 0: + index_array.put(0, index) + assert_raises(IndexError, ta.take, index_array, + mode=mode, axis=1) + elif index_array.size != 0: + index_array.put(0, index) + res = ta.take(index_array, mode=mode, axis=1) + assert_array_equal(res, tresult[real_index]) + else: + res = ta.take(index_array, mode=mode, axis=1) + assert_(res.shape == (2,) + index_array.shape) + +if __name__ == "__main__": + run_module_suite() |