code/run_apub_applications.py

#!/usr/bin/env python3
"""
run_apub_applications.py -- independent reproduction of Section 5 (+ Appendix D)
of "Minimizing Upper Confidence Bounds: A Data-Driven Framework for Stochastic
Programming" (arXiv 2403.08966), orid eXLcL70GXO, claim 6.

Runs the paper's OWN application instances at the paper's OWN dimensions and
Appendix-C parameters:

  A. two-stage product mix, RANDOM recourse, |I| = 20 products, |J| = 8
     departments, two-regime Gumbel-copula labor uncertainty  (Section 5.1-5.2)
  B. two-stage product mix, FIXED recourse, |I| = 4, |J| = 2, Gaussian-mixture
     gamma -- the paper's own Wasserstein-DRO comparison instance (Section 5.3)
  C. 10-product newsvendor, Case I and Case II                 (Appendix D)

Each replication draws a fresh training sample of size N, solves SAA, APUB-M at
several nominal levels (1 - alpha) and (arms B/C) Wasserstein-1 DRO at several
radii, then evaluates the true expected cost of every resulting first-stage
decision on a large INDEPENDENT test set drawn from the true distribution, and
records the coverage indicator beta(theta_hat, x_hat) = 1{theta_hat >= mu(x_hat)}
from eq. (beta).

Usage:  python3 code/run_apub_applications.py --arm A
        python3 code/run_apub_applications.py --arm BC
"""
import argparse
import json
import os
import sys
import time
import warnings

import numpy as np

warnings.filterwarnings("ignore", category=RuntimeWarning)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import apub_paper_models as P  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RESDIR = os.path.join(ROOT, "results")
T0 = time.time()

# nominal confidence levels (1 - alpha); 0.0 is the SAA model (Remark 3.1)
CONFS = [0.5, 0.8, 0.9, 0.95, 0.99]
TEST_N = 20000          # paper uses 5000; a larger test set tightens mu(x)


def log(msg):
    print(f"[{time.time() - T0:7.1f}s] {msg}", flush=True)


class Budget:
    """Wall-clock budget per experiment block.  The machine this ran on was
    heavily contended, so each (N, method) cell runs as many replications as fit
    in its budget (never fewer than MIN_REPS) and the actually-completed count is
    reported with every number."""
    MIN_REPS = 2

    def __init__(self, seconds):
        self.deadline = time.time() + seconds

    def stop(self, done):
        return done >= self.MIN_REPS and time.time() > self.deadline


def summarize(vals):
    v = np.asarray(vals, dtype=float)
    return {"mean": float(v.mean()), "std": float(v.std(ddof=1)) if v.size > 1 else 0.0,
            "min": float(v.min()), "max": float(v.max()), "n": int(v.size)}


# =============================================================================
# ARM A -- random-recourse product mix, |I| = 20, |J| = 8
# =============================================================================
def arm_a(reps_small, reps_large, M, budget_s):
    rng = np.random.default_rng(770301)
    test = P.sample_product_mix_xi(TEST_N, rng)

    def true_cost(x):
        """mu(x) = c'x + E[Q(x, xi)] on an independent test set.  The
        sum_j y_j <= h2 cap is suppressed exactly as Appendix C states; we also
        record the fraction of test scenarios at which it WOULD have bound."""
        cf, _ = P.recourse_pm_closed_form(x, test, h2_cap=False)
        _, feas = P.recourse_pm_closed_form(x, test, h2_cap=True)
        return float(P.PM_C @ x + cf.mean()), float(1.0 - feas.mean())

    out = {"model": "two-stage product mix, RANDOM recourse (Section 5.1-5.2)",
           "n_products": 20, "n_departments": 8, "M_bootstrap": M,
           "test_set_size": TEST_N,
           "h2_cap": "suppressed, per Appendix C ('Without loss, we suppress "
                     "the constraint on the capacity of total outsourced labor "
                     "by setting a large value for h2'); the rate at which it "
                     "would have bound out of sample is reported per method",
           "cells": []}

    plan = [(120, reps_small, CONFS, 0.55), (480, reps_large, [0.5, 0.9, 0.95], 0.45)]
    for N, reps, confs, share in plan:
        bud = Budget(budget_s * share)
        rec = {"N": N, "reps": reps,
               "SAA": {"true": [], "obj": [], "x": [], "h2_bind": []}}
        for c in confs:
            rec[f"APUB_{c}"] = {"true": [], "obj": [], "x": [], "h2_bind": []}
        for r in range(reps):
            xi = P.sample_product_mix_xi(N, rng)
            V = P.bootstrap_multiplicities(N, M, rng)
            for key, sol in [("SAA", P.solve_pm_random_recourse(
                    xi, 1.0, None, h2_cap=False))] + \
                    [(f"APUB_{c}", P.solve_pm_random_recourse(
                        xi, 1.0 - c, V, h2_cap=False)) for c in confs]:
                tc, bind = true_cost(sol["x"])
                rec[key]["true"].append(tc)
                rec[key]["obj"].append(sol["obj"])
                rec[key]["x"].append(sol["x"].tolist())
                rec[key]["h2_bind"].append(bind)
            log(f"  armA N={N} rep {r + 1}/{reps} done")
            if bud.stop(r + 1):
                log(f"  armA N={N} budget reached after {r + 1} reps")
                break
        out["cells"].append(finalize_cell(rec))
        log(f"armA N={N} complete")

    # the paper's "N=120 APUB beats N=240 SAA" claim needs an N=240 SAA arm
    saa240 = []
    for r in range(min(reps_small, 10)):
        xi = P.sample_product_mix_xi(240, rng)
        s = P.solve_pm_random_recourse(xi, 1.0, None, h2_cap=False)
        saa240.append(true_cost(s["x"])[0])
    out["SAA_N240_true_cost"] = summarize(saa240)
    log("armA N=240 SAA reference complete")
    return out


