code/evidence2.py

"""Sections D and E: structural-precondition tests on real solver output and
measured pseudo-dimension of real tuning classes against the theorem bounds."""
import itertools, json, math, sys, time
import numpy as np
from scipy import linalg
from evidence import stated_bound, measure_pdim_pool

RNG = np.random.default_rng(4242)


# --------------------------------------------------------------- ElasticNet
def enet_solve(A, b, a1, a2, iters=6000):
    """Cyclic coordinate descent for min ||A th - b||^2 + a1|th|_1 + a2|th|^2."""
    n, d = A.shape
    th = np.zeros(d)
    col = (A ** 2).sum(0)
    r = -b.copy()
    for _ in range(iters):
        mx = 0.0
        for k in range(d):
            r -= A[:, k] * th[k]
            rho = -2.0 * A[:, k] @ r
            den = 2.0 * col[k] + 2.0 * a2
            new = np.sign(rho) * max(abs(rho) - a1, 0.0) / den
            mx = max(mx, abs(new - th[k]))
            th[k] = new
            r += A[:, k] * th[k]
        if mx < 1e-14:
            break
    return th


def cramer_residual(A, b, a1, a2, th, tol=1e-9):
    """On the active set S with signs s, stationarity says
    (2 A_S'A_S + 2 a2 I) th_S = 2 A_S'b - a1 s, whose solution is a rational
    function of (a1, a2) by Cramer's rule.  Return the residual norm."""
    S = np.where(np.abs(th) > tol)[0]
    if S.size == 0:
        return 0.0, 0
    As = A[:, S]
    G = 2.0 * As.T @ As + 2.0 * a2 * np.eye(S.size)
    rhs = 2.0 * As.T @ b - a1 * np.sign(th[S])
    return float(np.linalg.norm(G @ th[S] - rhs)), int(S.size)


def group_lasso_solve(A, b, alphas, groups, iters=40000):
    """Proximal gradient for min ||A th - b||^2 + sum_i alpha_i ||th_{G_i}||_2."""
    d = A.shape[1]
    th = np.zeros(d)
    Lc = 2.0 * np.linalg.norm(A, 2) ** 2
    for _ in range(iters):
        g = 2.0 * A.T @ (A @ th - b)
        z = th - g / Lc
        new = z.copy()
        for gi, G in enumerate(groups):
            nz = np.linalg.norm(z[G])
            sc = max(0.0, 1.0 - alphas[gi] / (Lc * max(nz, 1e-300)))
            new[G] = z[G] * sc
        if np.max(np.abs(new - th)) < 1e-15:
            th = new
            break
        th = new
    return th


