How to multiply two 2D RFFT arrays (FFTPACK) to be compatible with NumPy's FFT?

Question

I\'m trying to multiply two 2D arrays that were transformed with fftpack_rfft2d() (SciPy\'s FFTPACK RFFT) and the result is not compatible with what I get from

Answer 1

In addition to @CrisLuengo answer (https://stackoverflow.com/a/61873672/501852).

Performance test

Test fftpack.FFT vs fftpack.RFFT - 1D

# test data
sz =50000
sz = fftpack.next_fast_len(sz)
in1 = np.random.randn(sz)

print(f"Input (len = {len(in1)}):", sep='\n')

rep = 1000

tic = time.perf_counter()
for i in range(rep):
    spec1 = fftpack.fft(in1,axis=0)
toc = time.perf_counter()
print("", f"Spectrum FFT (len = {len(spec1)}):",
      f"spec1 takes {10**6*((toc - tic)/rep):0.4f} us", sep="\n")

sz2 = sz//2 + 1
spec2 = np.empty(sz2, dtype=np.complex128)

tic = time.perf_counter()
for i in range(rep):
    tmp = fftpack.rfft(in1)

    assert  tmp.dtype == np.dtype('float64')

    if not sz & 0x1:
        end = -1 
        spec2[end] = tmp[end]
    else:
        end = None

    spec2[0] = tmp[0]
    spec2[1:end] = tmp[1:end].view(np.complex128)

toc = time.perf_counter()
print("", f"Spectrum RFFT (len = {len(spec2)}):",
      f"spec2 takes {10**6*((toc - tic)/rep):0.4f} us", sep="\n")

Results are

Input (len = 50000):

Spectrum FFT (len = 50000):
spec1 takes 583.5880 us

Spectrum RFFT (len = 25001):
spec2 takes 476.0843 us

• So, using fftpack.rfft() with further casting its output into complex view is ~15-20% faster, than fftpack.fft() for big arrays.

Test fftpack.FFT vs fftpack.FFT2 - 2D

Similar test for the 2D case:

# test data
sz = 5000
in1 = np.random.randn(sz, sz)

print(f"Input (len = {len(in1)}):", sep='\n')

rep = 1

tic = time.perf_counter()
for i in range(rep):
    spec1 = np.apply_along_axis(fftpack.fft, 0, in1)
    spec1 = np.apply_along_axis(fftpack.fft, 1, spec1)
toc = time.perf_counter()
print("", f"2D Spectrum FFT with np.apply_along_axis (len = {len(spec1)}):",
      f"spec1 takes {10**0*((toc - tic)/rep):0.4f} s", sep="\n")


tic = time.perf_counter()
for i in range(rep):
    spec2 = fftpack.fft(in1,axis=0)
    spec2 = fftpack.fft(spec2,axis=1)
toc = time.perf_counter()
print("", f"2D Spectrum 2xFFT (len = {len(spec2)}):",
      f"spec2 takes {10**0*((toc - tic)/rep):0.4f} s", sep="\n")

tic = time.perf_counter()
for i in range(rep):
    spec3 = fftpack.fft2(in1)
toc = time.perf_counter()
print("", f"2D Spectrum FFT2 (len = {len(spec3)}):",
      f"spec3 takes {10**0*((toc - tic)/rep):0.4f} s", sep="\n")

# compare
print('\nIs spec1 equivalent to the spec2?', np.allclose(spec1, spec2))
print('\nIs spec2 equivalent to the spec3?', np.allclose(spec2, spec3), '\n')

Results for matrix of size = 5x5

Input (len = 5):

2D Spectrum FFT with np.apply_along_axis (len = 5):
spec1 takes 0.000183 s

2D Spectrum 2xFFT (len = 5):
spec2 takes 0.000010 s

2D Spectrum FFT2 (len = 5):
spec3 takes 0.000012 s

Is spec1 equivalent to the spec2? True

Is spec2 equivalent to the spec3? True

Results for matrix of size = 500x500

Input (len = 500):

2D Spectrum FFT with np.apply_along_axis (len = 500):
spec1 takes 0.017626 s

2D Spectrum 2xFFT (len = 500):
spec2 takes 0.005324 s

2D Spectrum FFT2 (len = 500):
spec3 takes 0.003528 s

Is spec1 equivalent to the spec2? True

Is spec2 equivalent to the spec3? True 

Results for matrix of size = 5000x5000

Input (len = 5000):

2D Spectrum FFT with np.apply_along_axis (len = 5000):
spec1 takes 2.538471 s

2D Spectrum 2xFFT (len = 5000):
spec2 takes 0.846661 s

2D Spectrum FFT2 (len = 5000):
spec3 takes 0.574397 s

Is spec1 equivalent to the spec2? True

Is spec2 equivalent to the spec3? True

Conclusions

From the tests above, it seems, that use of fftpack.fft2() is more efficient for bigger matrices.

Use of np.apply_along_axis() is the most slow method.