1 files changed, 51 insertions, 0 deletions
diff --git a/logs/mech.py b/logs/mech.py
new file mode 100644
index 0000000..6762777
--- /dev/null
+++ b/logs/mech.py
@@ -0,0 +1,51 @@
+import sys, numpy as np, ot, torch
+sys.path.insert(0,'/home/yurenh2/emm')
+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")
+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
+st=torch.load("/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt",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)
+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])
+    acc=lambda c: float((hid[c]==np.arange(n)).mean())
+    def refine(c,it=2000):
+        f=fast_pair_descent(Sg,Vg,torch.from_numpy(np.ascontiguousarray(c)).to(DEV),it).cpu().numpy()
+        return float(en.energy(torch.from_numpy(np.ascontiguousarray(f)).to(DEV)[None])[0]), acc(f)
+    M=tlb(V,S); M=M/M.max()
+    _,tlbcols=linear_sum_assignment(M)
+    def vert(cols):
+        P=np.zeros((n,n)); P[np.arange(n),cols]=1.0/n; return P
+    starts={"tlb-vertex":tlbcols,
+            "random-vertex":np.random.default_rng(50+trial).permutation(n),
+            "grampa-vertex":None}
+    from worldalign.spectral_match import grampa
+    starts["grampa-vertex"]=grampa(V,S,1.0)
+    print(f"--- trial {trial}  E_truth={Et:.6f}")
+    for name,cols in starts.items():
+        # (i) pure power step: LSAP on gradient V P S  (anchor-propagation, one round)
+        P=vert(cols); grad=V@P@S
+        _,c1=linear_sum_assignment(-grad)
+        e1,a1=refine(c1)
+        # (ii) 5 rounds of power step
+        c=cols.copy()
+        for _ in range(5):
+            P=vert(c); _,c=linear_sum_assignment(-(V@P@S))
+        e2,a2=refine(c)
+        # (iii) warm FGW at alpha=0.1 then 0.2 from this vertex
+        G=vert(cols)
+        for a in (0.1,0.2,0.5):
+            G=ot.gromov.fused_gromov_wasserstein(M,V,S,w,w,"square_loss",alpha=a,G0=G,max_iter=200,tol_rel=1e-9)
+        _,c3=linear_sum_assignment(-G); e3,a3=refine(c3)
+        print("  %-14s start=%.4f | 1 power step ->%.4f (E=%.4f) | 5 steps ->%.4f (E=%.4f) | warm FGW ->%.4f (E=%.4f)"
+              %(name,acc(cols),a1,e1,a2,e2,a3,e3), flush=True)