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

#!/usr/bin/env python3
# pyright: reportArgumentType=false, reportCallIssue=false
"""Curated open-LIWC-22 appendix artifacts (offline; not yet wired into the PDF).

From the wide per-bin profile (open LIWC-22 schema) and the existing influence
top-bin selection, emit the curated paper artifacts:

  fig-lex-summary-vars.pdf  per-benchmark open summary variables
                            (Analytic-open / Tone-open / Clout-approx / Authentic-approx)
  fig-liwc-families.pdf     benchmark-group x LIWC-family mean-z heatmap
  tab-liwc-families.tex     benchmark-group x LIWC-family mean-z table (booktabs)

All z-scores are full-grid (576-bin) densities averaged over each benchmark's
top-N influential bins; group columns pool a group's benchmarks' top bins,
de-duplicated, matching the headline group contrast. The four headline
composites are unchanged --- these dimensions are additive.
"""

from __future__ import annotations

import argparse
import sys
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

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

SUMMARY_VARS = [
    ("analytic_open", "Analytic-open"),
    ("tone_open", "Tone-open"),
    ("clout_approx", "Clout-approx"),
    ("authentic_approx", "Authentic-approx"),
]
SUMMARY_COLORS = ["#000000", "#0072B2", "#009E73", "#CC79A7"]

FAMILIES: dict[str, list[str]] = {
    "Function words": ["fw_function_z"],
    "Cognition": ["cog_cogproc_z"],
    "Affect": ["empath_affect_z"],
    "Social": ["empath_social_z"],
    "Drives": [
        "liwc_affiliation_z",
        "liwc_achieve_z",
        "liwc_power_z",
        "liwc_reward_z",
        "liwc_allure_z",
    ],
    "Perception": [
        "liwc_motion_z",
        "liwc_auditory_z",
        "perc_visual_z",
        "perc_feeling_z",
        "perc_space_z",
        "perc_attention_z",
    ],
}
GROUP_LABEL = {
    "social_reasoning": "Social reasoning",
    "commonsense_reasoning": "Commonsense reasoning",
    "knowledge_recall": "Knowledge recall",
}
GROUP_ORDER = ["social_reasoning", "commonsense_reasoning", "knowledge_recall"]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--per-bin", required=True)
    p.add_argument("--top-bins", required=True)
    p.add_argument("--roster", required=True)
    p.add_argument("--out-dir", required=True)
    return p.parse_args()


def add_family_means(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    for fam, cols in FAMILIES.items():
        present = [c for c in cols if c in out.columns]
        out[f"fam::{fam}"] = (
            out[present].mean(axis=1, skipna=True) if present else np.nan
        )
    return out


def plot_summary_vars(top: pd.DataFrame, roster, out: Path) -> None:
    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.2 * len(bids) + 2), 4.4))
    width = 0.8 / len(SUMMARY_VARS)
    x = np.arange(len(bids))
    for i, (col, lab) in enumerate(SUMMARY_VARS):
        means = [top[top.benchmark == b][col].mean(skipna=True) for b in bids]
        ax.bar(x + i * width, means, width, label=lab, color=SUMMARY_COLORS[i])
    ax.axhline(50, color="#D55E00", ls="--", lw=1, label="neutral (50)")
    ax.set_xticks(x + width * (len(SUMMARY_VARS) - 1) / 2)
    ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8)
    ax.set_ylabel("open summary score (0--100; 50 = corpus-neutral)")
    ax.set_title(
        "Open LIWC-22 summary variables per benchmark\n(mean over top-20 high-influence bins; approximations, not LIWC-identical)"
    )
    ax.legend(fontsize=8, ncol=5, loc="upper center", bbox_to_anchor=(0.5, -0.16))
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def group_family_means(top: pd.DataFrame, roster) -> pd.DataFrame:
    rows: dict[str, dict[str, float]] = {}
    for g in GROUP_ORDER:
        ids = [b for b in roster.ids_in_group(g) if b in set(top["benchmark"])]
        if not ids:
            continue
        pooled = top[top.benchmark.isin(ids)].drop_duplicates(
            subset=["bin_topic", "bin_format"]
        )
        rows[g] = {
            fam: float(pooled[f"fam::{fam}"].mean(skipna=True)) for fam in FAMILIES
        }
    return pd.DataFrame(rows).T.reindex([g for g in GROUP_ORDER if g in rows])


def benchmark_family_means(top: pd.DataFrame, roster) -> pd.DataFrame:
    bids = [b for b in roster.ids if b in set(top["benchmark"])]
    rows = {
        b: {
            fam: float(top[top.benchmark == b][f"fam::{fam}"].mean(skipna=True))
            for fam in FAMILIES
        }
        for b in bids
    }
    return pd.DataFrame(rows).T.reindex(bids)


