add tsp

examples/tsp.py

@@ -0,0 +1,53 @@
+# Adapted from: 
+# https://github.com/anirudhgha/p-bit/tree/master/Research/Travelling%20Salesman%20Problem
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from p_kit.library.tsp import TSP
+from p_kit.solver.annealing import constant, execute, linear
+# from p_kit.solver.cpsl_solver import CpslSolver
+from p_kit.solver.csd_solver import CaSuDaSolver
+from p_kit.visualization.tsp_graph import visualize_tsp_route
+from p_kit.visualization.utils import tsp_hist
+
+# Define city graph (distance matrix)
+# city_graph = np.array([[0, 510, 480],
+#                       [510, 0, 240],
+                    #   [480, 240, 0]])
+
+city_graph = np.array([[0, 510, 480, 490],
+              [510, 0, 240, 370],
+              [480, 240, 0, 220],
+              [490, 370, 220, 0]])
+# city_graph = np.array([[  0.,1.,1.,1.,-10.,-10.,1.,-15.,-15.],
+#  [  1.,0.,1.,-10.,1.,-10.,-15.,1.,-15.],
+#  [  1.,1.,0.,-10.,-10.,1.,-15.,-15.,1.],
+#  [  1.,-10.,-10.,0.,1.,1.,-20.,-20.,-20.],
+#  [-10.,1.,-10.,1.,0.,1.,-20.,-20.,-20.],
+#  [-10.,-10.,1.,1.,1.,0.,-20.,-20.,-20.],
+#  [  1.,-15.,-15.,-20.,-20.,-20.,0.,1.,1.],
+#  [-15.,1.,-15.,-20.,-20.,-20.,1.,0.,1.],
+#  [-15.,-15.,1.,-20.,-20.,-20.,1.,1.,0.]])
+
+# Convert to negative weights for optimization
+city_graph *= -1
+
+# Initialize p-circuit
+circuit = TSP(city_graph=city_graph, tsp_modifier=0.33)
+
+# Simulated Annealing Schedule
+solver = CaSuDaSolver(Nt=int(1e5), dt=1 / (2 * len(city_graph)), i0=0, expected_mean=0, seed=None)
+
+# Increase n_shots to have more reliable results
+n_shots = 200
+samples = execute(solver, circuit, linear, n_shots, n_last_samples=50, n_jobs=-1)
+
+# Visualize Best path
+hist = tsp_hist(samples, city_graph)
+plt.bar(hist.keys(), hist.values())
+plt.show()
+
+best_path = max(hist, key=hist.get)
+
+visualize_tsp_route(best_path, city_graph)

p_kit/library/__init__.py

@@ -0,0 +1,3 @@
+from . import tsp
+
+__all__ = ["tsp"]

p_kit/library/tsp.py

@@ -0,0 +1,72 @@
+# Adapted from: https://github.com/anirudhgha/p-bit/blob/master/Research/Travelling%20Salesman%20Problem/tsp.py
+
+import numpy as np
+from p_kit.psl.p_circuit import PCircuit
+
+
+class TSP(PCircuit):
+
+    """ Build a J matrix for some travelling salesman problem graph. 
+    Designing a travelling salesman J:
+    Rule 1: 1 between pbits of same city
+    Rule 2: 1 between pbits of same order
+    Rule 3: negative distances as weights between rows (ex. all p-bits in city-1-row to all cities in city-3-row)
+    Rule 4: 0 connections from city_n-order_n to itself
+
+    Parameters
+    ----------
+    city_graph: np.array((n_cities, n_cities))
+        The city graph.
+    tsp_modifier: float, default: 1
+        Value between pbits of same city
+
+    Attributes
+    ----------
+    h : np.array((n_pbits, 1))
+        biases
+    J : np.array((n_pbits, n_pbits))
+        weights
+
+    Notes
+    -----
+    .. versionadded:: 0.0.1
+
+    """
+
+    def __init__(self, city_graph=None, tsp_modifier=None):
+        PCircuit.__init__(self, len(city_graph[0]) ** 2)
+        self.city_graph = np.asarray(city_graph)
+        self.tsp_modifier = 1 if tsp_modifier is None else tsp_modifier
+        self._init_city_graph()
+
+
+    def _init_city_graph(self):
+            city_graph = np.divide(self.city_graph, np.amax(np.abs(self.city_graph)))
+            n_cities = len(city_graph[0])
+
+            # Rule 3: negative distances from one city to another
+            for i in range(n_cities):
+                for j in range(n_cities):
+                    #if order(i) is one away from order(j)
+                    if i == j:
+                        continue
+                    off_diag = np.ones(n_cities-1) * city_graph[j, i]
+                    #set both off diagonals (one to the left and right of main diagonal) of weights to off_diag
+                    weights_i_j = np.diag(off_diag, 1) + np.diag(off_diag, -1)
+
+                    self.J[j * n_cities: j * n_cities + n_cities, i * n_cities: i * n_cities + n_cities] = weights_i_j[:,:]
+
+            # Rule 1 - 1 between pbits of same city
+            for i in range(n_cities):
+                self.J[i * n_cities:i * n_cities + n_cities, i * n_cities:i * n_cities + n_cities] = self.tsp_modifier # dif
+
+            # Rule 2 - 1 between pbits of same order
+            for i in range(n_cities ** 2):
+                for j in range(n_cities ** 2):
+                    if i == j % n_cities or j == i % n_cities:
+                        self.J[i, j] = self.tsp_modifier
+
+            # Rule 4: 0s on the diagonal
+            np.fill_diagonal(self.J, 0)
+
+

p_kit/visualization/__init__.py

@@ -1,7 +1,17 @@
-from .utils import m_to_string
+from .utils import m_to_string, tsp_hist
 from .histplot import histplot, energyplot
 from .vin_vout import vin_vout
 from .plot3d import plot3d
+from .tsp_graph import visualize_tsp_route
 from .heatmap import heatmap
 
-__all__ = ["m_to_string", "histplot", "energyplot", "vin_vout", "plot3d", "heatmap"]
+__all__ = [
+    "heatmap",
+    "m_to_string",
+    "tsp_hist",
+    "histplot",
+    "energyplot",
+    "vin_vout",
+    "plot3d",
+    "visualize_tsp_route"
+]

p_kit/visualization/tsp_graph.py

@@ -0,0 +1,38 @@
+"""Plot TSP graph"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import networkx as nx
+
+# Graph-Based TSP Path Visualization
+def visualize_tsp_route(path, city_graph):
+    """ Visualizes the best TSP path using NetworkX. """
+
+    if(type(path).__name__ == "str"):
+        path = [int(v) for v in path]
+
+    plt.figure(figsize=(6, 6))
+    n_cities = len(city_graph)
+    G = nx.DiGraph()
+
+    # Define city positions in a circular layout
+    pos = {i: (np.cos(2 * np.pi * i / n_cities), np.sin(2 * np.pi * i / n_cities)) for i in range(n_cities)}
+
+    # Add nodes (cities)
+    for i in range(n_cities):
+        G.add_node(i, pos=pos[i])
+
+    # Add edges based on best path
+    for i in range(len(path) - 1):
+        G.add_edge(path[i], path[i + 1], weight=-city_graph[path[i]][path[i + 1]])
+
+    # Add edge to complete the tour (last city → first city)
+    G.add_edge(path[-1], path[0], weight=-city_graph[path[-1]][path[0]], alpha=0.1)
+
+    # Draw the graph
+
+    labels = nx.get_edge_attributes(G, 'weight')
+    nx.draw(G, pos, with_labels=True, node_color='lightblue', edge_color='r', node_size=1000, font_size=15, arrows=True)
+    nx.draw_networkx_edge_labels(G, pos, edge_labels=labels)
+    plt.title("TSP Solution Path")
+    plt.show()

p_kit/visualization/utils.py

@@ -1,8 +1,32 @@
 """Utils method for visualization"""
+import numpy as np
 
 
 def m_to_string(outputs):
     ret = ""
     for output in outputs:
         ret += "1" if output == 1 else "0"
     return ret
+
+def extract_tsp_path(sample_matrix):
+    """ Extracts the TSP path from a binary matrix. """
+    path = []
+    for order in range(sample_matrix.shape[1]): 
+        step = (sample_matrix[:, order])
+        maxes = np.argwhere(step == np.max(step)).flatten()
+        city = np.random.choice(maxes)
+        path.append(city)
+    return path
+
+def tsp_hist(samples, city_graph):  
+    hist = {}
+    for i in range(len(samples)):
+        s = samples[i, :].reshape((len(city_graph), len(city_graph[0])))
+        path = extract_tsp_path(s)
+        key = "".join([str(p) for p in path])
+        if key in hist:
+            hist[key] = hist[key] + 1
+        else:
+            hist[key] = 1
+
+    return hist