scripts/summarize_conditional_probe.py

#!/usr/bin/env python3
"""Summarize Linear/Ridge and nonlinear conditional probe fold results."""

from __future__ import annotations

import argparse
import csv
import itertools
import json
import math
from collections import defaultdict
from pathlib import Path
from typing import Any

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


FAMILIES = ("self_forcing", "causal_forcing", "hy_worldplay")
ROLES = ("early", "middle", "late", "final")


def read_csv(path: Path) -> list[dict[str, Any]]:
    with path.open(newline="", encoding="utf-8") as handle:
        rows = list(csv.DictReader(handle))
    for row in rows:
        for key in ("target_step", "held_out_prompt", "seed", "layer_index", "test_tokens"):
            if key in row:
                row[key] = int(row[key])
        for key in ("mse", "nMSE", "nRMSE", "r2", "cosine"):
            row[key] = float(row[key])
    return rows


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def bootstrap(values: list[float], seed: int, rounds: int = 10000):
    values = np.asarray(values, dtype=np.float64)
    rng = np.random.default_rng(seed)
    if values.size == 0:
        return float("nan"), float("nan"), float("nan")
    indices = rng.integers(0, values.size, size=(rounds, values.size))
    means = values[indices].mean(axis=1)
    return float(values.mean()), float(np.quantile(means, 0.025)), float(np.quantile(means, 0.975))


def exact_signflip(values: list[float]) -> float:
    values = np.asarray(values, dtype=np.float64)
    values = values[np.isfinite(values)]
    if not values.size:
        return float("nan")
    observed = abs(float(values.mean()))
    exceed = 0
    total = 1 << int(values.size)
    for mask in range(total):
        signed = np.asarray(
            [value if (mask >> index) & 1 else -value for index, value in enumerate(values)]
        )
        if abs(float(signed.mean())) >= observed - 1e-15:
            exceed += 1
    return float((exceed + 1) / (total + 1))


