(Anthropic) Claude caught a bug, suggested an optimization

Author: WrathfulSpatula

pyqrackising/spin_glass_solver.py

@@ -108,7 +108,8 @@ def pick_gray_seeds(best_theta, thread_count, gray_seed_multiple, G_m, n, is_spi
         energies[s] = compute_cut_diff_between(best_theta, seed, G_m, n)
         seeds[s] = seed
 
-    indices = np.argsort(energies)[::-1]
+    indices = np.argpartition(energies, -thread_count)[-thread_count:]
+    indices = indices[np.argsort(energies[indices])[::-1]]
     best_seeds = np.empty((thread_count, n), dtype=np.bool_)
     best_energies = np.empty(thread_count, dtype=dtype)
     for i in prange(thread_count):
@@ -295,7 +296,7 @@ def run_gray_search_opencl(
     blocks = (n + 63) // 64
     offset = best_x * blocks
     for b in range(blocks):
-        s = max_theta_host[best_x]
+        s = max_theta_host[offset + b]
         b_offset = b << 6
         for i in range(64):
             j = b_offset + i

pyqrackising/spin_glass_solver_sparse.py

@@ -142,7 +142,8 @@ def pick_gray_seeds(best_theta, thread_count, gray_seed_multiple, G_data, G_rows
             energies[s] = compute_cut_sparse(seed, G_data, G_rows, G_cols, n)
             seeds[s] = seed
 
-    indices = np.argsort(energies)[::-1]
+    indices = np.argpartition(energies, -thread_count)[-thread_count:]
+    indices = indices[np.argsort(energies[indices])[::-1]]
     best_seeds = np.empty((thread_count, n), dtype=np.bool_)
     best_energies = np.empty(thread_count, dtype=dtype)
     for i in prange(thread_count):
@@ -362,7 +363,7 @@ def run_gray_search_opencl(
     blocks = (n + 63) // 64
     offset = best_x * blocks
     for b in range(blocks):
-        s = max_theta_host[best_x]
+        s = max_theta_host[offset + b]
         b_offset = b << 6
         for i in range(64):
             j = b_offset + i

pyqrackising/spin_glass_solver_streaming.py

@@ -99,7 +99,8 @@ def pick_gray_seeds(best_theta, thread_count, gray_seed_multiple, G_func, nodes,
         energies[s] = compute_cut_diff_between_streaming(best_theta, seed, G_func, nodes, n)
         seeds[s] = seed
 
-    indices = np.argsort(energies)[::-1]
+    indices = np.argpartition(energies, -thread_count)[-thread_count:]
+    indices = indices[np.argsort(energies[indices])[::-1]]
     best_seeds = np.empty((thread_count, n), dtype=np.bool_)
     best_energies = np.empty(thread_count, dtype=dtype)
     for i in prange(thread_count):

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name="pyqrackising",
-    version="9.11.2",
+    version="9.11.3",
     author="Dan Strano",
     author_email="stranoj@gmail.com",
     description="Fast MAXCUT, TSP, and sampling heuristics from near-ideal transverse field Ising model (TFIM)",