Solve particle-in-a-box problem

particleinbox.py

import numpy as np
import matplotlib.pyplot as plt

# line + scatter plot of multiple data sets a
def linescatter(a, titles=["", "", ""], labels=[], styles=[], fpath=""):
    f = plt.figure(figsize=(10, 10))
    plt.rcParams.update({'font.size': 22})
    colors = plt.cm.plasma(np.linspace(0, 1, len(a)))
    plt.rcParams['axes.prop_cycle'] = plt.cycler(color=colors)

    i = 0
    for ai in a:
        alabel = str(i+1)
        amarker = ""
        alinestyle = "-"
        
        if (len(labels) > 0):
            alabel = labels[i]
            
        if (len(styles) > 0):
            if (styles[i] == 1):
                amarker = "o"
                alinestyle = ""
                
        plt.plot(ai[:, 0], ai[:, 1], label=alabel, marker=amarker, linestyle=alinestyle)
        i += 1
        
    plt.title(titles[0])
    plt.xlabel(titles[1])
    plt.ylabel(titles[2])
    
    plt.grid(True)
    plt.legend(fontsize='xx-small')

    if (len(fpath) > 0):
        f.savefig(fpath, bbox_inches='tight') # save to file

    plt.show()
    plt.close()

# Solve m x m eigenvalue problem for particle in box [a, b] with potential v0 (1 - x^2).
# Returns eigenenergies and eigenfunctions.
def solve(a, b, v0, m):
    h = (b - a) / m # step
    d = np.zeros([m, m]) # finite difference matrix D

    diag = np.arange(1, m-1)
    d[diag, diag + 1] = -1. / h**2 / 2. # diagonal to right
    d[diag, diag - 1] = -1. / h**2 / 2. # diagonal to left
    d[diag, diag] = 1. / h**2 # diagonal

    for i in range(m): # add potential
        x = a + h * i
        d[i, i] += v0 * (1. - x**2)

    #d[0, 0] = 0. # boundaries
    #d[m - 1, m - 1] = 0.

    e, psi = np.linalg.eig(d) # solve eigenproblem
    e = e[:-2] # drop zeros
    psi = psi[:, :-2]
    
    ind = e.argsort() # sort by eigenvalue
    e = e[ind]
    psi = psi[:, ind]
    
    psi = psi / np.sqrt(h) # normalize
    
    return e, psi

# Solve particle-in-a-box problem for potential v0 with m points.
# Returns first n eigenstates.
def box(a, b, v0, m, n):
    e, psi = solve(a, b, v0, m) # solve energies and wavefunctions

    # plotting
    x = np.linspace(a, b, m)
    sol = np.zeros([n, m, 2])
    for i in range(n):
        sol[i, :, 0] = x
        sol[i, :, 1] = psi[:, i]

    print("E:", e[0:n])
    linescatter(sol, ["Wave function", "x", "$\psi$"])
    
    sol = np.zeros([n, m, 2])
    for i in range(n):
        sol[i, :, 0] = np.linspace(-1, 1, m)
        sol[i, :, 1] = [e[i]] * m
    
    linescatter(sol, ["Energy spectrum", "", "E"])

box(-1, 1, 0, 500, 5)
box(-1, 1, 1, 1000, 5)
box(-1, 1, 10, 1000, 5)
box(-1, 1, 30, 10000, 5)