plot_surface.py

import numpy as np 
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


N = 256 # Controls the grid size.
L = 8
X, Y = np.meshgrid(np.linspace(-L/2, L/2, N, dtype=float),
                   np.linspace(-L/2, L/2, N, dtype=float))
DX = X[0, 1] - X[0, 0]
HBAR = 1.0
M_E = 1.0


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')


from scipy.signal import square
# V = 200.0*(1.0 - square(2.1*np.pi*np.sqrt((X/L)**2 + (Y/L)**2)))

s = 0.12
# V = 50.0*(1.0 - np.exp(-0.5*((X/L)**2 + (Y/L - 0.25)**2)/s**2)
#           - np.exp(-0.5*((X/L)**2 + (Y/L + 0.25)**2)/s**2)
#           - np.exp(-0.5*((X/L - 0.25)**2 + (Y/L)**2)/s**2)
#           - np.exp(-0.5*((X/L + 0.25)**2 + (Y/L)**2)/s**2))
V = 50.0*(np.sqrt((X/L)**2 + (Y/L)**2))

ax.plot_surface(X, Y, V)
plt.show()