如何在 numpy/numba 中迭代两个数组（zip 抛出错误）

Question

我有以下代码用于计算给定距离矩阵的路径的距离。

dist_matrix = np.array(
    [
        [0.0, 0.5, 1.0, 1.41421356, 1.0],
        [0.5, 0.0, 0.5, 1.11803399, 1.11803399],
        [1.0, 0.5, 0.0, 1.0, 1.41421356],
        [1.41421356, 1.11803399, 1.0, 0.0, 1.0],
        [1.0, 1.11803399, 1.41421356, 1.0, 0.0],
    ]
)

@jit(nopython=True)
def calc_dist(tour):
    return np.sum(np.array([dist_matrix[i, j] for i, j in zip([tour[0:-1]], tour[1:])]))

tour = [0, 1, 2, 3, 4]
print(calc_dist(tour))

预期输出：2.118

但它抛出以下错误：

numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<intrinsic range_iter_len>) with argument(s) of type(s): (zip(iter(list(reflected list(int64))), iter(reflected list(int64))))```

我知道我可以通过设置nopython=False 来消除错误，但我的理解是它并不值得使用 numba，除非你可以将它与nopython=True 一起使用。但是我无法弄清楚如何在我的calc_distance 函数中替换zip。用 numpy/numba 替换 zip 的最佳方法是什么？

Answer 1

如果你有 numpy 数组，你可以使用 dstack():

import numpy as np

a = np.array([1,2,3,4,5])
b = np.array([6,7,8,9,10])

c = np.dstack((a,b))
#or
d = np.column_stack((a,b))

>>> c
array([[[ 1,  6],
        [ 2,  7],
        [ 3,  8],
        [ 4,  9],
        [ 5, 10]]])
>>> d
array([[ 1,  6],
       [ 2,  7],
       [ 3,  8],
       [ 4,  9],
       [ 5, 10]])

>>> c.shape
(1, 5, 2)
>>> d.shape
(5, 2)