您可以只使用jit编译器
您的解决方案一点也不差。我唯一更改的是索引和变量循环范围。如果您有numpy数组和过多的循环,则可以使用编译器(Numba),这确实很简单。
import numba as nbimport numpy as np#The function is compiled only at the first call (with using same datatypes)@nb.njit(cache=True) #set cache to false if copying the function to a command windowdef almost_your_solution(matrix1,matrix2): lis = np.zeros(matrix1.shape,np.float64) for i in range(matrix2.shape[0]): for j in range(matrix2.shape[1]): for k in range(matrix2.shape[2]): for l in range(matrix2.shape[3]): lis[i,j] += matrix1[k,l] * (2 * matrix2[i,j,k,l] - matrix2[i,k,j,l]) return lis
关于代码的简单性,与上面显示的解决方案相比,我更希望使用hpaulj的einsum解决方案。我认为,tensordot解决方案不是那么容易理解。但这是一个品味问题。
性能比较
我用于比较的hpaulj函数:
def hpaulj_1(matrix1,matrix2): matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3) return np.einsum('kl,ijkl->ij', matrix1, matrix3)def hpaulj_2(matrix1,matrix2): matrix3 = 2*matrix2-matrix2.transpose(0,2,1,3) (matrix1*matrix3).sum(axis=(2,3)) return np.tensordot(matrix1, matrix3, [[0,1],[2,3]])非常短的数组给出:
matrix1=np.random.rand(6,6)matrix2=np.random.rand(6,6,6,6)Original solution: 2.6 msCompiled solution: 2.1 µsEinsum solution: 8.3 µsTensordot solution: 36.7 µs
较大的数组可得出:
matrix1=np.random.rand(60,60)matrix2=np.random.rand(60,60,60,60)Original solution: 13,3 sCompiled solution: 18.2 msEinsum solution: 115 msTensordot solution: 180 ms
结论
编译将计算速度提高了大约3个数量级,并且比所有其他解决方案都快了许多。



