From d2769fabc12079014e7acce9312b92f3891806c0 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Fri, 7 Feb 2025 23:21:36 +0100
Subject: [PATCH 01/13] add tsp

---
 examples/tsp.py             | 75 +++++++++++++++++++++++++++++++++++++
 p_kit/core/__init__.py      |  2 +-
 p_kit/core/p_circuit.py     | 57 ++++++++++++++++++++++++++--
 p_kit/solver/base_solver.py |  9 ++++-
 p_kit/solver/csd_solver.py  |  5 +--
 setup.py                    |  2 +-
 6 files changed, 141 insertions(+), 9 deletions(-)
 create mode 100644 examples/tsp.py

diff --git a/examples/tsp.py b/examples/tsp.py
new file mode 100644
index 0000000..9878caf
--- /dev/null
+++ b/examples/tsp.py
@@ -0,0 +1,75 @@
+from p_kit.core import convert_city_graph
+import matplotlib.pyplot as plt
+import numpy as np
+import networkx as nx
+
+from p_kit.solver.csd_solver import CaSuDaSolver
+
+np.set_printoptions(linewidth=300)
+
+# Define city graph (distance matrix)
+city_graph = np.array([[0, 510, 480],
+                      [510, 0, 240],
+                      [480, 240, 0]])
+
+# Convert to negative weights for optimization
+city_graph *= -1
+
+# Initialize p-circuit
+myp = convert_city_graph(city_graph=city_graph, tsp_modifier=0.33)
+
+# Get interaction weights and biases
+J, h = myp.J, myp.h
+print(J.shape, h.shape)
+print("Interaction Matrix (J):\n", J)
+
+# Simulated Annealing Schedule
+solver = CaSuDaSolver(Nt=10000, dt=0.1667, i0=0.9, expected_mean=0, seed=930)
+_, samples, _ = solver.solve(myp)
+
+# 🏆 Extract Best Path from Binary Matrix
+def extract_tsp_path(sample_matrix):
+    """ Extracts the TSP path from a binary matrix. """
+    path = []
+    for order in range(sample_matrix.shape[1]):  # Iterate through tour steps
+        city = np.argmax(sample_matrix[:, order])  # Find city with value '1'
+        path.append(city)
+    return path
+
+print(samples)
+final_sample = samples[-1, :].reshape((len(city_graph), len(city_graph[0])))
+best_path = extract_tsp_path(final_sample)
+print("Best TSP Path Found:", best_path)
+
+
+# 📌 Graph-Based TSP Path Visualization
+def visualize_tsp_route(path, city_graph):
+    """ Visualizes the best TSP path using NetworkX. """
+    Nm = len(city_graph)
+    G = nx.DiGraph()
+
+    # Define city positions in a circular layout
+    pos = {i: (np.cos(2 * np.pi * i / Nm), np.sin(2 * np.pi * i / Nm)) for i in range(Nm)}
+
+    # Add nodes (cities)
+    for i in range(Nm):
+        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]])
+
+    # Draw the graph
+    plt.figure(figsize=(6, 6))
+    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()
+
+# 🔹 Visualize Best Path
+visualize_tsp_route(best_path, city_graph)
\ No newline at end of file
diff --git a/p_kit/core/__init__.py b/p_kit/core/__init__.py
index a563ace..7e6eba0 100644
--- a/p_kit/core/__init__.py
+++ b/p_kit/core/__init__.py
@@ -1,3 +1,3 @@
-from .p_circuit import PCircuit
+from .p_circuit import PCircuit, convert_city_graph
 
 __all__ = ["PCircuit"]
diff --git a/p_kit/core/p_circuit.py b/p_kit/core/p_circuit.py
index c34b58b..67df5c1 100644
--- a/p_kit/core/p_circuit.py
+++ b/p_kit/core/p_circuit.py
@@ -25,12 +25,63 @@ class PCircuit():
 
     def __init__(self, n_pbits):
         self.n_pbits = n_pbits
-        self.h = np.zeros((n_pbits, 1))
+        self.h = np.zeros((n_pbits,))
         self.J = np.zeros((n_pbits, n_pbits))
     
     def set_weight(self, from_pbit, to_pbit, weight, sym=True):
         self.J[from_pbit, to_pbit] = weight
         if(sym):
             self.J[to_pbit, from_pbit] = weight
-    
-                    
\ No newline at end of file
+
+
+def convert_city_graph(city_graph=None, tsp_modifier=None):
+            """
+            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
+            
+            See excel sheet (building_J_for_tsp in shark tank 2020 purdue onedrive)
+            """
+
+            city_graph = np.asarray(city_graph)
+
+            city_graph = np.divide(city_graph, np.amax(np.abs(city_graph)))
+
+            if tsp_modifier is None:
+                tsp_modifier = 1
+
+            Nm_cities = len(city_graph[0])
+
+            circuit = PCircuit(Nm_cities ** 2)
+
+            # Rule 3: negative distances from one city to another
+            for i in range(Nm_cities):
+                for j in range(Nm_cities):
+                    #if order(i) is one away from order(j)
+                    if i == j:
+                        continue
+                    off_diag = np.ones(Nm_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)
+
+                    circuit.J[j * Nm_cities: j * Nm_cities + Nm_cities, i * Nm_cities: i * Nm_cities + Nm_cities] = weights_i_j[:,:]
+
+            # Rule 1 - 1 between pbits of same city
+            for i in range(Nm_cities):
+                circuit.J[i * Nm_cities:i * Nm_cities + Nm_cities, i * Nm_cities:i * Nm_cities + Nm_cities] = tsp_modifier # dif
+
+            # Rule 2 - 1 between pbits of same order
+            for i in range(Nm_cities ** 2):
+                for j in range(Nm_cities ** 2):
+                    if i == j % Nm_cities or j == i % Nm_cities:
+                        circuit.J[i, j] = tsp_modifier
+
+            # Rule 4: 0s on the diagonal
+            np.fill_diagonal(circuit.J, 0)
+
+            return circuit
+
diff --git a/p_kit/solver/base_solver.py b/p_kit/solver/base_solver.py
index ab36374..d08a005 100644
--- a/p_kit/solver/base_solver.py
+++ b/p_kit/solver/base_solver.py
@@ -1,9 +1,16 @@
+import random
+
 class Solver:
-    def __init__(self, Nt, dt, i0, expected_mean=0) -> None:
+    def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
         self.Nt = Nt
         self.dt = dt
         self.i0 = i0
         self.expected_mean = expected_mean
+        print(seed)
+        self._random_gen = random.Random(seed)
+    
+    def random(self):
+        return self._random_gen.random()
 
     def solve(self):
         raise NotImplementedError()
diff --git a/p_kit/solver/csd_solver.py b/p_kit/solver/csd_solver.py
index 53505fe..3f3d515 100644
--- a/p_kit/solver/csd_solver.py
+++ b/p_kit/solver/csd_solver.py
@@ -1,6 +1,5 @@
 from p_kit.core.p_circuit import PCircuit
 from .base_solver import Solver
-from random import random
 import numpy as np
 
 class CaSuDaSolver(Solver):
@@ -15,7 +14,7 @@ def solve(self, c: PCircuit):
         all_m = [[]] * self.Nt
         E = [0] * self.Nt
 
-        m = [np.sign(0.5 - random()) for _ in indices]
+        m = [np.sign(0.5 - self.random()) for _ in indices]
 
         for run in range(self.Nt):
 
@@ -27,7 +26,7 @@ def solve(self, c: PCircuit):
             s = [np.exp(-1 * self.dt * np.exp(-1 * m[i] * (I[i] + threshold))) for i in indices]
 
             # compute new output
-            m = [m[i] * np.sign(s[i] - random()) for i in indices]
+            m = [m[i] * np.sign(s[i] - self.random()) for i in indices]
             
             all_I[run] = [_ for _ in I]
             all_m[run] = [_ for _ in m]
diff --git a/setup.py b/setup.py
index 1515ba0..ca4f80f 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,7 @@
                         'matplotlib==3.9.3'
                         ],
       extras_require={
-                      'tests': ['pytest', 'seaborn', 'flake8'],
+                      'tests': ['pytest', 'seaborn', 'flake8', 'networkx'],
                       },
       zip_safe=False,
 )

From 711a403c34b26bdac50e8a00bfbf49d1def38282 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Thu, 13 Feb 2025 09:07:40 +0100
Subject: [PATCH 02/13] add annealing improve visualization

---
 examples/tsp.py             | 65 ++++++++++++++++++++++++++++++-------
 p_kit/solver/base_solver.py |  2 +-
 p_kit/solver/csd_solver.py  |  8 +++--
 3 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/examples/tsp.py b/examples/tsp.py
index 9878caf..ef744ef 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -4,6 +4,7 @@
 import networkx as nx
 
 from p_kit.solver.csd_solver import CaSuDaSolver
+from p_kit.visualization import histplot, plot3d
 
 np.set_printoptions(linewidth=300)
 
@@ -12,6 +13,11 @@
                       [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]])
+
 # Convert to negative weights for optimization
 city_graph *= -1
 
@@ -23,24 +29,54 @@
 print(J.shape, h.shape)
 print("Interaction Matrix (J):\n", J)
 
+Nt = int(1e5)
+def annealing(i0, run):
+    v = (0 - i0) / Nt * (run - 1) + i0
+    # v = i0 * (1.0001)**(run - 1)
+    # print(v)
+    return v
+
 # Simulated Annealing Schedule
-solver = CaSuDaSolver(Nt=10000, dt=0.1667, i0=0.9, expected_mean=0, seed=930)
-_, samples, _ = solver.solve(myp)
+end_samples = []
+solver = CaSuDaSolver(Nt=Nt, dt=1 / (2 * len(city_graph)), i0=0.005)
+
+for i in range(100):
+    _, samples, _ = solver.solve(myp, annealing_func=annealing)
+    end_samples.append(samples[-1, :])
+
+samples = np.array(end_samples)
+# plot3d(samples, A=[0, 1, 2], B=[3, 4, 5])
 
 # 🏆 Extract Best Path from Binary Matrix
 def extract_tsp_path(sample_matrix):
     """ Extracts the TSP path from a binary matrix. """
     path = []
     for order in range(sample_matrix.shape[1]):  # Iterate through tour steps
-        city = np.argmax(sample_matrix[:, order])  # Find city with value '1'
+        #print(sample_matrix[:, order])
+        step = (sample_matrix[:, order])
+        # city = np.argmax(step)  # Find city with value '1'
+
+        maxes = np.argwhere(step == np.max(step)).flatten()
+        city = np.random.choice(maxes)
         path.append(city)
     return path
 
 print(samples)
-final_sample = samples[-1, :].reshape((len(city_graph), len(city_graph[0])))
-best_path = extract_tsp_path(final_sample)
-print("Best TSP Path Found:", best_path)
-
+hist = {}
+for i in range(len(samples)):
+    s = samples[i, :].reshape((len(city_graph), len(city_graph[0])))
+    path = extract_tsp_path(s)
+    #print(path)
+    key = "".join([str(p) for p in path])
+    if key in hist:
+        hist[key] = hist[key] + 1
+    else:
+        hist[key] = 1
+
+print(hist)
+plt.bar(hist.keys(), hist.values())
+# plt.show()
+print(extract_tsp_path(np.array([s[-1, :]])))
 
 # 📌 Graph-Based TSP Path Visualization
 def visualize_tsp_route(path, city_graph):
@@ -60,16 +96,21 @@ def visualize_tsp_route(path, city_graph):
         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]])
+    G.add_edge(path[-1], path[0], weight=-city_graph[path[-1]][path[0]], alpha=0.1)
 
     # Draw the graph
-    plt.figure(figsize=(6, 6))
+    
     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()
+   
+    
 
 # 🔹 Visualize Best Path
-visualize_tsp_route(best_path, city_graph)
\ No newline at end of file
+# plt.figure(figsize=(6, 6))
+# for path, count in hist.items():
+#     visualize_tsp_route([int(path[0]), int(path[1]), int(path[2])], city_graph)
+
+# plt.title("TSP Solution Path")
+plt.show()
\ No newline at end of file
diff --git a/p_kit/solver/base_solver.py b/p_kit/solver/base_solver.py
index d08a005..8ad166e 100644
--- a/p_kit/solver/base_solver.py
+++ b/p_kit/solver/base_solver.py
@@ -12,5 +12,5 @@ def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
     def random(self):
         return self._random_gen.random()
 
-    def solve(self):
+    def solve(self, annealing_func):
         raise NotImplementedError()
diff --git a/p_kit/solver/csd_solver.py b/p_kit/solver/csd_solver.py
index 3f3d515..cec3406 100644
--- a/p_kit/solver/csd_solver.py
+++ b/p_kit/solver/csd_solver.py
@@ -5,7 +5,11 @@
 class CaSuDaSolver(Solver):
     # K. Y. Camsari, B. M. Sutton, and S. Datta, ‘p-bits for probabilistic spin logic’, Applied Physics Reviews, vol. 6, no. 1, p. 011305, Mar. 2019, doi: 10.1063/1.5055860.
 
-    def solve(self, c: PCircuit):
+    def solve(self, c: PCircuit, annealing_func=None):
+
+        if annealing_func == None:
+            annealing_func = lambda i0, run: i0
+    
         # credit: https://www.purdue.edu/p-bit/blog.html
         n_pbits = c.n_pbits
         indices = range(n_pbits)
@@ -19,7 +23,7 @@ def solve(self, c: PCircuit):
         for run in range(self.Nt):
 
             # compute input biases
-            I = [self.i0 * (np.dot(m, c.J[i]) + c.h[i]) for i in indices]
+            I = [annealing_func(self.i0, run) * (np.dot(m, c.J[i]) + c.h[i]) for i in indices]
             
             # apply S(input)
             threshold = np.arctanh(self.expected_mean)

From a4230a400243c20ca6437f672e392203fc12c959 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Thu, 13 Feb 2025 09:58:56 +0100
Subject: [PATCH 03/13] add annealing module improve example

---
 examples/tsp.py            | 18 ++++------
 p_kit/core/__init__.py     |  2 +-
 p_kit/core/p_circuit.py    | 51 ---------------------------
 p_kit/library/__init__.py  |  3 ++
 p_kit/library/tsp.py       | 70 ++++++++++++++++++++++++++++++++++++++
 p_kit/solver/__init__.py   |  3 +-
 p_kit/solver/annealing.py  |  6 ++++
 p_kit/solver/csd_solver.py |  8 ++---
 8 files changed, 91 insertions(+), 70 deletions(-)
 create mode 100644 p_kit/library/__init__.py
 create mode 100644 p_kit/library/tsp.py
 create mode 100644 p_kit/solver/annealing.py

diff --git a/examples/tsp.py b/examples/tsp.py
index ef744ef..0511cda 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -1,8 +1,9 @@
-from p_kit.core import convert_city_graph
 import matplotlib.pyplot as plt
 import numpy as np
 import networkx as nx
 
+from p_kit.library import TSP
+from p_kit.solver.annealing import linear
 from p_kit.solver.csd_solver import CaSuDaSolver
 from p_kit.visualization import histplot, plot3d
 
@@ -22,26 +23,19 @@
 city_graph *= -1
 
 # Initialize p-circuit
-myp = convert_city_graph(city_graph=city_graph, tsp_modifier=0.33)
+myp = TSP(city_graph=city_graph, tsp_modifier=0.33)
 
 # Get interaction weights and biases
 J, h = myp.J, myp.h
 print(J.shape, h.shape)
 print("Interaction Matrix (J):\n", J)
 
-Nt = int(1e5)
-def annealing(i0, run):
-    v = (0 - i0) / Nt * (run - 1) + i0
-    # v = i0 * (1.0001)**(run - 1)
-    # print(v)
-    return v
-
 # Simulated Annealing Schedule
 end_samples = []
-solver = CaSuDaSolver(Nt=Nt, dt=1 / (2 * len(city_graph)), i0=0.005)
+solver = CaSuDaSolver(Nt=int(1e4), dt=1 / (2 * len(city_graph)), i0=0)
 
-for i in range(100):
-    _, samples, _ = solver.solve(myp, annealing_func=annealing)
+for i in range(200):
+    _, samples, _ = solver.solve(myp, annealing_func=linear)
     end_samples.append(samples[-1, :])
 
 samples = np.array(end_samples)
diff --git a/p_kit/core/__init__.py b/p_kit/core/__init__.py
index 7e6eba0..a563ace 100644
--- a/p_kit/core/__init__.py
+++ b/p_kit/core/__init__.py
@@ -1,3 +1,3 @@
-from .p_circuit import PCircuit, convert_city_graph
+from .p_circuit import PCircuit
 
 __all__ = ["PCircuit"]
diff --git a/p_kit/core/p_circuit.py b/p_kit/core/p_circuit.py
index 67df5c1..1250eca 100644
--- a/p_kit/core/p_circuit.py
+++ b/p_kit/core/p_circuit.py
@@ -34,54 +34,3 @@ def set_weight(self, from_pbit, to_pbit, weight, sym=True):
             self.J[to_pbit, from_pbit] = weight
 
 
-def convert_city_graph(city_graph=None, tsp_modifier=None):
-            """
-            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
-            
-            See excel sheet (building_J_for_tsp in shark tank 2020 purdue onedrive)
-            """
-
-            city_graph = np.asarray(city_graph)
-
-            city_graph = np.divide(city_graph, np.amax(np.abs(city_graph)))
-
-            if tsp_modifier is None:
-                tsp_modifier = 1
-
-            Nm_cities = len(city_graph[0])
-
-            circuit = PCircuit(Nm_cities ** 2)
-
-            # Rule 3: negative distances from one city to another
-            for i in range(Nm_cities):
-                for j in range(Nm_cities):
-                    #if order(i) is one away from order(j)
-                    if i == j:
-                        continue
-                    off_diag = np.ones(Nm_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)
-
-                    circuit.J[j * Nm_cities: j * Nm_cities + Nm_cities, i * Nm_cities: i * Nm_cities + Nm_cities] = weights_i_j[:,:]
-
-            # Rule 1 - 1 between pbits of same city
-            for i in range(Nm_cities):
-                circuit.J[i * Nm_cities:i * Nm_cities + Nm_cities, i * Nm_cities:i * Nm_cities + Nm_cities] = tsp_modifier # dif
-
-            # Rule 2 - 1 between pbits of same order
-            for i in range(Nm_cities ** 2):
-                for j in range(Nm_cities ** 2):
-                    if i == j % Nm_cities or j == i % Nm_cities:
-                        circuit.J[i, j] = tsp_modifier
-
-            # Rule 4: 0s on the diagonal
-            np.fill_diagonal(circuit.J, 0)
-
-            return circuit
-
diff --git a/p_kit/library/__init__.py b/p_kit/library/__init__.py
new file mode 100644
index 0000000..089146a
--- /dev/null
+++ b/p_kit/library/__init__.py
@@ -0,0 +1,3 @@
+from .tsp import TSP
+
+__all__ = ["TSP"]
diff --git a/p_kit/library/tsp.py b/p_kit/library/tsp.py
new file mode 100644
index 0000000..a59a1e7
--- /dev/null
+++ b/p_kit/library/tsp.py
@@ -0,0 +1,70 @@
+import numpy as np
+from p_kit.core 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
+    ----------
+    n_pbits: string
+        Identifier of the pipeline (for log purposes).
+
+    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.tsp_modifier = 1 if tsp_modifier is None else tsp_modifier
+        self._init_city_graph(city_graph)
+    
+
+    def _init_city_graph(self, city_graph=None):
+
+            city_graph = np.asarray(city_graph)
+            city_graph = np.divide(city_graph, np.amax(np.abs(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)
+
+
diff --git a/p_kit/solver/__init__.py b/p_kit/solver/__init__.py
index 1ef5810..9f0e979 100644
--- a/p_kit/solver/__init__.py
+++ b/p_kit/solver/__init__.py
@@ -1,3 +1,4 @@
 from .csd_solver import CaSuDaSolver
+from . import annealing
 
-__all__ = ["CaSuDaSolver"]
+__all__ = ["CaSuDaSolver", "annealing"]
diff --git a/p_kit/solver/annealing.py b/p_kit/solver/annealing.py
new file mode 100644
index 0000000..9ce10ed
--- /dev/null
+++ b/p_kit/solver/annealing.py
@@ -0,0 +1,6 @@
+
+def constant(solver, _run):
+    return solver.i0
+
+def linear(solver, run):
+    return (0 - solver.i0) / solver.Nt * (run - 1) + solver.i0
diff --git a/p_kit/solver/csd_solver.py b/p_kit/solver/csd_solver.py
index cec3406..4e8e2db 100644
--- a/p_kit/solver/csd_solver.py
+++ b/p_kit/solver/csd_solver.py
@@ -1,14 +1,12 @@
 from p_kit.core.p_circuit import PCircuit
+from p_kit.solver.annealing import constant
 from .base_solver import Solver
 import numpy as np
 
 class CaSuDaSolver(Solver):
     # K. Y. Camsari, B. M. Sutton, and S. Datta, ‘p-bits for probabilistic spin logic’, Applied Physics Reviews, vol. 6, no. 1, p. 011305, Mar. 2019, doi: 10.1063/1.5055860.
 
-    def solve(self, c: PCircuit, annealing_func=None):
-
-        if annealing_func == None:
-            annealing_func = lambda i0, run: i0
+    def solve(self, c: PCircuit, annealing_func=constant):
     
         # credit: https://www.purdue.edu/p-bit/blog.html
         n_pbits = c.n_pbits
@@ -23,7 +21,7 @@ def solve(self, c: PCircuit, annealing_func=None):
         for run in range(self.Nt):
 
             # compute input biases
-            I = [annealing_func(self.i0, run) * (np.dot(m, c.J[i]) + c.h[i]) for i in indices]
+            I = [annealing_func(self, run) * (np.dot(m, c.J[i]) + c.h[i]) for i in indices]
             
             # apply S(input)
             threshold = np.arctanh(self.expected_mean)

From 2bdc549a4930ca6aa96704fb0293b33f26a430a6 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Thu, 13 Feb 2025 11:30:57 +0100
Subject: [PATCH 04/13] extract usefull methods from example to separate
 component for reusability

---
 examples/tsp.py                  | 97 ++++++--------------------------
 p_kit/library/__init__.py        |  4 +-
 p_kit/library/tsp.py             | 12 ++--
 p_kit/solver/annealing.py        |  9 +++
 p_kit/solver/base_solver.py      |  1 -
 p_kit/visualization/__init__.py  | 13 ++++-
 p_kit/visualization/tsp_graph.py | 38 +++++++++++++
 p_kit/visualization/utils.py     | 24 ++++++++
 setup.py                         |  5 +-
 9 files changed, 109 insertions(+), 94 deletions(-)
 create mode 100644 p_kit/visualization/tsp_graph.py

diff --git a/examples/tsp.py b/examples/tsp.py
index 0511cda..da2c088 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -1,11 +1,14 @@
+# Adapted from: 
+# https://github.com/anirudhgha/p-bit/tree/master/Research/Travelling%20Salesman%20Problem
+
 import matplotlib.pyplot as plt
 import numpy as np
-import networkx as nx
 
-from p_kit.library import TSP
-from p_kit.solver.annealing import linear
+from p_kit.library.tsp import TSP
+from p_kit.solver.annealing import execute, linear
 from p_kit.solver.csd_solver import CaSuDaSolver
-from p_kit.visualization import histplot, plot3d
+from p_kit.visualization.tsp_graph import visualize_tsp_route
+from p_kit.visualization.utils import tsp_hist
 
 np.set_printoptions(linewidth=300)
 
@@ -23,88 +26,20 @@
 city_graph *= -1
 
 # Initialize p-circuit
-myp = TSP(city_graph=city_graph, tsp_modifier=0.33)
-
-# Get interaction weights and biases
-J, h = myp.J, myp.h
-print(J.shape, h.shape)
-print("Interaction Matrix (J):\n", J)
+circuit = TSP(city_graph=city_graph, tsp_modifier=0.33)
 
 # Simulated Annealing Schedule
-end_samples = []
 solver = CaSuDaSolver(Nt=int(1e4), dt=1 / (2 * len(city_graph)), i0=0)
 
-for i in range(200):
-    _, samples, _ = solver.solve(myp, annealing_func=linear)
-    end_samples.append(samples[-1, :])
-
-samples = np.array(end_samples)
-# plot3d(samples, A=[0, 1, 2], B=[3, 4, 5])
-
-# 🏆 Extract Best Path from Binary Matrix
-def extract_tsp_path(sample_matrix):
-    """ Extracts the TSP path from a binary matrix. """
-    path = []
-    for order in range(sample_matrix.shape[1]):  # Iterate through tour steps
-        #print(sample_matrix[:, order])
-        step = (sample_matrix[:, order])
-        # city = np.argmax(step)  # Find city with value '1'
+# Increase n_shots to have more reliable results
+n_shots = 200
+samples = execute(solver, circuit, linear, n_shots)
 
-        maxes = np.argwhere(step == np.max(step)).flatten()
-        city = np.random.choice(maxes)
-        path.append(city)
-    return path
-
-print(samples)
-hist = {}
-for i in range(len(samples)):
-    s = samples[i, :].reshape((len(city_graph), len(city_graph[0])))
-    path = extract_tsp_path(s)
-    #print(path)
-    key = "".join([str(p) for p in path])
-    if key in hist:
-        hist[key] = hist[key] + 1
-    else:
-        hist[key] = 1
-
-print(hist)
+# Visualize Best path
+hist = tsp_hist(samples, city_graph)
 plt.bar(hist.keys(), hist.values())
-# plt.show()
-print(extract_tsp_path(np.array([s[-1, :]])))
-
-# 📌 Graph-Based TSP Path Visualization
-def visualize_tsp_route(path, city_graph):
-    """ Visualizes the best TSP path using NetworkX. """
-    Nm = len(city_graph)
-    G = nx.DiGraph()
-
-    # Define city positions in a circular layout
-    pos = {i: (np.cos(2 * np.pi * i / Nm), np.sin(2 * np.pi * i / Nm)) for i in range(Nm)}
+plt.show()
 
-    # Add nodes (cities)
-    for i in range(Nm):
-        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
+best_path = max(hist, key=hist.get)
     
-    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)
-
-   
-    
-
-# 🔹 Visualize Best Path
-# plt.figure(figsize=(6, 6))
-# for path, count in hist.items():
-#     visualize_tsp_route([int(path[0]), int(path[1]), int(path[2])], city_graph)
-
-# plt.title("TSP Solution Path")
-plt.show()
\ No newline at end of file
+visualize_tsp_route(best_path, city_graph)
diff --git a/p_kit/library/__init__.py b/p_kit/library/__init__.py
index 089146a..15e8e39 100644
--- a/p_kit/library/__init__.py
+++ b/p_kit/library/__init__.py
@@ -1,3 +1,3 @@
-from .tsp import TSP
+from . import tsp
 
-__all__ = ["TSP"]
+__all__ = ["tsp"]
diff --git a/p_kit/library/tsp.py b/p_kit/library/tsp.py
index a59a1e7..b22981b 100644
--- a/p_kit/library/tsp.py
+++ b/p_kit/library/tsp.py
@@ -1,3 +1,5 @@
+# Adapted from: https://github.com/anirudhgha/p-bit/blob/master/Research/Travelling%20Salesman%20Problem/tsp.py
+
 import numpy as np
 from p_kit.core import PCircuit
 
@@ -31,15 +33,13 @@ class TSP(PCircuit):
 
     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(city_graph)
+        self._init_city_graph()
     
 
-    def _init_city_graph(self, city_graph=None):
-
-            city_graph = np.asarray(city_graph)
-            city_graph = np.divide(city_graph, np.amax(np.abs(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
diff --git a/p_kit/solver/annealing.py b/p_kit/solver/annealing.py
index 9ce10ed..7369a9a 100644
--- a/p_kit/solver/annealing.py
+++ b/p_kit/solver/annealing.py
@@ -1,6 +1,15 @@
 
+import numpy as np
+
 def constant(solver, _run):
     return solver.i0
 
 def linear(solver, run):
     return (0 - solver.i0) / solver.Nt * (run - 1) + solver.i0
+
+def execute(solver, circuit, annealing, n_shots):
+    samples = []
+    for _ in range(n_shots):
+        _, time_points, _ = solver.solve(circuit, annealing_func=annealing)
+        samples.append(time_points[-1, :])
+    return np.array(samples)
\ No newline at end of file
diff --git a/p_kit/solver/base_solver.py b/p_kit/solver/base_solver.py
index 8ad166e..6fe92d5 100644
--- a/p_kit/solver/base_solver.py
+++ b/p_kit/solver/base_solver.py
@@ -6,7 +6,6 @@ def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
         self.dt = dt
         self.i0 = i0
         self.expected_mean = expected_mean
-        print(seed)
         self._random_gen = random.Random(seed)
     
     def random(self):
diff --git a/p_kit/visualization/__init__.py b/p_kit/visualization/__init__.py
index 7c64f35..d9c39c2 100644
--- a/p_kit/visualization/__init__.py
+++ b/p_kit/visualization/__init__.py
@@ -1,6 +1,15 @@
-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
 
-__all__ = ["m_to_string", "histplot", "energyplot", "vin_vout", "plot3d"]
+__all__ = [
+    "m_to_string",
+    "tsp_hist",
+    "histplot",
+    "energyplot",
+    "vin_vout",
+    "plot3d",
+    "visualize_tsp_route"
+]
diff --git a/p_kit/visualization/tsp_graph.py b/p_kit/visualization/tsp_graph.py
new file mode 100644
index 0000000..a6b2a04
--- /dev/null
+++ b/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()
\ No newline at end of file
diff --git a/p_kit/visualization/utils.py b/p_kit/visualization/utils.py
index c582ed7..ef47ca3 100644
--- a/p_kit/visualization/utils.py
+++ b/p_kit/visualization/utils.py
@@ -1,7 +1,31 @@
 """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
\ No newline at end of file
diff --git a/setup.py b/setup.py
index ca4f80f..ec71f0c 100644
--- a/setup.py
+++ b/setup.py
@@ -38,10 +38,11 @@
                         'cython==3.0.11',
                         'cvxpy==1.6.0',
                         'scipy==1.13.1',
-                        'matplotlib==3.9.3'
+                        'matplotlib==3.9.3',
+                        'networkx'
                         ],
       extras_require={
-                      'tests': ['pytest', 'seaborn', 'flake8', 'networkx'],
+                      'tests': ['pytest', 'seaborn', 'flake8'],
                       },
       zip_safe=False,
 )

From 6f5cdded7068edc3707d0a393f70757dcfe65076 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Thu, 13 Feb 2025 11:33:15 +0100
Subject: [PATCH 05/13] improve documentation

---
 p_kit/library/tsp.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/p_kit/library/tsp.py b/p_kit/library/tsp.py
index b22981b..d257e5f 100644
--- a/p_kit/library/tsp.py
+++ b/p_kit/library/tsp.py
@@ -15,8 +15,10 @@ class TSP(PCircuit):
 
     Parameters
     ----------
-    n_pbits: string
-        Identifier of the pipeline (for log purposes).
+    city_graph: np.array((n_cities, n_cities))
+        The city graph.
+    tsp_modifier: float, default: 1
+        Value between pbits of same city
 
     Attributes
     ----------

From b6a532d44c3828507eb7a5691e0a5a9c976e7680 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Thu, 13 Feb 2025 15:29:21 +0100
Subject: [PATCH 06/13] threaded version of execute annealing

---
 examples/tsp.py              |  8 +++-----
 p_kit/core/p_circuit.py      |  5 +++++
 p_kit/solver/annealing.py    | 13 ++++++++-----
 p_kit/solver/base_solver.py  |  6 ++++++
 p_kit/solver/csd_solver.py   |  3 +++
 p_kit/visualization/utils.py |  2 +-
 setup.py                     |  3 ++-
 7 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/examples/tsp.py b/examples/tsp.py
index da2c088..e222b39 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -10,8 +10,6 @@
 from p_kit.visualization.tsp_graph import visualize_tsp_route
 from p_kit.visualization.utils import tsp_hist
 
-np.set_printoptions(linewidth=300)
-
 # Define city graph (distance matrix)
 city_graph = np.array([[0, 510, 480],
                       [510, 0, 240],
@@ -29,11 +27,11 @@
 circuit = TSP(city_graph=city_graph, tsp_modifier=0.33)
 
 # Simulated Annealing Schedule
-solver = CaSuDaSolver(Nt=int(1e4), dt=1 / (2 * len(city_graph)), i0=0)
+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_shots = 100
+samples = execute(solver, circuit, linear, n_shots, n_jobs=-1)
 
 # Visualize Best path
 hist = tsp_hist(samples, city_graph)
diff --git a/p_kit/core/p_circuit.py b/p_kit/core/p_circuit.py
index 1250eca..865198a 100644
--- a/p_kit/core/p_circuit.py
+++ b/p_kit/core/p_circuit.py
@@ -34,3 +34,8 @@ def set_weight(self, from_pbit, to_pbit, weight, sym=True):
             self.J[to_pbit, from_pbit] = weight
 
 
+    def copy(self):
+        new_circuit = PCircuit(self.n_pbits)
+        new_circuit.J = self.J
+        new_circuit.h = self.h
+        return new_circuit
\ No newline at end of file
diff --git a/p_kit/solver/annealing.py b/p_kit/solver/annealing.py
index 7369a9a..09fbf3d 100644
--- a/p_kit/solver/annealing.py
+++ b/p_kit/solver/annealing.py
@@ -1,5 +1,6 @@
 
 import numpy as np
+from joblib import Parallel, delayed
 
 def constant(solver, _run):
     return solver.i0
@@ -7,9 +8,11 @@ def constant(solver, _run):
 def linear(solver, run):
     return (0 - solver.i0) / solver.Nt * (run - 1) + solver.i0
 
-def execute(solver, circuit, annealing, n_shots):
-    samples = []
-    for _ in range(n_shots):
-        _, time_points, _ = solver.solve(circuit, annealing_func=annealing)
-        samples.append(time_points[-1, :])
+def execute(solver, circuit, annealing, n_shots, n_jobs=-1):
+    def job():
+        _, time_points, _ = solver.copy().solve(circuit.copy(), annealing_func=annealing)
+        return time_points[-1, :]
+
+    samples = Parallel(n_jobs=n_jobs)(delayed(job)() for _ in range(n_shots))
+
     return np.array(samples)
\ No newline at end of file
diff --git a/p_kit/solver/base_solver.py b/p_kit/solver/base_solver.py
index 6fe92d5..a272555 100644
--- a/p_kit/solver/base_solver.py
+++ b/p_kit/solver/base_solver.py
@@ -1,4 +1,5 @@
 import random
+import time
 
 class Solver:
     def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
@@ -6,6 +7,8 @@ def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
         self.dt = dt
         self.i0 = i0
         self.expected_mean = expected_mean
+        # self.seed = time.time() if seed is None else seed
+        self.seed = seed
         self._random_gen = random.Random(seed)
     
     def random(self):
@@ -13,3 +16,6 @@ def random(self):
 
     def solve(self, annealing_func):
         raise NotImplementedError()
+    
+    def copy(self):
+        raise NotImplementedError()
diff --git a/p_kit/solver/csd_solver.py b/p_kit/solver/csd_solver.py
index 4e8e2db..852be8d 100644
--- a/p_kit/solver/csd_solver.py
+++ b/p_kit/solver/csd_solver.py
@@ -36,3 +36,6 @@ def solve(self, c: PCircuit, annealing_func=constant):
             E[run] = self.i0 * (np.dot(m, c.h) + np.multiply(0.5, np.dot(np.dot(m, c.J), m)))
 
         return np.array(all_I), np.array(all_m), E
+
+    def copy(self):
+        return CaSuDaSolver(self.Nt, self.dt, self.i0, self.expected_mean, self.seed)
\ No newline at end of file
diff --git a/p_kit/visualization/utils.py b/p_kit/visualization/utils.py
index ef47ca3..669a2d9 100644
--- a/p_kit/visualization/utils.py
+++ b/p_kit/visualization/utils.py
@@ -28,4 +28,4 @@ def tsp_hist(samples, city_graph):
         else:
             hist[key] = 1
 
-    return hist
\ No newline at end of file
+    return hist
diff --git a/setup.py b/setup.py
index ec71f0c..9806b26 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,8 @@
                         'cvxpy==1.6.0',
                         'scipy==1.13.1',
                         'matplotlib==3.9.3',
-                        'networkx'
+                        'networkx',
+                        'joblib'
                         ],
       extras_require={
                       'tests': ['pytest', 'seaborn', 'flake8'],

From 767ce9729582848171e4c17680644df6afcd3e79 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Sat, 15 Feb 2025 22:25:18 +0100
Subject: [PATCH 07/13] improve implementation of CSD solver

---
 p_kit/core/p_circuit.py     |  1 -
 p_kit/solver/base_solver.py |  8 ++++----
 p_kit/solver/csd_solver.py  | 31 ++++++++++++++++---------------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/p_kit/core/p_circuit.py b/p_kit/core/p_circuit.py
index 865198a..c841f1b 100644
--- a/p_kit/core/p_circuit.py
+++ b/p_kit/core/p_circuit.py
@@ -33,7 +33,6 @@ def set_weight(self, from_pbit, to_pbit, weight, sym=True):
         if(sym):
             self.J[to_pbit, from_pbit] = weight
 
-
     def copy(self):
         new_circuit = PCircuit(self.n_pbits)
         new_circuit.J = self.J
diff --git a/p_kit/solver/base_solver.py b/p_kit/solver/base_solver.py
index a272555..20528f3 100644
--- a/p_kit/solver/base_solver.py
+++ b/p_kit/solver/base_solver.py
@@ -1,5 +1,6 @@
 import random
 import time
+import numpy as np
 
 class Solver:
     def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
@@ -7,12 +8,11 @@ def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
         self.dt = dt
         self.i0 = i0
         self.expected_mean = expected_mean
-        # self.seed = time.time() if seed is None else seed
         self.seed = seed
-        self._random_gen = random.Random(seed)
+        self._random_gen = np.random.default_rng(self.seed)
     
-    def random(self):
-        return self._random_gen.random()
+    def random(self, n_pbits):
+        return self._random_gen.random(n_pbits)
 
     def solve(self, annealing_func):
         raise NotImplementedError()
diff --git a/p_kit/solver/csd_solver.py b/p_kit/solver/csd_solver.py
index 852be8d..31ed4e4 100644
--- a/p_kit/solver/csd_solver.py
+++ b/p_kit/solver/csd_solver.py
@@ -3,6 +3,7 @@
 from .base_solver import Solver
 import numpy as np
 
+
 class CaSuDaSolver(Solver):
     # K. Y. Camsari, B. M. Sutton, and S. Datta, ‘p-bits for probabilistic spin logic’, Applied Physics Reviews, vol. 6, no. 1, p. 011305, Mar. 2019, doi: 10.1063/1.5055860.
 
@@ -10,32 +11,32 @@ def solve(self, c: PCircuit, annealing_func=constant):
     
         # credit: https://www.purdue.edu/p-bit/blog.html
         n_pbits = c.n_pbits
-        indices = range(n_pbits)
 
-        all_I = [[]] * self.Nt
-        all_m = [[]] * self.Nt
-        E = [0] * self.Nt
+        all_I = np.zeros((self.Nt, n_pbits))
+        all_m = np.zeros((self.Nt, n_pbits))
+        E = np.zeros(self.Nt)
+
+        m = np.sign(0.5 - self.random(n_pbits))
 
-        m = [np.sign(0.5 - self.random()) for _ in indices]
+        threshold = np.arctanh(self.expected_mean)
 
         for run in range(self.Nt):
 
             # compute input biases
-            I = [annealing_func(self, run) * (np.dot(m, c.J[i]) + c.h[i]) for i in indices]
-            
+            I = annealing_func(self, run) * (np.dot(m, c.J) + c.h)
+
             # apply S(input)
-            threshold = np.arctanh(self.expected_mean)
-            s = [np.exp(-1 * self.dt * np.exp(-1 * m[i] * (I[i] + threshold))) for i in indices]
+            s = np.exp(-self.dt * np.exp(-m * (I + threshold)))
 
             # compute new output
-            m = [m[i] * np.sign(s[i] - self.random()) for i in indices]
+            m = m * np.sign(s - self.random(n_pbits))
             
-            all_I[run] = [_ for _ in I]
-            all_m[run] = [_ for _ in m]
+            all_I[run] = I
+            all_m[run] = m
 
-            E[run] = self.i0 * (np.dot(m, c.h) + np.multiply(0.5, np.dot(np.dot(m, c.J), m)))
+            E[run] = self.i0 * (np.dot(m, c.h) + 0.5 * np.dot(np.dot(m, c.J), m))
 
-        return np.array(all_I), np.array(all_m), E
+        return all_I, all_m, E
 
     def copy(self):
-        return CaSuDaSolver(self.Nt, self.dt, self.i0, self.expected_mean, self.seed)
\ No newline at end of file
+        return CaSuDaSolver(self.Nt, self.dt, self.i0, self.expected_mean, self.seed)

From 21390a9df3f07ad48124af1ab67b4f9ccdf37821 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Sat, 15 Feb 2025 22:42:45 +0100
Subject: [PATCH 08/13] improve annealing method to return more samples

---
 examples/tsp.py           | 31 ++++++++++++++++++++-----------
 p_kit/solver/annealing.py |  6 +++---
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/examples/tsp.py b/examples/tsp.py
index e222b39..5573114 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -5,20 +5,29 @@
 import numpy as np
 
 from p_kit.library.tsp import TSP
-from p_kit.solver.annealing import execute, linear
+from p_kit.solver.annealing import constant, execute, linear
 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, 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
@@ -27,11 +36,11 @@
 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)
+solver = CaSuDaSolver(Nt=int(1e6), dt=1 / (2 * len(city_graph)), i0=0.005, expected_mean=0, seed=None)
 
 # Increase n_shots to have more reliable results
 n_shots = 100
-samples = execute(solver, circuit, linear, n_shots, n_jobs=-1)
+samples = execute(solver, circuit, constant, n_shots, n_last_samples=50, n_jobs=-1)
 
 # Visualize Best path
 hist = tsp_hist(samples, city_graph)
diff --git a/p_kit/solver/annealing.py b/p_kit/solver/annealing.py
index 09fbf3d..bd45eea 100644
--- a/p_kit/solver/annealing.py
+++ b/p_kit/solver/annealing.py
@@ -8,11 +8,11 @@ def constant(solver, _run):
 def linear(solver, run):
     return (0 - solver.i0) / solver.Nt * (run - 1) + solver.i0
 
-def execute(solver, circuit, annealing, n_shots, n_jobs=-1):
+def execute(solver, circuit, annealing, n_shots, n_last_samples=50, n_jobs=-1):
     def job():
         _, time_points, _ = solver.copy().solve(circuit.copy(), annealing_func=annealing)
-        return time_points[-1, :]
+        return time_points[-n_last_samples-1:-1, :]
 
     samples = Parallel(n_jobs=n_jobs)(delayed(job)() for _ in range(n_shots))
 
-    return np.array(samples)
\ No newline at end of file
+    return np.array(samples).reshape((n_shots * n_last_samples, circuit.n_pbits))
\ No newline at end of file

From 1e04ff512ca7cf266237ded889f75777b49bc9e8 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Sat, 15 Feb 2025 23:05:17 +0100
Subject: [PATCH 09/13] small changes to example

---
 examples/tsp.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/examples/tsp.py b/examples/tsp.py
index 5573114..4f05266 100644
--- a/examples/tsp.py
+++ b/examples/tsp.py
@@ -6,6 +6,7 @@
 
 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
@@ -13,7 +14,7 @@
 # Define city graph (distance matrix)
 # city_graph = np.array([[0, 510, 480],
 #                       [510, 0, 240],
-#                       [480, 240, 0]])
+                    #   [480, 240, 0]])
 
 city_graph = np.array([[0, 510, 480, 490],
               [510, 0, 240, 370],
@@ -36,11 +37,11 @@
 circuit = TSP(city_graph=city_graph, tsp_modifier=0.33)
 
 # Simulated Annealing Schedule
-solver = CaSuDaSolver(Nt=int(1e6), dt=1 / (2 * len(city_graph)), i0=0.005, expected_mean=0, seed=None)
+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 = 100
-samples = execute(solver, circuit, constant, n_shots, n_last_samples=50, n_jobs=-1)
+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)

From db59759907652b09dab9870fa21d2f33a5aa91fa Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Mon, 24 Feb 2025 16:53:38 +0100
Subject: [PATCH 10/13] fix: - import error - array broadcasting

---
 p_kit/library/tsp.py   | 2 +-
 p_kit/psl/p_circuit.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/p_kit/library/tsp.py b/p_kit/library/tsp.py
index d257e5f..7b70557 100644
--- a/p_kit/library/tsp.py
+++ b/p_kit/library/tsp.py
@@ -1,7 +1,7 @@
 # Adapted from: https://github.com/anirudhgha/p-bit/blob/master/Research/Travelling%20Salesman%20Problem/tsp.py
 
 import numpy as np
-from p_kit.core import PCircuit
+from p_kit.psl.p_circuit import PCircuit
 
 
 class TSP(PCircuit):
diff --git a/p_kit/psl/p_circuit.py b/p_kit/psl/p_circuit.py
index 120ed2f..4b50783 100644
--- a/p_kit/psl/p_circuit.py
+++ b/p_kit/psl/p_circuit.py
@@ -27,7 +27,7 @@ def __init__(self, n_pbits: int, ports: Dict[str, Any] = None):
         self.n_pbits = n_pbits
         self.ports = ports #Kept for copy behavior
 
-        self.h = np.zeros((n_pbits, 1))
+        self.h = np.zeros((n_pbits,))
         self.J = np.zeros((n_pbits, n_pbits))
         self._connections = {}
 

From e19552ec0a32bfda9518639a75e8669fddf5d51c Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Tue, 25 Feb 2025 11:04:54 +0100
Subject: [PATCH 11/13] fix quotes for version?

---
 p_kit/_version.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/p_kit/_version.py b/p_kit/_version.py
index 3dc1f76..b794fd4 100644
--- a/p_kit/_version.py
+++ b/p_kit/_version.py
@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = '0.1.0'

From 07dc996c69e012aef3d87daee59b627775f28500 Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Tue, 25 Feb 2025 11:24:23 +0100
Subject: [PATCH 12/13] fix setup.py

---
 p_kit/_version.py | 2 +-
 setup.py          | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/p_kit/_version.py b/p_kit/_version.py
index b794fd4..3dc1f76 100644
--- a/p_kit/_version.py
+++ b/p_kit/_version.py
@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = "0.1.0"
diff --git a/setup.py b/setup.py
index 9806b26..7fbb2a4 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
 with open(op.join('p_kit', '_version.py'), 'r') as fid:
     for line in (line.strip() for line in fid):
         if line.startswith('__version__'):
-            version = line.split('=')[1].strip().strip('\'')
+            version = line.split('=')[1].strip().strip('"')
             break
 if version is None:
     raise RuntimeError('Could not determine version')
@@ -29,7 +29,7 @@
       project_urls={
           'Documentation': 'https://github.com/IBM/p-kit/wiki',
           'Source': 'https://github.com/IBM/p-kit',
-          'Tracker': 'https://github.com/IBM/p-kitissues/',
+          'Tracker': 'https://github.com/IBM/p-kit/issues/',
       },
       platforms='any',
       python_requires=">=3.9",

From 5d83b4d09cc6fcab6b7b76bd16eb605aeaeb12fa Mon Sep 17 00:00:00 2001
From: Gregoire Cattan <gregoire.cattan@ibm.com>
Date: Tue, 25 Feb 2025 11:31:45 +0100
Subject: [PATCH 13/13] add back heatmap

---
 p_kit/visualization/__init__.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/p_kit/visualization/__init__.py b/p_kit/visualization/__init__.py
index d9c39c2..969bdbc 100644
--- a/p_kit/visualization/__init__.py
+++ b/p_kit/visualization/__init__.py
@@ -3,8 +3,10 @@
 from .vin_vout import vin_vout
 from .plot3d import plot3d
 from .tsp_graph import visualize_tsp_route
+from .heatmap import heatmap
 
 __all__ = [
+    "heatmap",
     "m_to_string",
     "tsp_hist",
     "histplot",