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不适用于后两次减薪？

好吧，我们需要保持dinMat1 / dinMat2与temp_rho / np.conj(temp_rho)的第一个轴对齐的第一个轴，这是不可能的tensordot因为不是sum-reduced的轴沿着两个单独的轴进行元素相乘。因此，当与np.tensordot一起使用时，我们最终将得到两个长度为N的轴，这两个轴对应于来自两个输入的第一个轴。

调整矩阵数组乘法使用Numpy Tensordot

1 个答案: