#!/usr/bin/python -Wall # ================================================================ # John Kerl # kerl.john.r@gmail.com # 2005-01-11 # (Ported to Python 2006-03-06) # # This is a radix-2 fast Fourier transform. # This file contains the library routines. The I/O wrapper is the file pyfft. # ---------------------------------------------------------------- # This software is released under the terms of the GNU GPL. # Please see LICENSE.txt in the same directory as this file. # ================================================================ from __future__ import division # 1/2 = 0.5, not 0. import sys import copy import re from math import * # ---------------------------------------------------------------- def half_trig_table(n, forward): W = [] for i in range(0, int(n/2)): arg = 2.0 * pi * i / n if (forward): arg = -arg W.append(cos(arg) + sin(arg)*1j) return W # ---------------------------------------------------------------- def full_trig_table(n, forward): W = [] for i in range(0, n): arg = 2.0 * pi * i / n if (forward): arg = -arg W.append(cos(arg) + sin(arg)*1j) return W # ---------------------------------------------------------------- # u is input # v is output def fft(u, fold_in, fold_out, forward, scale): v = copy.copy(u) n = len(v) n2 = n / 2 log2n = log2(n) sqrt_recip_2 = sqrt(0.5) if (not is_a_power_of_two(n)): print "pyfft_m: input length", n, "is not a power of two." W = half_trig_table(n, forward) # Output folding. if (fold_out): j = 1 for i in range(0, n2): v[j] = -v[j] j += 2 # Bit-reverse stage. for i in range(0, n): ir = bit_reverse(i, log2n) if (i < ir): temp = v[i] v[i] = v[ir] v[ir] = temp # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - # Radix-two stages. # # y00 \ / @ \ / @ \ / @ \ / Y00 # y08 / \ @ \ / @ \ / @ \ / Y01 # y04 \ / @ W0 / \ @ \ / @ \ / Y02 # y12 / \ @ W4 / \ @ \ / @ \ / Y03 # # y02 \ / @ \ / @ W0 / \ @ \ / Y04 # y10 / \ @ \ / @ W2 / \ @ \ / Y05 # y06 \ / @ W0 / \ @ W4 / \ @ \ / Y06 # y14 / \ @ W4 / \ @ W6 / \ @ \ / Y07 # # y01 \ / @ \ / @ \ / @ W0 / \ Y08 # y09 / \ @ \ / @ \ / @ W1 / \ Y09 # y05 \ / @ W0 / \ @ \ / @ W2 / \ Y10 # y13 / \ @ W4 / \ @ \ / @ W3 / \ Y11 # # y03 \ / @ \ / @ W0 / \ @ W4 / \ Y12 # y11 / \ @ \ / @ W2 / \ @ W5 / \ Y13 # y07 \ / @ W0 / \ @ W4 / \ @ W6 / \ Y14 # y15 / \ @ W4 / \ @ W6 / \ @ W7 / \ Y15 # # a \ / c --> c := a + b # b / \ d --> d := a - b # # a \ / c --> c := a + w * b # b w / \ d --> d := a - w * b # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - num_flocks = int(n/2) flock_sep = 2 bperflock = 1 wingspan = 1 twiddle_stride = int(n/2) stageno = 0 fno = 0 bno = 0 twiddleidx = 0 num_stages = log2n for stageno in range(0, num_stages): for fno in range(0, num_flocks): twiddleidx = 0 for bno in range(0, bperflock): upperidx = fno * flock_sep + bno loweridx = upperidx + wingspan a = v[upperidx] b = v[loweridx] w = W[twiddleidx] wb = w * b c = a + wb d = a - wb if (scale): c *= sqrt_recip_2 d *= sqrt_recip_2 v[upperidx] = c v[loweridx] = d twiddleidx += twiddle_stride num_flocks >>= 1 flock_sep <<= 1 bperflock <<= 1 wingspan <<= 1 twiddle_stride >>= 1 # Input folding. if (fold_in): j = 1 for i in range(0, n2): v[j] = -v[j] j += 2 return v # ---------------------------------------------------------------- # u is input # t is temp input # v is output def dft(u, fold_in, fold_out, forward, scale): t = copy.copy(u) v = copy.copy(u) n = len(v) n2 = int(n/2) W = full_trig_table(n, forward) # Output folding. if (fold_out): if (n & 1): print >> sys.stderr, "pyfft_m.dft: cannot fold output with odd vector length", n sys.exit(1) j = 1 for i in range(0, n2): t[j] = -t[j] j += 2 # DFT per se for k in range(0, n): sum = 0 for j in range(0, n): sum += W[(k*j) % n] * t[j] v[k] = sum if (scale): s = 1.0 / sqrt(n) for k in range(0, n): v[k] *= s # Input folding. if (fold_in): if (n & 1): print >> sys.stderr, "pyfft_m.dft: cannot fold input with odd vector length", n sys.exit(1) j = 1 for i in range(0, n2): v[j] = -v[j] j += 2 return v # ---------------------------------------------------------------- def bit_reverse(input, num_bits): output = input output = ((output & 0xaaaaaaaa) >> 1) | ((output & 0x55555555) << 1) output = ((output & 0xcccccccc) >> 2) | ((output & 0x33333333) << 2) output = ((output & 0xf0f0f0f0) >> 4) | ((output & 0x0f0f0f0f) << 4) output = ((output & 0xff00ff00) >> 8) | ((output & 0x00ff00ff) << 8) output = ((output & 0xffff0000) >> 16) | ((output & 0x0000ffff) << 16) output >>= (32 - num_bits) return output & ((1 << num_bits) - 1) # ---------------------------------------------------------------- def log2(n): rv = 0 nsave = n if (n == 0): print >> sys.stderr, "log2: Can't take logarithm of zero." sys.exit(1) while (n != 1): n >>= 1 rv += 1 if (nsave != (1 << rv)): print >> sys.stderr, "Argument", nsave, "is not a power of two." sys.exit(1) return rv # ---------------------------------------------------------------- def is_a_power_of_two(n): if (n == 0): return 0 if ((n & (n - 1)) == 0): return 1 return 0