def prompt_metric_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Average seeds within each prompt, keeping prompt as statistical unit."""
    grouped: dict[tuple, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        grouped[
            (
                row["method"], row["model_family"], row["layer_role"],
                row["target_step"], row["probe"], row["held_out_prompt"],
            )
        ].append(row)
    result = []
    for key, values in sorted(grouped.items(), key=lambda item: tuple(map(str, item[0]))):
        method, family, role, step, probe, prompt = key
        item = {
            "method": method,
            "model_family": family,
            "layer_role": role,
            "target_step": step,
            "probe": probe,
            "held_out_prompt": prompt,
            "seed_count": len(values),
        }
        for metric in ("mse", "nMSE", "nRMSE", "r2", "cosine"):
            item[metric] = float(np.mean([float(value[metric]) for value in values]))
        result.append(item)
    return result


def gain_rows(prompt_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    baseline_name = {"linear": "within_affine", "nonlinear": "step_only"}
    grouped = defaultdict(dict)
    for row in prompt_rows:
        grouped[
            (row["method"], row["model_family"], row["layer_role"], row["target_step"], row["held_out_prompt"])
        ][row["probe"]] = row
    result = []
    for key, probes in sorted(grouped.items(), key=lambda item: tuple(map(str, item[0]))):
        method, family, role, step, prompt = key
        baseline = probes.get(baseline_name[method])
        if baseline is None:
            continue
        for probe, current in probes.items():
            reference_mse = float(baseline["mse"])
            current_mse = float(current["mse"])
            result.append({
                "method": method,
                "model_family": family,
                "layer_role": role,
                "target_step": step,
                "held_out_prompt": prompt,
                "baseline_probe": baseline_name[method],
                "probe": probe,
                "baseline_mse": reference_mse,
                "probe_mse": current_mse,
                "gain": (reference_mse - current_mse) / max(reference_mse, 1e-12),
                "delta_r2": float(current["r2"]) - float(baseline["r2"]),
            })
    return result


def summarize_gains(gains: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped = defaultdict(list)
    for row in gains:
        grouped[(row["method"], row["model_family"], row["layer_role"], row["target_step"], row["probe"])].append(row)
    result = []
    for key, values in sorted(grouped.items(), key=lambda item: tuple(map(str, item[0]))):
        method, family, role, step, probe = key
        values = sorted(values, key=lambda row: row["held_out_prompt"])
        numbers = [float(row["gain"]) for row in values]
        mean, low, high = bootstrap(numbers, 1000 + sum(ord(ch) for ch in str(key)))
        result.append({
            "method": method,
            "model_family": family,
            "layer_role": role,
            "target_step": step,
            "probe": probe,
            "baseline_probe": values[0]["baseline_probe"],
            "prompt_count": len(numbers),
            "gain_mean": mean,
            "gain_ci95_low": low,
            "gain_ci95_high": high,
            "wins": int(sum(number > 0 for number in numbers)),
            "signflip_p": exact_signflip(numbers),
        })
    return result


def summarize_layer_gains(gains: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Average target steps inside each prompt, then summarize across prompts."""
    prompt_groups = defaultdict(list)
    for row in gains:
        prompt_groups[
            (
                row["method"], row["model_family"], row["layer_role"],
                row["probe"], row["held_out_prompt"], row["baseline_probe"],
            )
        ].append(float(row["gain"]))
    groups = defaultdict(list)
    for key, values in prompt_groups.items():
        method, family, role, probe, _prompt, baseline = key
        groups[(method, family, role, probe, baseline)].append(float(np.mean(values)))
    result = []
    for key, values in sorted(groups.items(), key=lambda item: tuple(map(str, item[0]))):
        method, family, role, probe, baseline = key
        mean, low, high = bootstrap(values, 17000 + sum(ord(ch) for ch in str(key)))
        result.append({
            "method": method,
            "model_family": family,
            "layer_role": role,
            "probe": probe,
            "baseline_probe": baseline,
            "target_step_count": 3,
            "prompt_count": len(values),
            "gain_mean": mean,
            "gain_ci95_low": low,
            "gain_ci95_high": high,
            "wins": int(sum(value > 0 for value in values)),
            "signflip_p": exact_signflip(values),
        })
    return result


def summarize_metrics(prompt_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped = defaultdict(list)
    for row in prompt_rows:
        grouped[(row["method"], row["model_family"], row["layer_role"], row["target_step"], row["probe"])].append(row)
    result = []
    for key, values in sorted(grouped.items(), key=lambda item: tuple(map(str, item[0]))):
        method, family, role, step, probe = key
        item = {
            "method": method,
            "model_family": family,
            "layer_role": role,
            "target_step": step,
            "probe": probe,
            "prompt_count": len(values),
        }
        for metric_index, metric in enumerate(("mse", "nMSE", "nRMSE", "r2", "cosine")):
            numbers = [float(row[metric]) for row in values]
            mean, low, high = bootstrap(numbers, 9000 + metric_index + sum(ord(ch) for ch in str(key)))
            item[f"{metric}_mean"] = mean
            item[f"{metric}_ci95_low"] = low
            item[f"{metric}_ci95_high"] = high
        result.append(item)
    return result


def plot_primary(gains: list[dict[str, Any]], output: Path) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(14, 5), sharey=True)
    for ax, method, title in zip(axes, ("linear", "nonlinear"), ("Linear/Ridge", "Nonlinear MLP")):
        selected = [
            row for row in gains
            if row["method"] == method and row["probe"] in ({"fusion_same"} if method == "linear" else {"both_correct"})
        ]
        x = np.arange(len(ROLES))
        width = 0.24
        for family_index, family in enumerate(FAMILIES):
            values = []
            lows = []
            highs = []
            for role in ROLES:
                group = [row for row in selected if row["model_family"] == family and row["layer_role"] == role]
                nums = [float(row["gain"]) for row in group]
                mean, low, high = bootstrap(nums, 1234 + family_index * 100 + ROLES.index(role))
                values.append(mean)
                lows.append(mean - low)
                highs.append(high - mean)
            pos = x + (family_index - (len(FAMILIES) - 1) / 2) * width
            ax.bar(pos, values, width, yerr=[lows, highs], capsize=3, label=family)
        ax.axhline(0, color="black", linewidth=0.8)
        ax.set_xticks(x, ROLES)
        ax.set_ylabel("Gain vs step-only baseline" if method == "nonlinear" else "Gain vs within-affine baseline")
        ax.set_title(title)
        ax.grid(axis="y", alpha=0.25)
    axes[1].legend(fontsize=9)
    fig.tight_layout()
    fig.savefig(output, dpi=180)
    plt.close(fig)


