In [1]:
%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 [2]:
import math as math
In [3]:
def well2d(x, y, nx, ny, L=1.0):
"""Compute the 2d quantum well wave function."""
wave_funct = (2/L)*np.sin((nx*math.pi*x)/L)*np.sin((ny*math.pi*y)/L)
return wave_funct
In [26]:
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 [47]:
X, Y = np.meshgrid(np.linspace(0,1,10), np.linspace(0,1,10))
psi = well2d(X, Y, 3, 2 ,1)
In [70]:
f = plt.figure(figsize=(7,5))
plt.contour(psi, cmap='hsv')
plt.xlim(0,1)
plt.ylim(0,1)
plt.title('Contour Plot of 2d Quantum Well Wave Function')
plt.xlabel('x')
plt.ylabel('y')
plt.tick_params(right=False,top=False)
plt.colorbar()
Out[70]:
In [ ]:
assert True # use this cell for grading the contour plot
Next make a visualization using one of the pcolor functions:
In [81]:
f = plt.figure(figsize=(7,5))
plt.pcolor(psi, cmap='jet')
plt.xlabel('x')
plt.ylabel('y')
plt.xlim(0,4)
plt.ylim(0,10)
plt.title('PseudoColor Plot of 2d Quantum Well Wave Function')
plt.tick_params(right=False,top=False)
plt.colorbar()
Out[81]:
In [ ]:
assert True # use this cell for grading the pcolor plot