numba.guvectorize 数组之间的映射值

Question

我是 numba 的新手。我有一个 n 大小的“索引”数组 (configuration)，其中包含 [0,m) 中的整数和一个大小为 m 的“映射”数组 (phase_values)。我想要的是使用 numba.guvectorize.

将 configuration 数组的每个元素映射到 phase_values 中的相应值

在 numpy 中，函数如下所示：

def configuration_to_phase_shifts_numpy(configuration, phase_values):
    map_states_to_phase_shifts = np.vectorize(lambda s: phase_values[s])
    return map_states_to_phase_shifts(configuration)

根据 numba 的文档，我理解我需要的等效矢量化代码是：

@numba.guvectorize([(numba.int64[:], numba.complex128[:], numba.complex128[:])], '(n),(m)->(n)',nopython=True)
def configuration_to_phase_shifts(configuration, phase_values, phase):
    for i in range(configuration.shape[0]):
        phase[i] = phase_values[configuration[i]]

但是，在第一次调用函数期间进行编译时：

phase = configuration_to_phase_shifts(configuration, phase_values)

numba 抛出以下错误：

Untyped global name 'configuration_to_phase_shifts': cannot determine Numba type of <class 'numpy.ufunc'>

所以，（除了如何让它工作这个明显的问题），我不明白为什么我的函数是“无类型的”——当我使用 guvectorize 装饰器时，声明了一个输出参数并且是否定义了符号输出布局？

根据@aerobiomat 的回答进行编辑

当我在 @numba.njit 修饰函数中调用向量化函数时出现上述错误，例如：

@numba.njit
def foo():
    phase_values  = np.array([1,2,3], dtype=complex)
    configuration = np.array([0,0,1,1,2,2])
    phase         = configuration_to_phase_shifts(configuration_to_phase_shifts)
    do_stuff(phase)

>>>foo()

Untyped global name 'configuration_to_phase_shifts': cannot determine Numba type of <class 'numpy.ufunc'>
File "<ipython-input>", line 5:
def foo():
    <source elided>
    configuration = np.array([0,0,1,1,2,2])
    phase         = configuration_to_phase_shifts(configuration_to_phase_shifts)
    ^

答案中的 numpy 索引方法解决了问题

Answer 1

在这个例子中使用 Numba 0.53 和 Python 3.8，您的代码似乎没有任何更改就可以正常工作：

>>> configuration = np.array([1,2,0], dtype=int)
>>> phase_values = np.array([4,6,8,1], dtype=complex)
>>> configuration_to_phase_shifts(configuration, phase_values)
array([6.+0.j, 8.+0.j, 4.+0.j])

然而，使用纯 numpy 可以获得相同的结果，不需要任何进一步的优化：

>>> phase_values[configuration]
array([6.+0.j, 8.+0.j, 4.+0.j])

这应该是评论而不是答案，但评论中的格式更难。

numba.guvectorize 数组之间的映射值

Mapping values between arrays with numba.guvectorize

python

numpy

vectorization

numba

根据@aerobiomat 的回答进行编辑