如何在numpy中矢量化這個循環差異？

Question

我覺得應該有一個快速的方法來加速這段代碼。我認爲答案是here，但我似乎無法以這種格式解決問題。我試圖解決的根本問題是找到平行和垂直分量方面的明顯差異，並創建這些差異的二維直方圖。

out = np.zeros((len(rpbins)-1,len(pibins)-1)) 
tmp = np.zeros((len(x),2)) 
for i in xrange(len(x)): 
    tmp[:,0] = x - x[i] 
    tmp[:,1] = y - y[i] 

    para = np.sum(tmp**2,axis=-1)**(1./2) 
    perp = np.abs(z - z[i]) 

    H, _, _ = np.histogram2d(para, perp, bins=[rpbins, pibins]) 
    out += H 

Answer 1

這樣向量化的東西是棘手的，因爲擺脫環比n元素，你必須構造的陣列，因此對於大的投入，你很可能得到一個更壞的性能比Python的循環。不過，這是可以做到：

mask = np.triu_indices(x.shape[0], 1) 
para = np.sqrt((x[:, None] - x)**2 + (y[:, None] - y)**2) 
perp = np.abs(z[:, None] - z) 
hist, _, _ = np.histogram2d(para[mask], perp[mask], bins=[rpbins, pibins]) 

的mask是爲了避免計數兩次，每次距離。我還將對角線偏移設置爲1，以避免在直方圖中包括每個點與其自身的距離0。但是如果你沒有索引para和perp，你會得到與你的代碼完全相同的結果。

有了這個樣本數據：

items = 100 
rpbins, pibins = np.linspace(0, 1, 3), np.linspace(0, 1, 3) 
x = np.random.rand(items) 
y = np.random.rand(items) 
z = np.random.rand(items) 

我得到這個我hist和你out：

>>> hist 
array([[ 1795., 651.], 
     [ 1632., 740.]]) 
>>> out 
array([[ 3690., 1302.], 
     [ 3264., 1480.]]) 

和out[i, j] = 2 * hist[i, j]除了i = j = 0，因爲0距離各項目到哪裏out[0, 0] = 2 * hist[0, 0] + items本身。

---

編輯試圖tcaswell的評論後，以下：

items = 1000 
rpbins, pibins = np.linspace(0, 1, 3), np.linspace(0, 1, 3) 
x, y, z = np.random.rand(3, items) 

def hist1(x, y, z, rpbins, pibins) : 
    mask = np.triu_indices(x.shape[0], 1) 
    para = np.sqrt((x[:, None] - x)**2 + (y[:, None] - y)**2) 
    perp = np.abs(z[:, None] - z) 
    hist, _, _ = np.histogram2d(para[mask], perp[mask], bins=[rpbins, pibins]) 
    return hist 

def hist2(x, y, z, rpbins, pibins) : 
    mask = np.triu_indices(x.shape[0], 1) 
    para = np.sqrt((x[:, None] - x)[mask]**2 + (y[:, None] - y)[mask]**2) 
    perp = np.abs((z[:, None] - z)[mask]) 
    hist, _, _ = np.histogram2d(para, perp, bins=[rpbins, pibins]) 
    return hist 

def hist3(x, y, z, rpbins, pibins) : 
    mask = np.triu_indices(x.shape[0], 1) 
    para = np.sqrt(((x[:, None] - x)**2 + (y[:, None] - y)**2)[mask]) 
    perp = np.abs((z[:, None] - z)[mask]) 
    hist, _, _ = np.histogram2d(para, perp, bins=[rpbins, pibins]) 
    return hist 

In [10]: %timeit -n1 -r10 hist1(x, y, z, rpbins, pibins) 
1 loops, best of 10: 289 ms per loop 

In [11]: %timeit -n1 -r10 hist2(x, y, z, rpbins, pibins) 
1 loops, best of 10: 294 ms per loop 

In [12]: %timeit -n1 -r10 hist3(x, y, z, rpbins, pibins) 
1 loops, best of 10: 278 ms per loop 

看來，大部分的時間都花在實例化新的陣列，而不是做實際的計算，因此，儘管有一些效率刮關閉，真的不多。