%pylab inline
#plt.style.use('dark_background')
# these are modules to deal with sparse matrices
import scipy.sparse as ss
import scipy.sparse.linalg as sl
# widges for sliders galore
from ipywidgets import *
L=20 # define rectangular sample
W=30
y,x=meshgrid(range(W),range(L))
x=x.flatten() # generate coordinates
y=y.flatten()
t=linspace(0,2*pi,314)
plot(x,y,'r+-')
# come up with the shape of the potato
px=8*L/20*cos(t)+L/2
py=10*W/30*sin(t)+sin(2*t)+W/2
plot(px,py)
plt.axes().set_aspect('equal')
# define Path object which can be used to check if points are inside
krumpli=mpl.path.Path(array([px,py]).T)
# find indeces that are in the potato
benne_van=krumpli.contains_points(list(map(lambda i:(i[0],i[1]),array([x,y]).T)))
# take inverse if you want
#benne_van=~benne_van # uncomment this line if you want the inverse of the shape
# look at who is inside and who is outside of the potato
plot(x,y,'r+')
plot(px,py)
plot(x[benne_van],y[benne_van],'bs',ms=3)
plt.axes().set_aspect('equal')
# define some helper matrices to be used in hamiltonian construction
idL=ss.eye(L)
idW=ss.eye(W)
odL=ss.diags(ones(L-1),1,(L,L))
odW=ss.diags(ones(W-1),1,(W,W))
# define slices of the hamiltonian
h0=-odW-odW.T
h1=-idW
# build hamiltonian corresponding to big square
H=(ss.kron(idL,h0)+ss.kron(odL,h1)+ss.kron(odL,h1).T)
# cut region of interest
Hbenne=H[:,benne_van][benne_van,:]
# look at the structure of the hamiltonian
spy(Hbenne,ms=1)
# solve sparse eigen problem, look for states near the lower edge of the spectrum i.e. around -4
va,ve=sl.eigsh(Hbenne,30,sigma=-4)
# these are the eigenvalues we found
va
# look at the first eigenvalue found
plot(x,y,'.',alpha=0.1)
plot(px,py)
tripcolor(x[benne_van],y[benne_van],ve[:,0]**2,cmap='magma_r')
plt.axes().set_aspect('equal')
# make a widget to explore wavefunctions
@interact(i=(0,len(va)-1))
def play(i=0):
tripcolor(x[benne_van],y[benne_van],ve[:,i]**2,cmap='magma_r') # density assotiated to i-th wavefunction
plot(x[~benne_van],y[~benne_van],'ws') # blanket out area not calculated
plot(px,py) # draw edge of the potato
plt.axes().set_aspect('equal')
axis('off')