{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Understanding the FFT Algorithm" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "*This notebook first appeared as a post by Jake Vanderplas on [Pythonic Perambulations](http://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/). The notebook content is BSD-licensed.*" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "The Fast Fourier Transform (FFT) is one of the most important algorithms in signal processing and data analysis. I've used it for years, but having no formal computer science background, It occurred to me this week that I've never thought to ask *how* the FFT computes the discrete Fourier transform so quickly. I dusted off an old algorithms book and looked into it, and enjoyed reading about the deceptively simple computational trick that JW Cooley and John Tukey outlined in their classic [1965 paper](http://www.ams.org/journals/mcom/1965-19-090/S0025-5718-1965-0178586-1/) introducing the subject." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The goal of this post is to dive into the Cooley-Tukey FFT algorithm, explaining the symmetries that lead to it, and to show some straightforward Python implementations putting the theory into practice. My hope is that this exploration will give data scientists like myself a more complete picture of what's going on in the background of the algorithms we use.\n", "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The Discrete Fourier Transform" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The FFT is a fast, $\\mathcal{O}[N\\log N]$ algorithm to compute the Discrete Fourier Transform (DFT), which\n", "naively is an $\\mathcal{O}[N^2]$ computation. The DFT, like the more familiar continuous version of the Fourier transform, has a forward and inverse form which are defined as follows:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Forward Discrete Fourier Transform (DFT):**\n", "$$X_k = \\sum_{n=0}^{N-1} x_n \\cdot e^{-i~2\\pi~k~n~/~N}$$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Inverse Discrete Fourier Transform (IDFT):**\n", "$$x_n = \\frac{1}{N}\\sum_{k=0}^{N-1} X_k e^{i~2\\pi~k~n~/~N}$$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The transformation from $x_n \\to X_k$ is a translation from configuration space to frequency space, and can be very useful in both exploring the power spectrum of a signal, and also for transforming certain problems for more efficient computation. For some examples of this in action, you can check out Chapter 10 of our upcoming Astronomy/Statistics book, with figures and Python source code available [here](http://www.astroml.org/book_figures/chapter10/). For an example of the FFT being used to simplify an otherwise difficult differential equation integration, see my post on [Solving the Schrodinger Equation in Python](http://jakevdp.github.io/blog/2012/09/05/quantum-python/).\n", "\n", "Because of the importance of the FFT in so many fields, Python contains many standard tools and wrappers to compute this. Both NumPy and SciPy have wrappers of the extremely well-tested FFTPACK library, found in the submodules ``numpy.fft`` and ``scipy.fftpack`` respectively. The fastest FFT I am aware of is in the [FFTW](http://www.fftw.org/) package, which is also available in Python via the [PyFFTW](https://pypi.python.org/pypi/pyFFTW) package.\n", "\n", "For the moment, though, let's leave these implementations aside and ask how we might compute the FFT in Python from scratch." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Computing the Discrete Fourier Transform" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "For simplicity, we'll concern ourself only with the forward transform, as the inverse transform can be implemented in a very similar manner. Taking a look at the DFT expression above, we see that it is nothing more than a straightforward linear operation: a matrix-vector multiplication of $\\vec{x}$,\n", "\n", "$$\\vec{X} = M \\cdot \\vec{x}$$\n", "\n", "with the matrix $M$ given by\n", "\n", "$$M_{kn} = e^{-i~2\\pi~k~n~/~N}.$$\n", "\n", "With this in mind, we can compute the DFT using simple matrix multiplication as follows:" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n", "def DFT_slow(x):\n", " \"\"\"Compute the discrete Fourier Transform of the 1D array x\"\"\"\n", " x = np.asarray(x, dtype=float)\n", " N = x.shape[0]\n", " n = np.arange(N)\n", " k = n.reshape((N, 1))\n", " M = np.exp(-2j * np.pi * k * n / N)\n", " return np.dot(M, x)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can double-check the result by comparing to numpy's built-in FFT function:" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false }, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "x = np.random.random(1024)\n", "np.allclose(DFT_slow(x), np.fft.fft(x))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Just to confirm the sluggishness of our algorithm, we can compare the execution times\n", "of these two approaches:" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "10 loops, best of 3: 75.4 ms per loop\n", "10000 loops, best of 3: 25.5 µs per loop\n" ] } ], "source": [ "%timeit DFT_slow(x)\n", "%timeit np.fft.fft(x)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We are over 1000 times slower, which is to be expected for such a simplistic implementation. But that's not the worst of it. For an input vector of length $N$, the FFT algorithm scales as $\\mathcal{O}[N\\log N]$, while our slow algorithm scales as $\\mathcal{O}[N^2]$. That means that for $N=10^6$ elements, we'd expect the FFT to complete in somewhere around 50 ms, while our slow algorithm would take nearly 20 hours!\n", "\n", "So how does the FFT accomplish this speedup? The answer lies in exploiting symmetry." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Symmetries in the Discrete Fourier Transform" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "One of the most important tools in the belt of an algorithm-builder is to exploit symmetries of a problem. If you can show analytically that one piece of a problem is simply related to another, you can compute the subresult\n", "only once and save that computational cost. Cooley and Tukey used exactly this approach in deriving the FFT.\n", "\n", "We'll start by asking what the value of $X_{N+k}$ is. From our above expression:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "$$\n", "\\begin{align*}\n", "X_{N + k} &= \\sum_{n=0}^{N-1} x_n \\cdot e^{-i~2\\pi~(N + k)~n~/~N}\\\\\n", " &= \\sum_{n=0}^{N-1} x_n \\cdot e^{- i~2\\pi~n} \\cdot e^{-i~2\\pi~k~n~/~N}\\\\\n", " &= \\sum_{n=0}^{N-1} x_n \\cdot e^{-i~2\\pi~k~n~/~N}\n", "\\end{align*}\n", "$$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "where we've used the identity $\\exp[2\\pi~i~n] = 1$ which holds for any integer $n$.\n", "\n", "The last line shows a nice symmetry property of the DFT:\n", "\n", "$$X_{N+k} = X_k.$$\n", "\n", "By a simple extension,\n", "\n", "$$X_{k + i \\cdot N} = X_k$$\n", "\n", "for any integer $i$. As we'll see below, this symmetry can be exploited to compute the DFT much more quickly." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## DFT to FFT: Exploiting Symmetry" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Cooley and Tukey showed that it's possible to divide the DFT computation into two smaller parts. From\n", "the definition of the DFT we have:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "$$\n", "\\begin{align}\n", "X_k &= \\sum_{n=0}^{N-1} x_n \\cdot e^{-i~2\\pi~k~n~/~N} \\\\\n", " &= \\sum_{m=0}^{N/2 - 1} x_{2m} \\cdot e^{-i~2\\pi~k~(2m)~/~N} + \\sum_{m=0}^{N/2 - 1} x_{2m + 1} \\cdot e^{-i~2\\pi~k~(2m + 1)~/~N} \\\\\n", " &= \\sum_{m=0}^{N/2 - 1} x_{2m} \\cdot e^{-i~2\\pi~k~m~/~(N/2)} + e^{-i~2\\pi~k~/~N} \\sum_{m=0}^{N/2 - 1} x_{2m + 1} \\cdot e^{-i~2\\pi~k~m~/~(N/2)}\n", "\\end{align}\n", "$$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We've split the single Discrete Fourier transform into two terms which themselves look very similar to smaller Discrete Fourier Transforms, one on the odd-numbered values, and one on the even-numbered values. So far, however, we haven't saved any computational cycles. Each term consists of $(N/2)*N$ computations, for a total of $N^2$.\n", "\n", "The trick comes in making use of symmetries in each of these terms. Because the range of $k$ is $0 \\le k 0:\n", " raise ValueError(\"size of x must be a power of 2\")\n", " elif N 0:\n", " raise ValueError(\"size of x must be a power of 2\")\n", "\n", " # N_min here is equivalent to the stopping condition above,\n", " # and should be a power of 2\n", " N_min = min(N, 32)\n", " \n", " # Perform an O[N^2] DFT on all length-N_min sub-problems at once\n", " n = np.arange(N_min)\n", " k = n[:, None]\n", " M = np.exp(-2j * np.pi * n * k / N_min)\n", " X = np.dot(M, x.reshape((N_min, -1)))\n", "\n", " # build-up each level of the recursive calculation all at once\n", " while X.shape[0]