1 files changed, 172 insertions, 0 deletions
diff --git a/worldalign/spectral_match.py b/worldalign/spectral_match.py
new file mode 100644
index 0000000..9e1a0fe
--- /dev/null
+++ b/worldalign/spectral_match.py
@@ -0,0 +1,172 @@
+"""Spectral matching of relation fields: Umeyama and GRAMPA.
+
+Relation fields are N x N regardless of embedding dimension, so the two
+modalities need no common representation dimension. What they do need is
+comparable spectra: eigenvector matching degrades when effective ranks
+differ or when eigenvalues cluster. GRAMPA is built for that regime -- it
+weights every pair of eigenvectors by 1 / ((lambda_i - mu_j)^2 + eta^2)
+instead of pairing them one to one -- so both solvers are provided along
+with the spectral compatibility diagnostic that predicts whether either
+can work.
+
+No local search: these are polynomial-time solvers that sidestep the
+glassy landscape entirely. Hidden pairs score the output only.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from scipy.optimize import linear_sum_assignment
+
+from .common import read_json, seed_everything, write_json
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fields", help="Saved .pt with visual/text fields.")
+    parser.add_argument("--data-dir", default="artifacts/synth_v0")
+    parser.add_argument("--split", choices=["val", "test"], default="test")
+    parser.add_argument("--samples", type=int, default=256)
+    parser.add_argument("--merge-distance", type=float, default=30.0)
+    parser.add_argument("--vision-views", type=int, default=4)
+    parser.add_argument("--eta", default="0.05,0.1,0.2,0.5")
+    parser.add_argument("--seed", type=int, default=20260731)
+    parser.add_argument(
+        "--output", default="artifacts/synth_v0/spectral_match.json"
+    )
+    parser.add_argument("--fields-output", default="")
+    return parser.parse_args()
+
+
+def spectral_profile(matrix: np.ndarray) -> dict:
+    values = np.linalg.eigvalsh(matrix)[::-1]
+    magnitude = np.abs(values)
+    total = magnitude.sum()
+    cumulative = np.cumsum(magnitude) / max(total, 1e-12)
+    participation = (magnitude.sum() ** 2) / max((magnitude**2).sum(), 1e-12)
+    gaps = np.abs(np.diff(values))
+    return {
+        "top_eigenvalues": values[:12].tolist(),
+        "effective_rank_participation": float(participation),
+        "rank_for_90_percent": int(np.searchsorted(cumulative, 0.90) + 1),
+        "rank_for_99_percent": int(np.searchsorted(cumulative, 0.99) + 1),
+        "median_relative_gap": float(
+            np.median(gaps) / max(magnitude.max(), 1e-12)
+        ),
+        "min_relative_gap_top20": float(
+            gaps[:20].min() / max(magnitude.max(), 1e-12)
+        ),
+    }
+
+
+def umeyama(visual: np.ndarray, text: np.ndarray) -> np.ndarray:
+    """Classic eigenvector-magnitude matching, sign ambiguity absorbed."""
+    _, u = np.linalg.eigh(visual)
+    _, v = np.linalg.eigh(text)
+    score = np.abs(u) @ np.abs(v).T
+    rows, cols = linear_sum_assignment(-score)
+    return cols
+
+
+def grampa(visual: np.ndarray, text: np.ndarray, eta: float) -> np.ndarray:
+    """Pairwise eigen-alignment similarity, robust to clustered spectra."""
+    lam, u = np.linalg.eigh(visual)
+    mu, v = np.linalg.eigh(text)
+    ones = np.ones(len(visual))
+    left = u.T @ ones  # [N]
+    right = v.T @ ones
+    weight = np.outer(left, right) / ((lam[:, None] - mu[None, :]) ** 2 + eta**2)
+    similarity = u @ weight @ v.T
+    rows, cols = linear_sum_assignment(-similarity)
+    return cols
+
+
+def score_assignment(
+    assignment: np.ndarray, hidden: np.ndarray, size: int
+) -> dict:
+    """assignment[i] is the shuffled-space index matched to visual row i.
+
+    Shuffled index j denotes original node hidden[j], and visual row i
+    denotes original node i, so the match is correct when
+    hidden[assignment[i]] == i.
+    """
+    return {
+        "accuracy": float((hidden[assignment] == np.arange(size)).mean()),
+        "chance": 1.0 / size,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    if args.fields:
+        state = torch.load(args.fields, map_location="cpu", weights_only=False)
+        visual_field = state["visual_field"]
+        text_field = state["text_field"]
+    else:
+        from .synth_triangle_gate import build_fields
+
+        manifest = read_json(Path(args.data_dir, "manifest.json"))
+        rows = manifest[args.split][: args.samples]
+        visual_field, text_field = build_fields(args, rows, manifest)
+        if args.fields_output:
+            torch.save(
+                {"visual_field": visual_field, "text_field": text_field},
+                args.fields_output,
+            )
+
+    size = len(visual_field)
+    visual = visual_field.double().numpy()
+    text = text_field.double().numpy()
+    # Center and scale: spectral methods compare shapes, not offsets.
+    mask = ~np.eye(size, dtype=bool)
+    for matrix in (visual, text):
+        values = matrix[mask]
+        matrix -= values.mean()
+        matrix /= values.std()
+        np.fill_diagonal(matrix, 0.0)
+
+    generator = np.random.default_rng(args.seed)
+    hidden = generator.permutation(size)
+    text_shuffled = text[np.ix_(hidden, hidden)]
+
+    report = {
+        "protocol": (
+            "Polynomial-time spectral solvers on N x N relation fields; "
+            "embedding dimensions are irrelevant by construction. The "
+            "hidden shuffle is applied to the text field and used only to "
+            "score the returned assignment."
+        ),
+        "samples": size,
+        "spectra": {
+            "visual": spectral_profile(visual),
+            "text": spectral_profile(text_shuffled),
+        },
+        "solvers": {},
+    }
+    # Both fields are built over the same row list, so they are aligned at
+    # the truth without any permutation.
+    correlation = float(np.corrcoef(visual[mask], text[mask])[0, 1])
+    report["field_correlation_at_truth"] = correlation
+
+    assignment = umeyama(visual, text_shuffled)
+    report["solvers"]["umeyama"] = score_assignment(assignment, hidden, size)
+    print(json.dumps({"umeyama": report["solvers"]["umeyama"]}))
+
+    for eta in (float(value) for value in args.eta.split(",")):
+        assignment = grampa(visual, text_shuffled, eta)
+        result = score_assignment(assignment, hidden, size)
+        report["solvers"][f"grampa_eta{eta}"] = result
+        print(json.dumps({f"grampa_eta{eta}": result}))
+
+    write_json(args.output, report)
+    print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()