Advanced Topics in Stochastic Programming

# Install course dependencies before importing Pyomo on Colab.
import sys

if "google.colab" in sys.modules:
    !wget "https://raw.githubusercontent.com/ndcbe/optimization/main/notebooks/helper.py"
    import helper

    helper.easy_install()
else:
    sys.path.insert(0, "../")
    import helper

helper.set_plotting_style()

import numpy as np

lower, upper = -1.0, 1.0


def integrand(x):
    return 7 * x**3 - 8 * x**2 - 3 * x + 3


# Odd powers integrate to zero; integrate -8*x**2 + 3 analytically.
exact = 2 / 3
trapezoid = (upper - lower) * (integrand(lower) + integrand(upper)) / 2
# Two nodes integrate polynomials through degree 3 exactly on [-1, 1].
nodes, weights = np.polynomial.legendre.leggauss(2)
gauss_legendre = sum(weights * integrand(nodes))

print(f"exact: {exact:.12f}")
print(f"trapezoid: {trapezoid:.12f}")
print(f"two-point Gauss--Legendre: {gauss_legendre:.12f}")
assert abs(gauss_legendre - exact) < 1e-12

# The scenario structure is built with Pyomo Blocks: the stage 1 acreage lives on
# the model, and each sampled realization gets its own Block holding that
# sample's stage 2 purchases and sales. This is the same pattern used in
# "Stochastic Programming" (notebook 4/SP); the only difference is that the
# blocks are now indexed by a sample number instead of a named scenario, and
# there can be hundreds of them.
import numpy as np
import pyomo.environ as pyo
from pyomo.environ import units as u
from pyomo.util.check_units import assert_units_consistent
from scipy import stats

# Pyomo's unit library has acres and metric tons but no money, so declare one.
if not hasattr(u, "USD"):
    u.load_definitions_from_strings(["USD = [currency]"])

USD_PER_ACRE = u.USD / u.acre
USD_PER_TON = u.USD / u.metric_ton
TON_PER_ACRE = u.metric_ton / u.acre

# Birge & Louveaux, 2nd ed., Table 1, p. 5: mean yields [T/acre], planting costs
# [USD/acre], purchase and selling prices [USD/T], and the cattle feed
# requirements [T]. The 6000 T quota caps sales at the favorable beet price.
CROPS = ["WHEAT", "CORN", "BEETS"]
FEED_CROPS = ["WHEAT", "CORN"]
SALE_TYPES = ["WHEAT", "CORN", "BEETS_FAVORABLE", "BEETS_UNFAVORABLE"]

MEAN_YIELD = {"WHEAT": 2.5, "CORN": 3.0, "BEETS": 20.0}
PLANT_COST = {"WHEAT": 150.0, "CORN": 230.0, "BEETS": 260.0}
BUY_PRICE = {"WHEAT": 238.0, "CORN": 210.0}
SELL_PRICE = {
    "WHEAT": 170.0,
    "CORN": 150.0,
    "BEETS_FAVORABLE": 36.0,
    "BEETS_UNFAVORABLE": 10.0,
}
FEED_NEED = {"WHEAT": 200.0, "CORN": 240.0}
TOTAL_LAND = 500.0
BEET_QUOTA = 6000.0


