10

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 scipy_rfft2d() (SciPy's FFT RFFT).

The image below shares the output of the script, which displays:

  • The initialization values of both input arrays;
  • Both arrays after they were transform with SciPy's FFT implementation for RFFT using scipy_rfft2d(), followed by the output of the multiplication after its transformed backwards with scipy_irfft2d();
  • The same things using SciPy's FFTPACK implementation for RFFT with fftpack_rfft2d() and fftpack_irfft2d();
  • The result of a test with np.allclose() that checks if the result of both multiplications are the same after they were transformed back with their respective implementations for IRFFT.

Just to be clear, the red rectangles display the multiplication result after the inverse transform IRFFT: the rectangle on the left uses SciPy's FFT IRFFT; the rectangle on the right, SciPy's FFTPACK IRFFT. They should present the same data when the multiplication with the FFTPACK version is fixed.

I think the multiplication result with the FFTPACK version is not correct because scipy.fftpack returns the real and imaginary parts in the resulting RFFT array differently than the RFFT from scipy.fft:

  • I believe that RFFT from scipy.fftpack returns an array where one element contains the real part and the next element holds its imaginary counterpart;
  • In RFFT from scipy.fft, each element is a complex number and therefore is able to hold the real and imaginary parts simultaneously;

Please correct me if I'm wrong! I would also like to point out that since scipy.fftpack doesn't provide functions for transforming 2D arrays like rfft2() and irfft2(), I'm providing my own implementations in the code below:

import numpy as np
from scipy import fftpack as scipy_fftpack
from scipy import fft as scipy_fft

# SCIPY RFFT 2D
def scipy_rfft2d(matrix):
    fftRows = [scipy_fft.rfft(row) for row in matrix]
    return np.transpose([scipy_fft.fft(row) for row in np.transpose(fftRows)])

# SCIPY IRFFT 2D
def scipy_irfft2d(matrix, s):
    fftRows = [scipy_fft.irfft(row) for row in matrix]
    return np.transpose([scipy_fft.ifft(row) for row in np.transpose(fftRows)])

# FFTPACK RFFT 2D
def fftpack_rfft2d(matrix):
    fftRows = [scipy_fftpack.rfft(row) for row in matrix]
    return np.transpose([scipy_fftpack.rfft(row) for row in np.transpose(fftRows)])

# FFTPACK IRFFT 2D
def fftpack_irfft2d(matrix):
    fftRows = [scipy_fftpack.irfft(row) for row in matrix]
    return np.transpose([scipy_fftpack.irfft(row) for row in np.transpose(fftRows)])


print('\n####################     INPUT DATA     ###################\n')

# initialize two 2D arrays with random data for testing
in1 = np.array([[0,   0,   0,   0], \
                [0, 255, 255,   0], \
                [0,   0, 255, 255], \
                [0,   0,   0,   0]])

print('\nin1 shape=', in1.shape, '\n', in1)

in2 = np.array([[0,   0,   0,   0], \
                [0,   0, 255,   0], \
                [0, 255, 255,   0], \
                [0, 255,   0,   0]])

print('\nin2 shape=', in2.shape, '\n', in2)

print('\n###############    SCIPY: 2D RFFT (MULT)    ###############\n')

# transform both inputs with SciPy RFFT for 2D
scipy_rfft1 = scipy_fft.rfftn(in1)
scipy_rfft2 = scipy_fft.rfftn(in2)

print('* Output from scipy_fft.rfftn():')
print('scipy_fft1 shape=', scipy_rfft1.shape, '\n', scipy_rfft1.real)
print('\nscipy_fft2 shape=', scipy_rfft2.shape, '\n', scipy_rfft2.real)

# perform multiplication between two 2D arrays from SciPy RFFT
scipy_rfft_mult = scipy_rfft1 * scipy_rfft2

# perform inverse RFFT for 2D arrays using SciPy
scipy_data = scipy_fft.irfftn(scipy_rfft_mult, in1.shape) # passing shape guarantees the output will have the original data size
print('\n* Output from scipy_fft.irfftn():')
print('scipy_data shape=', scipy_data.shape, '\n', scipy_data)

print('\n###############   FFTPACK: 2D RFFT (MULT)   ###############\n')

# transform both inputs with FFTPACK RFFT for 2D
fftpack_rfft1 = fftpack_rfft2d(in1)
fftpack_rfft2 = fftpack_rfft2d(in2)
print('* Output from fftpack_rfft2d():')
print('fftpack_rfft1 shape=', fftpack_rfft1.shape, '\n', fftpack_rfft1)
print('\nfftpack_rfft2 shape=', fftpack_rfft2.shape, '\n', fftpack_rfft2)

# TODO: perform multiplication between two 2D arrays from FFTPACK RFFT
fftpack_rfft_mult = fftpack_rfft1 * fftpack_rfft2 # this doesn't work

# perform inverse RFFT for 2D arrays using FFTPACK
fftpack_data = fftpack_irfft2d(fftpack_rfft_mult)
print('\n* Output from fftpack_irfft2d():')
print('fftpack_data shape=', fftpack_data.shape, '\n', fftpack_data)

print('\n#####################      RESULT     #####################\n')

# compare FFTPACK result with SCIPY
print('\nIs fftpack_data equivalent to scipy_data?', np.allclose(fftpack_data, scipy_data), '\n')

Assuming my guess is correct, what would be the correct implementation for a function that multiplies two 2D arrays that were generated from fftpack_rfft2d()? Remember: the resulting array must be able to be transformed back with fftpack_irfft2d().

Only answers that address the problem in 2-dimensions are invited. Those interested in how to multiply 1D FFTPACK arrays can check this thread.

karlphillip
  • 92,053
  • 36
  • 243
  • 426
  • I suspect the implementation of `fftpack_irfft2d()` is wrong. Any suggestions? – karlphillip May 17 '20 at 20:02
  • It looks like `fftpack` supports arrays, you can specify along which dimension to perform the operation. Thus there is no need to explicitly loop over rows and columns like you do. – Cris Luengo May 18 '20 at 01:04
  • The question is rather difficult, since your 2D transform has groups of two columns that together form a complex value. The complex multiplication would be non-trivial in this case (i.e. element-wise multiplication is incorrect). I think the simplest approach would be to convert the rows to complex values as in the linked Q&A. Also, if you do this before the FFT of columns, you’ll save about half the computation time (fewer columns to transform). – Cris Luengo May 18 '20 at 01:08
  • 1
    The `scipy.fft` output is much more useful. No wonder `fftpack` is deprecated in favor of it. – Cris Luengo May 18 '20 at 01:11
  • Thanks for the suggestion, I understand that. However, I'm a hostage of FFTPACK at the moment. – karlphillip May 18 '20 at 10:11

5 Answers5

4

Correct functions:

import numpy as np
from scipy import fftpack as scipy_fftpack
from scipy import fft as scipy

# FFTPACK RFFT 2D
def fftpack_rfft2d(matrix):
    fftRows = scipy_fftpack.fft(matrix, axis=1)
    fftCols = scipy_fftpack.fft(fftRows, axis=0)

    return fftCols

# FFTPACK IRFFT 2D
def fftpack_irfft2d(matrix):
    ifftRows = scipy_fftpack.ifft(matrix, axis=1)
    ifftCols = scipy_fftpack.ifft(ifftRows, axis=0)

    return ifftCols.real

You calculated the 2D FFT in wrong way. Yes, the first FFT (by columns in your case) can be calculated using rfft(), but the second FFT calculation must be provided on the complex output of the first FFT (by columns), so the output of the rfft() must be converted into true complex spectrum. Moreover, this mean, that you must use fft() instead of rfft() for the second FFT by rows. Consiquently, it is more convenient to use fft() in both calculations.

Moreover, you have input data as a numpy 2D arrays, why do you use list comprehension? Use fftpack.fft() directly, this is much faster.

  • If you already have only 2D arrays calculated by wrong functions and need multiply them: then, my opinion, to try reconstruct the input data from the wrong 2D FFT using the same 'wrong' way and then calculate correct 2D FFT

================================================================

The full testing code with new functions version:

import numpy as np
from scipy import fftpack as scipy_fftpack
from scipy import fft as scipy_fft


# FFTPACK RFFT 2D
def fftpack_rfft2d(matrix):
    fftRows = scipy_fftpack.fft(matrix, axis=1)
    fftCols = scipy_fftpack.fft(fftRows, axis=0)

    return fftCols

# FFTPACK IRFFT 2D
def fftpack_irfft2d(matrix):
    ifftRows = scipy_fftpack.ifft(matrix, axis=1)
    ifftCols = scipy_fftpack.ifft(ifftRows, axis=0)

    return ifftCols.real

print('\n####################     INPUT DATA     ###################\n')

# initialize two 2D arrays with random data for testing
in1 = np.array([[0,   0,   0,   0], \
                [0, 255, 255,   0], \
                [0,   0, 255, 255], \
                [0,   0,   0,   0]])

print('\nin1 shape=', in1.shape, '\n', in1)

in2 = np.array([[0,   0,   0,   0], \
                [0,   0, 255,   0], \
                [0, 255, 255,   0], \
                [0, 255,   0,   0]])

print('\nin2 shape=', in2.shape, '\n', in2)

print('\n###############    SCIPY: 2D RFFT (MULT)    ###############\n')

# transform both inputs with SciPy RFFT for 2D
scipy_rfft1 = scipy_fft.fftn(in1)
scipy_rfft2 = scipy_fft.fftn(in2)

print('* Output from scipy_fft.rfftn():')
print('scipy_fft1 shape=', scipy_rfft1.shape, '\n', scipy_rfft1)
print('\nscipy_fft2 shape=', scipy_rfft2.shape, '\n', scipy_rfft2)

# perform multiplication between two 2D arrays from SciPy RFFT
scipy_rfft_mult = scipy_rfft1 * scipy_rfft2

# perform inverse RFFT for 2D arrays using SciPy
scipy_data = scipy_fft.irfftn(scipy_rfft_mult, in1.shape) # passing shape guarantees the output will
                                                          # have the original data size
print('\n* Output from scipy_fft.irfftn():')
print('scipy_data shape=', scipy_data.shape, '\n', scipy_data)

print('\n###############   FFTPACK: 2D RFFT (MULT)   ###############\n')

# transform both inputs with FFTPACK RFFT for 2D
fftpack_rfft1 = fftpack_rfft2d(in1)
fftpack_rfft2 = fftpack_rfft2d(in2)
print('* Output from fftpack_rfft2d():')
print('fftpack_rfft1 shape=', fftpack_rfft1.shape, '\n', fftpack_rfft1)
print('\nfftpack_rfft2 shape=', fftpack_rfft2.shape, '\n', fftpack_rfft2)

# TODO: perform multiplication between two 2D arrays from FFTPACK RFFT
fftpack_rfft_mult = fftpack_rfft1 * fftpack_rfft2 # this doesn't work

# perform inverse RFFT for 2D arrays using FFTPACK
fftpack_data = fftpack_irfft2d(fftpack_rfft_mult)
print('\n* Output from fftpack_irfft2d():')
print('fftpack_data shape=', fftpack_data.shape, '\n', fftpack_data)

print('\n#####################      RESULT     #####################\n')

# compare FFTPACK result with SCIPY
print('\nIs fftpack_data equivalent to scipy_data?', np.allclose(fftpack_data, scipy_data), '\n')

The output is:

####################     INPUT DATA     ###################


in1 shape= (4, 4) 
 [[  0   0   0   0]
 [  0 255 255   0]
 [  0   0 255 255]
 [  0   0   0   0]]

in2 shape= (4, 4) 
 [[  0   0   0   0]
 [  0   0 255   0]
 [  0 255 255   0]
 [  0 255   0   0]]

###############    SCIPY: 2D RFFT (MULT)    ###############

* Output from scipy_fft.rfftn():
scipy_fft1 shape= (4, 4) 
 [[1020.  -0.j -510.  +0.j    0.  -0.j -510.  -0.j]
 [-510.-510.j    0.  +0.j    0.  +0.j  510.+510.j]
 [   0.  -0.j    0.+510.j    0.  -0.j    0.-510.j]
 [-510.+510.j  510.-510.j    0.  -0.j    0.  -0.j]]

scipy_fft2 shape= (4, 4) 
 [[1020.  -0.j -510.-510.j    0.  -0.j -510.+510.j]
 [-510.  +0.j  510.+510.j    0.-510.j    0.  -0.j]
 [   0.  -0.j    0.  +0.j    0.  -0.j    0.  -0.j]
 [-510.  -0.j    0.  +0.j    0.+510.j  510.-510.j]]

* Output from scipy_fft.irfftn():
scipy_data shape= (4, 4) 
 [[130050.  65025.  65025. 130050.]
 [ 65025.      0.      0.  65025.]
 [ 65025.      0.      0.  65025.]
 [130050.  65025.  65025. 130050.]]

###############   FFTPACK: 2D RFFT (MULT)   ###############

* Output from fftpack_rfft2d():
fftpack_rfft1 shape= (4, 4) 
 [[1020.  -0.j -510.  +0.j    0.  -0.j -510.  +0.j]
 [-510.-510.j    0.  +0.j    0.  +0.j  510.+510.j]
 [   0.  +0.j    0.+510.j    0.  +0.j    0.-510.j]
 [-510.+510.j  510.-510.j    0.  +0.j    0.  +0.j]]

fftpack_rfft2 shape= (4, 4) 
 [[1020.  -0.j -510.-510.j    0.  -0.j -510.+510.j]
 [-510.  +0.j  510.+510.j    0.-510.j    0.  +0.j]
 [   0.  +0.j    0.  +0.j    0.  +0.j    0.  +0.j]
 [-510.  +0.j    0.  +0.j    0.+510.j  510.-510.j]]

* Output from fftpack_irfft2d():
fftpack_data shape= (4, 4) 
 [[130050.+0.j  65025.+0.j  65025.+0.j 130050.+0.j]
 [ 65025.+0.j      0.+0.j      0.+0.j  65025.+0.j]
 [ 65025.+0.j      0.+0.j      0.+0.j  65025.+0.j]
 [130050.+0.j  65025.+0.j  65025.-0.j 130050.+0.j]]

#####################      RESULT     #####################


Is fftpack_data equivalent to scipy_data? True 
Andrei Krivoshei
  • 715
  • 7
  • 16
  • Ok, but you went with `scipy_fftpack` FFT/IFFT instead of `RFFT/IRFFT` which is the focus of the question. Can you test and adjust your answer accordingly? – karlphillip May 18 '20 at 10:06
  • Or are you saying that it can't be done with RFFT/IRFFT? (both functions were completely removed from your code). – karlphillip May 18 '20 at 10:35
  • It can't be solved using only RFFT/IRFFT. You can use RFFT only for transforming the first axis of the input matrix and convert the result of this transformation into true complex spectra (column arrays in your case). All other transformations must be done on the **complex** arrays, thus using FFT. But the IRFFT you cannot use at all, because you have complex arrays, but the IRFFT needs real-valued arrays at its input. Moreover, the fftpack is depricated (like @Cris Luengo already mentioned) and it is better to use *scipy.fft* and *scipy.ifft* – Andrei Krivoshei May 18 '20 at 11:32
  • Moreover, the fftpack is depricated (like @Cris Luengo already mentioned) and it is better to use *scipy.fft()* and *scipy.ifft()* – Andrei Krivoshei May 18 '20 at 11:39
  • As you asked __*"for a function that performs the correct multiplication of two 2D arrays transformed with FFTPACK RFFT"*__, **my answer**, that you don't need separate function for multiplication, because for multiplication element-wise for complex valued matrices is embedded into the numpy multiplication operation. Just use correct functions for transformation into 2D spectrum and back (inverse transform). Better, if you replace *scipy.fftpack* by *scipy.fft". – Andrei Krivoshei May 18 '20 at 12:07
  • Your answer is clear. You will receive the bounty when the time runs out. Thanks for your time. – karlphillip May 18 '20 at 12:19
  • Note that you can use (according to the docs) `scipy.fftpack.fft(x, axis=0)`. Also, if you don’t use `rfft`, then you can use `fft2` directly. – Cris Luengo May 18 '20 at 13:36
  • Although this answer gives the simple solution, it is not true that this cannot be solved using `rfft`. In fact, `rfft` can save you about half the computations to compute the transforms. – Cris Luengo May 18 '20 at 13:38
  • Just found similar question with similar answer: https://stackoverflow.com/a/11333566/501852 – Andrei Krivoshei May 18 '20 at 13:52
  • @CrisLuengo, if it is very critical, then the ```rfft``` of the first axis can be done with converting its result into true **cmplex** spectra of the columns (rows). But all other transformations must be done with **complex** inputs using ```fft``` and ```ifft```. – Andrei Krivoshei May 18 '20 at 14:03
  • @CrisLuengo, about the axis parameter you are right. But 2D FFT operates by complex values, however rfft needs real valued data. – Andrei Krivoshei May 18 '20 at 14:06
  • For the ```scipy.rfftn``` (https://docs.scipy.org/doc/numpy/reference/generated/numpy.fft.rfftn.html#numpy.fft.rfftn): "*By default, all axes are transformed, with the real transform performed over the last axis, while the remaining transforms are complex*", thus only **one** axis, last one, is calculated by the ```rfft```. So, if the author of the question @karlphillip can use of ```scipy.fft``` instead of ```scipy.fftpack```, then it is better variant to use ```scipy.rfft2``` or ```scipy.rfftn``` and inverse version ```scipy.irfft2``` or ```scipy.irfft2```. – Andrei Krivoshei May 18 '20 at 14:12
  • `fftpack.fft2` works perfectly well with a real-valued input. And yes, only the first transform is `rfft`, but this one produces half the values that `fft` produces, and so there will be half the number of columns to apply `fft` to in the 2nd set of transforms. You basically end up with only half of the 2D FFT result, because the other half is redundant. – Cris Luengo May 18 '20 at 14:42
3

Your hypothesis is correct. FFTPACK returns all coefficients in a single real vector in the format

[y(0),Re(y(1)),Im(y(1)),...,Re(y(n/2))]              if n is even
[y(0),Re(y(1)),Im(y(1)),...,Re(y(n/2)),Im(y(n/2))]   if n is odd

where scipy.rfft returns a complex vector

[y(0),Re(y(1)) + 1.0j*Im(y(1)),...,Re(y(n/2) + 1.0j*Im(y(n/2)))]

so you need to form a vector using the proper stride, as follows:

y_fft = np.cat([y_fftpack[0], y_fftpack[1:2:] + 1.0j*y_fftpack[2:2:]])
jfsantos
  • 835
  • 5
  • 20
  • That's kinda helpful, thank you. However, if you could demonstrate how to perform the actual multiplication of those 2D arrays mentioned in the question I would most certainly be happy to award you a bounty since it fully answers the question. – karlphillip May 11 '20 at 08:56
  • I think your current answer is more or less what was discussed [here](https://stackoverflow.com/a/39627135/176769). The challenging part has been to perform the multiplication for 2D arrays (not 1D) without changing it's shape or at least returning to the original shape so that a subsequent call to `fftpack_irfft2d()` would be able to transform it back. – karlphillip May 11 '20 at 09:54
  • NumPy does element-wise multiplication and I've shared the code that does it [here](https://stackoverflow.com/a/61737230/176769). Do you see a way to make this happen without stripping data out of the FFTPACK RRFT matrix? I need to be able to do this to perform the backwards transform. Any tips? – karlphillip May 11 '20 at 20:42
2

@Andrei is right: it's much simpler to just use the complex-valued FFT (though his implementation is unnecessarily complicated, just use scipy.fftpack.fft2). As I said in a comment, the best option is to switch over to scipy.fft, which is simpler to use; fftpack is deprecated in favour of it.

However, if you do need to use fftpack, and you do want to save some computational time by using the rfft function, then this is the right way to do so. It requires converting the real-valued output of the rfft function to a complex-valued array before computing the fft along the other dimension. With this solution, fftpack_rfft2d below outputs half the 2D FFT of its input, with the other half being redundant.

import numpy as np
from scipy import fftpack

# FFTPACK RFFT 2D
def fftpack_rfft1d(matrix):
    assert not (matrix.shape[1] & 0x1)
    tmp = fftpack.rfft(matrix, axis=1)
    assert  tmp.dtype == np.dtype('float64')
    return np.hstack((tmp[:, [0]], np.ascontiguousarray(tmp[:, 1:-1]).view(np.complex128), tmp[:, [-1]]))

def fftpack_rfft2d(matrix):
    return fftpack.fft(fftpack_rfft1d(matrix), axis=0)

# FFTPACK IRFFT 2D
def fftpack_irfft1d(matrix):
    assert  matrix.dtype == np.dtype('complex128')
    tmp = np.hstack((matrix[:, [0]].real, np.ascontiguousarray(matrix[:, 1:-1]).view(np.float64), matrix[:, [-1]].real))
    return fftpack.irfft(tmp, axis=1)

def fftpack_irfft2d(matrix):
    return fftpack_irfft1d(fftpack.ifft(matrix, axis=0))

######

# test data
in1 = np.random.randn(256,256)
in2 = np.random.randn(256,256)

# fftpack.fft2
gt_result = fftpack.ifft2(fftpack.fft2(in1) * fftpack.fft2(in2)).real

# fftpack_rfft2d
our_result = fftpack_irfft2d(fftpack_rfft2d(in1) * fftpack_rfft2d(in2) )

# compare
print('\nIs our result equivalent to the ground truth?', np.allclose(gt_result, our_result), '\n')

[This code only works for even-sized images, I didn't bother making it generic, see here for how to do so).

Nonetheless, as this solution requires copies of the data, it's actually slower than just using a normal, complex-valued FFT (fftpack.fft2), even though it does fewer computations:

import time

tic = time.perf_counter()
for i in range(100):
   fftpack.fft(in1)
toc = time.perf_counter()
print(f"fftpack.fft() takes {toc - tic:0.4f} seconds")

tic = time.perf_counter()
for i in range(100):
   fftpack_rfft2d(in1)
toc = time.perf_counter()
print(f"fftpack_rfft2d() takes {toc - tic:0.4f} seconds")

outputs:

fftpack.fft() takes 0.0442 seconds
fftpack_rfft2d() takes 0.0664 seconds

So, indeed, stick to fftpack.fft (or rather scipy.fft.fft if you can).

Cris Luengo
  • 55,762
  • 10
  • 62
  • 120
1

To multiply 2 arrays of complex coefficients, you have to do a complex multiplication.

See the Multiplication in the Operations section of https://en.m.wikipedia.org/wiki/Complex_number

You can’t just multiply the real components, and then the imaginary components separately or split element wise, which may be why your fftpack matrix mul produces garbage.

hotpaw2
  • 70,107
  • 14
  • 90
  • 153
  • 1
    Last time I needed to do something like this when I only had real numbers, I used an FFT instead of an rfft so the data was wasn’t packed too weirdly, and used 3 nested loops to do the 2D matrix Complex multiply explicitly (Crossing the real and imaginary components). – hotpaw2 May 18 '20 at 00:59
  • Excellent, that makes sense. – karlphillip May 18 '20 at 10:02
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.

Andrei Krivoshei
  • 715
  • 7
  • 16
  • 1
    You should do the "Test fftpack.FFT vs fftpack.RFFT" with a more meaningful transform size. A transform of 5 samples is trivial and not worth using an FFT for (actually, the FFT implants a transform of size 5 as a plain DFT, not recursively dividing the array by half). Try `sz=512` or `sz=2048` or something like that. – Cris Luengo May 18 '20 at 21:57