1 files changed, 138 insertions, 0 deletions
diff --git a/worldalign/world_match.py b/worldalign/world_match.py
new file mode 100644
index 0000000..e719a9f
--- /dev/null
+++ b/worldalign/world_match.py
@@ -0,0 +1,138 @@
+"""World matching: spectral initialisation plus energy refinement.
+
+The two solvers are complementary. A spectral solver reads the coarse
+correspondence out of the eigenstructure of two relation fields in
+polynomial time and without any search, but its rounding is noisy. Exact
+steepest descent on the closed-form energy repairs the rounding but
+cannot find the basin from a random start. Composed, they recover the
+hidden assignment.
+
+Neither stage sees a pair. The relation fields are built independently
+per modality; the hidden permutation is applied to the text field and
+read back only to score.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .common import seed_everything, write_json
+from .spectral_match import grampa, spectral_profile, umeyama
+from .synth_fast_gate import ClosedFormEnergy, all_swaps, steepest_descent
+from .synth_triangle_gate import standardized
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fields", required=True)
+    parser.add_argument("--eta", type=float, default=1.0)
+    parser.add_argument("--pair-weight", type=float, default=1.0)
+    parser.add_argument("--triangle-weight", type=float, default=1.0)
+    parser.add_argument("--refine-steps", type=int, default=4000)
+    parser.add_argument("--chunk", type=int, default=256)
+    parser.add_argument("--restarts", type=int, default=1)
+    parser.add_argument("--device", default="cuda:3")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output", default="")
+    return parser.parse_args()
+
+
+def normalise(matrix: np.ndarray) -> np.ndarray:
+    size = len(matrix)
+    mask = ~np.eye(size, dtype=bool)
+    values = matrix[mask]
+    out = (matrix - values.mean()) / values.std()
+    np.fill_diagonal(out, 0.0)
+    return out
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    state = torch.load(args.fields, map_location="cpu", weights_only=False)
+    visual = normalise(state["visual_field"].double().numpy().copy())
+    text = normalise(state["text_field"].double().numpy().copy())
+    size = len(visual)
+
+    generator = np.random.default_rng(args.seed)
+    hidden = generator.permutation(size)
+    shuffled = text[np.ix_(hidden, hidden)]
+    truth = np.arange(size)
+
+    def accuracy(assignment: np.ndarray) -> float:
+        return float((hidden[assignment] == truth).mean())
+
+    device = torch.device(args.device)
+    text_tensor = standardized(torch.from_numpy(shuffled).to(device)).float()
+    visual_tensor = standardized(torch.from_numpy(visual).to(device)).float()
+    energy = ClosedFormEnergy(
+        text_tensor,
+        visual_tensor,
+        args.pair_weight,
+        args.triangle_weight,
+        args.chunk,
+    )
+    swaps = all_swaps(size, device)
+    truth_permutation = torch.from_numpy(np.argsort(hidden)).to(device)
+    true_energy = float(energy.energy(truth_permutation[None])[0])
+
+    report = {
+        "protocol": (
+            "Fields are built per modality without pairs; the hidden "
+            "permutation is applied to the text field and used only to "
+            "score. Spectral initialisation is followed by exact "
+            "steepest descent on the closed-form energy."
+        ),
+        "samples": size,
+        "field_correlation_at_truth": float(
+            np.corrcoef(
+                visual[~np.eye(size, dtype=bool)],
+                text[~np.eye(size, dtype=bool)],
+            )[0, 1]
+        ),
+        "true_energy": true_energy,
+        "chance": 1.0 / size,
+        "spectra": {
+            "visual": spectral_profile(visual),
+            "text": spectral_profile(shuffled),
+        },
+        "stages": {},
+    }
+
+    spectral = grampa(visual, shuffled, args.eta)
+    report["stages"]["spectral"] = {"accuracy": accuracy(spectral)}
+    initial = torch.from_numpy(spectral.copy()).to(device)
+    refined, value = steepest_descent(energy, initial, swaps, args.refine_steps)
+    report["stages"]["spectral_plus_refinement"] = {
+        "accuracy": accuracy(refined.cpu().numpy()),
+        "energy": value,
+        "energy_over_true": value / abs(true_energy) - true_energy / abs(true_energy),
+    }
+
+    quenches = []
+    for restart in range(args.restarts):
+        start = torch.from_numpy(generator.permutation(size)).to(device)
+        final, final_value = steepest_descent(
+            energy, start, swaps, args.refine_steps
+        )
+        quenches.append(
+            {
+                "accuracy": accuracy(final.cpu().numpy()),
+                "energy": final_value,
+            }
+        )
+    report["stages"]["random_start_refinement"] = quenches
+
+    print(json.dumps({key: value for key, value in report["stages"].items()}))
+    if args.output:
+        write_json(args.output, report)
+        print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()