Numpy按轴加权平均多维数据

Question

我有一个名为 a 的 numpy 数组，形状是 (48,90,144)。我想沿着第一个轴计算 a 的加权平均值，权重来自另一个数组 b，它的形状是 (90,144)。所以我希望得到的结果是一个形状为 (48,) 的 numpy 数组。

我知道可以用列表推导式来实现这个功能：

np.array([np.average(a[i], weights=b) for i in range(48)])

但我想避免将列表再转换回 numpy 数组。

有没有人能帮帮我？我相信可以用 numpy 的函数和切片来做到这一点，但我现在卡住了。谢谢！

Answer 1

这是我能想到的最快的方法：

(a * b).mean(-1).mean(-1) * (b.size / b.sum())

这个方法可以适用于任何数量的输入和输出维度。

重新调整形状和求平均值的操作并没有进一步加快速度：

(a * b).reshape(len(a), -1).mean(-1) * (b.size / b.sum())

Answer 2

一句话搞定：

np.average(a.reshape(48, -1), weights=b.ravel()), axis=1)

你可以用下面的方式来测试：

a = np.random.rand(48, 90, 144)
b = np.random.rand(90,144)
np.testing.assert_almost_equal(np.average(a.reshape(48, -1),
                                          weights=b.ravel(), axis=1),
                               np.array([np.average(a[i],
                                                    weights=b) for i in range(48)]))