import numpy as np
from scipy.integrate import odeint
from scipy.misc import derivative
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
global r, n, Psi, Fi, X, XX

LST = open("schrodinger-2b.txt", "wt")

def U(x):
return float(U0 if abs(x) < L else W)

def q(e, x):
return 2.0*(e-U(x))

def system1(cond1, X):
global eee
Y0, Y1 = cond1[0], cond1[1]
dY0dX = Y1
dY1dX = -q(eee, X)*Y0
return [dY0dX, dY1dX]

def system2 (cond2, XX):
global eee
Z0, Z1 = cond2[0], cond2[1]
dZ0dX = Z1
dZ1dX = -q(eee, XX) *Z0
return [dZ0dX, dZ1dX]

def f_fun (e):
global r, n, Psi, Fi, X, XX, eee
eee = e
cond1 = [0.0, 1.0]
sol1 = odeint(system1, cond1, X)
Psi, Psi1 = sol1[:, 0], sol1[:, 1]
cond2 = [0.0, 1.0]
sol2 = odeint(system2, cond2, XX)
Fi, Fi1 = sol2[:, 0], sol2[:, 1]
p1 = np. abs(Psi).max()
p2 = np.abs(Psi).min()
big = p1 if p1 > p2 else p2
Psi[:] = Psi[:]/big
coef = Psi[r]/Fi[rr]
Fi[:] = coef * Fi[:]
curve1 = interp1d(XX, Psi, kind='cubic')
curve2 = interp1d(XX, Fi, kind='cubic')
der1 = derivative(curve1, X[r], dx=1.e-6)
der2 = derivative(curve2, XX[rr], dx=1.e-6)
f = der1 - der2
return f

def m_bis(x1, x2, tol):
global r, n
if f_fun(e=x2) * f_fun(e=x1) > 0.0:
print("ERROR: no root!!!")
print("x1=", x1)
print("x2=", x2)
print("f_fun(e=x1, r=r, n=n)=", f_fun(e=x1))
print("f_fun(e=x2, r=r, n=n)=", f_fun(e=x2))
exit()
while abs(x2-x1) > tol:
xr = (x1 + x2)/2.0
if f_fun(e=x2) * f_fun(e=xr) < 0.0:
x1 = xr
else:
x2 = xr
if f_fun (e=x1) * f_fun (e=xr) < 0.0:
x2 = xr
else:
x1 = xr
return (x1 + x2)/2.0

def plotting_f():
plt.axis([U0, e2, fmin, fmax])
ZeroE = np.zeros(ne, dtype=float)
plt.plot(ee, ZeroE, 'k-', linewidth=1.0)
plt.plot(ee, af, 'bo', markersize=2)
for i in range (nroots):
plt.plot([energy[i]], [0.0], color='red', marker='*', markersize=10)
plt.xlabel("E", fontsize=18, color="k")
plt.ylabel("f(E)", fontsize=18, color="k")
plt.grid(True)
plt.savefig('schrodinger-2b-f.pdf', dpi=300)
plt.show()

def plotting_wf(e):
global r, n, Psi, Fi, X, XX
ff = f_fun(e)
plt.axis([A, B, U0, W])
Upot = np.array([U(X[i]) for i in np.arange(n)])
plt.plot(X, Upot, 'g-', linewidth=6.0, label="U(x)")
Zero = np.zeros(n, dtype=float)
plt.plot(X, Zero, 'k-', linewidth=1.0)
plt.plot(X, Psi, 'r-', linewidth=2.0, label="Psi(x)")
plt.plot(XX, Fi, 'b-', linewidth=2.0, label="Fi(x)")
plt.xlabel("X", fontsize=18, color="k")
plt.ylabel("Psi(x), Fi(x), U(x)", fontsize=18, color="k")
plt.grid(True)
plt.legend(fontsize=16, shadow=True, fancybox=True, loc='upper right')
plt.plot([X[r]], [Psi[r]], color='red', marker='o', markersize=7)
string1 = "E = " + format(e, "10.7f")
string2 = "f(E) = " + format(ff, "10.3e")
plt.text(-1.5, 2.7, string1, fontsize=14, color='black')
plt.text(-1.5, 2.3, string2, fontsize=14, color='black')
name = f"wavefunction_{e:.5f}.png"
plt.savefig(name, dpi=300)
plt.show()

L = 2.0
A = -L
B = +L
n = 1001
print("n=", n)
print("n=", n, file=LST)
U0 = -1.0
W = 4.0
X = np.linspace(A, B, n)
XX = np.linspace(B, A, n)
r = (n-1)*3//4
rr = n-r-1
print("r", r)
print("r", r, file=LST)
print("rr=", rr)
print("rr", rr, file=LST)
print("X[r]=", X[r])
print("X[r]=", X[r], file=LST)
print("XX[rr]=", XX[rr])
print("XX[rr]=", XX[rr], file=LST)
e1 = U0+0.05
e2 = 3.0
print("e1=", e1, " e2=", e2)
print("e1=", e1, " e2=", e2, file=LST)
ne = 101
print("ne=", ne)
print("ne=", ne, file=LST)
ee = np.linspace(e1, e2, ne)
af = np.zeros(ne, dtype=float)
porog = 5.0
tol = 1.0e-7
energy = []
ngr = 0
for i in np.arange(ne):
e = ee[i]
af[i] = f_fun(e)
stroka = "i = {:3d} e = {:8.5f} f[e] = {:12.5e}"
print(stroka.format(i, e, af[i]))
print(stroka.format(i, e, af[i]), file=LST)
if i > 0:
Log1 = af[i] * af[i-1] < 0.0
Log2 = np.abs(af[i] - af[i-1]) < porog
if Log1 and Log2:
energy1 = ee[i-1]
energy2 = ee[i]
eval = m_bis(energy1, energy2, tol)
print("eval = {:12.5e}".format(eval))
dummy = plotting_wf(eval)
energy.append(eval)
ngr += 1
nroots = len(energy)
print("nroots =", nroots)
print("nroots =", nroots, file=LST)
for i in np.arange(nroots):
stroka = "i = {:1d} energy[i] = {:12.5e}"
print(stroka.format(i, energy [i]))
print(stroka.format(i, energy [i]), file=LST)
fmax = +10.0
fmin = -10.0
dummy = plotting_f()

переписать код так чтобы потенциал от -L до L был не прямой а шёл сначала от (-L, 0) до (0, -1) и от (0, -1) до (L, 0)