#!/usr/bin/python
#Make sure that scipy and pylab are installed.
from scipy import zeros, cos,sin,arange,exp
from scipy.fftpack import fftshift,fftfreq,fft
import pylab

#This program generates a power spectrum of the function y[] and writes the result to a file.

#Parameters:
xstep=0.05
xmax=2**8.
N=int(xmax/xstep)

#Definition of y
x=arange(float(N))*xstep
y=sin(2.*x)+cos(0.5*x)

#The fourier transform and the corresping frequencies are calculated:
ffty=fft(y)
freq=fftfreq(N,xstep)

#The output is written:
OutputFile=open('powerspectrum.ascii','w')
OutputFile.write('#x\ty\t1/x\tfft(y)\n')
for i in range(N):
    OutputFile.write(str(i*xstep) + '\t' + str(y[i]) + '\t' + str(freq[i]) + '\t' + str(abs(ffty[i])**2)+'\n')

OutputFile.close()

#plotting
pylab.xlabel('x')
pylab.ylabel('f(x)')
pylab.plot(x,y)
pylab.show()

pylab.xlabel('Frequency')
pylab.ylabel('Power Spectrum')
pylab.plot(freq,abs(ffty)**2)
pylab.show()