# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 15:07:27 2020

@author: jbobowsk
"""

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

# In this example, we take data from a series LRC circuit.  For the y-axis
# we will have the ratio of the magntidue of the voltage across the
# resistance to that supplied by the function generator driving the
# circuit.  On the x-axis we will have frequency.
import numpy as np

# Enter the data.
Vratio= np.array([0.198e-1, 0.67e-1, .117, .185, .331, .450, .573, .689, .718,\
                  .714, .704, .670, .631, .549, .412, .318, .249, .186, .138])
errVratio= np.array([0.1e-2, 0.2e-2, 0.3e-2, 0.3e-2, 0.5e-2, 0.6e-2, 0.7e-2,\
                     0.8e-2, 0.8e-2, 0.8e-2, 0.8e-2, 0.8e-2, 0.8e-2, 0.7e-2,\
                     0.6e-2, 0.4e-2, 0.5e-2, 0.3e-2, 0.3e-2])
f = np.array([3000, 10000, 15000, 20000, 25000, 28000, 30000, 32000, 33000,\
              34000, 35000, 36000, 37000, 39000, 43000, 47000, 52000, 60000,\
              70000])
import math
omega = (2*math.pi*f)

# Plot the data using errorbar(x,y,e).
import matplotlib.pyplot as plt
plt.errorbar(omega, Vratio, errVratio, fmt = 'ko', markersize = 5,\
                 linewidth = 1.8,\
                 markeredgecolor = 'b',\
                 markerfacecolor = (.49, 1, .63),\
                 capsize = 5)
plt.xlabel('Angular Frequency (rad/s)')
plt.ylabel('Voltage Ratio')

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

# The first step is to define a function for the model that we will fit our
# data to.  In this case, the modeil is a quadratic.
def LRCFunc(x, A, width, w0):
    y = A/np.sqrt(1+(x/width)**2*(1-(w0/x)**2)**2)
    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.
from scipy.optimize import curve_fit

# To give the fit a chance at being successful, we have to provide reasonable
# initial guesses for the fit parameters.  We use the option 'p0' for the 
# list of starting parameters.  The order must be in the order of the parameters
# defined in the function 'LRCFunc()'.
start = (0.7, 0.5e5, 2e5)
a_fit, cov = curve_fit(LRCFunc, omega, Vratio, p0 = start)

print('The best-fit parameters are:\n A =', a_fit[0], '\n',\
      'width =', a_fit[1], '\n',\
      'w0 =', a_fit[2])

# 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 (Note the use of the unicode
# character \u0394 to print $\Delta$):
print('The errors in the parameters are:\n \u0394A =', np.sqrt(np.diag(cov))[0],\
      '\n', '\u0394width =', np.sqrt(np.diag(cov))[1],\
      '\n', '\u0394w0 =', np.sqrt(np.diag(cov))[2])

# Use the function that we defined to generate the best-fit curve.
xx = np.arange(1, 8e5, 100)
plt.plot(xx, LRCFunc(xx, a_fit[0], a_fit[1], a_fit[2]), 'k-')
plt.axis((0, 4.6e5, 0, 0.75))

# 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.  We will
# also specify that the errors are absolute errors with the option
# 'absolute_sigma=True'.
a_fit, cov = curve_fit(LRCFunc, omega, Vratio, sigma = errVratio,\
                       p0 = start, absolute_sigma = True)

# Extract the fit parameters and their uncertainties.
print('The best-fit parameters are:\
      \n A \u00B1 \u0394A =', a_fit[0], '\u00B1', np.sqrt(np.diag(cov))[0],\
     '\n width \u00B1 \u0394width =', a_fit[1], '\u00B1', np.sqrt(np.diag(cov))[1],\
     '\n w0 \u00B1 \u0394w0 =', a_fit[2], '\u00B1', np.sqrt(np.diag(cov))[2])

# Plot the data.
plt.figure()
plt.errorbar(omega, Vratio, errVratio, fmt = 'ko', markersize = 5,\
                 linewidth = 1.8,\
                 markeredgecolor = 'b',\
                 markerfacecolor = 'r',\
                 capsize = 5)
plt.xlabel('Angular Frequency (rad/s)')
plt.ylabel('Voltage Ratio')

# Plot the best-fit line.
xx = np.arange(1, 8e5, 100)
plt.plot(xx, LRCFunc(xx, a_fit[0], a_fit[1], a_fit[2]), 'k-')
plt.axis((0, 4.6e5, 0, 0.75))