2

我目前正在使用 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 获得了一些很大的帮助，使它们变得非常好。

我现在想更新它以进行等高线绘图。我过去这样做的缓慢方法是使用 linspace() 创建两个一维向量，使用 meshgrid() 将它们网格化为二维网格，然后迭代所有 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() 的东西：我要么有 xvals 和 yvals 的 1d 网格、2D 版本 X、Y 和单个 z 值。

谁能想到一个不错的pythonic方式来做到这一点？

4

1 回答 1

1

如果你可以表达f为 and 的函数X，Y你可以这样避免 Python for-loops：

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

产量

在此处输入图像描述

于 2013-07-02T16:10:29.090 回答