c@409: #!/usr/bin/env python
c@409: 
c@409: import math
c@409: import sys
c@409: import os
c@409: import random
c@409: import struct
c@409: import popen2
c@409: import getopt
c@409: import numpy
c@409: 
c@409: pi=math.pi
c@409: e=math.e
c@409: j=complex(0,1)
c@409: 
c@409: doreal=0
c@409: 
c@409: datatype = os.environ.get('DATATYPE','float')
c@409: 
c@409: util = '../tools/fft_' + datatype
c@409: minsnr=90
c@409: if datatype == 'double':
c@409:     fmt='d'
c@409: elif datatype=='int16_t':
c@409:     fmt='h'
c@409:     minsnr=10
c@409: elif datatype=='int32_t':
c@409:     fmt='i'
c@409: elif datatype=='simd':
c@409:     fmt='4f'
c@409:     sys.stderr.write('testkiss.py does not yet test simd')
c@409:     sys.exit(0)
c@409: elif datatype=='float':
c@409:     fmt='f'
c@409: else:
c@409:     sys.stderr.write('unrecognized datatype %s\n' % datatype)
c@409:     sys.exit(1)
c@409:  
c@409: 
c@409: def dopack(x,cpx=1):
c@409:     x = numpy.reshape( x, ( numpy.size(x),) )
c@409:     
c@409:     if cpx:
c@409:         s = ''.join( [ struct.pack(fmt*2,c.real,c.imag) for c in x ] )
c@409:     else:
c@409:         s = ''.join( [ struct.pack(fmt,c.real) for c in x ] )
c@409:     return s
c@409: 
c@409: def dounpack(x,cpx):
c@409:     uf = fmt * ( len(x) / struct.calcsize(fmt) )
c@409:     s = struct.unpack(uf,x)
c@409:     if cpx:
c@409:         return numpy.array(s[::2]) + numpy.array( s[1::2] )*j
c@409:     else:
c@409:         return numpy.array(s )
c@409: 
c@409: def make_random(dims=[1]):
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:             if doreal:
c@409:                 res.append( r )
c@409:             else:
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 numpy.array(res)
c@409: 
c@409: def flatten(x):
c@409:     ntotal = numpy.size(x)
c@409:     return numpy.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,5) )
c@409:         if doreal and i==(ndims-1):
c@409:             curdim = int(curdim/2)*2 # force even last dimension if real
c@409:         dims.append( curdim )
c@409:     return make_random(dims )
c@409: 
c@409: def test_fft(ndims):
c@409:     x=randmat( ndims )
c@409: 
c@409: 
c@409:     if doreal:
c@409:         xver = numpy.fft.rfftn(x)
c@409:     else:
c@409:         xver = numpy.fft.fftn(x)
c@409:     
c@409:     open('/tmp/fftexp.dat','w').write(dopack( flatten(xver) , True ) )
c@409: 
c@409:     x2=dofft(x,doreal)
c@409:     err = xver - x2
c@409:     errf = flatten(err)
c@409:     xverf = flatten(xver)
c@409:     errpow = numpy.vdot(errf,errf)+1e-10
c@409:     sigpow = numpy.vdot(xverf,xverf)+1e-10
c@409:     snr = 10*math.log10(abs(sigpow/errpow) )
c@409:     print 'SNR (compared to NumPy) : %.1fdB' % float(snr)
c@409: 
c@409:     if snr<minsnr:
c@409:         print 'xver=',xver
c@409:         print 'x2=',x2
c@409:         print 'err',err
c@409:         sys.exit(1)
c@409:  
c@409: def dofft(x,isreal):
c@409:     dims=list( numpy.shape(x) )
c@409:     x = flatten(x)
c@409: 
c@409:     scale=1
c@409:     if datatype=='int16_t':
c@409:         x = 32767 * x
c@409:         scale = len(x) / 32767.0
c@409:     elif datatype=='int32_t':
c@409:         x = 2147483647.0 * x
c@409:         scale = len(x) / 2147483647.0
c@409: 
c@409:     cmd='%s -n ' % util
c@409:     cmd += ','.join([str(d) for d in dims])
c@409:     if doreal:
c@409:         cmd += ' -R '
c@409: 
c@409:     print cmd
c@409:     p = popen2.Popen3(cmd )
c@409: 
c@409:     open('/tmp/fftin.dat','w').write(dopack( x , isreal==False ) )
c@409: 
c@409:     p.tochild.write( dopack( x , isreal==False ) )
c@409:     p.tochild.close()
c@409: 
c@409:     res = dounpack( p.fromchild.read() , 1 )
c@409:     open('/tmp/fftout.dat','w').write(dopack( flatten(res) , True ) )
c@409:     if doreal:
c@409:         dims[-1] = int( dims[-1]/2 ) + 1
c@409: 
c@409:     res = scale * res
c@409: 
c@409:     p.wait()
c@409:     return numpy.reshape(res,dims)
c@409: 
c@409: def main():
c@409:     opts,args = getopt.getopt(sys.argv[1:],'r')
c@409:     opts=dict(opts)
c@409: 
c@409:     global doreal
c@409:     doreal = opts.has_key('-r')
c@409: 
c@409:     if doreal:
c@409:         print 'Testing multi-dimensional real FFTs'
c@409:     else:
c@409:         print 'Testing multi-dimensional FFTs'
c@409: 
c@409:     for dim in range(1,4):
c@409:         test_fft( dim )
c@409: 
c@409: if __name__ == "__main__":
c@409:     main()
c@409: