# coding: UTF-8
from scipy import *
from scipy import fftpack
import pylab

def test_fft1():
    """Test for Numpy-FFT against analytic results, first IFFT then FFT"""
    
    def fx(x, a):
        """original function, Lorentz function, time domain"""
        return a/(x**2+a**2)
    
    def ft(t, a):
        """analytic result of FFT, frequency domain"""
        return exp(-a*abs(t)) * sqrt(pi/2.0)
    
    vecfx = vectorize(fx)
    vecft = vectorize(ft)
    
    a = 1.0
    
    (xrange,dxrange) = linspace(-500,500,2**14,endpoint=True,retstep=True)
    
    data = vecft(xrange, a)
    
    # FFT to time domain
    ifftdata = fftpack.ifftshift(data)
    ifftdata = fftpack.ifft(ifftdata)

    ifftscale = fftpack.fftfreq(len(xrange), dxrange)
    ifftscale = fftpack.fftshift(ifftscale)
    difftscale= ifftscale[1] - ifftscale[0]

    # Calculate exact result
    fftexact = vecfx(ifftscale, a)
    
    # Re-Transform to frequency domain
    fftdata = fftpack.fft(ifftdata)
    fftdata = fftpack.fftshift(ifftdata)
    
    fftscale = fftpack.fftfreq(len(ifftscale), difftscale)
    fftscale = fftpack.fftshift(fftscale)
    
    pylab.subplot(211)
    pylab.plot(xrange, data, 'x', label='data')
    pylab.plot(fftscale, fftdata.real, label='Re FFT(IFFT(data))')
    pylab.plot(fftscale, fftdata.imag, label='Im FFT(IFFT(data))')
    pylab.legend()
    pylab.subplot(212)
    pylab.plot(ifftscale, ifftdata.real, label='Re IFFT')
    pylab.plot(ifftscale, ifftdata.imag, label='Im IFFT')
    pylab.plot(ifftscale, fftexact, '.-', label='exact inverse FT (real only)')
    pylab.legend()
    pylab.show()
    pylab.close()
    return 0

def test_fft2():
    """Test for Numpy-FFT against analytic results, first FFT then IFFT"""
    import pylab
    
    def ft(t, a):
        """original function in the time domain"""
        return exp(-a*abs(t)) * sqrt(pi/2.0)

    def ff(x, a):
        """analytic result of FFT, frequency domain"""
        return a/(x**2+a**2)
    
    vecft = vectorize(ft)
    vecff = vectorize(ff)
    
    a = 1.0
    
    (trange,dtrange) = linspace(-500,500,2**14,endpoint=True,retstep=True)
    
    data = vecft(trange, a)
    
    # FFT to frequency domain
    fftdata = fftpack.fft(data)
    # Shift 
    fftdata = fftpack.fftshift(fftdata)

    # Determine frequency scale
    fftscale = fftpack.fftfreq(len(trange), dtrange)
    fftscale = fftpack.fftshift(fftscale)
    dfftscale= fftscale[1]-fftscale[0]

    # Calculate exact result in frequency domain
    fftexact = vecff(fftscale, a)
    
    # Re-Transform to time domain
    ifftdata = fftpack.ifftshift(fftdata)
    ifftdata = fftpack.ifft(ifftdata)
    
    ifftscale = fftpack.fftfreq(len(fftscale), dfftscale)
    ifftscale = fftpack.fftshift(ifftscale)
    
    pylab.subplot(211)
    pylab.plot(trange, data, 'x', label='data')
    pylab.plot(ifftscale, ifftdata.real, label='Re IFFT(FFT(data))')
    pylab.plot(ifftscale, ifftdata.imag, label='Im IFFT(FFT(data))')
    pylab.legend()
    pylab.subplot(212)
    pylab.plot(fftscale, fftdata.real, label='Re FFT')
    pylab.plot(fftscale, fftdata.imag, label='Im FFT')
    pylab.plot(fftscale, fftexact, '.-', label='exact FT (real only)')
    pylab.legend()
    pylab.show()
    pylab.close()
    return 0
    
if __name__ == "__main__":
    test_fft1()
    test_fft2()