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

#!/usr/bin/env python3
from __future__ import annotations

"""RQ4 paper figures: lexical feature heatmaps and profile comparison.

Generates publication-quality Plotly figures following the existing
distribution_report style conventions. Outputs PNG, PDF, and HTML.

Usage:
    python scripts/analysis/rq4_figures.py \
        --features data/outputs/rq4_bin_features.csv \
        --scores-dir artifacts/influence_bin_scores \
        --output-dir artifacts/paper_figures
"""

import argparse
import csv
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "src"))

from dolma.constants import FORMATS, TOPICS
from dolma.distribution_report.data_loader import short_label
from dolma.distribution_report.style import (
    COLOR_EMPTY_CELL,
    COLOR_PRIMARY,
    DIVERGING_COLORSCALE,
    DPI_SCALE,
    FIGURE_HEIGHT_HEATMAP_PX,
    FIGURE_HEIGHT_SIDE_BY_SIDE_PX,
    FIGURE_WIDTH_HEATMAP_PX,
    FIGURE_WIDTH_SIDE_BY_SIDE_PX,
    paper_layout,
    plotly_font,
    save_figure,
)

FORMATS_ORDER = list(FORMATS)
TOPICS_ORDER = list(TOPICS)


# Local alias so existing call sites stay readable; delegates to the
# canonical short_label registry for visual consistency with Fig 4 and
# the rest of the distribution_report pipeline.
def display_label(label: str) -> str:
    return short_label(label)


def pivot_feature(df: pd.DataFrame, col: str) -> np.ndarray:
    matrix = df.pivot(index="bin_topic", columns="bin_format", values=col).fillna(0)
    matrix = matrix.reindex(index=TOPICS_ORDER, columns=FORMATS_ORDER, fill_value=0)
    return matrix.values.astype(float)


def zscore_influence(df: pd.DataFrame) -> pd.DataFrame:
    """Z-score standardize mean_score within a single benchmark (over 576 bins)."""
    out = df.copy()
    mu = out["mean_score"].mean()
    sigma = out["mean_score"].std()
    out["mean_score"] = (out["mean_score"] - mu) / sigma
    return out


def fig_feature_heatmap(
    df: pd.DataFrame,
    value_col: str,
    title: str,
    colorbar_title: str,
    output_dir: Path,
    filename: str,
    formats: tuple[str, ...],
) -> None:
    values = pivot_feature(df, value_col)
    x_labels = [display_label(f) for f in FORMATS_ORDER]
    y_labels = [display_label(t) for t in TOPICS_ORDER]

    fig = go.Figure(
        go.Heatmap(
            z=values,
            x=x_labels,
            y=y_labels,
            colorscale="Viridis",
            zmin=0,
            xgap=1,
            ygap=1,
            colorbar=dict(title=colorbar_title,
                          tickfont=plotly_font("COLORBAR_TICK")),
        )
    )
    fig.update_layout(
        **paper_layout(title, plot_bgcolor=COLOR_EMPTY_CELL),
        xaxis=dict(tickangle=45, tickfont=plotly_font("TICK")),
        yaxis=dict(tickfont=plotly_font("TICK")),
        width=FIGURE_WIDTH_HEATMAP_PX,
        height=FIGURE_HEIGHT_HEATMAP_PX,
    )
    save_figure(fig, output_dir, filename, formats,
                FIGURE_WIDTH_HEATMAP_PX, FIGURE_HEIGHT_HEATMAP_PX)
    print(f"  {filename}")


