#406 Nonlinear least squares fitting using leastsq from scipy.optimize in python

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import leastsq

A, freq, tau = 5, 4, 0.5

def f(t, A, freq, tau):
    return A * np.exp(-t/tau) * np.cos(3*np.pi * freq * t)

tmax, dt = 1, 0.01
t = np.arange(0, tmax, dt)
yexact = f(t, A, freq, tau)
y = yexact + np.random.randn(len(yexact))*2

def residuals(p, y, t):
    A, freq, tau = p
    return y - f(t, A, freq, tau)

# nonlinear least squares fitting routine
p0 = 5, 5, 1
plsq = leastsq(residuals, p0, args=(y, t))
plt.plot(t, y, 'o', c='black', label='Data')
plt.plot(t, yexact, c='red', label='Exact')
pfit = plsq[0]
plt.plot(t, f(t, *pfit), c='green', label='Fit')
plt.legend()
plt.savefig("ex406.png", dpi=100)
plt.show()