import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
plt.style.use('ff.mplstyle')

x = [3.05, 3.1, 3.15, 3.2, 3.25, 3.31, 3.36, 3.4, 3.41, 3.46, 3.5, 3.52, 3.55, 3.56, 3.61, 3.65, 3.66, 3.7, 3.71, 3.76, 3.81, 3.81, 3.85, 3.91, 4.01]
y = [0.16, 0.17, 0.21, 0.26, 0.29, 0.31, 0.47, 0.55, 0.62, 0.78, 0.85, 0.88, 0.64, 0.64, 0.39, 0.34, 0.32, 0.25, 0.28, 0.22, 0.17, 0.22, 0.18, 0.14, 0.12]
# x, кГц; y, см

plt.scatter(x, y, color='k')

def amp(w, w0, delta, f0):
    return f0 / np.sqrt( (w0**2 - w**2)**2 + 4 * delta**2 * w**2 )

popt, pcov = curve_fit(amp, x, y, p0=(3.5, 0.5, 1))
w0, delta, f0 = popt
Sw0, Sdelta, Sf0 = np.sqrt(np.diag(pcov))
print(f'{   w0 = : f} +- {Sw0:f}')
print(f'{delta = : f} +- {Sdelta:f}')
print(f'{   f0 = : f} +- {Sf0:f}')

xx = np.linspace(3, 4, 200)
plt.plot(xx, amp(xx, *popt), color='r')

plt.xlabel(r'Частота $\omega$, Гц')
plt.ylabel('Амплитуда А, см')
plt.title('Амплитудно-частотная характеристика пружинного маятника')
plt.savefig('example_7_nonlinear.pdf', bbox_inches='tight')