python - itertools.combinations 的 numba 安全版本？

Question

我有一些代码循环遍历一大组 itertools.combinations ,
现在是性能瓶颈。我正在尝试求助于 numba的 @jit(nopython=True)以加快速度，但我遇到了一些问题。

首先，numba 似乎无法处理 itertools.combinations本身，根据这个小例子:

import itertools
import numpy as np
from numba import jit

arr = [1, 2, 3]
c = 2

@jit(nopython=True)
def using_it(arr, c):
    return itertools.combinations(arr, c)

for i in using_it(arr, c):
    print(i)

抛出错误: numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'combinations' of type Module(<module 'itertools' (built-in)>)
经过一番谷歌搜索，我找到了 this github issue提问者提出了这个 numba-safe 函数来计算排列:

@jit(nopython=True)
def permutations(A, k):
    r = [[i for i in range(0)]]
    for i in range(k):
        r = [[a] + b for a in A for b in r if (a in b)==False]
    return r

利用这一点，我可以轻松地筛选出组合:

@jit(nopython=True)
def combinations(A, k):
    return [item for item in permutations(A, k) if sorted(item) == item]

现在我可以运行 combinations函数没有错误并得到正确的结果。但是，现在使用 @jit(nopython=True) 速度要慢得多。比没有它。运行这个计时测试:

A = list(range(20))  # numba throws 'cannot determine numba type of range' w/o list
k = 2
start = pd.Timestamp.utcnow()
print(combinations(A, k))
print(f"took {pd.Timestamp.utcnow() - start}")

用 numba @jit(nopython=True) 计时为 2.6 秒装饰者，并在不到 1/000 秒的时间内将其注释掉。所以这对我来说也不是一个真正可行的解决方案。

Answer 1

在这种情况下，使用 Numba 并没有什么好处，因为 itertools.combinations 是用C写的。

如果你想对其进行基准测试，这里是一个 Numba/Python 实现 itertools.combinatiions做:

@jit(nopython=True)
def using_numba(pool, r):
    n = len(pool)
    indices = list(range(r))
    empty = not(n and (0 < r <= n))

    if not empty:
        result = [pool[i] for i in indices]
        yield result

    while not empty:
        i = r - 1
        while i >= 0 and indices[i] == i + n - r:
            i -= 1
        if i < 0:
            empty = True
        else:
            indices[i] += 1
            for j in range(i+1, r):
                indices[j] = indices[j-1] + 1

            result = [pool[i] for i in indices]
            yield result

在我的机器上，这比 itertools.combinations 慢了大约 15 倍。 .获取排列和过滤组合肯定会更慢。