hamiltonian2d.ipynb

hamiltonian2d.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from numpy.linalg import eigh, eig

n = 50
length = 50.0
hbar = 1.0
m = 1.0
psi = np.zeros([n, n], np.complex)
x = np.array([np.linspace(-length/2.0, length/2.0-length/n, n)
              for i in range(n)])
y = x.T
dx = length/n
V = 50.0*((x/length)**2 + (y/length)**2)
V = np.zeros([n, n])
V[n//3:n//3+2, :] = 1.0
V[n//3:n//3+2, 3*n//8 : 3*n//8+4] = 0.0
V[n//3:n//3+2, 5*n//8-4 : 5*n//8] = 0.0
H = np.zeros([n, n, n, n])
 

for i in range(n):
    for j in range(n):
        if i+1 < n: H[i+1, j, i, j] = -hbar**2/(2*m*dx**2)
        if i-1 >= 0: H[i-1, j, i, j] = -hbar**2/(2*m*dx**2)
        H[i, j, i, j] = 4.0*hbar**2/(2*m*dx**2) + V[i, j]
        if j+1 < n: H[i, j+1, i, j] = -hbar**2/(2*m*dx**2)
        if j-1 >= 0: H[i, j-1, i, j] = -hbar**2/(2*m*dx**2)


H = H.reshape(n*n, n*n)
eigvals, eigvects = eigh(H)


psi0 = np.exp((-(x+1)**2 - (y+20)**2)/20)*np.exp(-20.0j*np.pi*y/length)
coeffs = np.dot(psi0.reshape(n*n), eigvects)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
image1 = ax.imshow(V, cmap='Greys_r')
image2 = ax.imshow(2*np.real(np.dot(coeffs, eigvects.T).reshape(n, n)), 
                   cmap='Greys_r', alpha=0.7)
ax.set_ylabel('y')
ax.set_xlabel('x')
ax.set_title('Double Slit')
image2.coeffs = coeffs

def animation_func(*args):
    image2.coeffs *= np.exp(1.0j*eigvals*0.2/hbar)
    image2.set_data(2*np.abs(np.dot(image2.coeffs, eigvects.T).reshape(n, n)))
    return [image2]

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