def section_D():
    out = {}
    # ---- D1  Assumption 7.1: is the ElasticNet path piecewise rational?
    d, n = 6, 25
    A = RNG.standard_normal((n, d)); b = RNG.standard_normal(n)
    res, sizes = [], []
    grid = [(a1, a2) for a1 in np.linspace(0.05, 20.0, 40)
            for a2 in np.linspace(0.05, 5.0, 40)]
    for a1, a2 in grid:
        th = enet_solve(A, b, a1, a2)
        r, sz = cramer_residual(A, b, a1, a2, th)
        res.append(r); sizes.append(sz)
    out["enet"] = dict(solves=len(grid), max_cramer_residual=float(np.max(res)),
                       mean_cramer_residual=float(np.mean(res)),
                       active_sizes=[int(min(sizes)), int(max(sizes))])
    # rational vs equal-capacity polynomial fit inside one active-set region
    a1f = 3.0
    pts = []
    for a2 in np.linspace(0.20, 3.00, 60):
        th = enet_solve(A, b, a1f, a2)
        pts.append((a2, th))
    S0 = tuple(np.where(np.abs(pts[0][1]) > 1e-9)[0])
    same = [(a, t) for a, t in pts if tuple(np.where(np.abs(t) > 1e-9)[0]) == S0]
    xa = np.array([a for a, _ in same]); ya = np.array([t[S0[0]] for _, t in same])
    # inside a fixed active set S the exact solution is
    #   th_S(a2) = (G0 + 2 a2 I)^{-1} r , a ratio of a degree-(|S|-1) numerator
    #   to a degree-|S| denominator in a2.  Fit that form, and an
    #   equal-parameter-count polynomial, and compare residuals.
    nS = len(S0)
    xs_ = (xa - xa.mean()) / xa.std()
    Mr = np.column_stack([xs_ ** k for k in range(nS)] +
                         [-(xs_ ** k) * ya for k in range(1, nS + 1)])
    cr, *_ = np.linalg.lstsq(Mr, ya, rcond=None)
    num = sum(cr[k] * xs_ ** k for k in range(nS))
    den = 1.0 + sum(cr[nS + k - 1] * xs_ ** k for k in range(1, nS + 1))
    fit_r = num / den
    npar = Mr.shape[1]
    Mp = np.column_stack([xs_ ** k for k in range(npar)])
    cp, *_ = np.linalg.lstsq(Mp, ya, rcond=None)
    fit_p = Mp @ cp
    out["enet_fit"] = dict(region_points=len(same),
                           rational_rel_residual=float(np.linalg.norm(fit_r - ya) /
                                                       max(np.linalg.norm(ya), 1e-30)),
                           polynomial_rel_residual=float(np.linalg.norm(fit_p - ya) /
                                                         max(np.linalg.norm(ya), 1e-30)),
                           active_set_size=nS, free_parameters_each=int(npar))
    # control: group LASSO violates the same stationarity identity
    groups = [np.arange(0, 3), np.arange(3, 6)]
    viol = []
    for a in np.linspace(0.5, 3.0, 12):
        th = group_lasso_solve(A, b, np.array([a, a]), groups)
        S = np.where(np.abs(th) > 1e-9)[0]
        if S.size == 0:
            continue
        As = A[:, S]
        G = 2.0 * As.T @ As
        rhs = 2.0 * As.T @ b - a * np.sign(th[S])
        viol.append(float(np.linalg.norm(G @ th[S] - rhs)))
    out["group_lasso_control"] = dict(n=len(viol), max_violation=float(np.max(viol)),
                                      min_violation=float(np.min(viol)))

    # ---- D2  Theorem 8.1 premise: ||v||_2 is not piecewise polynomial
    NS = 6000
    v = RNG.standard_normal((NS, 4))
    nrm = np.linalg.norm(v, axis=1)
    sq = nrm ** 2
    feats = []
    for deg in range(9):
        for e in itertools.combinations_with_replacement(range(4), deg):
            c = np.ones(NS)
            for i in e:
                c = c * v[:, i]
            feats.append(c)
    F = np.column_stack(feats)
    def relres(y):
        c, *_ = np.linalg.lstsq(F, y, rcond=None)
        return float(np.linalg.norm(F @ c - y) / np.linalg.norm(y))
    out["nonpoly"] = dict(deg8_residual_norm=relres(nrm),
                          deg8_residual_normsq=relres(sq),
                          n_features=F.shape[1])
    # nu-encoding on real proximal-gradient solves
    enc = []
    for _ in range(48):
        Ai = RNG.standard_normal((30, 8)); bi = RNG.standard_normal(30)
        gs = [np.arange(0, 4), np.arange(4, 8)]
        al = RNG.uniform(0.3, 3.0, 2)
        th = group_lasso_solve(Ai, bi, al, gs)
        nu = np.array([np.linalg.norm(th[G]) for G in gs])
        e1 = max(abs(nu[i] ** 2 - float(th[gs[i]] @ th[gs[i]])) for i in range(2))
        # KKT residual of the group-LASSO stationarity condition
        gsm = 2.0 * Ai.T @ (Ai @ th - bi)
        kkt = 0.0
        for i, G in enumerate(gs):
            if nu[i] > 1e-10:
                kkt = max(kkt, float(np.linalg.norm(gsm[G] + al[i] * th[G] / nu[i])))
            else:
                kkt = max(kkt, max(0.0, float(np.linalg.norm(gsm[G])) - al[i]))
        enc.append((e1, kkt, float(nu.min())))
    out["nu_encoding"] = dict(solves=len(enc),
                              max_nu_sq_error=float(max(e[0] for e in enc)),
                              max_kkt_residual=float(max(e[1] for e in enc)),
                              min_nu=float(min(e[2] for e in enc)))

    # ---- D3  Proposition G.1 for weighted fused LASSO
    def fused(y, w):
        d = y.size
        D = np.zeros((d - 1, d))
        for i in range(d - 1):
            D[i, i] = -1.0; D[i, i + 1] = 1.0
        H = D @ D.T
        u = np.zeros(d - 1)
        Lc = np.linalg.norm(H, 2)
        for _ in range(30000):
            g = H @ u - D @ y
            un = np.clip(u - g / Lc, -w, w)
            if np.max(np.abs(un - u)) < 1e-15:
                u = un; break
            u = un
        th = y - D.T @ u
        prim = 0.5 * float((th - y) @ (th - y)) + float(w @ np.abs(D @ th))
        dual = -0.5 * float(u @ H @ u) + float(u @ (D @ y))
        return th, u, prim, dual, H
    gaps, eigs, regions = [], [], []
    for d in [4, 6, 8, 10, 12]:
        y = RNG.standard_normal(d)
        seen = set()
        for _ in range(400):
            w = RNG.uniform(0.05, 1.5, d - 1)
            th, u, prim, dual, H = fused(y, w)
            gaps.append(abs(prim - dual))
            seen.add(tuple(np.sign(np.round(np.diff(th), 9)).astype(int)))
        eigs.append(float(np.min(np.linalg.eigvalsh(H))))
        regions.append(dict(d=d, regions=len(seen), cap=3 ** (d - 1)))
    out["fused"] = dict(max_duality_gap=float(np.max(gaps)),
                        min_dual_hessian_eig=float(np.min(eigs)),
                        regions=regions,
                        all_under_cap=all(r["regions"] < r["cap"] for r in regions))
    # piecewise-affine check: within one region theta*(w) is affine in w
    y = RNG.standard_normal(8)
    base = RNG.uniform(0.4, 0.6, 7)
    th0, u0, *_ = fused(y, base)
    sgn0 = tuple(np.sign(np.round(np.diff(th0), 9)).astype(int))
    W, T = [], []
    for _ in range(120):
        w = base + RNG.uniform(-0.02, 0.02, 7)
        th, u, *_ = fused(y, w)
        if tuple(np.sign(np.round(np.diff(th), 9)).astype(int)) == sgn0:
            W.append(np.concatenate([[1.0], w])); T.append(th)
    W = np.array(W); T = np.array(T)
    c, *_ = np.linalg.lstsq(W, T, rcond=None)
    inreg = float(np.linalg.norm(W @ c - T) / max(np.linalg.norm(T), 1e-30))
    W2, T2 = [], []
    for _ in range(160):
        w = RNG.uniform(0.05, 1.5, 7)
        th, u, *_ = fused(y, w)
        W2.append(np.concatenate([[1.0], w])); T2.append(th)
    W2 = np.array(W2); T2 = np.array(T2)
    c2, *_ = np.linalg.lstsq(W2, T2, rcond=None)
    across = float(np.linalg.norm(W2 @ c2 - T2) / max(np.linalg.norm(T2), 1e-30))
    out["fused_affine"] = dict(in_region_points=len(T), in_region_rel_residual=inreg,
                               across_region_points=len(T2),
                               across_region_rel_residual=across,
                               separation=across / max(inreg, 1e-300))
    # rank-deficient control breaks Prop G.1's precondition
    Dbad = np.zeros((3, 4)); Dbad[0] = [-1, 1, 0, 0]; Dbad[1] = [-1, 1, 0, 0]; Dbad[2] = [0, 0, -1, 1]
    Hb = Dbad @ Dbad.T
    out["fused_control"] = dict(min_eig_rank_deficient=float(np.min(np.linalg.eigvalsh(Hb))),
                                rank=int(np.linalg.matrix_rank(Dbad)), rows=3)
    return out


