CODEWITHSUNDEEP
pythonnumpymatrix-indexing
Thrastylon
Thrastylon

Reputation: 990

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]])

Upvotes: 3

Views: 1199

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114528

The error tells you what you need to do: the mask dimensions need to broadcast together. You can fix this at the source:

m0 = (X>0).all(axis=1, keepdims=True)
m1 = (X<3).any(axis=0, keepdims=True)

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

You only really need to apply keepdims to m0, so you can leave the masks as 1D:

>>> X[m0[:, None] & m1]
array([ 4,  5,  6,  8,  9, 10])

You can reshape to the desired shape:

>>> X[m0[:, None] & m1].reshape(np.count_nonzero(m0), np.count_nonzero(m1))
array([[ 4,  5,  6],
       [ 8,  9, 10]])

Another option is to convert the masks to indices:

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

Upvotes: 2

not_speshal
not_speshal

Reputation: 23166

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]])

Upvotes: 3

Related Questions