Python Forum

hi I'm trying to run the following code to solve Schrodinger equation in a infinite well with a barrier in themiddle but when I try to run the code I keep getting this error:

Error:
Traceback (most recent call last):
  File "C:/.../Desktop/S2.py", line 45, in <module>
    psi_b.append(Wave_function(e1))     # for each energy e1 find the the psi(x) at x = b
  File "C:/Users/Daniel/Desktop/S2.py", line 28, in Wave_function
    psi = odeint(SE, psi0, x)
  File "C:...\Python36-32\lib\site-packages\scipy\integrate\odepack.py", line 233, in odeint
    int(bool(tfirst)))
RuntimeError: The size of the array returned by func (2) does not match the size of y0 (4)
.

can someone tell me where I am getting it wronge on this script? the original code that I am trying to modify is here https://helentronica.com/2014/09/04/quan...he-python/

this is the part that I am trying to run

from pylab import *
from scipy.integrate import odeint
from scipy.optimize import brentq

def V(x):
    """
    Potential function in the finite square well. Width is L and value is global variable Vo
    """
    L = 1
    if abs(x) > L:
        return 0
    else:
        return Vo
    
def SE(psi, x):
    state0 = psi[1]
    state1 = 2.0*(V(x) - E)*psi[0]
    return array([state0, state1])

def Wave_function(energy):
    """
    Calculates wave function psi for the given value
    of energy E and returns value at point b
    """
    global psi
    global E
    E = energy
    psi = odeint(SE, psi0, x)
    return psi[-3,-2,-1,0]


N = 1000                  # number of points to take
psi = np.array([0, 0, 0, 0, 0, 0])  
E = 0.0                   # global variable Energy  needed for Sch.Eq, changed in function "Wave function"
psi0 = np.array([0,1,1,0])
b = 0                     # point outside of well where we need to check if the function diverges
a = 1
a2 = a*2
ba = a2+1
x = linspace(-b, ba, N)    # x-axis
Vo = 20

en = linspace(0, Vo, 100)   # 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
    E_zeroes = find_all_zeroes(en, psi_b)   # now find the energies where psi(b) = 0
    

According to error message, it seems like the array sizes aren't matching

# line 18:
return array([state0, state1]) # array of size 2

# line 35:
psi0 = np.array([0,1,1,0]) # array of size 4

# line 28:
psi = odeint(SE, psi0, x) 
# array returned by SE (line 18) and psi0 should probably be of same size

And looking at the original code, the author has psi0 of size 2.