Copyright 2019 Allen Downey, MIT License
import numpy as np
Let's start with a known result. The DFT of an impulse is a constant.
N = 4
x = [1, 0, 0, 0]
np.fft.fft(x)
array([1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j])
pi = np.pi
exp = np.exp
k = 0
sum(x[n] * exp(-2j * pi * k * n / N) for n in range(N))
(1+0j)
Wrapping this code in a function makes the roles of k
and n
clearer, where k
is a free parameter and n
is the bound variable of the summation.
def dft_k(x, k):
return sum(x[n] * exp(-2j * pi * k * n / N) for n in range(N))
Of course, we usually we compute $X$ all at once, so we can wrap this function in another function:
def dft(x):
N = len(x)
X = [dft_k(x, k) for k in range(N)]
return X
dft(x)
[(1+0j), (1+0j), (1+0j), (1+0j)]
And the results check out.
It is also common to express the DFT as a matrix operation:
$ X = W x $
with
$ W_{j, k} = \omega^{j k} $
and
$ \omega = e^{-2 \pi i / N}$
If we recognize the construction of $W$ as an outer product, we can use np.outer
to compute it.
Here's an implementation of DFT using outer product to construct the DFT matrix, and dot product to compute the DFT.
def dft(x):
N = len(x)
ks = range(N)
args = -2j * pi * np.outer(ks, ks) / N
W = exp(args)
X = W.dot(x)
return X
And the results check out.
dft(x)
array([1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j])
Finally, we can implement the FFT by translating from math notation:
$ X_k = E_k + e^{-2 \pi i k / N} O_k $
Where $E$ is the FFT of the even elements of $x$ and $O$ is the DFT of the odd elements of $x$.
Here's what that looks like in code.
def fft(x):
N = len(x)
if N == 1:
return x
E = fft(x[::2])
O = fft(x[1::2])
ks = np.arange(N)
args = -2j * pi * ks / N
W = np.exp(args)
return np.tile(E, 2) + W * np.tile(O, 2)
The length of $E$ and $O$ is half the length of $W$, so I use np.tile
to double them up.
And the results check out.
fft(x)
array([1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j])