What is the most efficient way to get length of path from adjacency matrix using numpy?

Question

The problem I am solving is optimizing a genetic algorithm for the Traveling Salesman Problem. Calculating the path takes the most time. Here is the current code I am working on:

from itertools import pairwise
import numpy as np
from random import shuffle

def get_path_len(adj_mat: np.ndarray, path: np.ndarray) -> float:
  return sum(adj_mat[i, j] for i, j in pairwise(path)) + adj_mat[path[-1], path[0]]

mat = np.random.randint(1, 1000, (100, 100))
path = np.asarray(list(range(100)))
shuffle(path)

print(get_path_len(mat, path))

Answer 1

Here is an example of how to use Numba to do that efficiently:

import numba as nb

# Pre-compile the code for a int32/int64 adj_mat 2D array
# and a int64 path 1D array
@nb.njit(['(int32[:,:], int64[:])', '(int64[:,:], int64[:])'])
def get_path_len(adj_mat, path):
    s = 0
    for i in range(path.size-1):
        # Assume path contains integers less than min(adj_mat.shape)
        s += adj_mat[path[i], path[i+1]]
    return s + adj_mat[path[-1], path[0]]

Here is performance results on my i5-9600KF CPU and Numpy 2.1.3:

Naive unvectorised code:   28.5 µs
Numba code:                 0.6 µs