scripts/claim4_knapsack.py

#!/usr/bin/env python3
"""Claim 4 (CROWN JEWEL) - HybridFlow ICML 2026 (#9046, arXiv 2512.22137).

Reproduces Table 3 (routing-strategy ablation on GPQA): the normalized cost
`c` and the "unified utility" `u` (benefit-cost ratio) for every row, using the
formulas RECOVERED from the paper (the paper's own reference for u is broken:
"see Sec. ??").  HybridFlow attains the highest utility u = 0.7940.

Recovered formulas (verified exactly against every Table-3 and Table-6 row):
    c_row = 1/2 * ((Latency - Latency_edge)/l_max) + 1/2 * (API_Cost / k_max)
    u_row = (Acc - Acc_edge) / (100 * c_row)
with Latency_edge = 11.99 s, Acc_edge = 25.54 %, l_max = 10 s, k_max = 0.02 $.

Also demonstrates the underlying MECHANISM on a small synthetic subtask set:
  * the exact 0-1 knapsack (max sum r_i*dq_i s.t. sum r_i*c_i <= C_max) via DP;
  * the Lagrangian / shadow-price rule (offload iff dq_i/c_i > lambda) that the
    paper's learned router approximates.

CPU-only, no API needed.  Writes outputs/claim4.json and figs/claim4_*.png.
"""
import json
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(HERE, "outputs")
FIG = os.path.join(HERE, "figs")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)

# Constants for the Claim-4 cost normalization (Eq 24).
ACC_EDGE = 25.54  # Edge (Llama3.2-3B) accuracy, %
LAT_EDGE = 11.99  # Edge latency, s
L_MAX = 10.0  # latency normalization scale, s  (NOT the 20 s threshold budget)
K_MAX = 0.02  # API-cost normalization scale, $


def norm_cost(latency, api_cost):
    """Normalized total cost c (Table-3/6 column, Eq 1/24)."""
    return 0.5 * ((latency - LAT_EDGE) / L_MAX) + 0.5 * (api_cost / K_MAX)


def utility(acc, c):
    """Unified utility u = benefit(over edge) / cost (Def 3.2 at row level)."""
    return (acc - ACC_EDGE) / (100.0 * c)


# ---------------------------------------------------------------------------
# Part A - reproduce every Table-3 row exactly.
# Columns: offload_rate%, accuracy%, latency s, api_cost $, reported c, reported u
# ---------------------------------------------------------------------------
TABLE3 = [
    # name,                       offload, acc,   lat,   api,    c_rep,  u_rep
    ("Edge (Llama3.2-3B)", 0.0, 25.54, 11.99, 0.0000, None, None),
    ("Cloud (GPT-4.1)", 100.0, 57.28, 18.26, 0.0185, 0.7760, 0.4090),
    ("Random (L3B + GPT-4.1)", 42.1, 46.00, 15.15, 0.0075, 0.3455, 0.5922),
    ("Fixed Threshold (t0=0.5)", 41.18, 51.62, 15.88, 0.0088, 0.4145, 0.6292),
    ("HybridFlow-Chain", 40.81, 50.62, 16.12, 0.0082, 0.4115, 0.6095),
    ("HybridFlow (Ours)", 40.48, 53.33, 15.24, 0.0075, 0.3500, 0.7940),
]

print("=" * 78)
print("Part A: Table 3 (GPQA) - reproduce normalized cost c and utility u")
print("=" * 78)
print(f"{'Method':<28}{'c_calc':>9}{'c_rep':>9}{'u_calc':>9}{'u_rep':>9}  ok")
table3_rows = []
all_ok = True
for name, offload, acc, lat, api, c_rep, u_rep in TABLE3:
    if name.startswith("Edge"):
        # Edge is the reference row: c = 0, u undefined.
        c_calc, u_calc = 0.0, None
        ok = True
        print(f"{name:<28}{'0':>9}{'-':>9}{'N/A':>9}{'-':>9}  {'OK' if ok else 'FAIL'}")
    else:
        c_calc = norm_cost(lat, api)
        u_calc = utility(acc, c_calc)
        c_ok = abs(round(c_calc, 4) - c_rep) <= 1e-3
        u_ok = abs(round(u_calc, 4) - u_rep) <= 1e-3
        ok = c_ok and u_ok
        all_ok = all_ok and ok
        assert c_ok, f"c mismatch {name}: {c_calc:.4f} vs {c_rep}"
        assert u_ok, f"u mismatch {name}: {u_calc:.4f} vs {u_rep}"
        print(
            f"{name:<28}{c_calc:>9.4f}{c_rep:>9.4f}{u_calc:>9.4f}{u_rep:>9.4f}  "
            f"{'OK' if ok else 'FAIL'}"
        )
    table3_rows.append(
        {
            "method": name,
            "offload_rate": offload,
            "accuracy": acc,
            "latency": lat,
            "api_cost": api,
            "c_calc": round(c_calc, 4),
            "c_reported": c_rep,
            "u_calc": None if u_calc is None else round(u_calc, 4),
            "u_reported": u_rep,
        }
    )

