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

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false, reportCallIssue=false
"""Visualization suite for RQ4 lexical results (appendix-bound).

Roster-driven; reads per_bin_profiles + top_bins + the bootstrap CSVs and the
format-cluster YAML. Emits several plot types for review:

  fig-lex-cluster-composition.pdf   stacked bars: top-N format-cluster mix / benchmark
  fig-lex-group-contrast-forest.pdf CI forest of pooled group deltas / composite
  fig-lex-group-profile.pdf         grouped bars: mean composite z per roster group
  fig-lex-benchmark-composite-dots  per-benchmark mean composite z (dot plot)
  fig-lex-scatter-words-mental.pdf  per-bin words vs mental-state, coloured by cluster

All plots use the Okabe-Ito colourblind-safe palette. Nothing is benchmark-
hardcoded; adding a benchmark to the roster extends every plot automatically.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
matplotlib.rcParams["pdf.fonttype"] = 42
matplotlib.rcParams["ps.fonttype"] = 42
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

import sys  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
from data_attribution.rq4_lexical.roster import load_roster  # noqa: E402

OKABE_ITO = {
    # roster groups
    "social_reasoning": "#009E73",  # green
    "commonsense_reasoning": "#CC79A7",  # purple
    "knowledge_recall": "#E69F00",  # orange
}
# Colours for the 6 interactional format groups (Okabe-Ito), interactional -> expository.
GROUP_COLORS = {
    "dialogic": "#0072B2",  # blue
    "personal": "#56B4E9",  # sky blue
    "expository": "#D55E00",  # vermillion
    "structured": "#E69F00",  # orange
    "news": "#009E73",  # green
    "boilerplate": "#999999",  # grey
}
GROUP_LEGEND = {
    "dialogic": "dialogic",
    "personal": "personal",
    "expository": "expository",
    "structured": "structured",
    "news": "news",
    "boilerplate": "boilerplate",
}
COMPOSITES = [
    ("mean_words_per_doc", "mean words/doc"),
    ("mental_state_per_1k", "mental-state /1k"),
    ("dialogue_composite_z", "dialogue (z)"),
    ("empath_social_z", "social (z)"),
    ("empath_affect_z", "affect (z)"),
]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--per-bin", required=True)
    p.add_argument(
        "--top-bins", required=True, help="top_bins_lexical_profiles.parquet"
    )
    p.add_argument("--group-contrast", required=True)
    p.add_argument(
        "--per-benchmark-sep",
        required=True,
        help="cluster_separation_bootstrap_all_benchmarks.csv",
    )
    p.add_argument("--format-clusters", required=True)
    p.add_argument("--roster", required=True)
    p.add_argument("--out-dir", required=True)
    return p.parse_args()


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

    return load_format_clusters(path)


def plot_cluster_composition(top: pd.DataFrame, clusters, roster, out: Path) -> None:
    """6-way stacked bars: count of each benchmark's top-N bins per format group."""
    fmt_to_group = clusters.format_to_group()
    groups = clusters.ordered_groups()
    bids = [b for b in roster.ids if b in set(top["benchmark"])]
    labels = [roster.shorts.get(b, b) for b in bids]
    counts = {g: [] for g in groups}
    for b in bids:
        sub = top[top.benchmark == b]
        grp_of = sub["bin_format"].map(fmt_to_group)
        for g in groups:
            counts[g].append(int((grp_of == g).sum()))
    fig, ax = plt.subplots(figsize=(max(7, 1.0 * len(labels) + 2), 4.2))
    x = np.arange(len(labels))
    bottom = np.zeros(len(labels))
    for g in groups:
        vals = np.array(counts[g], dtype=float)
        ax.bar(
            x,
            vals,
            bottom=bottom,
            color=GROUP_COLORS.get(g, "#777777"),
            label=GROUP_LEGEND.get(g, g),
        )
        bottom += vals
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
    ax.set_ylabel("top-20 bins")
    ax.set_title("Format-cluster composition of top-20 high-influence bins")
    ax.legend(fontsize=8, loc="upper right")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def plot_group_contrast_forest(gc: pd.DataFrame, out: Path) -> None:
    pairs = gc[["group_a", "group_b"]].drop_duplicates().itertuples(index=False)
    pairs = list(pairs)
    fig, axes = plt.subplots(
        1, len(COMPOSITES), figsize=(3.0 * len(COMPOSITES), 3.2), sharey=True
    )
    for ax, (feat, fl) in zip(axes, COMPOSITES):
        sub = gc[gc.feature == feat]
        ys = np.arange(len(pairs))
        for y, (ga, gb) in zip(ys, pairs):
            row = sub[(sub.group_a == ga) & (sub.group_b == gb)]
            if row.empty:
                continue
            r = row.iloc[0]
            sig = (r.ci_low > 0) or (r.ci_high < 0)
            color = "#000000" if sig else "#999999"
            ax.plot([r.ci_low, r.ci_high], [y, y], color=color, lw=2)
            ax.plot(r.delta_point, y, "o", color=color, ms=5)
        ax.axvline(0, color="#D55E00", ls="--", lw=1)
        ax.set_title(fl, fontsize=9)
        ax.set_yticks(ys)
        ax.set_yticklabels([f"{a[:10]}\nvs {b[:10]}" for a, b in pairs], fontsize=7)
    fig.suptitle(
        "Pooled group contrast (Δ = mean A − mean B), 95% bootstrap CI", fontsize=10
    )
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def plot_benchmark_profile(top: pd.DataFrame, roster, out: Path) -> None:
    """Per-benchmark grouped bars: one cluster per benchmark, four composite bars.

    Composites are z-scores against all 576 bins, averaged over that benchmark's
    top-20 most-influential bins (0 = corpus-average bin)."""
    composites = [
        ("mental_state_z", "mental-state (z)"),
        ("dialogue_composite_z", "dialogue (z)"),
        ("empath_social_z", "social (z)"),
        ("empath_affect_z", "affect (z)"),
    ]
    comp_colors = ["#000000", "#0072B2", "#009E73", "#CC79A7"]
    bids = [b for b in roster.ids if b in set(top["benchmark"])]
    labels = [roster.shorts.get(b, b) for b in bids]
    fig, ax = plt.subplots(figsize=(max(7, 1.1 * len(bids) + 2), 4.2))
    width = 0.8 / len(composites)
    x = np.arange(len(bids))
    for i, (feat, fl) in enumerate(composites):
        means = [top[top.benchmark == b][feat].mean(skipna=True) for b in bids]
        ax.bar(x + i * width, means, width, label=fl, color=comp_colors[i])
    ax.axhline(0, color="black", lw=0.8)
    ax.set_xticks(x + width * (len(composites) - 1) / 2)
    ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8)
    ax.set_ylabel("mean $z$ over top-20 bins (0 = corpus-average bin)")
    ax.set_title(
        "Lexical profile per benchmark (averaged over its top-20 high-influence bins)"
    )
    ax.legend(fontsize=8, ncol=4, loc="upper center", bbox_to_anchor=(0.5, -0.18))
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def plot_benchmark_composite_dots(top: pd.DataFrame, roster, out: Path) -> None:
    feats = [
        ("dialogue_composite_z", "dialogue z"),
        ("empath_social_z", "social z"),
        ("empath_affect_z", "affect z"),
    ]
    bids = [b for b in roster.ids if b in set(top["benchmark"])]
    fig, ax = plt.subplots(figsize=(7, max(4, 0.4 * len(bids) + 1)))
    ys = np.arange(len(bids))
    markers = {"dialogue z": "o", "social z": "s", "affect z": "^"}
    for feat, fl in feats:
        means = [top[top.benchmark == b][feat].mean(skipna=True) for b in bids]
        grp_colors = [OKABE_ITO.get(roster.groups.get(b), "#000000") for b in bids]
        ax.scatter(means, ys, c=grp_colors, marker=markers[fl], s=40, label=fl)
    ax.axvline(0, color="#D55E00", ls="--", lw=1)
    ax.set_yticks(ys)
    ax.set_yticklabels([roster.shorts.get(b, b) for b in bids], fontsize=8)
    ax.set_xlabel("mean composite z over top-20 bins")
    ax.set_title(
        "Composite lexical loadings per benchmark\n(colour = group, marker = composite)"
    )
    ax.legend(fontsize=8, loc="lower right")
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def plot_scatter_words_mental(top: pd.DataFrame, clusters, out: Path) -> None:
    fig, ax = plt.subplots(figsize=(6.5, 5))
    fmt_to_group = clusters.format_to_group()
    grp_of = top["bin_format"].map(fmt_to_group)
    for g in clusters.ordered_groups():
        sub = top[grp_of == g]
        if sub.empty:
            continue
        ax.scatter(
            sub["mean_words_per_doc"],
            sub["mental_state_per_1k"],
            c=GROUP_COLORS.get(g, "#777777"),
            label=GROUP_LEGEND.get(g, g),
            alpha=0.7,
            s=28,
            edgecolors="none",
        )
    ax.set_xscale("log")
    ax.set_xlabel("mean words / doc (log)")
    ax.set_ylabel("mental-state verbs / 1k words")
    ax.set_title("Top-20 bins (all benchmarks): length vs mental-state density")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def main() -> None:
    args = parse_args()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    roster = load_roster(Path(args.roster))
    top = pd.read_parquet(args.top_bins)
    per_bin = pd.read_parquet(args.per_bin)
    if "mental_state_z" in per_bin.columns and "mental_state_z" not in top.columns:
        top = top.merge(
            per_bin[["bin_topic", "bin_format", "mental_state_z"]],
            on=["bin_topic", "bin_format"],
            how="left",
        )
    gc = pd.read_csv(args.group_contrast)
    sep = pd.read_csv(args.per_benchmark_sep)
    clusters = _clusters(Path(args.format_clusters))

    del sep  # composition is now computed from top-bins + taxonomy, not the sep CSV
    plot_cluster_composition(
        top, clusters, roster, out_dir / "fig-lex-cluster-composition.pdf"
    )
    plot_group_contrast_forest(gc, out_dir / "fig-lex-group-contrast-forest.pdf")
    plot_benchmark_profile(top, roster, out_dir / "fig-lex-benchmark-profile.pdf")
    plot_benchmark_composite_dots(
        top, roster, out_dir / "fig-lex-benchmark-composite-dots.pdf"
    )
    plot_scatter_words_mental(
        top, clusters, out_dir / "fig-lex-scatter-words-mental.pdf"
    )


if __name__ == "__main__":
    main()