这个循环如何向量化？

Question

我有一个 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，它们附加到列表 lst_a 和 lst_b 中，稍后在 for 循环中使用。

在循环中，从 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 达到数组的长度时，整个循环结束。

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

Answer 1

好吧，您正在建立一个状态（lst_a 和 lst_b）并使用它来检查一个新条目是否应该属于该状态。因此，我认为它不能完全矢量化。

您可以尝试加快速度的方法是使用增量计算（实际上是矢量化的逆运算，如果您愿意的话）。

例如，mean_a 和 mean_b 当前在每一步都重新计算，即使它没有改变，当它改变时，你可以只更新它而不是完全重新计算它。

还可以更新 _a 和 _b 而不是重新计算。它将保存数组切片和一些计算（如果loc_start 不是太远，可能不值得）。 尝试并比较。

Answer 2

您无法对其进行矢量化（因为您正在迭代地填充列表并在计算中使用它们），但您可以预先进行大量矢量化计算。也许吧。

def loop(array):
    loc_start = 100
    A = np.cumsum(array[:, 0] * array[:, 1])
    B = np.cumsum(array[:, 2])
    a, b = A[loc_start], B[loc_start]
    idx = []
    lst_a, lst_b = [a], [b]
    thr = np.abs(a)/ b
    for step in range(n_init + 1, array.shape[0]):
        _a = A[step] - a
        _b = B[step] - b
        if np.abs(_a) * _b >= thr:
            loc_start = step
            a, b = a_, b_
            lst_a.append(a)
            lst_b.append(b)
            thr = np.abs(np.mean(lst_a)) / np.mean(lst_b)
            idx.append(step)
    return idx

只要您的 cumsum 数组不溢出（检查 A[-1] 和 B[-1] 的值以查看），这应该会快很多。