Nonuniform FFT in Python
Approximates the NUDFT in near-linear time by gridding, convolving and running an FFT.
- File
nonuniform_fft.py, 115 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/nonuniform-fft/. The authors are listed in prov/KGXG3C.jsonl.
"""Nonuniform FFT of types 1 and 2 in one dimension, with a Gaussian kernel.
Type 1: X_k = sum_j c_j exp(-2*pi*i*k*t_j) for 0 <= k < n.
Type 2: y_j = sum_k f_k exp(-2*pi*i*k*t_j) for 0 <= j < m.
With sign=+1 the exponent has the opposite sign. In exact arithmetic the
error is at most eps * ||input||_1 in every output (see proofs/correctness.md).
"""
from __future__ import annotations
import cmath
import math
from typing import Sequence
TWO_PI = 2.0 * math.pi
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 fft(a: list[complex]) -> list[complex]:
"""Iterative radix-2 FFT, A_r = sum_l a_l exp(-2*pi*i*l*r/N), N a power of two."""
a = list(a)
N = len(a)
j = 0
for i in range(1, N):
bit = N >> 1
while j & bit:
j ^= bit
bit >>= 1
j |= bit
if i < j:
a[i], a[j] = a[j], a[i]
size = 2
while size <= N:
half = size // 2
tw = [_unit(-r, size) for r in range(half)]
for start in range(0, N, size):
for r in range(half):
u, v = a[start + r], tw[r] * a[start + r + half]
a[start + r], a[start + r + half] = u + v, u - v
size *= 2
return a
def parameters(n: int, eps: float) -> tuple[int, int, float]:
"""Return the half-width M, the grid size N and the kernel parameter tau."""
if not 0 < eps < 1:
raise ValueError("eps must lie in (0, 1)")
M = max(2, math.ceil(3.0 / TWO_PI * math.log(4.0 / eps)))
N = 1
while N < max(2 * n, 2 * M):
N *= 2
tau = math.pi * M / (N * (N - n / 2.0))
return M, N, tau
def nufft1(t: Sequence[float], c: Sequence[complex], n: int, eps: float = 1e-12, sign: int = -1) -> list[complex]:
"""Type 1: X_k = sum_j c_j exp(sign*2*pi*i*k*t_j), 0 <= k < n."""
if sign == 1:
return [z.conjugate() for z in nufft1(t, [complex(v).conjugate() for v in c], n, eps)]
if n < 1 or len(t) != len(c):
raise ValueError("need n >= 1 and one weight per point")
M, N, tau = parameters(n, eps) # @line L1
h, s = TWO_PI / N, n // 2
v = [0j] * N # @line L2
for tj, cj in zip(t, c): # @line L3
u = tj % 1.0 # @line L4
x = TWO_PI * u
b = complex(cj) * cmath.exp(-TWO_PI * 1j * s * u)
l0 = math.floor(x / h)
for L in range(l0 - M + 1, l0 + M + 1): # @line L5
v[L % N] += b * math.exp(-((L * h - x) ** 2) / (4.0 * tau))
V = fft(v) # @line L6
scale = math.sqrt(math.pi / tau) / N
X = []
for k in range(n): # @line L7
kappa = k - s # @line L8
X.append(scale * math.exp(tau * kappa * kappa) * V[kappa % N])
return X # @line L9
def nufft2(t: Sequence[float], f: Sequence[complex], eps: float = 1e-12, sign: int = -1) -> list[complex]:
"""Type 2: y_j = sum_k f_k exp(sign*2*pi*i*k*t_j); the transpose of type 1."""
if sign == 1:
return [z.conjugate() for z in nufft2(t, [complex(v).conjugate() for v in f], eps)]
n = len(f)
if n < 1:
raise ValueError("need at least one frequency")
M, N, tau = parameters(n, eps)
h, s = TWO_PI / N, n // 2
scale = math.sqrt(math.pi / tau) / N
a = [0j] * N
for k in range(n): # deconvolve and place on the grid
kappa = k - s
a[kappa % N] += scale * math.exp(tau * kappa * kappa) * complex(f[k])
U = fft(a)
y = []
for tj in t: # interpolate with the same kernel
u = tj % 1.0
x = TWO_PI * u
l0 = math.floor(x / h)
acc = 0j
for L in range(l0 - M + 1, l0 + M + 1):
acc += U[L % N] * math.exp(-((L * h - x) ** 2) / (4.0 * tau))
y.append(acc * cmath.exp(-TWO_PI * 1j * s * u))
return y
test_nonuniform_fft.py 131 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/nonuniform-fft/. The authors are listed in prov/KGXG3C.jsonl.
import cmath
import math
import random
import pytest
from nonuniform_fft import fft, nufft1, nufft2, parameters
U = 2.0**-53
def e(t):
return cmath.exp(-2j * math.pi * t)
def direct1(t, c, n, sign=-1):
return [sum(cj * e(-sign * k * tj) for tj, cj in zip(t, c)) for k in range(n)]
def direct2(t, f, sign=-1):
return [sum(fk * e(-sign * k * tj) for k, fk in enumerate(f)) for tj in t]
def tol(n, eps, v):
N = parameters(n, eps)[1]
return (eps + 256 * (N.bit_length() + 1) * U) * max(sum(abs(z) for z in v), 1e-300)
def err(a, b):
return max((abs(x - y) for x, y in zip(a, b)), default=0.0)
def rand_c(rng, m):
return [complex(rng.gauss(0, 1), rng.gauss(0, 1)) for _ in range(m)]
@pytest.mark.parametrize("eps", [1e-2, 1e-5, 1e-9, 1e-12, 1e-14])
def test_type1_against_direct(eps):
rng = random.Random(int(-math.log10(eps)))
for n in (1, 2, 3, 8, 37, 128):
for m in (1, 5, 60):
t = [rng.uniform(-2, 2) for _ in range(m)]
c = rand_c(rng, m)
assert err(nufft1(t, c, n, eps), direct1(t, c, n)) <= tol(n, eps, c)
@pytest.mark.parametrize("eps", [1e-2, 1e-5, 1e-9, 1e-12, 1e-14])
def test_type2_against_direct(eps):
rng = random.Random(100 + int(-math.log10(eps)))
for n in (1, 2, 3, 8, 37, 128):
for m in (1, 5, 60):
t = [rng.random() for _ in range(m)]
f = rand_c(rng, n)
assert err(nufft2(t, f, eps), direct2(t, f)) <= tol(n, eps, f)
def test_sign_plus():
rng = random.Random(5)
t = [rng.random() for _ in range(30)]
c = rand_c(rng, 30)
assert err(nufft1(t, c, 20, 1e-10, sign=1), direct1(t, c, 20, sign=1)) <= tol(20, 1e-10, c)
f = rand_c(rng, 20)
assert err(nufft2(t, f, 1e-10, sign=1), direct2(t, f, sign=1)) <= tol(20, 1e-10, f)
def test_transpose_identity():
# sum_k X_k f_k = sum_j c_j y_j, since type 2 is the transpose of type 1
rng = random.Random(9)
for n, m in ((16, 40), (33, 7), (1, 3)):
t = [rng.random() for _ in range(m)]
c, f = rand_c(rng, m), rand_c(rng, n)
X, y = nufft1(t, c, n, 1e-12), nufft2(t, f, 1e-12)
lhs = sum(a * b for a, b in zip(X, f))
rhs = sum(a * b for a, b in zip(c, y))
scale = sum(map(abs, c)) * sum(map(abs, f))
assert abs(lhs - rhs) <= 1e-11 * scale
def test_grid_points_give_dft_and_inverse():
rng = random.Random(4)
for n in (1, 2, 4, 6, 16, 50):
t = [j / n for j in range(n)]
x = rand_c(rng, n)
X = nufft1(t, x, n, 1e-12)
D = [sum(x[j] * e(j * k / n) for j in range(n)) for k in range(n)]
assert err(X, D) <= tol(n, 1e-12, x)
back = [v / n for v in nufft2(t, X, 1e-12, sign=1)]
assert err(back, x) <= tol(n, 1e-12, X) / n + 1e-12
def test_edge_cases():
assert nufft1([], [], 5) == [0j] * 5
assert nufft2([], [1, 2, 3]) == []
# an impulse at a single point
X = nufft1([0.3], [1.0], 9, 1e-12)
assert err(X, [e(0.3 * k) for k in range(9)]) <= 2e-12
# zeros
assert err(nufft2([0.1, 0.7], [0, 0, 0]), [0, 0]) == 0.0
# a constant series is constant at every point
y = nufft2([0.0, 0.123, 0.999, 0.5], [2.5], 1e-12)
assert err(y, [2.5] * 4) <= 1e-11
# periodicity of the points
rng = random.Random(1)
c = rand_c(rng, 4)
t = [rng.random() for _ in range(4)]
assert err(nufft1(t, c, 10), nufft1([v + 3 for v in t], c, 10)) <= 1e-11
with pytest.raises(ValueError):
nufft1([0.1], [1.0], 0)
with pytest.raises(ValueError):
parameters(4, 0.0)
def test_inner_fft():
rng = random.Random(2)
for N in (1, 2, 4, 8, 64):
a = rand_c(rng, N)
D = [sum(a[l] * e(l * r / N) for l in range(N)) for r in range(N)]
assert err(fft(a), D) <= 64 * (N.bit_length() + 1) * U * math.sqrt(N) * sum(map(abs, a))
def test_larger_case():
rng = random.Random(77)
n, m = 512, 400
t = [rng.random() for _ in range(m)]
c = rand_c(rng, m)
X = nufft1(t, c, n, 1e-9)
for k in rng.sample(range(n), 20):
assert abs(X[k] - sum(cj * e(k * tj) for tj, cj in zip(t, c))) <= tol(n, 1e-9, c)
What the next level needs
- the tests did not run in this build