def build_saa_model(multipliers, weights=None, check_units=False):
    """Extensive form over a list of yield multipliers, one Block per sample.

    Arguments:
        multipliers: yield multipliers, one per sample. Each scales all three
            mean yields.
        weights: probability weight of each sample. Defaults to 1/N, which is
            the sample average approximation. Passing quadrature weights
            instead turns the same builder into a deterministic integration
            rule, which is how the reference answer below is computed.
        check_units: run assert_units_consistent. Left off by default because
            it walks every constraint on every block, and these models are
            built hundreds of times; it is exercised on a small model below.

    Return:
        A Pyomo model whose objective is the weighted average cost in USD.
    """
    multipliers = np.asarray(multipliers, dtype=float)
    n = len(multipliers)
    if weights is None:
        weights = np.full(n, 1.0 / n)
    weights = np.asarray(weights, dtype=float)
    assert len(weights) == n, "one weight per sample"

    m = pyo.ConcreteModel()
    m.CROPS = pyo.Set(initialize=CROPS)
    m.FEED_CROPS = pyo.Set(initialize=FEED_CROPS)
    m.SALE_TYPES = pyo.Set(initialize=SALE_TYPES)
    m.SAMPLES = pyo.Set(initialize=range(n))

    # Stage 1: acres of each crop, chosen before any yield is observed [acre]
    m.acreage = pyo.Var(m.CROPS, domain=pyo.NonNegativeReals, units=u.acre)

    # Stage 1 cost [USD]
    m.planting_cost = sum(PLANT_COST[c] * USD_PER_ACRE * m.acreage[c] for c in CROPS)

    # Stage 1 constraint: the farm is 500 acres [acre]
    m.total_land_constraint = pyo.Constraint(
        expr=pyo.summation(m.acreage) <= TOTAL_LAND * u.acre
    )

    def scenario_block_rule(b, i):
        """Stage 2 for sample i. `b` is the block; `m` comes from the closure."""
        # Stage 2 decisions, made after the yield is observed [T]
        b.purchases = pyo.Var(
            m.FEED_CROPS, domain=pyo.NonNegativeReals, units=u.metric_ton
        )
        b.sales = pyo.Var(m.SALE_TYPES, domain=pyo.NonNegativeReals, units=u.metric_ton)

        b.purchase_cost = sum(
            BUY_PRICE[c] * USD_PER_TON * b.purchases[c] for c in FEED_CROPS
        )
        b.sales_revenue = sum(
            SELL_PRICE[s] * USD_PER_TON * b.sales[s] for s in SALE_TYPES
        )

        # Yields realized in this sample [T/acre]
        y = {c: MEAN_YIELD[c] * multipliers[i] for c in CROPS}

        # Feed requirements: grow it, buy it, or do not sell it [T]
        b.wheat_constraint = pyo.Constraint(
            expr=y["WHEAT"] * TON_PER_ACRE * m.acreage["WHEAT"]
            + b.purchases["WHEAT"]
            - b.sales["WHEAT"]
            >= FEED_NEED["WHEAT"] * u.metric_ton
        )
        b.corn_constraint = pyo.Constraint(
            expr=y["CORN"] * TON_PER_ACRE * m.acreage["CORN"]
            + b.purchases["CORN"]
            - b.sales["CORN"]
            >= FEED_NEED["CORN"] * u.metric_ton
        )
        # Beet sales cannot exceed beet production [T]
        b.beets_constraint = pyo.Constraint(
            expr=y["BEETS"] * TON_PER_ACRE * m.acreage["BEETS"]
            - b.sales["BEETS_FAVORABLE"]
            - b.sales["BEETS_UNFAVORABLE"]
            >= 0 * u.metric_ton
        )
        b.sales["BEETS_FAVORABLE"].setub(BEET_QUOTA)

    # One Block per sample: identical structure, different yield data
    m.scenarios = pyo.Block(m.SAMPLES, rule=scenario_block_rule)

    # Weighted average cost [USD]
    m.objective = pyo.Objective(
        expr=m.planting_cost
        + sum(
            weights[i] * (m.scenarios[i].purchase_cost - m.scenarios[i].sales_revenue)
            for i in range(n)
        ),
        sense=pyo.minimize,
    )

    if check_units:
        assert_units_consistent(m)
    return m


solver = pyo.SolverFactory("appsi_highs")
assert solver.available(), "HiGHS is required for this section"
print("HiGHS version:", solver.version())


def solve_model(m):
    """Solve and return the objective value, refusing anything but optimal."""
    results = solver.solve(m)
    assert pyo.check_optimal_termination(
        results
    ), f"Solve failed: termination={results.termination_condition}"
    return pyo.value(m.objective)


