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