Stochastic Programming

# Imports
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
import pyomo.environ as pyo
from pyomo.environ import units as u
from pyomo.util.check_units import assert_units_consistent

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
CROPS = ["WHEAT", "CORN", "BEETS"]
FEED_CROPS = ["WHEAT", "CORN"]
SALE_TYPES = ["WHEAT", "CORN", "BEETS_FAVORABLE", "BEETS_UNFAVORABLE"]
MEAN_YIELD = dict(zip(CROPS, [2.5, 3.0, 20.0]))  # T/acre
PLANT_COST = dict(zip(CROPS, [150.0, 230.0, 260.0]))  # $/acre
BUY_PRICE = dict(zip(FEED_CROPS, [238.0, 210.0]))  # $/T
SELL_PRICE = dict(zip(SALE_TYPES, [170.0, 150.0, 36.0, 10.0]))  # $/T
FEED_NEED = dict(zip(FEED_CROPS, [200.0, 240.0]))  # T
TOTAL_LAND, BEET_QUOTA = 500.0, 6000.0  # acres, T

def build_farmer(
    scenarios, weights=None, plant_cost=None, risk=None, check_units=False
):
    """Return a fresh farmer model; scenarios contain yields and optional prices.

    scenarios: nonempty list of dictionaries; 'yield' maps crops to T/acre.
    Optional 'buy' and 'sell' mappings give $/T. weights sum to one.
    plant_cost maps crops to $/acre. risk selects the optional Lecture 10 policy.
    """
    n = len(scenarios)
    if n == 0:
        raise ValueError("At least one scenario is required")
    w = np.full(n, 1 / n) if weights is None else np.asarray(weights, float)
    if w.shape != (n,) or not np.all(np.isfinite(w)) or np.any(w <= 0):
        raise ValueError("Retained scenarios need positive finite probabilities")
    if not np.isclose(w.sum(), 1.0):
        raise ValueError("Scenario probabilities must sum to one")
    costs = PLANT_COST if plant_cost is None else plant_cost
    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.SCENARIOS = pyo.RangeSet(0, n - 1)
    m.weight = pyo.Param(m.SCENARIOS, initialize=dict(enumerate(w)))
    # Plant once, before learning the yields: no scenario index on acreage.
    m.acreage = pyo.Var(m.CROPS, domain=pyo.NonNegativeReals, units=u.acre)
    m.land_limit = pyo.Constraint(expr=sum(m.acreage.values()) <= TOTAL_LAND * u.acre)
    m.planting_cost = pyo.Expression(
        expr=sum(costs[c] * USD_PER_ACRE * m.acreage[c] for c in CROPS)
    )
    # The helper below attaches recourse to each Block, sharing root acreage.
    m.scenarios = pyo.Block(
        m.SCENARIOS, rule=lambda b, s: add_farmer_scenario(b, m, scenarios[s])
    )
    m.expected_profit = pyo.Expression(
        expr=sum(m.weight[s] * m.scenarios[s].profit for s in m.SCENARIOS)
    )
    m.objective = pyo.Objective(expr=m.expected_profit, sense=pyo.maximize)
    if risk:
        add_farmer_risk(m, risk)  # Optional policies introduced in Lecture 10.
    if check_units:
        assert_units_consistent(m)
    return m

def add_farmer_scenario(b, m, scenario):
    """Attach one scenario's recourse, using the root model's acreage."""
    crop_yield = scenario["yield"]
    buy = scenario.get("buy", BUY_PRICE)
    sell = scenario.get("sell", SELL_PRICE)
    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.sales["BEETS_FAVORABLE"].setub(BEET_QUOTA)

    # Harvest + purchases - sales must cover animal feed in this scenario.
    @b.Constraint(m.FEED_CROPS)
    def feed_balance(block, c):
        return (
            crop_yield[c] * TON_PER_ACRE * m.acreage[c]
            + block.purchases[c]
            - block.sales[c]
            >= FEED_NEED[c] * u.metric_ton
        )

    b.beets_constraint = pyo.Constraint(
        expr=b.sales["BEETS_FAVORABLE"] + b.sales["BEETS_UNFAVORABLE"]
        <= crop_yield["BEETS"] * TON_PER_ACRE * m.acreage["BEETS"]
    )
    # Profit includes planting cost; weights summing to one count it once.
    b.profit = pyo.Expression(
        expr=sum(sell[k] * USD_PER_TON * b.sales[k] for k in SALE_TYPES)
        - sum(buy[c] * USD_PER_TON * b.purchases[c] for c in FEED_CROPS)
        - m.planting_cost
    )

