極座標のpcolormesh

Question

私は、matplotlibでH原子の波動関数の確率密度をプロットしようとしています。私はデカルト座標でそれを行うことができましたが、極座標でpsiを指定すると、後の計算ではより良い結果が得られます。今はプロットを作ろうとしていますが、結果が間違っています（プロットの軸線はデカルト座標でなければなりません）。どのようにこれを修正するための任意のアイデア？

import numpy as np 
from matplotlib import pyplot as plt 
import matplotlib.cm as cm 
from scipy import integrate 


Z = 1 
a_0 = 1 
pi = np.pi 


n = 300 
r = np.linspace(-10, 10, n) 
theta = np.linspace(0, 2*pi, n) 
R, Theta = np.meshgrid(r, theta) 


def psi(r,theta): 
    return 1/(4*sqrt(2*pi))*(Z/a_0)**(3/2) * Z*r/a_0*np.exp(-Z*r/(2*a_0))*np.cos(theta) 


X1 = R*np.cos(Theta) 
X2 = R*np.sin(Theta) 

plt.pcolormesh(X1,X2,psi(R,Theta)**2) 
plt.axis('equal') 
plt.show() 

誤った出力：私はデカルト座標でそれを計算した場合

私が欲しいものを得る：

import numpy as np 
from matplotlib import pyplot as plt 
import matplotlib.cm as cm 
from scipy import integrate 


Z = 1 
a_0 = 1 
pi = np.pi 


n = 300 
x1 = np.linspace(-10, 10, n) 
x2 = np.linspace(-10,10, n) 
X1, X2 = np.meshgrid(x1,x2) 


def r(x,y): 
    return sqrt(x**2 + y**2) 

def psi(x,y): 
    return 1/(4*sqrt(2*pi))*(Z/a_0)**(3/2) * Z*r(x,y)/a_0*np.exp(-Z*r(x,y)/(2*a_0))*x/r(x,y) 



plt.pcolormesh(X1,X2,psi(X1,X2)**2) 

plt.axis('equal') 
plt.show() 

出力：

Answer 1

極座標の半径は負であるため、間違っています。これは極座標でのpsiの計算を失敗させます。単にあなたの問題を解決するために

r = np.linspace(0, 10, n) 

に

r = np.linspace(-10, 10, n) 

を変更。