DFT Example

The second edition of Think DSP is not for sale yet, but if you would like to support this project, you can buy me a coffee.

DFT Example#

Click here to run this notebook on Colab.

Hide code cell content

 
import importlib, sys
sys.modules["imp"] = importlib

%load_ext autoreload
%autoreload 2

Hide code cell content

 
# download thinkdsp.py

from os.path import basename, exists

def download(url):
    filename = basename(url)
    if not exists(filename):
        from urllib.request import urlretrieve
        local, _ = urlretrieve(url, filename)
        print('Downloaded ' + local)
        
download("https://github.com/AllenDowney/ThinkDSP2/raw/main/soln/thinkdsp.py")

Downloaded thinkdsp.py

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])

Literal translation#

The usual way the DFT is expressed is as a summation. Following the notation on Wikipedia:

\( X_k = \sum_{n=0}^N x_n \cdot e^{-2 \pi i n k / N} \)

Here’s a straightforward translation of that formula into Python.

pi = np.pi
exp = np.exp

k = 0
sum(x[n] * exp(-2j * pi * k * n / N) for n in range(N))

np.complex128(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)

[np.complex128(1+0j),
 np.complex128(1+0j),
 np.complex128(1+0j),
 np.complex128(1+0j)]

And the results check out.

DFT as a matrix operation#

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])

Implementing FFT#

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])

Think DSP: Digital Signal Processing in Python, 2nd Edition

Copyright 2024 Allen B. Downey

License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International

Contents