code/evidence3.py

"""Section F: certified MAXIMUM shattered-set size for real tuning classes at
small p, by exhaustive enumeration of the sign-pattern arrangement in alpha.

For p <= 2 the parameter space is low dimensional, so the arrangement induced
by the N boundary curves { alpha : l_alpha(x_i) = t_i } can be enumerated on a
fine mesh.  The number of realised sign patterns is then an exact count (up to
mesh resolution), which gives an upper bound on the shattered-set size, not
merely a lower bound.  This is the experiment the "you only show lower bounds"
objection asks for.

Section G: greedy instance selection to certify larger pseudo-dimension lower
bounds on the same real classes at higher p.
"""
import itertools, json, math, sys, time
import numpy as np
from evidence import stated_bound

RNG = np.random.default_rng(90210)


def ridge_loss_matrix(insts, alphas, p, d, groups_list, bilevel=True):
    """L[a, i] = loss of instance i at hyperparameter alphas[a]."""
    out = np.empty((alphas.shape[0], len(insts)))
    if p == 1:
        # theta(alpha) = V diag(1/(lam+alpha)) V' A'b  -- vectorised over alpha
        for i, (A, b, Ap, bp, gs) in enumerate(insts):
            AtA = A.T @ A
            lam, V = np.linalg.eigh(AtA)
            c = V.T @ (A.T @ b)
            al = alphas[:, 0][:, None]
            TH = ((c[None, :] / (lam[None, :] + al)) @ V.T)          # (nalpha, d)
            if bilevel:
                R = TH @ Ap.T - bp[None, :]
                out[:, i] = np.einsum("ij,ij->i", R, R)
            else:
                R = TH @ A.T - b[None, :]
                out[:, i] = np.einsum("ij,ij->i", R, R) + alphas[:, 0] * np.einsum("ij,ij->i", TH, TH)
        return out
    for i, (A, b, Ap, bp, gs) in enumerate(insts):
        AtA = A.T @ A
        Atb = A.T @ b
        for a in range(alphas.shape[0]):
            av = np.zeros(d)
            for k, G in enumerate(gs):
                av[G] = alphas[a, k]
            th = np.linalg.solve(AtA + np.diag(av), Atb)
            if bilevel:
                r = Ap @ th - bp
                out[a, i] = float(r @ r)
            else:
                r = A @ th - b
                out[a, i] = float(r @ r + av @ (th * th))
    return out


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 section_F():
    """Exhaustive arrangement enumeration at p = 1 and p = 2."""
    rows = []
    for p, d, mesh in [(1, 6, 400000), (1, 10, 400000), (2, 6, 420), (2, 10, 420)]:
        N = 12
        insts = [make_instance(40, d, p, 7000 + 13 * p + 3 * d + k) for k in range(N)]
        if p == 1:
            alphas = np.exp(np.linspace(math.log(1e-6), math.log(1e6), mesh))[:, None]
        else:
            g = np.exp(np.linspace(math.log(1e-5), math.log(1e5), mesh))
            alphas = np.array(np.meshgrid(g, g)).reshape(2, -1).T
        L = ridge_loss_matrix(insts, alphas, p, d, None, bilevel=True)
        ts = np.median(L, axis=0)
        S = (L >= ts[None, :]).astype(np.int8)
        # exact realised-pattern count for every prefix size
        detail = []
        maxN = 0
        for k in range(1, N + 1):
            cnt = len(set(map(tuple, S[:, :k].tolist())))
            detail.append(dict(k=k, realised=cnt, needed=2 ** k))
            if cnt == 2 ** k:
                maxN = k
        # total sign changes along the mesh (p=1 gives the exact cell count)
        flips = int(np.sum(np.abs(np.diff(S, axis=0)))) if p == 1 else None
        bound = stated_bound(p, (d, d), 6 * d + 64, 4)
        rows.append(dict(p=p, d=d, mesh=alphas.shape[0], instances=N,
                         max_shattered=maxN, detail=detail,
                         total_sign_changes=flips, thm61_bound=bound))
        print("F p=%d d=%d mesh=%d  max shattered = %d   patterns at k=%d: %d/%d"
              % (p, d, alphas.shape[0], maxN, min(maxN + 1, N),
                 detail[min(maxN, N - 1)]["realised"], 2 ** min(maxN + 1, N)), flush=True)
    return dict(rows=rows)


def section_G():
    """Greedy instance selection to certify larger Pdim lower bounds."""
    rows = []
    for p, d in [(2, 8), (3, 12), (4, 16), (6, 24), (8, 32)]:
        pool = np.exp(RNG.uniform(math.log(1e-4), math.log(1e4), (60000, p)))
        cand = [make_instance(50, d, p, 30000 + 71 * p + 5 * d + k) for k in range(30)]
        L = ridge_loss_matrix(cand, pool, p, d, None, bilevel=True)
        Ls = ridge_loss_matrix(cand, pool, p, d, None, bilevel=False)
        best = 0; chosen = []
        S_all = None
        for _ in range(min(2 * p + 4, 14)):
            bestk, bestcnt, bestS = None, -1, None
            for j in range(len(cand)):
                if j in chosen:
                    continue
                for q in [0.25, 0.5, 0.75]:
                    t = float(np.quantile(L[:, j], q))
                    col = (L[:, j] >= t).astype(np.int8)[:, None]
                    Snew = col if S_all is None else np.hstack([S_all, col])
                    cnt = len(set(map(tuple, Snew.tolist())))
                    if cnt > bestcnt:
                        bestcnt, bestk, bestS = cnt, (j, q), Snew
            if bestcnt < 2 ** (len(chosen) + 1):
                break
            chosen.append(bestk[0]); S_all = bestS; best = len(chosen)
        # same for the single-level control
        best_s = 0; chosen_s = []; S_s = None
        for _ in range(min(2 * p + 4, 14)):
            bestcnt, bestk, bestS = -1, None, None
            for j in range(len(cand)):
                if j in chosen_s:
                    continue
                for q in [0.25, 0.5, 0.75]:
                    t = float(np.quantile(Ls[:, j], q))
                    col = (Ls[:, j] >= t).astype(np.int8)[:, None]
                    Snew = col if S_s is None else np.hstack([S_s, col])
                    cnt = len(set(map(tuple, Snew.tolist())))
                    if cnt > bestcnt:
                        bestcnt, bestk, bestS = cnt, (j, q), Snew
            if bestcnt < 2 ** (len(chosen_s) + 1):
                break
            chosen_s.append(bestk[0]); S_s = bestS; best_s = len(chosen_s)
        ub61 = stated_bound(p, (d, d), 6 * d + 64, 4)
        ub51 = stated_bound(p, (d,), 3 * d + 64 + d, 4)
        rows.append(dict(p=p, d=d, pdim_lb_bilevel=best, pdim_lb_single=best_s,
                         patterns=2 ** best, thm61_bound=ub61, thm51_bound=ub51,
                         under_bound=bool(best <= ub61)))
        print("G p=%d d=%d  bilevel Pdim >= %d (%d/%d patterns)  single-level >= %d"
              % (p, d, best, 2 ** best, 2 ** best, best_s), flush=True)
    return dict(rows=rows)


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "fg"
    out = {}
    if "f" in which:
        t = time.time(); out["F"] = section_F(); out["F"]["secs"] = round(time.time() - t, 1)
    if "g" in which:
        t = time.time(); out["G"] = section_G(); out["G"]["secs"] = round(time.time() - t, 1)
    with open("evidence_%s.json" % which, "w") as f:
        json.dump(out, f, indent=1)
    print(json.dumps(out)[:1500])