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.