def finalize_cell(rec):
    cell = {"N": rec["N"], "reps": rec["reps"], "methods": {}}
    for k, v in rec.items():
        if not isinstance(v, dict) or "true" not in v:
            continue
        tr = np.array(v["true"])
        ob = np.array(v["obj"])
        cell["methods"][k] = {
            "true_out_of_sample_cost": summarize(tr),
            "model_optimal_value": summarize(ob),
            "coverage_probability": float((ob >= tr - 1e-9).mean()),
            "mean_solution": np.mean(np.array(v["x"]), axis=0).tolist(),
        }
        if "h2_bind" in v:
            cell["methods"][k]["h2_cap_would_bind_rate"] = \
                float(np.mean(v["h2_bind"]))
    return cell


# =============================================================================
# ARM B -- fixed-recourse product mix + Wasserstein DRO (Section 5.3)
# =============================================================================
DRO_EPS = [0.01, 0.05, 0.1, 0.3, 1.0, 3.0]


def arm_b(plan, M, budget_s):
    rng = np.random.default_rng(880412)
    test = P.sample_fixed_recourse_gamma(TEST_N, rng)

    def true_cost(x):
        return float(P.FR_C @ np.asarray(x) + P.fr_recourse(x, test).mean())

    out = {"model": "two-stage product mix, FIXED recourse (Section 5.3, the "
                    "paper's own WassDRO comparison instance)",
           "n_products": 4, "n_departments": 2, "M_bootstrap": M,
           "test_set_size": TEST_N,
           "dro": "exact Wasserstein-1 reformulation, l1 ground metric: "
                  "sup_{W1(P,Phat)<=eps} E_P[F(x,.)] = SAA(x) + eps*Lip(x) with "
                  "Lip(x) = (12/0.9)(sum(x)/4 + 500), DECISION-DEPENDENT",
           "cells": []}
    for N, reps, confs, share in plan:
        bud = Budget(budget_s * share)
        rec = {"N": N, "reps": reps, "SAA": {"true": [], "obj": [], "x": []}}
        for c in confs:
            rec[f"APUB_{c}"] = {"true": [], "obj": [], "x": []}
        for e in DRO_EPS:
            rec[f"WassDRO_eps{e}"] = {"true": [], "obj": [], "x": []}
        for r in range(reps):
            gam = P.sample_fixed_recourse_gamma(N, rng)
            V = P.bootstrap_multiplicities(N, M, rng)
            for key, sol in [("SAA", P.solve_fr(gam, 1.0, None))] + \
                    [(f"APUB_{c}", P.solve_fr(gam, 1.0 - c, V)) for c in confs] + \
                    [(f"WassDRO_eps{e}", P.solve_fr(gam, 1.0, None, dro_eps=e))
                     for e in DRO_EPS]:
                rec[key]["true"].append(true_cost(sol["x"]))
                rec[key]["obj"].append(sol["obj"])
                rec[key]["x"].append(sol["x"].tolist())
            log(f"  armB N={N} rep {r + 1}/{reps} done")
            if bud.stop(r + 1):
                log(f"  armB N={N} budget reached after {r + 1} reps")
                break
        out["cells"].append(finalize_cell(rec))
        log(f"armB N={N} complete")
    return out


