import numpy as np
from scipy.integrate import odeint
from scipy.misc import derivative
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

# Открываем файл для записи результатов
LST = open("schrodinger-2b.txt", "wt")

# Потенциал U(x)
def potential(x):
    if abs(x) < L:
        return -V0 * np.sin(np.pi * x / (0.1 * L)) * np.exp(-((x + L) / 10) ** 2)
    else:
        return W

# Функция для вычисления q(e, x)
def q(e, x):
    return 2.0 * (e - potential(x))

# Обыкновенные дифференциальные уравнения для первой системы
def system1(cond1, x):
    global energy
    psi, psi_prime = cond1[0], cond1[1]
    dpsi_dx = psi_prime
    dpsi_prime_dx = -q(energy, x) * psi
    return [dpsi_dx, dpsi_prime_dx]

# Обыкновенные дифференциальные уравнения для второй системы
def system2(cond2, xx):
    global energy
    fi, fi_prime = cond2[0], cond2[1]
    dfi_dx = fi_prime
    dfi_prime_dx = -q(energy, xx) * fi
    return [dfi_dx, dfi_prime_dx]

# Функция для поиска корня
def find_root(e):
    global r, n, psi_values, fi_values, x_values, xx_values, energy
    energy = e

    # Решение для первой системы (Psi)
    cond1 = [0.0, 1.0]
    sol1 = odeint(system1, cond1, x_values)
    psi_values, psi_prime_values = sol1[:, 0], sol1[:, 1]

    # Решение для второй системы (Fi)
    cond2 = [0.0, 1.0]
    sol2 = odeint(system2, cond2, xx_values)
    fi_values, fi_prime_values = sol2[:, 0], sol2[:, 1]

    # Нормировка Psi
    max_psi, min_psi = np.abs(psi_values).max(), np.abs(psi_values).min()
    normalization_factor = max_psi if max_psi > min_psi else min_psi
    psi_values /= normalization_factor

    # Приведение Fi к тому же масштабу, что и Psi
    coef = psi_values[r] / fi_values[rr]
    fi_values *= coef

    # Интерполяция для Psi и Fi
    psi_curve = interp1d(xx_values, psi_values, kind='cubic')
    fi_curve = interp1d(xx_values, fi_values, kind='cubic')

    # Производные для вычисления f
    psi_prime_at_r = derivative(psi_curve, x_values[r], dx=1.e-6)
    fi_prime_at_rr = derivative(fi_curve, xx_values[rr], dx=1.e-6)
    f_value = psi_prime_at_r - fi_prime_at_rr

    return f_value

# Метод бисекции для поиска корня
def bisection_method(x1, x2, tolerance):
    if find_root(x2) * find_root(x1) > 0.0:
        print("ERROR: No root found!")
        print("x1 =", x1, " x2 =", x2)
        print("find_root(x1) =", find_root(x1))
        print("find_root(x2) =", find_root(x2))
        exit()

    while abs(x2 - x1) > tolerance:
        x_mid = (x1 + x2) / 2.0
        if find_root(x2) * find_root(x_mid) < 0.0:
            x1 = x_mid
        else:
            x2 = x_mid

        if find_root(x1) * find_root(x_mid) < 0.0:
            x2 = x_mid
        else:
            x1 = x_mid

    return (x1 + x2) / 2.0

# Функция для построения графика зависимости f(E)
def plot_f():
    plt.axis([V0, e2, f_min, f_max])
    zero_line = np.zeros(ne, dtype=float)
    plt.plot(e_values, zero_line, 'k-', linewidth=1.0)
    plt.plot(e_values, f_values, color='steelblue', marker='o', markersize=2)

    # Отображаем найденные корни
    for i in range(nroots):
        plt.plot([energy_values[i]], [0.0], color='palevioletred', marker='*', markersize=10)

    plt.ylim(min(f_values) - 1, max(f_values) + 1)
    plt.xlabel("E", fontsize=13, color="k")
    plt.ylabel("f(E)", fontsize=13, color="k")
    plt.grid(True)
    plt.savefig('schrodinger-2b-f.pdf', dpi=300)
    plt.show()

