social-data-attribution-w2/scripts/visualization/influence_correlation.py

#!/usr/bin/env python3
"""Influence-unlearning correlation scaffold.

Accepts TrackStar influence scores as input and correlates them with
unlearning gamma values from Experiment 1 (single-bin).

Usage (once influence data is available):
    python -m scripts.visualization.influence_correlation \
        --influence-scores path/to/influence_scores.csv
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
import pandas as pd

from scripts.visualization._shared import (
    BENCH_COLORS, MAIN_METRICS, METRIC_LABELS,
    paper_rc, save_fig,
)
from scripts.visualization.experiment1_single_bin import build_dataframe as build_exp1_df

EXPECTED_INFLUENCE_COLS = ["topic", "influence_score"]


def load_influence_scores(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    missing = [c for c in EXPECTED_INFLUENCE_COLS if c not in df.columns]
    if missing:
        raise ValueError(f"Influence CSV missing columns: {missing}")
    return df


def fig_rank_heatmap(merged: pd.DataFrame, out: Path) -> None:
    metrics = MAIN_METRICS
    topics = merged.sort_values("influence_rank")["topic"].values

    rank_cols = ["influence_rank"] + [f"{m}_gamma_rank" for m in metrics]
    col_labels = ["Influence\nRank"] + [METRIC_LABELS[m] for m in metrics]

    ordered = merged.set_index("topic").loc[topics]
    data = ordered[rank_cols].values.astype(float)

    fig, ax = plt.subplots(figsize=(6, 8))
    import seaborn as sns
    sns.heatmap(
        data, ax=ax,
        xticklabels=col_labels,
        yticklabels=topics,
        cmap="YlOrRd_r", annot=True, fmt=".0f",
        annot_kws={"fontsize": 7},
        linewidths=0.4, linecolor="white",
        cbar_kws={"label": "Rank (1 = highest)", "shrink": 0.4},
    )
    ax.set_yticklabels(ax.get_yticklabels(), fontsize=8, rotation=0)
    ax.tick_params(axis="both", length=0)
    fig.subplots_adjust(left=0.22)
    save_fig(fig, out, "fig_influence_rank_heatmap")


def fig_correlation_scatter(merged: pd.DataFrame, out: Path) -> None:
    metrics = MAIN_METRICS
    fig, axes = plt.subplots(1, len(metrics), figsize=(4 * len(metrics), 4))

    for idx, m in enumerate(metrics):
        ax = axes[idx]
        x = merged["influence_score"].values
        y = merged[f"{m}_gamma"].values

        ax.scatter(x, y, color=BENCH_COLORS[m], s=30, edgecolors="black", linewidth=0.3, zorder=3)

        if len(x) > 2:
            z = np.polyfit(x, y, 1)
            p = np.poly1d(z)
            x_line = np.linspace(x.min(), x.max(), 50)
            ax.plot(x_line, p(x_line), "--", color="#333", linewidth=0.8, alpha=0.6)

            corr = np.corrcoef(x, y)[0, 1]
            rank_corr = _spearman(x, y)
            ax.text(
                0.05, 0.95, f"r={corr:.3f}\nρ={rank_corr:.3f}",
                transform=ax.transAxes, fontsize=7, va="top",
            )

        ax.axhline(0, color="gray", linestyle="--", linewidth=0.4, alpha=0.5)
        ax.set_xlabel("Influence score")
        ax.set_ylabel(r"$\gamma$")
        ax.set_title(METRIC_LABELS[m], fontweight="bold")
        ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
        ax.grid(True, alpha=0.2, linewidth=0.3)

    fig.tight_layout()
    save_fig(fig, out, "fig_influence_correlation_scatter")


def _spearman(x: np.ndarray, y: np.ndarray) -> float:
    from scipy.stats import spearmanr
    return spearmanr(x, y).correlation


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--influence-scores", type=Path, default=None)
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/paper_main"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    paper_rc()

    if args.influence_scores is None:
        print("No --influence-scores provided. Scaffold only, no figures generated.")
        print("Pass a CSV with columns: topic, influence_score")
        sys.exit(0)

    influence = load_influence_scores(args.influence_scores)
    df_unlearn = build_exp1_df()

    merged = df_unlearn.merge(influence, on="topic", how="inner")
    if len(merged) == 0:
        print("No matching topics between influence scores and unlearning data.")
        sys.exit(1)

    merged["influence_rank"] = merged["influence_score"].rank(ascending=False)
    for m in MAIN_METRICS:
        col = f"{m}_gamma"
        merged[f"{col}_rank"] = merged[col].rank(ascending=True)

    print(f"Matched {len(merged)} topics. Generating correlation figures...")
    fig_rank_heatmap(merged, args.output_dir)
    fig_correlation_scatter(merged, args.output_dir)
    print(f"Correlation figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()