numpy的累积加法

时间:2015-08-19 07:27:44

标签: python numpy

如何让循环更快?

import numpy as np

# naively small input data
image = np.array( [[2,2],[2,2]] )
polarImage = np.array( [[0,0],[0,0]] )
a = np.array( [[0,0],[0,1]] )
r = np.array( [[0,0],[0,1]] )

# TODO - this loop is too slow
it = np.nditer(image, flags=['multi_index'])
while not it.finished:
  polarImage[ a[it.multi_index],r[it.multi_index] ] += it[0]
  it.iternext()

print polarImage

# this is fast but doesn't cumulate the results!
polarImage = np.array( [[0,0],[0,0]] )
polarImage[a,r]+= image

print polarImage

第一次打印返回:

[[6 0]
 [0 2]]

第二个:

[[2 0]
 [0 2]]

通过累积添加,我的意思是有时必须将图像中的两个或更多值一起添加到 polarImage

的一个单元格中

1 个答案:

答案 0 :(得分:1)

在这种情况下,使用nditer会遮挡过程,而不会提高速度。我们更习惯于看到双循环:

In [670]: polarImage=np.zeros_like(image)
In [671]: for i in range(2):
    for j in range(2):
        polarImage[a[i,j],r[i,j]] += image[i,j]

In [672]: polarImage
Out[672]: 
array([[6, 0],
       [0, 2]])
由于缓冲问题,

polarImage[a,r]+= image无法正常工作。 (0,0)索引对使用了3次。专门针对此案例的ufunc方法at。它执行无缓冲的操作;很可能使用你的第一个例子的nditer,但是在编译的代码中。

In [676]: polarImage=np.zeros_like(image)
In [677]: np.add.at(polarImage, (a,r), image)
In [678]: polarImage
Out[678]: 
array([[6, 0],
       [0, 2]])