#! /usr/bin/env python
#
# Simple python script fitting a gaussian to noisy data using mpfit
# Brian R. Kent
# National Radio Astronomy Observatory
#

import pylab
import string
import math
import os
import sys
import time
import getopt
import numpy
import scipy.stats
import datetime
from math import sqrt
import mpfit

#
# append path
#
path=os.getcwd()
sys.path.append(path)
files = os.listdir(path)
#

def peval(x, p):
       # The model function with parameters p
       return (1./sqrt(2*numpy.pi*p[1]**2))*numpy.exp(-(x-p[0])**2/(2*p[1]**2))

def myfunct(p, fjac=None, x=None, y=None, err=None ):
       # Function that return the weighted deviates
       model = peval(x, p)
       status = 0
       return([status, (y-model)/err])

# Generate model data for a Gaussian with param mu and sigma and add noise
x=numpy.arange(-10.,10.,20./1000)
preal=[-2, .5]
y_true=peval(x,preal)
mu,sigma=0,0.7
y      = y_true + 0.06 * numpy.random.normal(mu,sigma, len(x) )
err    = 1.0 + 0.01 * numpy.random.normal(mu,sigma, len(x) )
# Initial estimates for MPFIT
p0 = [-0.5, 0.5]
fa = {'x':x, 'y':y, 'err':err}

# Call MPFIT with user defined function 'myfunct'
m = mpfit.mpfit( myfunct, p0, functkw=fa )

print "status: ", m.status
if (m.status <= 0): 
   print 'error message = ', m.errmsg
else:
   print "Iterations: ", m.niter
   print "Fitted pars: ", m.params
   print "Uncertainties: ", m.perror 

# Plot the result with Matplotlib
pylab.clf()
pylab.plot(x,y,'r', label="Noisy data")
pylab.plot( x,peval(x,m.params), label="Fit" )
pylab.plot( x,y_true, 'g', label="True data" )
pylab.xlabel( "X" )
pylab.ylabel( "Measurement data" )
pylab.title( "Least-squares fit to noisy data using MPFIT" )
pylab.legend()
pylab.show()