scripts/common.py

"""
Common machinery for the reproduction of
  "Provably Data-driven Lagrangian Relaxation for Mixed Integer Linear Programming"
  (arXiv 2605.19052, OpenReview OwLuqetJuB).

Everything is CPU-only and uses scipy / HiGHS.

Design
------
A MILP instance is  P = (c, A, b, C, d) with

    OPT(P) = min c'x  s.t.  x in R^m_+ x {0,1}^p,  Ax >= b (s coupling rows),  Cx >= d
    u(pi,P)= min c'x + pi'(b - Ax)  s.t. x in X = {x in R^m_+ x {0,1}^p : Cx >= d}

We use *block-decomposable local constraints*: C is block diagonal in
(y continuous, z binary), i.e. C_y y >= d_y and C_z z >= d_z, plus a box 0 <= y <= ymax.
This is exactly the structure the paper motivates (independent sub-problems linked
only by the s coupling rows Ax >= b).

Because the objective c'x + pi'(b-Ax) is linear in x, its minimum over the (non-convex)
set X is attained on the finite set

    Cand(P) = vert(Y) x {z in {0,1}^p : C_z z >= d_z},

where vert(Y) is the vertex set of the (bounded) polytope Y.  We enumerate that finite
set once per instance, so u(pi,P) and its subgradient become exact O(1) linear algebra
for *any* pi -- no solver in the inner loop.  `verify_against_highs` cross-checks the
enumeration against `scipy.optimize.milp` (HiGHS) on random (pi,P) pairs.

For the finite candidate representation we store, per instance i and candidate j,
    a[i,j]   = c' x_j                     (constant part)
    G[i,j,:] = b - A x_j                  (subgradient / linear part)
so that   u(pi,P_i) = min_j ( a[i,j] + G[i,j,:] @ pi ).
"""

import itertools
import numpy as np
from scipy.optimize import linprog, milp, LinearConstraint, Bounds
import scipy.sparse as sp

# --------------------------------------------------------------------------------------
# vertex enumeration of a small bounded polytope Y = {y : 0 <= y <= ymax, Cy y >= dy}
# --------------------------------------------------------------------------------------


def polytope_vertices(Cy, dy, ymax, tol=1e-9):
    """Brute-force vertex enumeration of {y in R^m : Cy y >= dy, 0 <= y <= ymax}."""
    m = ymax.shape[0]
    # write everything as  H y <= h
    H = np.vstack([-Cy, -np.eye(m), np.eye(m)])
    h = np.concatenate([-dy, np.zeros(m), ymax])
    verts = []
    n = H.shape[0]
    for idx in itertools.combinations(range(n), m):
        Hs, hs = H[list(idx)], h[list(idx)]
        if abs(np.linalg.det(Hs)) < 1e-8:
            continue
        y = np.linalg.solve(Hs, hs)
        if np.all(H @ y <= h + 1e-7):
            verts.append(y)
    if not verts:
        return np.zeros((0, m))
    V = np.unique(np.round(np.array(verts), 9), axis=0)
    return V


# --------------------------------------------------------------------------------------
# instance generation
# --------------------------------------------------------------------------------------


def gen_instance(rng, s, m=2, p=6, n_local_y=2, n_local_z=2, ymax_val=1.0):
    """Draw one random MILP instance with s coupling constraints. Returns a dict."""
    # ---- local (sub-problem) constraints -------------------------------------------
    Cy = rng.uniform(0.2, 1.0, size=(n_local_y, m))
    dy = rng.uniform(0.1, 0.5, size=n_local_y) * m * 0.5
    ymax = np.full(m, ymax_val)
    V = polytope_vertices(Cy, dy, ymax)
    if V.shape[0] == 0:
        return None

    Cz = rng.uniform(0.2, 1.0, size=(n_local_z, p))
    dz = rng.uniform(0.1, 0.4, size=n_local_z) * p * 0.5
    Zall = np.array(list(itertools.product([0, 1], repeat=p)), dtype=float)
    Zfeas = Zall[np.all(Zall @ Cz.T >= dz - 1e-12, axis=1)]
    if Zfeas.shape[0] == 0:
        return None

    # ---- objective and coupling constraints -----------------------------------------
    c = rng.uniform(0.2, 1.2, size=m + p)
    A = rng.uniform(0.0, 1.0, size=(s, m + p))
    # candidate set
    ny, nz = V.shape[0], Zfeas.shape[0]
    X = np.concatenate(
        [np.repeat(V, nz, axis=0), np.tile(Zfeas, (ny, 1))], axis=1
    )  # (ny*nz, m+p)
    AX = X @ A.T  # (ncand, s)
    # choose b so that the coupling constraints are active-ish (b = median of AX rows)
    b = np.quantile(AX, 0.6, axis=0)

    inst = dict(
        c=c,
        A=A,
        b=b,
        Cy=Cy,
        dy=dy,
        Cz=Cz,
        dz=dz,
        ymax=ymax,
        m=m,
        p=p,
        s=s,
        X=X,
        a=X @ c,
        G=b[None, :] - AX,
    )
    # exact B of Assumption 4.1 for this instance
    inst["B"] = float(max(np.max(np.abs(b)), np.max(np.abs(AX))))
    return inst