def plot_families_heatmap(
    mat_df: pd.DataFrame, row_labels: list[str], title: str, out: Path
) -> None:
    fams = list(FAMILIES)
    mat = mat_df[fams].to_numpy(dtype=float)
    fig, ax = plt.subplots(figsize=(1.2 * len(fams) + 2, 0.6 * len(mat_df) + 2))
    vmax = float(np.nanmax(np.abs(mat))) or 1.0
    im = ax.imshow(mat, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
    ax.set_xticks(range(len(fams)))
    ax.set_xticklabels(fams, rotation=30, ha="right", fontsize=8)
    ax.set_yticks(range(len(mat_df)))
    ax.set_yticklabels(row_labels, fontsize=8)
    for i in range(mat.shape[0]):
        for j in range(mat.shape[1]):
            ax.text(
                j,
                i,
                f"{mat[i, j]:+.2f}",
                ha="center",
                va="center",
                fontsize=7,
                color="black" if abs(mat[i, j]) < 0.6 * vmax else "white",
            )
    fig.colorbar(im, ax=ax, shrink=0.8, label="mean full-grid $z$")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def plot_summary_vars_bygroup(top: pd.DataFrame, roster, out: Path) -> None:
    groups = [g for g in GROUP_ORDER if roster.ids_in_group(g)]
    fig, ax = plt.subplots(figsize=(max(6, 1.4 * len(groups) + 2), 4.4))
    width = 0.8 / len(SUMMARY_VARS)
    x = np.arange(len(groups))
    for i, (col, lab) in enumerate(SUMMARY_VARS):
        means = []
        for g in groups:
            ids = [b for b in roster.ids_in_group(g) if b in set(top["benchmark"])]
            pooled = top[top.benchmark.isin(ids)].drop_duplicates(
                subset=["bin_topic", "bin_format"]
            )
            means.append(pooled[col].mean(skipna=True))
        ax.bar(x + i * width, means, width, label=lab, color=SUMMARY_COLORS[i])
    ax.axhline(50, color="#D55E00", ls="--", lw=1, label="neutral (50)")
    ax.set_xticks(x + width * (len(SUMMARY_VARS) - 1) / 2)
    ax.set_xticklabels([GROUP_LABEL[g] for g in groups], fontsize=9)
    ax.set_ylabel("open summary score (0--100; 50 = corpus-neutral)")
    ax.set_title(
        "Open LIWC-22 summary variables by benchmark group\n(mean over each group's pooled top-20 bins)"
    )
    ax.legend(fontsize=8, ncol=5, loc="upper center", bbox_to_anchor=(0.5, -0.12))
    fig.tight_layout()
    fig.savefig(out, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}", flush=True)


def write_families_table(grp: pd.DataFrame, out: Path) -> None:
    fams = list(FAMILIES)
    lines = [
        "% Auto-generated by scripts/analysis/rq4_liwc_appendix_artifacts.py. Do not edit by hand.",
        "\\begin{table}[t]",
        "  \\centering",
        "  \\small",
        "  \\begin{tabular}{l" + "r" * len(fams) + "}",
        "    \\toprule",
        "    Benchmark group & " + " & ".join(fams) + " \\\\",
        "    \\midrule",
    ]
    for g in grp.index:
        cells = " & ".join(f"${grp.loc[g, fam]:+.2f}$" for fam in fams)
        lines.append(f"    {GROUP_LABEL.get(g, g)} & {cells} \\\\")
    lines += [
        "    \\bottomrule",
        "  \\end{tabular}",
        "  \\caption{Open LIWC-22 family loadings by benchmark group: mean full-grid "
        "$z$-score (against all 576 WebOrganizer bins) over each group's pooled top-20 "
        "high-influence bins. Positive values indicate the group's high-influence "
        "training text is denser in that family than the corpus average. Families are "
        "open re-implementations (Empath plus in-repo closed-class lexicons); see "
        "Appendix~\\ref{app:rq4-lexical-enriched}.}",
        "  \\label{tab:liwc-families}",
        "\\end{table}",
        "",
    ]
    out.write_text("\n".join(lines))
    print(f"wrote {out}", flush=True)


def main() -> None:
    args = parse_args()
    per_bin = pd.read_parquet(args.per_bin)
    top_sel = pd.read_parquet(args.top_bins)[
        ["benchmark", "bin_topic", "bin_format", "rank"]
    ]
    roster = load_roster(Path(args.roster))
    top = add_family_means(
        top_sel.merge(per_bin, on=["bin_topic", "bin_format"], how="left")
    )

    out_dir = Path(args.out_dir)
    fig_dir = out_dir / "figures"
    tab_dir = out_dir / "tables"
    fig_dir.mkdir(parents=True, exist_ok=True)
    tab_dir.mkdir(parents=True, exist_ok=True)

    top.to_parquet(out_dir / "top_bins_liwc_profiles.parquet", index=False)

    # Summary variables: per-benchmark and per-group variants.
    plot_summary_vars(top, roster, fig_dir / "fig-lex-summary-vars.pdf")
    plot_summary_vars_bygroup(top, roster, fig_dir / "fig-lex-summary-vars-bygroup.pdf")

    # LIWC families: per-group (table + heatmap) and per-benchmark heatmap variant.
    grp = group_family_means(top, roster)
    grp.to_csv(out_dir / "liwc_families_by_group.csv")
    plot_families_heatmap(
        grp,
        [str(GROUP_LABEL.get(g, g)) for g in grp.index],
        "LIWC-family loadings by benchmark group\n(mean $z$ over pooled top-20 bins)",
        fig_dir / "fig-liwc-families.pdf",
    )
    bench = benchmark_family_means(top, roster)
    plot_families_heatmap(
        bench,
        [str(roster.shorts.get(b, b)) for b in bench.index],
        "LIWC-family loadings by benchmark\n(mean $z$ over top-20 bins)",
        fig_dir / "fig-liwc-families-bybenchmark.pdf",
    )
    write_families_table(grp, tab_dir / "tab-liwc-families.tex")


if __name__ == "__main__":
    main()