如何在Numpy中使用基於列表的索引添加分配

Question

我目前正試圖向我介紹一些使用Python中的大循環編寫的代碼。向量化代碼如下：

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構成了指示矩陣rho的哪個元素被修改的座標列表（pi,pj）。所述修飾包括添加(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中用索引數組進行了簡要描述（請參閱布爾索引前的最後三個示例）。不幸的是，他們沒有提供解決方案。

有沒有辦法做到這一點，而不訴諸for循環？我正試圖優化性能，並希望儘可能避免使用循環。參考文獻：此代碼構成二維粒子追蹤算法的一部分，其中來自每個粒子的電荷基於體積分數被添加到圍繞粒子位置的網格的四個相鄰節點。

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