|
| 1 | +import math |
| 2 | + |
| 3 | +def complex_dft(xr, xi, n): |
| 4 | + pi = 3.141592653589793 |
| 5 | + rex = [0] * n |
| 6 | + imx = [0] * n |
| 7 | + for k in range(0, n): # exclude n |
| 8 | + rex[k] = 0 |
| 9 | + imx[k] = 0 |
| 10 | + for k in range(0, n): # for each value in freq domain |
| 11 | + for i in range(0, n): # correlate with the complex sinusoid |
| 12 | + sr = math.cos(2 * pi * k * i / n) |
| 13 | + si = -math.sin(2 * pi * k * i / n) |
| 14 | + rex[k] += xr[i] * sr - xi[i] * si |
| 15 | + imx[k] += xr[i] * si + xi[i] * sr |
| 16 | + return rex, imx |
| 17 | + |
| 18 | +# FFT version based on the original BASIC program |
| 19 | +def fft_basic(rex, imx, n): |
| 20 | + pi = 3.141592653589793 |
| 21 | + m = int(math.log(n, 2)) # float to int |
| 22 | + j = n / 2 |
| 23 | + |
| 24 | + # bit reversal sorting |
| 25 | + for i in range(1, n - 1): # [1,n-2] |
| 26 | + if i >= j: |
| 27 | + # swap i with j |
| 28 | + print "swap %d with %d"%(i, j) |
| 29 | + rex[i], rex[j] = rex[j], rex[i] |
| 30 | + imx[i], imx[j] = imx[j], imx[i] |
| 31 | + k = n / 2 |
| 32 | + while (1): |
| 33 | + if k > j: |
| 34 | + break |
| 35 | + j -= k |
| 36 | + k /= 2 |
| 37 | + j += k |
| 38 | + |
| 39 | + for l in range(1, m + 1): # each stage |
| 40 | + le = int(math.pow(2, l)) # 2^l |
| 41 | + le2 = le / 2 |
| 42 | + ur = 1 |
| 43 | + ui = 0 |
| 44 | + sr = math.cos(pi / le2) |
| 45 | + si = -math.sin(pi / le2) |
| 46 | + for j in range(1, le2 + 1): # [1, le2] sub DFT |
| 47 | + for i in xrange(j - 1, n - 1, le): # for butterfly |
| 48 | + ip = i + le2 |
| 49 | + tr = rex[ip] * ur - imx[ip] * ui |
| 50 | + ti = rex[ip] * ui + imx[ip] * ur |
| 51 | + rex[ip] = rex[i] - tr |
| 52 | + imx[ip] = imx[i] - ti |
| 53 | + rex[i] += tr |
| 54 | + imx[i] += ti |
| 55 | + tr = ur |
| 56 | + ur = tr * sr - ui * si |
| 57 | + ui = tr * si + ui * sr |
| 58 | + |
| 59 | +def print_list(l): |
| 60 | + n = len(l) |
| 61 | + print "[%d]: {"%(n) |
| 62 | + for i in xrange(0, n): |
| 63 | + print l[i], |
| 64 | + print "}" |
| 65 | + |
| 66 | + |
| 67 | +if __name__ == "__main__": |
| 68 | + print "hello,world." |
| 69 | + pi = 3.1415926 |
| 70 | + x = [] |
| 71 | + n = 64 |
| 72 | + for i in range(0, n): |
| 73 | + p = math.sin(2 * pi * i / n) |
| 74 | + x.append(p) |
| 75 | + |
| 76 | + xr = x[:] |
| 77 | + xi = x[:] |
| 78 | + rex, imx = complex_dft(xr, xi, n) |
| 79 | + print "complet_dft(): n=", n |
| 80 | + print "rex: " |
| 81 | + print_list([int(e) for e in rex]) |
| 82 | + print "imx: " |
| 83 | + print_list([int(e) for e in imx]) |
| 84 | + |
| 85 | + fr = x[:] |
| 86 | + fi = x[:] |
| 87 | + |
| 88 | + fft_basic(fr, fi, n) |
| 89 | + print "fft_basic(): n=", n |
| 90 | + print "rex: " |
| 91 | + print_list([int(e) for e in fr]) |
| 92 | + print "imx: " |
| 93 | + print_list([int(e) for e in fi]) |
0 commit comments