lab1 / main

"""
Программа принимает на вход файл с данными dataset.json.
Переменные названы в соответствие с тем, как они выглядят.
Например, w - похожа на омегу, а v - на скорость.
s в названиях переменных означает, что это массив величин.
Например, vs - массив величин v (скорость), s - окончание, как в английском языке.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import *
from scipy.constants import c
from scipy.optimize import bisect
import json

np.seterr(divide='ignore', invalid='ignore')

with open('dataset.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

Rw = data["Большой радиус"]
Rc = data["Малый радиус"]
e = data["Диэлектрическая проницаемость"]
ws_step = data["Размер шага по фазовой частоте"]
ws_range = data["Границы фазовых частот"]
vs_step = data["Размер шага по фазовой скорости"]
vs_range = data["Границы фазовых скоростей"]
root_count = data["Количество корней дисперсионного уравнения"]


def get_d11(r, k, x1, nu):
    return jvp(nu, x1 * k * r, 1) * yvp(nu, x1 * k * Rw, 1) \
           - jvp(nu, x1 * k * Rw, 1) * yvp(nu, x1 * k * r, 1)


def get_d12(r, k, x1, nu):
    return jv(nu, x1 * k * r) * yvp(nu, x1 * k * Rw, 1) \
           - jvp(nu, x1 * k * Rw, 1) * yv(nu, x1 * k * r)


def get_d21(r, k, x1, nu):
    return jvp(nu, x1 * k * r, 1) * yv(nu, x1 * k * Rw) \
           - jv(nu, x1 * k * Rw) * yvp(nu, x1 * k * r, 1)


def get_d22(r, k, x1, nu):
    return jv(nu, x1 * k * r) * yv(nu, x1 * k * Rw) \
           - jv(nu, x1 * k * Rw) * yv(nu, x1 * k * r)


def get_m22(k, x0, x1, nu=0):
    return ivp(nu, x0 * k * Rc, 1) / x0 / iv(nu, x0 * k * Rc) \
           + get_d11(Rc, k, x1, nu) / x1 / get_d12(Rc, k, x1, nu)


def get_m11(k, x0, x1, nu=0):
    return ivp(0, x0 * k * Rc, 1) / x0 / iv(nu, x0 * k * Rc) \
           + e * get_d21(Rc, k, x1, nu) / x1 / get_d22(Rc, k, x1, nu)


def get_x1(b): return ((e * b * b - 1) ** 0.5) / b


def get_x0(b): return ((1 - b * b) ** 0.5) / b


def get_disp(v, w):
    b = v / c
    k_z = w / c
    x0 = get_x0(b)
    x1 = get_x1(b)

    return get_m11(k_z, x0, x1) * get_m22(k_z, x0, x1)


vs = np.arange(vs_range[0], vs_range[1], vs_step)
ws = np.arange(ws_range[0], ws_range[1], ws_step)

root_counter = 0
sols_and_ws = []

j = 1
for w in ws:
    row = []
    v_prev = vs[0]
    disp_prev = get_disp(v_prev, w)

    for v in vs[1:]:
        disp = get_disp(v, w)

        if root_counter > root_count:
            break

        if disp_prev * disp <= 0:
            sol = bisect(get_disp, v_prev, v, w)
            row += [(sol, w)]
        disp_prev = disp
        v_prev = v
    sols_and_ws += [row]

    print(f'{j / ws.shape[0] * 100:.1f} % calculations has been performed.')
    j += 1

h = len(sols_and_ws)
w = max([len(sols_and_ws[i]) for i in range(h)])
tmp = np.zeros((h, w, 2))

for i in range(h):
    size = len(sols_and_ws[i])
    while size < w:
        sols_and_ws[i] += [0]
        size += 1
    for j in range(w):
        tmp[i, j] = sols_and_ws[i][j]

for row in tmp:
    for cell in row:
        print(cell, end=' ')
    print()
sols_and_ws = tmp.transpose((1, 0, 2))
sols_and_ws /= np.array([c, 1]) # нормализация на скорость света

for row in sols_and_ws:
    if len(row) > 0:
        sols_, ws_ = zip(*list(filter(lambda x: x[0] > 0 and x[1] > 0, list(row))))
        vs_gr = [(ws_[i + 1] - ws_[i]) / (ws_[i + 1] / sols_[i + 1] - ws_[i] / sols_[i]) for i in range(0, len(ws_) - 1)]
        plt.plot(ws_, sols_, 'blue')
        plt.plot(ws_[1:], vs_gr, 'red')

plt.ylim(0, 1)
plt.xlabel('частота (w)')
plt.ylabel('относительная скорость (v/c)')
plt.legend(['фазовая скорость', 'групповая скорость'])
plt.savefig('grapics.png')
plt.show()