Как перебрать два массива в numpy / numba (ошибка выброса zip) - PullRequest
Новая
       13

Как перебрать два массива в numpy / numba (ошибка выброса zip)

0 голосов
/ 17 апреля 2020

У меня есть следующий код для расчета расстояния пути с учетом матрицы расстояний. 

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. Но я не могу понять, как заменить zip в моей функции calc_distance. Какой лучший способ заменить zip на numpy / numba?

1 Ответ

0 голосов
/ 17 апреля 2020

Если у вас есть 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)
...