def first_stage(m):
    """The stage 1 acreage as a plain array, in the order of CROPS."""
    return np.array([pyo.value(m.acreage[c]) for c in CROPS])


def saa_optimum(multipliers, weights=None):
    """Solve the sampled problem: returns (objective, acreage)."""
    m = build_saa_model(multipliers, weights)
    return solve_model(m), first_stage(m)


def evaluate_candidate(acreage, multipliers, weights=None):
    """Weighted average cost of a FIXED stage 1 plan, recourse still optimized.

    Fixing the stage 1 variables on the same model is all that separates
    "choose a plan" from "grade a plan". Every block still solves its own
    recourse problem.
    """
    m = build_saa_model(multipliers, weights)
    for c, acres in zip(CROPS, acreage):
        m.acreage[c].fix(acres)
    return solve_model(m)

# Birge & Louveaux Table 5, p. 8 -- the RP solution quoted in Lecture 10.
book_cost, book_acres = saa_optimum([0.8, 1.0, 1.2])
print(f"expected profit = ${-book_cost:,.0f}   (book: $108,390)")
print(
    "acres  wheat {:.1f}  corn {:.1f}  beets {:.1f}   (book: 170, 80, 250)".format(
        *book_acres
    )
)
assert abs(-book_cost - 108390.0) < 1.0
assert np.allclose(book_acres, [170.0, 80.0, 250.0], atol=1e-6)

# Exercise the unit declarations once, on a small model.
assert_units_consistent(build_saa_model([0.8, 1.0, 1.2], check_units=True))
print("units consistent")

def trapezoid_grid(n, lo=0.8, hi=1.2):
    """n-point trapezoid rule on [lo, hi] for a Uniform(lo, hi) expectation.

    The uniform density 1/(hi-lo) cancels the (hi-lo) in the quadrature
    weights, so the weights below sum to one and are directly comparable with
    the 1/N sampling weights.
    """
    nodes = np.linspace(lo, hi, n)
    weights = np.full(n, 1.0 / (n - 1))
    weights[0] = weights[-1] = 0.5 / (n - 1)
    return nodes, weights


print("grid      reference cost        wheat     corn    beets")
for n_grid in (201, 401, 801, 1601, 3201):
    nodes, weights = trapezoid_grid(n_grid)
    z_ref, x_ref = saa_optimum(nodes, weights)
    print(
        f"{n_grid:5d}   {z_ref:16,.4f}   {x_ref[0]:7.2f}  {x_ref[1]:6.2f}  {x_ref[2]:7.2f}"
    )

# Keep the finest grid as the reference for the rest of the section.
REF_NODES, REF_WEIGHTS = trapezoid_grid(3201)
Z_STAR, X_STAR = saa_optimum(REF_NODES, REF_WEIGHTS)
print(f"\nreference z* = {Z_STAR:,.2f} USD  (expected profit ${-Z_STAR:,.2f})")

# One independent, seeded draw per sample size.
rng = np.random.default_rng(20260920)

print("    N       z_N          q_ref(x_N)     true gap     wheat    corn   beets")
sweep = []
for N in (5, 10, 25, 50, 100, 200, 500):
    xi = rng.uniform(0.8, 1.2, N)
    z_N, x_N = saa_optimum(xi)
    q_ref = evaluate_candidate(x_N, REF_NODES, REF_WEIGHTS)
    sweep.append((N, z_N, q_ref))
    print(
        f"{N:5d}  {z_N:12,.2f}  {q_ref:14,.2f}  {q_ref - Z_STAR:10,.2f}   "
        f"{x_N[0]:7.1f} {x_N[1]:6.1f} {x_N[2]:7.1f}"
    )

# Each replication solves the sampled problem and evaluates the reference plan
# on the SAME sample, so the shared sampling noise cancels in the difference.
rng_bias = np.random.default_rng(424242)
T_bias = 100
sizes = (5, 10, 20, 40, 80, 160)
bias_mean, bias_se = [], []

