annotate src/ext/kissfft/test/fft.py @ 196:da283326bcd3 tip master

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