Numpy矩阵复运算优化

Question

我有一个正在尝试优化的 function。

def mul_spectrums_with_conj(x: ndarray, y: ndarray) -> ndarray:
    lst = np.empty((x.shape[0], x.shape[1]), dtype=np.complex64)
    for kx in range(x.shape[0]):
        for ky in range(x.shape[1]):
            acc0 = x.real[kx, ky] * y.real[kx, ky] + x.imag[kx, ky] * y.imag[kx, ky]
            acc1 = x.imag[kx, ky] * y.real[kx, ky] - x.real[kx, ky] * y.imag[kx, ky]
            lst[kx][ky] = complex(acc0, acc1)
    return lst

我已经实现了我需要的逻辑。 但是，我确定有一种优化的方式来编写这个逻辑。 有人可以帮忙吗？

Answer 1

您所拥有的是一种非常手动、冗长的方法，将x的每个元素乘以y的相应元素的复共轭。 你不需要那样写出来。 NumPy 已经可以自己进行复共轭和复数相乘。

NumPy 也支持 arrays 上的这些操作，因此您不必编写显式循环。 If you let NumPy handle it instead of looping explicitly, NumPy can do the work in fast C loops over the underlying data buffers, instead of going through all the dynamic dispatch and wrapper objects involved in looping over an array explicitly at Python level:

def mul_spectrums_with_conj(x: ndarray, y: ndarray) -> ndarray:
    return x * y.conj()

而且代码这么简单，你还不如直接写x * y.conj()来做这个操作，而不是调用function。