Milk Pooling and Blending

# This code cell installs packages 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()

# `helper` also provides the extract / archive / figure plumbing used below.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pyomo.environ as pyo
from pyomo.environ import units as u
from pyomo.util.check_units import assert_units_consistent

# The source data is scale-free: milk is measured in "units" and priced per unit.
# Pyomo's unit library has neither, so declare both. Fat content stays
# dimensionless -- it is a mass fraction, kg of fat per kg of milk.
u.load_definitions_from_strings(["USD = [currency]", "milk = [milk]"])

customers = pd.DataFrame(
    {
        "Customer 1": {"min_fat": 0.045, "price": 52.0, "demand": 6000.0},
        "Customer 2": {"min_fat": 0.030, "price": 48.0, "demand": 2500.0},
        "Customer 3": {"min_fat": 0.040, "price": 50.0, "demand": 4000.0},
    }
).T

suppliers = pd.DataFrame(
    {
        "Farm A": {"fat": 0.045, "cost": 45.0, "location": "local"},
        "Farm B": {"fat": 0.030, "cost": 42.0, "location": "local"},
        "Farm C": {"fat": 0.033, "cost": 37.0, "location": "remote"},
        "Farm D": {"fat": 0.050, "cost": 45.0, "location": "remote"},
    },
).T

# Farm and customer names
local_farms = suppliers.index[suppliers["location"] == "local"].tolist()
remote_farms = suppliers.index[suppliers["location"] == "remote"].tolist()
all_farms = local_farms + remote_farms
customer_names = customers.index.tolist()

# Problem data as plain dictionaries, keyed by farm or customer name

# Fat fraction of the milk from farm s, phi_s [dimensionless]
fat = suppliers["fat"].astype(float).to_dict()

# Purchase cost per unit from farm s, kappa_s [USD/milk]
cost = suppliers["cost"].astype(float).to_dict()

# Price paid per unit by customer k, pi_k [USD/milk]
price = customers["price"].astype(float).to_dict()

# Maximum demand of customer k, d_k [milk]
demand = customers["demand"].astype(float).to_dict()

# Minimum fat fraction customer k accepts, phi_min_k [dimensionless]
min_fat = customers["min_fat"].astype(float).to_dict()

print("Customers")
display(customers)

print("Suppliers")
display(suppliers)

