Good-Thomas prime-factor FFT in Python
Factors n into coprime lengths and uses the Chinese remainder theorem to index the data, avoiding twiddle factors.
- File
chinese_remainder_theorem_fft.py, 71 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/chinese-remainder-theorem-fft/. The authors are listed in prov/XHJRTQ.jsonl.
import math
from itertools import product
from typing import TypeVar
S = TypeVar("S", int, complex)
def _power(w: S, e: int, p: int | None) -> S:
return pow(w, e, p) if p is not None else w**e
def good_thomas(x: list[S], factors: list[int], w: S, p: int | None = None) -> list[S]:
"""Return the DFT X[k] = sum of x[j] * w**(j*k) by the Good-Thomas prime-factor algorithm.
Precondition: the factors are pairwise coprime positive integers whose
product n equals len(x), and w is a principal n-th root of unity, a
complex number such as exp(-2j*pi/n) when p is None, or an integer
modulo the prime p (then all arithmetic is modulo p).
The data are re-indexed by the Chinese remainder theorem so that the DFT
becomes a multidimensional DFT of shape factors, with no twiddle factors.
Each short DFT of length m is done directly here. Cost: n * sum(factors)
ring multiplications and n * sum(factors - 1) additions, plus O(n r)
index operations for r factors.
"""
n = math.prod(factors)
if n != len(x) or any(m < 1 for m in factors):
raise ValueError("the product of the factors must equal len(x)")
for i, a in enumerate(factors):
for b in factors[i + 1:]:
if math.gcd(a, b) != 1:
raise ValueError("the factors must be pairwise coprime")
big = [n // m for m in factors] # @line L1
e = [pow(big[i], -1, m) if m > 1 else 0 for i, m in enumerate(factors)]
a = [x[sum(big_i * j_i for big_i, j_i in zip(big, js)) % n] for js in product(*map(range, factors))] # @line L2
stride = n
for i, m in enumerate(factors): # @line L3
stride //= m
if m == 1:
continue
r = [_power(w, big[i] * q, p) for q in range(m)]
for hi in range(0, n, m * stride):
for lo in range(stride):
base = hi + lo
line = [a[base + t * stride] for t in range(m)]
for k in range(m): # @line L4
acc = line[0]
for j in range(1, m):
acc += line[j] * r[(j * k) % m]
a[base + k * stride] = acc % p if p is not None else acc
out: list[S] = [a[0]] * n
for flat, ks in enumerate(product(*map(range, factors))): # @line L5
out[sum(big_i * e_i * k_i for big_i, e_i, k_i in zip(big, e, ks)) % n] = a[flat]
return out # @line L6
def inverse_good_thomas(big_x: list[S], factors: list[int], w: S, p: int | None = None) -> list[S]:
"""Return x with good_thomas(x, factors, w, p) == big_x.
Same preconditions; n must be a unit (always true for complex numbers,
and for a prime p not dividing n). Cost: one good_thomas and n scalings.
"""
n = len(big_x)
if p is not None:
y = good_thomas(big_x, factors, pow(w, -1, p), p)
inv_n = pow(n, -1, p)
return [v * inv_n % p for v in y]
y = good_thomas(big_x, factors, 1 / w)
return [v / n for v in y]
test_chinese_remainder_theorem_fft.py 124 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/chinese-remainder-theorem-fft/. The authors are listed in prov/XHJRTQ.jsonl.
import cmath
import math
import random
import pytest
from chinese_remainder_theorem_fft import good_thomas, inverse_good_thomas
P = 998244353
PRIMES = {17: 3, 97: 5, P: 3}
def naive_dft(x: list, w, p: int | None = None) -> list:
n = len(x)
if p is not None:
table = [pow(w, m, p) for m in range(n)]
return [sum(x[j] * table[(j * k) % n] for j in range(n)) % p for k in range(n)]
return [sum(x[j] * w ** ((j * k) % n) for j in range(n)) for k in range(n)]
def root(n: int, p: int) -> int:
assert (p - 1) % n == 0
return pow(PRIMES[p], (p - 1) // n, p)
def close(a: list, b: list, scale: float) -> bool:
n = max(len(a), 1)
tol = 1e-12 * n * max(scale, 1.0)
return len(a) == len(b) and all(abs(u - v) <= tol for u, v in zip(a, b))
def test_worked_example() -> None:
x = [1, 2, 3, 4, 5, 6]
assert good_thomas(x, [2, 3], 3, 7) == [0, 3, 6, 4, 2, 5] == naive_dft(x, 3, 7)
assert inverse_good_thomas([0, 3, 6, 4, 2, 5], [2, 3], 3, 7) == x
MOD_SHAPES = {
17: [[16], [1, 16], [2], [1], [4, 1]],
97: [[3, 2], [3, 4], [3, 8], [3, 32], [32, 3], [1, 3]],
P: [[7, 8], [8, 7], [7, 16], [17, 8], [7, 17], [4, 7, 17], [7, 17, 2], [1, 1, 7]],
}
def test_modular_against_oracle() -> None:
rng = random.Random(1958)
cases = 0
for p, shapes in MOD_SHAPES.items():
for shape in shapes:
n = math.prod(shape)
w = root(n, p)
for _ in range(40 if n <= 64 else 5):
x = [rng.randrange(p) for _ in range(n)]
y = good_thomas(x, shape, w, p)
assert y == naive_dft(x, w, p)
assert inverse_good_thomas(y, shape, w, p) == x
cases += 1
assert cases >= 450
COMPLEX_SHAPES = [[2, 3], [3, 2], [2, 5], [3, 4], [4, 3], [3, 5], [5, 7], [2, 3, 5], [3, 5, 7], [4, 9], [5, 8, 9], [11], [2, 3, 5, 7]]
def test_complex_against_oracle() -> None:
rng = random.Random(1963)
cases = 0
for shape in COMPLEX_SHAPES:
n = math.prod(shape)
w = cmath.exp(-2j * math.pi / n)
for _ in range(20 if n <= 40 else 3):
x = [complex(rng.uniform(-5, 5), rng.uniform(-5, 5)) for _ in range(n)]
y = good_thomas(x, shape, w)
scale = 5 * math.sqrt(n) * len(shape)
assert close(y, naive_dft(x, w), scale)
assert close(inverse_good_thomas(y, shape, w), x, scale)
cases += 1
assert cases >= 150
def test_edge_cases() -> None:
assert good_thomas([5], [1], 1, 17) == [5]
assert good_thomas([3, 4], [2], 16, 17) == [7, 16]
for p, shape in ((97, [3, 4]), (P, [7, 8]), (17, [16])):
n = math.prod(shape)
w = root(n, p)
assert good_thomas([0] * n, shape, w, p) == [0] * n
assert good_thomas([1] + [0] * (n - 1), shape, w, p) == [1] * n
assert good_thomas([2] * n, shape, w, p) == [2 * n % p] + [0] * (n - 1)
for s in range(n):
imp = [0] * n
imp[s] = 1
assert good_thomas(imp, shape, w, p) == naive_dft(imp, w, p)
w = cmath.exp(-2j * math.pi / 15)
assert close(good_thomas([1.0] * 15, [3, 5], w), [15.0] + [0.0] * 14, 15)
def test_index_maps_are_bijections() -> None:
for shape in ([2, 3], [3, 5, 7], [4, 9], [5, 8, 9]):
n = math.prod(shape)
big = [n // m for m in shape]
e = [pow(b, -1, m) for b, m in zip(big, shape)]
ins = set()
outs = set()
for flat in range(n):
digits, rest = [], flat
for m in reversed(shape):
digits.append(rest % m)
rest //= m
digits.reverse()
ins.add(sum(b * d for b, d in zip(big, digits)) % n)
k = sum(b * f * d for b, f, d in zip(big, e, digits)) % n
assert all(k % m == d for m, d in zip(shape, digits))
outs.add(k)
assert ins == outs == set(range(n))
def test_bad_input() -> None:
with pytest.raises(ValueError):
good_thomas([1, 2, 3, 4], [2, 2], 13, 17)
with pytest.raises(ValueError):
good_thomas([1, 2, 3], [2, 3], 1, 17)
What the next level needs
- the tests did not run in this build