为 3D numpy 数组使用掩码 -> 返回 2D 数组答案

【问题标题】：Using a mask for a 3D numpy array -> return 2D array为 3D numpy 数组使用掩码 -> 返回 2D 数组
【发布时间】：2021-09-20 11:10:03
【问题描述】：

我有什么：

import numpy as np
np.random.seed(42)
dlen = 250000
data = np.random.rand(dlen, 3, 3)
mask = np.random.choice([0, 1, 2], dlen)

我想得到什么：

[[0.37454012 0.95071431 0.73199394], 
 [0.83244264 0.21233911 0.18182497], 
 [0.13949386 0.29214465 0.36636184], 
 [0.94888554 0.96563203 0.80839735], 
 [0.44015249 0.12203823 0.49517691],
 ....
(250000, 3)

我尝试使用什么：

data[:,mask,:]

{MemoryError}Unable to allocate 1.36 TiB for an array with shape (250000, 250000, 3) and data type float64

什么给出了正确的结果但看起来很奇怪：

data[np.arange(data.shape[0]), mask, :]

那么这个面具的正确使用方法是什么？

更新：掩码应选择具有指定索引的列。形状为 [2,3,3] 的数组示例：

array = [[[5 6 7], [7 8 9], [2 3 4]],
         [[2 1 0], [7 6 5], [7 6 5]]]
mask = [1 0]
result = [[7 8 9], 
          [2 1 0]]

【问题讨论】：

你能用文字解释一下你的面具应该完成什么吗？
试图在问题中解释。
您给定的示例数组只有 2 个轴而不是 3 个
..gives the correct result but looks strange - 之所以有效，是因为您使用 Index array
这些是现在例子中的三个轴。

标签： python arrays numpy slice mask

【解决方案1】：

data[np.arange(data.shape[0]), mask, :]

因为它是multi-dimensional index array，所以有效

当我在这里使用术语 mask 时，我想到了布尔索引。您的整数掩码可以转换为布尔掩码，以按照您想要的方式使用它。

>>> data.shape                 
(250000, 3, 3)
>>> mask.shape
(250000,)
>>> q = mask[:,None] == [0,1,2]
>>> q.shape
(250000, 3)
>>> q[:5]        
array([[ True, False, False],
       [False,  True, False],
       [False,  True, False],
       [False, False,  True],
       [False,  True, False]])
>>> r = data[q]
>>> r.shape
(250000, 3)
>>> r[:10]
array([[0.37454012, 0.95071431, 0.73199394],
       [0.83244264, 0.21233911, 0.18182497],
       [0.13949386, 0.29214465, 0.36636184],
       [0.94888554, 0.96563203, 0.80839735],
       [0.44015249, 0.12203823, 0.49517691],
       [0.66252228, 0.31171108, 0.52006802],
       [0.59789998, 0.92187424, 0.0884925 ],
       [0.14092422, 0.80219698, 0.07455064],
       [0.00552212, 0.81546143, 0.70685734],
       [0.31098232, 0.32518332, 0.72960618]])
>>>

你可以使用第二个维度 length 来让它更通用一点：

q = mask[:,None] == np.arange(data.shape[1])
>>> q[:5]                                        
array([[ True, False, False], 
       [False,  True, False], 
       [False,  True, False], 
       [False, False,  True], 
       [False,  True, False]])

如果您控制掩码的构造，您可能希望将其构造为布尔数组。

如果这是新代码，您可能需要升级到兼容版本的 Numpy 并使用 new random generator。

【讨论】：

非常有帮助，谢谢。对于面具的构造，我使用mask = np.argmin(some_array, axis=1) 而不是我可以这样做：boolean_mask = np.min(some_array, axis=1)[:, None] == some_array 只是给我一个注释。