wave3/lse_core_wave3.py

#!/usr/bin/env python3
"""Claim-matched numerical audit for Improved Stochastic Optimization of LogSumExp.

This script uses only finite distributions.  It checks the Safe-KL LogSumExp
approximation, its exact gradients and stochastic oracle, curvature bounds,
CVaR limits, mini-batch bias, and projected-SGD convergence scaling.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path

import numpy as np


def logmeanexp(values: np.ndarray) -> float:
    maximum = float(np.max(values))
    return maximum + math.log(float(np.mean(np.exp(values - maximum))))


def sigmoid(values: np.ndarray) -> np.ndarray:
    positive = values >= 0
    output = np.empty_like(values, dtype=float)
    output[positive] = 1.0 / (1.0 + np.exp(-values[positive]))
    exp_values = np.exp(values[~positive])
    output[~positive] = exp_values / (1.0 + exp_values)
    return output


def safe_weights(values: np.ndarray, alpha: float, rho: float) -> np.ndarray:
    return sigmoid(values - alpha + math.log(rho)) / rho


def solve_alpha(values: np.ndarray, rho: float) -> float:
    lower = float(np.min(values)) - 80.0
    upper = float(np.max(values)) + 80.0
    for _ in range(160):
        middle = 0.5 * (lower + upper)
        if float(np.mean(safe_weights(values, middle, rho))) > 1.0:
            lower = middle
        else:
            upper = middle
    return 0.5 * (lower + upper)


def safe_lse(values: np.ndarray, rho: float) -> tuple[float, float, np.ndarray]:
    alpha = solve_alpha(values, rho)
    softplus = np.logaddexp(0.0, values - alpha + math.log(rho))
    objective = alpha - 1.0 + float(np.mean(softplus)) / rho
    weights = safe_weights(values, alpha, rho)
    return objective, alpha, weights


def finite_difference_gradient(values: np.ndarray, rho: float, step: float) -> np.ndarray:
    gradient = np.empty_like(values)
    for index in range(values.size):
        plus = values.copy()
        minus = values.copy()
        plus[index] += step
        minus[index] -= step
        gradient[index] = (safe_lse(plus, rho)[0] - safe_lse(minus, rho)[0]) / (2.0 * step)
    return gradient


def approximation_audit(rng: np.random.Generator) -> tuple[dict, list[dict]]:
    rhos = [0.3, 0.1, 0.03, 0.01, 0.003, 0.001]
    rows: list[dict] = []
    upper_violations = []
    proposition_lower_violations = []
    maximum_based_lower_violations = []
    proposition_lower_cells = 0
    maximum_based_lower_cells = 0
    monotonic_violations = 0
    gradient_errors = []
    gradient_sum_errors = []
    maximum_density_ratios = []

    for trial in range(24):
        values = rng.normal(0.0, 2.0, size=37)
        values[0] += 5.0
        exact = logmeanexp(values)
        kappa = math.exp(logmeanexp(2.0 * values) - 2.0 * exact)
        maximum = float(np.max(values))
        trial_approximations = []
        for rho in rhos:
            approximation, alpha, weights = safe_lse(values, rho)
            trial_approximations.append(approximation)
            upper_violations.append(max(0.0, approximation - exact))
            proposition_lower_bound = None
            maximum_based_lower_bound = None
            if rho < 1.0 / kappa:
                proposition_lower_bound = exact + rho / 2.0 + math.log(1.0 - rho * kappa)
                proposition_lower_violations.append(
                    max(0.0, proposition_lower_bound - approximation)
                )
                proposition_lower_cells += 1
            if rho < math.exp(exact - maximum):
                maximum_based_lower_bound = exact - rho * math.exp(maximum - exact)
                maximum_based_lower_violations.append(
                    max(0.0, maximum_based_lower_bound - approximation)
                )
                maximum_based_lower_cells += 1
            maximum_density_ratios.append(float(np.max(weights) * rho))
            rows.append(
                {
                    "trial": trial,
                    "rho": rho,
                    "logmeanexp": exact,
                    "safe_lse": approximation,
                    "gap": exact - approximation,
                    "gap_over_rho": (exact - approximation) / rho,
                    "kappa": kappa,
                    "proposition_lower_bound": proposition_lower_bound,
                    "maximum_based_lower_bound": maximum_based_lower_bound,
                    "alpha": alpha,
                    "weight_sum": float(np.mean(weights)),
                    "max_density": float(np.max(weights)),
                    "density_cap": 1.0 / rho,
                }
            )
        # F_rho increases when rho decreases.
        if any(
            trial_approximations[index + 1] + 1e-12 < trial_approximations[index]
            for index in range(len(rhos) - 1)
        ):
            monotonic_violations += 1

    check_values = rng.normal(0.0, 1.5, size=31)
    check_values[:2] += np.array([4.0, -3.0])
    for rho in [0.2, 0.05, 0.01]:
        _, _, weights = safe_lse(check_values, rho)
        analytic = weights / check_values.size
        numeric = finite_difference_gradient(check_values, rho, 1e-5)
        gradient_errors.append(float(np.max(np.abs(analytic - numeric))))
        gradient_sum_errors.append(abs(float(np.sum(analytic)) - 1.0))

    median_gap_by_rho = {
        str(rho): float(np.median([row["gap"] for row in rows if row["rho"] == rho]))
        for rho in rhos
    }
    rate = float(
        np.polyfit(
            np.log(np.asarray(rhos[-4:])),
            np.log(np.asarray([median_gap_by_rho[str(rho)] for rho in rhos[-4:]])),
            1,
        )[0]
    )
    summary = {
        "cells": len(rows),
        "rho_values": rhos,
        "maximum_upper_bound_violation": max(upper_violations),
        "proposition_lower_bound_cells": proposition_lower_cells,
        "maximum_proposition_lower_bound_violation": max(
            proposition_lower_violations, default=0.0
        ),
        "maximum_based_lower_bound_cells": maximum_based_lower_cells,
        "maximum_maximum_based_lower_bound_violation": max(
            maximum_based_lower_violations, default=0.0
        ),
        "monotonicity_violations": monotonic_violations,
        "median_gap_by_rho": median_gap_by_rho,
        "small_rho_loglog_gap_rate": rate,
        "maximum_finite_difference_gradient_error": max(gradient_errors),
        "maximum_gradient_sum_error": max(gradient_sum_errors),
        "maximum_density_to_cap_ratio": max(maximum_density_ratios),
    }
    return summary, rows


def curvature_audit() -> tuple[dict, list[dict]]:
    rows = []
    for rho in [0.5, 0.2, 0.1, 0.05, 0.01]:
        t = np.linspace(1e-5, 1.0 / rho - 1e-5, 200_001)
        f_second = 1.0 / (t * (1.0 - rho * t))
        s = np.linspace(-16.0, 16.0, 200_001)
        exp_s = np.exp(s)
        conjugate_second = exp_s / (1.0 + rho * exp_s) ** 2
        rows.append(
            {
                "rho": rho,
                "minimum_safe_kl_curvature": float(np.min(f_second)),
                "claimed_strong_convexity_floor": rho,
                "minimum_over_claimed_floor": float(np.min(f_second) / rho),
                "maximum_conjugate_curvature": float(np.max(conjugate_second)),
                "claimed_smoothness_ceiling": 1.0 / rho,
                "maximum_over_claimed_ceiling": float(np.max(conjugate_second) * rho),
            }
        )
    summary = {
        "cells": len(rows),
        "all_strong_convexity_checks_pass": all(
            row["minimum_safe_kl_curvature"] >= row["claimed_strong_convexity_floor"]
            for row in rows
        ),
        "all_smoothness_checks_pass": all(
            row["maximum_conjugate_curvature"] <= row["claimed_smoothness_ceiling"]
            for row in rows
        ),
        "observed_minimum_curvature_ratio": min(
            row["minimum_over_claimed_floor"] for row in rows
        ),
        "observed_maximum_smoothness_ratio": max(
            row["maximum_over_claimed_ceiling"] for row in rows
        ),
    }
    return summary, rows


def cvar_audit(rng: np.random.Generator) -> tuple[dict, list[dict]]:
    rho = 0.1
    values = rng.normal(0.0, 1.0, size=1000)
    values[-20:] += 5.0
    top_count = int(round(rho * values.size))
    cvar = float(np.mean(np.sort(values)[-top_count:]))
    rows = []
    for temperature in [1.0, 0.5, 0.2, 0.1, 0.05, 0.02]:
        scaled, _, _ = safe_lse(values / temperature, rho)
        adjusted = temperature * scaled - temperature * (math.log(rho) - 1.0)
        rows.append(
            {
                "temperature": temperature,
                "cvar": cvar,
                "adjusted_safe_objective": adjusted,
                "gap": adjusted - cvar,
                "proposition_gap_ceiling": temperature / rho,
                "inside_bounds": cvar - 1e-10 <= adjusted <= cvar + temperature / rho + 1e-10,
            }
        )
    summary = {
        "rho": rho,
        "cells": len(rows),
        "cvar": cvar,
        "all_proposition_bounds_pass": all(row["inside_bounds"] for row in rows),
        "gap_at_temperature_1": rows[0]["gap"],
        "gap_at_temperature_0.02": rows[-1]["gap"],
        "gap_reduction_factor": rows[0]["gap"] / rows[-1]["gap"],
    }
    return summary, rows


def oracle_audit(rng: np.random.Generator) -> tuple[dict, list[dict]]:
    population = rng.normal(0.0, 1.0, size=4096)
    population[-40:] = rng.normal(6.0, 0.2, size=40)
    theta = 0.0
    losses = 0.5 * (theta - population) ** 2
    loss_gradients = theta - population
    rho = 0.05
    _, alpha, weights = safe_lse(losses, rho)
    contributions = weights * loss_gradients
    exact_safe_gradient = float(np.mean(contributions))
    exact_lse_gradient = float(
        np.sum(np.exp(losses - np.max(losses)) * loss_gradients)
        / np.sum(np.exp(losses - np.max(losses)))
    )

    rows = []
    for batch_size in [4, 16, 64, 256]:
        repetitions = 4000
        indices = rng.integers(0, population.size, size=(repetitions, batch_size))
        proposed = np.mean(contributions[indices], axis=1)
        batch_losses = losses[indices]
        batch_gradients = loss_gradients[indices]
        shifted = batch_losses - np.max(batch_losses, axis=1, keepdims=True)
        batch_probabilities = np.exp(shifted)
        baseline = np.sum(batch_probabilities * batch_gradients, axis=1) / np.sum(
            batch_probabilities, axis=1
        )
        rows.append(
            {
                "batch_size": batch_size,
                "repetitions": repetitions,
                "exact_safe_gradient": exact_safe_gradient,
                "proposed_mean": float(np.mean(proposed)),
                "proposed_absolute_bias": abs(float(np.mean(proposed)) - exact_safe_gradient),
                "proposed_rmse": float(
                    np.sqrt(np.mean((proposed - exact_safe_gradient) ** 2))
                ),
                "exact_logsumexp_gradient": exact_lse_gradient,
                "batch_logsumexp_mean": float(np.mean(baseline)),
                "batch_logsumexp_absolute_bias": abs(
                    float(np.mean(baseline)) - exact_lse_gradient
                ),
            }
        )

    sample_sizes = np.asarray([32, 128, 512, 2048, 8192])
    rmse_rows = []
    for sample_size in sample_sizes:
        repetitions = 500
        estimates = np.empty(repetitions)
        for repetition in range(repetitions):
            indices = rng.integers(0, population.size, size=int(sample_size))
            estimates[repetition] = float(np.mean(contributions[indices]))
        rmse_rows.append(
            {
                "samples": int(sample_size),
                "repetitions": repetitions,
                "rmse": float(
                    np.sqrt(np.mean((estimates - exact_safe_gradient) ** 2))
                ),
            }
        )
    slope = float(
        np.polyfit(
            np.log(sample_sizes),
            np.log([row["rmse"] for row in rmse_rows]),
            1,
        )[0]
    )
    summary = {
        "rho": rho,
        "alpha": alpha,
        "exact_safe_gradient": exact_safe_gradient,
        "exact_logsumexp_gradient": exact_lse_gradient,
        "batch_comparison_cells": len(rows),
        "oracle_scaling_cells": len(rmse_rows),
        "oracle_rmse_loglog_slope": slope,
        "smallest_proposed_bias": min(row["proposed_absolute_bias"] for row in rows),
        "largest_batch_logsumexp_bias": max(
            row["batch_logsumexp_absolute_bias"] for row in rows
        ),
        "oracle_scaling": rmse_rows,
    }
    return summary, rows + rmse_rows


def minimized_safe_objective(theta: float, population: np.ndarray, rho: float) -> float:
    losses = 0.5 * (theta - population) ** 2
    return safe_lse(losses, rho)[0]


def golden_section_minimize(population: np.ndarray, rho: float) -> tuple[float, float]:
    left, right = -2.0, 2.0
    ratio = (math.sqrt(5.0) - 1.0) / 2.0
    x1 = right - ratio * (right - left)
    x2 = left + ratio * (right - left)
    f1 = minimized_safe_objective(x1, population, rho)
    f2 = minimized_safe_objective(x2, population, rho)
    for _ in range(100):
        if f1 < f2:
            right, x2, f2 = x2, x1, f1
            x1 = right - ratio * (right - left)
            f1 = minimized_safe_objective(x1, population, rho)
        else:
            left, x1, f1 = x1, x2, f2
            x2 = left + ratio * (right - left)
            f2 = minimized_safe_objective(x2, population, rho)
    theta = 0.5 * (left + right)
    return theta, minimized_safe_objective(theta, population, rho)


def convergence_audit(rng: np.random.Generator) -> tuple[dict, list[dict]]:
    population = rng.normal(-0.2, 0.8, size=2048)
    population[-30:] = rng.normal(3.5, 0.15, size=30)
    rho = 0.1
    theta_star, objective_star = golden_section_minimize(population, rho)
    rows = []
    budgets = [256, 1024, 4096, 16384]
    repetitions = 64
    for budget in budgets:
        theta = np.full(repetitions, 1.5)
        initial_losses = 0.5 * (1.5 - population) ** 2
        alpha = np.full(repetitions, solve_alpha(initial_losses, rho))
        theta_sum = np.zeros(repetitions)
        alpha_sum = np.zeros(repetitions)
        step = 0.45 / math.sqrt(budget)
        for _ in range(budget):
            indices = rng.integers(0, population.size, size=repetitions)
            samples = population[indices]
            losses = 0.5 * (theta - samples) ** 2
            weights = sigmoid(losses - alpha + math.log(rho)) / rho
            gradient_theta = weights * (theta - samples)
            gradient_alpha = 1.0 - weights
            theta = np.clip(theta - step * gradient_theta, -2.0, 2.0)
            alpha = np.clip(alpha - step * gradient_alpha, -5.0, 20.0)
            theta_sum += theta
            alpha_sum += alpha
        average_theta = theta_sum / budget
        excess = np.asarray(
            [
                minimized_safe_objective(value, population, rho) - objective_star
                for value in average_theta
            ]
        )
        rows.append(
            {
                "iterations": budget,
                "repetitions": repetitions,
                "step_size": step,
                "median_surrogate_excess": float(np.median(excess)),
                "mean_surrogate_excess": float(np.mean(excess)),
                "q90_surrogate_excess": float(np.quantile(excess, 0.9)),
                "median_theta_error": float(np.median(np.abs(average_theta - theta_star))),
            }
        )
    slope = float(
        np.polyfit(
            np.log(np.asarray(budgets)),
            np.log(np.asarray([row["median_surrogate_excess"] for row in rows])),
            1,
        )[0]
    )
    summary = {
        "rho": rho,
        "theta_star": theta_star,
        "objective_star": objective_star,
        "cells": len(rows),
        "repetitions_per_cell": repetitions,
        "median_excess_loglog_slope": slope,
        "median_excess_reduction_factor": (
            rows[0]["median_surrogate_excess"] / rows[-1]["median_surrogate_excess"]
        ),
        "all_median_excess_values_decrease": all(
            rows[index + 1]["median_surrogate_excess"]
            < rows[index]["median_surrogate_excess"]
            for index in range(len(rows) - 1)
        ),
    }
    return summary, rows


def write_csv(path: Path, rows: list[dict]) -> None:
    keys: list[str] = []
    for row in rows:
        for key in row:
            if key not in keys:
                keys.append(key)
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=keys)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("wave3-output"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(20260726)

    approximation, approximation_rows = approximation_audit(rng)
    curvature, curvature_rows = curvature_audit()
    cvar, cvar_rows = cvar_audit(rng)
    oracle, oracle_rows = oracle_audit(rng)
    convergence, convergence_rows = convergence_audit(rng)
    results = {
        "paper": {
            "title": "Improved Stochastic Optimization of LogSumExp",
            "openreview": "TzQElzflxR",
            "arxiv": "2509.24894",
            "official_repository": "https://github.com/egorgladin/logsumexp-approx",
        },
        "approximation_and_gradient": approximation,
        "curvature": curvature,
        "cvar": cvar,
        "stochastic_oracle": oracle,
        "projected_sgd": convergence,
        "limitations": [
            "Finite-distribution mechanism audit, not a proof of the paper's theorems.",
            "Does not rerun the paper's neural OT, California Housing, or MNIST training.",
        ],
    }
    (args.output_dir / "results.json").write_text(
        json.dumps(results, indent=2, sort_keys=True) + "\n"
    )
    write_csv(args.output_dir / "approximation_sweep.csv", approximation_rows)
    write_csv(args.output_dir / "curvature_sweep.csv", curvature_rows)
    write_csv(args.output_dir / "cvar_sweep.csv", cvar_rows)
    write_csv(args.output_dir / "oracle_sweep.csv", oracle_rows)
    write_csv(args.output_dir / "convergence_sweep.csv", convergence_rows)
    print(json.dumps(results, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()