Fast math operations on an array in python

Question

I have a fairly simple math operation that I would like to perform on an array. Let me write an 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) 

Basically, there is a comparison between (i, j) and the three values ​​of this value (this is for visualization). I would like to collapse this and quote it somehow, but I'm not sure how and if I can. Thanks.

Answer 1

Here is the vector 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 in the last coordinates 1 and 0:

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

Answer 2

You need to change simple_funcN so that they take arrays as input and create arrays as output. After that, you can look at the function numpy.meshgrid () here or for Cartesian () here to create arrays of coordinates. After that, you can use the coordinate array to fill A with one layer.