code/audit_claims.py

#!/usr/bin/env python3
"""Independent audit of the five major claims in arXiv:2606.15589.

The script intentionally uses only NumPy plus the Python standard library. It
reads the official repository's checked-in JSONL artifacts, reconstructs the
paper's 40-task analysis set, recomputes the paired route statistics, audits
the Figure 5 reconstruction files, and numerically checks the linear
decision-theoretic covariance/risk statement.
"""

from __future__ import annotations

import argparse
import csv
import glob
import json
import math
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterable

import numpy as np


MODELS = (
    "anthropic/claude-haiku-4.5",
    "google/gemini-2.0-flash-001",
    "google/gemini-2.5-flash",
    "openai/gpt-4o-mini",
    "mistralai/codestral-2508",
    "mistralai/mixtral-8x22b-instruct",
)

# Reconstructed from the 48 checked-in task kinds using the paper's Table 1
# class/task/instance counts and Table 2 per-model results. The current
# repository's route-table script instead retains 44 tasks.
EXCLUDED_FROM_PAPER_40 = {
    "bsp",
    "msp",
    "gcp_d",
    "tsp_d",
    "segments_intersect",
    "knap",
    "gcp",
    "spp",
}

ROUTES = {
    "route_1_nl": "nl_correct",
    "route_2_sim": "sim_correct",
    "route_3_exec": "code_correct",
}

TRANSLATION_FILES = {
    "Claude Haiku 4.5": "translation_claude-haiku-4.5_20260127_081757_trials.jsonl",
    "Gemini 2.5 Flash": "translation_gemini-2.5-flash_20260127_082539_trials.jsonl",
    "Mixtral 8x22B": "translation_mixtral_20260127_180139_trials.jsonl",
}

