Spectral Basin Allocation for Multimodal Neural Memories / run_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import numpy as np
  3from pathlib import Path
  4
  5SEED = 3036
  6rng = np.random.default_rng(SEED)
  7N = 12
  8K = 1.0
  9DT = 0.08
 10STEPS = 250
 11QMAX = 3
 12
 13
 14def ring_adjacency(w):
 15    """Undirected ring with coupling w[d-1] at cyclic distance d=1,2."""
 16    A = np.zeros((N, N))
 17    for i in range(N):
 18        for d, wd in enumerate(w, 1):
 19            for j in ((i + d) % N, (i - d) % N):
 20                A[i, j] = wd
 21    return A
 22
 23
 24def prototype(q):
 25    return 2 * np.pi * q * np.arange(N) / N
 26
 27
 28def spectral_rate(A, alpha, psi):
 29    C = A * np.cos(psi[None, :] - psi[:, None] - alpha)
 30    L = np.diag(C.sum(axis=1)) - C
 31    ev = np.linalg.eigvals(L)
 32    # For a connected symmetric graph there is one gauge eigenvalue at zero.
 33    non_gauge = np.sort(np.real(ev[np.abs(ev) > 1e-8]))
 34    return float(K * non_gauge[0]), L, ev
 35
 36
 37def integrate_batch(theta, A, alpha):
 38    theta = theta.copy()
 39    for _ in range(STEPS):
 40        diff = theta[:, None, :] - theta[:, :, None]  # j-i, indexed [batch,i,j] after transpose below
 41        # Direct broadcasting: theta[b,j]-theta[b,i]-alpha[i,j]
 42        diff = theta[:, None, :] - theta[:, :, None] - alpha[None, :, :]
 43        theta += DT * K * np.sum(A[None, :, :] * np.sin(diff), axis=2)
 44        theta = (theta + np.pi) % (2 * np.pi) - np.pi
 45    return theta
 46
 47
 48def classify(theta):
 49    """Gauge-invariant cosine similarity to q-state prototypes."""
 50    z = np.exp(1j * theta)
 51    scores = []
 52    for q in range(QMAX + 1):
 53        p = np.exp(1j * prototype(q))
 54        # maximize over global phase by absolute complex overlap
 55        scores.append(np.abs(np.sum(z * np.conj(p), axis=1)) / N)
 56    scores = np.stack(scores, axis=1)
 57    return np.argmax(scores, axis=1), np.max(scores, axis=1)
 58
 59
 60def basin_fractions(w, samples=3000):
 61    A = ring_adjacency(w)
 62    alpha = np.zeros((N, N))
 63    theta0 = rng.uniform(-np.pi, np.pi, size=(samples, N))
 64    final = integrate_batch(theta0, A, alpha)
 65    labels, confidence = classify(final)
 66    counts = np.bincount(labels, minlength=QMAX + 1)
 67    return (counts / samples).tolist(), float(np.mean(confidence)), labels
 68
 69
 70def math_check():
 71    w = np.array([1.0, 0.55])
 72    A = ring_adjacency(w)
 73    alpha = np.zeros((N, N))
 74    psi = prototype(1)
 75    r, L, ev = spectral_rate(A, alpha, psi)
 76    # Jacobian of phase perturbations is -K L; measure decay in a non-gauge eigenmode.
 77    vals, vecs = np.linalg.eigh(L)
 78    ix = np.where(vals > 1e-8)[0][0]
 79    u = vecs[:, ix]
 80    eps = 1e-4
 81    theta = psi + eps * u
 82    decay = []
 83    for _ in range(80):
 84        f = K * np.sum(A * np.sin(theta[None, :] - theta[:, None]), axis=1)
 85        theta = theta + 0.01 * f
 86        err = theta - psi
 87        err -= err.mean()
 88        decay.append(np.linalg.norm(err))
 89    fitted = np.polyfit(np.arange(20, 75) * 0.01, np.log(np.maximum(decay[20:75], 1e-30)), 1)[0]
 90    return {"r_formula": r, "r_jacobian": float(K * vals[ix]),
 91            "measured_log_slope": float(fitted),
 92            "relative_rate_error": float(abs(r - K * vals[ix]) / max(abs(r), 1e-12))}
 93
 94
 95def main():
 96    check = math_check()
 97    # The q=1 phase pattern benefits from short-range coupling because its
 98    # edge phase difference is smaller there. Keep total coupling fixed.
 99    baseline = np.array([0.50, 0.50])
100    shaped = np.array([0.90, 0.10])
101    alpha = np.zeros((N, N))
102    rates = {}
103    for name, w in [("baseline", baseline), ("spectral_shaped", shaped)]:
104        A = ring_adjacency(w)
105        rates[name] = [spectral_rate(A, alpha, prototype(q))[0] for q in range(QMAX + 1)]
106    b0, c0, _ = basin_fractions(baseline)
107    b1, c1, _ = basin_fractions(shaped)
108    # Across independent graphs, test the claimed positive rate/basin relation.
109    corrs = []
110    graph_rows = []
111    for g in range(20):
112        # random positive distance weights, normalized to the same total
113        w = rng.dirichlet([2.0, 2.0])
114        A = ring_adjacency(w)
115        rr = np.array([spectral_rate(A, alpha, prototype(q))[0] for q in range(QMAX + 1)])
116        bf, conf, _ = basin_fractions(w, samples=700)
117        valid = np.isfinite(rr) & (np.array(bf) > 0)
118        if valid.sum() > 1:
119            corrs.append(float(np.corrcoef(rr[valid], np.array(bf)[valid])[0, 1]))
120        graph_rows.append({"weights": w.tolist(), "rates": rr.tolist(), "basins": bf})
121    result = {"seed": SEED, "N": N, "steps": STEPS, "dt": DT,
122              "math_check": check,
123              "baseline": {"weights": baseline.tolist(), "rates": rates["baseline"], "basins": b0, "mean_similarity": c0},
124              "idea": {"weights": shaped.tolist(), "rates": rates["spectral_shaped"], "basins": b1, "mean_similarity": c1},
125              "target_q": 1, "cross_graph_rate_basin_correlations": corrs,
126              "cross_graph_mean_correlation": float(np.mean(corrs)),
127              "cross_graph_median_correlation": float(np.median(corrs)),
128              "graphs": graph_rows}
129    Path("results.json").write_text(json.dumps(result, indent=2))
130    print(json.dumps({k: result[k] for k in ["math_check", "baseline", "idea", "cross_graph_mean_correlation", "cross_graph_median_correlation"]}, indent=2))
131
132if __name__ == "__main__":
133    main()