def u_of(inst, pi):
    """Exact u(pi,P) and a subgradient g = b - A x*(pi,P) via the finite candidate set."""
    vals = inst["a"] + inst["G"] @ pi
    j = int(np.argmin(vals))
    return float(vals[j]), inst["G"][j].copy(), inst["X"][j].copy()


def opt_of(inst):
    """OPT(P): min c'x over candidates that additionally satisfy A x >= b."""
    ok = np.all(inst["G"] <= 1e-12, axis=1)  # b - Ax <= 0  <=>  Ax >= b
    if not np.any(ok):
        return np.inf
    return float(np.min(inst["a"][ok]))


# --------------------------------------------------------------------------------------
# HiGHS cross-check
# --------------------------------------------------------------------------------------


def u_highs(inst, pi):
    """u(pi,P) computed independently with scipy.optimize.milp (HiGHS)."""
    m, p = inst["m"], inst["p"]
    n = m + p
    obj = inst["c"] - inst["A"].T @ pi
    const = float(pi @ inst["b"])
    # local constraints only:  Cy y >= dy ,  Cz z >= dz
    rows = []
    lb = []
    Cy, Cz = inst["Cy"], inst["Cz"]
    for k in range(Cy.shape[0]):
        r = np.zeros(n)
        r[:m] = Cy[k]
        rows.append(r)
        lb.append(inst["dy"][k])
    for k in range(Cz.shape[0]):
        r = np.zeros(n)
        r[m:] = Cz[k]
        rows.append(r)
        lb.append(inst["dz"][k])
    A_ub = np.array(rows)
    cons = LinearConstraint(A_ub, np.array(lb), np.inf)
    integrality = np.concatenate([np.zeros(m), np.ones(p)])
    bounds = Bounds(np.zeros(n), np.concatenate([inst["ymax"], np.ones(p)]))
    res = milp(c=obj, constraints=cons, integrality=integrality, bounds=bounds)
    assert res.success, res.message
    return float(res.fun + const), res.x


def opt_highs(inst):
    """OPT(P) with HiGHS (coupling constraints included)."""
    m, p = inst["m"], inst["p"]
    n = m + p
    rows, lb = [], []
    Cy, Cz = inst["Cy"], inst["Cz"]
    for k in range(Cy.shape[0]):
        r = np.zeros(n)
        r[:m] = Cy[k]
        rows.append(r)
        lb.append(inst["dy"][k])
    for k in range(Cz.shape[0]):
        r = np.zeros(n)
        r[m:] = Cz[k]
        rows.append(r)
        lb.append(inst["dz"][k])
    for k in range(inst["s"]):
        rows.append(inst["A"][k])
        lb.append(inst["b"][k])
    cons = LinearConstraint(np.array(rows), np.array(lb), np.inf)
    integrality = np.concatenate([np.zeros(m), np.ones(p)])
    bounds = Bounds(np.zeros(n), np.concatenate([inst["ymax"], np.ones(p)]))
    res = milp(c=inst["c"], constraints=cons, integrality=integrality, bounds=bounds)
    if not res.success:
        return np.inf
    return float(res.fun)


# --------------------------------------------------------------------------------------
# pooled representation of a whole population of instances (padded to a common K)
# --------------------------------------------------------------------------------------


class Pool:
    """A finite population D = Uniform{P_1,...,P_M}. Exact population risks."""

    def __init__(self, insts, pimax):
        self.M = len(insts)
        self.s = insts[0]["s"]
        self.pimax = float(pimax)
        K = max(i["a"].shape[0] for i in insts)
        self.K = K
        self.a = np.full((self.M, K), 1e18)
        self.G = np.zeros((self.M, K, self.s))
        for i, inst in enumerate(insts):
            k = inst["a"].shape[0]
            self.a[i, :k] = inst["a"]
            self.G[i, :k] = inst["G"]
        self.B = float(max(i["B"] for i in insts))
        self.insts = insts

    def u_all(self, pi, idx=None):
        """Vector of u(pi,P_i) (and argmin candidate) for the given instance indices."""
        a = self.a if idx is None else self.a[idx]
        G = self.G if idx is None else self.G[idx]
        vals = a + G @ pi
        j = np.argmin(vals, axis=1)
        return vals[np.arange(vals.shape[0]), j], j

    def grad_all(self, pi, idx=None):
        vals, j = self.u_all(pi, idx)
        G = self.G if idx is None else self.G[idx]
        g = G[np.arange(G.shape[0]), j]
        return vals, g

    def F(self, pi, idx=None, w=None):
        """Weighted average utility (w defaults to uniform over idx)."""
        vals, _ = self.u_all(pi, idx)
        if w is None:
            return float(vals.mean())
        return float(vals @ w)