def fig_influence_vs_feature(
    features_df: pd.DataFrame,
    influence_df: pd.DataFrame,
    feature_col: str,
    feature_title: str,
    feature_colorbar: str,
    output_dir: Path,
    filename: str,
    formats: tuple[str, ...],
) -> None:
    x_labels = [display_label(f) for f in FORMATS_ORDER]
    y_labels = [display_label(t) for t in TOPICS_ORDER]

    infl_z = zscore_influence(influence_df)
    infl = infl_z.rename(columns={
        "topic_label": "bin_topic", "format_label": "bin_format",
    })
    infl_matrix = infl.pivot(
        index="bin_topic", columns="bin_format", values="mean_score"
    ).fillna(0).reindex(
        index=TOPICS_ORDER, columns=FORMATS_ORDER, fill_value=0
    ).values.astype(float)

    feat_matrix = pivot_feature(features_df, feature_col)

    finite = infl_matrix[np.isfinite(infl_matrix)]
    abs_max = float(max(np.percentile(np.abs(finite), 95), 1e-8))

    w, h = 1800, 950

    fig = make_subplots(
        rows=1, cols=2,
        subplot_titles=[
            "SocialIQA: Signed Influence (z-score)",
            feature_title,
        ],
        horizontal_spacing=0.14,
    )

    fig.add_trace(
        go.Heatmap(
            z=infl_matrix, x=x_labels, y=y_labels,
            colorscale=DIVERGING_COLORSCALE,
            zmid=0, zmin=-abs_max, zmax=abs_max,
            xgap=1, ygap=1,
            colorbar=dict(
                title=dict(text="z-score", side="right"),
                x=0.42, len=0.75, thickness=12,
            ),
        ),
        row=1, col=1,
    )
    fig.add_trace(
        go.Heatmap(
            z=feat_matrix, x=x_labels, y=y_labels,
            colorscale="Viridis", zmin=0,
            xgap=1, ygap=1,
            colorbar=dict(
                title=dict(text=feature_colorbar, side="right"),
                x=1.02, len=0.75, thickness=12,
            ),
        ),
        row=1, col=2,
    )

    layout = paper_layout("", plot_bgcolor=COLOR_EMPTY_CELL)
    layout["margin"] = dict(l=140, r=80, t=60, b=140)
    fig.update_layout(**layout, width=w, height=h)
    fig.update_annotations(font=plotly_font("SUBPLOT_TITLE"))
    fig.update_xaxes(tickangle=45, tickfont=plotly_font("TICK"))
    fig.update_yaxes(tickfont=plotly_font("TICK"))

    save_figure(fig, output_dir, filename, formats, w, h)
    print(f"  {filename}")


def fig_profile_comparison(
    features_df: pd.DataFrame,
    output_dir: Path,
    filename: str,
    formats: tuple[str, ...],
) -> None:
    interp_bins = [
        ("literature", "customer_support"),
        ("social_life", "q_a_forum"),
        ("social_life", "about_pers"),
    ]
    tech_bins = [
        ("industrial", "documentation"),
        ("health", "documentation"),
        ("industrial", "structured_data"),
        ("industrial", "legal_notices"),
        ("fashion_and_beauty", "documentation"),
        ("home_and_hobbies", "documentation"),
        ("science_math_and_technology", "documentation"),
    ]

    def weighted_avg(subset: pd.DataFrame, col: str) -> float:
        tw = subset["total_words"].sum()
        if tw == 0:
            return 0.0
        return float((subset[col] * subset["total_words"]).sum() / tw)

    interp = features_df[features_df.apply(
        lambda r: (r["bin_topic"], r["bin_format"]) in interp_bins, axis=1
    )]
    tech = features_df[features_df.apply(
        lambda r: (r["bin_topic"], r["bin_format"]) in tech_bins, axis=1
    )]

    metrics = ["first_person_per_1k", "second_person_per_1k",
               "third_person_per_1k", "mental_state_per_1k"]
    metric_labels = ["1st person", "2nd person", "3rd person", "Mental state"]

    corpus_vals = [weighted_avg(features_df, m) for m in metrics]
    interp_vals = [weighted_avg(interp, m) for m in metrics]
    tech_vals = [weighted_avg(tech, m) for m in metrics]

    colors = ["#2c7bb6", "#d7191c", "#999999"]

    fig = go.Figure()
    fig.add_trace(go.Bar(
        name="Interpersonal (3 bins)",
        x=metric_labels, y=interp_vals,
        marker_color=colors[0],
    ))
    fig.add_trace(go.Bar(
        name="Technical (7 bins)",
        x=metric_labels, y=tech_vals,
        marker_color=colors[1],
    ))
    fig.add_trace(go.Bar(
        name="Corpus average",
        x=metric_labels, y=corpus_vals,
        marker_color=colors[2],
    ))

    w, h = 550, 380
    layout = paper_layout("")
    layout["margin"] = dict(l=50, r=10, t=10, b=40)
    fig.update_layout(
        **layout,
        barmode="group",
        yaxis=dict(title=dict(text="Count per 1,000 words",
                              font=plotly_font("AXIS_TITLE")),
                   tickfont=plotly_font("TICK")),
        xaxis=dict(tickfont=plotly_font("TICK")),
        legend=dict(
            x=0.35, y=0.98,
            font=plotly_font("LEGEND"),
            bgcolor="rgba(255,255,255,0.8)",
            borderwidth=0,
        ),
        width=w, height=h,
        bargap=0.15, bargroupgap=0.05,
    )
    save_figure(fig, output_dir, filename, formats, w, h)
    print(f"  {filename}")


def write_contrastive_table(
    features_df: pd.DataFrame,
    socialiqa_df: pd.DataFrame,
    gsm8k_df: pd.DataFrame,
    output_dir: Path,
    top_n: int = 15,
) -> None:
    siq_z = zscore_influence(socialiqa_df)
    gsm_z = zscore_influence(gsm8k_df)

    siq = siq_z.rename(columns={
        "topic_label": "bin_topic", "format_label": "bin_format",
        "mean_score": "siq_score",
    })[["bin_topic", "bin_format", "siq_score"]]
    gsm = gsm_z.rename(columns={
        "topic_label": "bin_topic", "format_label": "bin_format",
        "mean_score": "gsm_score",
    })[["bin_topic", "bin_format", "gsm_score"]]

    merged = features_df.merge(siq, on=["bin_topic", "bin_format"], how="left")
    merged = merged.merge(gsm, on=["bin_topic", "bin_format"], how="left")
    merged["diff"] = merged["siq_score"] - merged["gsm_score"]

    top = merged.nlargest(top_n, "diff")

    out_path = output_dir / "table_rq4_contrastive_characterized.tex"
    with open(out_path, "w") as f:
        f.write("\\begin{tabular}{rlrrrr}\n")
        f.write("\\toprule\n")
        f.write(" & Bin & \\shortstack{Mean\\\\words} "
                "& \\shortstack{1st per.\\\\per 1k} "
                "& \\shortstack{Mental st.\\\\per 1k} "
                "& $\\Delta$ \\\\\n")
        f.write("\\midrule\n")
        for i, (_, row) in enumerate(top.iterrows(), 1):
            label = f"{display_label(row['bin_topic'])} / {display_label(row['bin_format'])}"
            label_tex = label.replace("&", "\\&")
            f.write(
                f"{i} & {label_tex} "
                f"& {row['mean_doc_words']:.0f} "
                f"& {row['first_person_per_1k']:.1f} "
                f"& {row['mental_state_per_1k']:.2f} "
                f"& {row['diff']:.2f} \\\\\n"
            )
        f.write("\\bottomrule\n")
        f.write("\\end{tabular}\n")
    print(f"  {out_path.name}")


def main():
    parser = argparse.ArgumentParser(description="RQ4 paper figures")
    parser.add_argument("--features", required=True)
    parser.add_argument("--scores-dir", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--format", default="png,pdf,html",
                        help="Comma-separated output formats")
    args = parser.parse_args()

    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    fmts = tuple(args.format.split(","))

    features = pd.read_csv(args.features)
    scores_dir = Path(args.scores_dir)
    socialiqa = pd.read_csv(scores_dir / "queries_socialiqa_bin_scores.csv")
    gsm8k = pd.read_csv(scores_dir / "queries_gsm8k_bin_scores.csv")

    print("Generating RQ4 figures...")

    fig_feature_heatmap(
        features, "mental_state_per_1k",
        "Mental-State Verb Density (per 1,000 words)",
        "per 1k words",
        out, "fig_rq4_mental_state_heatmap", fmts,
    )

    fig_feature_heatmap(
        features, "first_person_per_1k",
        "First-Person Pronoun Density (per 1,000 words)",
        "per 1k words",
        out, "fig_rq4_first_person_heatmap", fmts,
    )

    fig_influence_vs_feature(
        features, socialiqa,
        "mental_state_per_1k",
        "Mental-State Verb Density",
        "per 1k words",
        out, "fig_rq4_influence_vs_mental_state", fmts,
    )

    fig_profile_comparison(
        features, out, "fig_rq4_profile_comparison", fmts,
    )

    write_contrastive_table(features, socialiqa, gsm8k, out)

    print("Done.")


if __name__ == "__main__":
    main()