def create_milk_model(local, pooled):
    """Create the milk blending/pooling model in Pyomo

    One model covers all three options. The only thing that changes is which
    farms are reachable, and which of them must share the single tanker:

        Option 1  local=local_farms, pooled=[]            business as usual
        Option 2  local=all_farms,   pooled=[]            a second truck
        Option 3  local=local_farms, pooled=remote_farms  one truck, one pool

    Farms in `local` blend freely: their milk reaches each customer as its own
    stream z_{l,k}. Farms in `pooled` share one tanker, so their milk is first
    mixed into a single stream of composition p and only then blended. If
    `pooled` is empty there is no tanker, the pool variables are never
    declared, and every remaining term is linear -- the model is an LP.

    Arguments:
        local: list of farm names whose milk blends freely
        pooled: list of farm names whose milk must share one tanker

    Returns:
        m: Pyomo concrete model
    """
    m = pyo.ConcreteModel("Milk blending and pooling")

    # Freely blended farms, L in the notes
    m.L = pyo.Set(initialize=local)

    # Pooled farms, R in the notes. Empty for Options 1 and 2.
    m.R = pyo.Set(initialize=pooled)

    # All farms sourced from, L union R
    m.S = pyo.Set(initialize=list(local) + list(pooled))

    # Customers, K in the notes
    m.K = pyo.Set(initialize=customer_names)

    # Fat fraction of the milk from farm s, phi_s [dimensionless]
    m.fat = pyo.Param(m.S, initialize={s: fat[s] for s in m.S})

    # Purchase cost per unit from farm s, kappa_s [USD/milk]
    m.cost = pyo.Param(m.S, initialize={s: cost[s] for s in m.S}, units=u.USD / u.milk)

    # Price paid per unit by customer k, pi_k [USD/milk]
    m.price = pyo.Param(m.K, initialize=price, units=u.USD / u.milk)

    # Maximum demand of customer k, d_k [milk]
    m.demand = pyo.Param(m.K, initialize=demand, units=u.milk)

    # Minimum fat fraction customer k accepts, phi_min_k [dimensionless]
    m.min_fat = pyo.Param(m.K, initialize=min_fat)

    # Freely blended milk from farm l to customer k, z_{l,k} [milk]
    m.z = pyo.Var(m.L, m.K, domain=pyo.NonNegativeReals, units=u.milk)

    # THE POOL. Declared only when there is one.
    if len(m.R) > 0:
        # Milk bought from pooled farm r, x_r [milk]
        m.x = pyo.Var(m.R, domain=pyo.NonNegativeReals, units=u.milk)

        # Pooled milk delivered to customer k, y_k [milk]
        m.y = pyo.Var(m.K, bounds=lambda b, k: (0, b.demand[k]), units=u.milk)

        # Fat fraction of the pool, p [dimensionless]. A pool cannot be richer
        # than its richest input, nor leaner than its leanest.
        m.p = pyo.Var(bounds=(min(fat[r] for r in pooled), max(fat[r] for r in pooled)))

        # What is bought from the pooled farms is what leaves the pool [milk]
        @m.Constraint()
        def pool_balance(b):
            return sum(b.x[r] for r in b.R) == sum(b.y[k] for k in b.K)

        # Fat into the pool equals fat out. BILINEAR. [milk]
        #
        # The right-hand side is written p * sum_R x_r, while the written model
        # states it as p * sum_K y_k. Those are the same quantity ONLY because
        # pool_balance above forces the two sums to be equal -- what goes into
        # the pool is what leaves it. Delete that constraint and these two
        # forms stop agreeing.
        @m.Constraint()
        def pool_quality(b):
            return sum(b.fat[r] * b.x[r] for r in b.R) == b.p * sum(b.x[r] for r in b.R)

    # Helper expressions, written so the constraints below read the same way
    # whether or not a pool exists.
    def pooled_milk(b, k):
        return b.y[k] if len(b.R) > 0 else 0.0

    def pooled_fat(b, k):
        return b.p * b.y[k] if len(b.R) > 0 else 0.0

    # Maximize profit [USD]
    @m.Objective(sense=pyo.maximize)
    def profit(b):
        revenue_and_cost_local = sum(
            (b.price[k] - b.cost[l]) * b.z[l, k] for l in b.L for k in b.K
        )
        if len(b.R) == 0:
            return revenue_and_cost_local
        return (
            revenue_and_cost_local
            + sum(b.price[k] * b.y[k] for k in b.K)
            - sum(b.cost[r] * b.x[r] for r in b.R)
        )

    # Deliveries cannot exceed each customer's demand [milk]
    @m.Constraint(m.K)
    def demand_limit(b, k):
        return sum(b.z[l, k] for l in b.L) + pooled_milk(b, k) <= b.demand[k]

    # Each delivery meets its fat spec. BILINEAR when a pool exists. [milk]
    @m.Constraint(m.K)
    def blend_quality(b, k):
        fat_delivered = pooled_fat(b, k) + sum(b.fat[l] * b.z[l, k] for l in b.L)
        milk_delivered = sum(b.z[l, k] for l in b.L) + pooled_milk(b, k)
        return fat_delivered >= b.min_fat[k] * milk_delivered

    # Raises UnitsError if any constraint or the objective is inconsistent
    assert_units_consistent(m)

    return m

def model_size(m, label):
    """Count what Pyomo actually built, to check the analysis in Step 2.

    Arguments:
        m: a Pyomo concrete model
        label: name to print above the counts
    """
    constraints = list(m.component_data_objects(pyo.Constraint, active=True))
    n_var = sum(1 for _ in m.component_data_objects(pyo.Var, active=True))
    n_eq = sum(1 for c in constraints if c.equality)
    n_ineq = len(constraints) - n_eq

    print(label)
    print(f"  variables               {n_var:3d}")
    print(f"  equality constraints    {n_eq:3d}")
    print(f"  inequality constraints  {n_ineq:3d}")
    print(f"  degrees of freedom      {n_var - n_eq:3d}\n")

