odeint to solve Schrodinger equation

[kiyoshi7]

the condition was something I hastily added to be able to plot the potential function on multiple figures without getting doubles.

I almost got the new code to work it runs and finds energy, though it is supposed to get 7 energy levels and its not being able to do that, I am running it with finer steps so maybe it will get things right.

from pylab import *
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
import numpy as np
import itertools

a=1
B=4
L= B+a
Vmax= 50
Vpot = False

N = 1000                  # number of points to take
psi = np.zeros([N,2])     # Wave function values and its derivative (psi and psi')
psi0 = array([0,1])   # Wave function initial states
Vo = 50
E = 0.0                   # global variable Energy  needed for Sch.Eq, changed in function "Wave function"
b = L                     # point outside of well where we need to check if the function diverges
x = linspace(-B-a, L, N)    # x-axis
def V(x):
    '''
    #Potential function in the finite square well.
    '''
    if -a <=x <=a:
        val = Vo
    elif x<=-a-B:
        val = Vmax
    elif x>=L:
        val = Vmax
    else:
        val = 0
    if Vpot==True:
          if -a-B-(10/N) < x <= L+(1/N):
             Ypotential.append(val)
             Xpotential.append(x)
    return val
 
def SE(z, p):
    state0 = p[1]
    state1 = 1.0*(V(z) - E)*p[0]
    return array([state0, state1])
 
def Wave_function(energy):
    global E
    E = energy
    #        odeint(func, y0, t)
    #     solve_ivp(fun, t_span, y0)
    psi = solve_ivp(SE, [-B-a, L], psi0, max_step = ((B+a+L)/(N*2)))
    return psi.y[0,-1]
 
def find_all_zeroes(x,y):
    """
    Gives all zeroes in y = Psi(x)
    """
    all_zeroes = []
    s = np.sign(y)
    for i in range(len(y)-1):
        if s[i]+s[i+1] == 0:
            zero = brentq(Wave_function, x[i], x[i+1])
            all_zeroes.append(zero)
    return all_zeroes

        
def main():
    # main program        
 
    en = linspace(0, Vo, 100000)   # vector of energies where we look for the stable states
 
    psi_b = []      # vector of wave function at x = b for all of the energies in en
    
    for e1 in en:
        psi_b.append(Wave_function(e1))     # for each energy e1 find the the psi(x) at x = b
        
    print(len(psi_b))
    E_zeroes = find_all_zeroes(en, psi_b)   # now find the energies where psi(b) = 0 
 
    # Print energies for the bound states
    print ("Energies for the bound states are: ")
    for E in E_zeroes:
        print ("%.2f" %E)

    # Plot wave function values at b vs energy vector
    f1 = figure(1)
    plot(en/Vmax,psi_b)
    title('Values of the $\Psi(b)$ vs. Energy')
    xlabel('Energy, $E/V_0$')
    ylabel('$\Psi(x = b)$', rotation='horizontal')
    for E in E_zeroes:
        plot(E/Vmax, [0], 'go')
        annotate("E = %.2f"%E, xy = (E/Vmax, 0), xytext=(E/Vmax, 30))
    grid()
    f1.show()

    
#z = (Wave_function(2))  
    
if __name__ == "__main__":
    main()

* edit
on further inspection I found out that solve_ivp is not solving for the span -B-a to B+a, but 0 to B+a

---

I'm abandoning part of the code because I am switching solve_ivp to solve_bvp, this way I might be able to remove find_all_zeroes.

here is what I did so far (to get Wave_function to work). I'm still trying to understand solve_bvp, but I think it might be better.

from scipy.integrate import solve_bvp
from pylab import *
import numpy as np

## dimentions         _
## |       :       |   | Vmax
## |    ___:___    |   |        _
## |___|       |___|  _| Vmin   _|- Vb
##-B  -A   0   A   B*Alpha
## \______/|\______/
##     Po       Pf

Vmax = 55
Vmin = 0
Vb = 50
A = 1
B = 4
alpha = 1
Po = -A-B
Pf = A+B*alpha
psi_b = array([0,1])   # Wave function boundry [first dir, secound dir]

N = 1000 #steps
psi = np.zeros([N,2])
x = linspace(Po, Pf, N)    # x-axis


def V(z):
    '''
    #Potential function in the finite square well.
    '''
    val=[]
    for i in z:
        if -A <=i <= A:
            val.append( Vb)
        elif i<=Po:
            val.append( Vmax)
        elif i>=Pf:
            val.append(Vmax)
        else:
            val.append(Vmin)
    return val

def boundary_conditions(psi_Po, psi_pf):
    return (psi_b)

def SE(z, p, E):
    state0 = p[1]
    state1 = 1.0*(V(z) - E)*p[0]
    return np.array([state0, state1])

def Wave_function(energy,psi):
    y_ = np.zeros((2, x.size))
    psi = solve_bvp(SE(x, psi, y_), boundary_conditions, x, y_)
    print(len(psi.y))
    return psi.y

def pltPotentail():
    plot(x,V(x))    

def main():
    en = linspace(0, Vb, 10)   # vector of energies where we look for the stable states
    figure(1)
    pltPotentail()
    Wave_function(0,psi)
    show()
if __name__ == "__main__":
    main()

this is the plot I get, as you can see despite passing x to the solver it is not using the span but the number of elements and everything is flat.


edit/update. I broke the solver even more now because I'm trying to pass psi and energy to SE so I can remove that last global. progress is very slow because I still don't really understand how solve_bvp works. I'm asking more specifically how to use solve_bvp on stackoverflow (site isn't really helpful but its worth a shot). I'm starting to agree with the idea of writing a class.