pythonnumpydiscrete-mathematics

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


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))

Solution

  • 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