2014-06-05 61 views
6

考慮使用numpy的陣列,這是非常緩慢的下面的代碼:加速python中的一個numpy循環?

# Intersection of an octree and a trajectory 
def intersection(octree, trajectory): 
    # Initialize numpy arrays 
    ox = octree.get("x") 
    oy = octree.get("y") 
    oz = octree.get("z") 
    oe = octree.get("extent")/2 
    tx = trajectory.get("x") 
    ty = trajectory.get("y") 
    tz = trajectory.get("z") 
    result = np.zeros(np.size(ox)) 
    # Loop over elements 
    for i in range(0, np.size(tx)): 
     for j in range(0, np.size(ox)): 
      if (tx[i] > ox[j]-oe[j] and 
       tx[i] < ox[j]+oe[j] and 
       ty[i] > oy[j]-oe[j] and 
       ty[i] < oy[j]+oe[j] and 
       tz[i] > oz[j]-oe[j] and 
       tz[i] < oz[j]+oe[j]): 
       result[j] += 1 
    # Finalize 
    return result 

如何重寫功能,以加快計算? (np.size(tx) == 10000np.size(ox) == 100000

+0

你是否也考慮使用OpenCL? –

+0

我不需要完整的表現,我只是想加快速度。 – Vincent

+1

從點tx,ty,tz構建一個'scipy.spatial.KDTree',然後在ox,oy,oz中的每個點的無限範圍內使用最近鄰居查找來查看是否有足夠接近的點。 –

回答

6

您分配的大小以100000要做的第一件事10000名名單將停止使用range的嵌套循環j和使用發電機版本xrange代替。這將節省您分配所有這些列表的時間和空間。

下一個是使用矢量操作:

for i in xrange(0, np.size(tx)): 
    index = (ox-oe < tx[i]) & (ox+oe > tx[i]) & (oy-oe < ty[i]) & (oy+oe > ty[i]) & (oz-oe < tz[i]) & (oz+oe > tz[i]) 
    result[index] += 1 
+0

更改完成,但它不提供巨大的加速 – Vincent

+0

給我一秒,有下一個:) –

+0

好的,矢量化表單處於(可能需要檢查條件),但是當您根據我的原始答案修改了您的問題時,答案的一部分變得無關緊要,即使它在非顛簸的數據情況下仍然是一個很好的初始步驟:) –

0

我認爲這會產生相同的結果爲雙迴路和更快:

for j in xrange(np.size(ox)): 
    result[j] += sum(abs(tx-ox[j])<oe[j] & abs(ty-oy[j])<oe[j] & abs(tz-oz[j])<oe[j]) 

爲了得到這樣的:1)重新排序環路(即交換ŧ這是有效的,因爲循環內沒有任何變化; 2)在i環外拉result[j]; 3)將所有的t>ox-oe and t<ox+oe轉換爲abs(t-ox)<oe(儘管這可能不是一個巨大的加速,它更容易閱讀)。

由於您沒有可運行代碼,並且我也不想爲此構建測試,所以我不能100%確定這是正確的。