我目前正在尝试使用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,{ {1}}),(pj
,pi
+ 1)和(pj
+ 1,pi
+ 1)。列表中的每个坐标(pj
,pi
)都有唯一的pj
,dt
,i_frac
,j_frac
和ip1_frac
与之相关。
问题是列表可以拥有(并且始终具有)重复坐标。因此,不是每次在列表中遇到相同的坐标时连续添加jp1_frac
,它只添加对应于最后一个重复坐标的术语。使用索引数组的花式索引下的Tentative Numpy Tutorial中的示例简要描述了这个问题(参见布尔索引之前的最后三个示例)。不幸的是,他们没有提供解决方案。
有没有办法在不诉诸rho
循环的情况下执行此操作?我正在尝试优化性能,并希望尽可能取消循环。
仅供参考:此代码构成了二维粒子追踪算法的一部分,其中每个粒子的电荷基于体积分数被添加到围绕粒子位置的网格的四个相邻节点。
答案 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