如何使用numpy这个while+（2x for）嵌套循环进行矢量化？

Question

我正在努力提高我正在研究的算法的处理速度。在尝试使用多处理池和map在所有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。你知道吗

你能检查一下下列结果是否正确吗？你知道吗

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/magest w，这些指数一次性满足各种不等式v。这提供了几个掩码，我们将它们组合起来，以确定需要执行哪些v, w对do_something和do_something_else。最后两行中的参数是w和v的索引