适应矩阵阵列乘法使用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的轴，分别对应于两个输入中的第一个轴。