python - 是否有一个 numpy/scipy 点积，只计算结果的对角线条目？

Question

想象有 2 个 numpy 数组:

> A, A.shape = (n,p)
> B, B.shape = (p,p)

通常 p 是一个较小的数字 (p <= 200)，而 n 可以任意大。

我正在做以下事情:

result = np.diag(A.dot(B).dot(A.T))

如您所见，我只保留了 n 个对角线条目，但是有一个计算出的中间 (n x n) 数组，仅保留对角线条目。

我希望有一个像 diag_dot() 这样的函数，它只计算结果的对角线条目，而不分配完整的内存。

结果是:

> result = diag_dot(A.dot(B), A.T)

是否有这样的预制功能，是否可以在无需分配中间 (n x n) 数组的情况下高效完成？

Answer 1

我想我自己得到了它，但仍然会分享解决方案:

因为只得到矩阵乘法的对角线

> Z = N.diag(X.dot(Y))

相当于X的行和Y的列的标量积的个体总和，前面的语句相当于:

> Z = (X * Y.T).sum(-1)

对于原始变量，这意味着:

> result = (A.dot(B) * A).sum(-1)

如果我错了，请纠正我，但应该是这样......