python - 如何在 Numpy 中使用基于列表的索引的附加赋值

Question

这个问题已经有答案了:

Handling of duplicate indices in NumPy assignments (5 个回答)

已关闭 9 年前。

我目前正在尝试对我在 Python 中使用大型 for 循环编写的一些代码进行向量化。矢量化后的代码如下:

rho[pi,pj] += (rho_coeff*dt)*i_frac*j_frac
rho[pi+1,pj] += (rho_coeff*dt)*ip1_frac*j_frac
rho[pi,pj+1] += (rho_coeff*dt)*i_frac*jp1_frac
rho[pi+1,pj+1] += (rho_coeff*dt)*ip1_frac*jp1_frac

每个 pi 、 pj 、 dt 、 i_frac 、 j_frac 、 ip1_frac 、 jp1_frac 都是一个一维且长度相同的 numpy 数组。 rho 是一个二维 numpy 数组。 pi 和 pj 组成一个坐标列表( pi 、 pj )，它们指示矩阵 rho 的哪个元素被修改。修改涉及向 ( (rho_coeff*dt)*i_frac*j_frac , pi ) 元素添加 pj 项，以及向相邻元素添加类似项:( pi +1, pj )、( pi , pj + 1) 和 ( pi +1, pj +1)。列表中的每个坐标( pi 、 pj )都有一个唯一的与之关联的 dt 、 i_frac 、 j_frac 、 ip1_frac 和 jp1_frac 。

问题是列表可以有(并且总是会有)重复的坐标。因此，每次在列表中遇到相同的坐标时，它不会连续添加 rho，而是仅添加与最后一个重复坐标相对应的项。这个问题通过 Tentative Numpy Tutorial 中的带有索引数组的花哨索引中的示例进行了简要描述(请参阅 bool 索引之前的最后三个示例)。不幸的是他们没有提供解决方案。

有没有一种方法可以在不诉诸 for 循环的情况下执行此操作？我正在尝试优化性能，并希望在可能的情况下消除循环。

仅供引用:此代码构成 2D 粒子跟踪算法的一部分，其中每个粒子的电荷根据体积分数添加到围绕粒子位置的网格的四个相邻节点。

Answer 1

在更新数组之前，您必须找出重复的项目并将它们添加到一起。以下代码显示了首次更新时执行此操作的方法:

rows, cols = 100, 100
items = 1000

rho = np.zeros((rows, cols))
rho_coeff, dt, i_frac, j_frac = np.random.rand(4, items)
pi = np.random.randint(1, rows-1, size=(items,))
pj = np.random.randint(1, cols-1, size=(items,))

# The following code assumes pi and pj have the same dtype
pij = np.column_stack((pi, pj)).view((np.void,
                                      2*pi.dtype.itemsize)).ravel()

unique_coords, indices = np.unique(pij, return_inverse=True)
unique_coords = unique_coords.view(pi.dtype).reshape(-1, 2)
data = rho_coeff*dt*i_frac*j_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

我认为您可以将上面所有唯一的坐标查找重用于其他更新，因此以下方法可行:

ip1_frac, jp1_frac = np.random.rand(2, items)

unique_coords[:, 0] += 1
data =  rho_coeff*dt*ip1_frac*j_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

unique_coords[:, 1] += 1
data =  rho_coeff*dt*ip1_frac*jp1_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data

unique_coords[:, 0] -= 1
data =  rho_coeff*dt*i_frac*jp1_frac
binned_data = np.bincount(indices, weights=data)
rho[tuple(unique_coords.T)] += binned_data