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 [R G B] 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