scripts/blockfam.py

"""
Block-decomposable MILP family with a *properly normalised* number of coupling constraints.

Motivation.  In the flat family of common.py the number of decision variables is fixed while
s grows, so the optimal multipliers themselves shrink like 1/s and the measured s-exponent of
the excess risk is confounded.  Here the instance GROWS with s exactly the way the paper
motivates Lagrangian relaxation (Section 1 / Appendix A): s otherwise-independent
sub-problems linked by s coupling rows.

Instance with s blocks:
  * block B has p binary variables with a local constraint C_B x_B >= d_B  (this is the
    "Cx >= d" part, block diagonal -> the sub-problems are independent given pi);
  * coupling row k involves the variables of blocks k and (k+1) mod s, so the rows really
    couple the blocks:  A x >= b  with s rows.

Key identity used for exact, fast evaluation (this is *why* Lagrangian relaxation is used):
    u(pi,P) = pi'b + sum_B min_{x_B in X_B} (c_B - A_B' pi)' x_B
so the dual value decomposes into s independent tiny sub-problems for any fixed pi.  We
enumerate each block's feasible set exhaustively, so u and its subgradient are EXACT.
"""

import itertools
import numpy as np
from scipy.optimize import linprog
import scipy.sparse as sp


def gen_block_instance(rng, s, p=4, n_local=1):
    """Returns per-block arrays (a, w_prev, w_self) plus b, and the exact Assumption-4.1 B."""
    Xall = np.array(list(itertools.product([0, 1], repeat=p)), float)
    blocks = []
    Aself, Aprev = [], []  # coupling coefficients
    for B in range(s):
        Cl = rng.uniform(0.2, 1.0, size=(n_local, p))
        dl = rng.uniform(0.1, 0.5, size=n_local) * p * 0.5
        Xf = Xall[np.all(Xall @ Cl.T >= dl - 1e-12, axis=1)]
        if Xf.shape[0] == 0:
            return None
        c = rng.uniform(0.2, 1.2, size=p)
        a_self = rng.uniform(0.0, 1.0, size=p)  # row B, block B columns
        a_prev = rng.uniform(0.0, 1.0, size=p)  # row B-1, block B columns
        blocks.append(
            {
                "X": Xf,
                "a": Xf @ c,
                "w_self": -(Xf @ a_self),
                "w_prev": -(Xf @ a_prev),
                "c": c,
                "Cl": Cl,
                "dl": dl,
            }
        )
        Aself.append(a_self)
        Aprev.append(a_prev)
    # b_k : 0.6-quantile of (A x)_k over random feasible x, so the coupling rows bind
    bs = np.empty(s)
    Bconst = 0.0
    for k in range(s):
        v1 = -blocks[k]["w_self"]  # (A_k x_k) contributions
        v2 = -blocks[(k + 1) % s]["w_prev"]
        tot = (v1[:, None] + v2[None, :]).ravel()
        bs[k] = float(np.quantile(tot, 0.6))
        Bconst = max(Bconst, float(np.max(np.abs(tot))), abs(bs[k]))
    # keep the raw (c, A, C, d) so the decomposed oracle can be cross-checked with HiGHS
    n = s * p
    cfull = np.zeros(n)
    Afull = np.zeros((s, n))
    Cfull = np.zeros((s * n_local, n))
    dfull = np.zeros(s * n_local)
    for B in range(s):
        sl = slice(B * p, (B + 1) * p)
        cfull[sl] = blocks[B]["c"]
        Afull[B, sl] += Aself[B]
        Afull[(B - 1) % s, sl] += Aprev[B]
        Cfull[B * n_local : (B + 1) * n_local, sl] = blocks[B]["Cl"]
        dfull[B * n_local : (B + 1) * n_local] = blocks[B]["dl"]
    return {
        "blocks": blocks,
        "b": bs,
        "s": s,
        "B": Bconst,
        "p": p,
        "c": cfull,
        "A": Afull,
        "C": Cfull,
        "d": dfull,
    }


def u_highs_block(inst, pi):
    """Independent u(pi,P) for a block instance via scipy.optimize.milp (HiGHS)."""
    from scipy.optimize import milp, LinearConstraint, Bounds

    n = inst["c"].shape[0]
    obj = inst["c"] - inst["A"].T @ pi
    cons = LinearConstraint(inst["C"], inst["d"], np.inf)
    r = milp(
        c=obj,
        constraints=cons,
        integrality=np.ones(n),
        bounds=Bounds(np.zeros(n), np.ones(n)),
    )
    assert r.success, r.message
    return float(r.fun + pi @ inst["b"])


