python - 从 meshgrid 数据生成数据 (Numpy)

Question

我想问一下如何从网格中生成相应的值。我有一个函数“foo”，它接受一个长度为 2 的一维数组，并返回一些实数。

import numpy as np

def foo(X):  
    #this function takes a vector, e.g., np.array([2,3]), and returns a real number.
    return sum(X)**np.sin( sum(X) );

x = np.arange(-2, 1, 1)          # points in the x axis
y = np.arange( 3, 8, 1)          # points in the y axis
X, Y = np.meshgrid(x, y)         # X, Y : grid

我使用 meshgrid 生成 X 和 Y 网格。

那么，如何使用“foo”函数生成相应的 Z 值，以便在 3D 中绘制它们，例如，使用具有 X、Y、Z 值的 plot_surface 函数绘制？

这里的问题是如何使用“foo”函数生成与 X 和 Y 具有相同形状的 Z 值。由于我的“foo”函数只采用一维数组，我不知道如何将此函数与 X 和 Y 一起使用以生成相应的 Z 值。

Answer 1

使用 np.dstack 将两个 numpy 数组“深度”堆叠起来，然后修改您的 foo 函数，使其仅在堆叠数组的最后一个轴上运行。使用 np.sum 很容易做到这一点使用参数 axis=-1，而不是使用内置的 sum:

import numpy as np

def foo(xy):
    return np.sum(xy, axis=-1) ** np.sin(np.sum(xy, axis=-1))

x = np.arange(-2, 1, 1)          # points in the x axis
y = np.arange( 3, 8, 1)          # points in the y axis
X, Y = np.meshgrid(x, y)         # X, Y : grid
XY = np.dstack((X, Y))

现在，您应该得到:

>>> XY.shape
(5, 3, 2)
>>> foo(XY)
array([[ 1.        ,  1.87813065,  1.1677002 ],
       [ 1.87813065,  1.1677002 ,  0.35023496],
       [ 1.1677002 ,  0.35023496,  0.2136686 ],
       [ 0.35023496,  0.2136686 ,  0.60613935],
       [ 0.2136686 ,  0.60613935,  3.59102217]])

---

如果你想达到同样的效果，但不修改foo，那么你可以使用np.apply_along_axis ，它应该完全满足您的需求:

>>> np.apply_along_axis(foo, -1, XY)
array([[ 1.        ,  1.87813065,  1.1677002 ],
       [ 1.87813065,  1.1677002 ,  0.35023496],
       [ 1.1677002 ,  0.35023496,  0.2136686 ],
       [ 0.35023496,  0.2136686 ,  0.60613935],
       [ 0.2136686 ,  0.60613935,  3.59102217]])