Alternative for numpy.choose that allows an arbitrary or at least more than 32 arguments?

Question

With my code I'm running into an issue where the numpy.choose methods does not accept all the arguments since it is limited by NPY_MAXARGS (=32). Is there an alternative available, that allows an arbitrary number of argument arrays or at least more than 32 that is as fast as numpy.choose?

choices = [np.arange(0,100)]*100
selection = [0] * 100
np.choose(selection, choices)

>> ValueError: Need between 2 and (32) array objects (inclusive).

Any help would be appreciated... :)

Answer 1

The indices can be given as lists. Assuming that selections has the same length as choices:

b = numpy.array(choices)
result = b[range(len(selections)), selections]

will give the value in choices specified by the index in selections. See it in action:

numpy.random.seed(1)
b = numpy.random.randint(0,100,(5,10))
>>> array([[37, 12, 72,  9, 75,  5, 79, 64, 16,  1],
          [76, 71,  6, 25, 50, 20, 18, 84, 11, 28],
          [29, 14, 50, 68, 87, 87, 94, 96, 86, 13],
          [ 9,  7, 63, 61, 22, 57,  1,  0, 60, 81],
          [ 8, 88, 13, 47, 72, 30, 71,  3, 70, 21]])

selections = numpy.random.randint(0,10,5)
>>> array([1, 9, 3, 4, 8])

result = b[range(len(selections)),selections]
>>>> array([12, 28, 68, 22, 70])

Answer 2

choose has the 32 object limit because it broadcasts the arrays together. Consider the error messages for these two actions:

In [982]: np.arange(33).choose(np.ones((33,33)))
...
ValueError: Need at least 1 and at most 32 array objects.
In [983]: np.broadcast(*range(33))
...
ValueError: Need at least 1 and at most 32 array objects.

An example exploiting that broadcasting, picking values from a 2d array, 1d and scalar.

In [998]: np.diag([2,1,0]).choose((np.arange(9).reshape(3,3), 0,[.1,.2,.3]))
Out[998]: 
array([[ 0.1,  1. ,  2. ],
       [ 3. ,  0. ,  5. ],
       [ 6. ,  7. ,  8. ]])

As @Benjamin shows np.choose can be used to select items from the successive columns of a 2d array - provided there aren't more than 32 columns

In [1002]: M=np.arange(9).reshape(3,3)
In [1003]: np.array([2,0,1]).choose(M)
Out[1003]: array([6, 1, 5])
In [1004]: M[[2,0,1],[0,1,2]]
Out[1004]: array([6, 1, 5])

It was in just such a context that I recall first seeing this 32 array limit to choose, and one of the few where I've seen choose used in an answer.

It is a compiled function, PyArray_Choose and array_choose

https://github.com/numpy/numpy/blob/0b2e590ec18942f8f149ab2306b80da86b04eaeb/numpy/core/src/multiarray/item_selection.c

https://github.com/numpy/numpy/blob/945c308e96fb815729e8f8aeb0ad6b39b8bdf84a/numpy/core/src/multiarray/methods.c

I don't see any uses of this function in other compiled numpy code. And apart from testing, little use in rest of numpy.