What's the difference between numpy.take and numpy.choose?

Question

It seems that numpy.take(array, indices) and numpy.choose(indices, array) return the same thing: a subset of array indexed by indice

Answer 1

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.