print(f"replications T = {T_bias}")
print("    N     mean D (bias)      std error    worst-case D      N x |bias|")
for N in sizes:
    D = np.empty(T_bias)
    for j in range(T_bias):
        xi = rng_bias.uniform(0.8, 1.2, N)
        D[j] = saa_optimum(xi)[0] - evaluate_candidate(X_STAR, xi)
    assert (D <= 0).all(), "the inequality above is pathwise, not on average"
    bias_mean.append(D.mean())
    bias_se.append(D.std(ddof=1) / np.sqrt(T_bias))
    print(
        f"{N:5d}  {D.mean():14,.2f}  {bias_se[-1]:12,.2f}  {D.max():14,.2f}  "
        f"{N * abs(D.mean()):14,.0f}"
    )

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(5.5, 3.6))
ax.errorbar(
    sizes,
    [-b for b in bias_mean],
    yerr=bias_se,
    marker="o",
    linestyle="-",
    color="tab:blue",
    label=r"measured $|\mathbb{E}[z_N] - z^*|$",
)
reference = abs(bias_mean[0]) * sizes[0] / np.array(sizes, dtype=float)
ax.plot(
    sizes,
    reference,
    marker="s",
    linestyle="--",
    color="0.35",
    label=r"$1/N$ reference, anchored at $N=5$",
)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlabel("sample size $N$")
ax.set_ylabel("optimistic bias [USD]")
ax.set_title("Bias of the sampled optimum")
ax.legend(frameon=False, fontsize=8)
fig.tight_layout()
plt.show()

# Training: a deliberately small sample, so the candidate is visibly imperfect.
rng_gap = np.random.default_rng(31415)
xi_train = rng_gap.uniform(0.8, 1.2, 20)
z_train, x_hat = saa_optimum(xi_train)
q_hat = evaluate_candidate(x_hat, REF_NODES, REF_WEIGHTS)
true_gap = q_hat - Z_STAR

print(
    f"candidate from N = 20: wheat {x_hat[0]:.1f}, corn {x_hat[1]:.1f}, "
    f"beets {x_hat[2]:.1f} acres"
)
print(f"training objective z_20      = {z_train:12,.2f}  (optimistic, not a bound)")
print(f"true expected cost q(x_hat)  = {q_hat:12,.2f}")
print(f"TRUE optimality gap          = {true_gap:12,.2f}  <- known only because")
print("                                               the reference exists\n")


def gap_bound(x_candidate, T, N, delta, seed):
    """Multiple-replication optimality-gap bound for a fixed candidate.

    Returns (G, bound). Each batch is fresh and independent of the others;
    within a batch the candidate evaluation and the re-optimization share the
    realizations, which is the paired (common random numbers) estimator.
    """
    rng_local = np.random.default_rng(seed)
    G = np.empty(T)
    for j in range(T):
        xi = rng_local.uniform(0.8, 1.2, N)
        G[j] = evaluate_candidate(x_candidate, xi) - saa_optimum(xi)[0]
    t_crit = stats.t.ppf(1.0 - delta, T - 1)  # ONE-sided critical value
    bound = G.mean() + t_crit * G.std(ddof=1) / np.sqrt(T)
    return G, bound


print("  T     N      G_bar       s_G     one-sided 95% bound    covers true gap?")
for T, N, seed in ((30, 50, 9001), (30, 200, 9002), (100, 200, 9003)):
    G, bound = gap_bound(x_hat, T=T, N=N, delta=0.05, seed=seed)
    assert (G >= -1e-6).all(), "paired gaps are nonnegative by construction"
    print(
        f"{T:4d} {N:5d}  {G.mean():9,.2f}  {G.std(ddof=1):8,.2f}  "
        f"{bound:18,.2f}    {'yes' if bound >= true_gap else 'NO'}"
    )