Chris@184: #!/usr/bin/env python Chris@184: Chris@184: import math Chris@184: import sys Chris@184: import random Chris@184: Chris@184: pi=math.pi Chris@184: e=math.e Chris@184: j=complex(0,1) Chris@184: Chris@184: def fft(f,inv): Chris@184: n=len(f) Chris@184: if n==1: Chris@184: return f Chris@184: Chris@184: for p in 2,3,5: Chris@184: if n%p==0: Chris@184: break Chris@184: else: Chris@184: raise Exception('%s not factorable ' % n) Chris@184: Chris@184: m = n/p Chris@184: Fout=[] Chris@184: for q in range(p): # 0,1 Chris@184: fp = f[q::p] # every p'th time sample Chris@184: Fp = fft( fp ,inv) Chris@184: Fout.extend( Fp ) Chris@184: Chris@184: for u in range(m): Chris@184: scratch = Fout[u::m] # u to end in strides of m Chris@184: for q1 in range(p): Chris@184: k = q1*m + u # indices to Fout above that became scratch Chris@184: Fout[ k ] = scratch[0] # cuz e**0==1 in loop below Chris@184: for q in range(1,p): Chris@184: if inv: Chris@184: t = e ** ( j*2*pi*k*q/n ) Chris@184: else: Chris@184: t = e ** ( -j*2*pi*k*q/n ) Chris@184: Fout[ k ] += scratch[q] * t Chris@184: Chris@184: return Fout Chris@184: Chris@184: def rifft(F): Chris@184: N = len(F) - 1 Chris@184: Z = [0] * (N) Chris@184: for k in range(N): Chris@184: Fek = ( F[k] + F[-k-1].conjugate() ) Chris@184: Fok = ( F[k] - F[-k-1].conjugate() ) * e ** (j*pi*k/N) Chris@184: Z[k] = Fek + j*Fok Chris@184: Chris@184: fp = fft(Z , 1) Chris@184: Chris@184: f = [] Chris@184: for c in fp: Chris@184: f.append(c.real) Chris@184: f.append(c.imag) Chris@184: return f Chris@184: Chris@184: def real_fft( f,inv ): Chris@184: if inv: Chris@184: return rifft(f) Chris@184: Chris@184: N = len(f) / 2 Chris@184: Chris@184: res = f[::2] Chris@184: ims = f[1::2] Chris@184: Chris@184: fp = [ complex(r,i) for r,i in zip(res,ims) ] Chris@184: print 'fft input ', fp Chris@184: Fp = fft( fp ,0 ) Chris@184: print 'fft output ', Fp Chris@184: Chris@184: F = [ complex(0,0) ] * ( N+1 ) Chris@184: Chris@184: F[0] = complex( Fp[0].real + Fp[0].imag , 0 ) Chris@184: Chris@184: for k in range(1,N/2+1): Chris@184: tw = e ** ( -j*pi*(.5+float(k)/N ) ) Chris@184: Chris@184: F1k = Fp[k] + Fp[N-k].conjugate() Chris@184: F2k = Fp[k] - Fp[N-k].conjugate() Chris@184: F2k *= tw Chris@184: F[k] = ( F1k + F2k ) * .5 Chris@184: F[N-k] = ( F1k - F2k ).conjugate() * .5 Chris@184: #F[N-k] = ( F1kp + e ** ( -j*pi*(.5+float(N-k)/N ) ) * F2kp ) * .5 Chris@184: #F[N-k] = ( F1k.conjugate() - tw.conjugate() * F2k.conjugate() ) * .5 Chris@184: Chris@184: F[N] = complex( Fp[0].real - Fp[0].imag , 0 ) Chris@184: return F Chris@184: Chris@184: def main(): Chris@184: #fft_func = fft Chris@184: fft_func = real_fft Chris@184: Chris@184: tvec = [0.309655,0.815653,0.768570,0.591841,0.404767,0.637617,0.007803,0.012665] Chris@184: Ftvec = [ complex(r,i) for r,i in zip( Chris@184: [3.548571,-0.378761,-0.061950,0.188537,-0.566981,0.188537,-0.061950,-0.378761], Chris@184: [0.000000,-1.296198,-0.848764,0.225337,0.000000,-0.225337,0.848764,1.296198] ) ] Chris@184: Chris@184: F = fft_func( tvec,0 ) Chris@184: Chris@184: nerrs= 0 Chris@184: for i in range(len(Ftvec)/2 + 1): Chris@184: if abs( F[i] - Ftvec[i] )> 1e-5: Chris@184: print 'F[%d]: %s != %s' % (i,F[i],Ftvec[i]) Chris@184: nerrs += 1 Chris@184: Chris@184: print '%d errors in forward fft' % nerrs Chris@184: if nerrs: Chris@184: return Chris@184: Chris@184: trec = fft_func( F , 1 ) Chris@184: Chris@184: for i in range(len(trec) ): Chris@184: trec[i] /= len(trec) Chris@184: Chris@184: for i in range(len(tvec) ): Chris@184: if abs( trec[i] - tvec[i] )> 1e-5: Chris@184: print 't[%d]: %s != %s' % (i,tvec[i],trec[i]) Chris@184: nerrs += 1 Chris@184: Chris@184: print '%d errors in reverse fft' % nerrs Chris@184: Chris@184: Chris@184: def make_random(dims=[1]): Chris@184: import Numeric Chris@184: res = [] Chris@184: for i in range(dims[0]): Chris@184: if len(dims)==1: Chris@184: r=random.uniform(-1,1) Chris@184: i=random.uniform(-1,1) Chris@184: res.append( complex(r,i) ) Chris@184: else: Chris@184: res.append( make_random( dims[1:] ) ) Chris@184: return Numeric.array(res) Chris@184: Chris@184: def flatten(x): Chris@184: import Numeric Chris@184: ntotal = Numeric.product(Numeric.shape(x)) Chris@184: return Numeric.reshape(x,(ntotal,)) Chris@184: Chris@184: def randmat( ndims ): Chris@184: dims=[] Chris@184: for i in range( ndims ): Chris@184: curdim = int( random.uniform(2,4) ) Chris@184: dims.append( curdim ) Chris@184: return make_random(dims ) Chris@184: Chris@184: def test_fftnd(ndims=3): Chris@184: import FFT Chris@184: import Numeric Chris@184: Chris@184: x=randmat( ndims ) Chris@184: print 'dimensions=%s' % str( Numeric.shape(x) ) Chris@184: #print 'x=%s' %str(x) Chris@184: xver = FFT.fftnd(x) Chris@184: x2=myfftnd(x) Chris@184: err = xver - x2 Chris@184: errf = flatten(err) Chris@184: xverf = flatten(xver) Chris@184: errpow = Numeric.vdot(errf,errf)+1e-10 Chris@184: sigpow = Numeric.vdot(xverf,xverf)+1e-10 Chris@184: snr = 10*math.log10(abs(sigpow/errpow) ) Chris@184: if snr<80: Chris@184: print xver Chris@184: print x2 Chris@184: print 'SNR=%sdB' % str( snr ) Chris@184: Chris@184: def myfftnd(x): Chris@184: import Numeric Chris@184: xf = flatten(x) Chris@184: Xf = fftndwork( xf , Numeric.shape(x) ) Chris@184: return Numeric.reshape(Xf,Numeric.shape(x) ) Chris@184: Chris@184: def fftndwork(x,dims): Chris@184: import Numeric Chris@184: dimprod=Numeric.product( dims ) Chris@184: Chris@184: for k in range( len(dims) ): Chris@184: cur_dim=dims[ k ] Chris@184: stride=dimprod/cur_dim Chris@184: next_x = [complex(0,0)]*len(x) Chris@184: for i in range(stride): Chris@184: next_x[i*cur_dim:(i+1)*cur_dim] = fft(x[i:(i+cur_dim)*stride:stride],0) Chris@184: x = next_x Chris@184: return x Chris@184: Chris@184: if __name__ == "__main__": Chris@184: try: Chris@184: nd = int(sys.argv[1]) Chris@184: except: Chris@184: nd=None Chris@184: if nd: Chris@184: test_fftnd( nd ) Chris@184: else: Chris@184: sys.exit(0)