#!/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()