1 files changed, 84 insertions, 0 deletions
diff --git a/logs/lean_probe.py b/logs/lean_probe.py
new file mode 100644
index 0000000..f3b394c
--- /dev/null
+++ b/logs/lean_probe.py
@@ -0,0 +1,84 @@
+"""Lean version: 15-alpha grid, TLB cost, cold vs warm, omit-size then synth-full."""
+from __future__ import annotations
+import json, sys, time
+sys.path.insert(0, "/home/yurenh2/emm")
+import numpy as np, ot, torch
+from scipy.optimize import linear_sum_assignment
+from worldalign.synth_fast_gate import ClosedFormEnergy, fast_pair_descent
+from worldalign.synth_triangle_gate import standardized
+
+DEV = torch.device("cuda:1")
+ALPHAS = [0.0, 0.2, 0.4, 0.6, 0.75, 0.85, 0.88, 0.90, 0.92, 0.94, 0.95,
+          0.96, 0.97, 0.98, 0.99, 1.0]
+
+
+def standardise(m):
+    mask = ~np.eye(len(m), dtype=bool); v = m[mask]
+    o = (m - v.mean()) / v.std(); np.fill_diagonal(o, 0.0); return o
+
+
+def tlb(V, S):
+    n = len(V); a = np.sort(V, 1); b = np.sort(S, 1)
+    return (a * a).sum(1)[:, None] / n + (b * b).sum(1)[None, :] / n - 2 * (a @ b.T) / n
+
+
+def run(path, label, trial=0):
+    st = torch.load(path, map_location="cpu", weights_only=False)
+    V = standardise(st["visual_field"].double().numpy())
+    T = standardise(st["text_field"].double().numpy())
+    n = len(V); w = ot.unif(n)
+    rng = np.random.default_rng(trial); hid = rng.permutation(n)
+    S = T[np.ix_(hid, hid)]
+    Vg = standardized(torch.from_numpy(V).to(DEV)).double()
+    Sg = standardized(torch.from_numpy(S).to(DEV)).double()
+    en = ClosedFormEnergy(Sg, Vg, 1.0, 0.0, 64)
+    M = tlb(V, S); M = M / M.max()
+
+    def acc(c): return float((hid[c] == np.arange(n)).mean())
+
+    def refine(c):
+        f = fast_pair_descent(Sg, Vg, torch.from_numpy(np.ascontiguousarray(c)).to(DEV),
+                              600).cpu().numpy()
+        e = float(en.energy(torch.from_numpy(np.ascontiguousarray(f)).to(DEV)[None])[0])
+        return e, acc(f)
+
+    et = float(en.energy(torch.from_numpy(np.ascontiguousarray(np.argsort(hid))).to(DEV)[None])[0])
+    print(f"# {label} trial{trial}  E_truth={et:.4f}", flush=True)
+
+    out = {"label": label, "E_truth": et, "cold": [], "warm": []}
+    for mode in ("cold", "warm"):
+        G = np.outer(w, w)
+        prev = None
+        for a in ALPHAS:
+            t0 = time.time()
+            g0 = None if mode == "cold" else G
+            G2 = ot.gromov.fused_gromov_wasserstein(M, V, S, w, w, "square_loss",
+                                                    alpha=a, G0=g0, max_iter=200,
+                                                    tol_rel=1e-9)
+            if mode == "warm":
+                G = G2
+            _, cols = linear_sum_assignment(-G2)
+            e, ar = refine(cols)
+            same = (prev is not None and np.array_equal(cols, prev))
+            prev = cols
+            rec = {"alpha": a, "raw": acc(cols), "refined": ar, "E": e,
+                   "sec": round(time.time() - t0, 1), "same_as_prev": bool(same)}
+            out[mode].append(rec)
+            print(f"  {mode} a={a:<5} raw={rec['raw']:.4f} ref={ar:.4f} E={e:.4f}"
+                  f" {'FROZEN' if same else ''} ({rec['sec']}s)", flush=True)
+        good = out[mode]
+        best = min(good, key=lambda x: x["E"])
+        print(f"  == {mode}: blind-by-E refined={best['refined']:.4f} (alpha={best['alpha']}),"
+              f" oracle-best={max(x['refined'] for x in good):.4f}", flush=True)
+    return out
+
+
+if __name__ == "__main__":
+    base = "/home/yurenh2/emm/artifacts/synth_v1/"
+    res = []
+    for p, l in [("omit_size.pt", "omit-size"), ("omit_none.pt", "synth-full")]:
+        res.append(run(base + p, l))
+    scratch = ("/tmp/claude-1273071/-home-yurenh2-emm/"
+               "bb97201d-d9b5-46e6-82fc-981d37e7ab98/scratchpad/")
+    json.dump(res, open(scratch + "lean.json", "w"))
+    print(json.dumps({"done": True}))