Monday, June 23, 2014

Fast Fourier Transform (FFT) and Inverse Fast Fourier Transform in Python


This post is an implementation of the fast fourier transform (FFT) and its inverse (IFFT) as outlined in Ch. 32 of "Introduction to Algorithms' by Cormen, Leiserson, and Rivest. The Python source is here. There are four top-level functions: recursive_fft(), recursive_inverse_fft(), recursive_fft_split()recursive_inverse_fft_split(). Let us define a sample array.

>>> a = [8, 4, 8, 0]

Now we can call recursive_fft to obtain the fast fourier transform of this array. The second argument is a function that computes complex n-th roots of unity. There are four functions for computing complex roots of unity: omega_with_error(), omega2_with_error(), omega(), and omega2(). For example, omega_with_error() implements the formula for the n-th complex root of unity on p. 784 of Cormen's Algorithms book. The only addition is the error argument: when the real or imaginary parts of a complex number is less than error, it is set to 0.

def omega_with_error(k, n, error=0):
    if n <= 0:
        return None
        x = cmath.exp((2 * cmath.pi * 1j * k)/n)
        xr = x.real
        xi = x.imag
        if abs(xr) < error: xr = 0
        if abs(xi) < error: xi = 0
        return complex(xr, xi)

The function omega2_with_error() is another formula for the n-th complex root of unity. For example, this formula is used in this FFT calculator. Here is how you can obtain the FFT of the input array. The key "wf" can be read as "omega function."

>>> fft_a = recursive_fft(a, wf=omega_with_error, error=0.001)
>>> fft_a
[20, 4j, 12, -4j]

We can also split the output into the real and imaginary components.

>>> real, imag = recursive_fft_split(fft_a, wf=omega_with_error, error=0.001)
>>> real
[20, 0.0, 12, 0.0]
>>> imag
[0, 4.0, 0, -4.0]

To restore the original array, you can call recursive_inverse_fft():

>>> recursive_inverse_fft(ar, wf=omega_with_error, error=0.001)
[(8+0j), (4+0j), (8+0j), 0j]

The output of the inverse FFT can also be split into real and imaginary components:

>>> recursive_inverse_fft_split(ar, wf=omega_with_error, error=0.001)
([8.0, 4.0, 8.0, 0.0], [0.0, 0.0, 0.0, 0.0])