我怎样才能使用numpy而+(2x for)嵌套循环?

时间:2017-02-24 11:25:18

标签: python loops numpy vectorization

我正在努力提高我正在处理的算法的处理速度。在尝试使用多处理池和映射以在所有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大小总是固定的。

1 个答案:

答案 0 :(得分: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 / maximum w,它们满足所有v的一次性不等式。这给出了几个掩码,我们将这些掩码组合起来确定需要执行v, wdo_somethingdo_something_else。最后两行中的参数是w和v的索引。