Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: Allow 0-d indexes in np.take #2725

Merged
merged 1 commit into from
Jan 23, 2013
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/release/1.8.0-notes.rst
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ Changes
General
-------

The function np.take now allows 0-d arrays as indices.

Deprecations
============
Expand Down
9 changes: 9 additions & 0 deletions numpy/core/fromnumeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
4 changes: 1 addition & 3 deletions numpy/core/src/multiarray/item_selection.c
Original file line number Diff line number Diff line change
Expand Up @@ -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++) {
Expand Down
43 changes: 43 additions & 0 deletions numpy/core/tests/test_item_selection.py
Original file line number Diff line number Diff line change
@@ -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()