lab2 / lab2

import matplotlib.pyplot as plt
import numpy as np
from numpy import exp, sqrt
from scipy import optimize


N = 1000
d = 2  # толщина барьера
E_max_Ep = 0.5

A = []
B = []


def energy_func(t):
    return t * 0.1 + 1.9


# вычисление eps p
max_eps_val = energy_func((optimize.fminbound(lambda t: -energy_func(t), 0, d)))
eps = max_eps_val * E_max_Ep


def calc_hi(i):
    return (i * d) / (N - 1)


def calc_gamma(i):
    if i == 0 or i == N:
        return complex(0, sqrt(eps))
    else:
        return complex(sqrt(energy_func(calc_hi(i)) - eps), 0)


def trans_matrix(i):
    matrix = [[0, 0], [0, 0]]
    temp_gamma = calc_gamma(i + 1) / (2 * calc_gamma(i))
    hi = calc_hi(i)
    matrix[0][0] = (0.5 + temp_gamma) * exp((calc_gamma(i + 1) - calc_gamma(i)) * hi)
    matrix[0][1] = (0.5 - temp_gamma) * exp(-(calc_gamma(i + 1) + calc_gamma(i)) * hi)
    matrix[1][0] = (0.5 - temp_gamma) * exp((calc_gamma(i + 1) + calc_gamma(i)) * hi)
    matrix[1][1] = (0.5 + temp_gamma) * exp((calc_gamma(i + 1) - calc_gamma(i)) * hi)

    return matrix


def calc_Sn(n):
    s = trans_matrix(0)
    for j in range(1, n):
        s = np.dot(s, trans_matrix(j))
    return s


def calc_AB(a, b):
    a += [1]
    b += [B11]
    temp_ab = [1, B11]

    for i in range(N - 1):
        # решаем уравнение в матричном виде
        temp = np.dot(np.linalg.inv(trans_matrix(i)), temp_ab)
        temp_ab = temp
        a += [temp[0]]
        b += [temp[1]]


def psi_1(hi, b1):
    return exp(complex(0, sqrt(eps) * hi)) + b1 * exp(complex(0, (-sqrt(eps) * hi)))


def psi_2(i):
    hi = calc_hi(i)
    gamma = calc_gamma(i)
    a = A[i]
    b = B[i]
    return a * exp(gamma * hi) + b * exp(- gamma * hi)


def psi_3(hi, a3):
    return a3 * exp(complex(0, sqrt(eps) * (hi - d)))

# main

# нахождение Sn и вычисление B1 из условия A1 = 1, Bn = 0
S = calc_Sn(N)
B11 = S[1][0] / S[0][0]


# нахождение Ai и Bi для каждого интервала разбиения
calc_AB(A, B)


# нахождение A3  и коэф прохождения
A3 = psi_2(N-1)
D = abs(A3) ** 2
print("Коэф. прохождения:", D)

# Данные для графиков
coef = 5
int_1 = np.arange(-coef * d, 0, 0.01 * d)
int_2 = [calc_hi(x) for x in range(N)]
int_3 = np.arange(d, d * coef, 0.01 * coef)

val_1 = [psi_1(x, B11).real for x in int_1]
val_2 = [psi_2(x).real for x in range(N)]
val_3 = [psi_3(x, A3).real for x in int_3]

plt.plot(int_1, val_1, color="red")
plt.plot(int_2, val_2, color="green")
plt.plot(int_3, val_3, color="blue")


plt.plot([0, 0], [-1.5, 1.5], linestyle='--', linewidth=0.8, color='black')
plt.plot([d, d], [-1.5, 1.5], linestyle='--', linewidth=0.8, color='black')
plt.plot([-coef * d, coef * d], [0, 0], color="black")
plt.show()