適應矩陣陣列乘法使用numpy的Tensordot

Question

我試圖加快我的代碼來執行一些數值計算，我需要乘以3點矩陣與數組。問題的結構如下：

• 的陣列的形狀（N，10）
• 的第一矩陣是沿陣列的動態尺寸恆定，並且具有（10，10的形狀）
• 另外兩個矩陣沿所述陣列的所述第一尺寸改變，並且具有（N，10,10）形狀
• 計算的結果應與（N，形狀）的陣列

我已經使用for循環實現了一個解決方案是工作，但我想有更好的表現，所以我嘗試使用numpy的功能。我已經使用numpy.tensordot嘗試，但與陣列乘以動態矩陣時，我得到的（N，10，N）代替（N，10）的形狀

我的for循環如下：

res = np.zeros(temp_rho.shape, dtype=np.complex128) 
for i in range(temp_rho.shape[0]): 
    res[i] = np.dot(self.constMatrix, temp_rho[i]) 
    res[i] += np.dot(self.dinMat1[i], temp_rho[i]) 
    res[i] += np.dot(self.dinMat2[i], np.conj(temp_rho[i])) 
#temp_rho.shape = (N, 10) 
#res.shape = (N, 10) 
#self.constMatrix.shape = (10, 10) 
#self.dinMat1.shape = (N, 10, 10) 
#self.dinMat2.shape = (N, 10, 10) 

如何把這個代碼實現numpy的點積，返回正確尺寸是多少？

Answer 1

這是一個使用的np.dot組合和np.einsum的方法 -

parte1 = constMatrix.dot(temp_rho.T).T 
parte2 = np.einsum('ijk,ik->ij',dinMat1, temp_rho) 
parte3 = np.einsum('ijk,ik->ij',dinMat2, np.conj(temp_rho)) 
out = parte1 + parte2 + parte3 

替代的方式得到parte1將與np.tensordot -

parte1 = np.tensordot(temp_rho, constMatrix, axes=([1],[1])) 

爲什麼對後兩個不numpy.tensordot工作總和，降低？

那麼，我們需要保持對的temp_rho/np.conj(temp_rho)第一軸線，這是不可能的tensordot是，不sum-reduced被的elementwise沿着兩個單獨的軸相乘的軸對準dinMat1/dinMat2之間的第一軸線。因此，當與np.tensordot一起使用時，我們最終會得到兩個長度爲N的軸，分別對應於兩個輸入中的第一個軸。