CODEWITHSUNDEEP
pythonnumpyscipy
Nullpoet
Nullpoet

Reputation: 11269

loop variable aware numpy's slicing and vectorized calculations

How can I speedup following pure python code using numpy's slicing and vectorized (automatic looping) calculations 

def foo(x, i, j):
  return x + i + j % 255

h, w = img.shape[:2]   # img is a numpy array of shape (100,100,1)
out_img = img.copy()
for i in xrange(h):
  for j in xrange(w):
    out_img[i][j] = foo(img[i][j], i, j)

If foo is of the form 'foo(img[i][j])' (without loop variable as params), following works for me

out_img[0:,0:] = foo(img[0:,0:])

Note : Numpy gives ~70x speedup for above case compared to pure python.

I am not able to figure how to to get it working for function foo with loop variables as params.
Can someone help?

Upvotes: 2

Views: 76

Answers (1)

shx2
shx2

Reputation: 64328

You can make use of numpy.indices, like:

out_img = foo(img, *np.indices(shape=(h,w)))

Or slightly more explicit:

h_indices, w_indices = np.indices(shape=(h,w))
out_img = foo(img, h_indices, w_indices)

To illustrate what np.indices does:

np.indices((3,2))
=> 
array([[[0, 0],
        [1, 1],
        [2, 2]],

       [[0, 1],
        [0, 1],
        [0, 1]]])

Upvotes: 4

Related Questions