如何有效地计算 numpy 二维数组的块均值（不规则块）？

Posted 2023-03-24

【中文标题】如何有效地计算 numpy 二维数组的块均值（不规则块）？

【英文标题】：How to efficiently calculate block mean (irregular blocks) for numpy 2D array?

【发布时间】：2021-06-19 04:56:10

【问题描述】：

我的问题与Block mean of numpy 2D array 和block mean of 2D numpy array (in both dimensions) 有关（实际上只是更一般的情况）。我将通过简单的例子来解释这一点。

假设我们有6x6 2D 数组：

array([[7, 1, 6, 6, 4, 2],
       [8, 5, 5, 6, 3, 5],
       [3, 1, 7, 1, 3, 4],
       [6, 8, 3, 2, 3, 3],
       [8, 6, 7, 1, 1, 3],
       [8, 5, 4, 5, 1, 4]])

现在这个矩阵中的每一行（和每一列）都被分配给三个社区之一（社区可以有不同的大小），例如array([0, 0, 1, 1, 1, 2]) 将代表此分配。现在我需要根据这个分配分割这个矩阵并计算块（切片）的平均值。这将产生3x3 块均值矩阵。例如社区对 (0,0) 的块（或切片）是一个 2x2 数组：

array([[7, 1],
       [8, 5]])

意思是5.25。社区对 (0, 1) 的块是 2x3 数组：

array([[6, 6, 4],
       [5, 6, 3]])

意思是5，等等..

块均值的结果数组应如下所示：

array([[5.25      , 5.        , 3.5       ],
       [5.33333333, 3.11111111, 3.33333333],
       [6.5       , 3.33333333, 4.        ]])

我的问题是如何有效地计算这个。现在我正在使用 for 循环——对于每一对社区，我得到适当的切片，计算该切片的平均值并将其存储在单独的矩阵中。但是我需要多次执行此操作，并且需要很多时间。

我不能真正使用（或者我不知道如何使用）reshape 的方法，因为它需要假设块大小相等。

【问题讨论】：

您应该显示循环版本。针对工作解决方案测试改进要容易得多。如果您提供有效的解决方案，我们的工作量也会减少。我怀疑我们能做的最好的事情就是在边缘调整解决方案。正如您所说的 numpy 依赖于常规模式的解决方案不适用。

【参考方案1】：

如果你对其他包开放，Pandas 作为一个方便的groupby 功能：

out = (pd.Series(a.ravel(), 
                  index = pd.MultiIndex.from_product((pairs,pairs)))
         .groupby(level=(0,1)).mean()
         .unstack().to_numpy()
      )

输出：

array([[5.25      , 5.        , 3.5       ],
       [5.33333333, 3.11111111, 3.33333333],
       [6.5       , 3.33333333, 4.        ]])

【参考方案2】：

我能想象的最好的方法是尝试限制循环的数量。我在这里假设 6x6 2D 数组是 arr 并且社区定义是 coms = np.array([0, 0, 1, 1, 1, 2])。

我将首先计算每个社区的切片：

dcoms = k: slice(min(x), 1 + max(x)) for k in np.unique(coms)
         for x in (np.where(coms==k)[0],)

1 次循环 coms

然后我可以在dcoms 上使用 2 个循环直接计算得到的 ndarray：

resul = np.array([[arr[dcoms[i],dcoms[j]].mean() for j in dcoms] for i in dcoms])

它按预期给出：

array([[5.25      , 5.        , 3.5       ],
       [5.33333333, 3.11111111, 3.33333333],
       [6.5       , 3.33333333, 4.        ]])

【参考方案3】：

谢谢大家，我进行了一些基准测试来比较这些解决方案。

设置如下

np.random.seed(0)

mat = (np.random.random((500, 500)) - 0.5) * 100

# Create 5 communities of size 50 and 10 communities of size 15
comms = np.concatenate((np.repeat(np.arange(5), 50), np.repeat(np.arange(5, 15), 25)))

我使用 for 循环的原始解决方案：

unique_comms = np.unique(comms)
block = np.zeros((len(unique_comms), len(unique_comms)))
for i in unique_comms:
    for j in unique_comms:
        block[i, j] = np.mean(mat[comms == i][:, comms == j])

每个循环耗时 7.66 ms ± 43.3 µs（平均值 ± 标准偏差，7 次运行，每次 100 个循环）。

熊猫解决方案：

out = (pd.Series(mat.ravel(), 
                  index = pd.MultiIndex.from_product((comms, comms)))
         .groupby(level=(0,1)).mean()
         .unstack().to_numpy())

每个循环花费 22.1 ms ± 221 µs（平均值 ± 标准偏差，7 次运行，每次 100 个循环）

而切片解决方案：

dcoms = k: slice(min(x), 1 + max(x)) for k in np.unique(comms)
         for x in (np.where(comms==k)[0],)
resul = np.array([[mat[dcoms[i],dcoms[j]].mean() for j in dcoms] for i in dcoms])

每个循环花费 2.1 ms ± 61.7 µs（平均值 ± 标准偏差，7 次运行，每次 100 个循环）。

不幸的是，Pandas 解决方案要慢得多，而 slice 解决方案显然比常规 for 循环解决方案快 3-4 倍。切片解决方案的一个缺点是（我相信）当社区向量被洗牌时它不起作用（即np.array([1, 0, 1, 0, 2, 1]) 而不是np.array([0, 0, 1, 1, 1, 2])。