library(gplots)

# This version is for R, tested with R version 3.3.2 and several earlier versions.

# Example commands to fit a line to data.
# Once beyond Phys1001, your data points' errors vary from point to point
# So how do you fit data then?
# The example uses two artificial data sets designed ot have proper statistical properties
# You can put your own data in, or 
# likely want to read data in from a file instead.

xdata <- c(1.0, 2.0, 3.0, 5.0, 7.0, 10.0)
ydata <- c(1.64, 3.22, 3.87, 6.85, 10.96, 14.37)
sigma <- c(0.20, 0.40, 0.50, 0.60, 0.70, 0.80)

# Alternate, less precise data gives larger uncertainties on slope
# The ydata here was synthetically genreated with that sigma
# ydata <- c(1.43, 2.10, 3.20, 6.43, 12.33, 12.75)
# sigma <- c(0.40, 0.80, 1.00, 1.20, 1.40, 1.60)

plot.new()

#plot(xdata,ydata)  # this one won't draw error bars
# will use plotCI instead and plot the data first.
#notice the formatting options surrounding this.
plotCI(xdata,ydata,sigma, lwd=2, xlab="time (ms)", ylab="distance (um)",
       main="illustration of a weighted fit", cex.axis=1.5, cex.lab=1.5, bty="n")
box(lwd=2)

# the non-linear least squares fitter is called nls
# invoke nls with a y=mx+b function and estimated starting parameters
# this is the unweighted form, and probably not what you want.
# comment these out when you are tired of looking at them.
nlmod <- nls(ydata~A*xdata+B, start = list(A=1,B=1))

# draw this on the plot with the data
lines(xdata, predict(nlmod), col="gray", lwd=2)

print("The above is the incorrect fit and physics errors from an unweighted fit.")
print(" ")

print(" ")
print("Now start a correct weighted fit, printing as we go")
print(" ")
# here is the weighted form, and is probably what you want
# here the first argument to nls is any functional form with parameters
# and the last argument is a set of weights, which is just 1/(sigma)^2
wnlmod <- nls(ydata ~ A*xdata+B, start = list(A=1,B=1), weights = sigma^(-2))

lines(xdata, predict(wnlmod), col="blue", lwd=2)


# the covariance matrix doesn't know about our sigma error bars yet
# you can tell if you rerun this putting in the "wrong" sigma above.
print(" ")
print("The correct fit but incorrect error directly from covariance matrix")
print("Doesn't yet know we had assigned uncertainties to the data points")
print("The errors above are just sqrt(diagonal elements)")
sum <- summary(wnlmod)
print(sum)
print(par)
print("Dont use the above parameters yet")

print("printing unweighted covariance, note the error is sqrt(diagonal) terms")
cov <- vcov(wnlmod)
print(cov)


print(" ")
print("So fix up the weighting so it uses the sigma on the data")

# But the above is not what a physicist wants.  We have to do one more step.
# The above best fit is right, but the uncertainties are not right.
# you can figure that out by artificially making sigma smaller
# and contrary to what you expect, the parameter uncertainties don't shrink.
# nls is thinking of sigma as weights, where only relative size is important
# and not as error bars when absolute size is also important.
# this will scale everything to the right absolute size.

# by hand if you are suspicious how it works.  See your Data/Stats book.
# chi = (ydata - predict(wnlmod))/sigma
# chi2hand = sum(chi**2)
# print(chi2hand)

# the fit computes the interesting things inside of it, I'll use those.
# in particular, it returns to you the residuals from the fit
# that is exactly what we need to compute the chisquare
chi2 = sum(residuals(wnlmod)^2/(sigma^2))
print(chi2)
dof = df.residual(wnlmod)
perrscale =  1.0 / (sqrt(chi2 / dof))
print(perrscale)

# not sure this is the right syntax to get the table of parameters and errors
par = summary(wnlmod)$parameters
# make a new copy so I can overwrite with the correct physics errors
newpar <- par

newpar[3] = par[3]*perrscale
newpar[4] = par[4]*perrscale

# but ignore the t-value calculation in nepar.
print(newpar)
print('use the parameter and error but ignore the t value')
print(" ")
print("The correct fit and approximate physics errors from an unweighted fit to data")
print("Scaled the covariance matrix by this chisquare/dof based on our sigmas")
print("First one is b(1) the slope, second one is b(2) intercept")
print(paste("weighted chisquare",chi2,"for dof", dof,"probability",pchisq(chi2,dof)))

# would be cool to draw parameters with higher and lower values.
# there is still a funny thing here.  The two fit parameters are correlated
# some uses of fits to data need to incorporate that information too.