Convolution via FFT in Python
Computes convolution in O(n log n) by pointwise multiplication in the frequency domain.
- File
convolution_via_fft.py, 87 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/convolution-via-fft/. The authors are listed in prov/XX6TKR.jsonl.
import cmath
from typing import Callable, TypeVar
S = TypeVar("S", int, complex)
def _fft(x: list[S], w: S, red: Callable[[S], S]) -> None:
"""In-place iterative radix-2 FFT, X[k] = sum of x[j] * w**(j*k).
len(x) must be a power of two and w**(len(x)/2) must be -1 (or len(x) = 1).
Every ring result passes through red, the reduction of the ring.
"""
n = len(x)
t = n.bit_length() - 1
table = [red(w**0)] * max(n // 2, 1)
for k in range(1, n // 2):
table[k] = red(table[k - 1] * w) if isinstance(w, int) else w**k
rev = [0] * n
for j in range(1, n):
rev[j] = (rev[j >> 1] >> 1) | ((j & 1) << (t - 1))
for j in range(n):
if j < rev[j]:
x[j], x[rev[j]] = x[rev[j]], x[j]
h = 1
while h < n:
step = n // (2 * h)
for s in range(0, n, 2 * h):
for k in range(h):
u = x[s + k]
v = red(table[k * step] * x[s + k + h])
x[s + k] = red(u + v)
x[s + k + h] = red(u - v)
h *= 2
def _sizes(m: int, l: int) -> tuple[int, int]:
if m < 1 or l < 1:
raise ValueError("both inputs must be non-empty")
out = m + l - 1 # @line L1
size = 1 << (out - 1).bit_length() # @line L2
return out, size
def convolve_mod(a: list[int], b: list[int], p: int, g: int) -> list[int]:
"""Return the linear convolution c[k] = sum of a[i] * b[k - i], modulo p.
Precondition: p is a prime, g is a primitive root modulo p, and the
power of two N >= len(a) + len(b) - 1 divides p - 1 (for p = 998244353
and g = 3, N up to 2**23). Cost: three FFTs of length N and N products,
(3/2) N lg N + 2N multiplications and 3 N lg N additions modulo p.
"""
out, size = _sizes(len(a), len(b))
if (p - 1) % size:
raise ValueError("the transform length must divide p - 1")
w = pow(g, (p - 1) // size, p) # @line L3
red: Callable[[int], int] = lambda v: v % p
fa = [v % p for v in a] + [0] * (size - len(a)) # @line L4
fb = [v % p for v in b] + [0] * (size - len(b)) # @line L5
_fft(fa, w, red) # @line L6
_fft(fb, w, red)
fc = [x * y % p for x, y in zip(fa, fb)] # @line L7
_fft(fc, pow(w, -1, p), red) # @line L8
n_inv = pow(size, -1, p)
return [v * n_inv % p for v in fc[:out]] # @line L9
def convolve_complex(a: list[complex], b: list[complex]) -> list[complex]:
"""Return the linear convolution of a and b, computed in floating point.
Precondition: a and b are non-empty. The result has a rounding error
bounded as on the page; round it when the exact result is an integer
vector and the bound is below 1/2. Cost: three complex FFTs of length
N, the least power of two >= len(a) + len(b) - 1, and N products.
"""
out, size = _sizes(len(a), len(b))
w = cmath.exp(-2j * cmath.pi / size) # @line L3
red: Callable[[complex], complex] = lambda v: v
fa = [complex(v) for v in a] + [0j] * (size - len(a)) # @line L4
fb = [complex(v) for v in b] + [0j] * (size - len(b)) # @line L5
_fft(fa, w, red) # @line L6
_fft(fb, w, red)
fc = [x * y for x, y in zip(fa, fb)] # @line L7
_fft(fc, w.conjugate(), red) # @line L8
return [v / size for v in fc[:out]] # @line L9
test_convolution_via_fft.py 118 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/convolution-via-fft/. The authors are listed in prov/XX6TKR.jsonl.
import cmath
import random
import pytest
from convolution_via_fft import _fft, convolve_complex, convolve_mod
# (prime, primitive root, largest power of two dividing p - 1)
PRIMES = [(17, 3, 16), (97, 5, 32), (998244353, 3, 1 << 23)]
def naive_conv(a: list, b: list) -> list:
c = [0] * (len(a) + len(b) - 1)
for i, x in enumerate(a):
for j, y in enumerate(b):
c[i + j] += x * y
return c
def naive_dft_mod(x: list[int], w: int, p: int) -> list[int]:
n = len(x)
return [sum(x[j] * pow(w, j * k, p) for j in range(n)) % p for k in range(n)]
def close(u: list, v: list, scale: float) -> bool:
n = max(len(u), 1)
tol = 1e-12 * n * max(n.bit_length(), 1) * max(scale, 1.0)
return len(u) == len(v) and all(abs(x - y) <= tol for x, y in zip(u, v))
def test_examples() -> None:
p = 998244353
assert convolve_mod([1, 2, 3], [4, 5], p, 3) == [4, 13, 22, 15]
assert convolve_mod([1, 2, 3, 4], [5, 6, 7, 8], p, 3) == [5, 16, 34, 60, 61, 52, 32]
die = [0, 1, 1, 1, 1, 1, 1]
assert convolve_mod(die, die, 17, 3) == [0, 0, 1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1]
assert convolve_mod([3, 2, 1], [5, 4], 97, 5) == [15, 22, 13, 4]
c = convolve_complex([1, 2, 3], [4, 5])
assert [round(v.real) for v in c] == [4, 13, 22, 15]
def test_internal_fft_against_naive_dft() -> None:
rng = random.Random(30)
for p, g, top in PRIMES:
for t in range(0, 5):
n = 1 << t
w = pow(g, (p - 1) // n, p)
for _ in range(10):
x = [rng.randrange(p) for _ in range(n)]
y = list(x)
_fft(y, w, lambda v: v % p)
assert y == naive_dft_mod(x, w, p)
def test_edges() -> None:
p = 97
assert convolve_mod([7], [5], p, 5) == [35]
assert convolve_mod([0, 0, 0], [0, 0], p, 5) == [0, 0, 0, 0]
a = [3, 1, 4, 1, 5, 9, 2, 6]
assert convolve_mod(a, [1], p, 5) == a
assert convolve_mod(a, [0, 0, 1], p, 5) == [0, 0] + a
assert convolve_mod([2] * 5, [3] * 4, p, 5) == [v % p for v in naive_conv([2] * 5, [3] * 4)]
assert close(convolve_complex([2.5], [-4]), [-10], 10)
assert close(convolve_complex([0, 0], [0]), [0, 0], 1)
def test_errors() -> None:
with pytest.raises(ValueError):
convolve_mod([], [1], 17, 3)
with pytest.raises(ValueError):
convolve_complex([1], [])
with pytest.raises(ValueError):
convolve_mod([1] * 10, [1] * 10, 17, 3) # needs N = 32, which does not divide 16
def test_random_mod_against_naive() -> None:
rng = random.Random(1966)
for _ in range(600):
p, g, top = PRIMES[rng.randrange(len(PRIMES))]
limit = min(top, 64)
m = rng.randint(1, limit // 2)
l = rng.randint(1, limit - m + 1)
a = [rng.randrange(p) for _ in range(m)]
b = [rng.randrange(p) for _ in range(l)]
assert convolve_mod(a, b, p, g) == [v % p for v in naive_conv(a, b)]
def test_random_complex_against_naive() -> None:
rng = random.Random(1970)
for _ in range(500):
m, l = rng.randint(1, 40), rng.randint(1, 40)
scale = rng.choice([1.0, 100.0])
a = [complex(rng.uniform(-scale, scale), rng.uniform(-scale, scale)) for _ in range(m)]
b = [complex(rng.uniform(-scale, scale), rng.uniform(-scale, scale)) for _ in range(l)]
assert close(convolve_complex(a, b), naive_conv(a, b), scale * scale * (m + l))
def test_integer_rounding_exact() -> None:
rng = random.Random(1971)
for _ in range(100):
m, l = rng.randint(1, 200), rng.randint(1, 200)
a = [rng.randint(-999, 999) for _ in range(m)]
b = [rng.randint(-999, 999) for _ in range(l)]
c = convolve_complex(a, b)
assert [round(v.real) for v in c] == naive_conv(a, b)
def test_complex_impulse_and_constant() -> None:
n = 16
w = cmath.exp(-2j * cmath.pi / n)
imp = [1 + 0j] + [0j] * (n - 1)
_fft(imp, w, lambda v: v)
assert close(imp, [1] * n, 1)
c = convolve_complex([1] * 6, [1] * 6)
assert close(c, [1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1], 36)
What the next level needs
- the tests did not run in this build