不同形状数组的快速numpy协方差

Question

我需要获取一个形状为(M1, M2, N)的数组a和另一个形状为(N,)的数组b之间的协方差。

目前我的做法是使用for循环块：

import numpy as np

M1, M2, N = 23, 50, 117
a = np.random.uniform(-2., 2., (M1, M2, N))
b = np.random.uniform(-1., 1., N)

c = []
for i in range(M1):
    for j in range(M2):
        c.append(np.cov(a[i][j], b)[0, 1])

但是对于大型的 (M1, M2)，速度会有些慢。有没有办法加快速度？

- Gabriel

Answer 1

您可以手动计算协方差。以下是两个建议，分别使用dot和einsum。

import numpy as np

M1, M2, N = 23, 50, 117
a = np.random.uniform(-2., 2., (M1, M2, N))
b = np.random.uniform(-1., 1., N)

c = []
for i in range(M1):
    for j in range(M2):
        c.append(np.cov(a[i][j], b)[0, 1])

c1 = np.reshape(c, (M1, M2))

ac = a - a.mean(axis=-1, keepdims=True)
bc = (b - b.mean()) / (N-1)

c2 = np.dot(ac, bc)
c3 = np.einsum('ijk,k->ij', ac, bc)
print(np.allclose(c1, c2))
print(np.allclose(c1, c3))

打印

True
True

- Paul Panzer

---

很高兴看到你回来了！我一直在尝试这些东西，但是一直无法通过。在张量上使用tensordot比dot更好：np.tensordot(ac, bc, axes=((-1,-1)))。 - Divakar

所有三个选项（np.einsum，np.dot，np.tensordot）都比for块快得多。由于它们显然都同样快，我会选择np.dot，因为对我来说更容易理解。谢谢你们两个！ - Gabriel