Python绘制氢原子波函数轨道图与预期形状不符的问题求助

我最近在基于氢原子波函数绘制各电子轨道的电子云，计算部分用了Numpy和SciPy，绘图用了Plotly的Graph Objects。我写了一个retrieve_wavefunction函数，通过径向分量和角向分量来计算波函数，它接收球坐标网格的r、theta、phi值。下面是相关代码：

from scipy.special import sph_harm_y, eval_genlaguerre
import math
import numpy as np
import plotly.graph_objects as go

def retrieve_wavefunction(r,n,l,m,theta,phi,a):
    psi = radial_part(r,n,l,a) * angular_part(l,m,theta, phi)

    return psi


def radial_part(r,n,l,a):
    
    p = (2 * r / n * a)

    scaling_factor = 1/math.factorial(n - l + 2*l)
    normalised_part = math.sqrt((2 / n * a) ** 3 * (math.factorial(n - l - 1) / (2 * n * (math.factorial(n + 1)) **   3)))*scaling_factor
    radial_part = normalised_part * np.exp(0.5 * p) * (p) ** l * eval_genlaguerre((n + 1), (2 * l + 1), p)

    return radial_part


def angular_part(l,m,theta,phi):
    Sph1 = sph_harm_y(l, m, theta, phi)
    Sph2 = sph_harm_y(l, (-1 * m), theta, phi)
    norm_fac = np.sign(m) * np.sqrt(2) + ((np.sign(abs(m)) - 1) * 2)

    psi_sph = (Sph1 + Sph2) / norm_fac
    psi_sph = np.array(psi_sph)

    return psi_sph


def obtain_system_wavefunction_in_cart(N, r, n, l, m, a, theta, phi):
    #find the wave function of the system by inputting r,n,l,m,theta,phi
    psi = np.abs(retrieve_wavefunction(r,n,l, m, theta, phi,a))

    #Use the meshgrid in spherical coordinates to covert into cartesian coordinates
    x,y,z = sph2cart(psi,theta, phi)

    return (x,y,z)


def sph2cart(r,theta,phi):
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)

    return (x, y, z)


def simple_plot(x,y,z):
    fig = go.Figure(data=go.Scatter3d(
        x=x.flatten(), y=y.flatten(), z=z.flatten(),
        mode='markers', marker=dict(size=3, color=z.flatten(), colorscale='Viridis')
    ))

    fig.update_layout(title="3D Scatter Plot")
    fig.show()


a = 1
N = 60 #Number of grid points

r = np.linspace(0,50*a,N)
theta = np.linspace(0, np.pi, N)
phi = np.linspace(0, 2 * np.pi, N)
theta, phi,R = np.meshgrid(theta, phi,r)

#Define system parameters
n = 3 #principal quantum number
l = 1 #angular momentum quantum number
m = 0 #magnetic spin quantum number

x,y,z = obtain_system_wavefunction_in_cart(N, R, n, l, m, 1, theta, phi)

# Put the cartesian coordinates into a function that plots the atomic orbitals
# in 3 dimensions
simple_plot(x,y,z)

之后我把retrieve_wavefunction返回的psi转换成了笛卡尔坐标，再用Plotly绘图，得到的轨道形状如下：

但问题是，这个形状看起来是球形的，而实际预期的2p轨道应该是哑铃状的，我实在搞不懂哪里出了问题。下面是预期的2p轨道的二维示意图：

备注：内容来源于stack exchange，提问作者Thisanda Ayuka