1

我目前正在尝试对我使用 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 教程中用索引数组的花式索引下的一个示例进行了简要描述（请参阅布尔索引之前的最后三个示例）。不幸的是，他们没有为此提供解决方案。

有没有办法在不求助于for循环的情况下执行此操作？我正在尝试优化性能，并希望尽可能取消循环。

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

4

1 回答 1

1

You are going to have to figure out the repeated items and add them together before updating your array. The following code shows a way of doing that for your first update:

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

I think you can reuse all of the unique coordinate finding above for the other updates, so the following would work:

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
于 2013-07-22T17:44:48.273 回答