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