scripts/plot_self_causal_evidence.py

#!/usr/bin/env python3
"""Create compact publication-style figures for the Self/Causal evidence chain."""

from __future__ import annotations

import argparse
import colorsys
import csv
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D
from matplotlib.patches import Patch


SELF = "#0F766E"
CAUSAL = "#4F46E5"
INK = "#172033"
MUTED = "#64748B"
GRID = "#DCE3EC"
PALE = "#F1F5F9"
WARM = "#D97706"
MODEL_COLOR = {"self_forcing": SELF, "causal_forcing": CAUSAL}
MODEL_LABEL = {"self_forcing": "Self-Forcing", "causal_forcing": "Causal-Forcing"}
LAYERS = ("early", "middle", "late", "final")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--redundancy_csv", type=Path, required=True)
    parser.add_argument("--motion_bins_csv", type=Path, required=True)
    parser.add_argument("--motion_summary_csv", type=Path, required=True)
    parser.add_argument("--native_gain_csv", type=Path, required=True)
    parser.add_argument(
        "--causal_late_sensitivity_csv",
        type=Path,
        help="Optional sensitivity summary used to replace only Causal late Ridge.",
    )
    parser.add_argument("--sensitivity_scenario", default="exclude_two_folds")
    parser.add_argument("--aligned_probe_csv", type=Path, required=True)
    parser.add_argument(
        "--unified_aligned_csv",
        type=Path,
        help="Optional four-layer 64-D/common-mask Ridge summary used for panels D/E.",
    )
    parser.add_argument("--output_dir", type=Path, required=True)
    return parser.parse_args()


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def apply_causal_late_sensitivity(
    native_rows: list[dict[str, str]],
    sensitivity_path: Path | None,
    scenario: str,
) -> bool:
    """Replace only the Causal late Ridge aggregate with an explicit sensitivity result."""
    if sensitivity_path is None:
        return False
    sensitivity_rows = read_csv(sensitivity_path)
    selected = [row for row in sensitivity_rows if row["scenario"] == scenario]
    if len(selected) != 1:
        raise ValueError(f"Expected one sensitivity row for {scenario!r}, found {len(selected)}")
    source = selected[0]
    targets = [
        row
        for row in native_rows
        if row["method"] == "linear"
        and row["model_family"] == "causal_forcing"
        and row["layer_role"] == "late"
        and row["probe"] == "fusion_same"
    ]
    if len(targets) != 1:
        raise ValueError(f"Expected one Causal late Ridge row, found {len(targets)}")
    target = targets[0]
    for key in ("gain_mean", "gain_ci95_low", "gain_ci95_high", "wins", "signflip_p"):
        target[key] = source[key]
    target["prompt_count"] = source["prompt_count"]
    return True


def lighten(color: str, amount: float = 0.45) -> str:
    r, g, b = matplotlib.colors.to_rgb(color)
    h, l, s = colorsys.rgb_to_hls(r, g, b)
    return matplotlib.colors.to_hex(colorsys.hls_to_rgb(h, 1 - amount * (1 - l), s * 0.85))


def setup_style() -> None:
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.titlesize": 10.5,
        "axes.labelsize": 9,
        "axes.titleweight": "semibold",
        "axes.labelcolor": INK,
        "axes.edgecolor": GRID,
        "axes.linewidth": 0.8,
        "xtick.color": MUTED,
        "ytick.color": MUTED,
        "xtick.labelsize": 8.5,
        "ytick.labelsize": 8.5,
        "grid.color": GRID,
        "grid.linewidth": 0.7,
        "grid.alpha": 0.65,
        "legend.fontsize": 7.7,
        "legend.frameon": False,
        "figure.facecolor": "white",
        "axes.facecolor": "white",
        "savefig.facecolor": "white",
        "savefig.bbox": "tight",
    })


def polish(ax: plt.Axes, grid_axis: str = "y") -> None:
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis=grid_axis, zorder=0)
    ax.tick_params(length=0)


def panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(
        -0.13,
        1.08,
        label,
        transform=ax.transAxes,
        fontsize=12,
        fontweight="bold",
        color=INK,
        va="top",
    )


