Split-radix FFT in Python
Radix-2/radix-4 hybrid with a low arithmetic count for power-of-two lengths.
- File
split_radix_fft.py, 104 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/split-radix-fft/. The authors are listed in prov/7Q1636.jsonl.
"""Split-radix FFT: X_k = sum_j x_j w**(j*k) for n a power of two.
One half-length transform of the even samples and two quarter-length
transforms of the samples 1 and 3 modulo 4. Over C the root is
w = exp(-2*pi*i/n). Over another ring, pass a table of powers and the
ring operations; the code uses only +, -, * on the entries.
"""
from __future__ import annotations
import cmath
import math
from typing import Any, Callable, Sequence
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 * math.pi * r / n)
def _check(n: int) -> None:
if n < 1 or n & (n - 1):
raise ValueError("the length must be a power of two")
def split_radix_core(x: Sequence[Any], table: Sequence[Any]) -> list[Any]:
"""DFT of x with table[e] = w**e for 0 <= e < n, where w**(n/2) = -1.
Uses only the operators +, -, * of the entries, so it runs over any
commutative ring (and over a counting wrapper). The recursion depth is lg n.
"""
N = len(x)
def rec(start: int, stride: int) -> list[Any]:
n = N // stride # @line L1
if n == 1:
return [x[start]]
if n == 2: # @line L2
a, b = x[start], x[start + stride]
return [a + b, a - b]
U = rec(start, 2 * stride) # @line L3
Z = rec(start + stride, 4 * stride) # @line L4
Zp = rec(start + 3 * stride, 4 * stride) # @line L5
J = table[(n // 4) * stride] # @line L6 w_n**(n/4), -i over C
q = n // 4
out: list[Any] = [None] * n
for k in range(q): # @line L7
if k == 0: # @line L8
a, b = Z[0], Zp[0]
else:
a, b = table[k * stride] * Z[k], table[3 * k * stride] * Zp[k]
s, d = a + b, J * (a - b) # @line L9
out[k], out[k + 2 * q] = U[k] + s, U[k] - s # @line L10
out[k + q], out[k + 3 * q] = U[k + q] + d, U[k + q] - d # @line L11
return out # @line L12
return rec(0, 1)
def split_radix_fft(x: list, w: int | None = None, p: int | None = None, inverse: bool = False) -> list:
"""DFT of x (n a power of two) over C, or modulo a prime p with w**(n/2) = -1 (mod p).
With inverse=True, use w**(-1) and the factor 1/n.
"""
n = len(x)
_check(n)
if p is None:
sign = 1 if inverse else -1
table = [_unit(sign * e, n) for e in range(n)]
y = split_radix_core([complex(v) for v in x], table)
return [v / n for v in y] if inverse else y
if w is None or pow(w, n, p) != 1 or (n > 1 and pow(w, n // 2, p) != p - 1):
raise ValueError("w must satisfy w**(n/2) = -1 modulo p")
base = pow(w, -1, p) if inverse else w
table = [Mod(pow(base, e, p), p) for e in range(n)]
y = [v.v for v in split_radix_core([Mod(v, p) for v in x], table)]
if inverse:
ninv = pow(n, -1, p)
y = [v * ninv % p for v in y]
return y
class Mod:
"""A residue modulo p with +, - and *."""
__slots__ = ("v", "p")
def __init__(self, v: int, p: int) -> None:
self.v, self.p = v % p, p
def __add__(self, o: "Mod") -> "Mod":
return Mod(self.v + o.v, self.p)
def __sub__(self, o: "Mod") -> "Mod":
return Mod(self.v - o.v, self.p)
def __mul__(self, o: "Mod") -> "Mod":
return Mod(self.v * o.v, self.p)
test_split_radix_fft.py 140 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/split-radix-fft/. The authors are listed in prov/7Q1636.jsonl.
import cmath
import math
import random
import pytest
from split_radix_fft import split_radix_core, split_radix_fft
U = 2.0**-53
def dft_mod(x, w, p):
n = len(x)
return [sum(v * pow(w, j * k, p) for j, v in enumerate(x)) % p for k in range(n)]
def dft_c(x):
n = len(x)
return [sum(v * cmath.exp(-2j * math.pi * ((j * k) % n) / n) for j, v in enumerate(x)) for k in range(n)]
def root(p, g, n):
return pow(g, (p - 1) // n, p)
def norm(v):
return math.sqrt(sum(abs(t) ** 2 for t in v))
def tol(n, x):
return 64 * (n.bit_length() + 1) * U * math.sqrt(n) * max(norm(x), 1.0)
MODS = [(17, 3, 16), (97, 5, 32), (998244353, 3, 2**23)]
@pytest.mark.parametrize("p,g,nmax", MODS)
def test_exact_mod_p(p, g, nmax):
rng = random.Random(p)
n = 1
while n <= min(nmax, 256):
w = root(p, g, n)
for _ in range(5):
x = [rng.randrange(p) for _ in range(n)]
X = split_radix_fft(x, w, p)
assert X == dft_mod(x, w, p)
assert split_radix_fft(X, w, p, inverse=True) == x
n *= 2
def test_example_mod_17():
assert split_radix_fft([1, 2, 3, 4, 5, 6, 7, 8], 2, 17) == [2, 8, 14, 6, 13, 3, 12, 1]
def test_large_mod_p_round_trip():
p, g = 998244353, 3
rng = random.Random(7)
n = 2**12
w = root(p, g, n)
x = [rng.randrange(p) for _ in range(n)]
X = split_radix_fft(x, w, p)
assert split_radix_fft(X, w, p, inverse=True) == x
k = rng.randrange(n)
assert X[k] == sum(v * pow(w, j * k, p) for j, v in enumerate(x)) % p
def test_integer_ring():
# Over Z[i], with w = -i for n = 4 and exact Gaussian integers as complex numbers.
rng = random.Random(3)
for _ in range(50):
x = [rng.randint(-100, 100) for _ in range(4)]
assert split_radix_fft(x) == [complex(round(v.real), round(v.imag)) for v in dft_c(x)]
def test_complex_against_direct():
rng = random.Random(11)
for L in range(0, 9):
n = 2**L
for _ in range(4):
x = [complex(rng.uniform(-1, 1), rng.uniform(-1, 1)) for _ in range(n)]
X = split_radix_fft(x)
R = dft_c(x)
assert norm([a - b for a, b in zip(X, R)]) <= tol(n, x)
def test_complex_round_trip_large():
rng = random.Random(13)
for L in (10, 12, 14):
n = 2**L
x = [complex(rng.uniform(-1, 1), rng.uniform(-1, 1)) for _ in range(n)]
y = split_radix_fft(split_radix_fft(x), inverse=True)
assert norm([a - b for a, b in zip(x, y)]) <= tol(n, x)
def test_edge_cases():
assert split_radix_fft([5]) == [5]
assert split_radix_fft([3, 1], 16, 17) == [4, 2]
for n in (1, 2, 4, 8, 64):
assert split_radix_fft([0] * n) == [0] * n
imp = [1] + [0] * (n - 1)
assert split_radix_fft(imp) == [1] * n
const = split_radix_fft([2.5] * n)
assert abs(const[0] - 2.5 * n) <= tol(n, [2.5] * n)
assert norm(const[1:]) <= tol(n, [2.5] * n)
with pytest.raises(ValueError):
split_radix_fft([1, 2, 3])
with pytest.raises(ValueError):
split_radix_fft([1, 2, 3, 4], 2, 17) # 2 has order 8, not 4
C = {"a": 0, "m": 0}
class Count:
def __init__(self, v):
self.v = v
def __add__(self, o):
C["a"] += 1
return Count(self.v + o.v)
def __sub__(self, o):
C["a"] += 1
return Count(self.v - o.v)
def __mul__(self, o):
C["m"] += 1
return Count(self.v * o.v)
def test_operation_counts():
for L in range(0, 13):
n = 2**L
C["a"] = C["m"] = 0
split_radix_core([Count(1)] * n, [Count(1)] * n)
assert C["a"] == n * L
assert 2 * C["m"] == n * L - 2 * n + 2
What the next level needs
- the tests did not run in this build