%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt

"""
Explanation: Fitting Models Exercise 2
Imports
End of explanation
"""


data=np.load('decay_osc.npz')
tdata=data['tdata']
ydata=data['ydata']
dy=data['dy']

tdata,ydata,dy

plt.plot(tdata,ydata)

plt.errorbar?

plt.errorbar(tdata,ydata,dy,fmt='k.')

assert True # leave this to grade the data import and raw data plot

"""
Explanation: Fitting a decaying oscillation
For this problem you are given a raw dataset in the file decay_osc.npz. This file contains three arrays:

tdata: an array of time values
ydata: an array of y values
dy: the absolute uncertainties (standard deviations) in y

Your job is to fit the following model to this data:
$$ y(t) = A e^{-\lambda t} \cos{\omega t + \delta} $$
First, import the data using NumPy and make an appropriately styled error bar plot of the raw data.
End of explanation
"""


def model(t,A,o,l,d):
    return A*np.exp(-l*t)*np.cos(o*t)+d

theta_best,theta_cov=opt.curve_fit(model,tdata,ydata,np.array((6,1,1,0)),dy,absolute_sigma=True)
print('A = {0:.3f} +/- {1:.3f}'.format(theta_best[0], np.sqrt(theta_cov[0,0])))
print('omega = {0:.3f} +/- {1:.3f}'.format(theta_best[1], np.sqrt(theta_cov[1,1])))
print('lambda = {0:.3f} +/- {1:.3f}'.format(theta_best[2], np.sqrt(theta_cov[2,2])))
print('delta = {0:.3f} +/- {1:.3f}'.format(theta_best[3], np.sqrt(theta_cov[3,3])))

tfit=np.linspace(0,20,100)
A,o,l,d=theta_best
yfit=A*np.exp(-l*tfit)*np.cos(o*tfit)+d
plt.plot(tfit,yfit)
plt.plot(tdata,ydata,'k.')
plt.xlabel('time')
plt.ylabel('y')
plt.title('Decaying Oscillator')
plt.axhline(0,color='lightgray')

assert True # leave this cell for grading the fit; should include a plot and printout of the parameters+errors

"""
Explanation: Now, using curve_fit to fit this model and determine the estimates and uncertainties for the parameters:

Print the parameters estimates and uncertainties.
Plot the raw and best fit model.
You will likely have to pass an initial guess to curve_fit to get a good fit.
Treat the uncertainties in $y$ as absolute errors by passing absolute_sigma=True.
End of explanation
"""


