pcolormesh в полярных координатах

Question

Я пытаюсь построить плотность вероятности волновых функций для H-атома с matplotlib. Мне удалось сделать это в декартовых координатах, но для более поздних вычислений будет лучше, если я укажу 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(-10, 10, n) 

в

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

, чтобы решить вашу проблему.