code/normalized_dimension_mechanism_audit.py

#!/usr/bin/env python3
"""CPU scope expansion for adaptive compressed-PCA claims.

The original mechanism panel held the iteration budget fixed while changing d.
This producer keeps t/d^2 fixed, adds a dimension sweep, and also executes
fresh warmup and moving-eigenvector protection regimes for the held claims.
"""

from __future__ import annotations

import csv
import json
import math
import re
from pathlib import Path
import sys

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from reproduce import (  # noqa: E402
    adaptive_step,
    constants,
    error,
    fully_observed_step,
    gaussian_samples,
    initial_vectors,
    nonadaptive_step,
    rotate_truth,
)


def fit_power(rows: list[dict[str, float]], key: str) -> dict[str, float]:
    x = np.log(np.asarray([row["d"] for row in rows], dtype=float))
    y = np.log(np.asarray([row[key] for row in rows], dtype=float))
    slope, intercept = np.polyfit(x, y, 1)
    residual = y - (intercept + slope * x)
    r2 = 1.0 - float(np.sum(residual * residual) / np.sum((y - y.mean()) ** 2))
    return {"slope": float(slope), "r2": r2}


def normalized_mechanism() -> list[dict[str, float]]:
    rows: list[dict[str, float]] = []
    for d in (4, 8, 12, 16, 24, 32):
        trials = 16
        steps = int(round(700.0 * d * d))
        eta = 0.01 / d
        rng = np.random.default_rng(20260731 + d)
        initial = initial_vectors(rng, trials, d)
        adaptive, nonadaptive, full = initial.copy(), initial.copy(), initial.copy()
        for _ in range(steps):
            sample = gaussian_samples(rng, trials, d)
            adaptive = adaptive_step(adaptive, sample, eta, rng)
            nonadaptive = nonadaptive_step(nonadaptive, sample, eta, rng)
            full = fully_observed_step(full, sample, eta)
        adaptive_error = float(np.median(error(adaptive)))
        nonadaptive_error = float(np.median(error(nonadaptive)))
        full_error = float(np.median(error(full)))
        rows.append({
            "d": float(d), "steps": float(steps), "trials": float(trials),
            "normalized_t_over_d2": 700.0, "eta": eta,
            "adaptive_error": adaptive_error,
            "nonadaptive_error": nonadaptive_error,
            "fully_observed_error": full_error,
            "nonadaptive_over_adaptive": nonadaptive_error / adaptive_error,
            "adaptive_over_fully_observed": adaptive_error / full_error,
        })
        print(f"mechanism d={d} steps={steps}", flush=True)
    return rows


def warmup_protection() -> list[dict[str, float]]:
    rows: list[dict[str, float]] = []
    for d in (8, 16, 24, 32):
        trials = 24
        setup = constants(d)
        t0 = int(math.ceil(setup["t0"]))
        rng = np.random.default_rng(20260840 + d)
        u = initial_vectors(rng, trials, d)
        for step in range(1, t0 + 1):
            eta = setup["eta0"] if step <= t0 else 2.0 * (d - 1.0) / (setup["gap"] * (4.0 * setup["S"] + step - t0))
            u = adaptive_step(u, gaussian_samples(rng, trials, d), eta, rng)
        values = error(u)
        rows.append({
            "d": float(d), "t0": float(t0), "trials": float(trials),
            "mean_error_at_t0": float(values.mean()),
            "max_error_at_t0": float(values.max()),
        })
    return rows


def tracking_protection() -> list[dict[str, float]]:
    rows: list[dict[str, float]] = []
    for d, velocity in ((8, 2e-4), (8, 8e-4), (12, 2e-4), (12, 8e-4)):
        trials, steps = 24, 20_000
        setup = constants(d)
        eta_hat = math.sqrt(velocity / setup["S"])
        rng = np.random.default_rng(20260880 + d + int(velocity * 1e7))
        truth = np.zeros((trials, d)); truth[:, 0] = 1.0
        u = np.zeros((trials, d)); u[:, 0] = math.sqrt(0.1); u[:, 1] = math.sqrt(0.9)
        tail: list[float] = []
        for step in range(steps):
            truth = rotate_truth(truth, velocity, rng)
            u = adaptive_step(u, gaussian_samples(rng, trials, d, truth), eta_hat, rng)
            if step >= steps - 2_000:
                tail.append(float(error(u, truth).mean()))
        rows.append({
            "d": float(d), "velocity": velocity, "steps": float(steps),
            "trials": float(trials), "eta_hat": eta_hat,
            "x_star": velocity + math.sqrt(velocity * setup["S"]),
            "tail_mean_error": float(np.mean(tail)),
        })
    return rows


def tracking_formula_protection() -> list[dict[str, float]]:
    rows: list[dict[str, float]] = []
    for d in (8, 12, 16):
        setup = constants(d)
        for velocity in (1e-5, 1e-4, 1e-3):
            eta_hat = math.sqrt(velocity / setup["S"])
            rows.append({
                "d": float(d), "velocity": velocity, "eta_hat": eta_hat,
                "x_star": velocity + math.sqrt(velocity * setup["S"]),
                "first_derivative": 0.5 * setup["S"] - 0.5 * velocity / (eta_hat * eta_hat),
                "curvature": velocity / (eta_hat ** 3),
            })
    return rows


def source_figure_protection() -> dict[str, object]:
    text = Path("source/sections/experiments.tex").read_text()
    figure1 = re.search(r"25--75.*?50 trials.*?\$d=64", text, re.S)
    figure3 = re.search(r"20 trials.*?20--80.*?\$d=10", text, re.S)
    return {
        "figure1_metadata_mismatch": bool(figure1),
        "figure3_metadata_match": bool(figure3),
        "mismatch_fields": 3 if figure1 else 0,
        "source_sha256": __import__("hashlib").sha256(text.encode()).hexdigest(),
    }


def main() -> None:
    mechanism = normalized_mechanism()
    warmup = warmup_protection()
    tracking = tracking_protection()
    tracking_formula = tracking_formula_protection()
    source = source_figure_protection()
    result = {
        "schema": "normalized-dimension-mechanism-audit-v2",
        "mechanism": mechanism,
        "ratio_fit": fit_power(mechanism, "nonadaptive_over_adaptive"),
        "adaptive_error_fit": fit_power(mechanism, "adaptive_error"),
        "ratio_growth_d4_to_d32": mechanism[-1]["nonadaptive_over_adaptive"] / mechanism[0]["nonadaptive_over_adaptive"],
        "all_nonadaptive_worse": all(row["nonadaptive_over_adaptive"] > 1.0 for row in mechanism),
        "warmup": warmup,
        "warmup_all_means_below_half": all(row["mean_error_at_t0"] < 0.5 for row in warmup),
        "tracking": tracking,
        "tracking_formula": tracking_formula,
        "tracking_formula_max_abs_derivative": max(abs(row["first_derivative"]) for row in tracking_formula),
        "tracking_formula_all_curvatures_positive": all(row["curvature"] > 0 for row in tracking_formula),
        "source_figure": source,
    }
    Path("outputs/normalized_dimension_mechanism.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    with Path("outputs/normalized_dimension_mechanism.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=sorted(mechanism[0]))
        writer.writeheader(); writer.writerows(mechanism)
    print(json.dumps({k: v for k, v in result.items() if k != "mechanism"}, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()