repro_bundle/scripts/analyze_run.py

"""Post-hoc analysis for the in-weights path-star run: geometry heatmap + UMAP (Claim 3)."""
import argparse
import glob
import json
import os
import pickle
import re

import numpy as np
import torch


def find_final_checkpoint(run_root):
    candidates = glob.glob(os.path.join(run_root, "**", "*_final_model.pt"), recursive=True)
    if not candidates:
        candidates = glob.glob(os.path.join(run_root, "**", "*.pt"), recursive=True)
    candidates.sort(key=os.path.getmtime)
    return candidates[-1]


def find_run_dir(run_root):
    dirs = [d for d in glob.glob(os.path.join(run_root, "in_weights", "*")) if os.path.isdir(d)]
    dirs.sort(key=os.path.getmtime)
    return dirs[-1]


def parse_config(run_dir_name):
    m = re.search(r"star-d(\d+)-dt(\d+)-p(\d+)-n(\d+)", run_dir_name)
    star_degree, subtree_degree, path_length, total_nodes = map(int, m.groups())
    return star_degree, subtree_degree, path_length, total_nodes


def load_paths(dataset_dir, star_degree, path_length, total_nodes):
    """Parse `leaf=root,hop1,...,leaf` lines from all train/test split files."""
    pattern = os.path.join(
        dataset_dir,
        f"star_deg_{star_degree}_deg_tree_1_path_{path_length}_nodes_{total_nodes}_sd_*_fb_*_selfedge_0_*.txt",
    )
    files = [f for f in glob.glob(pattern) if "pretrain" not in f]
    leaf_ids, first_hop_ids, full_paths = [], [], []
    for fp in files:
        with open(fp) as fh:
            for line in fh:
                line = line.strip()
                if not line or "=" not in line:
                    continue
                leaf_str, rest = line.split("=", 1)
                nodes = rest.split(",")
                if len(nodes) < 2:
                    continue
                leaf_ids.append(int(leaf_str))
                first_hop_ids.append(int(nodes[1]))
                full_paths.append([int(n) for n in nodes])
    return np.array(leaf_ids), np.array(first_hop_ids), full_paths


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--experiment_log_root", required=True)
    ap.add_argument("--tag", required=True)
    ap.add_argument("--out_dir", required=True)
    args = ap.parse_args()
    os.makedirs(args.out_dir, exist_ok=True)

    run_dir = find_run_dir(args.experiment_log_root)
    ckpt_path = find_final_checkpoint(run_dir)
    print("Using checkpoint:", ckpt_path)

    star_degree, subtree_degree, path_length, total_nodes = parse_config(os.path.basename(run_dir))
    print("Graph config:", star_degree, subtree_degree, path_length, total_nodes)

    repo_root = os.environ.get("REPO_ROOT", "/workspace/repo")
    import sys
    sys.path.insert(0, repo_root)

    state_dict = torch.load(ckpt_path, map_location="cpu")
    embed_w = state_dict["embed_tokens.weight"]
    vocab_size, d_model = embed_w.shape
    print("Embedding matrix:", embed_w.shape)

    node_embs = embed_w[:total_nodes].numpy()

    dataset_dir = os.path.join(repo_root, "data/datasets/in_weights_graphs", "star_graphs_randomized")
    leaf_ids, first_hop_ids, full_paths = load_paths(dataset_dir, star_degree, path_length, total_nodes)
    print(f"Parsed {len(leaf_ids)} paths from dataset files (train+test combined).")

    n_show = min(300, len(leaf_ids))
    idx = np.linspace(0, len(leaf_ids) - 1, n_show).astype(int)
    leaf_sample = leaf_ids[idx]
    first_hop_sample = first_hop_ids[idx]

    leaf_emb = node_embs[leaf_sample]
    first_hop_emb = node_embs[first_hop_sample]

    def cosine_distance_matrix(a, b):
        a_n = a / (np.linalg.norm(a, axis=1, keepdims=True) + 1e-8)
        b_n = b / (np.linalg.norm(b, axis=1, keepdims=True) + 1e-8)
        return 1 - a_n @ b_n.T

    heat = cosine_distance_matrix(leaf_emb, first_hop_emb)
    diag_mean = np.mean(np.diag(heat))
    offdiag_mean = (np.sum(heat) - np.sum(np.diag(heat))) / (heat.size - len(heat))
    print(f"Diagonal mean cosine distance: {diag_mean:.4f}")
    print(f"Off-diagonal mean cosine distance: {offdiag_mean:.4f}")

    np.savez(
        os.path.join(args.out_dir, f"heatmap_{args.tag}.npz"),
        heat=heat, diag_mean=diag_mean, offdiag_mean=offdiag_mean,
    )

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(heat, cmap="viridis", aspect="auto")
    ax.set_xlabel("First-hop token of path j")
    ax.set_ylabel("Leaf token of path i")
    ax.set_title(f"Cosine distance heatmap ({args.tag})\ndiag={diag_mean:.3f} vs offdiag={offdiag_mean:.3f}")
    plt.colorbar(im, ax=ax)
    plt.tight_layout()
    plt.savefig(os.path.join(args.out_dir, f"heatmap_{args.tag}.png"), dpi=150)
    plt.close(fig)

    # ---- UMAP projection of node embeddings, colored by path/branch id ----
    try:
        import umap
        n_branches_show = min(40, len(full_paths))
        branch_idx = np.linspace(0, len(full_paths) - 1, n_branches_show).astype(int)
        sel_list, branch_id_list = [], []
        for color_id, pi in enumerate(branch_idx):
            for node_id in full_paths[pi]:
                sel_list.append(node_id)
                branch_id_list.append(color_id)
        sel = np.array(sel_list)
        branch_id = np.array(branch_id_list)
        reducer = umap.UMAP(n_components=2, random_state=0)
        proj = reducer.fit_transform(node_embs[sel])
        np.savez(os.path.join(args.out_dir, f"umap_{args.tag}.npz"), proj=proj, branch_id=branch_id, sel=sel)

        fig, ax = plt.subplots(figsize=(6, 6))
        sc = ax.scatter(proj[:, 0], proj[:, 1], c=branch_id, cmap="hsv", s=4, alpha=0.6)
        ax.set_title(f"UMAP of node embeddings ({args.tag})")
        plt.tight_layout()
        plt.savefig(os.path.join(args.out_dir, f"umap_{args.tag}.png"), dpi=150)
        plt.close(fig)
    except Exception as e:
        print("UMAP failed:", e)

    # ---- Parse test accuracy from train log ----
    log_path = os.path.join(os.path.dirname(args.experiment_log_root), f"train_log_{args.tag}.txt")
    if not os.path.exists(log_path):
        log_path = os.path.join("/data", f"train_log_{args.tag}.txt")
    best_acc = None
    if os.path.exists(log_path):
        with open(log_path) as f:
            for line in f:
                m = re.search(r"Test Acc: ([\d.]+)%", line)
                if m:
                    val = float(m.group(1))
                    best_acc = val if best_acc is None else max(best_acc, val)
    summary = {
        "tag": args.tag,
        "star_degree": star_degree,
        "path_length": path_length,
        "total_nodes": total_nodes,
        "diag_mean_cosine_distance": float(diag_mean),
        "offdiag_mean_cosine_distance": float(offdiag_mean),
        "best_test_accuracy_pct": best_acc,
    }
    with open(os.path.join(args.out_dir, f"summary_{args.tag}.json"), "w") as f:
        json.dump(summary, f, indent=2)
    print("SUMMARY:", json.dumps(summary))


if __name__ == "__main__":
    main()