import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt
np.random.seed(0)

# Our test function
def f(t, omega, phi):
    return np.cos(omega * t + phi)

def g(t, a, b, c, d, e, f):
    return f*t**5+e*t**4+d*t**3+a*t**2+b*t+c


# Our x and y data
x = np.linspace(0, 3, 50)
y = f(x, 1.5, 1) + .1*np.random.normal(size=50)
# Fit the model: the parameters omega and phi can be found in the # params vector
params, params_cov = optimize.curve_fit(g, x, y)
print(params)
# plot the data and the fitted curve
t = np.linspace(0, 3, 1000)
plt.figure(1)
plt.clf()
plt.plot(x, y, 'bx')
plt.plot(t, f(t, 1.5, 1), 'g-')
plt.plot(t, g(t, *params), 'r-')
plt.show()