social-data-attribution-w2/scripts/figure1/build_panel3_influence.py

"""Panel-3 thumbnail for Fig 1 (taxonomy-based influence).

Reproduces the 6-row × 4-benchmark mini-heatmap currently hand-coded in
``figure_1_hero.drawio``. Each cell shows the per-bin signed mean
influence z-score for one benchmark; the curated 6 bins span the
SocialIQA-positive end through the STEM-positive end so the sign flip
reads at a glance.

Usage:
    uv run python scripts/figure1/build_panel3_influence.py
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.colors import TwoSlopeNorm

from dolma.distribution_report.style import style_size

# Curated bins (topic_label, format_label, display_name) mirroring the
# drawio panel-3 selection.
_BINS: list[tuple[str, str, str]] = [
    ("social_life", "personal_blog",                "Soc. Life × Blog"),
    ("social_life", "q_a_forum",                    "Soc. Life × Q&A"),
    ("literature", "tutorial",                      "Lit. × Tutorial"),
    ("science_math_and_technology", "academic_writing", "Sci. × Academic"),
    ("science_math_and_technology", "tutorial",     "Sci. × Tutorial"),
    ("education_and_jobs", "knowledge_article",     "Edu. × Knowledge"),
]

_BENCHMARKS: list[tuple[str, str]] = [
    ("queries_socialiqa",            "SocialIQA"),
    ("queries_mmlu_social_science",  "MMLU\nSoc. Sci."),
    ("queries_arc_challenge",        "ARC-\nChallenge"),
    ("queries_mmlu_stem",            "MMLU\nSTEM"),
]


def _zscore(series: pd.Series) -> pd.Series:
    mu = series.mean()
    sd = series.std()
    if sd == 0 or pd.isna(sd):
        return pd.Series(np.zeros(len(series)), index=series.index)
    return (series - mu) / sd


def _load_bin_scores(bin_dir: Path, benchmark_key: str) -> pd.DataFrame:
    """Load aggregated topic×format bin scores; fall back to perquery file."""
    agg = bin_dir / f"{benchmark_key}_bin_scores.csv"
    if agg.exists():
        df = pd.read_csv(agg)
        return df[["topic_label", "format_label", "mean_score"]]
    perquery = bin_dir / f"{benchmark_key}_bin_scores_perquery.csv"
    df = pd.read_csv(perquery)
    return df[["topic_label", "format_label", "mean_influence"]].rename(
        columns={"mean_influence": "mean_score"}
    )


def _load_z_for_bin(bin_dir: Path, benchmark_key: str, topic: str, fmt: str) -> float:
    """Return the z-scored mean_score for one (topic, format) cell of one benchmark."""
    df = _load_bin_scores(bin_dir, benchmark_key)
    df["z_mean_score"] = _zscore(df["mean_score"])
    hit = df[(df["topic_label"] == topic) & (df["format_label"] == fmt)]
    if hit.empty:
        return float("nan")
    return float(hit["z_mean_score"].iloc[0])


def build_panel(
    bin_dir: Path,
    out_dir: Path,
    figsize: tuple[float, float] = (4.4, 2.6),
    dpi: int = 300,
) -> None:
    matrix = np.zeros((len(_BINS), len(_BENCHMARKS)))
    for i, (topic, fmt, _) in enumerate(_BINS):
        for j, (bench_key, _) in enumerate(_BENCHMARKS):
            matrix[i, j] = _load_z_for_bin(bin_dir, bench_key, topic, fmt)

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    abs_max = float(np.nanmax(np.abs(matrix)))
    norm = TwoSlopeNorm(vmin=-abs_max, vcenter=0.0, vmax=abs_max)
    ax.imshow(matrix, cmap="RdBu", norm=norm, aspect="auto")

    ax.set_xticks(range(len(_BENCHMARKS)))
    ax.set_xticklabels([disp for _, disp in _BENCHMARKS],
                        fontsize=style_size("TICK"), family="serif")
    ax.set_yticks(range(len(_BINS)))
    ax.set_yticklabels([disp for _, _, disp in _BINS],
                        fontsize=style_size("TICK"), family="serif")
    ax.tick_params(axis="both", which="both", length=0)
    for spine in ax.spines.values():
        spine.set_visible(False)

    for i in range(len(_BINS)):
        for j in range(len(_BENCHMARKS)):
            val = matrix[i, j]
            # Choose text color for contrast on the colored cell
            text_color = "white" if abs(val) > abs_max * 0.55 else "black"
            ax.text(j, i, f"{val:+.2f}", ha="center", va="center",
                    fontsize=style_size("ANNOTATION"), color=text_color, family="serif")

    # Subtle grid lines between cells
    ax.set_xticks(np.arange(len(_BENCHMARKS) + 1) - 0.5, minor=True)
    ax.set_yticks(np.arange(len(_BINS) + 1) - 0.5, minor=True)
    ax.grid(which="minor", color="white", linewidth=1.2)
    ax.tick_params(which="minor", length=0)

    fig.tight_layout(pad=0.4)
    out_dir.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_dir / "panel3_influence.pdf", bbox_inches="tight")
    fig.savefig(out_dir / "panel3_influence.png", bbox_inches="tight", dpi=dpi)
    plt.close(fig)
    print(f"Wrote {out_dir / 'panel3_influence.pdf'} (+ .png)")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bin-dir", type=Path, default=Path("artifacts/influence_bin_scores"))
    parser.add_argument("--out-dir", type=Path, default=Path("artifacts/figure_1"))
    args = parser.parse_args()
    build_panel(args.bin_dir, args.out_dir)


if __name__ == "__main__":
    main()