Nonlinear least squares fitting

import os
import numpy as np
import matplotlib.pyplot as plt

os.chdir(r'D:datascriptswordpressex54')
os.getcwd()

A,freq,tau=10,4,0.5

def f(t,A,freq,tau):
return A*np.exp(-t/tau)*np.cos(2*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
plt.plot(t,yexact)
plt.scatter(t,y,marker='o',facecolors='blue', edgecolors='black')

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

from scipy.optimize import leastsq
p0=5,5,1
plsq=leastsq(residuals,p0,args=(y,t))
plsq[0][1]
pfit = (plsq[0][0],plsq[0][1],plsq[0][2])

plt.plot(t,y,'o',c='k',label='Data')
plt.plot(t,yexact,c='green',label='Exact')
plt.plot(t,f(t,plsq[0][0],plsq[0][1],plsq[0][2]),c='red',label='Fit')
plt.legend()

plt.tight_layout()
plt.savefig("example54.png", dpi=100)
plt.show()
plt.close()