python - 跨轴应用蒙版后，Numpy 数组变形

Question

问题

我有形状相同的 np.array 和 mask。一旦我应用了 mask ，阵列就会失去它的形状并变成一维 - 扁平化一维。

问题

我想根据轴长为 1D 的掩码，减少我在某个轴上的数组。

如何应用 mask ，但保持数组的维度？

例子

代码中的一个小例子:

# data ...
>>> data = np.ones((4, 4))
>>> data.shape
(4, 4)

# mask ...
>>> mask = np.ones((4, 4), dtype=bool)
>>> mask.shape
(4, 4)

# apply mask ...
>>> data[mask].shape
(16,)

我理想的形状是(4, 4)。

跨轴数组降维的示例:

# data, mask ...
>>> data = np.ones((4, 4))
>>> mask = np.ones((4, 4), dtype=bool)

# remove last column from data ...
>>> mask[:, 3] = False 
>>> mask
array([[ True,  True,  True, False],
       [ True,  True,  True, False],
       [ True,  True,  True, False],
       [ True,  True,  True, False]])

# equivalent mask in 1D ...
>>> mask[0]
array([ True,  True,  True, False])

# apply mask ...
>>> data[mask].shape 
(12,)

数组的理想尺寸为(4, 3)，无需整形。

感谢帮助，谢谢!

Answer 1

实现目标的“正确”方法是不将蒙版扩展为二维。而是使用带有 1D 掩码的 [:, mask] 进行索引。这向 numpy 表明您希望轴 0 不变并且 mask 沿轴 1 应用。

a = np.arange(12).reshape(3, 4)
b = np.array((1,0,1,0),'?')
a
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]])
b
# array([ True, False,  True, False])
a[:, b]
# array([[ 0,  2],
#        [ 4,  6],
#        [ 8, 10]])

如果您的mask 已经是二维的，numpy 将不会检查它的所有行是否都相同，因为那样效率很低。但很明显，您可以在这种情况下使用 [:, mask[0]]。

如果您的 mask 是二维的，并且恰好在每一行中有相同数量的 True，那么要么使用@tel 的答案。或者创建一个索引数组:

B = b^b[:3, None]
B
# array([[False,  True, False,  True],
#        [ True, False,  True, False],
#        [False,  True, False,  True]])
J = np.where(B)[1].reshape(len(B), -1)

现在要么

np.take_along_axis(a, J, 1)
# array([[ 1,  3],
#        [ 4,  6],
#        [ 9, 11]])

或

I = np.arange(len(J))[:, None]
IJ = I, J
a[IJ]
# #array([[ 1,  3],
#         [ 4,  6],
#         [ 9, 11]])