pythonnumpymatrix-indexing

Apply multiple masks at once to a Numpy array


Is there a way to apply multiple masks at once to a multi-dimensional Numpy array?

For instance:

X = np.arange(12).reshape(3, 4)
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]])
m0 = (X>0).all(axis=1) # array([False,  True,  True])
m1 = (X<3).any(axis=0) # array([ True,  True,  True, False])

# In one step: error
X[m0, m1]
# IndexError: shape mismatch: indexing arrays could not 
#             be broadcast together with shapes (2,) (3,) 

# In two steps: works (but awkward)
X[m0, :][:, m1]
# array([[ 4,  5,  6],
#        [ 8,  9, 10]])


Solution

  • Try:

    >>> X[np.ix_(m0, m1)]
    array([[ 4,  5,  6],
           [ 8,  9, 10]])
    

    From the docs:

    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.

    Another solution (also straight from the docs but less intuitive IMO):

    >>> X[m0.nonzero()[0][:, np.newaxis], m1]
    array([[ 4,  5,  6],
           [ 8,  9, 10]])