Skip to content

Lab 20: A Complete Python FFT

Time: ~50 minutes | Prerequisites: Lab 19 | Hardware: Pico 2 (no microphone needed)

Assembly time — and I don't mean the language yet

Echo waving welcome Four labs of pieces, and today they click together into a real FFT. Then we do the part that actually matters: check it against the DFT you already proved correct, and measure what all that cleverness bought. Time to transform!

What You'll Build

A complete iterative radix-2 FFT in about twenty lines — validated against your Lab 15 test suite and benchmarked against the Lab 16 DFT.

Learning Objectives

  • Assemble bit reversal, twiddle tables and butterflies into a working FFT
  • Validate a new implementation against one already proven correct
  • Measure speedup across several sizes
  • Explain why the measured gain exceeds the operation-count prediction
  • Distinguish an algorithm problem from a language problem

Concepts Introduced

ID Concept
393 Iterative FFT
394 Algorithm Assembly
395 Reference Implementation
396 Cross Validation
397 Speedup Factor
398 Correctness Before Speed
399 Function Decomposition

Background

The whole algorithm

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def fft(re, im, rev, tw_re, tw_im):
    n = len(re)

    # Lab 18: reorder once, in place
    for i in range(n):
        j = rev[i]
        if j > i:
            re[i], re[j] = re[j], re[i]
            im[i], im[j] = im[j], im[i]

    # Lab 19: log2(n) stages of butterflies
    half = 1
    while half < n:
        step = n // (half * 2)
        k = 0
        while k < n:
            j = 0
            while j < half:
                wr = tw_re[j * step]
                wi = tw_im[j * step]
                i1, i2 = k + j, k + j + half
                tr = wr * re[i2] - wi * im[i2]
                ti = wr * im[i2] + wi * re[i2]
                ar, ai = re[i1], im[i1]
                re[i1], im[i1] = ar + tr, ai + ti
                re[i2], im[i2] = ar - tr, ai - ti
                j += 1
            k += half * 2
        half *= 2

Three nested loops — stage, block, butterfly — wrapped around the eight lines from Lab 19. That's it. That's the algorithm that made real-time signal processing possible.

Where does step come from?

Echo thinking Early stages use few distinct twiddles spaced far apart in the table; later stages use many, spaced closely. step = n // (half*2) picks the right stride so each stage reads exactly the twiddles it needs from one shared table.

Procedure

Step 1 — Check it against the DFT

Open 20-complete-python-fft.py and run it:

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
# Lab 20: A Complete Python FFT
#
# Time to put it together. Everything here you have already built:
#
#   Lab 17  split evens and odds, recombine with a twiddle
#   Lab 18  bit-reverse once, then work in place; precompute twiddles
#   Lab 19  the butterfly, and how stages are arranged
#
# The whole FFT is about twenty lines. Then we do the important part:
# check it against the DFT we already proved correct in Lab 15, and measure
# how much faster it is than the one we timed in Lab 16.
#
# Correctness first. Speed only counts if the answer is right.

import config
import math
import time


# =========================================================================
# The FFT
# =========================================================================
def make_tables(n):
    """Precompute the bit-reversal permutation and the twiddle factors."""
    bits = 0
    while (1 << bits) < n:
        bits += 1

    rev = []
    for i in range(n):
        r = 0
        x = i
        for _ in range(bits):
            r = (r << 1) | (x & 1)
            x >>= 1
        rev.append(r)

    half = n // 2
    tw_re = [0.0] * half
    tw_im = [0.0] * half
    for k in range(half):
        angle = -2 * math.pi * k / n
        tw_re[k] = math.cos(angle)
        tw_im[k] = math.sin(angle)

    return rev, tw_re, tw_im


def fft(re, im, rev, tw_re, tw_im):
    """In-place iterative radix-2 FFT. Modifies re and im directly."""
    n = len(re)

    # --- Lab 18: reorder once, in place -----------------------------------
    for i in range(n):
        j = rev[i]
        if j > i:
            re[i], re[j] = re[j], re[i]
            im[i], im[j] = im[j], im[i]

    # --- Lab 19: log2(n) stages of butterflies ----------------------------
    half = 1
    while half < n:
        step = n // (half * 2)      # how far apart the twiddles we need are
        k = 0
        while k < n:                # for each block
            j = 0
            while j < half:         # for each butterfly in the block
                wr = tw_re[j * step]
                wi = tw_im[j * step]
                i1 = k + j
                i2 = i1 + half

                tr = wr * re[i2] - wi * im[i2]
                ti = wr * im[i2] + wi * re[i2]

                ar, ai = re[i1], im[i1]
                re[i1] = ar + tr
                im[i1] = ai + ti
                re[i2] = ar - tr
                im[i2] = ai - ti
                j += 1
            k += half * 2
        half *= 2


