numpy ndarray切片与数组

Question

I am having trouble understanding the reason behind a numpy broadcasting error when trying to slice an ndarray along separate dimensions using multiple slicing arrays. I am trying to slice data ndarray (100, 306, 481) along the first and second dimensions using an index array picks eg, np.arange(2, 306, 3) and a boolean array mask , where mask.shape is (481,) of which 361 elements are True .

data[:, picks, mask] returns

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (102,) (361,)

However data[:, :, mask] , data[:, picks, :] , and data[:, :10, mask] work as expected.

How does broadcasting work in this case? and what is a pythonic way of doing this?

Answer 1

So

data[:, :, mask]  => (100, 306, 361) 
data[:, :10, mask] => (100, 10, 361)
data[:, picks, :] => (100, 102, 481)

If picks had (361,) elements then

data[:, picks, mask] => (100, 361)  # I think :)

Think of picks matching np.where(mask)

But to index in separate dimensions, picks has to be a column vector, so (102,1) broadcasts with (1, 361) to produce a (102,361) selection

data[:, picks[:,None], mask] => (100, 102, 361) # again I need to test

So creating some test arrays:

In [253]: data=np.ones((100,306,481))
In [254]: picks=np.arange(2,306,3)
In [255]: mask=np.zeros(481,bool)
In [256]: mask[:361]=True
In [257]: data[:, picks[:,None],mask].shape
Out[257]: (100, 102, 361)

the arange could be replaced with a slice

In [259]: data[:, 2::3, mask].shape
Out[259]: (100, 102, 361)

ix_ is handy in this case

In [268]: I,J=np.ix_(picks,mask)
In [269]: I.shape
Out[269]: (102, 1)
In [270]: J.shape
Out[270]: (1, 361)
In [271]: data[:,I,J].shape
Out[271]: (100, 102, 361)