def lookup_rows(rows: list[dict[str, str]], keys: tuple[str, ...]) -> dict[tuple[str, ...], dict[str, str]]:
    return {tuple(row[key] for key in keys): row for row in rows}


def plot_redundancy(
    ax: plt.Axes,
    rows: list[dict[str, str]],
    family: str,
    variant_50: str,
    panel: str,
) -> None:
    table = lookup_rows(rows, ("model_family", "model_variant", "layer_role", "comparison"))
    x = np.arange(len(LAYERS))
    color = MODEL_COLOR[family]
    series = [
        (variant_50, "within_adjacent", "50-step · denoising", "#334155", "-", "o"),
        (variant_50, "cross_boundary_to_all", "50-step · boundary", "#94A3B8", "--", "D"),
        ("dmd4", "within_adjacent", "4-step · denoising", color, "-", "o"),
        ("dmd4", "cross_boundary_to_all", "4-step · boundary", lighten(color), "--", "D"),
    ]
    for variant, comparison, label, line_color, linestyle, marker in series:
        values, lows, highs = [], [], []
        for role in LAYERS:
            row = table[(family, variant, role, comparison)]
            value = float(row["token_cosine_mean"])
            values.append(value)
            lows.append(value - float(row["ci95_low"]))
            highs.append(float(row["ci95_high"]) - value)
        ax.errorbar(
            x,
            values,
            yerr=np.asarray([lows, highs]),
            color=line_color,
            linestyle=linestyle,
            linewidth=1.65,
            marker=marker,
            markersize=4.3,
            markeredgewidth=0.7,
            capsize=2,
            label=label,
            zorder=3,
        )
    ax.axvspan(-0.35, 2.35, color=PALE, alpha=0.58, zorder=-2)
    ax.set_xticks(x, [value.title() for value in LAYERS])
    ax.set_ylim(0.62, 1.015)
    ax.set_yticks([0.65, 0.75, 0.85, 0.95, 1.00])
    ax.set_ylabel("Token cosine")
    ax.set_title(MODEL_LABEL[family], loc="left", pad=8)
    polish(ax)
    panel_label(ax, panel)


def plot_motion(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "action", "motion_bin"))
    bins = ("low", "medium", "high")
    x = np.arange(3)
    for family in ("self_forcing", "causal_forcing"):
        color = MODEL_COLOR[family]
        raw = [float(table[(family, "none", motion)]["raw_cosine"]) for motion in bins]
        flow = [float(table[(family, "none", motion)]["flow_aligned_cosine"]) for motion in bins]
        ax.plot(x, raw, color=color, linewidth=1.8, marker="o", markersize=4.5, zorder=3)
        ax.plot(x, flow, color=color, linewidth=1.45, linestyle="--", marker="^", markersize=4.5, zorder=3)
        ax.fill_between(x, raw, flow, color=color, alpha=0.09, zorder=1)
        ax.annotate(
            f"+{100 * (flow[-1] - raw[-1]):.2f} pts",
            (x[-1], flow[-1]),
            xytext=(5, 3),
            textcoords="offset points",
            color=color,
            fontsize=7.3,
            fontweight="semibold",
        )
    ax.set_xticks(x, ["Low", "Medium", "High"])
    ax.set_ylim(0.875, 0.942)
    ax.set_ylabel("Boundary cosine")
    ax.set_title("Motion exposes spatial mismatch", loc="left", pad=8)
    handles = [
        Line2D([], [], color=SELF, marker="o", label="Self · raw"),
        Line2D([], [], color=SELF, marker="^", linestyle="--", label="Self · flow"),
        Line2D([], [], color=CAUSAL, marker="o", label="Causal · raw"),
        Line2D([], [], color=CAUSAL, marker="^", linestyle="--", label="Causal · flow"),
    ]
    ax.legend(handles=handles, ncol=2, loc="lower left", columnspacing=0.8, handlelength=2.0)
    polish(ax)
    panel_label(ax, panel)


