numpy.choose — NumPy v2.3.dev0 Manual (original) (raw)

numpy.choose(a, choices, out=None, mode='raise')[source]#

Construct an array from an index array and a list of arrays to choose from.

First of all, if confused or uncertain, definitely look at the Examples - in its full generality, this function is less simple than it might seem from the following code description:

np.choose(a,c) == np.array([c[a[I]][I] for I in np.ndindex(a.shape)])

But this omits some subtleties. Here is a fully general summary:

Given an “index” array (a) of integers and a sequence of n arrays (choices), a and each choice array are first broadcast, as necessary, to arrays of a common shape; calling these Ba and Bchoices[i], i = 0,…,n-1 we have that, necessarily, Ba.shape == Bchoices[i].shapefor each i. Then, a new array with shape Ba.shape is created as follows:

Parameters:

aint array

This array must contain integers in [0, n-1], where n is the number of choices, unless mode=wrap or mode=clip, in which cases any integers are permissible.

choicessequence of arrays

Choice arrays. a and all of the choices must be broadcastable to the same shape. If choices is itself an array (not recommended), then its outermost dimension (i.e., the one corresponding tochoices.shape[0]) is taken as defining the “sequence”.

outarray, optional

If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. Note that out is always buffered if mode='raise'; use other modes for better performance.

mode{‘raise’ (default), ‘wrap’, ‘clip’}, optional

Specifies how indices outside [0, n-1] will be treated:

Returns:

merged_arrayarray

The merged result.

Raises:

ValueError: shape mismatch

If a and each choice array are not all broadcastable to the same shape.

Notes

To reduce the chance of misinterpretation, even though the following “abuse” is nominally supported, choices should neither be, nor be thought of as, a single array, i.e., the outermost sequence-like container should be either a list or a tuple.

Examples

import numpy as np choices = [[0, 1, 2, 3], [10, 11, 12, 13], ... [20, 21, 22, 23], [30, 31, 32, 33]] np.choose([2, 3, 1, 0], choices ... # the first element of the result will be the first element of the ... # third (2+1) "array" in choices, namely, 20; the second element ... # will be the second element of the fourth (3+1) choice array, i.e., ... # 31, etc. ... ) array([20, 31, 12, 3]) np.choose([2, 4, 1, 0], choices, mode='clip') # 4 goes to 3 (4-1) array([20, 31, 12, 3])

because there are 4 choice arrays

np.choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4) array([20, 1, 12, 3])

i.e., 0

A couple examples illustrating how choose broadcasts:

a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]] choices = [-10, 10] np.choose(a, choices) array([[ 10, -10, 10], [-10, 10, -10], [ 10, -10, 10]])

With thanks to Anne Archibald

a = np.array([0, 1]).reshape((2,1,1)) c1 = np.array([1, 2, 3]).reshape((1,3,1)) c2 = np.array([-1, -2, -3, -4, -5]).reshape((1,1,5)) np.choose(a, (c1, c2)) # result is 2x3x5, res[0,:,:]=c1, res[1,:,:]=c2 array([[[ 1, 1, 1, 1, 1], [ 2, 2, 2, 2, 2], [ 3, 3, 3, 3, 3]], [[-1, -2, -3, -4, -5], [-1, -2, -3, -4, -5], [-1, -2, -3, -4, -5]]])