# Функция для построения графика волновых функций Psi и Fi
def plot_wavefunction(e):
    global r, n, psi_values, fi_values, x_values, xx_values
    f_value = find_root(e)

    plt.axis([A, B, V0, W])
    potential_values = np.array([potential(x_values[i]) for i in range(n)])
    plt.plot(x_values, potential_values, color='darkslategray', linewidth=5.0, label="U(x)")

    zero_line = np.zeros(n, dtype=float)
    plt.plot(x_values, zero_line, 'k-', linewidth=1.0)
    plt.plot(x_values, psi_values, color='palevioletred', linewidth=2.0, label="Psi(x)")
    plt.plot(xx_values, fi_values, color='steelblue', linewidth=2.0, label="Fi(x)")

    min_val = min(min(psi_values), min(fi_values), min(potential_values))
    max_val = max(max(psi_values), max(fi_values), max(potential_values))
    plt.ylim(min_val - 1, max_val + 1)

    plt.xlabel("X", fontsize=13, color="k")
    plt.ylabel("Psi(x), Fi(x), U(x)", fontsize=13, color="k")
    plt.grid(True)
    plt.legend(fontsize=13, shadow=True, fancybox=True, loc='upper right')
    plt.plot([x_values[r]], [psi_values[r]], color='palevioletred', marker='o', markersize=7)

    # Отображение энергии и f(E)
    energy_string = f"E   = {e:.7f}"
    f_value_string = f"f(E) = {f_value:.3e}"
    plt.text(-1.5, 2.9, energy_string, fontsize=13, color='black')
    plt.text(-1.5, 2.2, f_value_string, fontsize=13, color='black')

    filename = f"wavefunction_{e:.5f}.png"
    plt.savefig(filename, dpi=300)
    plt.show()

# Функция для вычисления квантовых средних значений <x> и <x^2>
def calculate_expectation_values(x_values, psi_values):
    probability_density = psi_values**2  # Вероятностная плотность
    dx = x_values[1] - x_values[0]

    # Нормировка функции волновой функции Psi
    normalization_factor = np.sum(probability_density * dx)
    probability_density /= normalization_factor

    # Вычисление <x> и <x^2>
    x_mean = np.sum(x_values * probability_density * dx)
    x2_mean = np.sum(x_values**2 * probability_density * dx)

    return x_mean, x2_mean

# Инициализация параметров
L = 5 * 1.8897  # Размер области
A = -L
B = +L
n = 1001  # Число точек в сетке
V0 = 15 * 0.0367493
W = 10.0
x_values = np.linspace(A, B, n)
xx_values = np.linspace(B, A, n)
r = (n - 1) * 3 // 4
rr = n - r - 1

# Параметры для поиска корней
e1 = V0 + 0.05
e2 = 3.0
ne = 101
e_values = np.linspace(e1, e2, ne)
f_values = np.zeros(ne, dtype=float)

# Параметры для метода бисекции
tolerance = 1.0e-7
energy_values = []
nroots = 0

# Вычисление f(E) и поиск корней
for i in range(ne):
    e = e_values[i]
    f_values[i] = find_root(e)
    result_str = f"i = {i:3d}   e = {e:8.5f}  f[e] = {f_values[i]:12.5e}"
    print(result_str)
    print(result_str, file=LST)

    # Поиск корней методом бисекции
    if i > 0:
        if f_values[i] * f_values[i - 1] < 0.0 and abs(f_values[i] - f_values[i - 1]) < 5.0:
            root_energy = bisection_method(e_values[i - 1], e_values[i], tolerance)
            energy_values.append(root_energy)
            nroots += 1
            plot_wavefunction(root_energy)

            x_mean, x2_mean = calculate_expectation_values(X, Psi)
            print(f"Квантовомеханические средние значения для E = {eval:.7f}: <x> = {x_mean:.5f}, <x^2> = {x2_mean:.5f}")
            print(f"Квантовомеханические средние значения для E = {eval:.7f}: <x> = {x_mean:.5f}, <x^2> = {x2_mean:.5f}", file=LST)
            if len(energy_values) == 2:
              break

# Вывод найденных энергий
print("Найдено корней:", nroots)
for i in range(nroots):
    result_str = f"i = {i:1d}    energy = {energy_values[i]:12.5e}"
    print(result_str)
    print(result_str, file=LST)

# Построение графика f(E)
f_max = +10.0
f_min = -10.0
plot_f()

# Закрытие файла
LST.close()

NameError                                 Traceback (most recent call last)
<ipython-input-2-a514cd7c8db4> in <cell line: 187>()
    200             plot_wavefunction(root_energy)
    201 
--> 202             x_mean, x2_mean = calculate_expectation_values(X, Psi)
    203             print(f"Квантовомеханические средние значения для E = {eval:.7f}: <x> = {x_mean:.5f}, <x^2> = {x2_mean:.5f}")
    204             print(f"Квантовомеханические средние значения для E = {eval:.7f}: <x> = {x_mean:.5f}, <x^2> = {x2_mean:.5f}", file=LST)

NameError: name 'X' is not defined