3

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]])
Mad Physicist
  • 107,652
  • 25
  • 181
  • 264
Thrastylon
  • 853
  • 7
  • 20

2 Answers2

3

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]])
not_speshal
  • 22,093
  • 2
  • 15
  • 30
  • [`np.ix_`](https://numpy.org/doc/stable/reference/generated/numpy.ix_.html) is indeed what I was after, thanks. I would delete the second option to make the answer clearer so I can accept it. – Thrastylon Aug 19 '21 at 16:42
  • That's just an additional answer choice - also straight from the docs. Perhaps would make sense to someone else who stumbles upon this question. There is no need to delete another correct answer I think. I however edited to push it to the bottom. – not_speshal Aug 19 '21 at 17:31
2

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]])
Mad Physicist
  • 107,652
  • 25
  • 181
  • 264
  • Very nice! But I should have mentioned that the masks being 1-D is a desired feature, as I want to use them elsewhere (e.g. on vectors or matrices of different shape than `X`). – Thrastylon Aug 19 '21 at 16:40
  • @Thrastylon. I've added more options. `ix_` does the last one. – Mad Physicist Aug 19 '21 at 17:42
  • I'd still prefer using `keepdims` parameter.. If needed anywhere as 1D, array then `np.ravel` can be used. – ThePyGuy Aug 19 '21 at 17:54