What's the difference between numpy.take and numpy.choose?
numpy.take(array, indices)
and numpy.choose(indices, array)
behave similarly on 1-D arrays, but this is just coincidence. As pointed out by jonrsharpe, they behave differently on higher-dimensional arrays.
numpy.take
numpy.take(array, indices)
picks out elements from a flattened version of array
. (The resulting elements are of course not necessarily from the same row.)
For example,
numpy.take([[1, 2], [3, 4]], [0, 3])
returns
array([1, 4])
numpy.choose
numpy.choose(indices, set_of_arrays)
plucks out element 0 from array indices[0]
, element 1 from array indices[1]
, element 2 from array indices[2]
, and so on. (Here, array
is actually a set of arrays.)
For example
numpy.choose([0, 1, 0, 0], [[1, 2, 3, 4], [4, 5, 6, 7]])
returns
array([1, 5, 3, 4])
because element 0 comes from array 0, element 1 comes from array 1, element 2 comes from array 0, and element 3 comes from array 0.
More Information
These descriptions are simplified – full descriptions can be found here: numpy.take, numpy.choose. For example, numpy.take
and numpy.choose
behave similarly when indices
and array
are 1-D because numpy.choose
first broadcasts array
.
They are certainly not equivalent, as you can see by giving the same arguments (switched) to both methods:
>>> a = np.array([[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16]])
>>> np.choose([0, 2, 1, 3], a)
array([ 1, 10, 7, 16]) # one from each row
>>> np.take(a, [0, 2, 1, 3])
array([1, 3, 2, 4]) # all from same row
I suggest you read the documentation on take
and choose
.