PUBLISHED = {
    "route_1_accuracy_pct": 17.21,
    "route_2_accuracy_pct": 17.37,
    "route_3_accuracy_pct": 48.84,
    "route_2_minus_route_1_pp": 0.15,
    "route_2_minus_route_1_ci_pct": [-0.30, 0.61],
    "route_3_minus_route_2_pp": 31.47,
    "route_3_minus_route_2_ci_pct": [29.20, 33.71],
    "execution_win_mass_pct": 33.08,
    "recovery_mass_pct": 1.61,
    "models": 6,
    "seeds": 3,
    "tasks": 40,
    "instances": 1113,
    "rows": 20034,
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--toolproj", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--bootstrap", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--theory-samples", type=int, default=200_000)
    return parser.parse_args()


def instance_key(row: dict[str, Any]) -> tuple[str, int, int]:
    return (
        str(row["kind"]),
        int(row["digit"]),
        int(row["index_in_kind"]),
    )


def load_route_rows(toolproj: Path) -> tuple[list[dict[str, Any]], list[str]]:
    pattern = toolproj / "src/exps_performance/results/*/tb/run_*/res.jsonl"
    paths = sorted(glob.glob(str(pattern)))
    model_set = set(MODELS)
    rows: list[dict[str, Any]] = []
    used_paths: list[str] = []
    for path_str in paths:
        path = Path(path_str)
        accepted = 0
        with path.open(encoding="utf-8") as handle:
            for line in handle:
                if not line.strip():
                    continue
                row = json.loads(line)
                if row.get("model") not in model_set:
                    continue
                if row.get("kind") in EXCLUDED_FROM_PAPER_40:
                    continue
                rows.append(row)
                accepted += 1
        if accepted:
            used_paths.append(str(path))
    if not rows:
        raise RuntimeError(f"No route rows found under {pattern}")
    return rows, used_paths


def exact_binomial_two_sided(n01: int, n10: int) -> float:
    """Exact two-sided McNemar p-value via a stable binomial tail."""
    n = n01 + n10
    if n == 0:
        return 1.0
    tail_k = min(n01, n10)
    log_terms = [
        math.lgamma(n + 1)
        - math.lgamma(k + 1)
        - math.lgamma(n - k + 1)
        - n * math.log(2.0)
        for k in range(tail_k + 1)
    ]
    max_log = max(log_terms)
    tail = math.exp(max_log) * sum(math.exp(value - max_log) for value in log_terms)
    return min(1.0, 2.0 * tail)


def paired_counts(
    rows: Iterable[dict[str, Any]],
    first: str,
    second: str,
) -> dict[str, int]:
    counts = {"both_wrong": 0, "second_only": 0, "first_only": 0, "both_correct": 0}
    for row in rows:
        a = bool(row[first])
        b = bool(row[second])
        if not a and not b:
            counts["both_wrong"] += 1
        elif not a and b:
            counts["second_only"] += 1
        elif a and not b:
            counts["first_only"] += 1
        else:
            counts["both_correct"] += 1
    return counts


def cluster_bootstrap(
    rows: list[dict[str, Any]],
    first: str,
    second: str,
    *,
    repetitions: int,
    seed: int,
) -> tuple[float, tuple[float, float], np.ndarray]:
    grouped: dict[tuple[str, int, int], list[float]] = defaultdict(list)
    for row in rows:
        grouped[instance_key(row)].append(float(bool(row[second])) - float(bool(row[first])))
    cluster_deltas = np.asarray(
        [np.mean(grouped[key]) for key in sorted(grouped)],
        dtype=np.float64,
    )
    rng = np.random.RandomState(seed)
    n_clusters = len(cluster_deltas)
    indices = rng.choice(n_clusters, size=(repetitions, n_clusters), replace=True)
    bootstrap_pp = cluster_deltas[indices].mean(axis=1) * 100.0
    observed_pp = float(cluster_deltas.mean() * 100.0)
    ci = tuple(float(value) for value in np.percentile(bootstrap_pp, [2.5, 97.5]))
    return observed_pp, (ci[0], ci[1]), bootstrap_pp


def holm_two(p_values: tuple[float, float]) -> tuple[float, float]:
    order = sorted(range(2), key=lambda index: p_values[index])
    adjusted = [0.0, 0.0]
    running = 0.0
    for rank, index in enumerate(order):
        candidate = min(1.0, (2 - rank) * p_values[index])
        running = max(running, candidate)
        adjusted[index] = running
    return adjusted[0], adjusted[1]


def summarize_routes(
    rows: list[dict[str, Any]],
    bootstrap: int,
    seed: int,
) -> tuple[dict[str, Any], np.ndarray, np.ndarray]:
    n = len(rows)
    correct = {
        label: sum(bool(row[field]) for row in rows)
        for label, field in ROUTES.items()
    }
    accuracies = {label: count / n * 100.0 for label, count in correct.items()}

    r21 = paired_counts(rows, "nl_correct", "sim_correct")
    r32 = paired_counts(rows, "sim_correct", "code_correct")
    p21 = exact_binomial_two_sided(r21["second_only"], r21["first_only"])
    p32 = exact_binomial_two_sided(r32["second_only"], r32["first_only"])
    holm21, holm32 = holm_two((p21, p32))
    delta21, ci21, boot21 = cluster_bootstrap(
        rows,
        "nl_correct",
        "sim_correct",
        repetitions=bootstrap,
        seed=seed,
    )
    delta32, ci32, boot32 = cluster_bootstrap(
        rows,
        "sim_correct",
        "code_correct",
        repetitions=bootstrap,
        seed=seed,
    )
    result = {
        "correct_counts": correct,
        "accuracy_pct": accuracies,
        "route_2_vs_route_1": {
            "paired_table": r21,
            "delta_pp": delta21,
            "cluster_bootstrap_95_ci_pp": list(ci21),
            "bootstrap_repetitions": bootstrap,
            "bootstrap_seed": seed,
            "mcnemar_exact_two_sided_raw_p": p21,
            "mcnemar_holm_adjusted_p": holm21,
        },
        "route_3_vs_route_2": {
            "paired_table": r32,
            "delta_pp": delta32,
            "cluster_bootstrap_95_ci_pp": list(ci32),
            "bootstrap_repetitions": bootstrap,
            "bootstrap_seed": seed,
            "mcnemar_exact_two_sided_raw_p": p32,
            "mcnemar_holm_adjusted_p": holm32,
            "execution_win_mass_pct": r32["second_only"] / n * 100.0,
            "recovery_mass_pct": r32["first_only"] / n * 100.0,
        },
    }
    return result, boot21, boot32


def summarize_coverage(rows: list[dict[str, Any]], used_paths: list[str]) -> dict[str, Any]:
    model_values = sorted({str(row["model"]) for row in rows})
    seed_values = sorted({int(row["seed"]) for row in rows})
    task_values = sorted({str(row["kind"]) for row in rows})
    instance_values = sorted({instance_key(row) for row in rows})
    rows_per_model = {
        model: sum(row["model"] == model for row in rows)
        for model in model_values
    }
    rows_per_instance = defaultdict(int)
    for row in rows:
        rows_per_instance[instance_key(row)] += 1
    return {
        "models": model_values,
        "model_count": len(model_values),
        "seeds": seed_values,
        "seed_count": len(seed_values),
        "tasks": task_values,
        "task_count": len(task_values),
        "instances": len(instance_values),
        "rows": len(rows),
        "rows_per_model": rows_per_model,
        "rows_per_instance_min": min(rows_per_instance.values()),
        "rows_per_instance_max": max(rows_per_instance.values()),
        "source_files": used_paths,
        "source_file_count": len(used_paths),
        "excluded_from_checked_in_48": sorted(EXCLUDED_FROM_PAPER_40),
        "selection_note": (
            "The current repository does not encode the paper's 40-task filter. "
            "This exclusion set is the unique match found from the paper's "
            "Table 1 class/task/instance counts and Table 2 per-model values."
        ),
    }


def audit_translation(toolproj: Path) -> list[dict[str, Any]]:
    results_dir = toolproj / "src/translation_additivity/results"
    output: list[dict[str, Any]] = []
    conditions = ("x", "x_nl_native", "x_nl_translated")
    for model, filename in TRANSLATION_FILES.items():
        by_sample: dict[str, dict[str, bool]] = defaultdict(dict)
        trial_rows = 0
        with (results_dir / filename).open(encoding="utf-8") as handle:
            for line in handle:
                if not line.strip():
                    continue
                row = json.loads(line)
                trial_rows += 1
                by_sample[str(row["sample_id"])][str(row["condition"])] = bool(row["correct"])
        complete_ids = sorted(
            sample_id
            for sample_id, values in by_sample.items()
            if all(condition in values for condition in conditions)
        )
        arrays = {
            condition: [by_sample[sample_id][condition] for sample_id in complete_ids]
            for condition in conditions
        }
        native_vs_translated = paired_counts(
            [
                {
                    "native": arrays["x_nl_native"][index],
                    "translated": arrays["x_nl_translated"][index],
                }
                for index in range(len(complete_ids))
            ],
            "native",
            "translated",
        )
        output.append(
            {
                "model": model,
                "file": str(results_dir / filename),
                "trial_rows": trial_rows,
                "complete_unique_sample_ids": len(complete_ids),
                "accuracy_pct": {
                    condition: sum(arrays[condition]) / len(complete_ids) * 100.0
                    for condition in conditions
                },
                "native_vs_translated": {
                    "paired_table": native_vs_translated,
                    "mcnemar_exact_two_sided_p": exact_binomial_two_sided(
                        native_vs_translated["second_only"],
                        native_vs_translated["first_only"],
                    ),
                },
            }
        )
    return output


def numerical_theory_audit(samples: int, seed: int) -> dict[str, Any]:
    """Monte Carlo check of Proposition 4.2 and Theorem 4.3."""
    rng = np.random.RandomState(seed)
    d_core = 4
    d_nuisance = 5
    b = rng.normal(size=(samples, d_core))
    u_code = rng.normal(size=(samples, d_nuisance))
    # Anisotropic extra NL nuisance remains conditionally mean-zero.
    scales = np.asarray([0.25, 0.5, 0.75, 1.0, 1.25])
    eta_nl = rng.normal(size=(samples, d_nuisance)) * scales
    u_nl = u_code + eta_nl
    epsilon = rng.normal(scale=0.2, size=samples)
    theta = np.asarray([0.4, -0.7, 1.1, 0.2])
    y = b @ theta + epsilon

    sigma_code = u_code.T @ u_code / samples
    sigma_nl = u_nl.T @ u_nl / samples
    covariance_difference = sigma_nl - sigma_code
    eigenvalues = np.linalg.eigvalsh(covariance_difference)

    risk_differences: list[float] = []
    formula_residuals: list[float] = []
    for _ in range(1000):
        a = rng.normal(size=d_core)
        v = rng.normal(size=d_nuisance)
        pred_code = b @ a + u_code @ v
        pred_nl = b @ a + u_nl @ v
        risk_code = float(np.mean((pred_code - y) ** 2))
        risk_nl = float(np.mean((pred_nl - y) ** 2))
        empirical_difference = risk_code - risk_nl
        covariance_formula = float(v @ (sigma_code - sigma_nl) @ v)
        risk_differences.append(empirical_difference)
        formula_residuals.append(empirical_difference - covariance_formula)

    # Control: violate the shared-core/no-hidden-answer-channel assumption by
    # injecting label information into the NL-only nuisance. NL can then win.
    signal = y - float(np.mean(y))
    u_nl_control = u_nl.copy()
    u_nl_control[:, 0] += 2.0 * signal
    control_risks = []
    for coefficient in np.linspace(-2.0, 2.0, 801):
        prediction = coefficient * u_nl_control[:, 0]
        control_risks.append(float(np.mean((prediction - y) ** 2)))
    best_control_risk = min(control_risks)
    zero_predictor_risk = float(np.mean(y**2))

    return {
        "samples": samples,
        "seed": seed,
        "min_eigenvalue_sigma_nl_minus_sigma_code": float(eigenvalues.min()),
        "max_eigenvalue_sigma_nl_minus_sigma_code": float(eigenvalues.max()),
        "max_observed_code_minus_nl_risk": float(max(risk_differences)),
        "min_observed_code_minus_nl_risk": float(min(risk_differences)),
        "max_abs_risk_identity_residual": float(max(abs(x) for x in formula_residuals)),
        "control_assumption_relaxed": {
            "zero_predictor_risk": zero_predictor_risk,
            "best_nl_nuisance_predictor_risk": best_control_risk,
            "risk_reduction": zero_predictor_risk - best_control_risk,
            "interpretation": (
                "When NL-only nuisance is allowed to carry label signal, the "
                "non-inferiority conclusion need not hold."
            ),
        },
    }


def write_json(path: Path, payload: Any) -> None:
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def write_bootstrap_csv(path: Path, first: np.ndarray, second: np.ndarray) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["replicate", "route_2_minus_route_1_pp", "route_3_minus_route_2_pp"])
        for index, (value_1, value_2) in enumerate(zip(first, second)):
            writer.writerow([index, f"{value_1:.10f}", f"{value_2:.10f}"])


