ffttest.py

ファイルコンテンツ

import scipy
import scipy.fftpack as fft
import pylab

L = 128
U = 1./8
UL = int(L*U)

print 'FFT Demo'
print """Please select:
  1. One sided rectangular function
  2. Folded symmetrical rectangular function
  3. Symmetrical Trapezoid function
"""
ans = input()

X = range(L)
if ans == 1:
    f = scipy.array([1.0]*UL+[0.0]*(L-UL))
elif ans == 2: ### folding
    f = scipy.array([1.0]*(UL/2)+[0.0]*(L-UL+1)+[1.0]*(UL/2-1))
else: ### non-sharp cutoff
    S = [2./UL*i for i in range(UL/2)]
    SR = S[:]
    SR.reverse()
    f = scipy.array([1.0]*(UL/4)+SR+[0.0]*(L-UL-UL/2+1)+S+[1.0]*(UL/4-1))

pylab.subplot(211)
pylab.plot(X,f, '-', label='Original Function')
pylab.legend(('Original Function',))

pylab.subplot(212)
F = fft.fft(f)
pylab.plot(X,F.real, '-', label='FFT Real part')
pylab.plot(X,F.imag, '-', label='FFT Imaginary part')
pylab.legend(('Real Part', 'Imaginary Part'))
pylab.title('FFT')
pylab.show()