import numpy as np
import matplotlib.pyplot as plt
import math
from scipy.stats import chi2
# the chisq
def getChisq(y0,y,e2):
return ((y-y0)*(y-y0)/e2).sum()
# the function
def f(x,p):
return p[0]/(1+p[1]*x)
# These are the derivatives
# i = 0 gives df/dp0
# i = 1 gives df/dp1
def dydp(i, p, x):
temp = 1 + p[1]*x
if i == 0:
return 1./temp
elif i == 1:
return -p[0]*x/(temp*temp)
else:
print("Illegal call to dydp with i = ", i)
return -1
# The dataset : x, y, and the errors in y
# Note: y was generated as y=2/(1+x) + random_number_between(-0.5 and 0.5)
x = np.array([0., 0.5, 1., 1.5, 1.8, 2.1])
y = np.array([1.77, 1.59, 0.83, 0.56, 0.34, 1.11])
N = len(x) # number of points
e2 = np.full(N, 1.)/12. # variance of y
# number of iterations
niter = 10 # a lot!!
# The inverse covariance matrix of the y's
W = np.diag(1./e2)
# we will fit as y = p[0] / (1 + p[1]*x)
p = np.array([12., 10.]) # initial guess (way off!!!!)
y0 = f(x,p)
for iteration in range(niter):
At = np.array( [dydp(0,p,x), dydp(1,p,x)] ) # 2xN matrix
A = (At.T).copy() # Nx2 matrix
dy = (np.array( [(y-y0),] )).T # Nx1 column vector
# debug: look at these matrices on the 1st iteration
if iteration == 0:
print(" ")
print("-- At ---")
print(At)
print(" ")
print("-- A ---")
print(A)
print(" ")
print("-- W ---")
print(W)
print(" ")
print("-- dy ---")
print(dy)
# find the matrix to be inverted, and invert it
temp = np.matmul(At, W) # 2xN * NxN = 2xN
temp2 = np.matmul(temp, A) # 2xN * Nx2 = 2x2
temp3 = np.linalg.inv(temp2) # 2x2 ... this is the covariance matrix
# multiply again
temp4 = np.matmul(temp3, At) # 2x2 * 2xN = 2xN
temp5 = np.matmul(temp4, W) # 2xN * NxN = 2xN
dpar = np.matmul(temp5, dy) # 2xN * Nx1 = 2x1 column vector
# the new values of the parameters
p[0] = p[0] + dpar[0][0]
p[1] = p[1] + dpar[1][0]
# the currently fitted value of y
y0 = f(x,p)
# the chisq
chisq = getChisq(y, y0, e2)
# output stuff
print(" ")
print("Iteration no. ", iteration+1)
print("Chisq = ", chisq)
print("A = p[0] = ", p[0])
print("B = p[1] = ", p[1])
print("Covariance:")
print(temp3)
-- At --- [[ 1. 0.16666667 0.09090909 0.0625 0.05263158 0.04545455] [-0. -0.16666667 -0.09917355 -0.0703125 -0.0598338 -0.05206612]] -- A --- [[ 1. -0. ] [ 0.16666667 -0.16666667] [ 0.09090909 -0.09917355] [ 0.0625 -0.0703125 ] [ 0.05263158 -0.0598338 ] [ 0.04545455 -0.05206612]] -- W --- [[12. 0. 0. 0. 0. 0.] [ 0. 12. 0. 0. 0. 0.] [ 0. 0. 12. 0. 0. 0.] [ 0. 0. 0. 12. 0. 0.] [ 0. 0. 0. 0. 12. 0.] [ 0. 0. 0. 0. 0. 12.]] -- dy --- [[-10.23 ] [ -0.41 ] [ -0.26090909] [ -0.19 ] [ -0.29157895] [ 0.56454545]] Iteration no. 1 Chisq = 15.302731583955243 A = p[0] = 1.7691492997948757 B = p[1] = 2.175730103746452 Covariance: [[0.08332243 0.07966512] [0.07966512 1.78214122]] Iteration no. 2 Chisq = 336.0119810810809 A = p[0] = 1.7967300735822131 B = p[1] = -0.26941802164603246 Covariance: [[0.08207281 0.15432027] [0.15432027 0.95699781]] Iteration no. 3 Chisq = 43.58890636881732 A = p[0] = 1.1148650427694777 B = p[1] = -0.16716553241190188 Covariance: [[0.02409395 0.00344169] [0.00344169 0.00062615]] Iteration no. 4 Chisq = 16.482413819807512 A = p[0] = 1.3956871593105673 B = p[1] = 0.1308860859432799 Covariance: [[0.03481601 0.01229534] [0.01229534 0.00573888]] Iteration no. 5 Chisq = 7.66429797056098 A = p[0] = 1.7510115409801505 B = p[1] = 0.584555797763564 Covariance: [[0.05590782 0.03186749] [0.03186749 0.02678354]] Iteration no. 6 Chisq = 6.410809864219388 A = p[0] = 1.8506032711652856 B = p[1] = 0.8664517292972878 Covariance: [[0.07133768 0.05704988] [0.05704988 0.07996813]] Iteration no. 7 Chisq = 6.385058573379272 A = p[0] = 1.8472731978151837 B = p[1] = 0.9051073348295428 Covariance: [[0.0758692 0.07231552] [0.07231552 0.13516539]] Iteration no. 8 Chisq = 6.385028685203338 A = p[0] = 1.8457950049412064 B = p[1] = 0.903169790129984 Covariance: [[0.07632048 0.07490154] [0.07490154 0.14640934]] Iteration no. 9 Chisq = 6.385028540993686 A = p[0] = 1.845872135980709 B = p[1] = 0.903320916984551 Covariance: [[0.07629864 0.07483874] [0.07483874 0.14609163]] Iteration no. 10 Chisq = 6.385028540126873 A = p[0] = 1.8458661304818023 B = p[1] = 0.9033091988170872 Covariance: [[0.07630034 0.07484519] [0.07484519 0.14612245]]
# Nicely formatted output
print("FIT RESULTS")
print("-----------")
print("A = " + "%.2f " %p[0] + u"\u00B1" +
" %.2f" %math.sqrt(temp3[0][0]))
print("B = " + "%.2f " %p[1] + u"\u00B1" +
" %.2f" %math.sqrt(temp3[1][1]))
print("chisq = " + "%.2f " %chisq+ "for " +
str(len(x)-2) + " dof")
print("prob = %.2f" %(1-chi2.cdf(chisq, len(x)-2)))
FIT RESULTS ----------- A = 1.85 ± 0.28 B = 0.90 ± 0.38 chisq = 6.39 for 4 dof prob = 0.17
# Now plot it
a = plt.subplot(111)
xmin = x.min()-0.2
xmax = x.max()+0.2
a.errorbar(x, y, yerr=np.sqrt(e2), linestyle='none', color="red", marker='o')
a.set_xlim(xmin, xmax)
xpl = np.linspace(xmin, xmax, 1000)
ypl = f(xpl, p)
a.plot(xpl, ypl, color='blue', linestyle='solid')
a.set_xlabel('X')
a.set_ylabel('Y')
# txtRes = "A = " + str(round(p[0],2)) + "$\pm$"
txtRes = "A = " + "%.2f" %p[0] + " $\pm$ "
txtRes = txtRes + "%.2f" %math.sqrt(temp3[0][0]) + "\n"
txtRes = txtRes + "B = " + "%.2f" %p[1] + " $\pm$ "
txtRes = txtRes + "%.2f" %math.sqrt(temp3[1][1])
props = dict(boxstyle='round', facecolor='white', alpha=0.8)
a.text(0.96, 0.96, txtRes,
transform=a.transAxes,
bbox=props, fontsize='large',
horizontalalignment='right',
verticalalignment='top' )
a.grid()
# Cute addition to the simple plot. Add a 1 sigma band
# Code is not very pythonesque...
# Note: temp3 is the covariance matrix
dy = np.empty(len(xpl))
for ii in range(len(xpl)):
xx = xpl[ii]
Bt = np.array( [dydp(0,p,xx), dydp(1,p,xx)] ) # 2x1 matrix
B = (Bt.T).copy() # 1x2 matrix
blah = np.matmul(B, temp3) # 1x2 * 2x2 = 1x2
dy2 = np.matmul(blah, Bt) # 1x2 * 2x1 = 1x1
dy[ii] = math.sqrt(dy2)
a2 = plt.subplot(111)
a2.errorbar(x, y, yerr=np.sqrt(e2), linestyle='none', color="red", marker='o')
a2.set_xlim(xmin, xmax)
a2.plot(xpl, ypl, color='blue', linestyle='solid')
a2.plot(xpl, ypl-dy, color='blue', linestyle='dotted')
a2.plot(xpl, ypl+dy, color='blue', linestyle='dotted')
a2.fill_between(xpl, ypl+dy, ypl-dy, color='blue', alpha=.2)
a2.set_xlabel('X')
a2.set_ylabel('Y')
a2.text(0.96, 0.96, txtRes,
transform=a2.transAxes,
bbox=props, fontsize='large',
horizontalalignment='right',
verticalalignment='top' )
a2.grid()
# Now a scan of the chisq
# We calculate the chisq for various values of p[0] and p[1]
# Note: temp3[0][0] and temp3[1][1] are the diagonal
# elements of the covariance matrix, ie, the square of the errors
# on p[0] and [p1],
# We will scan out to 5 sigma
ax2 = plt.subplot(111)
p0min = p[0] - 5*np.sqrt(temp3[0][0])
p0max = p[0] + 5*np.sqrt(temp3[0][0])
p1min = p[1] - 5*np.sqrt(temp3[1][1])
p1max = p[1] + 5*np.sqrt(temp3[1][1])
nscan = 100
p0 = np.linspace(p0min, p0max, nscan)
p1 = np.linspace(p1min, p1max, nscan)
ax2.set_xlim(p0min,p0max)
ax2.set_ylim(p1min,p1max)
z = np.zeros( shape=(nscan,nscan) ) # an emptyt array for now
for i0 in range(0, nscan):
this_p0 = p0[i0]
for i1 in range(0, nscan):
this_p1 = p1[i1]
yy = f(x, [this_p0, this_p1])
z[i0][i1] = getChisq(y,yy,e2) - chisq
# "z" is a map of chisq-chisq_at_minimum
# Correspond to the 68% 95% 99% coverage for 2 parameters
# (PDG Table 38.2)
CS = ax2.contour(p0, p1, z, [2.30, 5.99, 9.21], colors=['red','blue','darkgreen'])
fmt = {}
strs = [ '68%', '95%', '99%' ]
for l,s in zip( CS.levels, strs ):
fmt[l] = s
ax2.clabel(CS, inline=True, fmt=fmt, fontsize=10)
# put a point at the best fit
ax2.plot(p[0], p[1], 'ko')
# label the axes
ax2.set_xlabel('A or p[0]')
ax2.set_ylabel('B or p[1]')
# The resut
ax2.text(1-0.96, 0.96, txtRes,
transform=ax2.transAxes,
bbox=props, fontsize='large',
horizontalalignment='left',
verticalalignment='top' )
# put a grid
ax2.grid()