由于循环，使用 np.meshgrid 绘制 3d 图时出错

Question

我有一个函数，其形式为：

def f(x, y):
    total = 0
    u = np.zeros(10)
    for i in range(0,10):
        u[i] = x * i + y* i
        if u[i] < 10:
            print('do something')
    total = total + u[i]        
    return total

当我尝试使用给定的 x 和 y 值时，此函数运行良好。

f(3,4)
Out[49]: 63.0

我想使用 matplotlib 创建 3d 等高线图。尝试过


x = np.linspace(-6, 6, 30)
y = np.linspace(-6, 6, 30)

X, Y = np.meshgrid(x, y)
Z = f(X, Y)

fig = plt.figure()
ax = plt.axes(projection='3d')
ax.contour3D(X, Y, Z, 50, cmap='binary')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z');

我必须为 3d 图创建网格。由于函数中的循环，我尝试使用此方法时出现错误。我得到一个错误

ValueError: setting an array element with a sequence.

如果我的函数有循环，如何绘制 3d 图形？

Answer 1

你需要np.vectorize:

# same setup as above, then
Z = np.vectorize(f)(X, Y)

import pylab as plt
plt.imshow(Z, extent=[x[0], x[-1], y[0], y[-1]])

（我用 imshow 检查过，但 contour3D 也可以。）

np.vectorize 将接受一个函数，该函数接受标量 (non-array) 参数并神奇地循环遍历数组。它名义上等同于：

Z2 = np.array([f(xx, yy) for xx in x for yy in y]).reshape(X.shape)
print(np.abs(Z - Z2).max()) # should print 0

但更快：在我删除 f 中的 print 之后：

In [47]: %timeit Z = np.vectorize(f)(X, Y)
6 ms ± 339 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [48]: %timeit Z2 = np.array([f(xx, yy) for xx in x for yy in y]).reshape(X.shape)
13.7 ms ± 310 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

（由于打印速度非常慢，我不得不删除打印的时间。）

由于循环，使用 np.meshgrid 绘制 3d 图时出错

error while using np.meshgrid to plot a 3d-plot because of loops

python

numpy

matplotlib

mplot3d