compute_bootstrap_final.py

#!/usr/bin/env python3
"""Compute bootstrap standard errors and confidence intervals for final sweeps."""

import argparse
import csv
import json
import math
from pathlib import Path
from typing import Dict, Iterable, List, Mapping, Sequence

import numpy as np
from tqdm import tqdm

from bootstrap_utils import bca_interval, cluster_bootstrap, jackknife_over_tasks


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Bootstrap accuracy estimates for all result files in a directory."
    )
    parser.add_argument(
        "--results-dir",
        type=Path,
        default=Path("sweep_results_final"),
        help="Directory containing per-run JSON files (default: sweep_results_final).",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Optional CSV destination. Defaults to <results-dir>/bootstrap_summary.csv.",
    )
    parser.add_argument(
        "--bootstrap-samples",
        type=int,
        default=5000,
        help="Number of bootstrap replicates to draw per file (default: 5000).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=123,
        help="Random seed for the bootstrap RNG (default: 123).",
    )
    return parser.parse_args()


def load_json(path: Path) -> Dict:
    try:
        return json.loads(path.read_text())
    except FileNotFoundError as exc:
        raise SystemExit(f"Result file not found: {path}") from exc
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Invalid JSON in result file: {path}") from exc


def extract_outcomes(task_entry: Mapping) -> np.ndarray:
    trials = task_entry.get("trials")
    if trials is None:
        raise ValueError("Task entry missing 'trials' key")
    outcomes: List[int] = []
    for trial in trials:
        if "correct" not in trial:
            raise ValueError("Trial entry missing 'correct' field")
        outcomes.append(1 if trial["correct"] else 0)
    if not outcomes:
        raise ValueError("Encountered task with zero trials; cannot bootstrap")
    return np.array(outcomes, dtype=float)


def load_task_matrix(result_path: Path) -> Sequence[np.ndarray]:
    data = load_json(result_path)
    tasks = data.get("tasks")
    if tasks is None:
        raise SystemExit(f"Result file missing 'tasks' list: {result_path}")
    outcomes = [extract_outcomes(task) for task in tasks]
    if not outcomes:
        raise SystemExit(f"Result file had no tasks: {result_path}")
    return outcomes


def point_accuracy(outcomes: Sequence[np.ndarray]) -> float:
    per_task = [task.mean() for task in outcomes]
    return float(np.mean(per_task))


def summarize_file(
    path: Path,
    *,
    n_bootstrap: int,
    seed: int,
) -> Dict[str, object]:
    data = load_json(path)
    config = data.get("config", {})
    outcomes = load_task_matrix(path)
    p_hat = point_accuracy(outcomes)

    reported = data.get("overall", {}).get("mean_accuracy")
    if reported is not None and not math.isclose(p_hat, reported, rel_tol=1e-9, abs_tol=1e-9):
        raise SystemExit(
            f"Computed accuracy {p_hat:.12f} disagrees with stored mean_accuracy "
            f"{reported:.12f} in {path}"
        )

    boot = cluster_bootstrap(outcomes, n_bootstrap=n_bootstrap, seed=seed)
    se = float(np.std(boot, ddof=1))
    pct_low, pct_high = [float(x) for x in np.quantile(boot, [0.025, 0.975])]
    jackknife = jackknife_over_tasks(outcomes)
    bca_low, bca_high = bca_interval(boot=boot, stat=p_hat, jackknife=jackknife)

    filename = path.name
    if filename.startswith("naive_modal_ic"):
        filename = filename.replace("naive_modal_", "naive_", 1)

    return {
        "filename": filename,
        "model": config.get("model"),
        "backend": config.get("backend"),
        "lmshuffle": config.get("lmshuffle"),
        "in_context_examples": config.get("in_context_examples"),
        "accuracy": p_hat,
        "se": se,
        "ci_low": pct_low,
        "ci_high": pct_high,
        "bca_low": bca_low,
        "bca_high": bca_high,
    }


def iter_result_files(results_dir: Path) -> Iterable[Path]:
    if not results_dir.exists():
        raise SystemExit(f"Results directory not found: {results_dir}")
    for path in sorted(results_dir.glob("*.json")):
        name = path.name
        if name == "sweep_summary.json":
            continue
        if name.startswith("naive_identity_ic"):
            continue
        if name.startswith("naive_ic"):
            continue
        yield path


def write_csv(output_path: Path, rows: Sequence[Dict[str, object]]) -> None:
    header = [
        "filename",
        "model",
        "backend",
        "lmshuffle",
        "in_context_examples",
        "accuracy",
        "se",
        "ci_low",
        "ci_high",
        "bca_low",
        "bca_high",
    ]
    with output_path.open("w", newline="", encoding="ascii") as handle:
        writer = csv.DictWriter(handle, fieldnames=header)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def main() -> None:
    args = parse_args()
    results_dir: Path = args.results_dir
    output_path = args.output or results_dir / "bootstrap_summary.csv"

    rows: List[Dict[str, object]] = []
    result_paths = list(iter_result_files(results_dir))
    for seed_offset, path in enumerate(
        tqdm(result_paths, desc="Bootstrapping", unit="file", leave=False)
    ):
        row = summarize_file(
            path,
            n_bootstrap=args.bootstrap_samples,
            seed=args.seed + seed_offset,
        )
        rows.append(row)

    write_csv(output_path, rows)
    print(f"Wrote bootstrap summary for {len(rows)} files to {output_path}")


if __name__ == "__main__":
    main()