mirror of
https://github.com/mborgerding/kissfft.git
synced 2025-05-27 21:20:27 -04:00
118 lines
2.7 KiB
Python
118 lines
2.7 KiB
Python
#!/usr/local/bin/python2.3
|
|
import math
|
|
import sys
|
|
|
|
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]
|
|
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 = []
|
|
for k in range(N):
|
|
F1k = Fp[k] + Fp[-k].conjugate()
|
|
F2k = -j*(Fp[k] - Fp[-k].conjugate())
|
|
print 'F2k[%d]=%s' % (k,F2k)
|
|
|
|
F.append( ( F1k + e ** ( -j*pi*k/N ) * F2k ) * .5 )
|
|
|
|
F.append( 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
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|