# =========================================================================
# The DFT from Lab 14, kept as our reference
# =========================================================================
def dft(signal):
    n = len(signal)
    real, imag = [], []
    for k in range(n):
        rr = ii = 0.0
        for t in range(n):
            angle = 2 * math.pi * k * t / n
            rr += signal[t] * math.cos(angle)
            ii -= signal[t] * math.sin(angle)
        real.append(rr)
        imag.append(ii)
    return real, imag


def magnitudes(real, imag):
    return [math.sqrt(real[k] ** 2 + imag[k] ** 2) for k in range(len(real))]


def run_fft(signal):
    n = len(signal)
    rev, tw_re, tw_im = make_tables(n)
    re = list(signal)
    im = [0.0] * n
    fft(re, im, rev, tw_re, tw_im)
    return re, im


# =========================================================================
# PART 1 -- does it agree with the DFT?
# =========================================================================
N = 64
RATE = config.SAMPLE_RATE
BIN = RATE / N
signal = [math.sin(2 * math.pi * 3 * BIN * (i / RATE))
          + 0.5 * math.sin(2 * math.pi * 7 * BIN * (i / RATE))
          for i in range(N)]

dft_mag = magnitudes(*dft(signal))
fft_mag = magnitudes(*run_fft(signal))

print("=== PART 1: FFT vs the DFT we already trust ===")
print("%4s %12s %12s %12s" % ("bin", "DFT", "FFT", "difference"))
worst = 0.0
for k in range(N // 2 + 1):
    d = abs(dft_mag[k] - fft_mag[k])
    if d > worst:
        worst = d
    if dft_mag[k] > 0.5:
        print("%4d %12.4f %12.4f %12.2e" % (k, dft_mag[k], fft_mag[k], d))
print()
print("largest difference across all bins: %.3e" % worst)
peak = max(dft_mag)
print("as a fraction of the peak (%.1f): %.2e" % (peak, worst / peak))
print()
print("AGREEMENT CONFIRMED" if worst / peak < 1e-3 else "*** MISMATCH ***")

# =========================================================================
# PART 2 -- the validation suite from Lab 15, now on the FFT
# =========================================================================
print()
print("=== PART 2: the Lab 15 tests, run against the FFT ===")
REL = 1e-3
checks = []

checks.append(("silence", [0.0] * N,
               lambda m: all(v < REL for v in m)))
checks.append(("constant 0.5", [0.5] * N,
               lambda m: abs(m[0] - 0.5 * N) < REL * 0.5 * N))
imp = [0.0] * N
imp[0] = 1.0
checks.append(("impulse", imp,
               lambda m: all(abs(v - 1.0) < REL for v in m)))
alt = [1.0 if i % 2 == 0 else -1.0 for i in range(N)]
checks.append(("alternating +/-1", alt,
               lambda m: abs(m[N // 2] - N) < REL * N))
checks.append(("sine at bin 5", [math.sin(2 * math.pi * 5 * i / N) for i in range(N)],
               lambda m: abs(m[5] - N / 2) < REL * N / 2))

passed = 0
for name, sig, check in checks:
    ok = check(magnitudes(*run_fft(sig)))
    print("  %-20s %s" % (name, "PASS" if ok else "*** FAIL ***"))
    passed += 1 if ok else 0
print("  %d/%d passed" % (passed, len(checks)))

# =========================================================================
# PART 3 -- how much faster?
# =========================================================================
print()
print("=== PART 3: speed ===")
print("%6s %12s %12s %10s" % ("N", "DFT (ms)", "FFT (ms)", "speedup"))

for n in (32, 64, 128, 256):
    sig = [math.sin(2 * math.pi * 5 * i / n) for i in range(n)]

    start = time.ticks_ms()
    dft(sig)
    dft_ms = time.ticks_diff(time.ticks_ms(), start)

    rev, twr, twi = make_tables(n)
    re = list(sig)
    im = [0.0] * n
    start = time.ticks_ms()
    fft(re, im, rev, twr, twi)
    fft_ms = time.ticks_diff(time.ticks_ms(), start)

    ratio = ("%.0fx" % (dft_ms / fft_ms)) if fft_ms > 0 else ">100x"
    print("%6d %12d %12d %10s" % (n, dft_ms, fft_ms, ratio))

# =========================================================================
# PART 4 -- the deadline, revisited
# =========================================================================
print()
print("=== PART 4: can we hit the deadline now? ===")
n = 512
rev, twr, twi = make_tables(n)
sig = [math.sin(2 * math.pi * 40 * i / n) for i in range(n)]

re = list(sig)
im = [0.0] * n
start = time.ticks_ms()
fft(re, im, rev, twr, twi)
ms = time.ticks_diff(time.ticks_ms(), start)

budget = 512 / RATE * 1000
print("512-point FFT in pure Python : %d ms" % ms)
print("real-time budget             : %.0f ms" % budget)
print("Lab 16's DFT estimate        : 21196 ms")
print()
print("improvement over the DFT     : %.0fx" % (21196 / ms))
if ms > budget:
    print("still over budget by         : %.1fx" % (ms / budget))
    print()
    print("Enormous progress -- but pure Python is still too slow for")
    print("real time. The algorithm is now right; what remains is the")
    print("cost of the LANGUAGE. That is Module 6 and Module 7.")
else:
    print("WE MADE IT with %.0f ms to spare." % (budget - ms))
1
2
3
4
5
6
7
8
 bin          DFT          FFT   difference
   3      32.0000      32.0000     1.91e-06
   7      16.0000      16.0000     1.91e-06

largest difference across all bins: 3.110e-05
as a fraction of the peak (32.0): 9.72e-07

AGREEMENT CONFIRMED

Two completely different algorithms, same answer to seven digits. That agreement is your evidence. A fast FFT that disagrees with a proven DFT is just a fast way to be wrong.

Step 2 — Re-run the Lab 15 suite

1
2
3
4
5
6
  silence              PASS
  constant 0.5         PASS
  impulse              PASS
  alternating +/-1     PASS
  sine at bin 5        PASS
  5/5 passed

The tests you wrote for the DFT work unchanged on the FFT, because both compute the same thing. That's what a good test suite gets you — it outlives the implementation it was written for.

Step 3 — Measure the speedup

1
2
3
4
5
     N     DFT (ms)     FFT (ms)    speedup
    32           60            5        12x
    64          277           11        25x
   128         1218           27        45x
   256         5340           71        75x

Notice the speedup grows with N: 12× at 32, 75× at 256. The DFT gets four times slower per doubling; the FFT only about twice. The gap widens forever.

Step 4 — The deadline, revisited

1
2
3
4
5
6
512-point FFT in pure Python : 145 ms
real-time budget             : 40 ms
Lab 16's DFT estimate        : 21196 ms

improvement over the DFT     : 146x
still over budget by         : 3.6x

From 530× too slow to 3.6× too slow. A 146× improvement from restructuring alone, without a faster chip or a single line of assembly.

And yet — still short.

Being 3.6× away is a completely different problem

Echo encouraging 530× short means your approach is wrong. 3.6× short means your approach is right and something else is costing you. That something is the language: every one of those multiplies is an interpreted MicroPython operation. Modules 6 and 7 go after it, and by Lab 31 this same algorithm runs in under a millisecond.

Step 5 — Predict, then measure

The operation count in Lab 17 predicted 57×. We measured 146×.

Prediction: why is the measured gain more than twice the predicted one?

Think about what else changed between the DFT and the FFT. (Hint: Lab 18 measured a separate 5.5× on one specific thing.)

Expected Output

See the tables above. Your milliseconds will vary; agreement with the DFT and the general shape of the speedup should not.

Troubleshooting

Symptom Likely cause Fix
Peaks in the wrong bins Bit reversal skipped or wrong Reorder before the stages, and use if j > i
Magnitudes right, spectrum scrambled Twiddle stride wrong step = n // (half*2)
Output equals the input Stage loop never ran half must start at 1 and double
Disagrees only at high bins Twiddle table too small You need N/2 entries
MemoryError at N=512 Too many lists alive Reuse buffers; don't rebuild tables per call

Challenges

  1. Inverse FFT. Flip the twiddle sign and divide by N. Transform a signal, invert it, and check you get the original back. Round-trip tests are powerful.
  2. Hoist the tables. The program rebuilds tables on every call. Build them once and reuse. How much does that save at N = 512?
  3. Find the crossover. At what N does the FFT first beat the DFT? Below that size, is the extra complexity worth it?

Check Your Understanding

  1. What are the three nested loops in the FFT, from outside in?
  2. Why validate the FFT against the DFT rather than just checking it looks right?
  3. Why does the speedup grow as N grows?
  4. The count predicted 57× and we measured 146×. Give one reason for the difference.
  5. We're 3.6× from real time. Is that an algorithm problem or something else?

Module 4 complete — you built a real FFT

Echo celebrating From "what's a frequency bin?" in Lab 14 to a working, validated, 146× faster FFT here. Next module points it at the microphone — and you finally get to whistle at your Pico and watch the peak move. That's the payoff.


Next: Lab 21: Spectrum of a Real Sound | Previous: Lab 19