logbook-files/exact_d3_certification.py

#!/usr/bin/env python3
"""Literal d=3 finite-cover certification for arXiv:2602.22130.

Every estimate reported as ``exact`` minimizes the printed Algorithm-1
tournament over every point in explicit Cartesian epsilon' and eta covers.
The covers are cached, but no candidate or active frequency is subsampled.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import platform
import resource
import time
from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from numpy.linalg import norm

from reproduce_full import (
    algorithm1_continuous,
    cartesian_ball,
    ecf_ratio,
    gaussian_cf,
    gaussian_witness_delta,
    minimax_scores,
    sample_contaminated,
)

SEED = 26022131


@dataclass
class ExactDesign:
    scale: float
    candidates: np.ndarray
    omega_cover: np.ndarray
    omega: np.ndarray
    alpha: float
    epsilon: float
    radius: float
    A: float
    delta: float
    B_delta: float
    epsilon_prime: float
    eta: float
    candidate_step: float
    frequency_step: float
    theorem_n_C1: int


def build_design(alpha: float, epsilon: float, radius: float, scale: float, d: int = 3) -> ExactDesign:
    """Build every cover point used by Lines 1 and 3--9 of Algorithm 1."""
    A, _, delta = gaussian_witness_delta(alpha, epsilon, scale)
    M1 = math.sqrt(2.0 / math.pi) / scale
    L = 2.0 * math.pi * scale * math.exp(-0.5)
    B_delta = math.sqrt(d) * M1 / (2.0 * math.pi * delta)
    epsilon_prime = min(alpha / (2.0 * (1.0 - alpha) * math.pi * B_delta), epsilon)
    eta = min(delta / (2.0 * L), A / (2.0 * math.pi * radius))
    candidates = cartesian_ball(radius, epsilon_prime, d)
    omega_cover = cartesian_ball(B_delta, eta, d)
    omega = omega_cover[np.abs(gaussian_cf(omega_cover, scale)) >= delta / 2.0]
    margin = ((1.0 - alpha) * A - 2.0 * alpha) * delta
    log_term = math.log(math.sqrt(d) * M1 * radius * L / (delta**2 * A))
    theorem_n_C1 = math.ceil(d * log_term / margin**2)
    return ExactDesign(
        scale, candidates, omega_cover, omega, alpha, epsilon, radius, A, delta,
        B_delta, epsilon_prime, eta, 2.0 * epsilon_prime / math.sqrt(d),
        2.0 * eta / math.sqrt(d), theorem_n_C1,
    )


def exact_estimate(x: np.ndarray, design: ExactDesign) -> tuple[np.ndarray, float]:
    """Exhaustively score all mean-cover points against all active frequencies."""
    psi = ecf_ratio(x, design.omega, design.scale)
    scores = minimax_scores(design.candidates, design.omega, psi, design.alpha)
    winner = int(np.argmin(scores))
    return design.candidates[winner], float(scores[winner])


def wilson(successes: int, total: int, z: float = 1.959963984540054) -> tuple[float, float]:
    p = successes / total
    den = 1.0 + z * z / total
    center = (p + z * z / (2.0 * total)) / den
    half = z * math.sqrt(p * (1.0 - p) / total + z * z / (4.0 * total * total)) / den
    return center - half, center + half


def one_trial(seed: int, n: int, design: ExactDesign, compare_continuous: bool) -> dict:
    rng = np.random.default_rng(seed)
    d = design.candidates.shape[1]
    direction = rng.normal(size=d)
    direction /= norm(direction)
    mu = 0.35 * direction
    outlier = mu + 15.0 * direction
    x = sample_contaminated(rng, n, mu, design.alpha, design.scale, outlier)
    started = time.perf_counter()
    estimate, score = exact_estimate(x, design)
    exact_seconds = time.perf_counter() - started
    exact_error = float(norm(estimate - mu))
    row = {
        "seed": seed, "d": d, "scale": design.scale, "delta": design.delta,
        "n": n, "q_n_delta2_over_d": n * design.delta**2 / d,
        "exact_error": exact_error, "exact_success": int(exact_error <= design.epsilon),
        "exact_winning_score": score, "exact_seconds": exact_seconds,
        "naive_mean_error": float(norm(x.mean(axis=0) - mu)),
        "coordinate_median_error": float(norm(np.median(x, axis=0) - mu)),
    }
    if compare_continuous:
        continuous, meta = algorithm1_continuous(
            x, design.alpha, design.epsilon, design.radius, design.scale,
            np.random.default_rng(seed + 900_000_000),
        )
        continuous_error = float(norm(continuous - mu))
        row.update({
            "continuous_error": continuous_error,
            "continuous_success": int(continuous_error <= design.epsilon),
            "exact_continuous_distance": float(norm(estimate - continuous)),
            "success_decisions_agree": int(
                (exact_error <= design.epsilon) == (continuous_error <= design.epsilon)
            ),
            "continuous_optimizer_success": int(meta["optimizer_success"]),
        })
    return row


def cover_record(design: ExactDesign) -> dict:
    d = design.candidates.shape[1]
    candidate_radius = design.candidate_step * math.sqrt(d) / 2.0
    frequency_radius = design.frequency_step * math.sqrt(d) / 2.0
    return {
        "scale": design.scale, "delta": design.delta,
        "candidate_count": len(design.candidates),
        "frequency_cover_count": len(design.omega_cover),
        "active_frequency_count": len(design.omega),
        "candidate_frequency_pairs": len(design.candidates) * len(design.omega),
        "epsilon_prime": design.epsilon_prime,
        "candidate_lattice_certified_radius": candidate_radius,
        "candidate_cover_radius_holds": candidate_radius <= design.epsilon_prime * (1 + 1e-12),
        "eta": design.eta,
        "frequency_lattice_certified_radius": frequency_radius,
        "frequency_cover_radius_holds": frequency_radius <= design.eta * (1 + 1e-12),
        "B_delta": design.B_delta,
        "theorem_n_hidden_C_equal_1": design.theorem_n_C1,
    }


def aggregate(rows: list[dict], keys: tuple[str, ...]) -> list[dict]:
    groups: dict[tuple, list[dict]] = {}
    for row in rows:
        groups.setdefault(tuple(row[k] for k in keys), []).append(row)
    output = []
    for values, group in sorted(groups.items()):
        successes = sum(int(x["exact_success"]) for x in group)
        lo, hi = wilson(successes, len(group))
        item = {key: value for key, value in zip(keys, values)}
        item.update({
            "trials": len(group), "exact_successes": successes,
            "exact_success_rate": successes / len(group),
            "wilson95_low": lo, "wilson95_high": hi,
            "median_exact_error": float(np.median([x["exact_error"] for x in group])),
            "median_naive_error": float(np.median([x["naive_mean_error"] for x in group])),
        })
        if "continuous_success" in group[0]:
            item.update({
                "continuous_success_rate": float(np.mean([x["continuous_success"] for x in group])),
                "success_decision_agreement": float(np.mean([x["success_decisions_agree"] for x in group])),
                "median_exact_continuous_distance": float(np.median([x["exact_continuous_distance"] for x in group])),
            })
        output.append(item)
    return output


def make_plot(q_summary: list[dict], delta_summary: list[dict], formula_summary: list[dict], path: Path) -> None:
    fig, axes = plt.subplots(1, 3, figsize=(14, 4.2))
    axes[0].plot([x["target_q"] for x in q_summary], [x["exact_success_rate"] for x in q_summary], "o-", label="literal cover")
    axes[0].plot([x["target_q"] for x in q_summary], [x["continuous_success_rate"] for x in q_summary], "s--", label="continuous diagnostic")
    axes[0].axhline(2 / 3, color="k", ls=":")
    axes[0].set(xscale="log", ylim=(-0.03, 1.03), xlabel=r"target $n\delta^2/d$", ylabel="success rate", title="Exact d=3 tournament")
    axes[0].legend()
    for scale in sorted({x["scale"] for x in delta_summary}):
        part = [x for x in delta_summary if x["scale"] == scale]
        axes[1].plot([x["target_q"] for x in part], [x["exact_success_rate"] for x in part], "o-", label=f"scale={scale}")
    axes[1].axhline(2 / 3, color="k", ls=":")
    axes[1].set(xscale="log", ylim=(-0.03, 1.03), xlabel=r"target $n\delta^2/d$", title="Literal-cover delta intervention")
    axes[1].legend(fontsize=8)
    axes[2].bar([str(x["scale"]) for x in formula_summary], [x["exact_success_rate"] for x in formula_summary])
    axes[2].axhline(2 / 3, color="k", ls=":")
    axes[2].set(ylim=(0, 1.03), xlabel="Gaussian scale", ylabel="success rate", title="Theorem formula, hidden C=1")
    fig.tight_layout(); fig.savefig(path, dpi=180); plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("results/exact_d3"))
    parser.add_argument("--primary-reps", type=int, default=16)
    parser.add_argument("--delta-reps", type=int, default=10)
    parser.add_argument("--formula-reps", type=int, default=12)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    started = time.perf_counter()
    designs = {scale: build_design(0.05, 0.50, 0.80, scale) for scale in (0.75, 1.00, 1.25, 1.50)}
    rows: list[dict] = []

    primary = designs[1.00]
    for qi, q in enumerate((0.25, 0.50, 1.0, 2.0, 4.0, 8.0, 16.0)):
        n = max(1, math.ceil(q * 3 / primary.delta**2))
        for rep in range(args.primary_reps):
            row = one_trial(SEED + qi * 10_000 + rep, n, primary, True)
            row.update({"experiment": "exact_q_sweep", "target_q": q, "rep": rep})
            rows.append(row)

    for si, (scale, design) in enumerate(designs.items()):
        for qi, q in enumerate((1.0, 2.0, 4.0, 8.0)):
            n = max(1, math.ceil(q * 3 / design.delta**2))
            for rep in range(args.delta_reps):
                row = one_trial(SEED + 1_000_000 + si * 100_000 + qi * 10_000 + rep, n, design, False)
                row.update({"experiment": "exact_delta_sweep", "target_q": q, "rep": rep})
                rows.append(row)

    for si, (scale, design) in enumerate(designs.items()):
        for rep in range(args.formula_reps):
            row = one_trial(SEED + 2_000_000 + si * 100_000 + rep, design.theorem_n_C1, design, False)
            row.update({"experiment": "theorem_formula_C1", "target_q": design.theorem_n_C1 * design.delta**2 / 3, "rep": rep})
            rows.append(row)

    q_summary = aggregate([r for r in rows if r["experiment"] == "exact_q_sweep"], ("target_q", "n"))
    delta_summary = aggregate([r for r in rows if r["experiment"] == "exact_delta_sweep"], ("scale", "delta", "target_q", "n"))
    formula_summary = aggregate([r for r in rows if r["experiment"] == "theorem_formula_C1"], ("scale", "delta", "n"))
    elapsed = time.perf_counter() - started
    summary = {
        "paper": "Sample Complexity Bounds for Robust Mean Estimation with Mean-Shift Contamination",
        "openreview": "https://openreview.net/forum?id=no9dQDBxsu", "seed": SEED,
        "scope": "Every exact result exhaustively scores all d=3 mean-cover candidates against all active points of the d=3 frequency cover; no cover subsampling.",
        "cover_certificates": [cover_record(x) for x in designs.values()],
        "exact_q_sweep": q_summary, "exact_delta_sweep": delta_summary,
        "theorem_formula_hidden_C1": formula_summary,
        "continuous_link": {
            "purpose": "Only a paired diagnostic at d=3; it does not prove equivalence in d>=4.",
            "overall_success_decision_agreement": float(np.mean([r["success_decisions_agree"] for r in rows if r["experiment"] == "exact_q_sweep"])),
            "median_estimate_distance": float(np.median([r["exact_continuous_distance"] for r in rows if r["experiment"] == "exact_q_sweep"])),
        },
        "execution": {"elapsed_seconds": elapsed, "python": platform.python_version(), "max_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, "total_exact_trials": len(rows)},
        "limitations": [
            "The literal exponential tournament is certified through d=3, not d>=4.",
            "The hidden universal constant in Theorem 3.2 is unknown; C=1 is an explicit diagnostic, not an estimate of the theorem constant.",
            "The d>=4 dense-direction/L-BFGS study remains a non-equivalent diagnostic and is never labeled literal Algorithm 1.",
        ],
    }
    fieldnames = sorted({key for row in rows for key in row})
    with (args.output / "exact_d3_trials.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames); writer.writeheader(); writer.writerows(rows)
    (args.output / "exact_d3_summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    make_plot(q_summary, delta_summary, formula_summary, args.output / "exact_d3_results.png")
    manifest_lines = []
    for path in (Path(__file__), args.output / "exact_d3_trials.csv", args.output / "exact_d3_summary.json", args.output / "exact_d3_results.png"):
        manifest_lines.append(f"{hashlib.sha256(path.read_bytes()).hexdigest()}  {path}")
    (args.output / "SHA256SUMS").write_text("\n".join(manifest_lines) + "\n")
    print("LITERAL_D3_COMPLETE")
    print(json.dumps({"elapsed_seconds": elapsed, "total_exact_trials": len(rows), "cover_certificates": summary["cover_certificates"], "exact_q_sweep": q_summary, "exact_delta_sweep": delta_summary, "theorem_formula_hidden_C1": formula_summary, "continuous_link": summary["continuous_link"]}, indent=2))


if __name__ == "__main__":
    main()