Iterative FFT in Python
Non-recursive fast Fourier transform that combines values in place with butterfly operations, in loops over stages, and moves the bit-reversal permutation to the input or the output or removes it.
- File
iterative_fft.py, 130 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/iterative-fft/. The authors are listed in prov/37TN4K.jsonl.
import cmath
from typing import TypeVar
S = TypeVar("S", int, complex)
def _check_length(n: int) -> int:
"""Return t = lg n, or raise ValueError when n is not a power of two."""
if n < 1 or n & (n - 1):
raise ValueError("the length must be a power of two")
return n.bit_length() - 1
def _unit(r: int, n: int) -> complex:
"""Return exp(2*pi*i*r/n), exact when r/n is a multiple of 1/4."""
r %= n
if (4 * r) % n == 0:
return (1 + 0j, 1j, -1 + 0j, -1j)[4 * r // n]
return cmath.exp(2j * cmath.pi * r / n)
def twiddles(w: S, n: int, p: int | None = None) -> list[S]:
"""Return the table [w**k for k < n/2], modulo p when p is given.
Over the complex numbers, w must be close to exp(2*pi*i*m/n) for an
integer m, and the entries are taken from the exact angle 2*pi*m*k/n
rather than by repeated multiplication, whose error grows with k.
"""
if p is not None:
return [pow(w, k, p) for k in range(n // 2)]
m = round(cmath.phase(w) * n / (2 * cmath.pi)) % n
if abs(w - _unit(m, n)) > 1e-9:
raise ValueError("w is not close to a root of unity of order dividing n")
return [_unit(m * k, n) for k in range(n // 2)]
def bit_reverse(x: list[S]) -> None:
"""Permute x in place so that x[j] and x[rev(j)] trade places."""
n = len(x)
t = _check_length(n)
rev = [0] * n
for j in range(1, n):
rev[j] = (rev[j >> 1] >> 1) | ((j & 1) << (t - 1))
for j in range(n): # @line L11
if j < rev[j]: # @line L12
x[j], x[rev[j]] = x[rev[j]], x[j]
def dif_stages(x: list[S], table: list[S], p: int | None = None) -> None:
"""Decimation in frequency: natural order in, DFT in bit-reversed order out."""
n = len(x)
_check_length(n)
h = n // 2 # @line L3
while h >= 1: # @line L4
step = n // (2 * h)
for s in range(0, n, 2 * h): # @line L5
for k in range(h): # @line L6
a, b = x[s + k], x[s + k + h] # @line L7
if p is None:
x[s + k] = a + b # @line L8
x[s + k + h] = table[k * step] * (a - b) # @line L9
else:
x[s + k] = (a + b) % p # @line L8
x[s + k + h] = table[k * step] * (a - b) % p # @line L9
h //= 2 # @line L10
def dit_stages(y: list[S], table: list[S], p: int | None = None) -> None:
"""Decimation in time without the permutation: bit-reversed order in, DFT out."""
n = len(y)
_check_length(n)
h = 1 # @line L16
while h < n: # @line L17
step = n // (2 * h)
for s in range(0, n, 2 * h): # @line L18
for k in range(h): # @line L19
u = y[s + k] # @line L20
v = table[k * step] * y[s + k + h] # @line L20
if p is not None:
v %= p
y[s + k], y[s + k + h] = (u + v) % p, (u - v) % p # @line L21
else:
y[s + k], y[s + k + h] = u + v, u - v # @line L21
h *= 2 # @line L22
def iterative_fft(x: list[S], w: S, p: int | None = None) -> None:
"""Replace x by its DFT, X[k] = sum of x[j] * w**(j*k), in place.
Precondition: n = len(x) is a power of two and w is a principal n-th root
of unity (complex, or an integer modulo the prime p). Cost: n lg n
additions and (n/2) lg n multiplications, n/2 powers for the table, and
O(n) index operations for the bit reversal.
"""
n = len(x) # @line L1
_check_length(n)
table = twiddles(w, n, p) # @line L2
dif_stages(x, table, p)
bit_reverse(x)
def cyclic_convolution(a: list[S], b: list[S], w: S, p: int | None = None) -> list[S]:
"""Return c with c[m] = sum of a[j] * b[(m - j) mod n], with no bit reversal.
Same preconditions as iterative_fft, and n must be a unit (always true
over the complex numbers, and modulo a prime p > n). Cost: 3 n lg n
additions and (3/2) n lg n + 2n multiplications.
"""
n = len(a)
if len(b) != n:
raise ValueError("the two inputs must have the same length")
_check_length(n)
fa, fb = list(a), list(b)
table = twiddles(w, n, p)
dif_stages(fa, table, p) # @line L13
dif_stages(fb, table, p) # @line L13
if p is None:
c = [u * v for u, v in zip(fa, fb)] # @line L14
inv_table = twiddles(1 / w, n) # @line L15
else:
c = [u * v % p for u, v in zip(fa, fb)] # @line L14
inv_table = twiddles(pow(w, -1, p), n, p) # @line L15
dit_stages(c, inv_table, p)
if p is None:
return [v / n for v in c] # @line L23
n_inv = pow(n, -1, p)
return [v * n_inv % p for v in c] # @line L23
test_iterative_fft.py 160 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/iterative-fft/. The authors are listed in prov/37TN4K.jsonl.
import cmath
import random
import pytest
from iterative_fft import bit_reverse, cyclic_convolution, dif_stages, dit_stages, iterative_fft, twiddles
# (prime, primitive root): 3 is a primitive root of 17 and 998244353, 5 of 97.
PRIMES = [(17, 3), (97, 5), (998244353, 3)]
U = 2.0**-53
def naive_dft(x: list, w, p: int | None = None) -> list:
n = len(x)
if p is None:
return [sum(x[j] * w ** (j * k) for j in range(n)) for k in range(n)]
return [sum(x[j] * pow(w, j * k, p) for j in range(n)) % p for k in range(n)]
def naive_conv(a: list, b: list, p: int | None = None) -> list:
n = len(a)
c = [sum(a[j] * b[(m - j) % n] for j in range(n)) for m in range(n)]
return c if p is None else [v % p for v in c]
def rev(j: int, t: int) -> int:
return int(format(j, f"0{t}b")[::-1], 2) if t else 0
def root_mod(n: int, p: int, g: int) -> int:
assert (p - 1) % n == 0
return pow(g, (p - 1) // n, p)
def root_c(n: int) -> complex:
return cmath.exp(-2j * cmath.pi / n)
def norm(v: list) -> float:
return sum(abs(c) ** 2 for c in v) ** 0.5
def close(a: list, b: list, scale: float = 1.0) -> bool:
n = max(len(a), 1)
tol = 64 * (n.bit_length() + 1) * U * max(norm(b), scale, 1.0)
return norm([u - v for u, v in zip(a, b)]) <= tol
def test_examples() -> None:
x = [1, 2, 3, 4, 5, 6, 7, 8]
dif_stages(x, twiddles(2, 8, 17), 17)
assert x == [2, 13, 14, 12, 8, 3, 6, 1]
bit_reverse(x)
assert x == [2, 8, 14, 6, 13, 3, 12, 1]
assert cyclic_convolution([1, 2, 0, 0], [3, 4, 0, 0], 4, 17) == [3, 10, 8, 0]
assert cyclic_convolution([1, 2, 3, 4], [1, 1, 0, 0], 4, 17) == [5, 3, 5, 7]
z: list[complex] = [1, 2, 3, 4]
iterative_fft(z, root_c(4))
assert z == [10, -2 + 2j, -2, -2 - 2j]
def test_orders() -> None:
# DIF writes X[rev(i)] at position i; DIT stages read that order.
rng = random.Random(4)
for p, g in PRIMES:
for t in range(0, 5):
n = 1 << t
w = root_mod(n, p, g)
x = [rng.randrange(p) for _ in range(n)]
ref = naive_dft(x, w, p)
y = list(x)
dif_stages(y, twiddles(w, n, p), p)
assert y == [ref[rev(i, t)] for i in range(n)]
z = [x[rev(i, t)] for i in range(n)]
dit_stages(z, twiddles(w, n, p), p)
assert z == ref
def test_edges() -> None:
for p, g in PRIMES:
for t in range(0, 5):
n = 1 << t
w = root_mod(n, p, g)
zeros = [0] * n
iterative_fft(zeros, w, p)
assert zeros == [0] * n
for pos in range(n):
imp = [int(j == pos) for j in range(n)]
iterative_fft(imp, w, p)
assert imp == [pow(w, pos * k, p) for k in range(n)]
const = [7] * n
iterative_fft(const, w, p)
assert const == [7 * n % p] + [0] * (n - 1)
delta = [1] + [0] * (n - 1)
b = [rng_val % p for rng_val in range(3, 3 + n)]
assert cyclic_convolution(delta, b, w, p) == b
for n in (1, 2, 4, 8):
c: list[complex] = [1] + [0] * (n - 1)
iterative_fft(c, root_c(n))
assert close(c, [1] * n)
def test_not_power_of_two() -> None:
for n in (0, 3, 6, 10):
with pytest.raises(ValueError):
iterative_fft([1] * n, 1, 17)
with pytest.raises(ValueError):
cyclic_convolution([1, 2], [1, 2, 3, 4], 4, 17)
def test_random_mod() -> None:
rng = random.Random(1965)
for _ in range(500):
p, g = PRIMES[rng.randrange(len(PRIMES))]
t = rng.randint(0, 6 if p > 97 else (5 if p == 97 else 4))
n = 1 << t
w = root_mod(n, p, g)
x = [rng.randrange(p) for _ in range(n)]
y = list(x)
iterative_fft(y, w, p)
assert y == naive_dft(x, w, p)
b = [rng.randrange(p) for _ in range(n)]
assert cyclic_convolution(x, b, w, p) == naive_conv(x, b, p)
def test_random_complex() -> None:
rng = random.Random(1966)
for _ in range(300):
n = 1 << rng.randint(0, 6)
s = rng.choice([1.0, 1e3, 1e-3])
x = [complex(rng.uniform(-s, s), rng.uniform(-s, s)) for _ in range(n)]
w = root_c(n)
y = list(x)
iterative_fft(y, w)
assert close(y, naive_dft(x, w))
b = [complex(rng.uniform(-1, 1), rng.uniform(-1, 1)) for _ in range(n)]
c = cyclic_convolution(x, b, w)
assert close(c, naive_conv(x, b), norm(x) * norm(b))
def test_integer_product() -> None:
# Polynomial product with zero padding, exact modulo 998244353.
p, g = 998244353, 3
rng = random.Random(30)
for _ in range(40):
da, db = rng.randint(0, 30), rng.randint(0, 30)
a = [rng.randrange(1000) for _ in range(da + 1)]
b = [rng.randrange(1000) for _ in range(db + 1)]
n = 1
while n < da + db + 1:
n *= 2
c = cyclic_convolution(a + [0] * (n - len(a)), b + [0] * (n - len(b)), root_mod(n, p, g), p)
ref = [0] * (da + db + 1)
for i, u in enumerate(a):
for j, v in enumerate(b):
ref[i + j] += u * v
assert c[: da + db + 1] == ref and all(v == 0 for v in c[da + db + 1 :])
What the next level needs
- the tests did not run in this build