In [9]:
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
The wavefunction of a 2d quantum well is:
$$ \psi_{n_x,n_y}(x,y) = \frac{2}{L} \sin{\left( \frac{n_x \pi x}{L} \right)} \sin{\left( \frac{n_y \pi y}{L} \right)} $$This is a scalar field and $n_x$ and $n_y$ are quantum numbers that measure the level of excitation in the x and y directions. $L$ is the size of the well.
Define a function well2d
that computes this wavefunction for values of x
and y
that are NumPy arrays.
In [10]:
def well2d(x, y, nx, ny, L=1.0):
sine1 = np.sin(nx*np.pi*x/L)
sine2 = np.sin(ny*np.pi*y/L)
eq = 2/L * sine1 * sine2
return(eq)
In [11]:
psi = well2d(np.linspace(0,1,10), np.linspace(0,1,10), 1, 1)
assert len(psi)==10
assert psi.shape==(10,)
The contour
, contourf
, pcolor
and pcolormesh
functions of Matplotlib can be used for effective visualizations of 2d scalar fields. Use the Matplotlib documentation to learn how to use these functions along with the numpy.meshgrid
function to visualize the above wavefunction:
First make a plot using one of the contour functions:
In [12]:
#Set given variables
nx = 3
ny = 2
L = 1
x = np.linspace(0,1,100)
y = np.linspace(0,1,100)
X, Y = np.meshgrid(x, y)
#We need a 2D matrix, so set a second function
psi2 = well2d(X, Y, nx, ny, L)
#Contour graph
plt.contour(X, Y, psi2)
#Style Graph
plt.xlabel("X")
plt.ylabel("Y")
plt.title("2-D Infinite Well")
ax = plt.gca()
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
In [13]:
assert True # use this cell for grading the contour plot
Next make a visualization using one of the pcolor functions:
In [17]:
res = 70
# Because pcolor resolution depends on the size of psi and psi2 arrays
# I label a new variable "res"
# Greater "res" allows for greater resolution
psi2 = well2d(X, Y, nx, ny, L)
#Create matrix for "z" argument
z = sum(np.meshgrid(psi, psi2))
#Set colors for max/min contours (note: for this problem, this is unnecessary. It defaults the same values.)
maxd = np.amax(psi2)
mind = np.amin(psi2)
#Call the graph x, y are the lengths of the box, z is the function matrix values
plt.pcolor(X, Y, psi2, cmap='RdBu', vmin=mind, vmax=maxd)
#Style Graph
plt.xlabel("X")
plt.ylabel("Y")
plt.title("2D Infinite Well")
ax = plt.gca()
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
plt.show()
In [46]:
assert True # use this cell for grading the pcolor plot
Used matplotlib.org for pcolor example
In [ ]: