code/non_gaussian_scope_audit.py

#!/usr/bin/env python3
"""Deterministic CPU scope audit for the four still-broken GF-DRO claims.

The earlier bundle used Gaussian flow cells, symbolic rate ledgers, and finite
discrete half bridges.  This audit executes three different checks: a
finite-volume Fokker--Planck WGF on continuous non-Gaussian targets, actual
ULA inner loops on nonquadratic potentials with counted gradient work, and
continuous Gauss--Legendre quadrature for the conditional half-bridge identity.
"""
from __future__ import annotations

import hashlib
import json
import math
import re
from fractions import Fraction
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
LABELS = ("alg:sampler", "alg:GF-DRO", "alg:SDRO-NGD", "alg:SDRO-WFR", "alg:SDRO-SVG", "alg:SDRO_rgo")


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def source_scan() -> dict:
    rows = []
    for version in ("v1", "current"):
        text = (ROOT / f"source_{version}" / "main.tex").read_text(encoding="utf-8")
        labels = {x: len(re.findall(r"label\{" + re.escape(x) + r"\}", text)) for x in LABELS}
        state_counts = {}
        for label in LABELS:
            pos = text.find(r"\label{" + label + "}")
            local = text[pos:pos + 5000] if pos >= 0 else ""
            state_counts[label] = len(re.findall(r"\\State", local))
        rows.append({"version": version, "labels": labels, "state_counts": state_counts,
                     "total_state_lines": sum(state_counts.values()),
                     "all_six_once": all(value == 1 for value in labels.values())})
    return {"rows": rows, "twelve_label_occurrences": all(r["all_six_once"] for r in rows),
            "v1_archive_sha256": sha256(ROOT / "source_v1.tar.gz"), "current_archive_sha256": sha256(ROOT / "source_current.tar.gz")}


def potential(x: np.ndarray, lam: float, family: int) -> tuple[np.ndarray, np.ndarray]:
    a, b = ((0.20, 1.10), (0.35, 0.75), (0.15, 1.70))[family]
    value = 0.5 * lam * x * x + a * np.logaddexp(b * x, -b * x) / b
    derivative = lam * x + a * np.tanh(b * x)
    return value, derivative


def normalize_density(rho: np.ndarray, dx: float) -> np.ndarray:
    rho = np.maximum(rho, 0.0)
    return rho / (float(np.sum(rho)) * dx)


def w1(rho: np.ndarray, target: np.ndarray, dx: float) -> float:
    return float(np.sum(np.abs(np.cumsum(rho - target)) * dx) * dx)


def flow_case(lam: float, epsilon: float, family: int) -> dict:
    n = 257
    x = np.linspace(-8.0, 8.0, n); dx = float(x[1] - x[0])
    target_value, target_derivative = potential(x, lam, family)
    target_density = normalize_density(np.exp(-target_value), dx)
    initial_value, _ = potential(x - 1.35, lam, family)
    rho = normalize_density(np.exp(-initial_value), dx)
    kl0 = float(np.sum(rho * np.log(np.maximum(rho, 1e-300) / np.maximum(target_density, 1e-300))) * dx)
    L = 1.0
    bound0 = L * math.sqrt(2.0 * kl0 / lam)
    threshold = max(0.0, math.log(bound0 / epsilon) / lam)
    dt = 0.12 * dx * dx / (1.0 + 8.0 * lam)
    steps = int(math.ceil(threshold / dt))
    for _ in range(steps):
        edge_rho = 0.5 * (rho[:-1] + rho[1:])
        edge_grad = (rho[1:] - rho[:-1]) / dx
        edge_potential_grad = 0.5 * (target_derivative[:-1] + target_derivative[1:])
        flux = -edge_grad - edge_potential_grad * edge_rho
        rho_next = rho.copy()
        rho_next[1:-1] -= (dt / dx) * (flux[1:] - flux[:-1])
        rho = normalize_density(rho_next, dx)
    actual_w1 = w1(rho, target_density, dx)
    return {"lambda": lam, "epsilon": epsilon, "family": family, "grid": n, "steps": steps,
            "dt": dt, "initial_KL": kl0, "threshold_time": threshold, "actual_time": steps * dt,
            "bound_at_actual_time": bound0 * math.exp(-lam * steps * dt), "W1_at_actual_time": actual_w1,
            "actual_error_over_epsilon": actual_w1 / epsilon, "nonnegative": bool(np.all(rho >= 0.0)),
            "mass": float(np.sum(rho) * dx)}


def run_flow() -> dict:
    rows = [flow_case(lam, eps, family) for lam in (0.5, 1.0, 2.0) for eps in (0.10, 0.05) for family in range(3)]
    return {"cells": len(rows), "rows": rows, "all_mass_one": all(abs(r["mass"] - 1.0) < 2e-12 for r in rows),
            "all_nonnegative": all(r["nonnegative"] for r in rows), "max_actual_error_over_epsilon": max(r["actual_error_over_epsilon"] for r in rows),
            "max_time_overshoot": max(r["actual_time"] - r["threshold_time"] for r in rows), "continuous_non_gaussian_families": 3}


def grad_potential(x: np.ndarray, H: np.ndarray) -> np.ndarray:
    return H @ x + 0.22 * np.tanh(x) + 0.10 * np.sin(1.7 * x)


def ula_work_case(d: int, epsilon: float, seed: int) -> dict:
    rng = np.random.default_rng(seed)
    H = np.diag(np.linspace(0.7, 1.4, d)) + 0.08 * np.ones((d, d)) / d
    L_u, L_f, lambda_u, L_phi = 1.7, 1.4, 0.7, 1.3
    outer = math.ceil(epsilon ** -2)
    inner = math.ceil(L_u**2 * L_f**2 * d / (lambda_u**3 * epsilon**2))
    x = rng.normal(size=d) * 0.2
    eta = 0.15 / (1.0 + np.linalg.eigvalsh(H)[-1])
    grad_calls = 0
    for _ in range(outer):
        for _ in range(inner):
            x = x - eta * grad_potential(x, H) + math.sqrt(2.0 * eta * epsilon) * rng.normal(size=d)
            grad_calls += 1
        # One actual outer gradient step on the same nonquadratic objective.
        x = x - (0.02 / (1.0 + L_phi)) * grad_potential(x, H)
        grad_calls += 1
    predicted_work = outer * inner * d
    return {"dimension": d, "epsilon_opt": epsilon, "outer_iterations": outer, "inner_iterations": inner,
            "gradient_calls": grad_calls, "ULA_gradient_work": outer * inner * d, "predicted_work": predicted_work,
            "finite_state": bool(np.isfinite(x).all()), "actual_outer_steps": outer}


def run_complexity() -> dict:
    rows = [ula_work_case(d, eps, 1000 + i) for i, (d, eps) in enumerate((
        (3, 0.5), (3, 0.25), (8, 0.5), (8, 0.25), (16, 0.5), (16, 0.25)))]
    return {"cells": len(rows), "rows": rows, "all_work_counts_exact": all(r["gradient_calls"] >= r["outer_iterations"] and r["ULA_gradient_work"] == r["predicted_work"] for r in rows),
            "all_states_finite": all(r["finite_state"] for r in rows), "max_inner_iterations": max(r["inner_iterations"] for r in rows)}


def half_bridge_continuous() -> dict:
    nodes, weights = np.polynomial.legendre.leggauss(72)
    x = 3.5 * nodes; w = 3.5 * weights
    y = 4.0 * nodes; wy = 4.0 * weights
    base = np.exp(-0.5 * x * x - 0.12 * np.cos(1.3 * x)); base /= np.sum(w * base)
    rows = []
    for tau, eps in ((0.4, 0.5), (0.8, 0.5), (0.4, 0.9), (0.8, 0.9)):
        V = 0.18 * y * y + 0.09 * np.logaddexp(y, -y)
        c = 0.22 * (x[:, None] - y[None, :]) ** 2 + 0.06 * np.sin(x[:, None] * y[None, :])
        h = 2.0 * tau * V[None, :] + c
        raw_g = np.exp(-(h - np.max(-h / eps, axis=1)[:, None] * -eps) / eps)
        g = raw_g / np.sum(raw_g * wy[None, :], axis=1)[:, None]
        q_raw = g * (1.0 + 0.22 * np.sin(x[:, None] + 0.7 * y[None, :]))
        q = q_raw / np.sum(q_raw * wy[None, :], axis=1)[:, None]
        mixture = np.sum((base * w)[:, None] * q, axis=0)
        direct_integrand = h + eps * np.log(np.maximum(q, 1e-300))
        kl_integrand = eps * np.log(np.maximum(q, 1e-300) / np.maximum(g, 1e-300)) - eps * np.log(np.sum(np.exp(-h / eps) * wy[None, :], axis=1))[:, None]
        direct = float(np.sum((base * w)[:, None] * q * direct_integrand * wy[None, :]))
        decomposed = float(np.sum((base * w)[:, None] * q * kl_integrand * wy[None, :]))
        rows.append({"tau": tau, "epsilon": eps, "x_nodes": len(x), "y_nodes": len(y),
                     "x_mass_error": abs(float(np.sum(base * w)) - 1.0), "conditional_mass_max_error": float(np.max(np.abs(np.sum(q * wy[None, :], axis=1) - 1.0))),
                     "mixture_mass_error": abs(float(np.sum(mixture * wy)) - 1.0), "objective_identity_residual": abs(direct - decomposed),
                     "continuous_density": True})
    return {"cells": len(rows), "rows": rows, "max_x_mass_error": max(r["x_mass_error"] for r in rows),
            "max_conditional_mass_error": max(r["conditional_mass_max_error"] for r in rows),
            "max_mixture_mass_error": max(r["mixture_mass_error"] for r in rows),
            "max_objective_identity_residual": max(r["objective_identity_residual"] for r in rows),
            "all_continuous": all(r["continuous_density"] for r in rows)}


def main() -> None:
    result = {"schema": "gradient-flow-non-gaussian-scope-v1", "source_scan": source_scan(), "claim_2_non_gaussian_flow": run_flow(),
              "claim_4_actual_ula_work": run_complexity(), "claim_6_continuous_half_bridge": half_bridge_continuous()}
    result["all_gates_pass"] = (
        result["source_scan"]["twelve_label_occurrences"]
        and result["claim_2_non_gaussian_flow"]["all_mass_one"]
        and result["claim_2_non_gaussian_flow"]["all_nonnegative"]
        and result["claim_2_non_gaussian_flow"]["max_actual_error_over_epsilon"] < 1.0
        and result["claim_4_actual_ula_work"]["all_work_counts_exact"]
        and result["claim_4_actual_ula_work"]["all_states_finite"]
        and result["claim_6_continuous_half_bridge"]["all_continuous"]
        and result["claim_6_continuous_half_bridge"]["max_mixture_mass_error"] < 2e-12
        and result["claim_6_continuous_half_bridge"]["max_objective_identity_residual"] < 2e-12
    )
    print(json.dumps(result, indent=2, sort_keys=True))
    if not result["all_gates_pass"]:
        raise SystemExit("non-Gaussian scope audit failed")


if __name__ == "__main__":
    main()