如何在Python中进行增量向量化

Question

我有一个二维数组，需要对其中的某些单元格进行数字相加操作。为了节省时间，我希望将操作向量化。但当我需要向同一单元格添加多个数字时，问题就出现了。此时，向量化代码只会加上最后一个数字。

'a'是我的数组，'x'和'y'是我想要增加的单元格的坐标，'z'包含我想要添加的数字。

import numpy as np

a=np.zeros((4,4))
x=[1,2,1]
y=[0,1,0]
z=[2,3,1]
a[x,y]+=z
print(a)

如您所见，a[1,0]应该被增加两次：一次增加2，一次增加1。因此，期望的数组应该是：

[[0. 0. 0. 0.]
 [3. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

但实际上我得到的是：

[[0. 0. 0. 0.]
 [1. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

使用 for 循环可以轻松解决这个问题，但我想知道是否可以正确地进行向量化操作。

- Antoine Belgodere

---

a[x,y] 更像是一个索引选择器，而不是迭代器，因此它不会将索引 (1,0) 相加。 - user7440787

请查看此帖子 https://dev59.com/sHDYa4cB1Zd3GeqPD8AT - astrosyam

Answer 1

使用 np.add.at 来实现此功能。

import numpy as np

a = np.zeros((4,4))
x = [1, 2, 1]
y = [0, 1, 0]
z = [2, 3, 1]
np.add.at(a, (x, y), z)
print(a)
# [[0. 0. 0. 0.]
#  [3. 0. 0. 0.]
#  [0. 3. 0. 0.]
#  [0. 0. 0. 0.]]

- jdehesa

Answer 2

当你执行 a[x,y]+=z 时，我们可以将操作分解为：

a[1, 0], a[2, 1], a[1, 0] = [a[1, 0] + 2, a[2, 1] + 3, a[1, 0] + 1]
# Equivalent to :
a[1, 0] = 2
a[2, 1] = 3
a[1, 0] = 1

这就是为什么它不起作用。 但是，如果你在每个维度上使用循环来增加数组的大小，它应该可以工作。

- Alex_6

Answer 3

你可以创建一个大小为3x4x4的多维数组，然后将z添加到所有3个不同的维度中，并将它们全部相加。

import numpy as np
x = [1,2,1]
y = [0,1,0]
z = [2,3,1]
a = np.zeros((3,4,4))
n = range(a.shape[0])
a[n,x,y] += z
print(sum(a))

这将导致

[[0. 0. 0. 0.]
 [3. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

- user7440787

Answer 4

方法一：基于计数的性能优化方法

我们可以使用np.bincount进行高效的基于计数的求和，这个方法基本上受到了这篇帖子的启发 -

def accumulate_arr(x, y, z, out):
    # Get output array shape
    shp = out.shape

    # Get linear indices to be used as IDs with bincount
    lidx = np.ravel_multi_index((x,y),shp)
    # Or lidx = coords[0]*(coords[1].max()+1) + coords[1]

    # Accumulate arr with IDs from lidx
    out += np.bincount(lidx,z,minlength=out.size).reshape(out.shape)
    return out

如果您正在使用零初始化的输出数组，请直接将输出形状馈入函数，并将bincount输出作为最终输出。

给定示例的输出 -

In [48]: accumulate_arr(x,y,z,a)
Out[48]: 
array([[0., 0., 0., 0.],
       [3., 0., 0., 0.],
       [0., 3., 0., 0.],
       [0., 0., 0., 0.]])

方法二：使用稀疏矩阵实现内存效率

In [54]: from scipy.sparse import coo_matrix

In [56]: coo_matrix((z,(x,y)), shape=(4,4)).toarray()
Out[56]: 
array([[0, 0, 0, 0],
       [3, 0, 0, 0],
       [0, 3, 0, 0],
       [0, 0, 0, 0]])

如果您可以接受稀疏矩阵，可以跳过.toarray()部分，以获得更节省内存的解决方案。

- Divakar

---

你可以使用coo_matrix代替csr_matrix来保存一次昂贵的转换，因为你传递的数据是以coo格式存储的。 - Paul Panzer

@PaulPanzer coo_matrix 应该更快吗？ - Divakar

coo -> dense 比 coo -> csr -> dense 更快，对于较大的示例，我看到了5倍的差距。 - Paul Panzer

@PaulPanzer 是的，看起来是这样！已编辑。据我记得，从 csr 访问元素更好。 - Divakar