nickpick - 9 months ago 101

Python Question

I have a numpy problem with choose. I would like to choose certain indices as decribed in array a from array b.

`a`

Out[54]:

array([[3, 2, 2],

[0, 0, 2]], dtype=int64)

b

Out[55]:

array([[[ 6., 1., 8., 9., 3., 8., 5.],

[ 6., 1., 5., 8., 2., 2., 10.],

[ 6., 1., 1., 0., 9., 3., 6.]],

[[ 11., 3., 8., 9., 3., 8., 5.],

[ 12., 7., 5., 8., 2., 2., 10.],

[ 8., 9., 1., 0., 9., 3., 6.]]])

np.choose(a,b)

ValueError: shape mismatch: objects cannot be broadcast to a single shap

In the documentation of numpy it says:

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 to choices.shape[0]) is taken as defining the “sequence”.

I see that it is not recommended to choose from a ndarray but is there an elegant numpy way to get this to work anyway? Any suggestions are appreciated.

Expected output is:

`[[9,5,1], [11,12,1]]`

Answer

It looks you want to use `choose`

to select values from the dimension of length 7 in `b`

(which is sized (2,3,7)). Your choosing array `a`

will work for this, but only if the sequence dimension is the outermost dimension (as you quoted). The outermost dimension in Numpy is the *first* dimension. What you need to do, then, is roll `b`

so that it has dimensions (7,2,3).

```
np.choose(a, np.rollaxis(b, 2, 0))
```