m1 = create_milk_model(local_farms, [])
model_size(m1, "Option 1: local farms only")

m2 = create_milk_model(all_farms, [])
model_size(m2, "Option 2: all farms, one truck each")

m3 = create_milk_model(local_farms, remote_farms)
model_size(m3, "Option 3: remote farms pooled")

solver = pyo.SolverFactory("ipopt")


def solve_and_report(m, title, opt=None):
    """Solve one option and print the supplier and customer tables.

    Arguments:
        m: a Pyomo concrete model from create_milk_model
        title: heading printed above the tables
        opt: solver to use; defaults to Ipopt with no options set

    Returns:
        the optimal profit
    """
    opt = solver if opt is None else opt
    results = opt.solve(m)
    assert pyo.check_optimal_termination(results), f"Solve failed for {title}"

    pooled = len(m.R) > 0

    print(f"{title}")
    if pooled:
        print(f"  pool composition p = {pyo.value(m.p):6.4f}")
    print(f"  profit             = {pyo.value(m.profit):,.2f}\n")

    # One row per farm: what its milk cost, and where it went.
    rows = []
    for s in m.S:
        row = {"fat": pyo.value(m.fat[s]), "cost": pyo.value(m.cost[s])}
        for k in m.K:
            row[k] = pyo.value(m.z[s, k]) if s in m.L else 0.0
        if pooled:
            row["Pool"] = pyo.value(m.x[s]) if s in m.R else 0.0
        rows.append(row)
    print("  Supplier report")
    display(pd.DataFrame(rows, index=list(m.S)).round(4))

    # One row per customer: what was delivered, and at what fat fraction.
    rows = []
    for k in m.K:
        row = {"min_fat": pyo.value(m.min_fat[k]), "demand": pyo.value(m.demand[k])}
        for l in m.L:
            row[l] = pyo.value(m.z[l, k])
        total = sum(pyo.value(m.z[l, k]) for l in m.L)
        fat_out = sum(pyo.value(m.z[l, k]) * pyo.value(m.fat[l]) for l in m.L)
        if pooled:
            row["Pool"] = pyo.value(m.y[k])
            total += pyo.value(m.y[k])
            fat_out += pyo.value(m.y[k]) * pyo.value(m.p)
        row["Total"] = total
        row["fat delivered"] = fat_out / total if total > 0 else np.nan
        rows.append(row)
    print("  Customer report")
    display(pd.DataFrame(rows, index=list(m.K)).round(4))

    return pyo.value(m.profit)

profit_1 = solve_and_report(m1, "Option 1, local farms only")

profit_2 = solve_and_report(m2, "Option 2, all farms, one truck each")

nlp_solver = pyo.SolverFactory("ipopt")

nlp_solver.options["bound_frac"] = 0.01
profit_3 = solve_and_report(m3, "Option 3, pooled, bound_frac = 0.01", nlp_solver)

nlp_solver.options["bound_frac"] = 0.5
m3_alt = create_milk_model(local_farms, remote_farms)
profit_3_alt = solve_and_report(
    m3_alt, "Option 3, pooled, bound_frac = 0.5", nlp_solver
)

# ---------- SOLVE ----------------------------------------------------------
# Fix p and every remaining term is linear: what is left is an LP, one per
# value of p. Build the Pyomo model ONCE and re-fix p between solves --
# rebuilding it at every grid point costs far more than the solves do.
#
# The grid is split at the two customer fat requirements that fall strictly
# inside the range of p, because the profit is not continuous there. Each
# segment stops just short of its right-hand breakpoint, so nothing draws a
# line across a jump that is not there.

# The model's own bounds on p: a pool is no richer than its richest input and
# no leaner than its leanest.
p_lo = min(fat[r] for r in remote_farms)  # 0.033
p_hi = max(fat[r] for r in remote_farms)  # 0.050

# Derived from the data, not hard-coded, so the sweep follows the instance if
# the instance ever changes.
breaks = sorted({v for v in min_fat.values() if p_lo < v < p_hi})  # 0.040, 0.045

m_sweep = create_milk_model(local_farms, remote_farms)


def profit_at_fixed_p(p_value):
    """Solve the pooling model with the pool composition p fixed. This is an LP."""
    m_sweep.p.fix(p_value)
    status = solver.solve(m_sweep)
    if not pyo.check_optimal_termination(status):
        return np.nan
    return pyo.value(m_sweep.profit)


eps = 1e-7
edges = [p_lo] + breaks + [p_hi]
segments = []
for lo, hi in zip(edges[:-1], edges[1:]):
    grid = np.linspace(lo, hi - eps, 400)
    segments.append((grid, np.array([profit_at_fixed_p(p) for p in grid])))

# The right-hand end of the last segment IS attained, so include it.
grid, vals = segments[-1]
segments[-1] = (np.append(grid, p_hi), np.append(vals, profit_at_fixed_p(p_hi)))

# The three local maxima: the lower bound on p, and each breakpoint.
maxima = [(p, profit_at_fixed_p(p)) for p in [p_lo] + breaks]

# The limit from the left at the jump -- a value the profit approaches and
# never attains.
jump_limit = profit_at_fixed_p(breaks[-1] - eps)


# Where did the two Ipopt runs above START? `bound_frac` pushes every variable
# off its bounds before the first iteration; read that pushed-in point back by
# solving with max_iter = 0. Pyomo warns that the solve did not converge --
# it was never meant to.
def initial_p(bound_frac):
    """Ipopt's pushed-in initial value of p, for a given bound_frac."""
    probe = pyo.SolverFactory("ipopt")
    probe.options["bound_frac"] = bound_frac
    probe.options["max_iter"] = 0
    m = create_milk_model(local_farms, remote_farms)
    probe.solve(m)
    return pyo.value(m.p)


ipopt_runs = [
    {
        "label": "bound_frac = 0.01",
        "bound_frac": 0.01,
        "p_start": initial_p(0.01),
        "p_land": pyo.value(m3.p),
        "profit_land": profit_3,
    },
    {
        "label": "bound_frac = 0.5",
        "bound_frac": 0.5,
        "p_start": initial_p(0.5),
        "p_land": pyo.value(m3_alt.p),
        "profit_land": profit_3_alt,
    },
]

print(f"{sum(len(g) for g, _ in segments) + 5} LP solves")
for p_star, f_star in maxima:
    print(f"  local maximum at p = {p_star:.3f}:  profit = {f_star:12,.2f}")
for run in ipopt_runs:
    print(
        f"  {run['label']:<20} starts at p = {run['p_start']:.6f}"
        f" and lands at p = {run['p_land']:.4f}"
    )

# ---------- EXTRACT --------------------------------------------------------
# Pyomo objects -> plain Python. After this cell nothing below touches a model
# or a solver: the plotting cell reads `results` and nothing else, so it can be
# re-run as many times as it takes to get the labels off each other without
# paying for 1205 solves again.

results = {
    "bounds": {"p_lo": float(p_lo), "p_hi": float(p_hi)},
    "breaks": [float(b) for b in breaks],
    # One table per continuous piece of the profit function.
    "segments": [
        helper.table(pd.DataFrame({"p": g, "profit": v})) for g, v in segments
    ],
    "maxima": helper.table(
        pd.DataFrame(maxima, columns=["p", "profit"])
    ),
    # The value approached from the left at the jump, and never attained.
    "jump_limit": {"p": float(breaks[-1]), "profit": float(jump_limit)},
    "ipopt_runs": helper.table(pd.DataFrame(ipopt_runs)),
}

