Fast math operations on an array in python

Question

I have a fairly simple math operation I'd like to perform on a array. Let me write out the example:

A = numpy.ndarray((255, 255, 3), dtype=numpy.single)
# ..
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        x = simple_func1(i)
        y = simple_func2(j)
        A[i, j] = (alpha * x * y + beta * x**2 + gamma * y**2, 1, 0)

So basically, there's a mapping between (i, j) and the 3 values of that value (this is for visualization). I'd like to roll this up and somehow vectorize this, but I'm not sure how to or if I can. Thanks.

Answer 1

You have to change simple_funcN so that they take arrays as input, and create arrays as output. After that, you could look into the numpy.meshgrid() or the cartesian() function here to build coordinate arrays. After that, you should be able to use the coordinate array(s) to fill A with a one-liner.

Answer 2

Here is the vectorized version:

i = arange(255)
j = arange(255)
x = simple_func1(i)
y = simple_func2(j)
y = y.reshape(-1,1)    

A = alpha * x * y + beta * x**2 + gamma * y**2 # broadcasting is your friend here

If you want to fill the last coordinates with 1 and 0:

B = empty(A.shape+(3,))
B[:,:,0] = A
B[:,:,1] = 1 # broadcasting again
B[:,:,2] = 0