scripts/claim2_strong_convexity.py

"""Claim 2 -- Lemma 2.7: f_rho is rho-strongly convex and its conjugate f_rho* is
(1/rho)-smooth.

Independent tests:
  T1  second derivative of f_rho on its domain: analytic minimiser t=1/(2rho) and the
      numerical minimum of f_rho''  (predicted exactly 4 rho).
  T2  defining chord inequality of m-strong convexity for m = rho (claim), m = 4 rho
      (tight), m = 4.05 rho (must fail) -- executable tightness audit.
  T3  (f_rho*)'' bound: analytic maximiser s = -log rho, numerical max (predicted
      1/(4 rho)); Lipschitz-gradient inequality with L = 1/rho (claim), L = 1/(4 rho)
      (tight) and L = 1/(4.05 rho) (must fail).
  T4  Legendre duality identity  (f*)''(s) = 1 / f''(t)  at  t = (f*)'(s)  -- this is the
      independent re-derivation of "m-strongly convex  <=>  (1/m)-smooth conjugate".
  T5  biconjugate check (f*)* = f.
  T6  multivariate transfer: Hessian eigenvalues of  v -> sum_i w_i f_rho*(v_i)  and of
      the parametric objective G_rho(theta,alpha) of Sec. 2.2 for a linear-model loss.
  T7  boundary audit: rho -> 1^- and rho >= 1 (outside the paper's 0 < rho < 1).
"""

import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (
    F_rho,
    d2_f_rho,
    d2_f_rho_star,
    d_f_rho,
    d_f_rho_star,
    f_rho,
    f_rho_star,
)

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
os.makedirs(OUT, exist_ok=True)
SEED = 20260725
rng = np.random.default_rng(SEED)
res = {"seed": SEED}
RHOS = [1e-4, 1e-3, 1e-2, 0.03, 0.1, 0.3, 0.5, 0.7, 0.9, 0.99]

# ------------------------------------------------------------------ T1
t1 = {}
for rho in RHOS:
    tt = np.linspace(1e-10, 1.0 / rho * (1 - 1e-10), 2000001)
    d2 = d2_f_rho(tt, rho)
    k = int(np.argmin(d2))
    t1[str(rho)] = {
        "min_f2": float(d2.min()),
        "argmin_t": float(tt[k]),
        "predicted_argmin_1_over_2rho": 1.0 / (2 * rho),
        "min_over_rho": float(d2.min() / rho),
        "rho_SC_holds(min>=rho)": bool(d2.min() >= rho - 1e-12),
    }
res["T1_second_derivative"] = t1


# ------------------------------------------------------------------ T2
def sc_violations(rho, m, n=20000):
    lo, hi = 1e-9, 1.0 / rho * (1 - 1e-9)
    t1_ = rng.uniform(lo, hi, n)
    t2_ = rng.uniform(lo, hi, n)
    lam = rng.uniform(0, 1, n)
    mid = lam * t1_ + (1 - lam) * t2_
    lhs = f_rho(mid, rho)
    rhs = (
        lam * f_rho(t1_, rho)
        + (1 - lam) * f_rho(t2_, rho)
        - 0.5 * m * lam * (1 - lam) * (t1_ - t2_) ** 2
    )
    viol = lhs > rhs + 1e-9
    return int(viol.sum()), float(np.max(lhs - rhs))


t2 = {}
for rho in RHOS:
    a, ma = sc_violations(rho, rho)
    b, mb = sc_violations(rho, 4 * rho)
    c, mc = sc_violations(rho, 4.05 * rho)
    t2[str(rho)] = {
        "m=rho": {"violations": a, "max_excess": ma},
        "m=4rho": {"violations": b, "max_excess": mb},
        "m=4.05rho": {"violations": c, "max_excess": mc},
    }
res["T2_strong_convexity_chords"] = t2

# ------------------------------------------------------------------ T3
t3 = {}
for rho in RHOS:
    ss = np.linspace(-40 - np.log(rho), 40 - np.log(rho), 2000001)
    d2s = d2_f_rho_star(ss, rho)
    k = int(np.argmax(d2s))
    out = {
        "max_f_star2": float(d2s.max()),
        "argmax_s": float(ss[k]),
        "predicted_argmax_minus_log_rho": float(-np.log(rho)),
        "predicted_max_1_over_4rho": 1.0 / (4 * rho),
        "smooth_1_over_rho_holds": bool(d2s.max() <= 1.0 / rho + 1e-9),
    }
    for name, L in [
        ("L=1/rho", 1.0 / rho),
        ("L=1/(4rho)", 1.0 / (4 * rho)),
        ("L=1/(4.05rho)", 1.0 / (4.05 * rho)),
    ]:
        s1 = rng.uniform(-30 - np.log(rho), 30 - np.log(rho), 20000)
        s2 = rng.uniform(-30 - np.log(rho), 30 - np.log(rho), 20000)
        lhs = np.abs(d_f_rho_star(s1, rho) - d_f_rho_star(s2, rho))
        rhs = L * np.abs(s1 - s2)
        out[name] = {
            "violations": int(np.sum(lhs > rhs * (1 + 1e-12) + 1e-15)),
            "max_ratio_lhs_over_L_times_ds": float(
                np.max(lhs / np.maximum(rhs, 1e-300))
            ),
        }
    t3[str(rho)] = out
res["T3_conjugate_smoothness"] = t3