# ---------- ARCHIVE --------------------------------------------------------
# figures/results/pooling-profit-vs-p.json, committed to the repo, so a change
# to the house style can redraw this figure with no solver installed.
# `source_tag` points at the model cell above, so scripts/check_results_fresh.py
# can tell you when the model has changed and these numbers have not. A no-op
# on Colab, where there is nothing to commit to.
helper.save_results(
    "pooling-profit-vs-p",
    results,
    notebook="notebooks/1-dev/Milk-Pooling.ipynb",
    source_tag="handout:milk-pooling-model",
    description="Profit of the milk pooling model as a function of the pool "
    "composition p, swept as an LP over three continuous pieces, with the "
    "three local maxima and the two Ipopt runs marked.",
    solver="Ipopt via Pyomo",
);

# The PLOTTING function: it takes the extracted results, not the Pyomo model.
#
# This cell is tagged `figure:pooling-profit-vs-p`, which makes it the single
# source of the figure in the lecture handout. See figures/README.md;
# figures/render_from_notebook.py re-runs exactly this cell against the
# archived JSON when the house style changes, so there is no second copy of the
# plotting code anywhere.


def plot_pooling_profit_vs_p(results):
    """Profit against the pool composition p: three local maxima, no smoothness.

    Two things the usual phrasing -- "three local maxima near 0.033, 0.040 and
    0.045" -- does not say, and the figure shows:

    * The BEST of the three sits ON the lower bound of p. It is not found by
      setting a derivative to zero.
    * The profit is neither smooth nor even continuous. The kink at p = 0.040
      is Customer 3's minimum fat fraction and the JUMP at p = 0.045 is
      Customer 1's: at p = 0.045 the pool on its own can finally satisfy the
      largest, best-paying customer, and 6,000 units of demand become servable
      in one step. It is drawn as a discontinuity -- two segments, an open
      circle for the limit from the left, a filled one for the value attained
      -- because joining them with a line would draw a claim that is false.

    The two Ipopt runs are marked where they START, on the axis. The second
    starts nearer the maximum at 0.040 than the one at 0.045 and walks to
    0.045 anyway: the starting point is twelve numbers, not one.

    Greyscale: the profit curve is the only series. Nothing is keyed by colour.
    """
    segments = [helper.as_dataframe(t) for t in results["segments"]]
    maxima = helper.as_dataframe(results["maxima"])
    runs = helper.as_dataframe(results["ipopt_runs"])
    p_lo = results["bounds"]["p_lo"]
    p_hi = results["bounds"]["p_hi"]
    jump = results["jump_limit"]

    blue = "#0072B2"  # Okabe-Ito

    fig, ax = plt.subplots(figsize=(7.2, 4.6))

    for seg in segments:
        ax.plot(seg["p"], seg["profit"], color=blue, linestyle="-", linewidth=2.4)

    # --- the three local maxima -------------------------------------------
    for _, row in maxima.iterrows():
        ax.plot([row["p"]], [row["profit"]], marker="o", markersize=9,
                color="black", zorder=5, linestyle="none")

    # Open circle: the limit from the left at the jump, which is NOT attained.
    ax.plot([jump["p"]], [jump["profit"]], marker="o", markersize=9,
            markerfacecolor="white", markeredgecolor="black",
            markeredgewidth=1.6, linestyle="none", zorder=5)

    # --- direct labels on the maxima --------------------------------------
    label_offsets = [(10, 4), (-4, 14), (8, 4)]
    label_align = ["left", "center", "left"]
    for (_, row), (dx, dy), ha in zip(maxima.iterrows(), label_offsets, label_align):
        ax.annotate(
            # {,} rather than a bare comma: mathtext sets a bare comma as
            # punctuation and inserts a space, giving "$102, 833".
            f"$p = {row['p']:.3f}$\n$\\${row['profit']:,.0f}$".replace(",", "{,}"),
            xy=(row["p"], row["profit"]),
            xytext=(dx, dy),
            textcoords="offset points",
            fontsize=11,
            ha=ha,
            va="bottom",
            linespacing=1.05,
        )

    ax.annotate(
        "the best of the three is\non the LOWER BOUND of $p$",
        xy=(p_lo, maxima["profit"].iloc[0]),
        xytext=(0.05, 0.60),
        textcoords="axes fraction",
        fontsize=11,
        ha="left",
        va="center",
        arrowprops=dict(arrowstyle="->", color="black", linewidth=1.1,
                        shrinkB=8.0),
    )

    # --- where the two Ipopt runs started, and where they went -------------
    y_lo, y_hi = ax.get_ylim()
    tick = y_lo + 0.035 * (y_hi - y_lo)
    for _, run in runs.iterrows():
        ax.plot([run["p_start"]], [tick], marker="^", markersize=8,
                color="black", linestyle="none", clip_on=False, zorder=6)
        ax.annotate("$p_0$", xy=(run["p_start"], tick), xytext=(0, 9),
                    textcoords="offset points", fontsize=10, ha="center",
                    va="bottom")
    ax.annotate(
        "starts here $\\rightarrow$ lands at $0.045$,\n"
        "though $0.040$ is nearer",
        xy=(runs["p_start"].iloc[1], tick),
        xytext=(0.05, 0.30),
        textcoords="axes fraction",
        fontsize=11,
        ha="left",
        va="center",
        arrowprops=dict(arrowstyle="->", color="black", linewidth=1.1,
                        shrinkB=6.0),
    )

    # Headroom at the top for the two labels that sit above their markers.
    y_lo2, y_hi2 = ax.get_ylim()
    ax.set_ylim(y_lo2, y_hi2 + 0.13 * (y_hi2 - y_lo2))

    ax.yaxis.set_major_formatter(
        plt.FuncFormatter(lambda v, _pos: f"{v:,.0f}")
    )
    ax.set_xlabel("pool composition $p$")
    ax.set_ylabel("profit [\\$]")
    ax.set_xlim(p_lo - 0.0004, p_hi + 0.0004)
    ax.set_xticks(np.arange(0.033, 0.0501, 0.003))

    fig.tight_layout()
    return fig


