矢量化版本的for-loop + numpy.where

时间:2014-12-03 10:26:50

标签: python numpy vectorization sparse-matrix

背景资料: 我在3D中有大量(N)粒子。对于具有某些属性的所有粒子 [i,j],我计算几何因子c [i,j]。然后我想总结所有对[i,j]对固定i的贡献并称之为c [i](并对所有粒子i重复此过程)。

通常,相关对的数量远小于N ^ 2,因此在位置[i,j]处具有相关信息的(N,N)维数组C和其他地方的许多零相当快numpy,但在内存使用方面也非常低效。 所以现在我只是存储相关对的C [i,j]和在一维数组中形成对的粒子。

这可能是一个例子中最好的例子: 比方说,我有两对由粒子(3,5)和(3,10)组成。原则上,我的变量看起来像这样(重复计算):

i = [3,3,5,10]  #list of particles i that form a pair
j = [5,10,3,3]  #corresponding particles j (not used in the later example) 
cij = [c35,c310,-c35,-c310] #(with actual numbers in reality)

现在它真的归结为找到一种有效的矢量化方法来重写以下for循环:

part_list=np.arange(N)
for a in part_list:
    cond = np.where(i == a)
    ci[a] = np.sum(cij[cond])

我曾想过的其他解决方案,但我们想避免:

a)并行化for循环:不可行b / c这是嵌入在我希望在某一点并行化的外部循环中。

b)在C中编写for循环并将其包装到Python中:对于这个(希望如此)相当简单的问题来说似乎有些过分。

1 个答案:

答案 0 :(得分:2)

您可以使用np.bincount获得所需内容。如果您的粒子从0开始按顺序编号,您可以简单地执行:

ci = np.bincount(i, weights= cij)

要了解这是做什么的:

>>> i = [3, 3, 5, 10]
>>> j = [5, 10, 3, 3]
>>> cij = [0.1, 0.2, -0.1, -0.2]
>>> np.bincount(i, weights= cij)
array([ 0. ,  0. ,  0. ,  0.3,  0. , -0.1,  0. ,  0. ,  0. ,  0. , -0.2])

如果您不想要所有这些额外的零,您可以执行以下操作:

>>> unq_i, inv_i = np.unique(i, return_inverse=True)
>>> unq_ci = np.bincount(inv_i, weights=cij)
>>> unq_i
array([ 3,  5, 10])
>>> unq_ci
array([ 0.3, -0.1, -0.2])

您可以稍后通过执行以下操作来分配这些唯一值:

ci[unq_i] = unq_ci