python - 根据另一个数组中的 ID 值高效处理一个数组中的值

Question

我有一个包含一系列值的二维 numpy 数组，以及一个标识单元所属组的匹配数组。我希望在每组的第一个数组中运行一些简单的操作。为了说明这一点，我可能有一个如下所示的数据和 ID 数组:

data = np.arange(9).reshape(3,3)
IDs = np.array([
    ["A", "A", "B"],
    ["A", "B", "C"],
    ["B", "C", "C"]
])

我可能想要获取所有 A、B 和 C 值的平均值和标准差，输出如下所示:

[['A', 1.333333333333333, 1.247219128924647],
['B', 4.0, 1.6329931618554521],
['C', 6.666666666666667, 1.247219128924647]]

现在，我用以下代码解决了这个问题:

unique_IDs = np.unique(IDs)
results = []
for ID in unique_IDs:
    group_data = data[IDs == ID]
    mean = np.mean(group_data)
    stdev = np.std(group_data)
    results.append([ID, mean, stdev])

我认为这很简单且易于理解，但速度不是很快。当使用我的真实数据时，数组通常约为 200 mb(5,000 * 10,000 而不是 3 * 3)，并且有数百个唯一 ID。每次迭代需要几秒钟，整个操作可能需要半个小时以上。

我的“直觉”是，如果我不必迭代每个唯一的 ID 值并重复查找和操作，这可以更快地完成，但我不知道这是否属实或我将如何进行它。

Answer 1

您可以使用np.unique获取与字符串 IDs 对应的数字 IDs。这些数字 ID 在mean 和std 计算中都很有用。

现在，对于平均值计算，这些数字 ID 可以用作“权重”，用于使用 np.bincount 进行分箱。 ，为我们提供与每个 ID 对应的数据元素的总和。接下来。使用每个 ID 的元素计数，就可以得到平均值。这是实现 -

_,idx,counts = np.unique(IDs,return_inverse=True,return_counts=True)  
mean_vals = np.bincount(idx,data.ravel())/counts

对于 std 计算，一种方法是对 ID 进行排序，以便连续放置相同的 ID。然后，重新排列数据形成另一个二维数组，使得相同ID对应的所有数据元素位于同一行，未填充的位置设置为NaN。这里的想法是以矢量化方式沿列执行np.std。需要对需要填补的地方进行掩蔽。请注意，如果存在计数相对较高的 ID，这可能会占用大量内存。该实现将重用之前代码中的 idx 和 counts ，看起来像这样 -

data_RO = np.empty((counts.size,counts.max()))
data_RO[:] = np.nan
data_RO[np.arange(counts.max()) < counts[:,None]] = data.ravel()[idx.argsort()]
std_vals = np.nanstd(data_RO,axis=1)

示例运行并验证输出 -

1) 输入:

In [51]: data
Out[51]: 
array([[0, 1, 6],
       [2, 5, 0],
       [6, 3, 6]])

In [52]: IDs
Out[52]: 
array([['A', 'A', 'B'],
       ['A', 'B', 'C'],
       ['B', 'A', 'C']], 
      dtype='|S1')

2)问题中列出的代码的输出:

In [53]: unique_IDs = np.unique(IDs)

In [54]: results = []

In [55]: for ID in unique_IDs:
   ....:         group_data = data[IDs == ID]
   ....:         mean = np.mean(group_data)
   ....:         stdev = np.std(group_data)
   ....:         results.append([ID, mean, stdev])
   ....:     

In [56]: results
Out[56]: 
[['A', 1.5, 1.1180339887498949],
 ['B', 5.666666666666667, 0.47140452079103168],
 ['C', 3.0, 3.0]]

3)建议解决方案的输出:

In [57]: _,idx,counts = np.unique(IDs,return_inverse=True,return_counts=True) 

In [58]: np.bincount(idx,data.ravel())/counts
Out[58]: array([ 1.5       ,  5.66666667,  3.        ])

In [59]: data_RO = np.empty((counts.size,counts.max()))

In [60]: data_RO[:] = np.nan

In [61]: mask = np.arange(counts.max()) < counts[:,None]

In [62]: data_RO[mask] = data.ravel()[idx.argsort()]

In [63]: np.nanstd(data_RO,axis=1)
Out[63]: array([ 1.11803399,  0.47140452,  3.        ])