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

Question

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

让我们考虑一个大小为 (mxn) 的第一个随机矩阵A，具有 na 个大值，以及一个大小为 (nxn)的第二个随机矩阵B。

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

我们要计算以下表达式：

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

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

标签： pythonnumpy

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)