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

"""Read metrics and compute gamma grids for N10 Figure-4-style tables."""

from __future__ import annotations

import json
from collections import defaultdict
from pathlib import Path
from typing import Any

from figure4_table_helpers import (
    ALL_TOPICS,
    GAMMA_EPS,
    METRICS_OF_INTEREST,
    NULL_LABEL,
    PROBES,
)


def _find_metrics(eval_dir: Path, task_name: str) -> Path | None:
    if not eval_dir.exists():
        return None
    candidates = list(eval_dir.glob("task-*-metrics.json"))
    if not candidates:
        return None
    if len(candidates) == 1:
        return candidates[0]
    for candidate in candidates:
        normalized_task = task_name.replace("::", "_").replace(":", "_")
        if normalized_task in candidate.name.replace(":", "_"):
            return candidate
    return candidates[0]


def _load_metrics(path: Path | None) -> dict[str, float]:
    if path is None or not path.exists():
        return {}
    try:
        data = json.loads(path.read_text())
    except Exception:
        return {}
    metrics = data.get("metrics", {})
    return {key: metrics.get(key) for key in METRICS_OF_INTEREST if key in metrics}


def collect_baselines(baseline_root: Path) -> dict[str, dict[str, float]]:
    out: dict[str, dict[str, float]] = {}
    for probe, _ in PROBES:
        matches = (
            list(baseline_root.glob(f"{probe}:mc_tom-rebuttal_*"))
            + list(baseline_root.glob(f"{probe}:cot::olmes*"))
            + list(baseline_root.glob(f"{probe}:cot_olmes_*"))
        )
        if not matches:
            out[probe] = {}
            continue
        latest = sorted(matches)[-1]
        out[probe] = _load_metrics(_find_metrics(latest, probe))
    return out


def collect_post(eval_grid: Path) -> dict[str, dict[str, dict[str, float]]]:
    out: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)
    for topic in ALL_TOPICS:
        topic_dir = eval_grid / topic
        if not topic_dir.exists():
            continue
        for probe, task_name in PROBES:
            eval_dir = topic_dir / f"eval_{probe}"
            out[probe][topic] = _load_metrics(_find_metrics(eval_dir, task_name))
    return out


def collect_null(eval_grid: Path) -> dict[str, dict[str, float]]:
    out: dict[str, dict[str, float]] = {}
    null_dir = eval_grid / NULL_LABEL
    for probe, task_name in PROBES:
        eval_dir = null_dir / f"eval_{probe}"
        out[probe] = _load_metrics(_find_metrics(eval_dir, task_name))
    return out


def _gamma(post_val: float | None, pre: float | None, mode: str) -> float | None:
    if pre is None or post_val is None:
        return None
    if mode == "relative":
        if abs(pre) < GAMMA_EPS:
            return None
        return (post_val - pre) / pre
    return post_val - pre


def compute_gamma(
    baselines: dict[str, dict[str, float]],
    post: dict[str, dict[str, dict[str, float]]],
    metric: str = "acc_per_char",
    mode: str = "absolute",
    null_post: dict[str, dict[str, float]] | None = None,
) -> dict[str, dict[str, dict[str, Any]]]:
    grid: dict[str, dict[str, dict[str, Any]]] = {}
    for probe, _ in PROBES:
        pre = baselines.get(probe, {}).get(metric)
        post_null_val = (null_post or {}).get(probe, {}).get(metric)
        gamma_null = _gamma(post_null_val, pre, mode)
        grid[probe] = {}
        for topic in ALL_TOPICS:
            post_val = post.get(probe, {}).get(topic, {}).get(metric)
            gamma = _gamma(post_val, pre, mode)
            net = (
                gamma - gamma_null
                if gamma is not None and gamma_null is not None
                else None
            )
            grid[probe][topic] = {
                "pre": pre,
                "post": post_val,
                "post_null": post_null_val,
                "gamma": gamma,
                "gamma_null": gamma_null,
                "net_gamma": net,
            }
    return grid