如何将这个循环向量化？

admin

我有一个numpy数组，如下所示：

array([[ -1. , 184. ,   0.5],
       [ -1. , 174. ,   1.0],
       [ -1. , 104. ,   0.5],
       [  1. ,  44. ,   0.5],
       [  1. ,  28. ,   0.5],
       [  1. ,  70. ,   0.5],
       [ -1. ,  34. ,   0.5],
       ...,
       [  1. ,  10. ,   0.5],
       [  1. ,  12. ,   0.5],
       [  1. ,  86. ,   1.0],
       [ -1. ,  36. ,   0.5],
       [  1. ,   2. ,   0.5],
       [ -1. ,  32. ,   1.5],
       [  1. ,  10. ,   0.5]])

我有一个函数来查找满足for循环中列出的条件的数组的索引：

def loop(array):
    n_init = 100
    a = np.dot(array[:n_init, 0], array[:n_init, 1])
    b = np.sum(array[:n_init, 2])
    loc_start = n_init
    idx = []
    lst_a, lst_b = [], []
    lst_a.append(a)
    lst_b.append(b)
    for step in range(n_init + 1, array.shape[0]):
        mean_a = np.mean(lst_a)
        mean_b = np.mean(lst_b)
        _a = np.dot(array[loc_start:step, 0], array[loc_start:step, 1])
        _b = np.sum(array[loc_start:step, 2])
        if np.abs(_a) * _b >= np.abs(mean_a) * mean_b:
            loc_start = step
            lst_a.append(_a)
            lst_b.append(_b)
            idx.append(step)
    return idx

该函数首先初始化n_init行，以便计算2个标量a和b，这两个标量被附加到稍后在for循环中使用的列表lst_a和lst_b中。

在循环中，从紧接n_init行之后的条目开始，计算lst_a和lst_b的平均值，以及_a (从最后一个idx到当前步骤的列0和1的点操作)和_b (从最后一个idx到当前步骤的列2的和)的值。

当_a乘以_b的绝对值大于或等于lst_a的平均值乘以lst_b的平均值的绝对值时，找到行号，即附加到idx的索引，并且当满足条件时，将当前的_a和_b分别附加到lst_a和lst_b。当step达到数组的长度时，整个循环结束。

因为数组有数百万行，所以需要很长时间才能完成。这个过程可以矢量化吗?如何实现？

转载请注明出处:http://www.intsu.net/article/20230502/1958886.html