利用matplotlib在pyquante2中绘制等高线轨道

Question

我目前正在使用matplotlib为我的PyQuante量子化学软件包编写直线和等高线绘制函数。我有一些很棒的函数，可以沿着(npts,3)点数组评估基集，例如

from somewhere import basisset, line
bfs = basisset(h2) # Generate a basis set
points = line((0,0,-5),(0,0,5)) # Create a line in 3d space
bfmesh = bfs.mesh(points)
for i in range(bfmesh.shape[1]):
    plot(bfmesh[:,i])

这很快，因为它一次计算所有的基函数，而且我从stackoverflow here和here那里得到了一些很好的帮助，让它们变得特别好。

我现在想要更新它来做等高线绘制。我在过去做的缓慢的方法是使用linspace()创建两个一维矢量，使用meshgrid()将它们网格到2D网格中，然后迭代所有xyz点并计算每个点：

f = np.empty((50,50),dtype=float)
xvals = np.linspace(0,10)
yvals = np.linspace(0,20)
z = 0
for x in xvals:
    for y in yvals:
        f = bf(x,y,z)
X,Y = np.meshgrid(xvals,yvals)
contourplot(X,Y,f)

(这不是真正的代码--可能做了一些愚蠢的事情)

我想要做的就是用我在等高线绘图示例中所做的大致相同的方法来生成网格，将其“解开”到一个(npts,3)点列表，使用我的新的快速例程来评估基函数，然后将其“重新拉开”到X，Y矩阵，以便使用contourplot绘图。

问题是，我没有任何可以简单地调用.ravel()的东西:我要么有一维的X和Y网格，要么是2D版本的X、Y和单个z值。

有没有人能想出一种很好的，巨蟒式的方法来做这件事？

Answer 1

如果您可以将f表示为X和Y的函数，则可以这样避免使用Python for-loop：

import matplotlib.pyplot as plt
import numpy as np


def bf(x, y):
    return np.sin(np.sqrt(x**2+y**2))

xvals = np.linspace(0,10)
yvals = np.linspace(0,20)
X, Y = np.meshgrid(xvals,yvals)
f = bf(X,Y)
plt.contour(X,Y,f)
plt.show()

收益率

​

​