def plot_native_gain(
    ax: plt.Axes,
    rows: list[dict[str, str]],
    panel: str,
    sensitivity_applied: bool = False,
) -> None:
    table = lookup_rows(rows, ("method", "model_family", "layer_role", "probe"))
    x = np.arange(len(LAYERS))
    offsets = {"self_forcing": -0.11, "causal_forcing": 0.11}
    for family in ("self_forcing", "causal_forcing"):
        color = MODEL_COLOR[family]
        for method, probe, marker, filled in (
            ("linear", "fusion_same", "o", True),
            ("nonlinear", "both_correct", "D", False),
        ):
            values, lows, highs = [], [], []
            for role in LAYERS:
                row = table[(method, family, role, probe)]
                value = 100 * float(row["gain_mean"])
                values.append(value)
                lows.append(value - 100 * float(row["gain_ci95_low"]))
                highs.append(100 * float(row["gain_ci95_high"]) - value)
            ax.errorbar(
                x + offsets[family],
                values,
                yerr=np.asarray([lows, highs]),
                linestyle="none",
                marker=marker,
                markersize=5.2 if method == "linear" else 4.6,
                markerfacecolor=color if filled else "white",
                markeredgecolor=color,
                markeredgewidth=1.25,
                ecolor=color,
                elinewidth=1.1,
                capsize=2.5,
                zorder=4,
            )
    ax.axhline(0, color=INK, linewidth=0.8, zorder=1)
    ax.set_xticks(x, [value.title() for value in LAYERS])
    ax.set_ylabel("Held-out MSE reduction (%)")
    ax.set_ylim(-2.8, 5.8)
    title = "Native previous-boundary adds predictive value"
    if sensitivity_applied:
        title += "†"
    ax.set_title(title, loc="left", pad=8)
    handles = [
        Line2D([], [], color=SELF, marker="o", linestyle="none", label="Self · Ridge"),
        Line2D([], [], color=SELF, marker="D", markerfacecolor="white", linestyle="none", label="Self · MLP"),
        Line2D([], [], color=CAUSAL, marker="o", linestyle="none", label="Causal · Ridge"),
        Line2D([], [], color=CAUSAL, marker="D", markerfacecolor="white", linestyle="none", label="Causal · MLP"),
    ]
    ax.legend(handles=handles, ncol=2, loc="lower left", columnspacing=1.0)
    ax.text(
        0.995,
        0.03,
        "95% prompt-bootstrap CI",
        transform=ax.transAxes,
        ha="right",
        color=MUTED,
        fontsize=7.3,
    )
    polish(ax)
    panel_label(ax, panel)


def plot_aligned_predictor(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "probe"))
    x = np.arange(2)
    width = 0.29
    for index, family in enumerate(("self_forcing", "causal_forcing")):
        color = MODEL_COLOR[family]
        for offset, probe, label, shade in (
            (-width / 2, "both_raw", "Raw boundary", lighten(color, 0.28)),
            (width / 2, "both_flow", "Flow-aligned", color),
        ):
            row = table[(family, probe)]
            value = 100 * float(row["mse_gain_vs_step_mean"])
            low = value - 100 * float(row["mse_gain_vs_step_ci95_low"])
            high = 100 * float(row["mse_gain_vs_step_ci95_high"]) - value
            ax.bar(
                x[index] + offset,
                value,
                width=width,
                color=shade,
                edgecolor=color if probe == "both_raw" else "white",
                linewidth=0.9,
                hatch="///" if probe == "both_raw" else None,
                zorder=3,
            )
            ax.errorbar(
                x[index] + offset,
                value,
                yerr=np.asarray([[low], [high]]),
                fmt="none",
                ecolor=INK,
                elinewidth=0.9,
                capsize=2.2,
                zorder=4,
            )
        flow = table[(family, "both_flow")]
        delta = 100 * float(flow["mse_gain_vs_raw_mean"])
        ymax = max(
            100 * float(table[(family, "both_raw")]["mse_gain_vs_step_ci95_high"]),
            100 * float(flow["mse_gain_vs_step_ci95_high"]),
        )
        ax.text(x[index], ymax + 0.32, f"+{delta:.2f} pp", ha="center", color=color, fontsize=8, fontweight="semibold")
    ax.axhline(0, color=INK, linewidth=0.8)
    ax.set_xticks(x, ["Self", "Causal"])
    ax.set_ylim(0, 9.3)
    ax.set_ylabel("MSE reduction vs step-only (%)")
    ax.set_title("Alignment further improves prediction", loc="left", pad=8)
    ax.legend(
        handles=[
            Patch(facecolor="white", edgecolor=MUTED, hatch="///", label="Raw boundary"),
            Patch(facecolor="#334155", edgecolor="white", label="Flow-aligned"),
        ],
        loc="lower right",
    )
    polish(ax)
    panel_label(ax, panel)


