schrodinger_infinite_well_integrate.py

def integrate(E, x_low, x_high, delta, V, plot=True):
    def f(x):
      return 2.0 * (V(x) - E)

    delta_sq = delta ** 2
    N = int((x_high - x_low) / delta)

    psi_left = 0.0
    # psi_right = 0.0
    psi_first = 1.0

    psi = [None for _ in range(N+1)]
    psi[0] = psi_left
    psi[1] = psi_first

    for i in range(2, N+1):
        x_next = i * delta
        x = x_next - delta
        x_prev = x - delta

        denominator = 1.0 - delta_sq * f(x_next)
        c1 = (2.0 - delta_sq * f(x)) / denominator
        c2 = (1.0 - delta_sq * f(x_prev)) / denominator

        psi[i] = c1 * psi[i-1] - c2 * psi[i - 2]