1from typing import Literal
2
3import matplotlib.pyplot as plt
4import networkx as nx
5import numpy as np
6from qiskit_algorithms import SamplingMinimumEigensolverResult
7
8from qlauncher import hampy, models
9from qlauncher.base import Problem
10from qlauncher.hampy import Equation
11from qlauncher.hampy.utils import shift_affected_qubits
12
13
[docs]
14class TSP(Problem):
15 """
16 Traveling Salesman Problem (TSP) definition.
17 """
18
19 def __init__(self, instance: nx.Graph, instance_name: str = 'unnamed'):
20 """
21 Args:
22 instance (nx.Graph): Graph representing the TSP instance.
23 instance_name (str): Name of the instance.
24 quadratic (bool): Whether to use quadratic constraints
25 """
26 super().__init__(instance=instance, instance_name=instance_name)
27
28 @property
29 def setup(self) -> dict:
30 return {'instance_name': self.instance_name}
31
32 def _get_path(self) -> str:
33 return f'{self.name}@{self.instance_name}'
34
35 def _solution_to_node_chain(self, solution: SamplingMinimumEigensolverResult | str) -> np.ndarray:
36 """
37 Converts the solution of the TSP problem to a chain of nodes (order of visiting)
38
39 Args:
40 solution: Solution of the TSP problem. If type is str, qiskit-style ordering is assumed (MSB first)
41
42 Returns:
43 np.ndarray: Solution chain of nodes to visit
44 """
45 bitstring = solution[::-1] if isinstance(solution, str) else solution.best_measurement['bitstring'][::-1] # type: ignore
46
47 node_count = int(len(bitstring) ** 0.5)
48 chain = []
49
50 for i in range(0, len(bitstring), node_count):
51 step_string = bitstring[i : i + node_count]
52 chosen_node = np.argmax([int(x) for x in step_string])
53 chain.append(chosen_node)
54
55 return np.array(chain)
56
57 def _calculate_solution_cost(self, solution: SamplingMinimumEigensolverResult | str | list[int]) -> float:
58 cost = 0
59 chain = solution if isinstance(solution, list) else self._solution_to_node_chain(solution)
60
61 for i in range(len(chain) - 1):
62 cost += self.instance[chain[i]][chain[i + 1]]['weight']
63 cost += self.instance[chain[-1]][chain[0]]['weight']
64 return cost
65
66 def _visualize_solution(self, solution: SamplingMinimumEigensolverResult | str | list[int]) -> None:
67 chain = solution if isinstance(solution, list) else self._solution_to_node_chain(solution)
68
69 import matplotlib.pyplot as plt
70
71 pos = nx.spring_layout(self.instance, weight=None, seed=42) # set seed for same node graphs in plt
72 problem_edges = list(self.instance.edges)
73
74 marked_edges = []
75 for i in range(len(chain) - 1):
76 marked_edges.append({chain[i], chain[i + 1]})
77 marked_edges.append({chain[-1], chain[0]})
78
79 draw_colors = [] # Limegreen for marked edges, gray for unmarked
80 draw_widths = [] # Draw marked edges thicker
81
82 path_cost = 0
83 for edge in problem_edges:
84 draw_colors.append('limegreen' if set(edge) in marked_edges else 'gray')
85 draw_widths.append(2 if set(edge) in marked_edges else 1)
86
87 path_cost += self.instance[edge[0]][edge[1]]['weight'] if set(edge) in marked_edges else 0
88
89 plt.figure(figsize=(8, 6))
90 nx.draw(
91 self.instance,
92 pos,
93 edge_color=draw_colors,
94 width=draw_widths,
95 with_labels=True,
96 node_color='skyblue',
97 node_size=500,
98 font_size=10,
99 font_weight='bold',
100 )
101
102 labels = nx.get_edge_attributes(self.instance, 'weight')
103 nx.draw_networkx_edge_labels(
104 self.instance,
105 pos,
106 edge_labels=labels,
107 rotate=False,
108 font_weight='bold',
109 label_pos=0.45,
110 )
111
112 plt.title('TSP Solution Visualization')
113 plt.suptitle(f'Path cost:{path_cost}')
114 plt.show()
115
116 def _visualize_problem(self) -> None:
117 """
118 Show plot of the TSP instance.
119 """
120
121 pos = nx.spring_layout(self.instance, weight=None, seed=42)
122 plt.figure(figsize=(8, 6))
123
124 nx.draw(
125 self.instance,
126 pos,
127 with_labels=True,
128 node_color='skyblue',
129 node_size=500,
130 edge_color='gray',
131 font_size=10,
132 font_weight='bold',
133 )
134
135 labels = nx.get_edge_attributes(self.instance, 'weight')
136 nx.draw_networkx_edge_labels(
137 self.instance,
138 pos,
139 edge_labels=labels,
140 rotate=False,
141 font_weight='bold',
142 node_size=500,
143 label_pos=0.45,
144 )
145
146 plt.title('TSP Instance Visualization')
147 plt.show()
148
[docs]
149 def visualize(self, solution: SamplingMinimumEigensolverResult | str | list[int] | None = None) -> None:
150 if solution is None:
151 self._visualize_problem()
152 else:
153 self._visualize_solution(solution)
154
[docs]
155 @staticmethod
156 def from_preset(instance_name: Literal['default'] = 'default', **kwargs) -> 'TSP':
157 """
158 Generate TSP instance from a preset name.
159
160 Args:
161 instance_name (str): Name of the preset instance
162 quadratic (bool, optional): Whether to use quadratic constraints. Defaults to False
163
164 Returns:
165 TSP: TSP instance
166 """
167 match instance_name:
168 case 'default':
169 edge_costs = np.array([[0, 1, 2, 3], [1, 0, 4, 5], [2, 4, 0, 6], [3, 5, 6, 0]])
170 case _:
171 raise ValueError('Unknown instance name')
172 G = nx.Graph()
173 n = edge_costs.shape[0]
174 for i in range(n):
175 for j in range(i + 1, n): # No connections to self
176 G.add_edge(i, j, weight=edge_costs[i, j])
177
178 return TSP(instance=G, instance_name=instance_name)
179
[docs]
180 @staticmethod
181 def generate_tsp_instance(num_vertices: int, min_distance: float = 1.0, max_distance: float = 10.0, **kwargs) -> 'TSP':
182 """
183 Generate a random TSP instance.
184
185 Args:
186 num_vertices (int): Number of vertices in the graph
187 min_distance (float, optional): Minimum distance between vertices. Defaults to 1.0
188 max_distance (float, optional): Maximum distance between vertices. Defaults to 10.0
189 quadratic (bool, optional): Whether to use quadratic constraints. Defaults to False
190
191 Returns:
192 TSP: TSP instance
193 """
194 if num_vertices < 2:
195 raise ValueError('num_vertices must be at least 2')
196
197 if min_distance <= 0:
198 raise ValueError('min_distance must be greater than 0')
199
200 g = nx.Graph()
201 rng = np.random.default_rng()
202 for i in range(num_vertices):
203 for j in range(i + 1, num_vertices):
204 g.add_edge(i, j, weight=int(rng.uniform(low=min_distance, high=max_distance)))
205
206 return TSP(instance=g, instance_name='generated')
207
208 def _make_non_collision_hamiltonian(self, node_count: int, quadratic: bool = False) -> Equation:
209 """
210 Creates a Hamiltonian representing constraints for the TSP problem. (Each node visited only once, one node per timestep)
211 Qubit mapping: [step1:[node_1, node_2,... node_n]],[step_2],...[step_n]
212
213 Args:
214 node_count: Number of nodes in the TSP problem
215 quadratic: Whether to encode as a QUBO problem
216
217 Returns:
218 np.ndarray: Hamiltonian representing the constraints
219 """
220
221 n = node_count**2
222 eq = Equation(n)
223
224 # hampy.one_in_n() takes a long time to execute with large qubit counts.
225 # We can exploit the nature of tsp constraints and create the operator for the first timestep and then shift it for the rest of the timesteps
226 # Same with nodes below.
227 # I'm pretty sure this works as intended...
228
229 # Ensure that at each timestep only one node is visited
230 t0_op = hampy.one_in_n(list(range(node_count)), eq.size, quadratic=quadratic)
231
232 for timestep in range(node_count):
233 shift = shift_affected_qubits(t0_op, timestep * node_count)
234 eq += shift
235
236 # Ensure that each node is visited only once
237 n0_op = hampy.one_in_n([timestep * node_count for timestep in range(node_count)], eq.size, quadratic=quadratic)
238
239 for node in range(node_count):
240 shift = shift_affected_qubits(n0_op, node)
241 eq += shift
242
243 return -1 * eq
244
245 def _make_connection_hamiltonian(self, edge_costs: np.ndarray, return_to_start: bool = True) -> Equation:
246 """
247 Creates a Hamiltonian that represents the costs of picking each path.
248
249 Args:
250 tsp_matrix: Edge cost matrix of the TSP problem
251
252 Returns:
253 np.ndarray: Optimal chain of nodes to visit
254 """
255 node_count = edge_costs.shape[0]
256 if len(edge_costs.shape) != 2 or edge_costs.shape[1] != node_count:
257 raise ValueError('edge_costs must be a square matrix')
258
259 n = node_count**2
260 eq = Equation(n)
261
262 for timestep in range(node_count - 1):
263 for node in range(node_count):
264 for next_node in range(node_count):
265 if node == next_node:
266 continue
267 and_hamiltonian = eq[node + timestep * node_count] & eq[next_node + (timestep + 1) * node_count]
268 eq += edge_costs[node, next_node] * and_hamiltonian
269
270 if not return_to_start:
271 return eq
272
273 # Add cost of returning to the first node
274 for node in range(node_count):
275 for node2 in range(node_count):
276 and_hamiltonian = eq[node + (node_count - 1) * node_count] & eq[node2]
277 eq += edge_costs[node, node2] * and_hamiltonian
278
279 return eq
280
[docs]
281 def to_hamiltonian(
282 self,
283 constraints_weight: int = 5,
284 costs_weight: int = 1,
285 return_to_start: bool = True,
286 onehot: Literal['exact', 'quadratic'] = 'exact',
287 ) -> models.Hamiltonian:
288 """
289 Creates a Hamiltonian for the TSP problem.
290
291 Args:
292 problem: TSP problem instance
293 quadratic: Whether to encode as a quadratic Hamiltonian
294 constraints_weight: Weight of the constraints in the Hamiltonian
295 costs_weight: Weight of the costs in the Hamiltonian
296
297 Returns:
298 np.ndarray: Hamiltonian representing the TSP problem
299 """
300 instance_graph = self.instance
301
302 edge_costs = nx.to_numpy_array(instance_graph)
303 # discourage breaking the constraints
304 edge_costs += np.eye(len(edge_costs)) * np.max(edge_costs)
305 scaled_edge_costs = edge_costs.astype(np.float32) / np.max(edge_costs)
306
307 node_count = len(instance_graph.nodes)
308
309 constraints = self._make_non_collision_hamiltonian(node_count, quadratic=(onehot == 'quadratic'))
310 costs = self._make_connection_hamiltonian(scaled_edge_costs, return_to_start=return_to_start)
311
312 return models.Hamiltonian(constraints * constraints_weight + costs * costs_weight)
313
[docs]
314 def to_qubo(self, constraints_weight: int = 5, costs_weight: int = 1, return_to_start: bool = True) -> models.QUBO:
315 return self.to_hamiltonian(constraints_weight, costs_weight, return_to_start, onehot='quadratic').to_qubo()