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

时间:2013-07-22 16:23:42

标签: python numpy indexing

我目前正在尝试使用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

pipjdti_fracj_fracip1_fracjp1_frac中的每一个都是一个numpy数组一个维度和所有相同的长度。 rho是一个二维numpy数组。 pipj组成了一个坐标列表(pipj),表明矩阵rho的哪个元素被修改。修改涉及将(rho_coeff*dt)*i_frac*j_frac术语添加到(pipj)元素以及向相邻元素添加类似术语:(pi + 1,{ {1}}),(pjpi + 1)和(pj + 1,pi + 1)。列表中的每个坐标(pjpi)都有唯一的pjdti_fracj_fracip1_frac与之相关。

问题是列表可以拥有(并且始终具有)重复坐标。因此,不是每次在列表中遇到相同的坐标时连续添加jp1_frac,它只添加对应于最后一个重复坐标的术语。使用索引数组的花式索引下的Tentative Numpy Tutorial中的示例简要描述了这个问题(参见布尔索引之前的最后三个示例)。不幸的是,他们没有提供解决方案。

有没有办法在不诉诸rho循环的情况下执行此操作?我正在尝试优化性能,并希望尽可能取消循环。

仅供参考:此代码构成了二维粒子追踪算法的一部分,其中每个粒子的电荷基于体积分数被添加到围绕粒子位置的网格的四个相邻节点。

1 个答案:

答案 0 :(得分: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