Question

您好我正在使用numpy + numba运行科学计算。我已经意识到，与matlab相比，就地添加numpy数组非常慢......

这是matlab代码：

tic;
% A,B are 2-d matrices, ind may not be distinct
for ii=1:N 
    A(ind(ii),:) =  A(ind(ii),:) +  B(ii,:);
end
toc;

这是numpy代码：

s = time.time()
# A,B are numpy.ndarray, ind may not be distinct
for k in xrange(N):
     A[ind[k],:] += B[k,:];
print time.time() - s

结果显示numpy代码比matlab慢10倍......这让我很困惑。

此外，当我从for循环中拉出添加，并且只是将单个矩阵添加与numpy.add进行比较时，numpy和matlab似乎在速度上具有可比性。

我知道的一个因素是matlab使用JIT版本＆gt; = 2012a来加速循环，但我尝试使用python代码的numba，它仍然没有加速甚至一点点。我认为这与numba根本没有触及numpy.add函数有关，因此性能根本没有变化。

我猜测matlab在这种情况下做了一些病态的缓存，因此它显着地胜过numpy。

关于如何加速numpy的任何建议？

Answer 1

尝试

A[ind] += B[:N]

即。没有任何循环。

如果ind可能包含重复元素，则可以使用np.add.at：

np.add.at(A, ind, B[:N])

Answer 2

这是一个使用点阵乘法的版本。它构造了ind的1和0矩阵。

def bar(A,B,ind):
    K,M =B.shape
    N,M =A.shape
    I = np.zeros((N,K))
    I[ind,np.arange(K)] = 1
    return A+np.dot(I,B)

对于像K,M,N = 30,14,15这样大小的问题，这大约快3倍。但对于像K,M,N = 300,100,150这样较大的那些，它会慢一些。