1 files changed, 222 insertions, 0 deletions
diff --git a/worldalign/vg_graph_diagnose.py b/worldalign/vg_graph_diagnose.py
new file mode 100644
index 0000000..3a0ad2e
--- /dev/null
+++ b/worldalign/vg_graph_diagnose.py
@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+from scipy import sparse
+from scipy.sparse.linalg import eigsh
+import torch
+
+from .common import normalized, write_json
+from .vg_diagnose import (
+    aligned_text_indices,
+    arbitrary_target_retrieval,
+    bundle_signature,
+    rank_normalize_columns,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--vision", default="artifacts/vg/vision_features.pt"
+    )
+    parser.add_argument("--text", default="artifacts/vg/text_features.pt")
+    parser.add_argument(
+        "--truth", default="artifacts/vg/ground_truth.private.jsonl"
+    )
+    parser.add_argument("--max-nodes", type=int, default=5_000)
+    parser.add_argument("--knn", type=int, default=32)
+    parser.add_argument("--eigenpairs", type=int, default=96)
+    parser.add_argument("--heat-scales", type=int, default=24)
+    parser.add_argument("--message-passing-steps", type=int, default=3)
+    parser.add_argument("--output", default="artifacts/vg/graph_diagnostics.json")
+    return parser.parse_args()
+
+
+def scene_features(state: dict, indices: torch.Tensor | None = None) -> torch.Tensor:
+    if "global_features" in state:
+        features = state["global_features"]
+    else:
+        features = state["region_features"].mean(1)
+    return features if indices is None else features[indices]
+
+
+def exact_rank_knn_graph(
+    features: torch.Tensor, neighbors: int
+) -> tuple[sparse.csr_matrix, torch.Tensor]:
+    """Build a symmetrized rank-weighted graph and retain density statistics."""
+    x = normalized(features.float())
+    similarity = x @ x.T
+    similarity.fill_diagonal_(-torch.inf)
+    k = min(neighbors, len(x) - 1)
+    values, indices = similarity.topk(k, dim=1)
+
+    # Rank weights avoid assuming that cosine scales agree across modalities.
+    rank_weight = np.exp(
+        -np.arange(k, dtype=np.float64) / max(k / 4.0, 1.0)
+    )
+    rows = np.repeat(np.arange(len(x), dtype=np.int64), k)
+    cols = indices.cpu().numpy().reshape(-1)
+    data = np.tile(rank_weight, len(x))
+    directed = sparse.csr_matrix(
+        (data, (rows, cols)), shape=(len(x), len(x))
+    )
+    adjacency = directed.maximum(directed.T)
+    adjacency.setdiag(0)
+    adjacency.eliminate_zeros()
+
+    density_quantiles = torch.quantile(
+        values,
+        torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0]),
+        dim=1,
+    ).T
+    return adjacency, density_quantiles
+
+
+def heat_kernel_signature(
+    adjacency: sparse.csr_matrix, eigenpairs: int, scales: int
+) -> tuple[torch.Tensor, np.ndarray]:
+    """Return basis-invariant diagonal heat-kernel signatures."""
+    degree = np.asarray(adjacency.sum(axis=1)).reshape(-1)
+    inv_sqrt_degree = 1.0 / np.sqrt(np.maximum(degree, 1e-12))
+    normalizer = sparse.diags(inv_sqrt_degree)
+    normalized_adjacency = normalizer @ adjacency @ normalizer
+    k = min(eigenpairs, adjacency.shape[0] - 2)
+    eigenvalues, eigenvectors = eigsh(
+        normalized_adjacency, k=k, which="LA", tol=1e-4
+    )
+    order = np.argsort(eigenvalues)[::-1]
+    eigenvalues = eigenvalues[order]
+    eigenvectors = eigenvectors[:, order]
+    laplacian_eigenvalues = np.clip(1.0 - eigenvalues, 0.0, None)
+
+    positive = laplacian_eigenvalues[laplacian_eigenvalues > 1e-8]
+    lower = 0.25 / max(float(positive.max(initial=1.0)), 1e-8)
+    upper = 8.0 / max(float(positive.min(initial=1.0)), 1e-8)
+    times = np.geomspace(lower, upper, scales)
+    heat = (eigenvectors**2) @ np.exp(
+        -laplacian_eigenvalues[:, None] * times[None, :]
+    )
+    return torch.from_numpy(heat).float(), degree
+
+
+def graph_signature(
+    features: torch.Tensor,
+    neighbors: int,
+    eigenpairs: int,
+    heat_scales: int,
+    message_passing_steps: int,
+) -> torch.Tensor:
+    adjacency, density = exact_rank_knn_graph(features, neighbors)
+    heat, degree = heat_kernel_signature(
+        adjacency, eigenpairs=eigenpairs, scales=heat_scales
+    )
+    base = rank_normalize_columns(
+        torch.cat(
+            [
+                density,
+                torch.from_numpy(np.log1p(degree)).float()[:, None],
+                heat,
+            ],
+            dim=1,
+        )
+    )
+    transition = sparse.diags(
+        1.0 / np.maximum(np.asarray(adjacency.sum(axis=1)).reshape(-1), 1e-12)
+    ) @ adjacency
+    current = base
+    for _ in range(message_passing_steps):
+        current_np = current.numpy()
+        neighbor_mean = torch.from_numpy(transition @ current_np).float()
+        neighbor_square_mean = torch.from_numpy(
+            transition @ np.square(current_np)
+        ).float()
+        neighbor_std = torch.sqrt(
+            torch.clamp(neighbor_square_mean - neighbor_mean.square(), min=0)
+        )
+        current = rank_normalize_columns(
+            torch.cat([base, neighbor_mean, neighbor_std], dim=1)
+        )
+    return current
+
+
+def main() -> None:
+    args = parse_args()
+    vision = torch.load(args.vision, map_location="cpu", weights_only=False)
+    text = torch.load(args.text, map_location="cpu", weights_only=False)
+    n = min(args.max_nodes, len(vision["node_ids"]), len(text["node_ids"]))
+    vision_ids = vision["node_ids"][:n]
+    target_full = aligned_text_indices(
+        vision_ids, text["node_ids"], args.truth
+    )
+    candidate_indices = torch.unique(target_full, sorted=False)
+    if len(candidate_indices) != n:
+        raise ValueError("Ground truth is not a one-to-one permutation")
+    candidate_lookup = {
+        int(old): new for new, old in enumerate(candidate_indices.tolist())
+    }
+    targets = torch.tensor(
+        [candidate_lookup[int(old)] for old in target_full.tolist()]
+    )
+
+    vision_scene = scene_features(vision)[:n]
+    text_scene = scene_features(text, candidate_indices)
+    vision_graph = graph_signature(
+        vision_scene,
+        args.knn,
+        args.eigenpairs,
+        args.heat_scales,
+        args.message_passing_steps,
+    )
+    text_graph = graph_signature(
+        text_scene,
+        args.knn,
+        args.eigenpairs,
+        args.heat_scales,
+        args.message_passing_steps,
+    )
+
+    vision_bundle = rank_normalize_columns(
+        bundle_signature(vision["region_features"][:n])
+    )
+    text_bundle = rank_normalize_columns(
+        bundle_signature(text["region_features"][candidate_indices])
+    )
+    combined_vision = torch.cat(
+        [normalized(vision_graph), normalized(vision_bundle)], dim=1
+    )
+    combined_text = torch.cat(
+        [normalized(text_graph), normalized(text_bundle)], dim=1
+    )
+
+    result = {
+        "nodes": n,
+        "knn": min(args.knn, n - 1),
+        "eigenpairs": min(args.eigenpairs, n - 2),
+        "heat_scales": args.heat_scales,
+        "message_passing_steps": args.message_passing_steps,
+        "graph_signature_retrieval": arbitrary_target_retrieval(
+            vision_graph, text_graph, targets
+        ),
+        "bundle_signature_retrieval": arbitrary_target_retrieval(
+            vision_bundle, text_bundle, targets
+        ),
+        "combined_signature_retrieval": arbitrary_target_retrieval(
+            combined_vision, combined_text, targets
+        ),
+        "evaluation_note": (
+            "Each signature is computed independently from one modality. "
+            "Ground truth is loaded only after graph construction to score "
+            "the hidden permutation."
+        ),
+    }
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    write_json(args.output, result)
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()