1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
|
# Copyright (c) 2019 Project Nayuki. (MIT License)
# https://www.nayuki.io/page/free-small-fft-in-multiple-languages
import math, cmath
def transform_radix2(vector, inverse):
# Returns the integer whose value is the reverse of the lowest 'bits' bits of the integer 'x'.
def reverse(x, bits):
y = 0
for i in range(bits):
y = (y << 1) | (x & 1)
x >>= 1
return y
# Initialization
n = len(vector)
levels = int(math.log2(n))
coef = (2 if inverse else -2) * cmath.pi / n
exptable = [cmath.rect(1, i * coef) for i in range(n // 2)]
vector = [vector[reverse(i, levels)] for i in range(n)] # Copy with bit-reversed permutation
# Radix-2 decimation-in-time FFT
size = 2
while size <= n:
halfsize = size // 2
tablestep = n // size
for i in range(0, n, size):
k = 0
for j in range(i, i + halfsize):
temp = vector[j + halfsize] * exptable[k]
vector[j + halfsize] = vector[j] - temp
vector[j] += temp
k += tablestep
size *= 2
return vector
###########################################################################
# Benchmark interface
bm_params = {
(50, 25): (2, 128),
(100, 100): (3, 256),
(1000, 1000): (20, 512),
(5000, 1000): (100, 512),
}
def bm_setup(params):
state = None
signal = [math.cos(2 * math.pi * i / params[1]) + 0j for i in range(params[1])]
fft = None
fft_inv = None
def run():
nonlocal fft, fft_inv
for _ in range(params[0]):
fft = transform_radix2(signal, False)
fft_inv = transform_radix2(fft, True)
def result():
nonlocal fft, fft_inv
fft[1] -= 0.5 * params[1]
fft[-1] -= 0.5 * params[1]
fft_ok = all(abs(f) < 1e-3 for f in fft)
for i in range(len(fft_inv)):
fft_inv[i] -= params[1] * signal[i]
fft_inv_ok = all(abs(f) < 1e-3 for f in fft_inv)
return params[0] * params[1], (fft_ok, fft_inv_ok)
return run, result
|
