from __future__ import print_function
from math import sqrt
from scipy import odr
def linear_fit(xdata, ydata, sigmay=None, sigmax=None):
"""
Performs a linear fit to data.
Parameters
----------
xdata : An array of length N.
ydata : An array of length N.
sigmay : None or an array of length N.
sigmax : None or an array of length N.
If one is provided, it is the standard deviation of ydata. Analytical linear regression used.
If both are provided, they are the standard deviations of ydata and xdata, respectively. ODR is used.
Returns
-------
a, b : Optimal parameter of linear fit (y = a*x + b)
sa, sb : Uncertainties of the parameters
"""
def lin_func(p, x):
a, b = p
return a*x + b
dof = len(ydata) - 2
if sigmax is None:
if sigmay is None:
w = ones(len(ydata))
else:
w=1.0/(sigmay**2)
sw = sum(w)
wx = w*xdata
swx = sum(wx)
swy = sum(w*ydata)
swxy = sum(wx*ydata)
swx2 = sum(wx*xdata)
a = (sw*swxy - swx*swy)/(sw*swx2 - swx*swx)
b = (swy*swx2 - swx*swxy)/(sw*swx2 - swx*swx)
sa = sqrt(sw/(sw*swx2 - swx*swx))
sb = sqrt(swx2/(sw*swx2 - swx*swx))
if sigmay is None:
chi2 = sum(((a*xdata + b)-ydata)**2)
else:
chi2 = sum((((a*xdata + b)-ydata)/sigmay)**2)
rchi2 = chi2/dof
else:
a0 = (ydata[-1]-ydata[0])/(xdata[-1]-xdata[0])
b0 = ydata[0]-xdata[0]*a0
model = odr.Model(lin_func)
data = odr.RealData(x=xdata, y=ydata, sx=sigmax, sy=sigmay)
od = odr.ODR(data, model, [a0,b0])
out = od.run()
a,b = out.beta
sa,sb = out.sd_beta
rchi2 = out.res_var
if rchi2 < 1.0 :
sa = sa/sqrt(rchi2)
sb = sb/sqrt(rchi2)
if sigmay is None:
print('no uncertainties provided, use uncertainties of slope and intercept with caution')
else:
print('results of linear_fit:')
print(' reduced chi squared = ', rchi2)
print(' degrees of freedom = ', dof)
return a, b, sa, sb, rchi2, dof
from scipy.optimize import curve_fit
from pylab import *
def general_fit(f, xdata, ydata, p0=None, sigmay=None, sigmax=None):
"""
Pass all arguments to curve_fit, which uses non-linear least squares
to fit a function, f, to data. Calculate the uncertaities in the
fit parameters from the covariance matrix.
"""
dof = len(ydata) - len(p0)
if sigmax is None:
popt, pcov = curve_fit(f, xdata, ydata, p0, sigmay)
if sigmay is None:
chi2 = sum(((f(xdata,*popt)-ydata))**2)
else:
chi2 = sum(((f(xdata,*popt)-ydata)/sigmay)**2)
rchi2 = chi2/dof
punc = zeros(len(popt))
for i in arange(0,len(popt)):
punc[i] = sqrt(pcov[i,i])
else:
def f_fixed(p,x):
return f(x,*p)
model = odr.Model(f_fixed)
data = odr.RealData(x=xdata, y=ydata, sx=sigmax, sy=sigmay)
od = odr.ODR(data, model, p0)
out = od.run()
popt = out.beta
punc = out.sd_beta
rchi2 = out.res_var
if rchi2 < 1.0 :
sa = sa/sqrt(rchi2)
sb = sb/sqrt(rchi2)
if sigmay is None:
print('results of general_fit: no uncertainties provided, so use with caution')
else:
print('results of general_fit:')
print(' degrees of freedom = ', dof)
print(' reduced chi squared = ', rchi2)
return popt, punc, rchi2, dof