diff options
Diffstat (limited to 'logs/control_probe.py')
| -rw-r--r-- | logs/control_probe.py | 93 |
1 files changed, 93 insertions, 0 deletions
diff --git a/logs/control_probe.py b/logs/control_probe.py new file mode 100644 index 0000000..170d731 --- /dev/null +++ b/logs/control_probe.py @@ -0,0 +1,93 @@ +"""Control: is the FGW warm-homotopy win driven by TLB *information*, or by any +linear term at all (generic symmetry-breaking)? Three linear costs: + + tlb -- the proposed distance-profile cost + tlbperm -- the same matrix with its columns randomly permuted: identical + value distribution, correspondence information destroyed + rand -- iid uniform noise cost + +Plus 3 trials, cold and warm sweeps, blind selection by E(P). +""" +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.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 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 main(): + path = "/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt" + 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) + rows = [] + for trial in range(3): + 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) + et = float(en.energy(torch.from_numpy( + np.ascontiguousarray(np.argsort(hid))).to(DEV)[None])[0]) + + 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) + + base = tlb(V, S) + shuf = np.random.default_rng(100 + trial).permutation(n) + variants = { + "tlb": base / base.max(), + "tlbperm": base[:, shuf] / base.max(), + "rand": np.random.default_rng(200 + trial).random((n, n)), + } + for vname, M in variants.items(): + for mode in ("cold", "warm"): + G = np.outer(w, w); recs = [] + for a in ALPHAS: + 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) + recs.append({"alpha": a, "raw": acc(cols), "refined": ar, "E": e}) + best = min(recs, key=lambda x: x["E"]) + row = {"trial": trial, "cost": vname, "mode": mode, + "E_truth": et, "blind_refined": best["refined"], + "blind_alpha": best["alpha"], "blind_E": best["E"], + "oracle_refined": max(x["refined"] for x in recs), + "curve": [round(x["refined"], 4) for x in recs]} + rows.append(row) + print(json.dumps(row), flush=True) + scratch = ("/tmp/claude-1273071/-home-yurenh2-emm/" + "bb97201d-d9b5-46e6-82fc-981d37e7ab98/scratchpad/") + json.dump(rows, open(scratch + "control.json", "w")) + print(json.dumps({"done": True})) + + +if __name__ == "__main__": + main() |