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

#!/usr/bin/env python3
# pyright: reportMissingImports=false, reportGeneralTypeIssues=false, reportAttributeAccessIssue=false, reportArgumentType=false, reportReturnType=false
"""Build cross-probe topic-format bin consensus for SocialTDA probes."""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OUT = REPO_ROOT / "artifacts/cross_probe_consensus"
REQUIRED_COLUMNS = {"topic_label", "format_label", "mean_influence", "doc_count"}
GRID_ROWS = 24 * 24
LABEL_OVERRIDES = {"art_and_design": "Art & Design", "crime_and_law": "Crime & Law", "home_and_hobbies": "Home & Hobbies", "q_a_forum": "Q&A Forum", "science_math_and_technology": "Science & Technology"}

@dataclass(frozen=True)
class Source:
    name: str
    display: str
    group: str
    path: Path

def sources() -> list[Source]:
    raw = DEFAULT_OUT / "holdout_raw"
    core = REPO_ROOT / "artifacts/influence_bin_scores"
    bbh = REPO_ROOT / "artifacts/influence_bbh_instruct"
    rows = [
        ("socialiqa", "SocialIQA", "headline", core / "queries_socialiqa_bin_scores_perquery.csv"),
        ("mmlu_social_science", "MMLU Social Sciences", "headline", core / "queries_mmlu_social_science_bin_scores_perquery.csv"),
        ("bbq", "BBQ", "headline", raw / "bbq_main_bin_scores_perquery.csv"),
        ("bbh_disambiguation_qa", "BBH Disambig. QA", "headline", raw / "bbh_disambiguation_qa_main_bin_scores_perquery.csv"),
        ("bbh_snarks", "BBH Snarks", "headline", bbh / "queries_bbh_snarks_bin_scores_perquery.csv"),
        ("tombench", "ToMBench", "headline", raw / "tombench_secondary_bin_scores_perquery.csv"),
        ("negotiationtom", "NegotiationToM", "headline", raw / "negotiationtom_main_bin_scores_perquery.csv"),
        ("pub_retained", "PUB retained", "headline", raw / "pub_main_excluding_pub_2_pub_3_bin_scores_perquery.csv"),
        ("simpletom_mental", "SimpleToM mental", "headline", raw / "simpletom_mental_state_bin_scores_perquery.csv"),
        ("mmlu_moral", "MMLU moral/hum.", "headline", raw / "mmlu_moral_main_bin_scores_perquery.csv"),
        ("ethics_non_justice", "ETHICS non-justice", "headline", raw / "ethics_main_non_justice_bin_scores_perquery.csv"),
        ("morables", "MORABLES", "headline", raw / "morables_secondary_bin_scores_perquery.csv"),
        ("moralexceptqa_rbqa", "MoralExceptQA/RBQA", "headline", raw / "moralexceptqa_rbqa_secondary_bin_scores_perquery.csv"),
        ("arc_challenge", "ARC-Challenge", "control", core / "queries_arc_challenge_bin_scores_perquery.csv"),
        ("mmlu_stem", "MMLU STEM", "control", core / "queries_mmlu_stem_bin_scores_perquery.csv"),
    ]
    return [Source(*row) for row in rows]

def pretty_label(value: str) -> str:
    return LABEL_OVERRIDES.get(value, value.replace("_", " ").title())

def latex_escape(value: object) -> str:
    text = str(value)
    return (
        text.replace("\\", r"\textbackslash{}")
        .replace("&", r"\&")
        .replace("%", r"\%")
        .replace("$", r"\$")
        .replace("#", r"\#")
        .replace("_", r"\_")
    )

def load_source(source: Source) -> pd.DataFrame:
    if not source.path.exists():
        raise FileNotFoundError(source.path)
    df = pd.read_csv(source.path)
    missing = REQUIRED_COLUMNS - set(df.columns)
    if missing:
        raise ValueError(f"{source.path} missing columns: {sorted(missing)}")
    if len(df) != GRID_ROWS:
        raise ValueError(f"{source.path} has {len(df)} rows, expected {GRID_ROWS}")
    sigma = df["mean_influence"].std()
    if sigma == 0 or pd.isna(sigma):
        raise ValueError(f"{source.path} has zero or NaN mean_influence std")
    out = df[["topic_label", "format_label", "mean_influence", "doc_count"]].copy()
    out["benchmark"] = source.name
    out["display"] = source.display
    out["group"] = source.group
    out["zscore"] = (out["mean_influence"] - out["mean_influence"].mean()) / sigma
    out["bin_label"] = out["topic_label"].map(pretty_label) + " / " + out["format_label"].map(pretty_label)
    return out

def consensus_table(long_df: pd.DataFrame) -> pd.DataFrame:
    headline = long_df[long_df["group"] == "headline"]
    controls = long_df[long_df["group"] == "control"]
    rows = []
    for (topic, fmt), group in headline.groupby(["topic_label", "format_label"]):
        ctrl = controls[(controls["topic_label"] == topic) & (controls["format_label"] == fmt)]
        mean_z = float(group["zscore"].mean())
        ctrl_z = float(ctrl["zscore"].mean()) if len(ctrl) else float("nan")
        rows.append(
            {
                "topic_label": topic,
                "format_label": fmt,
                "bin_label": group["bin_label"].iloc[0],
                "positive_probe_count": int((group["zscore"] > 0).sum()),
                "strong_positive_count": int((group["zscore"] >= 1).sum()),
                "mean_z": mean_z,
                "median_z": float(group["zscore"].median()),
                "control_mean_z": ctrl_z,
                "social_specificity": mean_z - ctrl_z,
            }
        )
    return pd.DataFrame(rows).sort_values(
        ["positive_probe_count", "strong_positive_count", "mean_z"],
        ascending=[False, False, False],
    )

