comparison ext/kissfft/test/fft.py @ 409:1f1999b0f577

Bring in kissfft into this repo (formerly a subrepo, but the remote is not responding)

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