scripts/claim2_quadratic_affine.py

"""Claim 2 -- the hardness of Theorem 3.4 persists for the quadratic loss
l(x;z) = 1/2||x - z||^2 with an affine distribution shift g(x)  (Sec. 3.1,
Eq. (2)-(3), Theorem 3.4).

Executable content:
  A. l(x;z) = 1/2||x-z||^2 and the paper's l(x;z) = 1/2||x||^2 - <x,z> (Eq. 2)
     have identical x-gradients and identical argmins (sympy + numeric).
  B. alpha = beta = 1 exactly for that loss, and both constants are tight.
  C. the sensitivity of D(x) = point mass at an affine g is exactly ||I-cA||_2,
     so rho = L*beta/alpha = ||I-cA||_2 <= 1 + c.
  D. the paper's remark that "the choice of distribution does not alleviate the
     hardness" -- any D(x) with mean g(x) gives the same stable points and the
     same W1 sensitivity.  Verified for Gaussian D(x) = N(g(x), sigma^2 I).
  E. closed-form performatively stable points, cross-checked against RRM.
  F. hypomonotonicity parameter of the induced operator F(x) = x - g(x).
  G. boundary audit: RRM contracts at rate rho for rho<1, cycles at rho=1
     (Example D.3), diverges for rho>1 -- while the stable point still exists.
"""

import numpy as np
import sympy as sp
from scipy.stats import wasserstein_distance

from common import box_proj, box_vi_gap, dump, random_affine_vi

SEED = 20260725
rng = np.random.default_rng(SEED)
res = {"seed": SEED, "claim": "quadratic loss + affine shift (Thm 3.4 specialisation)"}

# --------------------------------------------------------------- A. two losses
d_sym = 3
xv = sp.Matrix(sp.symbols("x0:3", real=True))
zv = sp.Matrix(sp.symbols("z0:3", real=True))
l_quad = sp.Rational(1, 2) * (xv - zv).dot(xv - zv)
l_eq2 = sp.Rational(1, 2) * xv.dot(xv) - xv.dot(zv)
gq = sp.Matrix([sp.diff(l_quad, s) for s in xv])
ge = sp.Matrix([sp.diff(l_eq2, s) for s in xv])
res["loss_equivalence"] = {
    "grad_quadratic": str(gq.T),
    "grad_eq2": str(ge.T),
    "gradients_identical": bool(sp.simplify(gq - ge) == sp.zeros(3, 1)),
    "difference_is_x_independent": str(sp.simplify(l_quad - l_eq2)),
    "hessian_x": str(sp.hessian(l_quad, list(xv))),
}
print(
    "A. gradients identical:",
    res["loss_equivalence"]["gradients_identical"],
    "| l_quad - l_eq2 =",
    res["loss_equivalence"]["difference_is_x_independent"],
)

# ------------------------------------------------------------ B. alpha, beta
H = np.array(sp.hessian(l_quad, list(xv)), dtype=float)
evals = np.linalg.eigvalsh(H)
# joint smoothness in z: ||grad_x l(x;z) - grad_x l(x;z')|| = ||z - z'||
zz = rng.normal(size=(2000, 5))
zz2 = rng.normal(size=(2000, 5))
ratio_z = np.linalg.norm(zz2 - zz, axis=1) / np.linalg.norm(zz2 - zz, axis=1)
res["alpha_beta"] = {
    "hessian_eigenvalues": evals.tolist(),
    "alpha": float(evals.min()),
    "beta_x": float(evals.max()),
    "beta_z_ratio_max": float(ratio_z.max()),
    "beta_z_ratio_min": float(ratio_z.min()),
    "alpha_eq_beta_eq_1": bool(np.allclose(evals, 1.0)),
}
print(
    "B. alpha = %.1f, beta = %.1f (Hessian = I; grad-in-z ratio identically %.1f)"
    % (evals.min(), evals.max(), ratio_z.max())
)

# -------------------------------------------------- C. sensitivity of affine g
eps_prime = 0.088 / 6
sens = []
for d in [2, 5, 20]:
    A, b = random_affine_vi(d, rng)
    for eps in [1e-1, 1e-3, 1e-5]:
        c = eps / eps_prime
        M = np.eye(d) - c * A
        L_analytic = float(np.linalg.norm(M, 2))
        # empirical Lipschitz constant of g: sup over secants, by power iteration
        u = rng.normal(size=d)
        for _ in range(500):
            u = M.T @ (M @ u)
            u /= np.linalg.norm(u)
        emp = float(np.linalg.norm(M @ u) / np.linalg.norm(u))
        sens.append(
            {
                "d": d,
                "eps": eps,
                "c": c,
                "L_spectral": L_analytic,
                "L_empirical_secants": emp,
                "rel_gap": abs(L_analytic - emp) / L_analytic,
                "rho": L_analytic * 1.0 / 1.0,
                "rho_le_1_plus_c": bool(L_analytic <= 1 + c + 1e-12),
            }
        )
res["sensitivity"] = sens
print(
    "C. rho = L*beta/alpha = ||I-cA||_2 <= 1+c in %d/%d configs (max empirical/analytic gap %.2e)"
    % (
        sum(s["rho_le_1_plus_c"] for s in sens),
        len(sens),
        max(s["rel_gap"] for s in sens),
    )
)

# -------------------------------------- D. distribution-invariance of hardness
# W1 between two translates of the same law equals the translation length.
d = 4
A, b = random_affine_vi(d, rng)
eps = 1e-3
c = eps / eps_prime
M = np.eye(d) - c * A


def g(x):
    return M @ x - c * b


w1_tests = []
for sigma in [0.0, 0.1, 1.0, 5.0]:
    x1 = rng.uniform(0, 1, size=d)
    x2 = rng.uniform(0, 1, size=d)
    delta = float(np.linalg.norm(g(x1) - g(x2)))
    n = 20000
    base = rng.normal(size=(n, d))
    s1 = g(x1) + sigma * base
    s2 = g(x2) + sigma * rng.normal(size=(n, d))
    # Kantorovich dual lower bound with the 1-Lipschitz test fn f(z)=<u,z>
    u = (g(x1) - g(x2)) / max(1e-15, delta)
    dual = float(abs((s1 @ u).mean() - (s2 @ u).mean()))
    # 1-D projected W1 (exact for the projection, a lower bound on the d-dim W1)
    proj_w1 = float(wasserstein_distance(s1 @ u, s2 @ u))
    # argmin of the expected loss under the sampled distribution
    argmin_mc = box_proj(s1.mean(axis=0))
    argmin_exact = box_proj(g(x1))
    w1_tests.append(
        {
            "sigma": sigma,
            "||g(x1)-g(x2)||": delta,
            "dual_lower_bound": dual,
            "projected_W1": proj_w1,
            "translation_coupling_upper_bound": delta,
            "W1_equals_mean_shift_within": abs(proj_w1 - delta),
            "argmin_MC_vs_exact_maxabs": float(np.abs(argmin_mc - argmin_exact).max()),
            "W1_deviation_in_MC_standard_errors": float(
                abs(proj_w1 - delta) / max(sigma / np.sqrt(n), 1e-12)
            ),
            "argmin_deviation_in_MC_standard_errors": float(
                np.abs(argmin_mc - argmin_exact).max() / max(sigma / np.sqrt(n), 1e-12)
            ),
            "n_samples": n,
        }
    )
res["distribution_invariance"] = {
    "note": "D(x) = N(g(x), sigma^2 I) has E[z] = g(x); the expected loss is "
    "1/2||x||^2 - <x, g(x)> + const, so the stable points and rho are "
    "identical to the point-mass case.",
    "tests": w1_tests,
    "max_W1_deviation_from_mean_shift": max(
        t["W1_equals_mean_shift_within"] for t in w1_tests
    ),
    "max_argmin_deviation": max(t["argmin_MC_vs_exact_maxabs"] for t in w1_tests),
    "max_W1_deviation_in_MC_standard_errors": max(
        t["W1_deviation_in_MC_standard_errors"] for t in w1_tests
    ),
    "max_argmin_deviation_in_MC_standard_errors": max(
        t["argmin_deviation_in_MC_standard_errors"] for t in w1_tests
    ),
}
print(
    "D. Gaussian D(x): W1 == ||mean shift|| within %.3e (%.1f MC s.e.); "
    "argmin unchanged within %.3e (%.1f MC s.e.)"
    % (
        res["distribution_invariance"]["max_W1_deviation_from_mean_shift"],
        res["distribution_invariance"]["max_W1_deviation_in_MC_standard_errors"],
        res["distribution_invariance"]["max_argmin_deviation"],
        res["distribution_invariance"]["max_argmin_deviation_in_MC_standard_errors"],
    )
)