# ------------------------------------------------------------------ T4 duality identity
t4 = {"max_rel_err": 0.0, "n": 0}
for rho in RHOS:
    ss = np.linspace(-25 - np.log(rho), 25 - np.log(rho), 2001)
    tstar = d_f_rho_star(ss, rho)  # t = (f*)'(s)
    inv = 1.0 / d2_f_rho(tstar, rho)
    got = d2_f_rho_star(ss, rho)
    rel = np.max(np.abs(inv - got) / np.maximum(np.abs(got), 1e-300))
    t4["max_rel_err"] = max(t4["max_rel_err"], float(rel))
    t4["n"] += ss.size
    # also: s = f'(t) inverts (f*)'
    tt = np.linspace(1e-6, 1.0 / rho * (1 - 1e-6), 2001)
    back = d_f_rho_star(d_f_rho(tt, rho), rho)
    t4[f"inverse_map_max_abs_err_rho{rho}"] = float(np.max(np.abs(back - tt)))
res["T4_legendre_identity"] = t4

# ------------------------------------------------------------------ T5 biconjugate
t5 = {"max_abs_err": 0.0}
for rho in [0.03, 0.1, 0.3, 0.9]:
    tt = np.linspace(1e-4, 1.0 / rho * (1 - 1e-4), 400)
    ss = np.linspace(-60, 60, 400001)
    fs = f_rho_star(ss, rho)
    bic = np.array([np.max(t * ss - fs) for t in tt])
    err = float(np.max(np.abs(bic - f_rho(tt, rho))))
    t5[f"rho{rho}"] = err
    t5["max_abs_err"] = max(t5["max_abs_err"], err)
res["T5_biconjugate"] = t5

# ------------------------------------------------------------------ T6 multivariate
t6 = {}
for rho in [1e-2, 0.1, 0.3]:
    n = 40
    w = rng.dirichlet(np.ones(n))
    v = rng.normal(0, 3, n)
    diag = w * d2_f_rho_star(v, rho)  # Hessian diagonal of sum_i w_i f*(v_i)
    ev = diag
    t6[f"separable_rho{rho}"] = {
        "max_curvature_per_unit_weight": float(np.max(diag / w)),
        "bound_1_over_4rho": 1.0 / (4 * rho),
        "bound_1_over_rho": 1.0 / rho,
    }
    # G_rho(theta,alpha) = mean_i [ alpha + (lam/rho) softplus((l_i(theta)-alpha)/lam + log rho) ]
    # with l_i(theta) = (y_i - x_i^T theta)^2 ; compare ||Hess|| to (1/rho) * L(loss)
    d = 5
    X = rng.normal(size=(200, d))
    y = rng.normal(size=200)
    theta = rng.normal(size=d) * 0.3
    lam = 1.0

    def G(z):
        th, al = z[:d], z[d]
        l = (y - X @ th) ** 2
        return np.mean(
            al + (lam / rho) * np.logaddexp(0.0, (l - al) / lam + np.log(rho))
        )

    z0 = np.concatenate([theta, [1.0]])
    eps = 1e-4
    H = np.zeros((d + 1, d + 1))
    for i in range(d + 1):
        for j in range(d + 1):
            zpp, zpm, zmp, zmm = z0.copy(), z0.copy(), z0.copy(), z0.copy()
            zpp[i] += eps
            zpp[j] += eps
            zpm[i] += eps
            zpm[j] -= eps
            zmp[i] -= eps
            zmp[j] += eps
            zmm[i] -= eps
            zmm[j] -= eps
            H[i, j] = (G(zpp) - G(zpm) - G(zmp) + G(zmm)) / (4 * eps**2)
    ev = np.linalg.eigvalsh(0.5 * (H + H.T))
    Lloss = 2 * np.max(
        np.linalg.eigvalsh(X.T @ X / len(X))
    )  # smoothness of mean squared loss
    t6[f"G_rho_rho{rho}"] = {
        "min_eig": float(ev.min()),
        "max_eig": float(ev.max()),
        "loss_smoothness_L": float(Lloss),
        "max_eig_over_L": float(ev.max() / Lloss),
        "one_over_rho": 1.0 / rho,
        "convex(min_eig>=-1e-6)": bool(ev.min() >= -1e-6),
    }
res["T6_multivariate"] = t6

# ------------------------------------------------------------------ T7 boundary
t7 = {}
phi = rng.normal(0, 1.5, 200)
w = np.full(200, 1 / 200)
for rho in [0.9, 0.99, 0.999, 0.9999]:
    v, a = F_rho(phi, rho, w)
    t7[f"F_rho_rho{rho}"] = {
        "F_rho": v,
        "limit_mean_phi_minus_1": float(np.sum(w * phi) - 1.0),
    }
t7["mean_phi_minus_1"] = float(np.sum(w * phi) - 1.0)
# rho >= 1: domain [0,1/rho] cannot contain a density with mean 1 unless rho == 1 exactly
t7["rho_ge_1_infeasible"] = {
    "explanation": "for rho > 1 the constraint dnu/dmu <= 1/rho < 1 is incompatible with "
    "int dnu/dmu dmu = 1, so D_rho == +inf and the variational problem (5) is vacuous",
    "f_rho(1)_at_rho1.5_is_inf": bool(np.isinf(f_rho(np.array([1.0]), 1.5)[0])),
}
# strong-convexity modulus 4rho -> 4 as rho -> 1
t7["min_f2_at_rho0.999"] = float(
    np.min(d2_f_rho(np.linspace(1e-9, 1 / 0.999 * (1 - 1e-9), 200001), 0.999))
)
res["T7_boundary"] = t7

with open(os.path.join(OUT, "claim2_strong_convexity.json"), "w") as fh:
    json.dump(res, fh, indent=2, default=float)
print(json.dumps(res, indent=2, default=float)[:4000])