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

#!/usr/bin/env python3
# Compute SocialIQA vs BBH-Snarks bin-level overlap (Pearson correlation
# of z-scored influence over 576 bins), and check the §H prose claims
# about Social Life / Literature topics and Q&A Forum / Creative Writing
# formats.

import csv
import os
import statistics
from collections import defaultdict
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]
DATA_ROOT = Path(
    os.environ.get(
        "SDA_ZSCORED_AGGREGATED_ROOT",
        str(REPO_ROOT / "artifacts/zscored_bin_scores/aggregated"),
    )
)


def read_zscores(path: Path) -> dict[tuple[str, str], float]:
    out = {}
    with path.open() as fh:
        for row in csv.DictReader(fh):
            out[(row["topic_label"], row["format_label"])] = float(row["zscore"])
    return out


def pearson(xs: list[float], ys: list[float]) -> float:
    mx = statistics.mean(xs)
    my = statistics.mean(ys)
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sx = (sum((x - mx) ** 2 for x in xs)) ** 0.5
    sy = (sum((y - my) ** 2 for y in ys)) ** 0.5
    if sx == 0 or sy == 0:
        return float("nan")
    return num / (sx * sy)


def topic_marg(z: dict[tuple[str, str], float]) -> dict[str, float]:
    s, c = defaultdict(float), defaultdict(int)
    for (t, _), v in z.items():
        s[t] += v
        c[t] += 1
    return {t: s[t] / c[t] for t in s}


def format_marg(z: dict[tuple[str, str], float]) -> dict[str, float]:
    s, c = defaultdict(float), defaultdict(int)
    for (_, f), v in z.items():
        s[f] += v
        c[f] += 1
    return {f: s[f] / c[f] for f in s}


def main() -> None:
    soc = read_zscores(DATA_ROOT / "zscored_socialiqa.csv")
    snarks_base = read_zscores(DATA_ROOT / "zscored_bbh_snarks_base.csv")
    snarks_inst = read_zscores(DATA_ROOT / "zscored_bbh_snarks_instruct.csv")

    bins = sorted(set(soc) & set(snarks_base) & set(snarks_inst))
    print(f"# Bins in all three: {len(bins)}")
    xs = [soc[b] for b in bins]
    ys_b = [snarks_base[b] for b in bins]
    ys_i = [snarks_inst[b] for b in bins]
    print(f"# Pearson r(SocialIQA, Snarks-base):     {pearson(xs, ys_b):+.3f}")
    print(f"# Pearson r(SocialIQA, Snarks-instruct): {pearson(xs, ys_i):+.3f}")
    print(f"# Pearson r(Snarks-base, Snarks-instruct): {pearson(ys_b, ys_i):+.3f}")
    print()

    # Topic marginals: where do Social Life and Literature land in Snarks?
    print("# SocialIQA vs Snarks-base topic-marginal z-scores")
    soc_topics = topic_marg(soc)
    snk_topics = topic_marg(snarks_base)
    for t in sorted(soc_topics):
        print(
            f"  {t:35s}  SocialIQA {soc_topics[t]:+.2f}   Snarks {snk_topics[t]:+.2f}"
        )
    print()
    print("# Format-marginal z-scores")
    soc_fmt = format_marg(soc)
    snk_fmt = format_marg(snarks_base)
    for f in sorted(soc_fmt):
        print(f"  {f:25s}  SocialIQA {soc_fmt[f]:+.2f}   Snarks {snk_fmt[f]:+.2f}")


if __name__ == "__main__":
    main()