pset5.py

# curvefit .py Problem Sheet 5 example 
import numpy as np
from scipy.optimize import curve_fit 
import matplotlib.pyplot as plt
def func(x, a, b, c):
    return a*np.exp(-b*x) + c # function to generate data for curve fit
# The first part of this program generates a set of data that follows a particular functional form that allows us to test the curvefit routine
x = np.linspace(0, 4, 50)
y = func(x, 3.0, 1.3, 5)
yn = y + 0.2*np.random.normal(size=len(x)) # adding some noise to the data points
# Now that we have our data we can attempt to fit to the curve
popt, pcov = curve_fit(func, x, yn) # performing curve fit, and returning parameters
print 'Parameters : ', popt
print 'Covariance : ', pcov
# graphical output of results
fig=plt.figure()
plt.scatter(x,y, label='data')
plt.scatter(x,yn, color='r', label='data + noise')
plt.plot(x, func(x, popt[0], popt[1], popt[2]), color='green', label='best fit') 
plt.legend()
plt.show()