#!/usr/bin/python
#By Samuel George; sgeorge@mrao.cam.ac.uk #Original 08/12/11
#Purpose: Read ascii data fit a Gaussian
#usage: simplegaussfit.py
#produces: plotfitting.png
#input - change the variable "data" to be you filename.
import warnings
warnings.simplefilter("ignore",DeprecationWarning)
import matplotlib
matplotlib.use('Agg') #avoid X11 issues.
import scipy, numpy, pylab, asciidata
from numpy import *
from pylab import *
from scipy import *
from scipy.optimize import leastsq

gauss_fit = lambda p, x: p[0]*(1/sqrt(2*pi*(p[2]**2)))*exp(-(x-p[1])**2/(2*p[2]**2)) #1d Gaussian func
e_gauss_fit = lambda p, x, y: (gauss_fit(p,x) -y) #1d Gaussian fit

data = "test_gauss.dat"  #your input data - should be x,y 
demo = asciidata.open(data)
x_pos = double(demo[0])  #this is expected to be pixel number
y_power = double(demo[1])

v0= [1.0,mean(x_pos),1.0] #inital guesses for Gaussian Fit. $just do it around the peak
out = leastsq(e_gauss_fit, v0[:], args=(x_pos, y_power), maxfev=100000,  full_output=1) #Gauss Fit
v = out[0] #fit parammeters out
covar = out[1] #covariance matrix output

xxx = arange(min(x_pos),max(x_pos),x_pos[1]-x_pos[0])
ccc = gauss_fit(v,xxx) # this will only work if the units are pixel and not wavelength

fig = figure(figsize=(9, 9)) #make a plot
ax1 = fig.add_subplot(111)
ax1.plot(x_pos,y_power,'gs') #spectrum
ax1.plot(xxx,ccc,'b--') #fitted spectrum
ax1.axvline(x=xxx[where(ccc == max(ccc))[0]][0],color='r') #max position in data
setp(gca(), ylabel="power", xlabel="pixel position")
pylab.savefig("plotfitting.png")

print "p[0], a: ", v[0]
print "p[1], mu: ", v[1]
print "p[2], sigma: ", v[2]