使用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矩阵，用于绘制轮廓图。

问题是我没有任何我可以简单地调用.ravel（）的东西：我要么有xd和yvals的1d网格，2D版本X，Y和单个z值。

有人能想出一个好的，pythonic的方法吗？

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()

产量