---

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import RegularGridInterpolator

hbar = 1.0545718e-34 
me = 9.1093837e-31
kB = 1.380649e-23

T = (hbar**2 / (2 * me)) * 4 / kB

mx = -4.0 * me
my = -me / 4.0

def energy(kx, ky):
    return (hbar**2 * kx**2) / (2 * mx) + (hbar**2 * ky**2) / (2 * my)

def distribution(kx, ky, T):
    return np.exp(energy(kx, ky) / (kB * T))

K = 5
dx = 1.0
dy = 1.0
kx_coarse = np.arange(-K, K + dx, dx)
ky_coarse = np.arange(-K, K + dy, dy)
KX_coarse, KY_coarse = np.meshgrid(kx_coarse, ky_coarse, indexing='ij')
n_coarse = distribution(KX_coarse, KY_coarse, T)


plt.figure(figsize=(7, 6))
im1 = plt.pcolormesh(kx_coarse, ky_coarse, n_coarse.T, shading='auto', cmap='hot')
plt.colorbar(im1, label='n(k)')
plt.title(f'Грубая сетка (шаг {dx}x{dy})')
plt.xlabel('kx')
plt.ylabel('ky')
plt.gca().set_aspect('equal')
plt.grid(True, alpha=0.3)
plt.show()

interp = RegularGridInterpolator((kx_coarse, ky_coarse), n_coarse, method='linear', bounds_error=False, fill_value=None)

dk_fine = 0.1
kx_fine = np.arange(-K, K + dk_fine, dx_fine)
ky_fine = np.arange(-K, K + dk_fine, dy_fine)
KX_fine, KY_fine = np.meshgrid(kx_fine, ky_fine, indexing='ij')

points_fine = np.array([KX_fine.ravel(), KY_fine.ravel()]).T
n_interpolated = interp(points_fine).reshape(KX_fine.shape)

plt.figure(figsize=(7, 6))
im2 = plt.pcolormesh(kx_fine, ky_fine, n_interpolated.T, shading='auto', cmap='hot')
plt.colorbar(im2, label='n(k)')
plt.title(f'Мелкая сетка (шаг {dk_fine})')
plt.xlabel('kx')
plt.ylabel('ky')
plt.gca().set_aspect('equal')
plt.grid(True, alpha=0.3)
plt.show()

n_fine_real = distribution(KX_fine, KY_fine, T)
difference = n_fine_real - n_interpolated

plt.figure(figsize=(7, 6))
im3 = plt.pcolormesh(kx_fine, ky_fine, difference.T, shading='auto', cmap='RdBu_r', vmin=-np.abs(difference).max(), vmax=np.abs(difference).max())
plt.colorbar(im3, label='Разница')
plt.title(f'Погрешность интерполяции')
plt.xlabel('kx')
plt.ylabel('ky')
plt.gca().set_aspect('equal')
plt.grid(True, alpha=0.3)
plt.show()

rmse = np.sqrt(np.mean(difference**2))
max_error = np.abs(difference).max()
mean_abs_error = np.mean(np.abs(difference))

print(f"Среднеквадратичное отклонение: {rmse:.6e}")