social-data-attribution-w2/scripts/diagnostics/check_influence_differentials.py

"""Diagnostic: verify influence differential calculations match source data.

Run from the repo root:
    uv run --no-sync python scripts/diagnostics/check_influence_differentials.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[2]
PERQUERY_DIR = REPO_ROOT / "artifacts" / "influence_bin_scores"
SPLIT_DIR = REPO_ROOT / "artifacts" / "influence_bin_scores_split"

BENCHMARKS = [
    "queries_socialiqa",
    "queries_gsm8k",
    "queries_mmlu_social_science",
    "queries_mmlu_stem",
]

CONTRASTIVE_PAIRS = [
    ("queries_socialiqa", "queries_gsm8k"),
    ("queries_mmlu_social_science", "queries_mmlu_stem"),
]

SCALE = 1e6


def _normalize(label: str) -> str:
    label = label.strip().lower().replace(" ", "_")
    prefix = "__label__"
    return label if label.startswith(prefix) else f"{prefix}{label}"


def load_perquery(key: str) -> pd.DataFrame:
    path = PERQUERY_DIR / f"{key}_bin_scores_perquery.csv"
    if not path.exists():
        path = PERQUERY_DIR / f"{key}_bin_scores.csv"
    df = pd.read_csv(path)
    df["topic_label"] = df["topic_label"].apply(_normalize)
    df["format_label"] = df["format_label"].apply(_normalize)
    if "median_influence" in df.columns:
        df = df.rename(columns={"median_influence": "mean_score"})
    return df


def load_split(key: str, kind: str) -> pd.DataFrame:
    path = SPLIT_DIR / f"{key}_bin_scores_{kind}.csv"
    df = pd.read_csv(path)
    df["topic_label"] = df["topic_label"].apply(_normalize)
    df["format_label"] = df["format_label"].apply(_normalize)
    if "median_influence" in df.columns:
        df = df.rename(columns={"median_influence": "mean_score"})
    return df


def pivot(df: pd.DataFrame, col: str = "mean_score") -> pd.DataFrame:
    return df.pivot(index="topic_label", columns="format_label", values=col).fillna(0)


def report_range(label: str, arr: np.ndarray) -> None:
    finite = arr[np.isfinite(arr)]
    p5, p95 = np.percentile(np.abs(finite), [5, 95])
    print(f"  {label}")
    print(f"    min/max: {finite.min():.4f} / {finite.max():.4f}")
    print(f"    |val| p5/p95: {p5:.4f} / {p95:.4f}")


def main() -> None:
    print("=== Per-benchmark signed influence (perquery median, ×10⁻⁶) ===")
    grids: dict[str, pd.DataFrame] = {}
    for key in BENCHMARKS:
        df = load_perquery(key)
        grids[key] = df
        vals = df["mean_score"].values * SCALE
        short = key.removeprefix("queries_")
        report_range(short, vals)
    print()

    print("=== Cross-benchmark contrastive differentials (A − B, same perquery data) ===")
    for key_a, key_b in CONTRASTIVE_PAIRS:
        m_a = pivot(grids[key_a])
        m_b = pivot(grids[key_b])
        common_topics = sorted(set(m_a.index) & set(m_b.index))
        common_formats = sorted(set(m_a.columns) & set(m_b.columns))
        a_aligned = m_a.reindex(index=common_topics, columns=common_formats, fill_value=0).values * SCALE
        b_aligned = m_b.reindex(index=common_topics, columns=common_formats, fill_value=0).values * SCALE
        diff = a_aligned - b_aligned
        short_a = key_a.removeprefix("queries_")
        short_b = key_b.removeprefix("queries_")
        report_range(f"{short_a} − {short_b}", diff)
        print(f"    (individual ranges: [{a_aligned.min():.4f}, {a_aligned.max():.4f}] vs [{b_aligned.min():.4f}, {b_aligned.max():.4f}])")
    print()

    print("=== Correctness differentials (correct − incorrect, split CSVs) ===")
    for key in BENCHMARKS:
        split_path_c = SPLIT_DIR / f"{key}_bin_scores_correct.csv"
        split_path_i = SPLIT_DIR / f"{key}_bin_scores_incorrect.csv"
        if not split_path_c.exists() or not split_path_i.exists():
            print(f"  {key.removeprefix('queries_')}: split files not found, skipping")
            continue
        df_c = load_split(key, "correct")
        df_i = load_split(key, "incorrect")
        vals_c = df_c["mean_score"].values * SCALE
        vals_i = df_i["mean_score"].values * SCALE
        diff = vals_c - vals_i
        short = key.removeprefix("queries_")
        report_range(f"{short} correct", vals_c)
        report_range(f"{short} incorrect", vals_i)
        report_range(f"{short} diff", diff)
        print()

    print("=== Column format check (split vs perquery) ===")
    for key in BENCHMARKS:
        pq_df = grids[key]
        split_path_c = SPLIT_DIR / f"{key}_bin_scores_correct.csv"
        if not split_path_c.exists():
            continue
        sp_df = pd.read_csv(split_path_c)
        pq_col = "median_influence" if "median_influence" in pd.read_csv(
            PERQUERY_DIR / f"{key}_bin_scores_perquery.csv", nrows=0
        ).columns else "mean_score"
        sp_col = "mean_score"
        short = key.removeprefix("queries_")
        print(f"  {short}: perquery uses '{pq_col}', split uses '{sp_col}'")
        pq_sample = pd.read_csv(PERQUERY_DIR / f"{key}_bin_scores_perquery.csv").head(3)[pq_col].values * SCALE
        sp_sample = sp_df.head(3)[sp_col].values * SCALE
        print(f"    perquery first 3 (×1e6): {np.round(pq_sample, 4)}")
        print(f"    split correct first 3 (×1e6): {np.round(sp_sample, 4)}")


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