# --------------------------------------------------- E: real tuning classes
def make_instance(n, d, p, seed):
    r = np.random.default_rng(seed)
    A = r.standard_normal((n, d)); b = r.standard_normal(n)
    Ap = r.standard_normal((n, d)); bp = r.standard_normal(n)
    groups = np.array_split(np.arange(d), p)
    return A, b, Ap, bp, groups


def ridge_theta(A, b, alpha_vec):
    G = A.T @ A + np.diag(alpha_vec)
    return np.linalg.solve(G, A.T @ b)


def section_E():
    out = {}
    # ---- f != g check on a real bi-level ridge instance
    A, b, Ap, bp, groups = make_instance(40, 8, 4, 11)
    al = np.array([0.7, 1.3, 0.2, 2.1])
    av = np.zeros(8)
    for gi, G in enumerate(groups):
        av[G] = al[gi]
    th = ridge_theta(A, b, av)
    gf = 2.0 * (A.T @ (A @ th - b) + av * th)
    gg = 2.0 * Ap.T @ (Ap @ th - bp)
    out["bilevel_fneqg"] = dict(train_stationarity=float(np.linalg.norm(gf)),
                                val_gradient_at_same_point=float(np.linalg.norm(gg)))

    # ---- certified pseudo-dimension lower bounds on real tuning classes
    def pdim_of(loss, xs, ts, alphas, kmax):
        return measure_pdim_pool(loss, xs, ts, alphas, nmax=kmax)

    rows = []
    for (p, d) in [(2, 8), (3, 12), (4, 16), (6, 24), (8, 32)]:
        A, b, Ap, bp, groups = make_instance(60, d, p, 100 + p * 7 + d)
        insts = [make_instance(60, d, p, 500 + p * 31 + d * 5 + k) for k in range(p + 4)]

        def loss_bi(alpha, inst):
            Ai, bi, Api, bpi, gi = inst
            av = np.zeros(d)
            for k, G in enumerate(gi):
                av[G] = alpha[k]
            t = ridge_theta(Ai, bi, av)
            return float(np.sum((Api @ t - bpi) ** 2))

        def loss_single(alpha, inst):
            Ai, bi, Api, bpi, gi = inst
            av = np.zeros(d)
            for k, G in enumerate(gi):
                av[G] = alpha[k]
            t = ridge_theta(Ai, bi, av)
            return float(np.sum((Ai @ t - bi) ** 2) + av @ (t * t))

        pool = np.exp(RNG.uniform(math.log(1e-3), math.log(1e3), (40000, p)))
        ts_bi = [np.median([loss_bi(a, inst) for a in pool[:400]]) for inst in insts]
        ts_si = [np.median([loss_single(a, inst) for a in pool[:400]]) for inst in insts]
        kb, _ = pdim_of(loss_bi, insts, ts_bi, pool, min(p + 3, 12))
        ks, _ = pdim_of(loss_single, insts, ts_si, pool, min(p + 3, 12))
        Lb = np.array([[loss_bi(a, x) for x in insts] for a in pool[:8000]])
        Sb = (Lb >= np.array(ts_bi)[None, :]).astype(np.int8)
        Ls = np.array([[loss_single(a, x) for x in insts] for a in pool[:8000]])
        Ss = (Ls >= np.array(ts_si)[None, :]).astype(np.int8)
        cells_bi = len(set(map(tuple, Sb.tolist())))
        cells_si = len(set(map(tuple, Ss.tolist())))
        Mtot, Dtot = 6 * d + 64, 4
        ub61 = stated_bound(p, (d, d), Mtot, Dtot)
        ub51 = stated_bound(p, (d,), 3 * d + 64 + d, 4)
        rows.append(dict(p=p, d=d, pdim_lb_bilevel=kb, pdim_lb_single=ks,
                         sign_cells_bilevel=cells_bi, sign_cells_single=cells_si,
                         cell_ratio=cells_bi / max(cells_si, 1),
                         thm61_bound=ub61, thm51_bound=ub51,
                         lb_under_bound=bool(kb <= ub61)))
        print("E p=%d d=%d pdim_lb bi=%d single=%d cells %d/%d bound61=%.1f"
              % (p, d, kb, ks, cells_bi, cells_si, ub61), flush=True)
    out["pdim_rows"] = rows

    # ---- precondition control: a non-semi-algebraic (sinusoidal) class
    xs = [np.array([1.0])] * 12
    pool1 = RNG.uniform(0.0, 200.0, (200000, 1))
    def sinloss(a, x):
        return float(np.sin(a[0] * (1.0 + 0.37 * x[0])))
    freqs = [np.array([1.0 + 0.31 * k]) for k in range(12)]
    def sinloss2(a, x):
        return float(np.sin(a[0] * x[0]))
    k_sin, _ = measure_pdim_pool(sinloss2, freqs, [0.0] * 12, pool1, nmax=12)
    def linloss(a, x):
        return float(a[0] * x[0])
    k_lin, _ = measure_pdim_pool(linloss, freqs, [0.0] * 12, pool1, nmax=12)
    out["precondition_control"] = dict(
        p=1, sinusoidal_pdim_lb=k_sin, affine_pdim_lb=k_lin,
        note="a single-parameter semi-algebraic class has Pdim 1; sin(alpha x) is not "
             "semi-algebraic and shatters far more points, so Theorem 4.1's polynomial-FOL "
             "hypothesis is doing real work")
    return out


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "de"
    out = {}
    if "d" in which:
        t = time.time(); out["D"] = section_D(); out["D"]["secs"] = round(time.time() - t, 1)
        print("D", json.dumps(out["D"])[:1500], flush=True)
    if "e" in which:
        t = time.time(); out["E"] = section_E(); out["E"]["secs"] = round(time.time() - t, 1)
        print("E", json.dumps(out["E"])[:1200], flush=True)
    with open("evidence_%s.json" % which, "w") as f:
        json.dump(out, f, indent=1)