I have a function that I am trying to optimize.

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

I have implemented the logic I needed. But, I am definite there is a optimized way to write this logic. Can someone help?

CodePudding user response：

What you have there is a very manual, lengthy way of multiplying each element of x by the complex conjugate of the corresponding element of y. You don't need to write it out like that. NumPy can already take complex conjugates and multiply complex numbers on its own.

NumPy supports those operations on arrays, too, so you don't have to write out explicit loops. 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()

And with the code being so simple, you might as well just write x * y.conj() directly when you want to do this operation, instead of calling a function.