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[0] == distance_matrix.shape[1]
    ), "distance_matrix is not square"

    num_points = distance_matrix.shape[0]

    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

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

    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(
            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.