From e17ee5768648d84979c54fba62a815e1796196dd Mon Sep 17 00:00:00 2001 From: fanninpm Date: Mon, 4 May 2020 18:50:09 -0400 Subject: [PATCH] Make fft.py conform to PEP8 --- contents/cooley_tukey/code/python/fft.py | 76 +++++++++++++----------- contents/cooley_tukey/cooley_tukey.md | 4 +- 2 files changed, 43 insertions(+), 37 deletions(-) diff --git a/contents/cooley_tukey/code/python/fft.py b/contents/cooley_tukey/code/python/fft.py index e7206e35f..6f4fe3735 100644 --- a/contents/cooley_tukey/code/python/fft.py +++ b/contents/cooley_tukey/code/python/fft.py @@ -2,56 +2,62 @@ from cmath import exp, pi from math import log2 + def dft(X): N = len(X) - temp = [0]*N + temp = [0] * N for i in range(N): for k in range(N): - temp[i] += X[k] * exp(-2.0j*pi*i*k/N) + temp[i] += X[k] * exp(-2.0j * pi * i * k / N) return temp + def cooley_tukey(X): - N = len(X) - if N <= 1: - return X - even = cooley_tukey(X[0::2]) - odd = cooley_tukey(X[1::2]) - - temp = [i for i in range(N)] - for k in range(N//2): - temp[k] = even[k] + exp(-2j*pi*k/N) * odd[k] - temp[k+N//2] = even[k] - exp(-2j*pi*k/N)*odd[k] - return temp - -def bitReverse(X): - N = len(X) - temp = [i for i in range(N)] - for k in range(N): - b = sum(1<<(int(log2(N))-1-i) for i in range(int(log2(N))) if k>>i&1) - temp[k] = X[b] - temp[b] = X[k] - return temp + N = len(X) + if N <= 1: + return X + even = cooley_tukey(X[0::2]) + odd = cooley_tukey(X[1::2]) + + temp = [i for i in range(N)] + for k in range(N // 2): + temp[k] = even[k] + exp(-2.0j * pi * k / N) * odd[k] + temp[k + N // 2] = even[k] - exp(-2.0j * pi * k / N) * odd[k] + return temp + + +def bit_reverse(X): + N = len(X) + temp = [i for i in range(N)] + for k in range(N): + b = sum(1 << int(log2(N)) - 1 - + i for i in range(int(log2(N))) if k >> i & 1) + temp[k] = X[b] + temp[b] = X[k] + return temp + def iterative_cooley_tukey(X): - N = len(X) + N = len(X) + + X = bit_reverse(X) - X = bitReverse(X) + for i in range(1, int(log2(N)) + 1): + stride = 2 ** i + w = exp(-2.0j * pi / stride) + for j in range(0, N, stride): + v = 1 + for k in range(stride // 2): + X[k + j + stride // 2] = X[k + j] - v * X[k + j + stride // 2] + X[k + j] -= X[k + j + stride // 2] - X[k + j] + v *= w + return X - for i in range(1, int(log2(N)) + 1): - stride = 2**i - w = exp(-2j*pi/stride) - for j in range(0, N, stride): - v = 1 - for k in range(stride//2): - X[k + j + stride//2] = X[k + j] - v*X[k + j + stride//2]; - X[k + j] -= (X[k + j + stride//2] - X[k + j]); - v *= w; - return X X = [] for i in range(64): - X.append(random()) + X.append(random()) Y = cooley_tukey(X) Z = iterative_cooley_tukey(X) diff --git a/contents/cooley_tukey/cooley_tukey.md b/contents/cooley_tukey/cooley_tukey.md index 1b7066063..b40923ca5 100644 --- a/contents/cooley_tukey/cooley_tukey.md +++ b/contents/cooley_tukey/cooley_tukey.md @@ -80,7 +80,7 @@ For some reason, though, putting code to this transformation really helped me fi {% sample lang="hs" %} [import:7-13, lang:"haskell"](code/haskell/fft.hs) {% sample lang="py" %} -[import:5-11, lang:"python"](code/python/fft.py) +[import:6-12, lang:"python"](code/python/fft.py) {% sample lang="scratch" %} [import:4-13, lang:"julia"](code/julia/fft.jl) {% sample lang="asm-x64" %} @@ -131,7 +131,7 @@ In the end, the code looks like: {% sample lang="hs" %} [import:15-28, lang:"haskell"](code/haskell/fft.hs) {% sample lang="py" %} -[import:13-24, lang:"python"](code/python/fft.py) +[import:15-26, lang:"python"](code/python/fft.py) {% sample lang="scratch" %} [import:16-32, lang:"julia"](code/julia/fft.jl) {% sample lang="asm-x64" %}