plt.imshow() & plt.pcolor()

import  numpy as np
import matplotlib.pyplot as plt

epsilon=10E-10;Z=1;N=3
def func_px(x,y,z):
    r=np.sqrt(x**2+y**2+z**2)
    if r<epsilon:
        return 0
    rho=2*Z*r/N
    R2p=(1/2/np.sqrt(6))*rho*Z**3/2*np.exp(-rho/2)
    Y2px=np.sqrt(3)*x/r/np.sqrt(4*np.pi)
    return R2p*Y2px

x_min=-10;x_max=10;y_min=-10;y_max=10;z=0

num_x=40;num_y=40
px=np.zeros([num_x,num_y],float)
xx=np.zeros([num_x,num_y],float)
yy=np.zeros([num_x,num_y],float)
for nx in range(num_y):
    for ny in range(num_x):
        xx[nx,ny]=nx/(num_x-1)*(x_max-x_min)+x_min
        yy[nx,ny]=ny/(num_y-1)*(y_max-y_min)+y_min
        px[nx,ny]=func_px(nx/(num_x-1)*(x_max-x_min)+x_min,ny/(num_y-1)*(y_max-y_min)+y_min,z)

#plt.imshow(px.T)
plt.pcolor(xx,yy,px)
plt.axis('square')
plt.ylim([-5,5])
plt.show()

plt.subplot(121)
plt.imshow(px.T)
plt.subplot(122)
plt.imshow(px.T,interpolation='sinc')

plt.subplot(121)
plt.pcolormesh(xx,yy,px,shading='flat')
plt.subplot(122)
plt.pcolormesh(xx,yy,px,shading='gouraud')

plt.subplot(121)
CS=plt.contour(xx,yy,px)
plt.clabel(CS,CS.levels,inline=False)
plt.subplot(122)
CS=plt.contour(xx,yy,px)
plt.clabel(CS,CS.levels,inline=True)