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

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false, reportArgumentType=false, reportCallIssue=false, reportGeneralTypeIssues=false
"""Bootstrap CIs, format-cluster separation per benchmark, GMM diagnostic.

Inputs:
  --per-doc            Glob to per_doc_features.parquet partial files
  --per-bin            Path to per_bin_profiles.parquet (from toolkit merge)
  --top-bins           Path to top_bins_lexical_profiles.parquet
  --format-clusters    YAML with interpersonal_dialogue / documentation_structural
                       lists (configs/rq4_format_clusters.yaml).
  --bin-scores-dir     Directory with zscored_<benchmark>.csv files.
  --benchmarks         Comma list of benchmark ids; default uses the listed map.
  --top-n              Top-N bins per benchmark for cluster intersection (default 20).
  --out                Output directory.
  --n-bootstrap        Bootstrap iterations (default 1000).
  --seed               Random seed (default 42).
  --gmm                Also run GMM(k=1..3) diagnostic on SocialIQA top-10 and top-25.

Writes:
  cluster_separation_bootstrap_all_benchmarks.csv
  cluster_separation_socialiqa_vs_others.csv
  gmm_bimodality_diagnostic.json
  top_bins_with_cis.parquet
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))

from data_attribution.rq4_lexical.bootstrap import (  # noqa: E402
    bootstrap_bin_densities,
    bootstrap_words_per_doc,
    gmm_bimodality,
)


# Any benchmark with a zscored_<id>.csv in --bin-scores-dir works even if it is
# not listed here (the loader falls back to the zscored_<id>.csv convention).
BENCHMARK_FILES = {
    "socialiqa": "zscored_socialiqa.csv",
    "mmlu_social_science": "zscored_mmlu_social_science.csv",
    "arc_challenge": "zscored_arc_challenge.csv",
    "mmlu_stem": "zscored_mmlu_stem.csv",
    "bbh_snarks_instruct": "zscored_bbh_snarks_instruct.csv",
    "bbh_causal_judgement_instruct": "zscored_bbh_causal_judgement_instruct.csv",
    "bbh_sports_understanding_instruct": "zscored_bbh_sports_understanding_instruct.csv",
}

FEATURES_FOR_CLUSTER = [
    "mean_words_per_doc",
    "mental_state_per_1k",
    "dialogue_composite_z",
    "empath_social_z",
    "empath_affect_z",
]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--per-doc", default=None)
    p.add_argument("--per-bin", required=True)
    p.add_argument("--top-bins", default=None)
    p.add_argument("--format-clusters", required=True)
    p.add_argument("--bin-scores-dir", required=True)
    p.add_argument(
        "--roster",
        default=None,
        help="configs/rq4_paper_benchmarks.yaml; drives benchmark "
        "set. Overridden by an explicit --benchmarks.",
    )
    p.add_argument("--benchmarks", default=None)
    p.add_argument("--top-n", type=int, default=20)
    p.add_argument("--out", required=True)
    p.add_argument("--n-bootstrap", type=int, default=1000)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--gmm", action="store_true")
    return p.parse_args()


def _load_per_doc(glob: str) -> pd.DataFrame:
    paths = sorted(Path().glob(glob)) if any(c in glob for c in "*?[") else [Path(glob)]
    dfs = [pd.read_parquet(p) for p in paths]
    return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame()


def _load_format_clusters(path: Path):
    from data_attribution.rq4_lexical.format_clusters import load_format_clusters

    return load_format_clusters(path)


def _load_top_bins(bin_scores_dir: Path, benchmark: str, top_n: int) -> pd.DataFrame:
    fname = BENCHMARK_FILES.get(benchmark)
    if fname is None:
        # holdout benchmark — assume zscored_<benchmark>.csv naming
        fname = f"zscored_{benchmark}.csv"
    df = pd.read_csv(bin_scores_dir / fname)
    df = df.rename(columns={"topic_label": "bin_topic", "format_label": "bin_format"})
    df = df.sort_values("zscore", ascending=False).head(top_n).reset_index(drop=True)
    df["rank"] = df.index + 1
    df["benchmark"] = benchmark
    return df


def _bootstrap_delta(
    inter_vals: np.ndarray,
    docs_vals: np.ndarray,
    n_boot: int,
    seed: int,
) -> tuple[float, float, float]:
    rng = np.random.default_rng(seed)
    delta_point = float(np.nanmean(inter_vals) - np.nanmean(docs_vals))
    if inter_vals.size == 0 or docs_vals.size == 0:
        return delta_point, float("nan"), float("nan")
    i_idx = rng.integers(0, inter_vals.size, size=(n_boot, inter_vals.size))
    d_idx = rng.integers(0, docs_vals.size, size=(n_boot, docs_vals.size))
    delta_boot = np.nanmean(inter_vals[i_idx], axis=1) - np.nanmean(
        docs_vals[d_idx], axis=1
    )
    lo, hi = np.nanpercentile(delta_boot, [2.5, 97.5])
    return delta_point, float(lo), float(hi)


def per_benchmark_cluster_separation(
    per_bin: pd.DataFrame,
    bin_scores_dir: Path,
    benchmarks: list[str],
    top_n: int,
    clusters,
    n_boot: int,
    seed: int,
) -> pd.DataFrame:
    """Within-benchmark interactional - expository separation (interactional axis).

    `clusters` is a FormatClusters; the binary contrast uses the
    interactional_axis collapse (interactional = dialogic+personal,
    expository = expository+structured; news/boilerplate excluded)."""
    axis = clusters.axis_format_sets()
    inter_fmts = axis.get("interactional", set())
    expo_fmts = axis.get("expository", set())
    rows: list[dict] = []
    for bench in benchmarks:
        top = _load_top_bins(bin_scores_dir, bench, top_n)
        joined = top.merge(per_bin, on=["bin_topic", "bin_format"], how="left")
        inter = joined[joined["bin_format"].isin(inter_fmts)]
        docs = joined[joined["bin_format"].isin(expo_fmts)]
        for feature in FEATURES_FOR_CLUSTER:
            inter_vals = (
                inter[feature].to_numpy(dtype=float)
                if feature in joined.columns
                else np.array([])
            )
            docs_vals = (
                docs[feature].to_numpy(dtype=float)
                if feature in joined.columns
                else np.array([])
            )
            delta, lo, hi = _bootstrap_delta(inter_vals, docs_vals, n_boot, seed)
            rows.append(
                {
                    "benchmark": bench,
                    "top_n": top_n,
                    "feature": feature,
                    "delta_point": delta,
                    "ci_low": lo,
                    "ci_high": hi,
                    "interactional_mean": float(np.nanmean(inter_vals))
                    if inter_vals.size
                    else float("nan"),
                    "expository_mean": float(np.nanmean(docs_vals))
                    if docs_vals.size
                    else float("nan"),
                    "n_interactional_bins": int(inter.shape[0]),
                    "n_expository_bins": int(docs.shape[0]),
                    "n_excluded_bins": int(
                        top.shape[0] - inter.shape[0] - docs.shape[0]
                    ),
                }
            )
    return pd.DataFrame(rows)


def cross_benchmark_contrast(
    per_bin: pd.DataFrame,
    bin_scores_dir: Path,
    benchmarks: list[str],
    top_n: int,
    n_boot: int,
    seed: int,
    target: str = "socialiqa",
    others: list[str] | None = None,
    label: str = "socialiqa_vs_others",
) -> pd.DataFrame:
    r"""One benchmark's top-N profile vs. the pooled top-N of a set of others.

    Tests whether ``target``'s high-influence bins are, as a group, more
    interpersonal / affective than the union of ``others``' top-N bins.
    ``others`` defaults to every other benchmark; for a meaningful contrast
    pass a specific comparison set (e.g. the knowledge-recall group), since
    pooling reasoning benchmarks dilutes the signal.
    """
    if target not in benchmarks:
        return pd.DataFrame()
    others = others if others is not None else [b for b in benchmarks if b != target]
    target_top = _load_top_bins(bin_scores_dir, target, top_n)
    target_joined = target_top.merge(
        per_bin, on=["bin_topic", "bin_format"], how="left"
    )
    other_frames = [_load_top_bins(bin_scores_dir, b, top_n) for b in others]
    other_top = pd.concat(other_frames, ignore_index=True).drop_duplicates(
        subset=["bin_topic", "bin_format"]
    )
    other_joined = other_top.merge(per_bin, on=["bin_topic", "bin_format"], how="left")

    rows: list[dict] = []
    for feature in FEATURES_FOR_CLUSTER:
        t_vals = (
            target_joined[feature].to_numpy(dtype=float)
            if feature in target_joined.columns
            else np.array([])
        )
        o_vals = (
            other_joined[feature].to_numpy(dtype=float)
            if feature in other_joined.columns
            else np.array([])
        )
        delta, lo, hi = _bootstrap_delta(t_vals, o_vals, n_boot, seed)
        rows.append(
            {
                "contrast": label,
                "target_benchmark": target,
                "other_benchmarks": ",".join(others),
                "top_n": top_n,
                "feature": feature,
                "delta_point": delta,
                "ci_low": lo,
                "ci_high": hi,
                "target_mean": float(np.nanmean(t_vals))
                if t_vals.size
                else float("nan"),
                "others_mean": float(np.nanmean(o_vals))
                if o_vals.size
                else float("nan"),
                "n_target_bins": int(target_joined.shape[0]),
                "n_other_bins": int(other_joined.shape[0]),
            }
        )
    return pd.DataFrame(rows)


def group_profile_contrast(
    per_bin: pd.DataFrame,
    bin_scores_dir: Path,
    groups: dict[str, list[str]],
    top_n: int,
    n_boot: int,
    seed: int,
) -> pd.DataFrame:
    """Pool each roster group's top-N bins and bootstrap pairwise group deltas.

    `groups` maps group name -> list of benchmark ids. Emits, for every ordered
    pair of groups (A, B) and every composite feature, the bootstrap delta
    mean(A) - mean(B) over the pooled, de-duplicated top-N bins of each group.
    Roster-driven, so new benchmarks/groups need no code change."""
    pooled: dict[str, pd.DataFrame] = {}
    for gname, bids in groups.items():
        if not bids:
            continue
        frames = [_load_top_bins(bin_scores_dir, b, top_n) for b in bids]
        pooled_top = pd.concat(frames, ignore_index=True).drop_duplicates(
            subset=["bin_topic", "bin_format"]
        )
        pooled[gname] = pooled_top.merge(
            per_bin, on=["bin_topic", "bin_format"], how="left"
        )

    rows: list[dict] = []
    gnames = list(pooled.keys())
    for i, a in enumerate(gnames):
        for b in gnames[i + 1 :]:
            for feature in FEATURES_FOR_CLUSTER:
                a_vals = (
                    pooled[a][feature].to_numpy(dtype=float)
                    if feature in pooled[a].columns
                    else np.array([])
                )
                b_vals = (
                    pooled[b][feature].to_numpy(dtype=float)
                    if feature in pooled[b].columns
                    else np.array([])
                )
                delta, lo, hi = _bootstrap_delta(a_vals, b_vals, n_boot, seed)
                rows.append(
                    {
                        "group_a": a,
                        "group_b": b,
                        "top_n": top_n,
                        "feature": feature,
                        "delta_point": delta,
                        "ci_low": lo,
                        "ci_high": hi,
                        "group_a_mean": float(np.nanmean(a_vals))
                        if a_vals.size
                        else float("nan"),
                        "group_b_mean": float(np.nanmean(b_vals))
                        if b_vals.size
                        else float("nan"),
                        "n_a_bins": int(pooled[a].shape[0]),
                        "n_b_bins": int(pooled[b].shape[0]),
                    }
                )
    return pd.DataFrame(rows)


def _gmm_socialiqa(
    per_bin: pd.DataFrame, bin_scores_dir: Path, seed: int
) -> dict[str, object]:
    soc = pd.read_csv(bin_scores_dir / "zscored_socialiqa.csv")
    soc = soc.rename(columns={"topic_label": "bin_topic", "format_label": "bin_format"})
    soc_sorted = soc.sort_values("zscore", ascending=False)
    out: dict[str, object] = {}
    gmm_features = [
        "mean_words_per_doc",
        "mental_state_per_1k",
        "dialogue_composite_z",
        "empath_social_z",
    ]
    for n in (10, 25):
        sub = soc_sorted.head(n).merge(per_bin, on=["bin_topic", "bin_format"])
        out[f"socialiqa_top{n}"] = gmm_bimodality(sub, gmm_features, seed=seed)
    return out


def top_bin_cis(
    per_doc: pd.DataFrame,
    top_bins: pd.DataFrame,
    n_boot: int,
    seed: int,
) -> pd.DataFrame:
    bin_keys = list(zip(top_bins["bin_topic"], top_bins["bin_format"]))
    unique_keys = list(dict.fromkeys(bin_keys))
    words_ci = bootstrap_words_per_doc(per_doc, unique_keys, n_boot, seed)
    density_features = [
        "first_person_count",
        "second_person_count",
        "third_person_count",
        "mental_state_count",
        "question_mark_count",
        "quote_mark_count",
        "speaker_turn_count",
    ]
    density_cis: dict[str, dict[tuple[str, str], tuple[float, float, float]]] = {}
    for col in density_features:
        density_cis[col] = bootstrap_bin_densities(
            per_doc,
            unique_keys,
            col,
            n_boot,
            seed,
        )
    rows: list[dict] = []
    for _, r in top_bins.iterrows():
        bin_key = (r["bin_topic"], r["bin_format"])
        row = dict(r)
        wp, wl, wh = words_ci[bin_key]
        row["mean_words_per_doc_point"] = wp
        row["mean_words_per_doc_ci_low"] = wl
        row["mean_words_per_doc_ci_high"] = wh
        for col, ci_map in density_cis.items():
            point, lo, hi = ci_map[bin_key]
            base = col[: -len("_count")]
            row[f"{base}_per_1k_point"] = point
            row[f"{base}_per_1k_ci_low"] = lo
            row[f"{base}_per_1k_ci_high"] = hi
        rows.append(row)
    return pd.DataFrame(rows)


def _resolve_benchmarks(args) -> list[str]:
    if args.benchmarks:
        return [b.strip() for b in args.benchmarks.split(",") if b.strip()]
    if args.roster:
        from data_attribution.rq4_lexical.roster import load_roster

        return load_roster(Path(args.roster)).ids
    return list(BENCHMARK_FILES.keys())


def main() -> None:
    args = parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    benchmarks = _resolve_benchmarks(args)
    roster_groups: dict[str, list[str]] = {}
    if args.roster:
        from data_attribution.rq4_lexical.roster import load_roster

        roster = load_roster(Path(args.roster))
        for gname in {e.group for e in roster.entries}:
            roster_groups[gname] = roster.ids_in_group(gname)

    per_bin = pd.read_parquet(args.per_bin)
    clusters = _load_format_clusters(Path(args.format_clusters))
    bin_scores_dir = Path(args.bin_scores_dir)

    per_bench = per_benchmark_cluster_separation(
        per_bin,
        bin_scores_dir,
        benchmarks,
        args.top_n,
        clusters,
        args.n_bootstrap,
        args.seed,
    )
    per_bench_path = out / "cluster_separation_bootstrap_all_benchmarks.csv"
    per_bench.to_csv(per_bench_path, index=False)
    print(f"wrote {per_bench_path}", flush=True)

    # Specific, non-diluted contrast: SocialIQA vs the knowledge-recall group.
    knowledge = roster_groups.get("knowledge_recall") or [
        b
        for b in benchmarks
        if b in ("mmlu_social_science", "arc_challenge", "mmlu_stem")
    ]
    cross = cross_benchmark_contrast(
        per_bin,
        bin_scores_dir,
        benchmarks,
        args.top_n,
        args.n_bootstrap,
        args.seed,
        target="socialiqa",
        others=knowledge,
        label="socialiqa_vs_knowledge_recall",
    )
    cross_path = out / "cluster_separation_socialiqa_vs_others.csv"
    cross.to_csv(cross_path, index=False)
    print(f"wrote {cross_path}", flush=True)

    # Roster-group pairwise profile contrast (reasoning vs knowledge etc.).
    if roster_groups:
        gc = group_profile_contrast(
            per_bin,
            bin_scores_dir,
            roster_groups,
            args.top_n,
            args.n_bootstrap,
            args.seed,
        )
        gc_path = out / "cluster_separation_group_contrast.csv"
        gc.to_csv(gc_path, index=False)
        print(f"wrote {gc_path}", flush=True)

    if args.gmm:
        gmm = _gmm_socialiqa(per_bin, bin_scores_dir, seed=args.seed)
        gmm_path = out / "gmm_bimodality_diagnostic.json"
        with gmm_path.open("w") as fh:
            json.dump(gmm, fh, indent=2)
        print(f"wrote {gmm_path}", flush=True)

    if args.per_doc and args.top_bins:
        per_doc = _load_per_doc(args.per_doc)
        if per_doc.empty:
            print("[bootstrap] no per_doc rows; skipping top-bin CIs", flush=True)
            return
        top = pd.read_parquet(args.top_bins)
        cis = top_bin_cis(per_doc, top, args.n_bootstrap, args.seed)
        out_top = out / "top_bins_with_cis.parquet"
        cis.to_parquet(out_top, index=False)
        print(f"wrote {out_top}", flush=True)


if __name__ == "__main__":
    main()