import itertools

import numpy as np
import pyomo.core as pyo

def tsp(distance_matrix: np.ndarray) -> pyo.ConcreteModel:
model = pyo.ConcreteModel()

assert (
distance_matrix.shape == distance_matrix.shape
), "distance_matrix is not square"

num_points = distance_matrix.shape

points = range(num_points)
idxs = list(itertools.product(points, repeat=2))
model.x = pyo.Var(
idxs, domain=pyo.Binary
)  # x[i, j] = 1 indicates that point i is visited at step j

@model.Constraint(points)
def each_step_visits_one_point_rule(model, ii):
return sum(model.x[ii, jj] for jj in range(num_points)) == 1

@model.Constraint(points)
def each_point_visited_once_rule(model, jj):
return sum(model.x[ii, jj] for ii in range(num_points)) == 1

def is_travel_between_2_points(point1: int, point2: int):
return sum(model.x[point1, kk] * model.x[point2, kk + 1] for kk in points[:-1])

model.cost = pyo.Objective(
expr=sum(
distance_matrix[point1, point2] * is_travel_between_2_points(point1, point2)
for point1, point2 in idxs
)
)

return model


This function generates a PYOMO model which formulates the Traveling Salesman Problem. Here, given a set of points and the distances between each pair, the problem searches the shortest possible route that visits each city exactly once.

The model consists of:

• Binary variables which determines for each point at what step it was visited.
• Constraints – each city is visited once, in each step we visit only one city.
• Objective rule – shortest path.