numpy - itertools.combinations 的矢量化 Numpy (1d) 版本

Question

我正在尝试完成 itertools.combinations 的 Vectoriced Numpy 版本，我还可以用 @jit (Numba) 对其进行装饰 - 以使其更快。 (我不确定这是否可以做到)

我正在处理的数据集是 1d - np.array，目标是获得 3 组的组合。

x = np.array([1,2,3,4,5])

set(combinations(x, 3))

Result:

{(1, 2, 3),
 (1, 2, 4),
 (1, 2, 5),
 (1, 3, 4),
 (1, 3, 5),
 (1, 4, 5),
 (2, 3, 4),
 (2, 3, 5),
 (2, 4, 5),
 (3, 4, 5)}

我一直在搜索堆栈和其他资源，我确实找到了很多关于这个主题的信息，但没有找到任何关于我的用例的有用信息。

N-D version of itertools.combinations in numpy

这个线程表明可能很难获得比 itertools.combinations 更快的东西:

numba-safe version of itertools.combinations?

---

我被要求详细说明用例:

1. 我有一个价格数据 np.array price= ([100,101,102,103,104,105 etc..])

2. 我使用 scipy.signal.argrelmax 来查找数组中的所有峰值。 (上图中的红点)

3. 我现在需要获取峰的所有组合。

4. 然后我将对所有组合运行简单线性回归，并寻找某个 r_val 阈值，以验证趋势线。(上图中的绿点)

(3 的组合是因为那时我知道趋势线有 3 次触动。在我发布的插图上它是 4 次触动，所以我同时处理 3 次或 4 次触动。)

(此外，我使用高于/低于趋势线的积分来过滤)

我不需要组合功能。为了返回一组元组，我编写的简单线性回归算法针对 numpy 进行了优化。

Answer 1

itertools.combinations 在 CPython 中是次优，因为它必须为每个组合调用 C 函数，还因为它必须创建一个新的元组。 在 CPython 中调用 C 函数和分配对象有点昂贵。此外，该代码是通用的，但它可以专门用于像您的情况一样想要生成三元组或四元组组合的情况。计算 set 肯定是不需要的，因为如果输入 x 也由排序的唯一项组成。二进制搜索可用于查找数组是否是另一个数组的排列。实际上，在这种情况下，只需对数组进行排序并进行比较即可。 Numba 函数可以生成 Numpy 数组中的所有组合。让我们首先编写一个优化的泛型函数:

import numba as nb
import numpy as np

@nb.njit('(int64, int64)')
def combCount(n, r):
    if r < 0:
        return 0
    res = 1
    if r > n - r:
        r = n - r
    for i in range(r):
        res *= (n - i)
        res //= (i + 1)
    return res

@nb.njit(inline='always')
def genComb_generic(arr, r):
    n = arr.size
    out = np.empty((combCount(n, r), r), dtype=arr.dtype)
    idx = np.empty(r, dtype=np.int32)

    for i in range(r):
        idx[i] = i

    i = r - 1

    cur = 0
    while idx[0] < n - r + 1:
        while i > 0 and idx[i] == n - r + i:
            i -= 1
        for j in range(r):
            out[cur, j] = arr[idx[j]]
        cur += 1
        idx[i] += 1
        while i < r - 1:
            idx[i + 1] = idx[i] + 1
            i += 1
    return out

@nb.njit
def genComb_x3(arr):
    return genComb_generic(arr, 3)

@nb.njit
def genComb_x4(arr):
    return genComb_generic(arr, 4)

x = np.array([1,2,3,4,5])
genComb_x3(x) # generate triplets based on `x`

genComb_generic 速度很快，尤其是当输出很大时，但是当 r 已知时可以做得更快。这是 r=3(三元组)的特定情况:

@nb.njit
def genComb_x3(arr):
    n = arr.size
    nComb = combCount(n, 3)
    out = np.empty((nComb, 3), dtype=arr.dtype)
    a, b, c = 0, 1, 2
    arr_a = arr[a]
    arr_b = arr[b]
    for cur in range(nComb):
        out[cur, 0] = arr_a
        out[cur, 1] = arr_b
        out[cur, 2] = arr[c]
        if c < n - 1:
            c += 1
        else:
            if b < n - 2:
                b, c = b + 1, b + 2
                arr_b = arr[b]
            else:
                a, b, c = a + 1, a + 2, a + 3
                arr_a = arr[a]
                arr_b = arr[b]
    return out

当输出相对较大时，与 itertools 相比，最后一个实现非常快。当 x.size 为 30 时，它比 itertools 快 68 倍。当 x.size 为 5 时，它只快了大约 2 倍。这是因为从 CPython 调用 Numba 函数的开销很大(例如 Numpy)，更不用说分配输出的时间了。在我的机器 (i5-9600KF) 上，这个开销大约是 800 ns。您可以通过从另一个 Numba 函数调用此类函数来减少很多开销，并在可能的情况下预分配输出。最后，在这种情况下，Numba 函数应该比 itertool 实现快大约 80 倍。

如果您计划生成大量组合，那么最好即时计算它们，而不是生成一个巨大的数组(因为 RAM 往往很慢)。这应该会更快。