Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
# DO NOT CHANGE THIS FILE!
2
#
3
# This file contains the functions linear_fit for fitting a straight
4
# line to data and general_fit for fitting any user-defined funciton
5
# to data.  To use either of them,  the first line of your program
6
# should be "from fitting import *".
7

8
from __future__ import print_function  # in case Python 2 is used
9

10
from math import sqrt
11
from scipy import odr
12

13

14

15
def linear_fit(xdata, ydata, sigmay=None, sigmax=None):
16
    """
17
    Performs a linear fit to data.
18

19
    Parameters
20
    ----------
21
    xdata : An array of length N.
22
    ydata : An array of length N.
23

24
    sigmay : None or an array of length N.
25
    sigmax : None or an array of length N.
26
        If one is provided, it is the standard deviation of ydata.  Analytical linear regression used.
27
        If both are provided, they are the standard deviations of ydata and xdata, respectively. ODR is used.
28

29
    Returns
30
    -------
31
    a, b   : Optimal parameter of linear fit (y = a*x + b)
32
    sa, sb : Uncertainties of the parameters
33
    """
34

35
    def lin_func(p, x):
36
        a, b = p
37
        return a*x + b
38
    
39
    dof = len(ydata) - 2
40

41
    if sigmax is None:
42
        if sigmay is None:
43
            w = ones(len(ydata)) # Each point is equally weighted.
44
        else:
45
            w=1.0/(sigmay**2)
46

47
        sw = sum(w)
48
        wx = w*xdata # this product gets used to calculate swxy and swx2
49
        swx = sum(wx)
50
        swy = sum(w*ydata)
51
        swxy = sum(wx*ydata)
52
        swx2 = sum(wx*xdata)
53

54
        a = (sw*swxy - swx*swy)/(sw*swx2 - swx*swx)
55
        b = (swy*swx2 - swx*swxy)/(sw*swx2 - swx*swx)
56
        sa = sqrt(sw/(sw*swx2 - swx*swx))
57
        sb = sqrt(swx2/(sw*swx2 - swx*swx))
58

59
        if sigmay is None:
60
            chi2 = sum(((a*xdata + b)-ydata)**2)
61
        else:
62
            chi2 = sum((((a*xdata + b)-ydata)/sigmay)**2)
63
        rchi2 = chi2/dof
64

65
    else:
66
        # make the initial guess a line passing through first and last points
67
        a0 = (ydata[-1]-ydata[0])/(xdata[-1]-xdata[0])
68
        b0 = ydata[0]-xdata[0]*a0
69
        model = odr.Model(lin_func)
70
        data = odr.RealData(x=xdata, y=ydata, sx=sigmax, sy=sigmay)
71
        od = odr.ODR(data, model, [a0,b0])
72
        out = od.run()
73
        a,b = out.beta
74
        sa,sb = out.sd_beta
75
        rchi2 = out.res_var
76

77
#    From: https://www.physics.utoronto.ca/~phy326/python/odr_fit_to_data.py
78
#    scipy.odr scales the parameter uncertainties by the reduced chi 
79
#    square (out.res_var).  If the fit is poor, i.e. reduced chisq is 
80
#    large, the uncertainties are scaled up, which makes sense. If the 
81
#    fit is "too good", i.e. reduced chisq << 1, it suggests that the 
82
#    uncertainties may have been overestimated, but it seems risky to 
83
#    scale down the uncertainties. 
84
        if rchi2 < 1.0 :
85
            sa = sa/sqrt(rchi2)
86
            sb = sb/sqrt(rchi2)
87

88
    if sigmay is None:
89
        print('results of linear_fit: no uncertainties provided, so use with caution')
90
    else:
91
        print('results of linear_fit:')
92

93
#    print('   chi squared = ', chi2)
94
    print('   reduced chi squared = ', rchi2)
95
    print('   degrees of freedom = ', dof)
96

97

98
    return a, b, sa, sb, rchi2, dof
99

100

101
from scipy.optimize import curve_fit
102
from pylab import *  # for array, zeros, arange
103

104

105
#def general_fit(f, xdata, ydata, p0=None, sigma=None, **kw):
106
def general_fit(f, xdata, ydata, p0=None, sigmay=None, sigmax=None):
107

108
    """
109
    Pass all arguments to curve_fit, which uses non-linear least squares
110
    to fit a function, f, to data.  Calculate the uncertaities in the
111
    fit parameters from the covariance matrix.
112
    """
113

114
    dof = len(ydata) - len(p0)
115

116
    if sigmax is None:
117
        popt, pcov = curve_fit(f, xdata, ydata, p0, sigmay)
118

119
        if sigmay is None:
120
            chi2 = sum(((f(xdata,*popt)-ydata))**2)
121
        else:
122
            chi2 = sum(((f(xdata,*popt)-ydata)/sigmay)**2)
123
#        dof = len(ydata) - len(popt)
124
        rchi2 = chi2/dof
125
        # The uncertainties are the square roots of the diagonal elements
126
        punc = zeros(len(popt))
127
        for i in arange(0,len(popt)):
128
            punc[i] = sqrt(pcov[i,i])
129
    else:
130
        # ODR expects a funciton with parameters, then x
131
        def f_fixed(p,x):
132
            return f(x,*p)
133
        
134
        model = odr.Model(f_fixed)
135
        data = odr.RealData(x=xdata, y=ydata, sx=sigmax, sy=sigmay)
136
        od = odr.ODR(data, model, p0)
137
        out = od.run()
138
        popt = out.beta
139
        punc = out.sd_beta
140
        rchi2 = out.res_var
141

142
#    From: https://www.physics.utoronto.ca/~phy326/python/odr_fit_to_data.py
143
#    scipy.odr scales the parameter uncertainties by the reduced chi 
144
#    square (out.res_var).  If the fit is poor, i.e. reduced chisq is 
145
#    large, the uncertainties are scaled up, which makes sense. If the 
146
#    fit is "too good", i.e. reduced chisq << 1, it suggests that the 
147
#    uncertainties may have been overestimated, but it seems risky to 
148
#    scale down the uncertainties. 
149
        if rchi2 < 1.0 :
150
            punc = punc/sqrt(rchi2)
151
    if sigmay is None:
152
        print('results of general_fit: no uncertainties provided, so use with caution')
153
    else:
154
        print('results of general_fit:')
155
#    print('   chi squared = ', chi2)
156
    print('   degrees of freedom = ', dof)
157
    print('   reduced chi squared = ', rchi2)
158

159
    return popt, punc, rchi2, dof
160