如何计算 3d numpy 数组的协方差矩阵？

Question

我有一个矩阵A，形状为(N, D, 4)。首先我计算A转置，A_t。我想计算A_t 乘以A 的乘积。我希望得到的矩阵形状为(D, D)，并且矩阵的乘积就像最后一个由 4 个分量组成的向量是一个数字。 （两个向量的点积是一个数。）

import numpy as np

N = 15
D = 98
A = np.random.random((N, D, 4))
A_t = np.zeros((D, N, 4))
for i in range(N):
    A_t[:, i] = A[i]

S = np.zeros((D, D))

for i in range(D):
    row = A_t[i]
    for j in range(D):
        col = A[:, j, :]
        val = 0
        for n in range(N):
            val += np.matmul(row[n], col[n])
        S[i][j] = val

print(A.shape)
print(A_t.shape)
print(S.shape)

Answer 1

让我们来看看您正在尝试的操作，看看我们可以做些什么来简化它们。对于初学者，你可以写

A_t = np.swapaxes(A, 0, 1)

这相当于

A_t = np.transpose(A, [0, 1, 2])

或

A_t = A.transpose([0, 1, 2])

碰巧的是，您当前的应用程序都不需要。要了解原因，让我们使用一个简化的示例：

np.random.seed(42)
N = 4
D = 3
K = 2
A = np.random.randint(0, 10, (N, D, K))

在你的外循环中，你有row = A_t[i]。但根据您的转置定义，这与row = A[:, i, :] 相同，让您的生活更轻松，转置是多余的。

内部循环对一些点积求和：

val = 0
for n in range(N):
    val += np.matmul(row[n], col[n])

如果你记得点积的定义，你会发现你正在做的等价于

np.sum(np.sum(row * col, axis=1), axis=0)

内部和是循环中的和积，而外部和是val 的计算。分别对两个维度求和与一次对整个缓冲区求和是一样的，所以我们可以立即用 just 替换内部循环

for i in range(D):
    for j in range(D):
        S[i][j] = np.sum(A[:, i, :] * A[:, j, :])

您可以使用np.dot、np.tensordot、np.einsum 或简单的广播来简化此操作。前两个是不必要的复杂，因为您实际上是同时在两个维度上相乘。 np.einsum 总体上提供了最直接的解决方案，但它对您的代码的翻译不太直接。

解决方案 1：广播

让我们从双循环的直接广播版本开始，然后再转向更惯用的解决方案：

S = (A[:, None, ...] * A[:, :, None, ...]).sum(axis=(0, -1))

或

S = np.sum(A[:, None, ...] * A[:, :, None, ...], axis=(0, -1))

这会分别创建A 形状的(N, 1, D, K) 和(N, D, 1, K) 的视图。乘法在每种情况下都将复制的 D 轴广播到 for 循环所做的事情，因此 N 和 K 轴上的最终总和与 S[i][j] = np.sum(A[:, i, :] * A[:, j, :]) 之前的行完全相同。

解决方案 2：np.einsum

此解决方案可让您将 sum-product 直接应用于您想要的任何轴：

S = np.einsum('ijk,ihk->jh', A, A)

请注意，您必须为第二个矩阵的第二个轴（j 和 h）使用不同的字母，以表明您不会在该轴上求和。 S 是对称的，但如果不是，您可以通过在结果中转置为 ->hj 来转置它。