在 numpy 网格中评估每个 position 并插入一个数组

Question

My goal is to evaluate a function over a 2D plane and return an RGB value for each point on the plane so that my final output is a nested array with an [RGB] for each pixel. Here was my first attempt:

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

which raises:

ValueError: setting an array element with a sequence.

I somewhat understand this error? I found other threads on here about this error being raised, but none of them used meshgrid. When I replaced the output of foo(x,y) with a single boolean rather than an RGB array, I was generating a reasonable output, but this isn't what I'm looking for.

My current solution is:

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

This works, but it is slow, and it doesn't seem pythonic. I'm creating an array and then iterating over it, which from what I've read is a no-no. Should I do something other than meshgrid? Is my current solution the best solution? Thanks

Answer 1

Vectorize is designed for convenience not necessarily for speed. It can be quite tricky to work with so if it is not convenient to adapt a function with vectorize it is often better to rewrite the function yourself to handle the indented full scale objects. That being said there are a few things going on here. First your output is a sequence object not a scalar. Numpy is expecting to create an array of scalars the same shape as the input array. The 'signature' keyword parameter will be helpful here, forcing numpy to respect the intended output shape. Also you are trying to vecorize over two arrays at once xx and yy. This is not really possible as far as I know. Put xx and yy into a three dimensional array and vectorize over this object. For example with a 5 by 5 grid of x, y values you have the following.

Input:

# 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

Output:

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.]]])