# Инициализация
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit

plt.style.use('ff.mplstyle')

# 1. Простейший график (раздел 3.3)

t = [1, 2, 3, 4, 5]                        # с
x = [0.2415, 0.276, 0.3335, 0.414, 0.5175] # м

plt.scatter(t, x)
plt.xlabel("t, c")
plt.ylabel("x, м")
plt.savefig("example_1_simple.pdf")

# 2. Простой график с погрешностями (раздел 3.4)

m = [0.003005, 0.00635, 0.01017, 0.01291, 0.01613, 0.02551] # кг
f = [0.029, 0.0303, 0.0309, 0.0319, 0.0328, 0.0344]         # Н
df = [0.0022, 0.0046, 0.003, 0.0025, 0.002, 0.0014]         # Н

m = np.array(m)
m = m * 1e3  # кг -> г

plt.errorbar(m, f, yerr=df, fmt='o', capsize=5)
plt.xlabel('m, г')
plt.ylabel('<f>, Н')
plt.title('Зависимость средней силы удара от массы пули')
plt.savefig('example_2_errors.pdf')

# 3. Несколько графиков на одних осях (раздел 3.5)

x = [0.2, 0.4, 0.6, 0.8, 1.0]                  # м
v2_1 = np.array([9, 18, 26, 29, 42]) * 1e-3    # м^2/с^2
v2_2 = np.array([15, 34, 47, 65, 85]) * 1e-3   # м^2/с^2
v2_3 = np.array([20, 49, 75, 95, 127]) * 1e-3  # м^2/с^2

plt.scatter(x, v2_1, label=r'$\alpha = 0.11^\circ$')
plt.scatter(x, v2_2, label=r'$\alpha = 0.23^\circ$')
plt.scatter(x, v2_3, label=r'$\alpha = 0.37^\circ$')
plt.xlim(xmin=0)
plt.ylim(ymin=1e-6)
plt.xlabel('$x$, м')
plt.ylabel(r'$v^2, \frac{м^2}{c^2}$')
plt.title('Зависимость квадрата скорости тележки $v^2$ \n от пройденного ей расстояния $x$. Случай тележки без груза')
plt.legend(title='Наклон скамьи')
plt.savefig('example_3_sameaxis.pdf')

# 4. Несколько графиков на разных осях (раздел 3.6)

I = [0.06, 0.1, 0.15, 0.21, 0.28, 0.35, 0.43, 0.5, 0.58, 0.65, 0.73, 0.81, 0.89, 0.96, 1.04, 1.11, 1.19, 1.27, 1.34] # mV
U = [320, 415, 490, 536, 557, 567, 573, 579, 583, 587, 590, 593, 595, 598, 600, 602, 604, 606, 607] # mA

fig, axes = plt.subplots(2, 1, figsize=(5, 8), sharex=True, gridspec_kw=dict(hspace=0.05))
ax1, ax2 = axes

ax1.scatter(U, I)
ax1.set_title('ВАХ полупроводникового диода')
ax1.set_ylabel('I, мА')
ax1.set_ylim(ymin=0, ymax=1.4)

I = np.array(I)
ax2.scatter(U, np.log(I))
ax2.set_xlabel('U, мВ')
ax2.set_ylabel('ln I')
ax2.set_xlim(xmin=300, xmax=650)
ax2.set_ylim(-3, 0.5)

fig.savefig('example_4_multipleaxes.pdf', bbox_inches='tight')

# 5. Случай пропорциональной зависимости (раздел 3.7.1)

x = np.array([0.2, 0.4, 0.6, 0.8, 1.0])       # м
v2_1 = np.array([9, 18, 26, 29, 42]) * 1e-3   # м^2/с^2
v2_2 = np.array([15, 34, 47, 65, 85]) * 1e-3  # м^2/с^2
v2_3 = np.array([20, 49, 75, 95, 127]) * 1e-3 # м^2/с^2
plt.scatter(x, v2_1, label=r'$\alpha = 0.11^\circ$')
plt.scatter(x, v2_2, label=r'$\alpha = 0.23^\circ$')
plt.scatter(x, v2_3, label=r'$\alpha = 0.37^\circ$')

xx = np.linspace(0, 1.1, 100)
for y in [v2_1, v2_2, v2_3]:
    s0 = 1.5e-7
    a = np.sum(x * y) / np.sum(x**2)
    Sa = np.sqrt(s0**2 / (np.sum(x**2))**2)
    print(a, Sa)
    
    ym = np.mean(y)
    SSres = np.sum((y - a*x)**2)
    SStot = np.sum((y - ym)**2)
    R2 = 1 - SSres / SStot
    print(R2)
    
    plt.plot(xx, a*xx, label=fr'{a =: .2e} $\pm$ {Sa:.2e};  $R^2$ = {R2:.3f}')

plt.xlim(xmin=0)
plt.ylim(ymin=1e-6)
plt.xlabel('$x$, м')
plt.ylabel(r'$v^2, \frac{м^2}{c^2}$')
plt.xticks([0, 0.25, 0.5, 0.75, 1])
plt.title('Зависимость квадрата скорости тележки $v^2$ \n от пройденного ей расстояния $x$. Случай тележки без груза')
plt.legend(title='Наклон скамьи')
plt.savefig('example_5_proportionality.pdf')

0.04081818181818182 6.818181818181817e-08
0.9663928082395666
0.08263636363636363 6.818181818181817e-08
0.9943104430301379
0.12345454545454544 6.818181818181817e-08
0.9926724230283033

# 6. Случай линейной зависимости (раздел 3.7.2)

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')

a = 0.232018 +- 0.105140
b = 0.028635 +- 0.001943
r = 0.988887

# 7. График с нелинейной регрессией (раздел 3.8)

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')

   w0 =  3.490206 +- 0.002606
delta =  0.072659 +- 0.003098
   f0 =  0.445524 +- 0.012464

# 8. График в полярных координатах (раздел 3.9)

alpha = np.arange(0, 180, 15) * 3.14 / 180  # рад
J =  [34.932, 35.605, 38.659, 43.843, 52.709, 62.682, 67.729, 66.948, 55.613, 45.787, 38.473, 36.451] # кг*м^2
Jt = [46.578, 47.585, 50.706, 56.572, 63.896, 72.126, 76.043, 72.126, 63.896, 56.172, 50.706, 47.585] # кг*м^2

fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})

ax.plot(alpha, J, 'o', color='C0', label='эксперимент')
ax.plot(alpha + 3.14, J, 'o', color='C0')
ax.plot(alpha, Jt, 'o', color='C1', label='расчёт')
ax.plot(alpha + 3.14, Jt, 'o', color='C1')

ax.set_rlim(0, 80)
ax.set_rticks([20, 40, 60])  # Меньше делений
ax.set_rlabel_position(-90)  # Угол для подписей к делениям
ax.grid(True)
ax.legend(loc=(0.8,0.94))

ax.set_title('Сечение эллипсоида инерции для тела, \nсостоящего из стержней и шаров')
fig.savefig('example_8_polar.pdf', bbox_inches='tight')