def plot_unified_raw_gain(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "layer_role", "probe"))
    x = np.arange(len(LAYERS))
    offsets = {"self_forcing": -0.08, "causal_forcing": 0.08}
    for family in ("self_forcing", "causal_forcing"):
        values, lows, highs = [], [], []
        for role in LAYERS:
            row = table[(family, role, "both_raw")]
            value = 100 * float(row["mse_gain_vs_step_mean"])
            values.append(value)
            lows.append(value - 100 * float(row["mse_gain_vs_step_ci95_low"]))
            highs.append(100 * float(row["mse_gain_vs_step_ci95_high"]) - value)
        ax.errorbar(
            x + offsets[family],
            values,
            yerr=np.asarray([lows, highs]),
            color=MODEL_COLOR[family],
            linestyle="none",
            marker="o",
            markeredgecolor="white",
            markeredgewidth=0.7,
            markersize=5.3,
            elinewidth=1.1,
            capsize=2.5,
            label=MODEL_LABEL[family],
            zorder=4,
        )
    ax.axhline(0, color=INK, linewidth=0.8)
    ax.set_xticks(x, [value.title() for value in LAYERS])
    ax.set_ylim(0, 7.25)
    ax.set_ylabel("MSE reduction vs step-only (%)")
    ax.set_title("Raw previous-boundary adds predictive value", loc="left", pad=8)
    ax.legend(ncol=2, loc="upper left", columnspacing=1.0)
    ax.text(
        0.995,
        0.03,
        "95% prompt-bootstrap CI",
        transform=ax.transAxes,
        ha="right",
        color=MUTED,
        fontsize=7.3,
    )
    polish(ax)
    panel_label(ax, panel)


def plot_unified_alignment_gain(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "layer_role", "probe"))
    x = np.arange(len(LAYERS))
    for family in ("self_forcing", "causal_forcing"):
        values, lows, highs = [], [], []
        for role in LAYERS:
            row = table[(family, role, "both_flow")]
            value = 100 * float(row["mse_gain_vs_raw_mean"])
            values.append(value)
            lows.append(value - 100 * float(row["mse_gain_vs_raw_ci95_low"]))
            highs.append(100 * float(row["mse_gain_vs_raw_ci95_high"]) - value)
        ax.errorbar(
            x,
            values,
            yerr=np.asarray([lows, highs]),
            color=MODEL_COLOR[family],
            linewidth=1.55,
            marker="o" if family == "self_forcing" else "D",
            markersize=4.7,
            markeredgecolor="white",
            markeredgewidth=0.6,
            capsize=2.2,
            label=MODEL_LABEL[family],
            zorder=4,
        )
    ax.axhline(0, color=INK, linewidth=0.8)
    ax.set_xticks(x, ["Early", "Mid", "Late", "Final"])
    ax.set_ylim(0, 3.0)
    ax.set_ylabel("Additional MSE reduction (%)")
    ax.set_title("Flow alignment adds beyond raw", loc="left", pad=8)
    ax.legend(loc="upper left")
    polish(ax)
    panel_label(ax, panel)


