Numba 将映射应用于 ndarray

Question

很简单的问题，我一直在寻找但没有找到明确的答案。我想有效地将映射应用于输入 ndarray，重新映射数组中的每个元素并返回修改后的数组。

使用 numpy 的简单版本类似于：

def remap(a, mapping):
    return np.vectorize(mapping.__getitem__)(a)

我很确定我不能使用 Numba vectorize 但我希望能够使用 guvectorize。但看起来我也无法将 Numba TypedDict 传递给 guvectorize 函数。欢迎提出任何建议

Answer 1

如果你知道目标数组的维数，那么你只需要一个简单的循环。据我所知，guvectorize 还不支持此 use-case。但是，您可以使用简单的重塑来解决这个问题。

import numba as nb

@nb.njit
def remap(a, mapping):
    # Linearize the array (can cause a copy of `a`)
    b = a.reshape(-1)

    for i in range(b.size):
        b[i] = mapping[b[i]]

    return b.reshape(a.shape)

请注意，如果 a 未连续存储在内存中并且 Numpy 无法在 1D 中找到统一步长，则第一次重塑可能会导致 a 的副本。因此，returned 数组可能是 a 的副本。第二次重塑保证不复制数组。如果你想像大多数 Numpy 函数一样总是 return 输入数组的副本（无突变），那么你可以使用 flatten 而不是第一个 reshape 调用。

Numba 将映射应用于 ndarray

Numba apply mapping to ndarray

numba