Numpy multi-dimensional slicing with multiple boolean arrays

Question

I'm trying to use individual 1-dimensional boolean arrays to slice a multi-dimension array. For some reason, this code doesn't work:

>>> a = np.ones((100, 200, 300, 2))
>>> a.shape
(100, 200, 300, 2)
>>> m1 = np.asarray([True]*200)
>>> m2 = np.asarray([True]*300)
>>> m2[-1] = False
>>> a[:,m1,m2,:]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (200,) (299,) 
>>> m2 = np.asarray([True]*300) # try again with all 300 dimensions True
>>> a[:,m1,m2,:]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (200,) (300,) 

But this works just fine:

>>> a = np.asarray([[[1, 2], [3, 4], [5, 6]], [[11, 12], [13, 14], [15, 16]]])
>>> a.shape
(2, 3, 2)
>>> m1 = np.asarray([True, False, True])
>>> m2 = np.asarray([True, False])
>>> a[:,m1,m2]
array([[ 1,  5],
       [11, 15]])

Any idea of what I might be doing wrong in the first example?

Answer 1

This might be a simple workaround:

a[:,m1,:,:][:,:,m2,:]

Answer 2

Short answer: The number of True elements in m1 and m2 must match, unless one of them has only one True term.

Also distinguish between 'diagonal' indexing and 'rectangular' indexing. This is about indexing, not slicing. The dimensions with : are just along for the ride.

Initial ideas

I can get your first case working with:

In [137]: a=np.ones((100,200,300,2))

In [138]: m1=np.ones((200,),bool)    
In [139]: m2=np.ones((300,),bool)
In [140]: m2[-1]=False

In [141]: I,J=np.ix_(m1,m2)

In [142]: a[:,I,J,:].shape
Out[142]: (100, 200, 299, 2)

np.ix_ turns the 2 boolean arrays into broadcastable index arrays

In [143]: I.shape
Out[143]: (200, 1)
In [144]: J.shape
Out[144]: (1, 299)

Note that this picks 200 'rows' in one dimension, and 299 in the other.

I'm not sure why this kind of reworking of the arrays is needed in this case, but not in the 2nd

In [154]: b=np.arange(2*3*2).reshape((2,3,2))

In [155]: n1=np.array([True,False,True])
In [156]: n2=np.array([True,False])

In [157]: b[:,n1,n2]
Out[157]: 
array([[ 0,  4],      # shape (2,2)
       [ 6, 10]])

Taking the same ix_ strategy produces the same values but a different shape:

In [164]: b[np.ix_(np.arange(b.shape[0]),n1,n2)]
# or I,J=np.ix_(n1,n2);b[:,I,J]
Out[164]: 
array([[[ 0],
        [ 4]],

       [[ 6],
        [10]]])

In [165]: _.shape
Out[165]: (2, 2, 1)

Both cases use all rows of the 1st dimension. The ix one picks 2 'rows' of the 2nd dim, and 1 column of the last, resulting the (2,2,1) shape. The other picks b[:,0,0] and b[0,2,0] terms, resulting (2,2) shape. (see my addenda as to why both are simply broadcasting). These are all cases of advanced indexing, with boolean and numeric indexes. One can study the docs, or one can play around. Sometimes it's more fun to do the later. :)

(I knew that ix_ was good for adding the necessary np.newaxis to arrays so can be broadcast together, but didn't realize that worked with boolean arrays as well - it uses np.nonzero() to convert boolean to indices.)

Resolution

Underlying this is, I think, a confusion over 2 modes of indexing. which might called 'diagonal' and 'rectangular' (or element-by-element selection versus block selection). To illustrate look at a small 2d array

In [73]: M=np.arange(6).reshape(2,3)
In [74]: M
Out[74]: 
array([[0, 1, 2],
       [3, 4, 5]])

and 2 simple numeric indexes

In [75]: m1=np.arange(2); m2=np.arange(2)

They can be used 2 ways:

In [76]: M[m1,m2]
Out[76]: array([0, 4])

and

In [77]: M[m1[:,None],m2]
Out[77]: 
array([[0, 1],
       [3, 4]])

The 1st picks 2 points, the M[0,0] and M[1,1]. This kind of indexing lets us pick out the diagonals of an array.

The 2nd picks 2 rows and from that 2 columns. This is the kind of indexing the np.ix_ produces. The 1st picks 2 points, the M[0,0] and M[1,1]. This a 'rectangular' form of indexing.

Change m2 to 3 values:

In [78]: m2=np.arange(3)
In [79]: M[m1[:,None],m2]   # returns a 2x3
Out[79]: 
array([[0, 1, 2],
       [3, 4, 5]])
In [80]: M[m1,m2]   # produces an error
...
ValueError: shape mismatch: objects cannot be broadcast to a single shape

But if m2 has just one element, we don't get the broadcast error - because the size 1 dimension can be expanded during broadcasting:

In [81]: m2=np.arange(1)
In [82]: M[m1,m2]
Out[82]: array([0, 3])

Now change the index arrays to boolean, each matching the length of the respective dimensions, 2 and 3.

In [91]: m1=np.ones(2,bool); m2=np.ones(3,bool)
In [92]: M[m1,m2]
...
ValueError: shape mismatch: objects cannot be broadcast to a single shape
In [93]: m2[2]=False  # m1 and m2 each have 2 True elements
In [94]: M[m1,m2]
Out[94]: array([0, 4])
In [95]: m2[0]=False   # m2 has 1 True element
In [96]: M[m1,m2]
Out[96]: array([1, 4])

With 2 and 3 True terms we get an error, but with 2 and 2 or 2 and 1 it runs - just as though we'd used the indices of the True elements: np.nonzero(m2).

To apply this to your examples. In the first, m1 and m2 have 200 and 299 True elements. a[:,m1,m2,:] fails because of a mismatch in the number of True terms.

In the 2nd, they have 2 and 1 True terms, with nonzero indices of [0,2] and [0], which can be broadcast to [0,0]. So it runs.

http://docs.scipy.org/doc/numpy-1.10.0/reference/arrays.indexing.html explains boolean array indexing in terms of nonzero and ix_.

Combining multiple Boolean indexing arrays or a Boolean with an integer indexing array can best be understood with the obj.nonzero() analogy. The function ix_ also supports boolean arrays and will work without any surprises.

Addenda

On further thought the distinction between 'diagonal' and 'block/rectangular' indexing might be more my mental construct that numpys. Underlying both is the concept of broadcasting.

Take the n1 and n2 booleans, and get their nonzero equivalents:

In [107]: n1
Out[107]: array([ True, False,  True], dtype=bool)
In [108]: np.nonzero(n1)
Out[108]: (array([0, 2], dtype=int32),)
In [109]: n2
Out[109]: array([ True, False], dtype=bool)
In [110]: np.nonzero(n2)
Out[110]: (array([0], dtype=int32),)

Now try broadcasting in the 'diagonal' and 'rectangular' modes:

In [105]: np.broadcast_arrays(np.array([0,2]),np.array([0]))
Out[105]: [array([0, 2]), 
           array([0, 0])]

In [106]: np.broadcast_arrays(np.array([0,2])[:,None],np.array([0]))
Out[106]: 
[array([[0],
        [2]]), 
 array([[0],
        [0]])]

One produces (2,) arrays, the other (2,1).