Goertzel algorithm in Python
Computes one DFT bin in O(n) time with a second-order recurrence, cheaper than a full FFT.
- File
goertzel_algorithm.py, 49 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/goertzel-algorithm/. The authors are listed in prov/R480RJ.jsonl.
"""One DFT bin X_k = sum_j x_j w**(j*k) by Goertzel's second-order recurrence."""
from __future__ import annotations
import cmath
import math
from typing import TypeVar
S = TypeVar("S", int, complex)
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 goertzel(x: list[S], k: int, w: int | None = None, p: int | None = None) -> S:
"""Return X_k = sum_j x[j] * w**(j*k) for one index k.
Over C (p is None) the root is w = exp(-2*pi*i/n) and w is ignored.
Modulo a prime p, w must be an element with w**n = 1 (mod p); the
entries of x are residues. Any length n >= 1 is allowed.
"""
n = len(x)
if n == 0:
raise ValueError("the input must not be empty")
if p is None:
z = _unit(k, n) # @line L1 z = w**(-k) = exp(2*pi*i*k/n)
c = 2.0 * z.real # @line L2 c = z + 1/z = 2 cos(2*pi*k/n), a real number
s1 = s2 = 0.0 # @line L3
for xj in x: # @line L4
s = xj + c * s1 - s2 # @line L5
s2, s1 = s1, s # @line L6
return z * s1 - s2 # @line L7
if w is None or pow(w, n, p) != 1:
raise ValueError("w must satisfy w**n = 1 modulo p")
z = pow(w, (-k) % n, p)
c = (z + pow(w, k % n, p)) % p
s1 = s2 = 0
for xj in x:
s = (xj + c * s1 - s2) % p
s2, s1 = s1, s
return (z * s1 - s2) % p
test_goertzel_algorithm.py 111 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/goertzel-algorithm/. The authors are listed in prov/R480RJ.jsonl.
import cmath
import math
import random
import pytest
from goertzel_algorithm import goertzel
U = 2.0**-53
# (prime, primitive root): 3 is a primitive root of 17 and 998244353, 5 of 97.
PRIMES = [(17, 3), (97, 5), (998244353, 3)]
def naive_mod(x, k, w, p):
return sum(xj * pow(w, j * k, p) for j, xj in enumerate(x)) % p
def naive_complex(x, k):
n = len(x)
re = math.fsum((xj * cmath.exp(-2j * math.pi * ((j * k) % n) / n)).real for j, xj in enumerate(x))
im = math.fsum((xj * cmath.exp(-2j * math.pi * ((j * k) % n) / n)).imag for j, xj in enumerate(x))
return complex(re, im)
def tolerance(x, k):
"""The proved first-order bound (13u + eta + mu)(n + 2) M ||x||_1, with eta + mu <= 7u."""
n = len(x)
s = abs(math.sin(2 * math.pi * (k % n) / n))
m = n if s == 0 else min(n, 1 / s)
return 20 * U * (n + 2) * m * max(sum(abs(t) for t in x), 1e-300)
def root(p, g, n):
assert (p - 1) % n == 0
return pow(g, (p - 1) // n, p)
def test_example_mod_17():
x = [1, 2, 3, 4, 5, 6, 7, 8]
assert goertzel(x, 1, 2, 17) == 8
assert [goertzel(x, k, 2, 17) for k in range(8)] == [2, 8, 14, 6, 13, 3, 12, 1]
def test_example_complex():
assert abs(goertzel([1, 2, 3, 4], 1) - (-2 + 2j)) < 1e-12
def test_edges_complex():
assert goertzel([5.0], 0) == 5
assert abs(goertzel([3.0, 4.0], 1) - (-1)) < 1e-15
for n in [1, 2, 3, 7, 8, 12]:
for k in range(n):
assert goertzel([0.0] * n, k) == 0
assert abs(goertzel([1.0] + [0.0] * (n - 1), k) - 1) < 1e-12
want = n * 2.5 if k == 0 else 0
assert abs(goertzel([2.5] * n, k) - want) < 1e-12 * n
def test_errors():
with pytest.raises(ValueError):
goertzel([], 0)
with pytest.raises(ValueError):
goertzel([1, 2, 3], 1, 2, 17)
def test_random_mod():
rng = random.Random(20261107)
for p, g in PRIMES:
lengths = [n for n in range(1, 50) if (p - 1) % n == 0]
for n in lengths:
w = root(p, g, n)
for _ in range(5):
x = [rng.randrange(p) for _ in range(n)]
for k in range(n):
assert goertzel(x, k, w, p) == naive_mod(x, k, w, p)
assert goertzel(x, k + 3 * n, w, p) == goertzel(x, k, w, p)
def test_inverse_mod():
rng = random.Random(5)
for p, g in PRIMES:
for n in [n for n in (1, 2, 3, 4, 6, 8, 12, 16) if (p - 1) % n == 0]:
w = root(p, g, n)
x = [rng.randrange(p) for _ in range(n)]
y = [goertzel(x, k, w, p) for k in range(n)]
ninv = pow(n, -1, p)
back = [goertzel(y, j, pow(w, -1, p), p) * ninv % p for j in range(n)]
assert back == x
def test_random_complex_any_length():
rng = random.Random(77)
for n in list(range(1, 40)) + [97, 128, 255, 256, 331]:
x = [complex(rng.uniform(-1, 1), rng.uniform(-1, 1)) for _ in range(n)]
for k in rng.sample(range(n), min(n, 6)) + [0, n - 1]:
assert abs(goertzel(x, k) - naive_complex(x, k)) <= tolerance(x, k)
y = [goertzel(x, k) for k in range(n)] if n <= 64 else None
if y is not None:
back = [goertzel([v.conjugate() for v in y], j).conjugate() / n for j in range(n)]
assert max(abs(a - b) for a, b in zip(back, x)) <= 1e-12 * n
def test_large_real():
rng = random.Random(3)
n = 4096
x = [rng.uniform(-1, 1) for _ in range(n)]
for k in [0, 1, 2, n // 4, n // 2 - 1, n // 2, n - 1]:
assert abs(goertzel(x, k) - naive_complex(x, k)) <= tolerance(x, k)
What the next level needs
- the tests did not run in this build