# In this script we will fit a linear function to a set of experimental
# data.  The fit will be weighted by the measurement uncertainties.

# Updated by Jordan Andrews on June 9, 2021 with the use of 
# np.polynomial.Polynomial([a, b, c]).


# In this example, our data will be the voltage across and the current through a resistor.
# The slope of our data will tell us the resistance of the resistor.
import numpy as np

# Enter the data.
Irms= np.array([0.00907, 0.00794, 0.00642, 0.00472, 0.00296, 0.001911, 0.000467, 0.0000469])
errIrms= np.array([1.9e-4, 1.8e-4, 1.6e-4, 1.5e-4, 1.3e-4, 1e-4, 5e-5, 1e-5])
Vp2p = np.array([5.62, 4.94, 4.02, 2.96, 1.86, 1.19, 0.302, 0.047])
errVp2p = np.array([0.04, 0.04, 0.04, 0.04, 0.04, 0.03, 0.01, 0.008])

# Since the current was measured as an rms value, convert the peak-to-peak
# voltage measurements to rms by dividing by 2*sqrt(2).
Vrms = Vp2p/(2*np.sqrt(2))
errVrms = errVp2p/(2*np.sqrt(2))

# Notice that the current data has larger relative errors than the
# voltage data (typically by a factor of two or more):
Irel = errIrms/Irms
Vrel = errVrms/Vrms
print(Irel/Vrel)

# Because the relative error in current is larger, we will plot the data as
# current vs voltage and show the error bars only along the y-axis.
# Use plt.errorbar(x,y,e).
import matplotlib.pyplot as plt
plt.errorbar(Vrms, Irms, errIrms, fmt = 'ko', markersize = 5,\
                 linewidth = 1.8,\
                markeredgecolor = 'b',\
                markerfacecolor = 'yellow',\
                capsize = 5)
plt.xlabel('RMS Voltage (volts)')
plt.ylabel('RMS Current (amps)')
plt.title('Unweighted Linear Fit')

# To do the actual fit, we will use the 'curve_fit()' function from, the
# 'SciPy' module.  This way of fitting is very nice because will be able to
# use it for all types of fit models (linear, polynomial, line-in-parameter fits,
# and nonlinear fit).  As we will see, it is capable of doing both unweighted
# and weighted fits and it will return uncertainties in the fit parameters.
from scipy.optimize import curve_fit

# The first step is to define a function for the model that we will fit our
# data to.  In this case, the modeil is just the equation of a staight a line.
def linearFunc(x, intercept, slope):
    y = slope*x + intercept
    return y

# Here is the actual command to execute the fit.  At a minimum, 'curve_fit()' 
# requires as inputs the function that defines the model, the x-data, and 
# the y-data.  The statement below tells 'curve_fit()' to return a list of the 
# the best-fit parameters and the covariance matrix which will be used to 
# determine the error in the fit parameters.
a_fit, cov = curve_fit(linearFunc, Vrms, Irms)

print('The best-fit parameters are: (1) slope =', a_fit[1], 'and (2) intercept',\
      a_fit[0])
print('Here is the covariance matrix:\n', cov)

# The uncertainties of the best-fit parameters are determined from the 
# square roots of the diagonal elements of the covariance matrix.  We
# can select the diagonal elements using:
print(np.diag(cov))
print('The error in the slop is:', np.sqrt(np.diag(cov))[1], '\n',\
      'The error in the intercept is:', np.sqrt(np.diag(cov))[0])

# NumPy  has a nice package for polynomials, called 'polynomial'. There are six
# different polynomial types in this package. For our case we are dealing with
# a simple power series. You will use the 'Polynomial' constructor for this.
# y = np.polynomial.Polynomial([a, b, c]) results in y = a + b*x + c*x**2. 
# Here's the best-fit fucntion obtained using 'a_fit' from 'curve_fit()' and 
# the built in polynomial package of NumPy.
fitFcn = np.polynomial.Polynomial(a_fit)

# We can then evaluate 'fitFcn(x)' at any value of x that we want.
print(fitFcn(0))
print(fitFcn(1))
    
# This gives us an easy way to plot the fit on top of the data.
xx = np.arange(-1, 10, 0.1)
plt.plot(xx, fitFcn(xx), 'k-')
plt.axis((0, 2.2, 0, 0.01))

# All of this has produced an "unweighted" fit to the data.  To include 
# weights, all we need to do is include another option in 'curve_fit()'.
# Everything else is exactly the same!  The new option is 'sigma' and it is
# simply a list of the errors in the y-values.  Note that many fitting
# routines require you to provide the actual weights as 1/sigma^2.  That is
# not the case here.  You just have to provide the absolute y-uncertainties.
# Here's a block of code that does the fit, extracts the best fit function, the
# best-fit parameters, the uncertainties, and then plots the result.
a_fit, cov = curve_fit(linearFunc, Vrms, Irms, sigma = errIrms)

# Extract the fit parameters and their uncertainties.
print('The best-fit parameters are: (1) slope =', a_fit[1], 'and (2) intercept',\
      a_fit[0])
print('The error in the slop is:', np.sqrt(np.diag(cov))[1], '\n',\
      'The error in the intercept is:', np.sqrt(np.diag(cov))[0])

# Plot the data.
plt.figure()
plt.errorbar(Vrms, Irms, errIrms, fmt = 'ko', markersize = 5,\
                 linewidth = 1.8,\
                markeredgecolor = 'b',\
                markerfacecolor = 'r',\
                capsize = 5)
plt.xlabel('RMS Voltage (volts)')
plt.ylabel('RMS Current (amps)')
plt.title('Weighted Linear Fit')

# Get the best-fit line.
fitFcn = np.polynomial.Polynomial(a_fit)

# Plot the best-fit line.
xx = np.arange(-1, 10, 0.1)
plt.plot(xx, fitFcn(xx), 'k-')
plt.axis((0, 2.2, 0, 0.01))