def plot_controls(gains: list[dict[str, Any]], output: Path) -> None:
    probes = ["fusion_same", "fusion_step_duplicate", "fusion_wrong_step", "fusion_distant", "fusion_batch_shuffle", "fusion_zero", "fusion_noise"]
    labels = {
        "fusion_same": "correct",
        "fusion_step_duplicate": "step duplicate",
        "fusion_wrong_step": "wrong step",
        "fusion_distant": "distant",
        "fusion_batch_shuffle": "other video",
        "fusion_zero": "zero",
        "fusion_noise": "noise",
    }
    fig, axes = plt.subplots(
        1,
        len(FAMILIES),
        figsize=(5.7 * len(FAMILIES), 5),
        sharey=True,
    )
    axes = np.atleast_1d(axes)
    for ax, family in zip(axes, FAMILIES):
        values = []
        errors = []
        for probe in probes:
            group = [row for row in gains if row["method"] == "linear" and row["model_family"] == family and row["layer_role"] == "final" and row["probe"] == probe]
            nums = [float(row["gain"]) for row in group]
            mean, low, high = bootstrap(nums, 4000 + probes.index(probe))
            values.append(mean)
            errors.append((mean - low, high - mean))
        y = np.arange(len(probes))
        ax.errorbar(values, y, xerr=np.asarray(errors).T, fmt="o", capsize=3)
        ax.axvline(0, color="black", linewidth=0.8)
        ax.set_yticks(y, [labels[p] for p in probes])
        ax.set_title(family)
        ax.grid(axis="x", alpha=0.25)
    axes[0].set_xlabel("Linear MSE gain")
    fig.tight_layout()
    fig.savefig(output, dpi=180)
    plt.close(fig)


