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

m = np.array([0.003005, 0.00635, 0.01017, 0.01291, 0.01613, 0.02551]) # кг
f = np.array([0.029, 0.0303, 0.0309, 0.0319, 0.0328, 0.0344]) # Н
df = np.array([0.0022, 0.0046, 0.003, 0.0025, 0.002, 0.0014]) # Н
plt.errorbar(m, f, yerr=df, fmt='o', capsize=5)

x, y, dy = m, f, df

def func(x, a, b): return a * x + b
popt, pcov = curve_fit(func, x, y, p0=[y[0]/x[0], y[0]], sigma=dy, absolute_sigma=1)
a, b = popt
Sa, Sb = np.sqrt(np.diag(pcov))
print(f'{a = :f} +- {Sa:f}')
print(f'{b = :f} +- {Sb:f}')

xx = np.linspace(0, 0.03, 100)
plt.plot(xx, a*xx+b, zorder=1000)

xm, ym = np.mean(x), np.mean(y)
r = np.sum( (x-xm)*(y-ym) ) / np.sqrt( np.sum((x-xm)**2) * np.sum((y-ym)**2) )
print(f'{r = :f}')

plt.xlabel('m, г')
plt.ylabel('<f>, Н')
plt.title('Зависимость средней силы удара от массы пули')
plt.savefig('example_6_linear.pdf', bbox_inches='tight')