winner = max(
    (r for r in table3_rows if r["u_calc"] is not None), key=lambda r: r["u_calc"]
)
assert winner["method"].startswith("HybridFlow (Ours)"), winner["method"]
assert abs(winner["u_calc"] - 0.7940) <= 1e-3
print(
    f"\nHighest utility: {winner['method']}  u = {winner['u_calc']:.4f}  "
    f"(expected 0.7940)  -> PASS"
)

# ---------------------------------------------------------------------------
# Part B - the underlying 0-1 knapsack + Lagrangian mechanism on synthetic data.
# item i = subtask; value dq_i = accuracy gain (cloud vs edge); weight c_i = cost.
# ---------------------------------------------------------------------------
print("\n" + "=" * 78)
print("Part B: 0-1 knapsack (max sum dq_i s.t. sum c_i <= C_max) + Lagrangian rule")
print("=" * 78)
# synthetic subtask set (dq = quality gain if offloaded to cloud; c = norm cost)
SUBTASKS = [
    {"id": "t1", "dq": 0.36, "c": 0.12},  # ratio 3.00
    {"id": "t2", "dq": 0.40, "c": 0.20},  # ratio 2.00
    {"id": "t3", "dq": 0.30, "c": 0.20},  # ratio 1.50
    {"id": "t4", "dq": 0.18, "c": 0.20},  # ratio 0.90
    {"id": "t5", "dq": 0.10, "c": 0.25},  # ratio 0.40
    {"id": "t6", "dq": 0.06, "c": 0.30},  # ratio 0.20
]
C_MAX = 0.52
SCALE = 1000  # scale costs to integers for the DP capacity axis


