Search code examples
pythonnumpyfftfftw

numpy.fft numpy.fft2 and FFTW for 2D arrays


I am trying to reproduce the output of numpy.fft.fft and numpy.fft.fft2 using C FFTW library.

>>> b
array([1, 2, 3, 4, 5, 6])
>>> type(b)
<class 'numpy.ndarray'>
>>> b.shape
(6,)
>>> np.fft.fft(b)
array([21.+0.j        , -3.+5.19615242j, -3.+1.73205081j, -3.+0.j        ,
       -3.-1.73205081j, -3.-5.19615242j])

This output can be obtained by :

int N = 10;
double in[] = {1,2,3,4,5,6,0,0,0,0};
fftw_complex *out;
fftw_plan p;
out = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * (N/2 +1));
p = fftw_plan_dft_r2c_1d(6, in, out, FFTW_ESTIMATE);
fftw_execute(p);
fftw_destroy_plan(p);
fftw_free(out);

Similarly, output of 2d array passed to numpy.fft.fft2 can be reproduced :

>>> a
array([[1, 2],
       [3, 4],
       [5, 6]])
>>> a.shape
(3, 2)
>>> np.fft.fft2(a)
array([[21.+0.j        , -3.+0.j        ],
       [-6.+3.46410162j,  0.+0.j        ],
       [-6.-3.46410162j,  0.+0.j        ]])

and the corresponding C++ code is (only one line change)

p = fftw_plan_dft_r2c_2d(3, 2, in, out, FFTW_ESTIMATE);

I have come across a Python code which passes a 2d array to the numpy.fft.fft

>>> a
array([[1, 2],
       [3, 4],
       [5, 6]])
>>> a.shape
(3, 2)
>>> np.fft.fft(a)
array([[ 3.+0.j, -1.+0.j],
       [ 7.+0.j, -1.+0.j],
       [11.+0.j, -1.+0.j]])

I am trying to find out how this can be achieved using FFTW APIs. Any clue on how to reproduce this ? or why does numpy allows 1D Fourier transformation of the matrix/2d array ?


Solution

  • Why does NumPy allow to pass 2-D arrays to the 1-dimensional FFT? The goal is to be able to calculate the FFT of multiple individual 1-D signals at the same time.

    If

    >>> a = np.array([[1, 2], [3, 4], [5, 6]])
    >>> A = np.fft.fft(a)
    

    then the first row of A will be the 1-D FFT of the first row of a. The second row of A is the 1-D FFT of the second row of a and so on.

    This can be verified with

    >>> np.fft.fft(a[0, :])
    array([ 3.+0.j, -1.+0.j])
    
    >>> A[0, :]
    array([ 3.+0.j, -1.+0.j])
    

    or

    >>> np.fft.fft(a[1, :])
    array([ 7.+0.j, -1.+0.j])
    
    >>> A[1, :]
    array([ 7.+0.j, -1.+0.j])
    

    To do the same in FFTW, you could either execute the plan multiple times for the different rows, or use fftw_plan fftw_plan_many_dft.