def plot_unified_combined_gain(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    """Grouped raw/flow bars under one directly comparable y-axis."""
    table = lookup_rows(rows, ("model", "layer_role", "probe"))
    positions = np.asarray([0, 1, 2, 3, 5, 6, 7, 8], dtype=float)
    groups = [
        (family, role)
        for family in ("self_forcing", "causal_forcing")
        for role in LAYERS
    ]
    width = 0.34
    ax.axvspan(-0.55, 3.55, color=SELF, alpha=0.035, zorder=-3)
    ax.axvspan(4.45, 8.55, color=CAUSAL, alpha=0.035, zorder=-3)
    for index, (family, role) in enumerate(groups):
        color = MODEL_COLOR[family]
        raw = table[(family, role, "both_raw")]
        flow = table[(family, role, "both_flow")]
        raw_value = 100 * float(raw["mse_gain_vs_step_mean"])
        flow_value = 100 * float(flow["mse_gain_vs_step_mean"])
        raw_error = np.asarray([[
            raw_value - 100 * float(raw["mse_gain_vs_step_ci95_low"])
        ], [
            100 * float(raw["mse_gain_vs_step_ci95_high"]) - raw_value
        ]])
        flow_error = np.asarray([[
            flow_value - 100 * float(flow["mse_gain_vs_step_ci95_low"])
        ], [
            100 * float(flow["mse_gain_vs_step_ci95_high"]) - flow_value
        ]])
        x = positions[index]
        ax.bar(
            x - width / 2,
            raw_value,
            width=width,
            facecolor=lighten(color, 0.27),
            edgecolor=color,
            linewidth=0.9,
            hatch="///",
            zorder=3,
        )
        ax.bar(
            x + width / 2,
            flow_value,
            width=width,
            color=color,
            edgecolor="white",
            linewidth=0.7,
            zorder=3,
        )
        ax.errorbar(
            x - width / 2,
            raw_value,
            yerr=raw_error,
            fmt="none",
            ecolor=INK,
            elinewidth=0.85,
            capsize=2.1,
            zorder=4,
        )
        ax.errorbar(
            x + width / 2,
            flow_value,
            yerr=flow_error,
            fmt="none",
            ecolor=INK,
            elinewidth=0.85,
            capsize=2.1,
            zorder=4,
        )
        delta = 100 * float(flow["mse_gain_vs_raw_mean"])
        annotation_y = 100 * float(flow["mse_gain_vs_step_ci95_high"]) + 0.18
        ax.text(
            x + width / 2,
            annotation_y,
            f"+{delta:.2f}%",
            ha="center",
            va="bottom",
            color=color,
            fontsize=7.0,
            fontweight="semibold",
        )
    ax.axhline(0, color=INK, linewidth=0.8)
    ax.axvline(4.0, color=GRID, linewidth=0.9)
    ax.set_xticks(
        positions,
        [
            "Early\nSelf", "Middle\nSelf", "Late\nSelf", "Final\nSelf",
            "Early\nCausal", "Middle\nCausal", "Late\nCausal", "Final\nCausal",
        ],
    )
    ax.set_xlim(-0.65, 8.65)
    ax.set_ylim(0, 10.5)
    ax.set_ylabel("MSE reduction vs step-only (%)")
    ax.set_title("Raw and flow-aligned boundaries add predictive value", loc="left", pad=8)
    ax.legend(
        handles=[
            Patch(facecolor=lighten(SELF, 0.27), edgecolor=SELF, hatch="///", label="Self · raw"),
            Patch(facecolor=SELF, edgecolor="white", label="Self · flow-aligned"),
            Patch(facecolor=lighten(CAUSAL, 0.27), edgecolor=CAUSAL, hatch="///", label="Causal · raw"),
            Patch(facecolor=CAUSAL, edgecolor="white", label="Causal · flow-aligned"),
        ],
        ncol=4,
        loc="upper left",
        columnspacing=1.0,
        handlelength=1.6,
    )
    polish(ax)
    panel_label(ax, panel)


def main_figure(args: argparse.Namespace, output: Path) -> None:
    redundancy = read_csv(args.redundancy_csv)
    motion = read_csv(args.motion_bins_csv)
    native = read_csv(args.native_gain_csv)
    sensitivity_applied = apply_causal_late_sensitivity(
        native,
        args.causal_late_sensitivity_csv,
        args.sensitivity_scenario,
    )
    aligned = read_csv(args.aligned_probe_csv)
    unified = read_csv(args.unified_aligned_csv) if args.unified_aligned_csv else None
    fig = plt.figure(figsize=(13.2, 7.35), constrained_layout=True)
    grid = fig.add_gridspec(2, 3, height_ratios=(1.0, 1.07), width_ratios=(1, 1, 1.03))
    ax_a = fig.add_subplot(grid[0, 0])
    ax_b = fig.add_subplot(grid[0, 1], sharey=ax_a)
    ax_c = fig.add_subplot(grid[0, 2])
    if unified is None:
        ax_d = fig.add_subplot(grid[1, :2])
        ax_e = fig.add_subplot(grid[1, 2])
    else:
        ax_d = fig.add_subplot(grid[1, :])
    plot_redundancy(ax_a, redundancy, "self_forcing", "wan14b50", "A")
    plot_redundancy(ax_b, redundancy, "causal_forcing", "ar50", "B")
    ax_b.set_ylabel("")
    ax_b.tick_params(labelleft=False)
    ax_a.legend(
        ncol=1,
        loc="center left",
        bbox_to_anchor=(0.015, 0.50),
        handlelength=2.0,
        labelspacing=0.35,
        frameon=True,
        facecolor="white",
        edgecolor="none",
        framealpha=0.88,
    )
    plot_motion(ax_c, motion, "C")
    if unified is None:
        plot_native_gain(ax_d, native, "D", sensitivity_applied=sensitivity_applied)
        plot_aligned_predictor(ax_e, aligned, "E")
    else:
        plot_unified_combined_gain(ax_d, unified, "D")
    fig.suptitle(
        "Previous-chunk boundaries retain useful information—but spatial alignment matters",
        x=0.012,
        ha="left",
        fontsize=14,
        fontweight="bold",
        color=INK,
    )
    footer = (
        "10 prompts; prompt-first aggregation. Error bars show 95% prompt-bootstrap CI where available. "
        "Boundary = previous chunk's final temporal slot broadcast to all current slots."
    )
    if unified is not None:
        footer += " D: 64-D full grid/common mask; bars vs step-only; labels = flow gain over raw; 9-train/1-test; no exclusions."
    elif sensitivity_applied:
        footer += " † Causal late Ridge excludes prompt 6/step 1 and prompt 9/step 2 only."
    fig.text(
        0.012,
        -0.012,
        footer,
        ha="left",
        fontsize=7.8,
        color=MUTED,
    )
    fig.savefig(output / "self_causal_evidence_overview.png", dpi=240)
    fig.savefig(output / "self_causal_evidence_overview.pdf")
    plt.close(fig)


def plot_alignment_controls(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "action"))
    controls = [
        ("global_gain", "Global shift"),
        ("negated_gain", "Negated flow"),
        ("shuffled_gain", "Shuffled flow"),
        ("flow_gain", "Correct flow"),
    ]
    y = np.arange(len(controls))
    ax.axhspan(2.55, 3.45, color="#ECFDF5", zorder=-2)
    values_by_family = {}
    for family, offset in (("self_forcing", -0.09), ("causal_forcing", 0.09)):
        row = table[(family, "none")]
        values = [100 * float(row[key]) for key, _ in controls]
        values_by_family[family] = values
        ax.scatter(values, y + offset, s=30, color=MODEL_COLOR[family], edgecolor="white", linewidth=0.6, zorder=3, label=MODEL_LABEL[family])
    for index in range(len(controls)):
        ax.plot(
            [values_by_family["self_forcing"][index], values_by_family["causal_forcing"][index]],
            [y[index] - 0.09, y[index] + 0.09],
            color="#AAB7C8",
            linewidth=1.0,
            zorder=1,
        )
    ax.set_yticks(y, [label for _, label in controls])
    ax.set_xlabel("Cosine recovery over raw (×100)")
    ax.set_xlim(0, 1.82)
    ax.set_title("Interpolation controls", loc="left", pad=8)
    ax.legend(loc="upper left")
    polish(ax, "x")
    panel_label(ax, panel)


