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