def knapsack_01(items, capacity, scale):
    """Exact 0-1 knapsack via DP. Returns (best_value, chosen_ids)."""
    W = int(round(capacity * scale))
    weights = [int(round(it["c"] * scale)) for it in items]
    values = [it["dq"] for it in items]
    n = len(items)
    dp = [[0.0] * (W + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        wi, vi = weights[i - 1], values[i - 1]
        for w in range(W + 1):
            dp[i][w] = dp[i - 1][w]
            if wi <= w and dp[i - 1][w - wi] + vi > dp[i][w]:
                dp[i][w] = dp[i - 1][w - wi] + vi
    # backtrack
    chosen, w = [], W
    for i in range(n, 0, -1):
        if dp[i][w] != dp[i - 1][w]:
            chosen.append(items[i - 1]["id"])
            w -= weights[i - 1]
    return dp[n][W], sorted(chosen)


best_val, chosen = knapsack_01(SUBTASKS, C_MAX, SCALE)
chosen_cost = sum(it["c"] for it in SUBTASKS if it["id"] in chosen)
print(f"Capacity C_max = {C_MAX}")
print(
    f"DP optimum: offload {chosen}  total_gain = {best_val:.2f}  "
    f"total_cost = {chosen_cost:.2f} (<= {C_MAX})"
)
assert chosen_cost <= C_MAX + 1e-9

# brute-force check that the DP is optimal
best_bf, best_set = -1.0, None
for mask in range(1 << len(SUBTASKS)):
    tot_c = sum(SUBTASKS[i]["c"] for i in range(len(SUBTASKS)) if mask & (1 << i))
    tot_v = sum(SUBTASKS[i]["dq"] for i in range(len(SUBTASKS)) if mask & (1 << i))
    if tot_c <= C_MAX + 1e-9 and tot_v > best_bf:
        best_bf, best_set = tot_v, [
            SUBTASKS[i]["id"] for i in range(len(SUBTASKS)) if mask & (1 << i)
        ]
assert abs(best_bf - best_val) <= 1e-9, f"DP {best_val} != brute {best_bf}"
print(f"Brute-force optimum matches DP: gain = {best_bf:.2f}  -> PASS")

# Lagrangian / shadow-price rule: offload iff dq_i/c_i > lambda.
# Sweep lambda; report the routing at a representative shadow price.
print("\nLagrangian rule: offload subtask iff (dq_i / c_i) > lambda (shadow price)")
ratios = {it["id"]: it["dq"] / it["c"] for it in SUBTASKS}
for it in sorted(SUBTASKS, key=lambda x: -ratios[x["id"]]):
    print(
        f"  {it['id']}: dq={it['dq']:.2f} c={it['c']:.2f} "
        f"benefit-cost ratio = {ratios[it['id']]:.3f}"
    )
lam = 2.5
lagr_offload = sorted([it["id"] for it in SUBTASKS if ratios[it["id"]] > lam])
lagr_cost = sum(it["c"] for it in SUBTASKS if it["id"] in lagr_offload)
print(
    f"At high lambda = {lam} (scarce budget): offload {lagr_offload}  "
    f"cost = {lagr_cost:.2f}  (conservative)"
)
# At a shadow price that makes the Lagrangian budget match, it recovers knapsack.
lam_star = 1.2
lagr_star = sorted([it["id"] for it in SUBTASKS if ratios[it["id"]] > lam_star])
lagr_star_cost = sum(it["c"] for it in SUBTASKS if it["id"] in lagr_star)
recovers = lagr_star == chosen
print(
    f"At lambda* = {lam_star} (tuned to budget): offload {lagr_star}  "
    f"cost = {lagr_star_cost:.2f}"
)
print(f"Lagrangian shadow-price rule recovers the knapsack optimum: {recovers}")
assert recovers, f"Lagrangian {lagr_star} != knapsack {chosen}"

# ---------------------------------------------------------------------------
# Figure 1: utility per routing strategy (bar chart, HybridFlow highest).
# ---------------------------------------------------------------------------
bar_rows = [r for r in table3_rows if r["u_calc"] is not None]
names = [
    r["method"]
    .replace(" (Ours)", "")
    .replace(" (L3B + GPT-4.1)", "")
    .replace(" (GPT-4.1)", "")
    .replace(" (t0=0.5)", " t0=.5")
    for r in bar_rows
]
uvals = [r["u_calc"] for r in bar_rows]
colors = [
    "#888888" if not r["method"].startswith("HybridFlow (Ours)") else "#d1495b"
    for r in bar_rows
]
fig, ax = plt.subplots(figsize=(7.5, 4.2))
bars = ax.bar(names, uvals, color=colors)
for b, v in zip(bars, uvals):
    ax.text(
        b.get_x() + b.get_width() / 2,
        v + 0.01,
        f"{v:.4f}",
        ha="center",
        va="bottom",
        fontsize=9,
    )
ax.set_ylabel("Unified utility u (higher is better)")
ax.set_title("Claim 4 - Table 3: benefit-cost utility by routing strategy (GPQA)")
ax.set_ylim(0, 0.9)
ax.axhline(0.7940, ls="--", lw=0.8, color="#d1495b", alpha=0.6)
plt.xticks(rotation=20, ha="right", fontsize=8)
plt.tight_layout()
fig.savefig(os.path.join(FIG, "claim4_utility_bars.png"), dpi=130)
plt.close(fig)

# ---------------------------------------------------------------------------
# Figure 2: knapsack benefit-cost ratios with the selected (offloaded) subtasks.
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(7.0, 4.0))
ids = [it["id"] for it in SUBTASKS]
rvals = [ratios[i] for i in ids]
kcolors = ["#2e86ab" if i in chosen else "#cccccc" for i in ids]
bars = ax.bar(ids, rvals, color=kcolors)
for b, v in zip(bars, rvals):
    ax.text(
        b.get_x() + b.get_width() / 2,
        v + 0.02,
        f"{v:.2f}",
        ha="center",
        va="bottom",
        fontsize=9,
    )
ax.set_ylabel("benefit-cost ratio  dq_i / c_i")
ax.set_title(
    f"Claim 4 - 0-1 knapsack: offloaded subtasks (blue), "
    f"C_max={C_MAX}, gain={best_val:.2f}"
)
plt.tight_layout()
fig.savefig(os.path.join(FIG, "claim4_knapsack.png"), dpi=130)
plt.close(fig)

# ---------------------------------------------------------------------------
# Dump results.
# ---------------------------------------------------------------------------
result = {
    "claim": "4 - Table 3 0-1 knapsack benefit-cost utility (u=0.7940)",
    "formulas": {
        "c_row": "0.5*((Lat-11.99)/10) + 0.5*(API/0.02)",
        "u_row": "(Acc-25.54)/(100*c_row)",
        "note": "paper's u reference is broken ('see Sec. ??'); recovered & verified",
    },
    "table3": table3_rows,
    "table3_all_rows_match": all_ok,
    "winner": {"method": winner["method"], "u": winner["u_calc"]},
    "knapsack_demo": {
        "subtasks": SUBTASKS,
        "C_max": C_MAX,
        "dp_optimum_offload": chosen,
        "dp_total_gain": round(best_val, 4),
        "dp_total_cost": round(chosen_cost, 4),
        "brute_force_matches_dp": abs(best_bf - best_val) <= 1e-9,
        "lagrangian_lambda_star": lam_star,
        "lagrangian_recovers_knapsack": recovers,
    },
    "figures": ["figs/claim4_utility_bars.png", "figs/claim4_knapsack.png"],
    "verdict": "PASS - all 5 finite Table-3 rows reproduced exactly; "
    "HybridFlow highest u=0.7940; knapsack DP == brute force.",
}
with open(os.path.join(OUT, "claim4.json"), "w") as f:
    json.dump(result, f, indent=2)
print(f"\nAll Table-3 rows match: {all_ok}")
print(f"Wrote {os.path.join(OUT, 'claim4.json')}")
print("VERDICT: PASS")