def write_summary_csv(
    path: Path,
    route: dict[str, Any],
    coverage: dict[str, Any],
    translation: list[dict[str, Any]],
) -> None:
    rows = [
        ("Claim 1", "Route 3 - Route 2 (pp)", route["route_3_vs_route_2"]["delta_pp"], 31.47),
        (
            "Claim 1",
            "CI lower (pp)",
            route["route_3_vs_route_2"]["cluster_bootstrap_95_ci_pp"][0],
            29.20,
        ),
        (
            "Claim 1",
            "CI upper (pp)",
            route["route_3_vs_route_2"]["cluster_bootstrap_95_ci_pp"][1],
            33.71,
        ),
        ("Claim 2", "Route 1 accuracy (%)", route["accuracy_pct"]["route_1_nl"], 17.21),
        ("Claim 2", "Route 2 accuracy (%)", route["accuracy_pct"]["route_2_sim"], 17.37),
        ("Claim 2", "Route 2 - Route 1 (pp)", route["route_2_vs_route_1"]["delta_pp"], 0.15),
        (
            "Claim 3",
            "Execution-win mass (%)",
            route["route_3_vs_route_2"]["execution_win_mass_pct"],
            33.08,
        ),
        (
            "Claim 3",
            "Recovery mass (%)",
            route["route_3_vs_route_2"]["recovery_mass_pct"],
            1.61,
        ),
        ("Claim 5", "Models", coverage["model_count"], 6),
        ("Claim 5", "Seeds", coverage["seed_count"], 3),
        ("Claim 5", "Tasks", coverage["task_count"], 40),
        ("Claim 5", "Instances", coverage["instances"], 1113),
        ("Claim 5", "Rows", coverage["rows"], 20034),
    ]
    for item in translation:
        rows.extend(
            [
                ("Claim 4", f"{item['model']} baseline (%)", item["accuracy_pct"]["x"], ""),
                ("Claim 4", f"{item['model']} native NL (%)", item["accuracy_pct"]["x_nl_native"], ""),
                (
                    "Claim 4",
                    f"{item['model']} translated NL (%)",
                    item["accuracy_pct"]["x_nl_translated"],
                    "",
                ),
            ]
        )
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["claim", "metric", "reproduced", "published"])
        writer.writerows(rows)


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rows, used_paths = load_route_rows(args.toolproj)
    route, boot21, boot32 = summarize_routes(rows, args.bootstrap, args.seed)
    coverage = summarize_coverage(rows, used_paths)
    translation = audit_translation(args.toolproj)
    theory = numerical_theory_audit(args.theory_samples, args.seed)
    payload = {
        "paper": "Is Code Better Than Language for Algorithmic Reasoning?",
        "paper_url": "https://arxiv.org/abs/2606.15589",
        "official_repo_commit": (
            "https://github.com/TerryTong-Git/ToolProj/tree/"
            "b1d171ac08d2333197807493136d8f82b470e10d"
        ),
        "published_claims": PUBLISHED,
        "route_audit": route,
        "coverage_audit": coverage,
        "translation_reconstruction_audit": translation,
        "linear_theory_numerical_audit": theory,
    }
    write_json(args.output_dir / "audit_summary.json", payload)
    write_json(args.output_dir / "coverage.json", coverage)
    write_json(args.output_dir / "translation_reconstruction.json", translation)
    write_json(args.output_dir / "linear_theory_audit.json", theory)
    write_bootstrap_csv(args.output_dir / "bootstrap_samples.csv", boot21, boot32)
    write_summary_csv(args.output_dir / "claim_summary.csv", route, coverage, translation)

    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()