是否有 numpy/scipy 点积，仅计算结果的对角线条目? [英] Is there a numpy/scipy dot product, calculating only the diagonal entries of the result?

问题描述

想象一下有 2 个 numpy 数组:

Imagine having 2 numpy arrays:

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

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

Typically p is a smaller number (p <= 200), while n can be arbitrarily large.

我正在做以下事情:

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

如您所见，我只保留了 n 个对角线条目，但是有一个中间 (n x n) 数组，从中只保留了对角线条目.

As you can see, I am keeping only the n diagonal entries, however there is an intermediate (n x n) array calculated from which only the diagonal entries are kept.

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

I wish for a function like diag_dot(), which only calculates the diagonal entries of the result and does not allocate the complete memory.

结果是:

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

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

Is there a premade functionality like this and can this be done efficiently without the need for allocating the intermediate (n x n) array?

推荐答案

我想我是自己搞定的，但还是会分享解决方案:

I think i got it on my own, but nevertheless will share the solution:

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

since getting only the diagonals of a matrix multiplication

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

相当于X的行和Y的列的标量积的单独和，前面的语句相当于:

is equivalent to the individual sum of the scalar product of rows of X and columns of Y, the previous statement is equivalent to:

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

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

For the original variables this means:

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

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

Please correct me if I am wrong but this should be it ...

这篇关于是否有 numpy/scipy 点积，仅计算结果的对角线条目?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！