# =============================================================================
# ARM C -- 10-product newsvendor (Appendix D)
# =============================================================================
def arm_c(plan, M, budget_s):
    out = {"model": "10-product newsvendor (Appendix D, params Appendix C.4)",
           "n_products": 10, "M_bootstrap": M, "test_set_size": TEST_N,
           "dro": "exact Wasserstein-1 reformulation: the Lipschitz modulus of "
                  "F(x,.) wrt xi is max(h,b)=9 for EVERY x, so the worst case is "
                  "SAA(x) + 9*eps, a CONSTANT shift -- Wasserstein DRO cannot "
                  "move the decision at all here, exactly the limitation the "
                  "paper cites from Mohajerin Esfahani & Kuhn (2018)",
           "cases": {}}
    for case in (1, 2):
        rng = np.random.default_rng(990523 + case)
        test = P.sample_newsvendor(TEST_N, rng, case)

        def true_cost(x):
            return float(P.nv_cost(x, test).mean())

        cells = []
        for N, reps, confs, share in plan:
            bud = Budget(budget_s * share)
            rec = {"N": N, "reps": reps, "SAA": {"true": [], "obj": [], "x": []}}
            for c in confs:
                rec[f"APUB_{c}"] = {"true": [], "obj": [], "x": []}
            for e in DRO_EPS:
                rec[f"WassDRO_eps{e}"] = {"true": [], "obj": [], "x": []}
            for r in range(reps):
                xi = P.sample_newsvendor(N, rng, case)
                V = P.bootstrap_multiplicities(N, M, rng)
                for key, sol in [("SAA", P.solve_nv(xi, 1.0, None))] + \
                        [(f"APUB_{c}", P.solve_nv(xi, 1.0 - c, V)) for c in confs] + \
                        [(f"WassDRO_eps{e}", P.solve_nv(xi, 1.0, None, dro_eps=e))
                         for e in DRO_EPS]:
                    rec[key]["true"].append(true_cost(sol["x"]))
                    rec[key]["obj"].append(sol["obj"])
                    rec[key]["x"].append(sol["x"].tolist())
                log(f"  armC case{case} N={N} rep {r + 1}/{reps} done")
                if bud.stop(r + 1):
                    log(f"  armC case{case} N={N} budget reached after {r + 1} reps")
                    break
            cells.append(finalize_cell(rec))
            log(f"armC case{case} N={N} complete")
        out["cases"][f"case{case}"] = cells

    # Appendix D.2: l2 shift of the recommended order vector from N=30 to N=120.
    # The paper reports 7.89 (SAA-M), 3.99 (APUB at 1-alpha=0.5), 3.36 (at 0.95).
    shifts = {}
    c1 = out["cases"]["case1"]
    by_n = {c["N"]: c for c in c1}
    if 30 in by_n and 120 in by_n:
        for key in ["SAA", "APUB_0.5", "APUB_0.95"]:
            if key in by_n[30]["methods"] and key in by_n[120]["methods"]:
                a = np.array(by_n[30]["methods"][key]["mean_solution"])
                b = np.array(by_n[120]["methods"][key]["mean_solution"])
                shifts[key] = float(np.linalg.norm(b - a))
    out["l2_solution_shift_N30_to_N120"] = {
        "ours": shifts,
        "paper_appendix_D2": {"SAA": 7.89, "APUB_0.5": 3.99, "APUB_0.95": 3.36},
    }
    return out


# =============================================================================
# M-convergence check (reproduces the Appendix "m_converge" experiment)
# =============================================================================
def m_convergence(n_sim=4, N=120, budget_s=420):
    rng = np.random.default_rng(660214)
    grid = [250, 500, 1000, 2000, 4000]
    bud = Budget(budget_s)
    vals = {str(m): [] for m in grid}
    for s in range(n_sim):
        gam = P.sample_fixed_recourse_gamma(N, rng)
        for m in grid:
            V = P.bootstrap_multiplicities(N, m, rng)
            vals[str(m)].append(P.solve_fr(gam, 0.1, V)["obj"])
        log(f"  M-convergence sim {s + 1}/{n_sim} done")
        if bud.stop(s + 1):
            break
    vals = {k: v for k, v in vals.items()}
    n_sim = len(vals[str(grid[0])])
    ref = np.array(vals[str(grid[-1])])
    return {"N": N, "n_simulations": n_sim, "nominal_level": 0.9,
            "instance": "fixed-recourse product mix (Section 5.3)",
            "optimal_values_by_M": {k: [float(x) for x in v] for k, v in vals.items()},
            "across_simulation_std_by_M": {k: float(np.std(v, ddof=1))
                                           for k, v in vals.items()},
            "mean_abs_rel_gap_to_M5000_by_M": {
                k: float(np.mean(np.abs(np.array(v) - ref) / np.abs(ref)))
                for k, v in vals.items()}}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=["A", "BC"])
    args = ap.parse_args()
    os.makedirs(RESDIR, exist_ok=True)
    meta = {"orid": "eXLcL70GXO", "arxiv": "2403.08966",
            "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "confidence_levels": CONFS}
    if args.arm == "A":
        res = {"meta": meta, "armA_product_mix_random_recourse":
               arm_a(reps_small=16, reps_large=4, M=1200, budget_s=3000)}
        path = os.path.join(RESDIR, "applications_armA.json")
    else:
        res = {"meta": meta}
        res["M_convergence"] = m_convergence()
        log("M-convergence done")
        res["armB_product_mix_fixed_recourse_vs_DRO"] = arm_b(
            plan=[(30, 25, CONFS, 0.25), (120, 15, CONFS, 0.35),
                  (480, 4, [0.5, 0.9, 0.95], 0.40)],
            M=1200, budget_s=1500)
        log("armB done")
        res["armC_newsvendor_10product"] = arm_c(
            plan=[(30, 20, CONFS, 0.25), (60, 12, CONFS, 0.35),
                  (120, 10, CONFS, 0.40)], M=1500, budget_s=1500)
        log("armC done")
        path = os.path.join(RESDIR, "applications_armBC.json")
    with open(path, "w") as fh:
        json.dump(res, fh, indent=1)
    log(f"wrote {path}")


if __name__ == "__main__":
    main()