lab2 / main

import matplotlib.pyplot as plt
from numpy import exp, dot, arange, linalg
from functools import reduce
import json


def get_χ(i): return i * d / (n - 1)


def get_γ(i): return ε ** .5 * 1j if i == 0 or i == n else (f2(get_χ(i)) - ε) ** .5


def get_transition_matrix(i):
    matrix = [[0, 0], [0, 0]]
    γ_i, γ_i_next, χ_i = get_γ(i), get_γ(i + 1), get_χ(i)
    γ_tmp = γ_i_next / 2 / γ_i

    matrix[0][0] = matrix[1][1] = (.5 + γ_tmp) * exp((γ_i_next - γ_i) * χ_i)
    matrix[0][1] = (.5 - γ_tmp) * exp(-(γ_i_next + γ_i) * χ_i)
    matrix[1][0] = matrix[0][1] * exp(2 * (γ_i_next + γ_i) * χ_i)
    
    return matrix


def get_sn(): return reduce(lambda a, b: dot(a, b), [get_transition_matrix(i) for i in range(n)])


def get_a_b():
    a = [1]
    b = [b11]
    input_data = [1, b11]
    for i in range(n - 1):
        res = linalg.inv(get_transition_matrix(i)) @ input_data
        input_data = res
        a += [res[0]]
        b += [res[1]]

    return a, b


def get_ψ1(x): return exp(ε ** .5 * x * 1j) + b11 * exp(- ε ** .5 * x * 1j)


def get_ψ2(i): return a[i] * exp(get_γ(i) * get_χ(i)) + b[i] * exp(- get_γ(i) * get_χ(i))


def get_ψ3(x): return a3 * exp(ε ** .5 * (x - d) * 1j)


if __name__ == '__main__':
    # извлечение данных из файла
    with open('dataset.json', 'r', encoding='utf-8') as f:
        data = json.load(f)

    exec(data['include_lib_for_f2(t)'])  # импортировать функцию для f2(t), если необходимо

    # данные
    f2 = eval(data['f2(t)'])
    d = eval(data['d'])
    numerator = eval(data['ε/max(ε_p)'])
    ε = max(list(map(f2, arange(0, d, 0.001)))) * numerator
    n = eval(data['n'])

    sn = get_sn()

    b11 = sn[1][0] / sn[0][0]

    a, b = get_a_b()
    a3 = get_ψ2(n - 1)
    passage = abs(a3) ** 2

    print('Коэффициент прохождения электрона сквозь барьер:', passage)

    x_scale = 2

    interval1 = arange(-x_scale * d, 0, 0.01)
    interval2 = [get_χ(i) for i in range(n)]
    interval3 = arange(d, x_scale * d, 0.01)

    interval1_val = [get_ψ1(x).real for x in interval1]
    interval2_val = [get_ψ2(i).real for i in range(n)]
    interval3_val = [get_ψ3(x).real for x in interval3]

    plt.plot(interval1, interval1_val, color='red', label='I')
    plt.plot(interval2, interval2_val, color='blue', label='II')
    plt.plot(interval3, interval3_val, color='green', label='III')
    plt.plot([0, 0], [-1.5, 1.5], linestyle='--', linewidth=1.0)
    plt.plot([d, d], [-1.5, 1.5], linestyle='--', linewidth=1.0)
    plt.plot([-x_scale * d, x_scale * d], [0, 0], color='black')
    plt.legend(loc='upper right')
    plt.show()