def plot_predictor_controls(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("method", "model_family", "layer_role", "probe"))
    controls = [
        ("fusion_same", "Correct boundary"),
        ("fusion_wrong_step", "Wrong timestep"),
        ("fusion_distant", "Distant boundary"),
        ("fusion_token_shuffle", "Spatial shuffle"),
        ("fusion_batch_shuffle", "Other video"),
        ("fusion_zero", "Zero"),
        ("fusion_noise", "Matched noise"),
    ]
    y = np.arange(len(controls))
    ax.axhspan(-0.43, 0.43, color="#ECFDF5", zorder=-2)
    for family, offset in (("self_forcing", -0.10), ("causal_forcing", 0.10)):
        values, lows, highs = [], [], []
        for probe, _ in controls:
            row = table[("linear", family, "final", probe)]
            value = 100 * float(row["gain_mean"])
            values.append(value)
            lows.append(value - 100 * float(row["gain_ci95_low"]))
            highs.append(100 * float(row["gain_ci95_high"]) - value)
        ax.errorbar(
            values,
            y + offset,
            xerr=np.asarray([lows, highs]),
            fmt="o",
            color=MODEL_COLOR[family],
            markeredgecolor="white",
            markeredgewidth=0.6,
            markersize=5.0,
            elinewidth=1.0,
            capsize=2,
            label=MODEL_LABEL[family],
            zorder=3,
        )
    ax.axvline(0, color=INK, linewidth=0.8)
    ax.set_yticks(y, [label for _, label in controls])
    ax.set_xlabel("Final-layer held-out MSE reduction (%)")
    ax.set_xlim(-0.45, 4.5)
    ax.set_title("Predictor controls", loc="left", pad=8)
    ax.legend(loc="lower right")
    ax.invert_yaxis()
    polish(ax, "x")
    panel_label(ax, panel)


