我該如何使用numpy進行矢量化，而+（2x for）嵌套循環？

Question

我正在嘗試改進我正在處理的算法的處理速度。在嘗試使用多處理池和映射在所有CPU核心上有效地分配工作負載之前，我想要矢量化（如果可能的話）此循環。

這裏是一個例子。

v = [1,2,3,4,5,6,7,8,9,10] 
w = [-3,-2,-1,0,1,2,3] 
u = sorted(w, reverse=True) 
i = 0 
check = 0 

while v[i] != v[-1]: 
    if check == 0: 
     for k in range(len(w)): 
      if (v[i] < w[k] & v[i+1] >= w[k]) or (v[i] > w[k] & v[i+1] <= w[k]): 
       do_somthing() 
       check = 1 
       break 
    i = i+1 
    if check == 1: 
     for k in range(len(u)): 
      if (v[i] <= u[k] & v[i-1] > u[k]) or (v[i] >= u[k] & v[i-1] < u[k]): 
       do_something_else() 
       check = 0 
       break 
    i = i+1  

該示例中的數組值完全是隨機的。 V至少包含2000個元素，而w的大小始終是固定的。

Answer 1

這是一種嘗試。我觀察到，在第一和第二for塊的條件是相同的，只有一個環挑選滿足它的最低w，其它採最高w。

能否請您檢查是否follwing給出正確的結果？

import numpy as n 

assert len(v) % 2 == 0 
sg = w[-1] + 1 # create numbers that are safely out of range 
sl = w[0] - 1 # 

# in the original code, do_something is executed at most once for 
# every pair of v depending on some conditions relating to w. 
# the next lines find the smallest w that is still larger than v and 
# similar things (in other words we are binning v with respect to w), 
# in a vectorised fashion, i.e. for all v in one go. 
# using this info we can avoid the loop and for each pair of v just 
# pick the first w if any that would have left the two v through in 
# the if inside the loop 
# the difference between bin_inds_left and _right is whether the bins 
# are l <= bin < r or l < bin <= r 
bin_inds_left = np.digitize(v, w) 
bin_inds_right = np.digitize(v, w, True) 

# in your loops, various conditions are applied to v[i] and v[i+1] 
# or similarly, in vectorised code this translates to slice offsets 
# the two following lines create the logical masks corresponding to 
# the conditions in the left and right halves of the if statement 
# in the first loop (IIRC they are swapped in the second loop) 
# these masks are True if there is any permissible w and otherwise 
# False 
mask_l = bin_inds_left[1::2] > bin_inds_left[::2] 
mask_r = bin_inds_right[1::2] < bin_inds_right[::2] 

mask = mask_l | mask_r 

# for each pair find the smallest w that satisfies the left condition 
k1 = bin_inds_left[::2][mask] 
k1[~mask_l[mask]] = sg # this marks places where there is no such w 

# and in the right condition 
k2 = bin_inds_right[1::2][mask] 
k2[~mask_r[mask]] = sg 

# since it is or gated the first (smaller) of the two w's wins 
kminw = np.minimum(k1, k2) 

# same for the second loop 
k1 = bin_inds_left[1::2][mask] 
k1[~mask_l[mask]] = sl 

k2 = bin_inds_right[::2][mask] 
k2[~mask_r[mask]] = sl 

# but since u is upside-down compared to w and we do all caluclations 
# in w coordinates we must take the maximum 
# and in the very last step we translate to u coordinates 
kminu = len(w) - 1 - np.maximum(k1, k2) 

do_something(kminw, 2*np.where(mask)[0]) 
do_something_else(kminu, 1 + 2*np.where(mask)[0]) 

說明：我們使用np.digitize找到指數smalles /最大w滿足一氣呵成所有v的各種不平等。這給了一對夫婦，我們共同決定爲其v, w雙do_something和do_something_else需要執行的口罩。在最後兩行的參數是索引到W和V。