python - 不计算整个矩阵的 numpy 矩阵的对角线

Question

我有一个简单的代数问题，我想用 numpy 解决它(当然我可以用 numba 轻松解决它，但这不是重点).

让我们考虑第一个随机矩阵 A，大​​小为 (m x n)，n 是一个大值，第二个随机矩阵 B 大小为 (n x n)。

A = np.random.random((1E6, 1E2))
B = np.random.random((1E2, 1E2))

我们要计算以下表达式:

np.diag(np.dot(np.dot(A,B),B.T))

问题是整个矩阵被加载到内存中，然后才提取对角线。是否有可能以更有效的方式进行此操作？

Answer 1

这就是我从您的起始表达式开始处理它的方式

np.diag(np.dot(np.dot(A,B),B.T))

您可以从分组术语开始:

np.diag(np.dot(A, np.dot(B,B.T)))

然后只使用 A 的第一个相关(正方形)部分:

np.diag(np.dot(A[:B.shape[0], :], np.dot(B,B.T)))

然后通过自己进行逐元素乘法来避免额外的乘法(那会落在对角线之外):

np.sum( np.multiply(A[:B.shape[0], :].T, np.dot(B,B.T)), 0)