path: root/experiments/analyze_kp_innovation_confirmation.py

#!/usr/bin/env python3
"""Audit the untouched MT-3 five-seed mixed-traffic test confirmation."""
import argparse
import json
import math
import os
import statistics


CONDITIONS = ("clean", "raw", "matched", "innovation")
SEEDS = tuple(range(10, 15))
T_CRITICAL_ONE_SIDED_95_DF4 = 2.131846786


def mean_early(values):
    count = max(1, len(values) // 3)
    return sum(float(value) for value in values[:count]) / count


def lower_confidence_bound(values):
    return (statistics.mean(values)
            - T_CRITICAL_ONE_SIDED_95_DF4
            * statistics.stdev(values) / math.sqrt(len(values)))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_dir", default="results/kp_innovation_confirmation")
    parser.add_argument(
        "--full_gate", default="results/kp_innovation_full_gate.json")
    parser.add_argument(
        "--out", default="results/kp_innovation_confirmation_gate.json")
    args = parser.parse_args()
    with open(args.full_gate) as handle:
        full_gate = json.load(handle)
    if (full_gate.get("protocol") != "kp_mixed_traffic_full_v1"
            or full_gate.get("status") != "passed"):
        raise ValueError("MT-2 did not pass")

    common = {
        "depth": 20, "width": 16, "batch_size": 128, "epochs": 200,
        "train_limit": 0, "val_examples": 0, "split_seed": 2027,
        "eval_split": "test", "eval_every": 0, "augment_train": 1,
        "lr": 0.1, "output_lr": 0.1, "lr_schedule": "step",
        "lr_milestones": "100,150", "lr_gamma": 0.1,
        "warmup_epochs": 0, "momentum": 0.9, "weight_decay": 1e-4,
        "normalization": "batchnorm", "a_scale": 1.0,
        "alignment_probe": 32,
    }
    traffic = {
        "traffic_ratio": 4.0, "traffic_calibration_examples": 64,
        "learn_P": 1, "eta_P": 0.1, "predictor_warmup_steps": 20,
        "predictor_every": 16,
    }
    records = {}
    source_commits = set()
    for seed in SEEDS:
        for condition in CONDITIONS:
            path = os.path.join(
                args.input_dir, f"seed{seed}_{condition}.json")
            with open(path) as handle:
                record = json.load(handle)
            records[(seed, condition)] = record
            expected = {**common, "seed": seed, "loader_seed": seed,
                        "mode": "kp" if condition == "clean" else "kp_traffic"}
            if condition != "clean":
                expected.update(traffic)
                expected.update({"traffic_rule": condition,
                                 "traffic_seed": 5000 + seed})
            for key, value in expected.items():
                if record["args"].get(key) != value:
                    raise ValueError(
                        f"MT-3 seed {seed} {condition} {key} drift")
            if record["provenance"]["git_tracked_dirty"]:
                raise ValueError(
                    f"tracked-dirty MT-3 seed {seed} {condition}")
            source_commits.add(record["provenance"]["git_commit"])
            expected_space = (
                "reciprocal_local_activity_products" if condition == "clean"
                else "reciprocal_local_activity_products_with_mixed_apical_traffic")
            if record.get("calibration_metric_space") != expected_space:
                raise ValueError(
                    f"MT-3 seed {seed} {condition} metric-space drift")
            if condition != "clean":
                warmup = record.get("predictor_warmup", {})
                if (warmup.get("instruction_present") is not False
                        or warmup.get("task_loader_state_restored") is not True):
                    raise ValueError(
                        f"MT-3 seed {seed} {condition} warmup invariant failed")
            split = record["split"]
            if not (split["validation_examples"] == 0
                    and split["validation_index_sha256"] is None
                    and split["train_examples"] == 50000
                    and split["test_examples"] == 10000):
                raise ValueError(
                    f"MT-3 seed {seed} {condition} split drift")
            evaluation = record["evaluation_protocol"]
            if not (evaluation["validation_evaluations"] == 0
                    and evaluation["test_evaluations"] == 1
                    and not evaluation["test_used_for_selection"]):
                raise ValueError(
                    f"MT-3 seed {seed} {condition} evaluation drift")
    if len(source_commits) != 1:
        raise ValueError("all MT-3 runs must share one source revision")

    accuracies = {
        condition: [float(records[(seed, condition)]["final"]["accuracy"])
                    for seed in SEEDS]
        for condition in CONDITIONS
    }
    raw_gains = [innovation - raw for innovation, raw in zip(
        accuracies["innovation"], accuracies["raw"])]
    matched_gains = [innovation - matched for innovation, matched in zip(
        accuracies["innovation"], accuracies["matched"])]
    clean_deficits = [clean - innovation for clean, innovation in zip(
        accuracies["clean"], accuracies["innovation"])]

    finite_values = []
    invariant_failures = []
    bp_macs_50k = int(
        full_gate["metrics"]["total_macs"]["innovation"]
        / full_gate["metrics"]["mac_ratio_to_bp"]["innovation"]
        * (50_000 / 45_000))
    for seed in SEEDS:
        for condition in CONDITIONS:
            record = records[(seed, condition)]
            if not record["final"]["finite"]:
                invariant_failures.append(f"seed{seed}_{condition}:nonfinite")
            finite_values.extend([
                float(record["final"]["accuracy"]),
                float(record["final"]["loss"]),
            ])
            if int(record["work"]["logical_batch_loss_queries"]) != 0:
                invariant_failures.append(f"seed{seed}_{condition}:queries")
            if int(record["work"]["total_macs_estimate"]) > 1.40 * bp_macs_50k:
                invariant_failures.append(f"seed{seed}_{condition}:cost")
            trajectory = record["epochs"]
            if len(trajectory) != 200:
                invariant_failures.append(f"seed{seed}_{condition}:trajectory")
                continue
            finite_values.extend(float(row["train_loss"]) for row in trajectory)
            tracking = [row.get("feedback_tracking") for row in trajectory]
            if any(value is None for value in tracking):
                invariant_failures.append(f"seed{seed}_{condition}:tracking")
                continue
            for value in tracking:
                finite_values.extend([
                    float(value["mean_feedback_forward_cosine"]),
                    float(value["mean_feedback_forward_relative_error"]),
                    float(value["min_feedback_forward_cosine"]),
                    float(value["max_feedback_forward_relative_error"]),
                ])
            final_cosine = float(
                record["diagnostics"]["mean_feedback_forward_cosine"])
            late_cosine = statistics.mean(float(
                value["mean_feedback_forward_cosine"]) for value in tracking[150:])
            finite_values.extend([final_cosine, late_cosine])
            if final_cosine < 0.95 or late_cosine < 0.95:
                invariant_failures.append(f"seed{seed}_{condition}:kp_tracking")
            if condition != "clean":
                ratios = record["traffic_calibration"][
                    "realized_traffic_instruction_rms_ratio"]
                finite_values.extend(float(value) for value in ratios)
                if max(abs(float(value) - 4.0) for value in ratios) > 1e-5:
                    invariant_failures.append(f"seed{seed}_{condition}:ratio")
                diagnostics = record["diagnostics"]
                predictor_ratio = float(diagnostics[
                    "predictor_traffic_residual_rms_ratio"])
                norm_error = float(diagnostics["max_norm_match_relative_error"])
                finite_values.extend([predictor_ratio, norm_error])
                for row in trajectory:
                    mixed = row.get("mixed_apical")
                    if mixed is None:
                        invariant_failures.append(
                            f"seed{seed}_{condition}:mixed_trajectory")
                        break
                    finite_values.extend(float(mixed[key]) for key in (
                        "teaching_rms", "instruction_rms", "raw_apical_rms",
                        "innovation_rms", "traffic_rms"))
                if predictor_ratio > 0.05:
                    invariant_failures.append(f"seed{seed}_{condition}:predictor")
                if condition == "matched" and norm_error > 1e-6:
                    invariant_failures.append(f"seed{seed}_{condition}:norm")
                if int(record["work"][
                        "elementwise_operations_estimate"]) <= 0:
                    invariant_failures.append(f"seed{seed}_{condition}:elementwise")
            if condition == "innovation":
                diagnostics = record["diagnostics"]
                used_early = float(diagnostics["early_third_mean"])
                raw_early = mean_early(
                    diagnostics["raw_negative_gradient_cosine"])
                finite_values.extend([used_early, raw_early])
                if used_early < 0.80 or used_early - raw_early < 0.15:
                    invariant_failures.append(f"seed{seed}:alignment")

    all_finite = all(math.isfinite(value) for value in finite_values)
    checks = {
        "all_records_trajectories_and_metrics_finite": all_finite,
        "mean_innovation_accuracy_at_least_0.88": (
            statistics.mean(accuracies["innovation"]) >= 0.88),
        "mean_paired_deficit_to_clean_at_most_0.03": (
            statistics.mean(clean_deficits) <= 0.03),
        "mean_gain_over_raw_at_least_0.05": (
            statistics.mean(raw_gains) >= 0.05),
        "mean_gain_over_matched_at_least_0.03": (
            statistics.mean(matched_gains) >= 0.03),
        "raw_gain_one_sided_95pct_lower_bound_above_zero": (
            lower_confidence_bound(raw_gains) > 0.0),
        "matched_gain_one_sided_95pct_lower_bound_above_zero": (
            lower_confidence_bound(matched_gains) > 0.0),
        "at_least_four_positive_raw_gains": (
            sum(value > 0 for value in raw_gains) >= 4),
        "at_least_four_positive_matched_gains": (
            sum(value > 0 for value in matched_gains) >= 4),
        "all_mechanism_query_tracking_and_cost_invariants": (
            not invariant_failures),
    }
    status = "passed" if all(checks.values()) else "failed"
    output = {
        "protocol": "kp_mixed_traffic_confirmation_v1", "status": status,
        "checks": checks,
        "metrics": {
            "accuracy_by_seed": accuracies,
            "mean_accuracy": {condition: statistics.mean(values)
                              for condition, values in accuracies.items()},
            "paired_clean_deficit": clean_deficits,
            "paired_gain_over_raw": raw_gains,
            "paired_gain_over_matched": matched_gains,
            "mean_clean_deficit": statistics.mean(clean_deficits),
            "mean_gain_over_raw": statistics.mean(raw_gains),
            "mean_gain_over_matched": statistics.mean(matched_gains),
            "raw_gain_one_sided_95pct_lower_bound": (
                lower_confidence_bound(raw_gains)),
            "matched_gain_one_sided_95pct_lower_bound": (
                lower_confidence_bound(matched_gains)),
            "bp_50k_mac_reference": bp_macs_50k,
            "invariant_failures": invariant_failures,
            "source_commit": next(iter(source_commits)),
        },
        "test_evaluations": 20,
        "review_score_before": 5,
        "review_score_after": 6 if status == "passed" else 5,
        "score_change_rule": (
            "only a complete untouched five-seed innovation confirmation "
            "establishes the weak-accept bar"),
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as handle:
        json.dump(output, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()