def write_latex_table(df: pd.DataFrame, path: Path, n_rows: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    rows = [
        r"\begin{table}[!htbp]",
        r"\centering",
        r"\scriptsize",
        r"\setlength{\tabcolsep}{3pt}",
        r"\begin{tabularx}{\linewidth}{@{}p{0.33\linewidth}rrrrr@{}}",
        r"\toprule",
        r"Topic-format bin & Pos. & Strong & Mean $z$ & Ctrl. $z$ & Spec. \\",
        r"\midrule",
    ]
    for row in df.head(n_rows).itertuples(index=False):
        rows.append(
            f"{latex_escape(row.bin_label)} & {row.positive_probe_count} & "
            f"{row.strong_positive_count} & {row.mean_z:+.2f} & "
            f"{row.control_mean_z:+.2f} & {row.social_specificity:+.2f} \\\\"
        )
    rows += [
        r"\bottomrule",
        r"\end{tabularx}",
        r"\caption{Topic-format bins with the strongest supportive consensus across the held-out and core social-domain probes. Scores are z-scored within each probe over the 576 topic-format bins. \textit{Pos.} counts headline probes with positive signed influence, \textit{Strong} counts headline probes with $z \geq 1$, and \textit{Spec.} subtracts the mean control-probe $z$ from the mean headline-probe $z$.}",
        r"\label{tab:heldout-supportive-consensus}",
        r"\end{table}",
    ]
    path.write_text("\n".join(rows) + "\n")

def write_figure(long_df: pd.DataFrame, ranking: pd.DataFrame, paper_root: Path, n_bins: int) -> None:
    scripts = paper_root / "scripts"
    sys.path.insert(0, str(scripts))
    from figure_style import APPENDIX_FIGURE_WIDTH_IN, DIVERGING_NEG, DIVERGING_NEUTRAL, DIVERGING_POS, apply_matplotlib_style, save_publication_figure

    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm

    apply_matplotlib_style()
    top = ranking.head(n_bins)
    bin_order = list(zip(top["topic_label"], top["format_label"]))
    bench_order = [source.display for source in sources()]
    pivot = long_df.pivot_table(index=["topic_label", "format_label"], columns="display", values="zscore")
    matrix = np.array([[pivot.loc[bin_key, bench] for bench in bench_order] for bin_key in bin_order])
    fig, ax = plt.subplots(figsize=(APPENDIX_FIGURE_WIDTH_IN, 3.85))
    limit = max(2.5, float(np.nanmax(np.abs(matrix))))
    cmap = LinearSegmentedColormap.from_list("socialtda_diverging", [DIVERGING_NEG, DIVERGING_NEUTRAL, DIVERGING_POS], N=256)
    image = ax.imshow(matrix, aspect="auto", cmap=cmap, norm=TwoSlopeNorm(vmin=-limit, vcenter=0, vmax=limit))
    ax.set_xticks(range(len(bench_order)))
    ax.set_xticklabels(bench_order, rotation=38, ha="right")
    ax.set_yticks(range(len(bin_order)))
    ax.set_yticklabels([top.iloc[i]["bin_label"] for i in range(len(top))])
    ax.set_xlabel("Probe")
    ax.set_ylabel("Topic-format bin")
    cbar = fig.colorbar(image, ax=ax, pad=0.012, fraction=0.035)
    cbar.set_label("Within-probe signed influence (z)")
    output_base = paper_root / "figures/holdout_figures/consensus/fig_cross_probe_supportive_consensus"
    save_publication_figure(fig, output_base)
    plt.close(fig)
    wrapper = paper_root / "figures/fig-heldout-supportive-consensus.tex"
    wrapper.write_text(
        "\\begin{figure}[!htbp]\n"
        "\\centering\n"
        "\\includegraphics[width=\\MainFigureWidth]{figures/holdout_figures/consensus/fig_cross_probe_supportive_consensus.pdf}\n"
        "\\caption{Cross-probe supportive consensus for the top topic-format bins in Table~\\ref{tab:heldout-supportive-consensus}. Each column is standardized within a probe, so the heatmap shows whether the same corpus bins are repeatedly supportive rather than comparing raw influence magnitudes across probes.}\n"
        "\\label{fig:heldout-supportive-consensus}\n"
        "\\end{figure}\n"
    )

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--paper-root", type=Path)
    parser.add_argument("--table-rows", type=int, default=8)
    parser.add_argument("--figure-bins", type=int, default=10)
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    long_df = pd.concat([load_source(source) for source in sources()], ignore_index=True)
    ranking = consensus_table(long_df)
    long_df.to_csv(args.out_dir / "bin_z_matrix_long.csv", index=False)
    ranking.to_csv(args.out_dir / "supportive_bin_consensus.csv", index=False)
    ranking.head(50).to_csv(args.out_dir / "top_supportive_bins.csv", index=False)
    if args.paper_root:
        write_latex_table(ranking, args.paper_root / "tables/tab-heldout-supportive-consensus.tex", args.table_rows)
        write_figure(long_df, ranking, args.paper_root, args.figure_bins)
    print(f"wrote {args.out_dir / 'supportive_bin_consensus.csv'}")

if __name__ == "__main__":
    main()