-
-
Save tianqig/0d5df70fd50a4f9139a26ff6406e0d04 to your computer and use it in GitHub Desktop.
[python]DFT(discrete fourier transform) and FFT
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""DFT and FFT""" | |
import math | |
def iexp(n): | |
return complex(math.cos(n), math.sin(n)) | |
def is_pow2(n): | |
return False if n == 0 else (n == 1 or is_pow2(n >> 1)) | |
def dft(xs): | |
"naive dft" | |
n = len(xs) | |
return [sum((xs[k] * iexp(-2 * math.pi * i * k / n) for k in range(n))) | |
for i in range(n)] | |
def dftinv(xs): | |
"naive dft" | |
n = len(xs) | |
return [sum((xs[k] * iexp(2 * math.pi * i * k / n) for k in range(n))) / n | |
for i in range(n)] | |
def fft_(xs, n, start=0, stride=1): | |
"cooley-turkey fft" | |
if n == 1: return [xs[start]] | |
hn, sd = n // 2, stride * 2 | |
rs = fft_(xs, hn, start, sd) + fft_(xs, hn, start + stride, sd) | |
for i in range(hn): | |
e = iexp(-2 * math.pi * i / n) | |
rs[i], rs[i + hn] = rs[i] + e * rs[i + hn], rs[i] - e * rs[i + hn] | |
pass | |
return rs | |
def fft(xs): | |
assert is_pow2(len(xs)) | |
return fft_(xs, len(xs)) | |
def fftinv_(xs, n, start=0, stride=1): | |
"cooley-turkey fft" | |
if n == 1: return [xs[start]] | |
hn, sd = n // 2, stride * 2 | |
rs = fftinv_(xs, hn, start, sd) + fftinv_(xs, hn, start + stride, sd) | |
for i in range(hn): | |
e = iexp(2 * math.pi * i / n) | |
rs[i], rs[i + hn] = rs[i] + e * rs[i + hn], rs[i] - e * rs[i + hn] | |
pass | |
return rs | |
def fftinv(xs): | |
assert is_pow2(len(xs)) | |
n = len(xs) | |
return [v / n for v in fftinv_(xs, n)] | |
if __name__ == "__main__": | |
wave = [0, 1, 2, 3, 3, 2, 1, 0] | |
dfreq = dft(wave) | |
ffreq = fft(wave) | |
dwave = dftinv(dfreq) | |
fwave= fftinv(ffreq) | |
print(dfreq) | |
print(ffreq) | |
print([v.real for v in dwave]) | |
print([v.real for v in fwave]) | |
pass |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment