Merge pull request #3 from vm6502q/belief_propagation

Belief propagation (by Anthropic Claude)

Author: WrathfulSpatula

README.md

@@ -73,10 +73,23 @@ def generate_spin_glass_graph(n_nodes=16, degree=3, seed=None):
 
 
 G = generate_spin_glass_graph(n_nodes=64, seed=42)
-solution_bit_string, cut_value, node_groups, energy = spin_glass_solver(G, quality=6, shots=None, anneal_t=8.0, anneal_h=8.0, is_spin_glass=True, best_guess=None, is_maxcut_gpu=True, gray_iterations=None, gray_seed_multiple=None)
+solution_bit_string, cut_value, node_groups, energy = spin_glass_solver(
+    G,
+    quality=6,
+    shots=None,
+    anneal_t=8.0,
+    anneal_h=8.0,
+    is_spin_glass=True,
+    best_guess=None,
+    is_maxcut_gpu=True,
+    gray_iterations=None,
+    gray_seed_multiple=None,
+    bp_scale=None,
+    bp_damping=0.5
+)
 # solution_bit_string, cut_value, node_groups, energy = spin_glass_solver(G, best_guess=maxcut_tfim(G, quality=6)[0])
 ```
-We also provide `spin_glass_solver_sparse(G)` and `spin_glass_solver_streaming(G_func, nodes)`. `best_guess` gives the option to seed the algorithm with a best guess as to the maximal cut (as an integer, binary string, or list of booleans). By default, `spin_glass_solver()` uses `maxcut_tfim(G)` with passed-through `quality` as `best_guess`, which typically works well, but it could be seeded with higher `maxcut_tfim()` `quality` or Goemans-Williamson, for example. `is_spin_glass` controls whether the solver optimizes for cut value or spin-glass energy. This function is designed with a sign convention for weights such that it can immediately be used as a MAXCUT solver itself: you might need to reverse the sign convention on your weights for spin glass graphs, but this is only convention. `gray_iterations` gives manual control over how many iterations are carried out of a parallel Gray-code search on `best_guess`. `gray_seed_multiple` controls how many parallel search seeds (as a multiple of your CPU thread count) are tested for the best parallel seeds, and a value of `1` will perfectly cover the search space without collision if your node count is a power of 2.
+We also provide `spin_glass_solver_sparse(G)` and `spin_glass_solver_streaming(G_func, nodes)`. `best_guess` gives the option to seed the algorithm with a best guess as to the maximal cut (as an integer, binary string, or list of booleans). By default, `spin_glass_solver()` uses `maxcut_tfim(G)` with passed-through `quality` as `best_guess`, which typically works well, but it could be seeded with higher `maxcut_tfim()` `quality` or Goemans-Williamson, for example. `is_spin_glass` controls whether the solver optimizes for cut value or spin-glass energy. This function is designed with a sign convention for weights such that it can immediately be used as a MAXCUT solver itself: you might need to reverse the sign convention on your weights for spin glass graphs, but this is only convention. `gray_iterations` gives manual control over how many iterations are carried out of a parallel Gray-code search on `best_guess`. `gray_seed_multiple` controls how many parallel search seeds (as a multiple of your CPU thread count) are tested for the best parallel seeds, and a value of `1` will perfectly cover the search space without collision if your node count is a power of 2. `bp_scale` controls opt-in and time complexity scaling for _belief propagation_ as a competing "warm start," with `None` disabling this option completely, `1.0` representing a budget of exactly `O(n^3)` time complexity, and lesser or greater values representing a linear multiplier on that budget. `bp_damping` similarly controls belief propagation damping.
 
 These solvers above can be used as a "warm state" for _MaxSAT-based_ solvers or _branch-and-bound_ solvers, potentially to reach _exact_ solutions to MAXCUT:
 

pyqrackising/spin_glass_solver.py

@@ -307,6 +307,120 @@ def run_gray_search_opencl(
     return energy, theta
 
 
+# Belief propagation was added by (Anthropic) Claude
+def belief_propagation_marginals_dense(G_m, n, bp_scale=1.0, damping=0.5):
+    """
+    Run loopy belief propagation on the dense (upper-triangular) adjacency
+    matrix to produce per-node marginals for warm-starting the partition.
+
+    The dense solver stores G_m as a full n×n matrix; we read only the
+    upper triangle (j > i) to match the sparse convention, so behaviour
+    is consistent across all three solver variants.
+
+    Parameters
+    ----------
+    G_m : ndarray, shape (n, n)
+        Dense adjacency/weight matrix. May be full or upper-triangular.
+    n : int
+        Number of nodes.
+    bp_scale : float
+        Iteration cap = int(bp_scale * n). Gives O(n * m) total work,
+        <= O(n^3) for any graph density. User-controlled.
+    damping : float
+        Message damping factor in [0, 1). Stabilises oscillating marginals
+        on dense or heavily frustrated graphs.
+
+    Returns
+    -------
+    marginals : ndarray, shape (n,)
+        Soft assignment in (-1, +1). Positive values favour partition 1,
+        negative values favour partition 0.
+    """
+    max_iterations = max(1, int(bp_scale * n))
+
+    # Collect undirected edges from upper triangle
+    us = []
+    vs = []
+    ws = []
+    for i in range(n):
+        for j in range(i + 1, n):
+            w = G_m[i, j]
+            if abs(w) > 1e-12:
+                us.append(i)
+                vs.append(j)
+                ws.append(w)
+
+    m = len(us)
+    if m == 0:
+        return np.zeros(n, dtype=np.float64)
+
+    us = np.array(us, dtype=np.int32)
+    vs = np.array(vs, dtype=np.int32)
+    ws = np.array(ws, dtype=np.float64)
+
+    # msg_fwd[k] = message from us[k] -> vs[k]
+    # msg_bwd[k] = message from vs[k] -> us[k]
+    msg_fwd = np.zeros(m, dtype=np.float64)
+    msg_bwd = np.zeros(m, dtype=np.float64)
+
+    # Build neighbour index: node -> list of (edge_idx, direction)
+    # direction=0: node is 'u' end, incoming = msg_bwd[k]
+    # direction=1: node is 'v' end, incoming = msg_fwd[k]
+    node_edges = [[] for _ in range(n)]
+    for k in range(m):
+        node_edges[us[k]].append((k, 0))
+        node_edges[vs[k]].append((k, 1))
+
+    for _ in range(max_iterations):
+        new_fwd = np.empty(m, dtype=np.float64)
+        new_bwd = np.empty(m, dtype=np.float64)
+
+        for k in range(m):
+            u = us[k]
+            v = vs[k]
+            w = ws[k]
+
+            # Message u -> v: aggregate incoming to u excluding v
+            h_u = 0.0
+            for (ek, direction) in node_edges[u]:
+                if ek == k:
+                    continue
+                h_u += msg_bwd[ek] if direction == 0 else msg_fwd[ek]
+
+            tanh_w = np.tanh(w)
+            product = tanh_w * np.tanh(h_u)
+            product = np.clip(product, -1.0 + 1e-9, 1.0 - 1e-9)
+            new_fwd[k] = damping * msg_fwd[k] + (1.0 - damping) * np.arctanh(product)
+
+            # Message v -> u: aggregate incoming to v excluding u
+            h_v = 0.0
+            for (ek, direction) in node_edges[v]:
+                if ek == k:
+                    continue
+                h_v += msg_bwd[ek] if direction == 0 else msg_fwd[ek]
+
+            product_v = tanh_w * np.tanh(h_v)
+            product_v = np.clip(product_v, -1.0 + 1e-9, 1.0 - 1e-9)
+            new_bwd[k] = damping * msg_bwd[k] + (1.0 - damping) * np.arctanh(product_v)
+
+        msg_fwd = new_fwd
+        msg_bwd = new_bwd
+
+    # Marginals: sum all incoming messages per node
+    marginals = np.zeros(n, dtype=np.float64)
+    for k in range(m):
+        marginals[us[k]] += msg_bwd[k]
+        marginals[vs[k]] += msg_fwd[k]
+
+    return np.tanh(marginals)
+
+
+def bp_warm_start_dense(G_m, n, bp_scale=1.0, damping=0.5):
+    """Partition bitstring from BP marginals over the dense adjacency matrix."""
+    marginals = belief_propagation_marginals_dense(G_m, n, bp_scale, damping)
+    return marginals > 0.0
+
+
 def spin_glass_solver(
     G,
     quality=None,
@@ -320,7 +434,23 @@ def spin_glass_solver(
     is_log=False,
     gray_iterations=None,
     gray_seed_multiple=None,
+    bp_scale=None,
+    bp_damping=0.5,
 ):
+    """
+    Dense spin glass / MaxCut solver with optional belief propagation warm-start.
+
+    Parameters
+    ----------
+    bp_scale : float or None
+        Controls BP iteration count as int(bp_scale * n_qubits).
+        Default None disables BP. Set to 1.0 to enable with O(n*m) cost
+        (<=O(n^3) for any graph density). Increase for more BP iterations
+        at proportionally higher cost.
+    bp_damping : float
+        Message damping in [0, 1). Default 0.5. Higher values help
+        convergence on dense or heavily frustrated graphs.
+    """
     dtype = opencl_context.dtype
     nodes = None
     n_qubits = 0
@@ -360,17 +490,49 @@ def spin_glass_solver(
     elif isinstance(best_guess, list):
         bitstring = "".join(["1" if b else "0" for b in best_guess])
     else:
-        bitstring, cut_value, _ = maxcut_tfim(
-            G_m,
-            quality=quality,
-            shots=shots,
-            is_spin_glass=is_spin_glass,
-            anneal_t=anneal_t,
-            anneal_h=anneal_h,
-            repulsion_base=repulsion_base,
-            is_maxcut_gpu=is_maxcut_gpu,
-            is_nested=True,
-        )
+        # BP warm-start: run before the sampling heuristic if bp_scale is set,
+        # so that the cascade local search begins from a sign-aware partition
+        # rather than the TFIM prior alone.
+        if bp_scale is not None and n_qubits >= heuristic_threshold:
+            bp_theta = bp_warm_start_dense(G_m, n_qubits, bp_scale, bp_damping)
+            bitstring, cut_value, _ = maxcut_tfim(
+                G_m,
+                quality=quality,
+                shots=shots,
+                is_spin_glass=is_spin_glass,
+                anneal_t=anneal_t,
+                anneal_h=anneal_h,
+                repulsion_base=repulsion_base,
+                is_maxcut_gpu=is_maxcut_gpu,
+                is_nested=True,
+            )
+            # Keep whichever of BP or sampling gave the better cut
+            bp_energy = (
+                compute_energy(bp_theta, G_m, n_qubits)
+                if is_spin_glass
+                else compute_cut(bp_theta, G_m, n_qubits)
+            )
+            sample_theta = np.array([b == "1" for b in list(bitstring)], dtype=np.bool_)
+            sample_energy = (
+                compute_energy(sample_theta, G_m, n_qubits)
+                if is_spin_glass
+                else compute_cut(sample_theta, G_m, n_qubits)
+            )
+            if bp_energy > sample_energy:
+                bitstring = "".join(["1" if b else "0" for b in bp_theta])
+                cut_value = bp_energy if not is_spin_glass else None
+        else:
+            bitstring, cut_value, _ = maxcut_tfim(
+                G_m,
+                quality=quality,
+                shots=shots,
+                is_spin_glass=is_spin_glass,
+                anneal_t=anneal_t,
+                anneal_h=anneal_h,
+                repulsion_base=repulsion_base,
+                is_maxcut_gpu=is_maxcut_gpu,
+                is_nested=True,
+            )
 
     best_theta = np.array([b == "1" for b in list(bitstring)], dtype=np.bool_)
     if is_spin_glass: