home *** CD-ROM | disk | FTP | other *** search
/ MacHack 2000 / MacHack 2000.toast / pc / The Hacks / MacHacksBug / Python 1.5.2c1 / Extensions / Numerical / Lib / FFT.py < prev    next >
Encoding:
Python Source  |  2000-06-23  |  2.1 KB  |  78 lines
  1. import Numeric, fftpack, copy
  2.  
  3. _fft_cache = {}
  4. _real_fft_cache = {}
  5.  
  6. def _raw_fft(a, n=None, axis=-1, init_function=fftpack.cffti, 
  7.         work_function=fftpack.cfftf, fft_cache = _fft_cache ):
  8.     a = Numeric.asarray(a)
  9.  
  10.     if n == None: n = a.shape[axis]
  11.  
  12.     try:
  13.         wsave = fft_cache[n]
  14.     except(KeyError):
  15.         wsave = init_function(n)
  16.         fft_cache[n] = wsave
  17.  
  18.     if a.shape[axis] != n:
  19.         s = list(a.shape)
  20.         if s[axis] > n:
  21.             index = [slice(None)]*len(s)
  22.             index[axis] = slice(0,n)
  23.             a = a[index]
  24.         else:    
  25.             index = [slice(None)]*len(s)
  26.             index[axis] = slice(0,s[axis])
  27.             s[axis] = n
  28.             z = Numeric.zeros(s, a.typecode())
  29.             z[index] = a
  30.             a = z
  31.  
  32.     if axis != -1: a = Numeric.swapaxes(a, axis, -1)    
  33.     r = work_function(a, wsave)
  34.     if axis != -1: r = Numeric.swapaxes(r, axis, -1)
  35.     return r
  36.  
  37. def fft(a, n=None, axis=-1): 
  38.     return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftf, _fft_cache)
  39.  
  40. def inverse_fft(a, n=None, axis=-1): 
  41.     if n == None: n = Numeric.shape(a)[axis]
  42.     return _raw_fft(a, n, axis, fftpack.cffti, fftpack.cfftb, _fft_cache)/n
  43.  
  44. def real_fft(a, n=None, axis=-1): 
  45.     return _raw_fft(a.astype(Numeric.Float), n, axis, fftpack.rffti, fftpack.rfftf, _real_fft_cache)
  46.  
  47. def inverse_real_fft(a, n=None, axis=-1): 
  48.     if n == None: n = Numeric.shape(a)[axis]
  49.     return _raw_fft(a.astype(Numeric.Float), n, axis, fftpack.rffti, fftpack.rfftb, _real_fft_cache)/n
  50.  
  51. def _raw_fft2d(a, s=None, axes=(-2,-1), function=fft):
  52.     a = Numeric.asarray(a)
  53.     if s == None: s = a.shape[-2:]
  54.     f1 = function(a, n=s[1], axis=axes[1])
  55.     return function(f1, n=s[0], axis=axes[0])
  56.  
  57. def fft2d(a, s=None, axes=(-2,-1)):
  58.     return _raw_fft2d(a,s,axes,fft)
  59.  
  60. def inverse_fft2d(a, s=None, axes=(-2,-1)):
  61.     return _raw_fft2d(a, s, axes, inverse_fft)
  62.  
  63. def real_fft2d(a, s=None, axes=(-2,-1)):
  64.     return _raw_fft2d(a, s, axes, real_fft)
  65.  
  66. def test():
  67.     print fft( (0,1)*4 )
  68.     print inverse_fft( fft((0,1)*4) )
  69.     print fft( (0,1)*4, n=16 )
  70.     print fft( (0,1)*4, n=4 )
  71.  
  72.     print fft2d( [(0,1),(1,0)] )
  73.     print inverse_fft2d (fft2d( [(0, 1), (1, 0)] ) )
  74.     print real_fft2d([(0,1),(1,0)] )
  75.     print real_fft2d([(1,1),(1,1)] )
  76.  
  77. if __name__ == '__main__': test()
  78.