评估 numpy 网格中的每个位置并插入一个数组

Question

我的目标是在 2D 平面上评估一个函数，并为平面上的每个点返回一个 RGB 值，这样我的最终输出是一个嵌套数组，每个像素都有一个 [R G B]。这是我的第一次尝试：

@np.vectorize
def foo(x,y):
     return [R,G,B]

x = np.linspace(-10,10)
y = np.linspace(-10,10)

xx, yy = np.meshgrid(x,y)

output_array = foo(xx,yy)

引发：

ValueError: setting an array element with a sequence.

我有点理解这个错误？我在这里发现了有关此错误的其他线程，但没有一个使用网格网格。当我用单个布尔值而不是 RGB 数组替换 foo(x,y) 的输出时，我生成了一个合理的输出，但这不是我想要的。

我目前的解决方案是：

def foo(x,y):
     return [R,G,B]

a = np.linspace(-10,10)
b = np.linspace(10,10)
arr = np.zeros((10,10,3), dtype = np.uint8)

for i,x in enumerate(a):
    for j,y in enumerate(b):
        arr[i][j] = foo(x,y)

output_array = arr

这行得通，但速度很慢，而且看起来不像 Python。我正在创建一个数组，然后对其进行迭代，从我读过的内容来看，这是一个禁忌。我应该做网格网格以外的事情吗？我目前的解决方案是最好的解决方案吗？谢谢

Answer 1

Vectorize 是为了方便而不一定是为了速度而设计的。使用它可能非常棘手，因此如果不方便使用矢量化调整函数，通常最好自己重写函数以处理缩进的全尺寸对象。话虽如此，这里发生了一些事情。首先，您的输出是序列对象而不是标量。 Numpy 期望创建一个与输入数组形状相同的标量数组。 'signature' 关键字参数在这里会很有帮助，强制 numpy 尊重预期的输出形状。此外，您正在尝试同时对两个数组进行 xx 和 yy 的 vecorize。据我所知，这实际上是不可能的。将 xx 和 yy 放入一个三维数组并在该对象上进行矢量化。例如，对于 5 x 5 的 x、y 值网格，您有以下内容。

输入：

# design foo to take x and y values in one sequence object
def foo(xy):
    x = xy[0]
    y = xy[1]

    # do stuff with x and y and get ouput RGB array
    return np.array([x+y, x*y, x*y])

# force numpy to respect input/ouput array shapes
foo = np.vectorize(foo, signature='(2)->(3)')

# x and y value grids
x = np.linspace(-10, 10, 5)
y = np.linspace(-10, 10, 5)
xx, yy  = np.meshgrid(x, y)

# 3 dimensional array with both x and y values at each grid point
xxyy = np.concatenate([np.expand_dims(xx, axis=2), np.expand_dims(yy, axis=2)], axis=2)

# call foo with only one array to vecorise over
output_array = foo(xxyy)
output_array

输出：

array([[[ -20.,  100.,  100.],
        [ -15.,   50.,   50.],
        [ -10.,   -0.,   -0.],
        [  -5.,  -50.,  -50.],
        [   0., -100., -100.]],

       [[ -15.,   50.,   50.],
        [ -10.,   25.,   25.],
        [  -5.,   -0.,   -0.],
        [   0.,  -25.,  -25.],
        [   5.,  -50.,  -50.]],

       [[ -10.,   -0.,   -0.],
        [  -5.,   -0.,   -0.],
        [   0.,    0.,    0.],
        [   5.,    0.,    0.],
        [  10.,    0.,    0.]],

       [[  -5.,  -50.,  -50.],
        [   0.,  -25.,  -25.],
        [   5.,    0.,    0.],
        [  10.,   25.,   25.],
        [  15.,   50.,   50.]],

       [[   0., -100., -100.],
        [   5.,  -50.,  -50.],
        [  10.,    0.,    0.],
        [  15.,   50.,   50.],
        [  20.,  100.,  100.]]])