class BlockPool:
    """Finite population D = Uniform{P_1..P_M} of block instances; exact population risks."""

    def __init__(self, insts, pimax):
        self.M = len(insts)
        self.s = insts[0]["s"]
        self.pimax = float(pimax)
        K = max(len(bl["a"]) for it in insts for bl in it["blocks"])
        self.K = K
        s = self.s
        self.AA = np.full((self.M, s, K), 1e18)
        self.W1 = np.zeros((self.M, s, K))  # multiplies pi[(B-1) % s]
        self.W2 = np.zeros((self.M, s, K))  # multiplies pi[B]
        self.bb = np.zeros((self.M, s))
        for i, it in enumerate(insts):
            self.bb[i] = it["b"]
            for B, bl in enumerate(it["blocks"]):
                k = len(bl["a"])
                self.AA[i, B, :k] = bl["a"]
                self.W1[i, B, :k] = bl["w_prev"]
                self.W2[i, B, :k] = bl["w_self"]
        self.prev = (np.arange(s) - 1) % s
        self.nxt = (np.arange(s) + 1) % s
        self.B = float(max(it["B"] for it in insts))

    def _block_vals(self, pi, idx):
        AA = self.AA[idx]
        W1 = self.W1[idx]
        W2 = self.W2[idx]
        return AA + W1 * pi[self.prev][None, :, None] + W2 * pi[None, :, None]

    def u_all(self, pi, idx=None):
        idx = np.arange(self.M) if idx is None else idx
        vals = self._block_vals(pi, idx)
        j = np.argmin(vals, axis=2)  # (n, s)
        mins = np.take_along_axis(vals, j[:, :, None], axis=2)[:, :, 0]
        return self.bb[idx] @ pi + mins.sum(axis=1), j

    def F(self, pi, idx=None, w=None):
        v, _ = self.u_all(pi, idx)
        return float(v.mean()) if w is None else float(v @ w)

    def subgrad(self, pi, idx):
        """g_k = b_k - (A x*)_k = b_k + W2[i,k,j_k] + W1[i,k+1,j_{k+1}]"""
        _, j = self.u_all(pi, idx)
        n = len(idx)
        r = np.arange(n)[:, None]
        cols = np.arange(self.s)[None, :]
        w2 = self.W2[idx][r, cols, j]
        jn = j[:, self.nxt]
        w1 = self.W1[idx][r, self.nxt[None, :], jn]
        return self.bb[idx] + w2 + w1


def maximize_weighted_block(pool, idx, w, tol=1e-9, max_iter=120):
    """Exact max_{pi in [0,pimax]^s} sum_i w_i u(pi,P_i) by Kelley cutting planes.
    LP variables: [t_{i,B} (n*s), pi (s)]."""
    n, s = len(idx), pool.s
    nt = n * s
    pimax = pool.pimax
    AA = pool.AA[idx]
    W1 = pool.W1[idx]
    W2 = pool.W2[idx]
    bb = pool.bb[idx]
    cuts = [[[] for _ in range(s)] for _ in range(n)]
    pi = np.full(s, pimax / 2.0)
    _, j = pool.u_all(pi, idx)
    for i in range(n):
        for B in range(s):
            cuts[i][B].append(int(j[i, B]))
    cobj = np.concatenate([-np.repeat(w, s), -(w @ bb)])
    bounds = [(None, None)] * nt + [(0.0, pimax)] * s
    for it in range(max_iter):
        ri, rj, dat, rhs = [], [], [], []
        r = 0
        for i in range(n):
            for B in range(s):
                pB = pool.prev[B]
                for jj in cuts[i][B]:
                    ri.append(r)
                    rj.append(i * s + B)
                    dat.append(1.0)
                    ri.append(r)
                    rj.append(nt + pB)
                    dat.append(-W1[i, B, jj])
                    ri.append(r)
                    rj.append(nt + B)
                    dat.append(-W2[i, B, jj])
                    rhs.append(AA[i, B, jj])
                    r += 1
        Aub = sp.csr_matrix((dat, (ri, rj)), shape=(r, nt + s))
        res = linprog(cobj, A_ub=Aub, b_ub=np.array(rhs), bounds=bounds, method="highs")
        assert res.success, res.message
        pi = res.x[nt:]
        t = res.x[:nt].reshape(n, s)
        vals = pool._block_vals(pi, idx)
        jj = np.argmin(vals, axis=2)
        mins = np.take_along_axis(vals, jj[:, :, None], axis=2)[:, :, 0]
        viol = np.argwhere(t - mins > tol)
        added = 0
        for i, B in viol:
            if int(jj[i, B]) not in cuts[i][B]:
                cuts[i][B].append(int(jj[i, B]))
                added += 1
        if added == 0:
            u = bb @ pi + mins.sum(axis=1)
            return pi, float(u @ w), it + 1
    u = bb @ pi + mins.sum(axis=1)
    return pi, float(u @ w), max_iter


def erm_block(pool, sample_idx):
    uniq, cnt = np.unique(sample_idx, return_counts=True)
    return maximize_weighted_block(pool, uniq, cnt / cnt.sum())


def sga_block(pool, stream_idx, eta, average=True):
    s = pool.s
    pi = np.zeros(s)
    acc = np.zeros(s)
    for i in stream_idx:
        acc += pi
        g = pool.subgrad(pi, np.array([i]))[0]
        pi = np.clip(pi + eta * g, 0.0, pool.pimax)
    return acc / len(stream_idx) if average else pi