fft.py
Upload User: ozl2332
Upload Date: 2009-12-28
Package Size: 38k
Code Size: 5k
Category:

Voice Compress

Development Platform:

C/C++

  1. #!/usr/bin/env python2.3
  2. import math
  3. import sys
  4. import random
  5. pi=math.pi
  6. e=math.e
  7. j=complex(0,1)
  8. def fft(f,inv):
  9.     n=len(f)
  10.     if n==1:
  11.         return f
  12.     for p in 2,3,5:
  13.         if n%p==0:
  14.             break
  15.     else:
  16.         raise Exception('%s not factorable ' % n)
  17.     m = n/p
  18.     Fout=[]
  19.     for q in range(p): # 0,1
  20.         fp = f[q::p]
  21.         Fp = fft( fp ,inv)
  22.         Fout.extend( Fp )
  23.     for u in range(m):
  24.         scratch = Fout[u::m] # u to end in strides of m
  25.         for q1 in range(p):
  26.             k = q1*m + u  # indices to Fout above that became scratch
  27.             Fout[ k ] = scratch[0] # cuz e**0==1 in loop below
  28.             for q in range(1,p):
  29.                 if inv:
  30.                     t = e ** ( j*2*pi*k*q/n )
  31.                 else:                    
  32.                     t = e ** ( -j*2*pi*k*q/n )
  33.                 Fout[ k ] += scratch[q] * t
  34.     return Fout
  35. def rifft(F):
  36.     N = len(F) - 1
  37.     Z = [0] * (N)
  38.     for k in range(N):
  39.         Fek = ( F[k] + F[-k-1].conjugate() )
  40.         Fok = ( F[k] - F[-k-1].conjugate() ) * e ** (j*pi*k/N)
  41.         Z[k] = Fek + j*Fok
  42.     fp = fft(Z , 1)
  43.     f = []
  44.     for c in fp:
  45.         f.append(c.real)
  46.         f.append(c.imag)
  47.     return f
  48. def real_fft( f,inv ):
  49.     if inv:
  50.         return rifft(f)
  51.     N = len(f) / 2
  52.     res = f[::2]
  53.     ims = f[1::2]
  54.     fp = [ complex(r,i) for r,i in zip(res,ims) ]
  55.     print 'fft input ', fp
  56.     Fp = fft( fp ,0 )
  57.     print 'fft output ', Fp
  58.     F = [ complex(0,0) ] * ( N+1 )
  59.     
  60.     F[0] = complex( Fp[0].real + Fp[0].imag , 0 ) 
  61.     for k in range(1,N/2+1):
  62.         tw = e ** ( -j*pi*(.5+float(k)/N ) )
  63.         
  64.         F1k = Fp[k] + Fp[N-k].conjugate()
  65.         F2k = Fp[k] - Fp[N-k].conjugate()
  66.         F2k *= tw
  67.         F[k] = ( F1k + F2k ) * .5
  68.         F[N-k] = ( F1k - F2k ).conjugate() * .5
  69.         #F[N-k] = ( F1kp + e ** ( -j*pi*(.5+float(N-k)/N ) ) * F2kp ) * .5
  70.         #F[N-k] = ( F1k.conjugate() - tw.conjugate() * F2k.conjugate() ) * .5
  71.     F[N] = complex( Fp[0].real - Fp[0].imag , 0 ) 
  72.     return F
  73. def main():
  74.     #fft_func = fft
  75.     fft_func = real_fft
  76.     tvec = [0.309655,0.815653,0.768570,0.591841,0.404767,0.637617,0.007803,0.012665]
  77.     Ftvec = [ complex(r,i) for r,i in zip(
  78.                 [3.548571,-0.378761,-0.061950,0.188537,-0.566981,0.188537,-0.061950,-0.378761],
  79.                 [0.000000,-1.296198,-0.848764,0.225337,0.000000,-0.225337,0.848764,1.296198] ) ]
  80.     F = fft_func( tvec,0 )
  81.     nerrs= 0
  82.     for i in range(len(Ftvec)/2 + 1):
  83.         if abs( F[i] - Ftvec[i] )> 1e-5:
  84.             print 'F[%d]: %s != %s' % (i,F[i],Ftvec[i])
  85.             nerrs += 1
  86.     print '%d errors in forward fft' % nerrs
  87.     if nerrs:
  88.         return
  89.     trec = fft_func( F , 1 )
  90.     for i in range(len(trec) ):
  91.         trec[i] /= len(trec)
  92.     for i in range(len(tvec) ):
  93.         if abs( trec[i] - tvec[i] )> 1e-5:
  94.             print 't[%d]: %s != %s' % (i,tvec[i],trec[i])
  95.             nerrs += 1
  96.     print '%d errors in reverse fft' % nerrs
  97. def make_random(dims=[1]):
  98.     import Numeric 
  99.     res = []
  100.     for i in range(dims[0]):
  101.         if len(dims)==1:
  102.             r=random.uniform(-1,1)
  103.             i=random.uniform(-1,1)
  104.             res.append( complex(r,i) )
  105.         else:
  106.             res.append( make_random( dims[1:] ) )
  107.     return Numeric.array(res)
  108. def flatten(x):
  109.     import Numeric
  110.     ntotal = Numeric.product(Numeric.shape(x))
  111.     return Numeric.reshape(x,(ntotal,))
  112. def randmat( ndims ):
  113.     dims=[]
  114.     for i in range( ndims ):
  115.         curdim = int( random.uniform(2,4) )
  116.         dims.append( curdim )
  117.     return make_random(dims )
  118. def test_fftnd(ndims=3):
  119.     import FFT
  120.     import Numeric
  121.     x=randmat( ndims )
  122.     print 'dimensions=%s' % str( Numeric.shape(x) )
  123.     #print 'x=%s' %str(x)
  124.     xver = FFT.fftnd(x)
  125.     x2=myfftnd(x)
  126.     err = xver - x2
  127.     errf = flatten(err)
  128.     xverf = flatten(xver)
  129.     errpow = Numeric.vdot(errf,errf)+1e-10
  130.     sigpow = Numeric.vdot(xverf,xverf)+1e-10
  131.     snr = 10*math.log10(abs(sigpow/errpow) )
  132.     if snr<80:
  133.         print xver
  134.         print x2
  135.     print 'SNR=%sdB' % str( snr )
  136.  
  137. def myfftnd(x):
  138.     import Numeric
  139.     xf = flatten(x)
  140.     Xf = fftndwork( xf , Numeric.shape(x) )
  141.     return Numeric.reshape(Xf,Numeric.shape(x) )
  142. def fftndwork(x,dims):
  143.     import Numeric
  144.     dimprod=Numeric.product( dims )
  145.     for k in range( len(dims) ):
  146.         cur_dim=dims[ k ]
  147.         stride=dimprod/cur_dim
  148.         next_x = [complex(0,0)]*len(x)
  149.         for i in range(stride):
  150.             next_x[i*cur_dim:(i+1)*cur_dim] = fft(x[i:(i+cur_dim)*stride:stride],0)
  151.         x = next_x
  152.     return x
  153. if __name__ == "__main__":
  154.     try:
  155.         nd = int(sys.argv[1])
  156.     except:
  157.         nd=None
  158.     if nd:    
  159.         test_fftnd( nd )
  160.     else:    
  161.         sys.exit(0)