# --------------------------------------------------------------------------------------
# exact maximisation of a weighted average of concave piecewise-linear u's  (cutting plane)
# --------------------------------------------------------------------------------------


def maximize_weighted(pool, idx, w, tol=1e-10, max_iter=200):
    """Exact  max_{pi in [0,pimax]^s}  sum_i w_i u(pi, P_{idx[i]})  by Kelley cutting planes.

    The objective is a weighted sum of concave piecewise-linear functions, hence concave
    piecewise linear; Kelley's method with an LP master terminates finitely and exactly.
    Returns (pi_star, value, n_iters).
    """
    n = len(idx)
    s = pool.s
    pimax = pool.pimax
    a = pool.a[idx]
    G = pool.G[idx]

    # active cut lists
    cuts = [[] for _ in range(n)]
    pi = np.full(s, pimax / 2.0)
    vals, j = pool.u_all(pi, idx)
    for i in range(n):
        cuts[i].append(j[i])

    # LP variables: [t_1..t_n, pi_1..pi_s]; maximise sum w_i t_i  ->  minimise -sum w_i t_i
    cobj = np.concatenate([-w, np.zeros(s)])
    bounds = [(None, None)] * n + [(0.0, pimax)] * s

    for it in range(max_iter):
        rows_i, rows_j, data, rhs = [], [], [], []
        r = 0
        for i in range(n):
            for jj in cuts[i]:
                # t_i - G[i,jj]@pi <= a[i,jj]
                rows_i.append(r)
                rows_j.append(i)
                data.append(1.0)
                for k in range(s):
                    if G[i, jj, k] != 0.0:
                        rows_i.append(r)
                        rows_j.append(n + k)
                        data.append(-G[i, jj, k])
                rhs.append(a[i, jj])
                r += 1
        Aub = sp.csr_matrix((data, (rows_i, rows_j)), shape=(r, n + s))
        res = linprog(cobj, A_ub=Aub, b_ub=np.array(rhs), bounds=bounds, method="highs")
        assert res.success, res.message
        pi = res.x[n:]
        t = res.x[:n]
        vals, j = pool.u_all(pi, idx)
        viol = t - vals  # >0 means the master over-estimates u_i
        added = 0
        for i in np.where(viol > tol)[0]:
            if j[i] not in cuts[i]:
                cuts[i].append(int(j[i]))
                added += 1
        if added == 0:
            return pi, float(vals @ w), it + 1
    return pi, float(vals @ w), max_iter


def erm(pool, sample_idx):
    """Exact ERM maximiser on a multiset `sample_idx` of pool instances."""
    uniq, cnt = np.unique(sample_idx, return_counts=True)
    w = cnt / cnt.sum()
    return maximize_weighted(pool, uniq, w)


def population_opt(pool):
    idx = np.arange(pool.M)
    w = np.full(pool.M, 1.0 / pool.M)
    return maximize_weighted(pool, idx, w)


# --------------------------------------------------------------------------------------
# Algorithm 1 : stochastic (sub)gradient ascent with iterate averaging
# --------------------------------------------------------------------------------------


def sga(pool, stream_idx, eta, average=True, pi0=None):
    """Algorithm 1 of the paper. Returns the averaged (or last) iterate."""
    s = pool.s
    pimax = pool.pimax
    pi = np.zeros(s) if pi0 is None else pi0.copy()
    acc = np.zeros(s)
    for t, i in enumerate(stream_idx):
        acc += pi
        vals = pool.a[i] + pool.G[i] @ pi
        j = int(np.argmin(vals))
        g = pool.G[i, j]  # unbiased stochastic supergradient
        pi = np.clip(pi + eta * g, 0.0, pimax)  # Proj_Pi ( pi + eta g )
    if average:
        return acc / len(stream_idx)
    return pi


# --------------------------------------------------------------------------------------
# log-log fitting helper
# --------------------------------------------------------------------------------------


def loglog_fit(x, y):
    x = np.asarray(x, float)
    y = np.asarray(y, float)
    ok = (x > 0) & (y > 0)
    A = np.vstack([np.log(x[ok]), np.ones(ok.sum())]).T
    coef, res, *_ = np.linalg.lstsq(A, np.log(y[ok]), rcond=None)
    pred = A @ coef
    ss_res = float(np.sum((np.log(y[ok]) - pred) ** 2))
    ss_tot = float(np.sum((np.log(y[ok]) - np.log(y[ok]).mean()) ** 2))
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    return float(coef[0]), float(np.exp(coef[1])), r2


def dump_json(path, obj):
    import json

    def default(o):
        if isinstance(o, (np.floating,)):
            return float(o)
        if isinstance(o, (np.integer,)):
            return int(o)
        if isinstance(o, np.ndarray):
            return o.tolist()
        raise TypeError(str(type(o)))

    with open(path, "w") as f:
        json.dump(obj, f, indent=2, default=default)
    print("wrote", path)