是否有一個numpy/scipy點積，只計算結果的對角線條目？

Question

想象具有2 numpy的陣列：

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

通常p是一個較小的數字（第< = 200），而n可以是任意大的。

我做了以下內容：

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

正如你所看到的，我只保留n個對角線項，但有計算的中間爲（N×N）陣列其中只有對角線條目保留。

我希望有一個像diag_dot（）這樣的函數，它只計算結果的對角條目並且不分配完整的內存。

甲結果將是：

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

是否有這樣的預製的功能，並且可以這樣無需用於分配中間（N×n）個陣列有效地進行？

Answer 1

我想我對我自己，但仍然將分享解決方案：

，因爲只得到一個矩陣乘法

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

的對角線相當於標產品的逐筆X和Y的列的行，前面的語句相當於：

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

對於原始變量，這意味着：

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

請糾正我，如果我錯了，但是這應該是吧...

Answer 2

你能得到你曾經與numpy.einsum夢想幾乎所有的東西。直到你開始了它的竅門，它基本上看起來像黑巫術......

>>> a = np.arange(15).reshape(5, 3) 
>>> b = np.arange(9).reshape(3, 3) 

>>> np.diag(np.dot(np.dot(a, b), a.T)) 
array([ 60, 672, 1932, 3840, 6396]) 
>>> np.einsum('ij,ji->i', np.dot(a, b), a.T) 
array([ 60, 672, 1932, 3840, 6396]) 
>>> np.einsum('ij,ij->i', np.dot(a, b), a) 
array([ 60, 672, 1932, 3840, 6396]) 

編輯其實你可以得到單杆整個事情，這是荒謬的......

>>> np.einsum('ij,jk,ki->i', a, b, a.T) 
array([ 60, 672, 1932, 3840, 6396]) 
>>> np.einsum('ij,jk,ik->i', a, b, a) 
array([ 60, 672, 1932, 3840, 6396]) 

編輯你不想讓它自己的數字太多，但...增加了OP的答案，以自己的問題進行比較。

n, p = 10000, 200 
a = np.random.rand(n, p) 
b = np.random.rand(p, p) 

In [2]: %timeit np.einsum('ij,jk,ki->i', a, b, a.T) 
1 loops, best of 3: 1.3 s per loop 

In [3]: %timeit np.einsum('ij,ij->i', np.dot(a, b), a) 
10 loops, best of 3: 105 ms per loop 

In [4]: %timeit np.diag(np.dot(np.dot(a, b), a.T)) 
1 loops, best of 3: 5.73 s per loop 

In [5]: %timeit (a.dot(b) * a).sum(-1) 
10 loops, best of 3: 115 ms per loop 

Answer 3

甲行人答案，這避免了大的中間陣列的結構是：

result=np.empty([n.], dtype=A.dtype) 
for i in xrange(n): 
    result[i]=A[i,:].dot(B).dot(A[i,:])