def build_report(metrics: list[dict[str, Any]], gains: list[dict[str, Any]], output: Path, config: dict[str, Any]) -> None:
    lines = [
        "# Conditional prediction and incremental chunk information",
        "",
        "This report uses 10 prompt-grouped held-out folds. The test prompt and its",
        "other-video donor are excluded from training; nonlinear seeds are averaged",
        "within prompt before confidence intervals are computed.",
        "",
        "The primary endpoint is `fusion_same` vs `within_affine` for Linear/Ridge",
        "and `both_correct` vs `step_only` for the nonlinear MLP.",
        "",
        "| method | family | role | step | probe | gain | 95% CI | wins | sign-flip p |",
        "|---|---|---|---:|---|---:|---|---:|---:|",
    ]
    primary = [
        row for row in gains
        if (row["method"] == "linear" and row["probe"] == "fusion_same")
        or (row["method"] == "nonlinear" and row["probe"] == "both_correct")
    ]
    for row in primary:
        lines.append(
            f"| {row['method']} | {row['model_family']} | {row['layer_role']} | {row['target_step']} | "
            f"{row['probe']} | {row['gain_mean']:.4f} | [{row['gain_ci95_low']:.4f}, {row['gain_ci95_high']:.4f}] | "
            f"{row['wins']}/{row['prompt_count']} | {row['signflip_p']:.4f} |"
        )
    lines += [
        "",
        "Interpretation: positive gain means that adding the auxiliary feature lowers held-out MSE.",
        "The exact sign-flip test treats prompt, not token, as the independent unit.",
        "Absolute MSE is not compared across model families because feature dimensions and",
        "conditioning paths differ.",
        "",
        "## Files",
        "",
        "- `probe_folds_unified.csv`",
        "- `probe_prompt_averaged.csv`",
        "- `probe_metrics_summary.csv`",
        "- `probe_gain_summary.csv`",
        "- `conditional_gain_by_layer.png`",
        "- `conditional_control_comparison.png`",
        "",
        "```json",
        json.dumps(config, indent=2, ensure_ascii=False),
        "```",
    ]
    output.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    global FAMILIES
    parser = argparse.ArgumentParser()
    parser.add_argument("--linear_csv", type=Path, required=True)
    parser.add_argument("--nonlinear_dir", type=Path, required=True)
    parser.add_argument("--output_dir", type=Path, required=True)
    parser.add_argument("--families", default=",".join(FAMILIES))
    parser.add_argument(
        "--chunk_pairing",
        choices=("matched_slot", "boundary_to_all"),
        default="matched_slot",
    )
    args = parser.parse_args()
    requested_families = tuple(
        value.strip() for value in args.families.split(",") if value.strip()
    )
    unknown = set(requested_families) - set(FAMILIES)
    if not requested_families or unknown:
        raise ValueError(f"Invalid families: {requested_families}; unknown={sorted(unknown)}")
    FAMILIES = requested_families
    args.output_dir.mkdir(parents=True, exist_ok=True)

    linear = read_csv(args.linear_csv)
    linear = [row for row in linear if row["model_family"] in FAMILIES]
    linear = [{**row, "method": "linear"} for row in linear]
    nonlinear = []
    for family in FAMILIES:
        path = args.nonlinear_dir / f"nonlinear_probe_{family}_folds.csv"
        rows = read_csv(path)
        nonlinear.extend({**row, "method": "nonlinear"} for row in rows)
    expected_linear = 1320 * len(FAMILIES)
    expected_nonlinear = 2160 * len(FAMILIES)
    if len(linear) != expected_linear:
        raise ValueError(f"Expected {expected_linear} linear rows, found {len(linear)}")
    if len(nonlinear) != expected_nonlinear:
        raise ValueError(f"Expected {expected_nonlinear} nonlinear rows, found {len(nonlinear)}")
    unified = linear + nonlinear
    prompt_rows = prompt_metric_rows(unified)
    gains = gain_rows(prompt_rows)
    metrics = summarize_metrics(prompt_rows)
    gain_summary = summarize_gains(gains)
    layer_gain_summary = summarize_layer_gains(gains)
    write_csv(args.output_dir / "probe_folds_unified.csv", unified)
    write_csv(args.output_dir / "probe_prompt_averaged.csv", prompt_rows)
    write_csv(args.output_dir / "probe_metrics_summary.csv", metrics)
    write_csv(args.output_dir / "probe_gain_by_prompt.csv", gains)
    write_csv(args.output_dir / "probe_gain_summary.csv", gain_summary)
    write_csv(args.output_dir / "probe_layer_gain_summary.csv", layer_gain_summary)
    plot_primary(gains, args.output_dir / "conditional_gain_by_layer.png")
    plot_controls(gains, args.output_dir / "conditional_control_comparison.png")
    config = {
        "linear_rows": len(linear),
        "nonlinear_rows": len(nonlinear),
        "families": list(FAMILIES),
        "chunk_pairing": args.chunk_pairing,
        "prompt_averaged_rows": len(prompt_rows),
        "gain_rows": len(gains),
        "prompt_count": 10,
        "target_chunks": [2, 3],
        "target_steps": [1, 2, 3],
        "linear_baseline": "within_affine",
        "nonlinear_baseline": "step_only",
        "nonlinear_seed_aggregation": "mean within held-out prompt",
        "outer_split": "held-out prompt plus its cyclic other-video donor",
    }
    (args.output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
    build_report(metrics, gain_summary, args.output_dir / "REPORT.md", config)
    print(f"[complete] {args.output_dir} unified={len(unified)} gains={len(gains)}", flush=True)


if __name__ == "__main__":
    main()