def plot_unified_predictor_controls(ax: plt.Axes, rows: list[dict[str, str]], panel: str) -> None:
    table = lookup_rows(rows, ("model", "layer_role", "probe"))
    controls = [
        ("both_raw", "Raw boundary"),
        ("both_global", "Global shift"),
        ("both_negated_flow", "Negated flow"),
        ("both_shuffled_flow", "Shuffled flow"),
        ("both_flow", "Correct flow"),
    ]
    y = np.arange(len(controls))
    ax.axhspan(3.57, 4.43, color="#ECFDF5", zorder=-2)
    for family, offset in (("self_forcing", -0.10), ("causal_forcing", 0.10)):
        values, lows, highs = [], [], []
        for probe, _ in controls:
            row = table[(family, "final", probe)]
            value = 100 * float(row["mse_gain_vs_step_mean"])
            values.append(value)
            lows.append(value - 100 * float(row["mse_gain_vs_step_ci95_low"]))
            highs.append(100 * float(row["mse_gain_vs_step_ci95_high"]) - value)
        ax.errorbar(
            values,
            y + offset,
            xerr=np.asarray([lows, highs]),
            fmt="o",
            color=MODEL_COLOR[family],
            markeredgecolor="white",
            markeredgewidth=0.6,
            markersize=5.0,
            elinewidth=1.0,
            capsize=2,
            label=MODEL_LABEL[family],
            zorder=3,
        )
    ax.axvline(0, color=INK, linewidth=0.8)
    ax.set_yticks(y, [label for _, label in controls])
    ax.set_xlabel("Final-layer MSE reduction vs step-only (%)")
    ax.set_xlim(0, 9.5)
    ax.set_title("Unified predictor controls", loc="left", pad=8)
    ax.legend(loc="upper left")
    ax.invert_yaxis()
    polish(ax, "x")
    panel_label(ax, panel)


def controls_figure(args: argparse.Namespace, output: Path) -> None:
    motion = read_csv(args.motion_summary_csv)
    native = read_csv(args.native_gain_csv)
    unified = read_csv(args.unified_aligned_csv) if args.unified_aligned_csv else None
    fig, axes = plt.subplots(1, 2, figsize=(11.4, 4.25), constrained_layout=True)
    plot_alignment_controls(axes[0], motion, "A")
    if unified is None:
        plot_predictor_controls(axes[1], native, "B")
    else:
        plot_unified_predictor_controls(axes[1], unified, "B")
    fig.suptitle(
        "Control experiments isolate correct spatial correspondence",
        x=0.012,
        ha="left",
        fontsize=13,
        fontweight="bold",
        color=INK,
    )
    fig.text(
        0.012,
        -0.02,
        "Green bands mark correct-flow conditions. Predictor error bars: 95% prompt-bootstrap CI.",
        ha="left",
        fontsize=7.8,
        color=MUTED,
    )
    fig.savefig(output / "self_causal_control_checks.png", dpi=240)
    fig.savefig(output / "self_causal_control_checks.pdf")
    plt.close(fig)


def main() -> None:
    args = parse_args()
    setup_style()
    output = args.output_dir.resolve()
    output.mkdir(parents=True, exist_ok=True)
    main_figure(args, output)
    controls_figure(args, output)
    print(f"[complete] {output}", flush=True)


if __name__ == "__main__":
    main()