python - 尝试生成等高线图时如何避免“TypeError: Input z must be at least a 2x2 array”

Question

我正在尝试生成 2 体系统的引力势等值线图。输入是 2 个物体的质量及其分离。我不断收到错误

TypeError：输入 z 必须至少是一个 2x2 数组。

我认为这是指 ax.contour(X,Y,PHI) 中的 PHI 术语。

我尝试更改 x、y 和 phi 以使它们匹配（我不明白为什么它们不匹配，因为 phi 是从 x 和 y 值生成的）。另外，我以前从未使用过等高线图，并且我的编程经验有限，所以请原谅我的无知。

import matplotlib.pyplot as plt
import numpy as np

n=100
#x and y evenly spaced values
x=np.arange(-n,n,0.01)
y=np.arange(-n,n,0.01)
X,Y=np.meshgrid(x,y)
r=np.array([X,Y])
#r is the scalar distance from the center that each test particle resides at
#r=np.array([(X**2+Y**2)**0.5])<--this was used to generate the output image

def Lagrange(m1,m2,a):
    mtot=m1+m2
    #the distance from the bodies to the center of mass of the system
    x1=-(m2/mtot)*a
    x2=(m1/mtot)*a
    omsq=mtot/(a**3) #omega squared term

    def phi(r):#gravitational potential function
        phi= -m1/abs(r-x1)-m2/abs(r-x2)-0.5*omsq*r**2
        return phi

#I also had a vector plot included in this code (among other details), 
#but decided to omit them as it's not relevant to the error I am receiving. 

    fig=plt.figure()
    ax=fig.add_subplot(111)
    PHI=np.meshgrid(phi(r))
    ax.contour(X,Y,PHI) #3rd dimension is the contour lines
    plt.show()

Lagrange(3.0,1.0,1.0)

我期望通过获取系统的输入 x 和 y 坐标（它们的长度匹配）并使用 phi 函数生成的输出值（它也应该匹配 x 和 y 的长度）来生成等高线图，其中等高线将代表由 phi 给出的引力势。

这是我在评论中提到的一对图像。python 代码（左）与工作 matlab 代码（右）

标签： pythonnumpymatplotlib