fig = plot_pooling_profit_vs_p(results)

# Write media/figures/pooling-profit-vs-p.{png,pdf} -- what the lecture handout
# \includegraphics. A no-op on Colab, where there is no repo to write to.
helper.save_figure(fig, "pooling-profit-vs-p");

summary = pd.DataFrame(
    [
        {
            "Option": "1. Local farms only",
            "L (blended)": ", ".join(local_farms),
            "R (pooled)": "(empty)",
            "Model class": "LP",
            "Profit": profit_1,
            "p": np.nan,
            "Catch": "leaves money on the table",
        },
        {
            "Option": "2. All farms, one truck each",
            "L (blended)": ", ".join(all_farms),
            "R (pooled)": "(empty)",
            "Model class": "LP",
            "Profit": profit_2,
            "p": np.nan,
            "Catch": "needs a second truck",
        },
        {
            "Option": "3. Remote farms pooled",
            "L (blended)": ", ".join(local_farms),
            "R (pooled)": ", ".join(remote_farms),
            "Model class": "bilinear NLP",
            "Profit": profit_3,
            "p": pyo.value(m3.p),
            "Catch": "nonconvex; bound_frac = 0.01",
        },
    ]
)
display(summary.round(4))

# Regression guard. These three numbers are what this notebook is for; if a
# change to the builder moves one of them, that is the thing to notice first.
# Option 3's value is the one reached from bound_frac = 0.01 -- the other local
# solution, 101,392.16, is equally valid output from the same model.
for got, want, label in [
    (profit_1, 81_000.00, "Option 1"),
    (profit_2, 122_441.18, "Option 2"),
    (profit_3, 102_833.33, "Option 3, bound_frac = 0.01"),
    (profit_3_alt, 101_392.16, "Option 3, bound_frac = 0.50"),
]:
    assert abs(got - want) < 0.01, f"{label}: {got:,.2f} != {want:,.2f}"
    print(f"{label:<28} {got:>12,.2f}")