social-data-attribution-w2/scripts/analysis/build_figure4_table.py

"""Build the Figure-4-style 24x24 γ table for the N10 held-out ToM analysis.

Reads cross-probe eval metrics from <eval-grid>/<topic>/eval_<probe>/ and
per-probe OLMES baselines from <baselines>/<probe>:mc_*/.

Computes γ = post_acc - pre_acc for each (probe, topic) cell.
Emits a JSON dump plus Markdown tables for 24 topics x 24 probes.

Usage (CPU-only, runs on macOS):
    python build_figure4_table.py \
        --eval-grid /path/to/n10b/eval_grid \
        --baselines /path/to/tom-evals-runs/parallel_5shot/allenai-Olmo-3-1025-7B \
        --out-dir notes/n10b_figure4/
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from figure4_table_helpers import (
    ALL_TOPICS,
    METRICS_OF_INTEREST,
    PROBES,
    render_baseline_md,
    render_md_table,
)
from figure4_table_io import (
    collect_baselines,
    collect_null,
    collect_post,
    compute_gamma,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--eval-grid", type=Path, required=True)
    parser.add_argument("--baselines", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--metrics", nargs="*", default=METRICS_OF_INTEREST)
    args = parser.parse_args(argv)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    baselines = collect_baselines(args.baselines)
    post = collect_post(args.eval_grid)
    null_post = collect_null(args.eval_grid)
    has_null = any(null_post.get(p, {}) for p, _ in PROBES)

    full: dict[str, Any] = {
        "baselines": baselines,
        "null_post": null_post,
        "post": post,
        "grid": {},
    }
    for metric in args.metrics:
        full["grid"][metric] = {
            mode: compute_gamma(baselines, post, metric, mode=mode, null_post=null_post)
            for mode in ("absolute", "relative")
        }

    (args.out_dir / "full.json").write_text(json.dumps(full, indent=2, default=str))

    md_parts = [render_baseline_md(baselines), ""]
    for metric in args.metrics:
        md_parts.append(
            render_md_table(
                full["grid"][metric]["absolute"], metric, "gamma", "absolute"
            )
        )
        md_parts.append("")
        if has_null:
            md_parts.append(
                render_md_table(
                    full["grid"][metric]["absolute"], metric, "net_gamma", "absolute"
                )
            )
            md_parts.append("")
            md_parts.append(
                render_md_table(
                    full["grid"][metric]["relative"], metric, "net_gamma", "relative"
                )
            )
            md_parts.append("")
    (args.out_dir / "figure4.md").write_text("\n".join(md_parts))

    # Coverage report
    n_total = len(ALL_TOPICS) * len(PROBES)
    cov_grid = full["grid"].get("primary_score", full["grid"][args.metrics[0]])[
        "absolute"
    ]
    n_filled = sum(
        1
        for probe, _ in PROBES
        for topic in ALL_TOPICS
        if cov_grid[probe][topic].get("gamma") is not None
    )
    print(
        f"Coverage: {n_filled} / {n_total} cells ({n_filled / n_total * 100:.1f}%)  null_row={'yes' if has_null else 'no'}"
    )
    print(f"Wrote {args.out_dir / 'full.json'}")
    print(f"Wrote {args.out_dir / 'figure4.md'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())