具有numpy的多轴上的布尔遮罩

Question

我想对行和列应用布尔掩码.

同

X = np.array([[1,2,3],[4,5,6]])
mask1 = np.array([True, True])
mask2 = np.array([True, True, False])
X[mask1, mask2]

我期待输出

array([[1,2],[4,5]])

代替

array([1,5])

众所周知

X[:, mask2]

可以在这里使用,但这不是一般情况的解决方案.

我想知道它是如何在引擎盖下工作的,为什么在这种情况下结果是array([1,5]).

Answer 1

X[mask1, mask2]在布尔数组索引文档中描述为等效的

In [249]: X[mask1.nonzero()[0], mask2.nonzero()[0]]
Out[249]: array([1, 5])
In [250]: X[[0,1], [0,1]]
Out[250]: array([1, 5])

实际上它给你X[0,0]和X[1,1](配对0和1).

你想要的是:

In [251]: X[[[0],[1]], [0,1]]
Out[251]: 
array([[1, 2],
       [4, 5]])

np.ix_ 是一个方便的工具,用于创建正确的维度组合

In [258]: np.ix_([0,1],[0,1])
Out[258]: 
(array([[0],
        [1]]), array([[0, 1]]))
In [259]: X[np.ix_([0,1],[0,1])]
Out[259]: 
array([[1, 2],
       [4, 5]])

这实际上是第一轴的列向量和第二轴的行向量,它们共同定义了所需的矩形值.

但尝试广播这样的布尔数组不起作用: X[mask1[:,None], mask2]

但该参考部分说:

使用obj.nonzero()类比可以最好地理解组合多个布尔索引数组或布尔与整数索引数组.函数ix_也支持布尔数组,并且可以毫无意外地工作.

In [260]: X[np.ix_(mask1, mask2)]
Out[260]: 
array([[1, 2],
       [4, 5]])
In [261]: np.ix_(mask1, mask2)
Out[261]: 
(array([[0],
        [1]], dtype=int32), array([[0, 1]], dtype=int32))

布尔部分ix_:

    if issubdtype(new.dtype, _nx.bool_):
        new, = new.nonzero()

所以它适用于像 X[np.ix_(mask1, [0,2])]