farmer_solver = pyo.SolverFactory("appsi_highs")
assert farmer_solver.available(), "HiGHS is required for the rest of this notebook"


def solve_farmer(m):
    """Solve, refusing anything but a proven optimum."""
    results = farmer_solver.solve(m)
    assert pyo.check_optimal_termination(
        results
    ), f"Solve failed: termination={results.solver.termination_condition}"
    return m


def plan_of(m):
    """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 grade_plan(plan, scenarios, weights=None, plant_cost=None):
    """Fix acreage in CROPS order; optimize trades for the supplied scenarios.

    Return expected profit and an array of scenario profits, both in dollars.
    """
    m = build_farmer(scenarios, weights, plant_cost)
    for crop, acres in zip(CROPS, plan):
        m.acreage[crop].fix(acres)
    solve_farmer(m)
    return (
        pyo.value(m.expected_profit),
        np.array([pyo.value(m.scenarios[i].profit) for i in m.SCENARIOS]),
    )

def correlated_scenario(multiplier):
    """One scenario of Section 1.1b: every crop scaled by the same factor."""
    return {"yield": {c: MEAN_YIELD[c] * multiplier for c in CROPS}}


# The three equally likely scenarios of Section 1.1b
THREE = [correlated_scenario(f) for f in (1.2, 1.0, 0.8)]

def add_farmer_cvar(m, risk):
    """Lower-tail profit CVaR: alpha is the confidence level."""
    w = m.weight
    n = len(m.SCENARIOS)
    # Rockafellar-Uryasev epigraph, B&L (9.8)-(9.10), p. 85, written for the
    # LOWER tail of profit: CVaR_a = max_nu { nu - E[(nu - profit)_+]/(1-a) }
    alpha = (risk.get("cvar") or risk.get("cvar_floor"))[0]
    m.nu = pyo.Var(domain=pyo.Reals, units=u.USD)
    m.shortfall = pyo.Var(m.SCENARIOS, domain=pyo.NonNegativeReals, units=u.USD)
    m.shortfall_con = pyo.Constraint(
        m.SCENARIOS,
        rule=lambda m, i: m.shortfall[i] >= m.nu - m.scenarios[i].profit,
    )
    m.cvar = m.nu - (1.0 / (1.0 - alpha)) * sum(w[i] * m.shortfall[i] for i in range(n))
    if "cvar_floor" in risk:
        m.cvar_con = pyo.Constraint(expr=m.cvar >= risk["cvar_floor"][1] * u.USD)

def add_farmer_risk(m, risk):
    """Attach the requested risk policy to a fresh farmer model."""
    w = m.weight
    n = len(m.SCENARIOS)
    m.del_component(m.objective)
    if "floor" in risk:
        # Exercise 7(b): a profit floor imposed in every scenario
        m.profit_floor = pyo.Constraint(
            m.SCENARIOS,
            rule=lambda m, i: m.scenarios[i].profit >= risk["floor"] * u.USD,
        )
    if "minimax" in risk:
        # B&L (9.15), p. 86, in profit form: maximize the worst scenario
        m.worst = pyo.Var(domain=pyo.Reals, units=u.USD)
        m.worst_con = pyo.Constraint(
            m.SCENARIOS, rule=lambda m, i: m.worst <= m.scenarios[i].profit
        )
    if "cvar" in risk or "cvar_floor" in risk:
        add_farmer_cvar(m, risk)
    if "downside" in risk:
        # Expected downside risk, B&L (5.1)-(5.3), p. 68
        target, level = risk["downside"]
        m.shortfall_below_target = pyo.Var(
            m.SCENARIOS, domain=pyo.NonNegativeReals, units=u.USD
        )
        m.downside_def = pyo.Constraint(
            m.SCENARIOS,
            rule=lambda m, i: m.shortfall_below_target[i]
            >= target * u.USD - m.scenarios[i].profit,
        )
        m.downside_limit = pyo.Constraint(
            expr=sum(w[i] * m.shortfall_below_target[i] for i in range(n))
            <= level * u.USD
        )

    if "minimax" in risk:
        m.objective = pyo.Objective(expr=m.worst, sense=pyo.maximize)
    elif "cvar" in risk:
        lam = risk["cvar"][1]
        m.objective = pyo.Objective(
            expr=(1 - lam) * m.expected_profit + lam * m.cvar, sense=pyo.maximize
        )
    else:
        m.objective = pyo.Objective(expr=m.expected_profit, sense=pyo.maximize)

def synthetic_farmer_yields(n, seed):
    """Draw synthetic joint yields, not observed agricultural data.

    Weather multiplier = 0.65 + 0.70 * Beta(2, 2), shared by all crops.
    Independent crop deviations are Uniform(-0.08, 0.08).
    All crop multipliers lie in [0.57, 1.43], with population mean 1.
    Dependence comes from shared weather; draws across scenarios are iid.
    """
    rng = np.random.default_rng(seed)
    weather = 0.65 + 0.70 * rng.beta(2, 2, size=n)
    multipliers = weather[:, None] + rng.uniform(-0.08, 0.08, size=(n, 3))
    return [{"yield": dict(zip(CROPS, row * np.array(list(MEAN_YIELD.values()))))}
            for row in multipliers]


SYNTHETIC_SEED = 20260927
TWENTY = synthetic_farmer_yields(20, SYNTHETIC_SEED)


def solution_record(m):
    """Full solved policy and recourse, with units explicit in field names."""
    return {
        "acreage": {c: float(pyo.value(m.acreage[c])) for c in CROPS},
        "expected_profit": float(pyo.value(m.expected_profit)),
        "scenarios": [
            {"id": int(i) + 1,
             "probability": float(pyo.value(m.weight[i])),
             "purchases_tons": {c: float(pyo.value(m.scenarios[i].purchases[c]))
                                for c in FEED_CROPS},
             "sales_tons": {c: float(pyo.value(m.scenarios[i].sales[c]))
                            for c in SALE_TYPES},
             "profit": float(pyo.value(m.scenarios[i].profit))}
            for i in m.SCENARIOS],
    }


def evaluate_solution(plan, scenarios, weights=None, plant_cost=None):
    """Return all efficient trades for fixed acreage, not arbitrary risk-solve trades."""
    m = build_farmer(scenarios, weights=weights, plant_cost=plant_cost)
    for c, acres in zip(CROPS, plan):
        m.acreage[c].fix(float(acres))
    return solution_record(solve_farmer(m))


def analytic_profits(plan, scenarios, plant_cost=None):
    """Independent closed-form recourse for these farmer data (dollars).

    Valid when sales prices are nonnegative, purchase prices exceed selling
    prices, and favorable beet revenue exceeds over-quota revenue. Feed
    shortages are bought, surpluses sold; beet sales fill the premium quota first.
    """
    costs = PLANT_COST if plant_cost is None else plant_cost
    acreage = dict(zip(CROPS, plan))
    planting = sum(costs[c] * acreage[c] for c in CROPS)
    profits = []
    for s in scenarios:
        sell, buy = s.get("sell", SELL_PRICE), s.get("buy", BUY_PRICE)
        assert all(buy[c] > sell[c] >= 0 for c in FEED_CROPS)
        assert sell["BEETS_FAVORABLE"] >= sell["BEETS_UNFAVORABLE"] >= 0
        profit = -planting
        for c in FEED_CROPS:
            surplus = s["yield"][c] * acreage[c] - FEED_NEED[c]
            profit += sell[c] * max(surplus, 0) - buy[c] * max(-surplus, 0)
        beets = s["yield"]["BEETS"] * acreage["BEETS"]
        profit += sell["BEETS_FAVORABLE"] * min(beets, BEET_QUOTA)
        profit += sell["BEETS_UNFAVORABLE"] * max(beets - BEET_QUOTA, 0)
        profits.append(profit)
    return np.asarray(profits)


def lower_tail_profit(profits, alpha, weights=None):
    """Average the worst mass 1-alpha, including fractional mass at ties."""
    values = np.asarray(profits, dtype=float)
    if not 0 <= alpha < 1:
        raise ValueError("alpha must be in [0, 1)")
    w = np.full(len(values), 1 / len(values)) if weights is None else np.asarray(weights)
    if len(values) == 0 or w.shape != values.shape or np.any(w < 0) or not np.isclose(w.sum(), 1):
        raise ValueError("Use nonempty outcomes with probabilities summing to one")
    order = np.argsort(values)
    remaining, total = 1 - alpha, 0.0
    for i in order:
        mass = min(float(w[i]), remaining)
        total += mass * values[i]
        remaining -= mass
        if remaining <= 1e-14:
            break
    return total / (1 - alpha)

mean_model = build_farmer([correlated_scenario(1.0)], check_units=True)
mean_model.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT)
solve_farmer(mean_model)
mean_plan = plan_of(mean_model)
print("Mean-yield acreage [wheat, corn, beets]:", mean_plan)
print(f"Profit: ${pyo.value(mean_model.expected_profit):,.2f}")
print(f"Land multiplier: ${mean_model.dual[mean_model.land_limit]:.2f}/acre")
for c in FEED_CROPS:
    con = mean_model.scenarios[0].feed_balance[c]
    print(f"Feed minimum multiplier, {c}: ${mean_model.dual[con]:.2f}/T")
assert np.allclose(mean_plan, [120, 80, 300])
assert np.isclose(pyo.value(mean_model.expected_profit), 118600)

perfect_models = [solve_farmer(build_farmer([scenario])) for scenario in THREE]

print("Known yield     wheat     corn    beets       profit ($)")
for label, model in zip(["above", "average", "below"], perfect_models):
    a = plan_of(model)
    print(
        f"{label:12s} {a[0]:8.1f} {a[1]:8.1f} {a[2]:8.1f}"
        f" {pyo.value(model.expected_profit):16,.2f}"
    )

# Validation first: reproduce Birge & Louveaux Table 5, p. 8.
RP = solve_farmer(build_farmer(THREE))
rp_plan = plan_of(RP)
rp_expected = pyo.value(RP.expected_profit)
rp_profits = np.array([pyo.value(RP.scenarios[i].profit) for i in RP.SCENARIOS])

print(
    f"acres  wheat {rp_plan[0]:.1f}  corn {rp_plan[1]:.1f}  beets {rp_plan[2]:.1f}"
    "   (book: 170, 80, 250)"
)
print(f"expected profit = ${rp_expected:,.0f}   (book: $108,390)")
print(
    "scenario profits  above ${:,.0f} | average ${:,.0f} | below ${:,.0f}".format(
        *rp_profits
    )
)

assert np.allclose(rp_plan, [170.0, 80.0, 250.0], atol=1e-6)
assert abs(rp_expected - 108390.0) < 1.0

# Exercise the unit declarations once, on this small model.
assert_units_consistent(build_farmer(THREE, check_units=True))
print("units consistent")

print("Scenario    wheat sold   corn bought   corn sold   premium beets   profit ($)")
for label, s in zip(["above", "average", "below"], RP.SCENARIOS):
    b = RP.scenarios[s]
    print(
        f"{label:8s} {pyo.value(b.sales['WHEAT']):13.1f}"
        f" {pyo.value(b.purchases['CORN']):13.1f} {pyo.value(b.sales['CORN']):11.1f}"
        f" {pyo.value(b.sales['BEETS_FAVORABLE']):15.1f} {pyo.value(b.profit):12,.0f}"
    )

WS = np.mean([pyo.value(m.expected_profit) for m in perfect_models])
EEV, ev_profits = grade_plan(mean_plan, THREE)

print("Policy                         expected profit ($)")
print(f"Perfect information (WS)        {WS:17,.2f}")
print(f"Shared stochastic plan (RP)     {rp_expected:17,.2f}")
print(f"Mean-yield plan evaluated (EEV) {EEV:17,.2f}")
print(f"EVPI = WS - RP: ${WS - rp_expected:,.2f}")
print(f"VSS  = RP - EEV: ${rp_expected - EEV:,.2f}")
print("EEV scenario profits [above, average, below]:", ev_profits)
assert np.isclose(EEV, 107240)
assert WS >= rp_expected >= EEV

import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import display

classical = {
    "known_mean": solution_record(mean_model),
    "perfect": [solution_record(m) for m in perfect_models],
    "rp": solution_record(RP),
    "eev": evaluate_solution(mean_plan, THREE),
    "inputs": THREE, "WS": float(WS), "RP": float(rp_expected), "EEV": float(EEV),
    "EVPI": float(WS-rp_expected), "VSS": float(rp_expected-EEV),
}

def full_solution_table(record):
    rows=[]
    for s in record["scenarios"]:
        rows.append({"Scenario":s["id"], "Probability":s["probability"],
            **{"Buy "+k+" (T)":v for k,v in s["purchases_tons"].items()},
            **{"Sell "+k+" (T)":v for k,v in s["sales_tons"].items()},
            "Profit ($)":s["profit"]})
    display(pd.Series(record["acreage"],name="Planted acres").to_frame().round(3))
    display(pd.DataFrame(rows).set_index("Scenario").round(2))
    print(f"Expected profit: ${record['expected_profit']:,.2f}")

for title,record in [("Mean yields",classical["known_mean"]),
    *[(name,r) for name,r in zip(["Known high","Known average","Known low"],classical["perfect"])],
    ("Shared stochastic plan",classical["rp"]),("Mean-yield plan evaluated",classical["eev"])]:
    print(title)
    full_solution_table(record)
results=classical
helper.save_results("farmer-information-values",results,
    notebook="notebooks/4-dev/SP.ipynb",source_tag="handout:farmer-root",
    solver="appsi_highs",description="Complete classical farmer solutions and information values",
    source_files=["notebooks/farmer.py"],
    source_cells=[("notebooks/4-dev/SP.ipynb", "provenance:farmer-information-solve")])

fig,axes=plt.subplots(1,2,figsize=(10,4),layout="constrained")
values=[results[k]/1000 for k in ["EEV","RP","WS"]]
colors=["#009E73","#0072B2","#D55E00"]
for i,(value,label,color) in enumerate(zip(values,["EEV","RP","WS"],colors)):
    axes[0].plot([i-.25,i+.25],[value,value],color=color,linewidth=3)
    axes[0].text(i,value+.4,f"{label}\n${value:,.2f}k",ha="center",fontsize=11)
for i,label in [(0,"VSS"),(1,"EVPI")]:
    axes[0].annotate("",xy=(i+.5,values[i+1]),xytext=(i+.5,values[i]),
                     arrowprops={"arrowstyle":"<->","color":"black"})
    axes[0].text(i+.57,(values[i]+values[i+1])/2,label,fontsize=10,va="center")
axes[0].set(xlim=(-.5,2.5),ylim=(105,119),ylabel="Expected profit ($1000)",xticks=[])
for label,key,marker,color in [("EEV","eev","s",colors[0]),("RP","rp","o",colors[1])]:
    p=[s["profit"]/1000 for s in results[key]["scenarios"]]
    axes[1].plot([0,1,2],p,marker=marker,label=label,color=color)
axes[1].plot([0,1,2],[r["expected_profit"]/1000 for r in results["perfect"]],
             marker="^",linestyle="--",color=colors[2],label="WS: perfect information")
axes[1].set(xticks=[0,1,2],xticklabels=["High","Average","Low"],
            xlabel="Harvest scenario",ylabel="Profit ($1000)")
fig.legend(*axes[1].get_legend_handles_labels(),loc="outside upper center",ncol=3,fontsize=11)
helper.save_figure(fig,"farmer-information-values")