Your First Optimization Problem

## Tip: Please put code like this at the top of your notebook.
# We want all of the module/package installations to start up front

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 pyomo.environ as pyo

# Create instance of concrete Pyomo model.
# concrete means all of the sets and model data are specified at the time of model construction.
# In this class, you'll use a concrete model.
m = pyo.ConcreteModel()

## Declare variables with initial values with bounds
m.x1 = pyo.Var(initialize=1, bounds=(-10, 10))
m.x2 = pyo.Var(initialize=1, bounds=(-10, 10))
m.x3 = pyo.Var(initialize=1, bounds=(-10, 10))

## Declare objective
m.OBJ = pyo.Objective(expr=m.x1**2 + 2 * m.x2**2 - m.x3, sense=pyo.minimize)

## Declare equality constraints
m.h1 = pyo.Constraint(expr=m.x1 + m.x2 == 1)

# Add your solution here

## Display model
m.pprint()

opt1 = pyo.SolverFactory("ipopt")
status1 = opt1.solve(m, tee=True)
assert pyo.check_optimal_termination(status1), (
    f"Solve failed: status={status1.solver.status}, "
    f"termination={status1.solver.termination_condition}"
)

## Return the solution
print("x1 = ", pyo.value(m.x1))
print("x2 = ", pyo.value(m.x2))
print("x3 = ", pyo.value(m.x3))
print("\n")

# ---------- EXTRACT --------------------------------------------------------
# Guess and check: sweep x3, solve the constraints for (x1, x2), evaluate the
# objective. Nothing is plotted here -- this cell only produces numbers.

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def constraints(x3):
    """Solve the linear constraints for x1 and x2, given a guess of x3.

    Args:
        x3: Value for the decision variable x3

    Returns:
        x1 and x2: Values calculated from the constraints
    """

    # Define the matrices in the above equations
    A = np.array([[1, 1], [1, 2]])
    b = np.array([1, 5 + x3])

    # Solve the linear system of equations
    z = np.linalg.solve(A, b)
    x1 = z[0]
    x2 = z[1]
    return x1, x2


def objective(x1, x2, x3):
    """Evaluate the objective function."""
    return x1**2 + 2 * x2**2 - x3


# Guess many values of x3.
x3_guesses = np.linspace(-10, 4, 29)
obj = [objective(*constraints(x3), x3) for x3 in x3_guesses]

# Everything the figure needs, as plain Python: the guess-and-check curve, and
# the point Pyomo/Ipopt found in the cells above.
results = {
    "sweep": helper.table(pd.DataFrame({"x3": x3_guesses, "obj": obj})),
    "solution": helper.extract(m, x1=m.x1, x2=m.x2, x3=m.x3, obj=m.OBJ),
}

# ---------- ARCHIVE --------------------------------------------------------
# figures/results/guess-and-check.json, committed to the repository, so the
# figure can be re-rendered without Pyomo or a solver. A no-op on Colab.
helper.save_results(
    "guess-and-check",
    results,
    notebook="notebooks/1-dev/Pyomo-Introduction.ipynb",
    description="Objective versus x3 along the one-dimensional feasible set of "
    "the first NLP, with the Pyomo/Ipopt solution marked.",
    solver="Ipopt via Pyomo",
)

print("x3 sweep:", len(x3_guesses), "points")
print("Pyomo solution:", results["solution"])

# ---------- PLOT -----------------------------------------------------------
# The plotting function takes the EXTRACTED results, not the Pyomo model. It
# needs no solver, so you can re-run this one cell as many times as it takes to
# get the labels where you want them.
#
# This cell is tagged `figure:guess-and-check`, which makes it the single
# source of the figure that also appears in the lecture handout: there is no
# second script re-deriving the curve somewhere else. See figures/README.md.


def plot_guess_and_check(results):
    """Objective versus x3 along the feasible curve, with the solver's answer marked."""
    sweep = helper.as_dataframe(results["sweep"])
    sln = results["solution"]

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

    ax.plot(sweep["x3"], sweep["obj"], zorder=2)
    ax.plot(
        [sln["x3"]],
        [sln["obj"]],
        marker="o",
        markersize=13,
        linestyle="",
        markerfacecolor="none",
        markeredgewidth=3,
        zorder=3,
    )

    # Direct labelling rather than a legend: a legend keyed only by colour is
    # unreadable in black and white, and both labels fit in empty parts of the
    # axes. See figures/README.md, "Course additions".
    ax.annotate("guess and check", xy=(-8.8, 122.0), fontsize=14, ha="left")
    ax.annotate(
        "solver\n" + f"$x_3 = {sln['x3']:.2f}$, $f = {sln['obj']:.2f}$",
        xy=(sln["x3"], sln["obj"]),
        xytext=(sln["x3"], 42.0),
        fontsize=14,
        ha="center",
        va="bottom",
        arrowprops=dict(arrowstyle="->", lw=1.5, color="black"),
    )

    ax.set_xlabel("$x_3$")
    ax.set_ylabel("$f(x)$")
    ax.set_xlim(-10.0, 4.0)
    ax.set_ylim(bottom=0.0)

    fig.tight_layout()
    return fig


fig = plot_guess_and_check(results)

# Writes media/figures/guess-and-check.{png,pdf} at 300 dpi. A no-op on Colab.
helper.save_figure(fig, "guess-and-check");

# Check the solver against the algebra above, not just against the picture.
x3_star = -7 / 2
assert abs(pyo.value(m.x3) - x3_star) < 1e-6, "x3 does not match the closed-form solution"
assert abs(pyo.value(m.x1) - 0.5) < 1e-6, "x1 does not match the closed-form solution"
assert abs(pyo.value(m.x2) - 0.5) < 1e-6, "x2 does not match the closed-form solution"
assert abs(pyo.value(m.OBJ) - 17 / 4) < 1e-6, "objective does not match 17/4"

print("Closed form:  x1 = 1/2, x2 = 1/2, x3 = -7/2, f = 17/4 = 4.25")
print(
    f"Ipopt:        x1 = {pyo.value(m.x1):.6f}, x2 = {pyo.value(m.x2):.6f}, "
    f"x3 = {pyo.value(m.x3):.6f}, f = {pyo.value(m.OBJ):.6f}"
)