scripts/claim1b_blocks.py

"""
CLAIM 1 / CLAIM 3, part B -- the s-exponent on a properly s-normalised MILP family.

The flat family of common.py holds the instance size fixed while s grows, so the optimal
multipliers shrink like 1/s and the s-exponent of the excess risk is confounded.  Here we use
the block-decomposable family of blockfam.py: the instance GROWS with s (s sub-problems, s
coupling rows, each row linking two adjacent blocks), so B, pi_max and the scale of pi* are
all s-independent -- exactly the regime Theorems 5.5/5.6/5.12 describe.

Everything is exact: u(pi,P) by exhaustive per-block enumeration (cross-checked against
scipy.optimize.milp / HiGHS), the population optimum and the ERM optimum by Kelley cutting
planes on HiGHS, and the population D is the uniform distribution on an explicit pool.
"""

import sys, os, time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np
from blockfam import (
    gen_block_instance,
    BlockPool,
    maximize_weighted_block,
    erm_block,
    sga_block,
    u_highs_block,
)
from common import loglog_fit, dump_json

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED = 918273
PIMAX = 1.0
M_POOL = 1500
S_LIST = [2, 4, 8, 16, 32]
N_LIST = [16, 32, 64, 128, 256]
TRIALS = 20
t0 = time.time()
res = {
    "seed": SEED,
    "pimax": PIMAX,
    "M_pool": M_POOL,
    "trials": TRIALS,
    "S_LIST": S_LIST,
    "N_LIST": N_LIST,
    "family": "block-decomposable MILP, s blocks / s coupling rows",
}

cells, pinfo, xchk = [], {}, []
for s in S_LIST:
    r = np.random.default_rng(SEED + 13 * s)
    insts = []
    while len(insts) < M_POOL:
        it = gen_block_instance(r, s)
        if it is not None:
            insts.append(it)
    pool = BlockPool(insts, PIMAX)
    # independent HiGHS cross-check of the decomposed dual oracle
    for k in range(6):
        pi = r.uniform(0, PIMAX, s)
        xchk.append(
            float(
                abs(pool.u_all(pi, np.array([k]))[0][0] - u_highs_block(insts[k], pi))
            )
        )
    tp = time.time()
    pi_pop, v_pop, its = maximize_weighted_block(
        pool, np.arange(pool.M), np.full(pool.M, 1.0 / pool.M)
    )
    pinfo[s] = {
        "B": pool.B,
        "K": pool.K,
        "pop_opt_value": v_pop,
        "cuts": its,
        "pop_pi_mean": float(pi_pop.mean()),
        "pop_pi_min": float(pi_pop.min()),
        "pop_pi_max": float(pi_pop.max()),
        "seconds": time.time() - tp,
    }
    print(
        "s=%2d  B=%.3f  pop opt=%.4f  pi* in [%.3f,%.3f] mean %.3f  (%.1fs)"
        % (
            s,
            pool.B,
            v_pop,
            pi_pop.min(),
            pi_pop.max(),
            pi_pop.mean(),
            time.time() - tp,
        )
    )
    for N in N_LIST:
        eta = PIMAX / (2 * pool.B * np.sqrt(N))
        e_erm, e_sga, e_last = [], [], []
        tc = time.time()
        for _ in range(TRIALS):
            idx = r.integers(0, pool.M, N)
            pih, _, _ = erm_block(pool, idx)
            e_erm.append(v_pop - pool.F(pih))
            e_sga.append(v_pop - pool.F(sga_block(pool, idx, eta, True)))
            e_last.append(v_pop - pool.F(sga_block(pool, idx, eta, False)))
        c = {
            "s": s,
            "N": N,
            "B": pool.B,
            "erm_excess_mean": float(np.mean(e_erm)),
            "erm_excess_sem": float(np.std(e_erm, ddof=1) / np.sqrt(TRIALS)),
            "sga_excess_mean": float(np.mean(e_sga)),
            "sga_excess_sem": float(np.std(e_sga, ddof=1) / np.sqrt(TRIALS)),
            "sga_last_excess_mean": float(np.mean(e_last)),
            "thm55_constant_bound": float(
                12 * np.sqrt(np.pi) * pool.B * PIMAX * s**1.5 / np.sqrt(N)
            ),
            "thm512_constant_bound": float(2 * pool.B * PIMAX * s / np.sqrt(N)),
            "seconds": time.time() - tc,
        }
        c["erm_within_thm55"] = bool(c["erm_excess_mean"] <= c["thm55_constant_bound"])
        c["sga_within_thm512"] = bool(
            c["sga_excess_mean"] <= c["thm512_constant_bound"]
        )
        cells.append(c)
        print(
            "   N=%4d  ERM=%.5f (<=%.2f:%s)  SGA=%.5f (<=%.3f:%s)  SGAlast=%.5f  [%.1fs]"
            % (
                N,
                c["erm_excess_mean"],
                c["thm55_constant_bound"],
                c["erm_within_thm55"],
                c["sga_excess_mean"],
                c["thm512_constant_bound"],
                c["sga_within_thm512"],
                c["sga_last_excess_mean"],
                c["seconds"],
            )
        )

res["oracle_vs_highs_max_abs_diff"] = float(max(xchk))
res["pool_info"] = pinfo
res["cells"] = cells
res["fits"] = {
    "erm_N_exponent_per_s": {
        str(s): loglog_fit(
            [c["N"] for c in cells if c["s"] == s],
            [c["erm_excess_mean"] for c in cells if c["s"] == s],
        )
        for s in S_LIST
    },
    "sga_N_exponent_per_s": {
        str(s): loglog_fit(
            [c["N"] for c in cells if c["s"] == s],
            [c["sga_excess_mean"] for c in cells if c["s"] == s],
        )
        for s in S_LIST
    },
    "erm_s_exponent_per_N": {
        str(N): loglog_fit(
            [c["s"] for c in cells if c["N"] == N],
            [c["erm_excess_mean"] for c in cells if c["N"] == N],
        )
        for N in N_LIST
    },
    "sga_s_exponent_per_N": {
        str(N): loglog_fit(
            [c["s"] for c in cells if c["N"] == N],
            [c["sga_excess_mean"] for c in cells if c["N"] == N],
        )
        for N in N_LIST
    },
}
res["thm55_never_violated"] = bool(all(c["erm_within_thm55"] for c in cells))
res["thm512_never_violated"] = bool(all(c["sga_within_thm512"] for c in cells))
res["max_erm_over_thm55_bound"] = float(
    max(c["erm_excess_mean"] / c["thm55_constant_bound"] for c in cells)
)
res["max_sga_over_thm512_bound"] = float(
    max(c["sga_excess_mean"] / c["thm512_constant_bound"] for c in cells)
)
res["wall_time_s"] = time.time() - t0

print("\noracle vs HiGHS max |diff| =", res["oracle_vs_highs_max_abs_diff"])
print(
    "ERM N-exponents:",
    {k: round(v[0], 3) for k, v in res["fits"]["erm_N_exponent_per_s"].items()},
)
print(
    "SGA N-exponents:",
    {k: round(v[0], 3) for k, v in res["fits"]["sga_N_exponent_per_s"].items()},
)
print(
    "ERM s-exponents:",
    {k: round(v[0], 3) for k, v in res["fits"]["erm_s_exponent_per_N"].items()},
)
print(
    "SGA s-exponents:",
    {k: round(v[0], 3) for k, v in res["fits"]["sga_s_exponent_per_N"].items()},
)
dump_json(os.path.join(OUT, "block_sweep.json"), res)
print("done in", round(time.time() - t0, 1), "s")