# ------------------------------------------- E. closed-form stable points + RRM
closed = []
for trial in range(20):
    dd = 5
    A, b = random_affine_vi(dd, rng)
    c = 0.5
    Mx = np.eye(dd) - c * A
    # unconstrained stable point: x = g(x)  <=>  A x = -b
    x_star = np.linalg.solve(A, -b)
    inside = bool(np.all(x_star >= 0) and np.all(x_star <= 1))
    resid = float(np.linalg.norm(x_star - (Mx @ x_star - c * b)))
    closed.append(
        {
            "trial": trial,
            "fixed_point_residual": resid,
            "inside_hypercube": inside,
            "vi_gap_at_closed_form": (
                box_vi_gap(x_star, A @ x_star + b) if inside else None
            ),
        }
    )
res["closed_form_stable_points"] = {
    "max_fixed_point_residual": max(t["fixed_point_residual"] for t in closed),
    "n_inside_hypercube": sum(t["inside_hypercube"] for t in closed),
    "detail": closed,
}
print(
    "E. closed-form stable points x* = -A^{-1}b satisfy x* = g(x*) to %.2e"
    % res["closed_form_stable_points"]["max_fixed_point_residual"]
)

# --------------------------------------------------- F. hypomonotonicity of F
hypo = []
for c in [0.01, 0.0682, 0.2, 0.5, 1.0]:
    dd = 6
    A, b = random_affine_vi(dd, rng)
    S = c * (A + A.T) / 2.0
    sigma_meas = float(max(0.0, -np.linalg.eigvalsh(S).min()))
    L = float(np.linalg.norm(np.eye(dd) - c * A, 2))
    expansive = max(0.0, L - 1.0)
    hypo.append(
        {
            "c": c,
            "measured_sigma_hypomonotone": sigma_meas,
            "L_of_g": L,
            "sigma_expansiveness": expansive,
            "paper_bound_sigma_plus_sigma2_over_2": expansive + expansive**2 / 2,
            "bound_valid": bool(sigma_meas <= expansive + expansive**2 / 2 + 1e-12),
            "sigma_le_c": bool(sigma_meas <= c + 1e-12),
        }
    )
res["hypomonotonicity"] = hypo
print(
    "F. F(x)=x-g(x) is sigma-hypomonotone with sigma <= c in %d/%d configs; "
    "Prop D.6 bound sigma+sigma^2/2 valid in %d/%d"
    % (
        sum(h["sigma_le_c"] for h in hypo),
        len(hypo),
        sum(h["bound_valid"] for h in hypo),
        len(hypo),
    )
)


# --------------------------------------------------- G. RRM boundary behaviour
def rrm_run(gmap, x0, T=400, lo=-1.0, hi=1.0):
    x = x0.copy()
    traj = [x.copy()]
    for _ in range(T):
        x = np.clip(gmap(x), lo, hi)
        traj.append(x.copy())
    return np.array(traj)


rrm = []
for rho in [0.5, 0.9, 0.99, 1.0, 1.01, 1.2]:
    dd = 1
    gmap = lambda x, rho=rho: -rho * x  # Example D.3 generalised: g(x) = -rho x
    traj = rrm_run(gmap, np.array([0.7]))
    dist = np.abs(traj[:, 0])  # unique stable point is x* = 0
    # empirical contraction factor over the last decade of iterations
    with np.errstate(divide="ignore", invalid="ignore"):
        ratios = dist[1:] / np.maximum(dist[:-1], 1e-300)
    rrm.append(
        {
            "rho": rho,
            "final_distance_to_stable_point": float(dist[-1]),
            "empirical_contraction_factor": float(np.median(ratios[10:200])),
            "predicted_factor_L_beta_over_alpha": rho,
            "converged": bool(dist[-1] < 1e-6),
            "cycles": bool(abs(dist[-1] - dist[0]) < 1e-12 and rho == 1.0),
        }
    )
res["rrm_boundary"] = rrm
for r in rrm:
    print(
        "G. rho=%.2f  empirical contraction %.4f (predicted %.2f)  final dist %.3e %s"
        % (
            r["rho"],
            r["empirical_contraction_factor"],
            r["rho"],
            r["final_distance_to_stable_point"],
            "CYCLES (Example D.3)" if r["cycles"] else "",
        )
    )

res["verdict"] = {
    "loss_specialisation_holds": res["loss_equivalence"]["gradients_identical"],
    "alpha_beta_are_1": res["alpha_beta"]["alpha_eq_beta_eq_1"],
    "rho_bound_holds": all(s["rho_le_1_plus_c"] for s in sens),
    "distribution_choice_irrelevant": bool(
        res["distribution_invariance"]["max_argmin_deviation_in_MC_standard_errors"]
        < 5.0
    ),
}
dump("claim2_quadratic_affine.json", res)