循环变量感知numpy的切片和矢量化计算

Question

如何使用numpy的切片和矢量化（自动循环）计算来加快纯python代码的速度

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)

如果foo的形式是“foo（img[i][j]）”（没有循环变量作为参数），那么下面的方法对我有效

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

注意：与纯python相比，Numpy为上述情况提供了约70倍的加速。你知道吗

我不知道如何让它在以循环变量为参数的函数foo中工作。
有人能帮忙吗？你知道吗

Answer 1

您可以使用^{}，例如：

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

或者更明确一点：

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

为了说明np.indices的作用：

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

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