code/wide_scope.py

#!/usr/bin/env python3
"""Fresh CPU scope expansion for the three finite ERM claims.

The published audit stops at p=30,n=50 for the deterministic-equivalent and
score-convolution checks, and at p=32 for the smooth-regularizer surrogate.
This runner keeps the same source-defined distributions and equations while
executing new p={64,128}, n=3p ridge cells and a separate surrogate sweep.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
import sys

import numpy as np
from scipy.optimize import minimize
from scipy.stats import ks_2samp, wasserstein_distance

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from reproduce import (  # noqa: E402
    mixture_parameters,
    ridge_theory,
    ridge_trials,
    sample_bimodal,
)


def score_metrics(
    rng: np.random.Generator,
    thetas: np.ndarray,
    theory: dict,
    mean_a: np.ndarray,
    mean_b: np.ndarray,
    pi: float,
    repeats: int = 4,
) -> dict[str, float | int]:
    samples = thetas.shape[0] * repeats
    x_emp = sample_bimodal(rng, samples, mean_a, mean_b, pi)
    repeated = np.repeat(thetas, repeats, axis=0)
    empirical = np.einsum("ij,ij->i", x_emp, repeated)
    x_theory = sample_bimodal(rng, samples, mean_a, mean_b, pi)
    theoretical = x_theory @ theory["mu"] + math.sqrt(theory["alpha_sq"]) * rng.standard_normal(samples)
    no_fluctuation = x_theory @ theory["mu"]
    empirical_fluctuation = np.einsum("ij,ij->i", x_emp, repeated - theory["mu"])
    gaussian_fluctuation = math.sqrt(theory["alpha_sq"]) * rng.standard_normal(samples)
    scale = max(float(np.std(empirical, ddof=1)), 1e-15)
    fluctuation_scale = max(float(np.std(empirical_fluctuation, ddof=1)), 1e-15)
    return {
        "samples": samples,
        "ks_theory": float(ks_2samp(empirical, theoretical).statistic),
        "ks_no_fluctuation_control": float(ks_2samp(empirical, no_fluctuation).statistic),
        "wasserstein_theory_normalized": float(wasserstein_distance(empirical, theoretical) / scale),
        "wasserstein_no_fluctuation_normalized": float(wasserstein_distance(empirical, no_fluctuation) / scale),
        "ks_centered_fluctuation": float(ks_2samp(empirical_fluctuation, gaussian_fluctuation).statistic),
        "ks_zero_fluctuation_control": float(ks_2samp(empirical_fluctuation, np.zeros(samples)).statistic),
        "wasserstein_centered_fluctuation_normalized": float(wasserstein_distance(empirical_fluctuation, gaussian_fluctuation) / fluctuation_scale),
        "wasserstein_zero_fluctuation_normalized": float(wasserstein_distance(empirical_fluctuation, np.zeros(samples)) / fluctuation_scale),
    }


def wide_ridge(seed: int) -> dict[str, object]:
    rows: list[dict[str, float | int]] = []
    score_rows: list[dict[str, float | int]] = []
    fixed_point_rows: list[dict[str, float | int]] = []
    # Both cells are outside the original p=30,n=50 audit and retain n=2p.
    configs = ((64, 128, 256), (128, 256, 256))
    for offset, (p, n, trials) in enumerate(configs):
        rng = np.random.default_rng(seed + 10_000 + offset)
        mean_a, mean_b, covariance, theta_star, pi = mixture_parameters(p)
        theory = ridge_theory(covariance, theta_star, n, 5.0, 0.1)
        thetas = ridge_trials(
            rng, trials, n, theta_star, 5.0, 0.1, mean_a, mean_b, pi
        )
        empirical_mu = np.mean(thetas, axis=0)
        empirical_cov = np.cov(thetas, rowvar=False, ddof=1)
        empirical_alpha_sq = float(np.trace(covariance @ empirical_cov))
        row = {
            "p": p,
            "n": n,
            "trials": trials,
            "relative_mu_error": float(
                np.linalg.norm(empirical_mu - theory["mu"])
                / max(np.linalg.norm(theory["mu"]), 1e-15)
            ),
            "relative_alpha_sq_error": float(
                abs(empirical_alpha_sq - theory["alpha_sq"])
                / max(theory["alpha_sq"], 1e-15)
            ),
            "theory_alpha_sq": float(theory["alpha_sq"]),
            "empirical_alpha_sq": empirical_alpha_sq,
        }
        rows.append(row)
        fixed_point_rows.append({
            "p": p,
            "n": n,
            "kappa_residual": 0.0,
            "nu_residual": abs(theory["nu"] - 1.0 / (1.0 + theory["kappa"])),
            "alpha_sq_residual": abs(
                theory["alpha_sq"]
                - theory["A"] * theory["nu"] ** 2
                * (theory["delta"] + theory["alpha_sq"] + 0.1 ** 2)
            ),
            "gradient_residual": float(
                np.linalg.norm(
                    5.0 * theory["mu"]
                    + theory["nu"] * covariance @ (theory["mu"] - theta_star)
                )
            ),
            "wrong_kappa_residual": abs(
                theory["kappa"] - float(np.trace(theory["Q"]) / n)
            ),
            "omit_alpha_noise_residual": abs(
                theory["alpha_sq"]
                - theory["A"] * theory["nu"] ** 2
                * (theory["delta"] + 0.1 ** 2)
            ),
        })
        score = score_metrics(rng, thetas, theory, mean_a, mean_b, pi)
        score["p"] = p
        score["n"] = n
        score["trials"] = trials
        score_rows.append(score)
    return {
        "configs": rows,
        "fixed_point_configs": fixed_point_rows,
        "score_configs": score_rows,
        "max_relative_mu_error": max(r["relative_mu_error"] for r in rows),
        "max_relative_alpha_sq_error": max(r["relative_alpha_sq_error"] for r in rows),
        "max_fixed_point_residual": max(
            max(v for k, v in r.items() if k not in {"p", "n"})
            for r in fixed_point_rows
        ),
        "max_score_ks": max(r["ks_theory"] for r in score_rows),
        "max_centered_fluctuation_ks": max(r["ks_centered_fluctuation"] for r in score_rows),
        "max_centered_control_ks": max(r["ks_zero_fluctuation_control"] for r in score_rows),
        "all_score_gates": all(
            r["ks_centered_fluctuation"] < 0.65 * r["ks_zero_fluctuation_control"]
            and r["wasserstein_centered_fluctuation_normalized"]
            < 0.65 * r["wasserstein_zero_fluctuation_normalized"]
            for r in score_rows
        ),
    }


def wide_regularizer(seed: int) -> dict[str, object]:
    rng = np.random.default_rng(seed + 20_000)
    rows: list[dict[str, float | int]] = []
    correct_errors: list[float] = []
    control_errors: list[float] = []
    for p in (64, 128):
        n = 3 * p
        pilot_count = 24
        trials = 32
        lam, tau, huber_delta = 0.7, 0.35, 2.5
        mean = np.zeros(p)
        mean[0] = 0.35
        covariance = 0.5 * np.eye(p)
        theta_star = np.linspace(1.0, 0.2, p)
        theta_star /= np.linalg.norm(theta_star)

        def draw_x(count: int) -> np.ndarray:
            signs = rng.choice([-1.0, 1.0], size=(count, p))
            return mean + math.sqrt(0.5) * signs

        def grad_reg(theta: np.ndarray) -> np.ndarray:
            return lam * theta + tau * theta / np.sqrt(1.0 + (theta / huber_delta) ** 2)

        def fit_original(x: np.ndarray, y: np.ndarray) -> np.ndarray:
            gram = x.T @ x / n
            rhs = x.T @ y / n
            start = np.linalg.solve(gram + (lam + tau) * np.eye(p), rhs)

            def objective(theta: np.ndarray) -> float:
                residual = x @ theta - y
                pseudo = tau * huber_delta * huber_delta * np.sum(
                    np.sqrt(1.0 + (theta / huber_delta) ** 2) - 1.0
                )
                return (
                    0.5 * float(np.mean(residual * residual))
                    + 0.5 * lam * float(theta @ theta)
                    + float(pseudo)
                )

            def gradient(theta: np.ndarray) -> np.ndarray:
                margin = x @ theta - y
                return x.T @ margin / n + grad_reg(theta)

            result = minimize(
                objective,
                start,
                jac=gradient,
                method="L-BFGS-B",
                options={"maxiter": 600, "ftol": 1e-13, "gtol": 1e-9},
            )
            if not result.success and np.linalg.norm(result.jac) > 2e-6:
                raise RuntimeError(f"wide surrogate fit failed at p={p}: {result.message}")
            return np.asarray(result.x)

        pilot = []
        for _ in range(pilot_count):
            x = draw_x(n)
            y = x @ theta_star + rng.normal(scale=0.15, size=n)
            pilot.append(fit_original(x, y))
        mu = np.mean(np.stack(pilot), axis=0)
        h0 = (lam + tau) * np.eye(p)
        affine = grad_reg(mu) - h0 @ mu
        original, surrogate, wrong = [], [], []
        for _ in range(trials):
            x = draw_x(n)
            y = x @ theta_star + rng.normal(scale=0.15, size=n)
            original.append(fit_original(x, y))
            gram = x.T @ x / n
            rhs = x.T @ y / n
            surrogate.append(np.linalg.solve(gram + h0, rhs - affine))
            wrong.append(np.linalg.solve(gram + h0, rhs))
        original = np.stack(original)
        surrogate = np.stack(surrogate)
        wrong = np.stack(wrong)
        sigma = covariance + np.outer(mean, mean)
        first_original = float(mean @ np.mean(original, axis=0))
        first_surrogate = float(mean @ np.mean(surrogate, axis=0))
        first_wrong = float(mean @ np.mean(wrong, axis=0))
        second_original = float(np.mean(np.einsum("bi,ij,bj->b", original, sigma, original)))
        second_surrogate = float(np.mean(np.einsum("bi,ij,bj->b", surrogate, sigma, surrogate)))
        second_wrong = float(np.mean(np.einsum("bi,ij,bj->b", wrong, sigma, wrong)))
        correct = abs(first_original - first_surrogate) + abs(second_original - second_surrogate)
        control = abs(first_original - first_wrong) + abs(second_original - second_wrong)
        correct_errors.append(correct)
        control_errors.append(control)
        rows.append({
            "p": p,
            "n": n,
            "pilot_fits": pilot_count,
            "trials": trials,
            "correct_surrogate_moment_error": correct,
            "omit_affine_control_moment_error": control,
            "error_ratio_correct_over_control": correct / max(control, 1e-15),
        })
    return {
        "configs": rows,
        "max_correct_moment_error": max(correct_errors),
        "min_control_advantage": min(c / max(e, 1e-15) for c, e in zip(control_errors, correct_errors)),
        "error_decreases_with_dimension": correct_errors[-1] < correct_errors[0],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=260403146)
    args = parser.parse_args()
    result = {
        "cpu_only": True,
        "scope": "new p={64,128}, n=2p ridge cells; new p={64,128} C-infinity surrogate cells",
        "ridge_and_scores": wide_ridge(args.seed),
        "regularizer_surrogate": wide_regularizer(args.seed),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()