1 files changed, 76 insertions, 0 deletions
diff --git a/worldalign/diagnose.py b/worldalign/diagnose.py
new file mode 100644
index 0000000..2d05fd9
--- /dev/null
+++ b/worldalign/diagnose.py
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from scipy.stats import spearmanr
+
+from .common import linear_cka, normalized, read_json, write_json
+from .io import load_feature_pair, select_rows
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--manifest", default="artifacts/manifest.json")
+    p.add_argument("--vision", default="artifacts/vision.pt")
+    p.add_argument("--text", default="artifacts/text.pt")
+    p.add_argument("--split", choices=["val", "test"], default="val")
+    p.add_argument("--max-samples", type=int, default=1_000)
+    p.add_argument("--permutations", type=int, default=20)
+    p.add_argument("--seed", type=int, default=20260728)
+    p.add_argument("--output", default="artifacts/diagnostics.json")
+    return p.parse_args()
+
+
+def upper_triangle(x: torch.Tensor) -> np.ndarray:
+    n = x.shape[0]
+    i, j = torch.triu_indices(n, n, offset=1)
+    return x[i, j].cpu().numpy()
+
+
+def main() -> None:
+    args = parse_args()
+    manifest = read_json(args.manifest)
+    vision, text, vlookup, tlookup = load_feature_pair(args.vision, args.text)
+    rows = manifest[args.split][: args.max_samples]
+    x = select_rows(vision["features"], vlookup, rows)
+    y = select_rows(text["features"], tlookup, rows)
+
+    gx = normalized(x) @ normalized(x).T
+    gy = normalized(y) @ normalized(y).T
+    gx_upper = upper_triangle(gx)
+    gy_upper = upper_triangle(gy)
+    rho = spearmanr(gx_upper, gy_upper).statistic
+    generator = torch.Generator().manual_seed(args.seed)
+    shuffled_rhos = []
+    for _ in range(args.permutations):
+        permutation = torch.randperm(len(y), generator=generator)
+        shuffled_gy = gy[permutation][:, permutation]
+        shuffled_rhos.append(
+            float(spearmanr(gx_upper, upper_triangle(shuffled_gy)).statistic)
+        )
+    shuffled_mean = float(np.mean(shuffled_rhos))
+    shuffled_std = float(np.std(shuffled_rhos))
+    result = {
+        "split": args.split,
+        "samples": len(rows),
+        "linear_cka": linear_cka(x, y),
+        "pairwise_cosine_spearman": float(rho),
+        "shuffled_spearman_mean": shuffled_mean,
+        "shuffled_spearman_std": shuffled_std,
+        "spearman_shuffle_z": float(
+            (rho - shuffled_mean) / max(shuffled_std, 1e-12)
+        ),
+        "permutations": args.permutations,
+        "vision_dim": x.shape[-1],
+        "text_dim": y.shape[-1],
+    }
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    write_json(args.output, result)
+    print(result)
+
+
+if __name__ == "__main__":
+    main()