hamiltonian1d.py

"""Plots and animations of single particle quantum mechanics phenomena in 1D
by discretizing the Hamiltonian. This method is described here:
https://wiki.physics.udel.edu/phys824/Discretization_of_1D_continuous_Hamiltonian
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from numpy.linalg import eigh, eig

n = 512
length = 64.0
x = np.linspace(-length/2, length/2 - length/(2*n), n)
dx = x[1] - x[0]
V = 50.0*(x/length)**2
hbar = 1.0
m = 1.0

H = np.zeros((n, n), np.complex128)
for i in range(n):
    if i-1 >= 0: H[i-1, i] = (-hbar**2/(2*m))/dx**2
    H[i, i] = V[i] - 2.0*(-hbar**2/(2*m))/dx**2
    if i+1 < n: H[i+1, i] = (-hbar**2/(2*m))/dx**2
# periodic conditions
# H[n-1, 0] = (-hbar**2/(2*m))/dx**2
# H[0, n-1] = (-hbar**2/(2*m))/dx**2

vals, vects = eigh(H)

for i in range(0, 3):
    plt.plot(x, np.real(vects.T[i]), label=r"$\psi_{%d}(x)$" % i)
plt.xlim(x[0], x[-1])
plt.title("Eigenstates")
plt.legend()
plt.grid()
plt.show()
plt.close()

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim(x[0], x[-1])
ax.set_ylim(-1.1, 1.1)
psi0 = np.exp(-((x+length/4.0)/length)**2/0.05**2) # *np.exp(30.0j*np.pi*x/length)
coeffs = np.dot(psi0, vects)
line1, = ax.plot(x, np.real(psi0), label="$Re(\psi(x))$")
line2, = ax.plot(x, np.imag(psi0), label="$Im(\psi(x))$")
line3, = ax.plot(x, np.abs(psi0), color="black", label="$|\psi(x)|$")
line4, = ax.plot(x, 10*V/50.0, color="gray", label="V(x)")
line1.t = 0.0

def animation_func(*args):
    line1.t += 0.01
    psi = np.dot(coeffs*np.exp(-1.0j*vals*line1.t/hbar), vects.T)
    line1.set_ydata(np.real(psi))
    line2.set_ydata(np.imag(psi))
    line3.set_ydata(np.abs(psi))
    return line1, line2, line3

main_animation = animation.FuncAnimation(fig, animation_func, blit=True, interval=1.0)
plt.legend()
plt.grid()
plt.show()