[Numpy-discussion] Use-case for np.choose

Anne Archibald peridot.faceted at gmail.com
Sun Nov 8 02:59:16 EST 2009

2009/11/7 David Goldsmith <d.l.goldsmith at gmail.com>:
> Thanks, Anne.
> On Sat, Nov 7, 2009 at 1:32 PM, Anne Archibald <peridot.faceted at gmail.com>
> wrote:
>> 2009/11/7 David Goldsmith <d.l.goldsmith at gmail.com>:
> <snip>
>> > Also, my experimenting suggests that the index array ('a', the first
>> > argument in the func. sig.) *must* have shape (choices.shape[-1],) -
>> > someone
>> > please let me know ASAP if this is not the case, and please furnish me
>> > w/ a
>> > counterexample because I was unable to generate one myself.
>> It seems like a and each of the choices must have the same shape
> So in essence, at least as it presently functions, the shape of 'a'
> *defines* what the individual choices are within 'choices`, and if 'choices'
> can't be parsed into an integer number of such individual choices, that's
> when an exception is raised?

Um, I don't think so.

Think of it this way: you provide np.choose with a selector array, a,
and a list (not array!) [c0, c1, ..., cM] of choices. You construct an
output array, say r, the same shape as a (no matter how many
dimensions it has). The (i0, i1, ..., iN) element of the output array
is obtained by looking at the (i0, i1, ..., iN) element of a, which
should be an integer no larger than M; say j. Then r[i0, i1, ..., iN]
= cj[i0, i1, ..., iN]. That is, each element of the selector array
determines which of the choice arrays to pull the corresponding
element from.

For example, suppose that you are processing an array C, and have
constructed a selector array A the same shape as C in which a value is
0, 1, or 2 depending on whether the C value is too small, okay, or too
big respectively. Then you might do something like:

C = np.choose(A, [-inf, C, inf])

This is something you might want to do no matter what shape A and C
have. It's important not to require that the choices be an array of
choices, because they often have quite different shapes (here, two are
scalars) and it would be wasteful to broadcast them up to the same
shape as C, just to stack them.


More information about the NumPy-Discussion mailing list