Algorithms Homework 5

import matplotlib.pyplot as plt
import numpy as np
from scipy import linalg
from scipy.optimize import minimize

# Seed every random number generator (course style guide, section 8).
rng = np.random.default_rng(seed=0)

# Stand-in for "no bound". Using a finite sentinel keeps every array
# a plain float array, which makes the linear algebra below simpler.
BIG = 1.0e20

def solve_kkt(H, g, A, r):
    """
    Solve the equality constrained QP

        min_p   g' p + 1/2 p' H p    s.t.   A p + r = 0

    by solving its first order KKT system

        [ H   A' ] [ p  ]     [ g ]
        [ A   0  ] [ mu ]  = -[ r ].

    This is Eq. (5.5) of Biegler (2010), p. 93 -- the same linear system
    the equality constrained Newton method solved -- and it is the only
    linear algebra in the whole method.

    Inputs:
        H : Hessian of the QP objective, (n, n)
        g : gradient of the QP objective, (n,)
        A : constraint Jacobian, (m, n). Our row convention: A = J_h.
        r : constraint residual, (m,)

    Outputs:
        p  : QP step, (n,)
        mu : multipliers for A p + r = 0, (m,)
    """
    # Build K = [[H, A'], [A, 0]] and rhs = -[g; r], then solve.
    # Return the first n entries as p and the last m as mu.
    # Add your solution here

p_test, mu_test = solve_kkt(
    H=np.eye(2),
    g=np.zeros(2),
    A=np.array([[1.0, 1.0]]),
    r=np.array([-2.0]),
)
print(f"p  = {p_test}   (expected [1. 1.])")
print(f"mu = {mu_test}   (expected [-1.])")

assert np.allclose(p_test, [1.0, 1.0]), "KKT step is wrong"
assert np.allclose(mu_test, [-1.0]), "KKT multiplier is wrong"
print("\nKKT solver passes.")

def solve_bound_qp(H, c, A, b, lo, hi, z0, max_iter=200, tol=1e-9):
    """
    Primal active set method for the bound constrained QP

        min_z   1/2 z' H z + c' z
        s.t.    A z = b
                lo <= z <= hi

    started from a *feasible* point z0.

    The working set W holds the indices currently fixed at a bound. On
    every pass we fix those variables, solve the resulting equality
    constrained QP with `solve_kkt`, and then either

      * move along the step, adding whichever bound blocks first, or
      * if the step is zero, inspect the bound multipliers and release
        the bound whose multiplier has the wrong sign.

    This is Algorithm 16.3 of Nocedal & Wright (2006), specialized to
    bounds. It is deliberately simple: no warm start, no cycling
    safeguard beyond `max_iter`, dense factorizations throughout.

    Inputs:
        H, c    : QP objective, (n, n) and (n,). H must be positive definite.
        A, b    : equality constraints A z = b, (m, n) and (m,)
        lo, hi  : bounds, (n,) each
        z0      : feasible starting point, (n,)

    Outputs:
        z    : QP solution, (n,)
        mu   : multipliers for A z = b, (m,)
        zeta : multipliers for the bounds, (n,).
               zeta_i > 0 at a lower bound, zeta_i < 0 at an upper bound.
        nit  : number of active set passes
    """
    n = len(z0)
    m = A.shape[0]
    z = np.array(z0, dtype=float)

    # Working set: indices sitting on a bound at the starting point.
    W = set(i for i in range(n) if abs(z[i] - lo[i]) <= tol or abs(z[i] - hi[i]) <= tol)
    mu = np.zeros(m)

    for it in range(max_iter):
        free = sorted(set(range(n)) - W)
        g = H @ z + c

        # Step in the free variables only; fixed variables do not move.
        p = np.zeros(n)
        if free:
            p_free, mu = solve_kkt(
                H[np.ix_(free, free)], g[free], A[:, free], np.zeros(m)
            )
            p[free] = p_free
        else:
            mu = np.zeros(m)

        if linalg.norm(p, np.inf) <= tol:
            # Stationary for the current working set. Check the bound
            # multipliers and release one if its sign is wrong.
            zeta = H @ z + c + A.T @ mu
            worst, worst_val = None, tol
            for i in W:
                if hi[i] - lo[i] <= tol:
                    continue  # genuinely fixed variable, never released
                at_lo = abs(z[i] - lo[i]) <= tol
                at_hi = abs(z[i] - hi[i]) <= tol
                if at_lo and -zeta[i] > worst_val:
                    worst, worst_val = i, -zeta[i]
                elif at_hi and zeta[i] > worst_val:
                    worst, worst_val = i, zeta[i]
            if worst is None:
                return z, mu, zeta, it
            W.discard(worst)
        else:
            # Ratio test: how far can we move before hitting a bound?
            alpha, blocking = 1.0, None
            for i in free:
                if p[i] > tol and hi[i] < BIG:
                    a = (hi[i] - z[i]) / p[i]
                elif p[i] < -tol and lo[i] > -BIG:
                    a = (lo[i] - z[i]) / p[i]
                else:
                    continue
                if a < alpha:
                    alpha, blocking = a, i
            z = z + alpha * p
            if blocking is not None:
                W.add(blocking)

    zeta = H @ z + c + A.T @ mu
    return z, mu, zeta, max_iter

def solve_sqp_subproblem(g, B, h, J, d_lo, d_hi, M=1.0e4):
    """
    Assemble and solve the relaxed SQP subproblem, Biegler (6.15)-(6.16),
    pp. 138-139:

        min_{d, xi}   g' d + 1/2 d' B d + M (xi + xi^2 / 2)
        s.t.          h (1 - xi) + J d = 0
                      d_lo <= d <= d_hi,   0 <= xi <= 1

    Stacking z = [d; xi] turns this into exactly the bound constrained QP
    `solve_bound_qp` already handles: one extra variable, one extra
    column -h in the Jacobian, and one extra diagonal entry M.

    Inputs:
        g, B         : objective gradient and Hessian (approximation) at x^k
        h, J         : equality constraint residual and Jacobian at x^k
        d_lo, d_hi   : x_L - x^k and x_U - x^k
        M            : elastic penalty weight

    Outputs:
        d       : primal step, (n,)
        xi      : relaxation, scalar in [0, 1]. xi = 0 means the
                  linearization was consistent; xi = 1 means it could
                  only be satisfied by giving up entirely.
        lam_bar : equality multipliers from the QP, (m,)
        zeta    : bound multipliers for d, (n,)
        nit     : active set passes
    """
    # Build H_qp, c_qp, A_qp, b_qp, lo, hi and the feasible point z0.
    # Add your solution here

    z, lam_bar, zeta, nit = solve_bound_qp(H_qp, c_qp, A_qp, b_qp, lo, hi, z0)
    return z[:n], z[n], lam_bar, zeta[:n], nit

# A consistent linearization: min d1 + d2 + 1/2 ||d||^2 s.t. 1 + d1 + d2 = 0,
# no bounds. The relaxation should not be used, so xi = 0, and the step should
# match the plain equality constrained KKT solve.
g_t = np.array([1.0, 1.0])
B_t = np.eye(2)
h_t = np.array([1.0])
J_t = np.array([[1.0, 1.0]])
d_lo_t = -BIG * np.ones(2)
d_hi_t = BIG * np.ones(2)

d_e, xi_e, lam_e, zeta_e, _ = solve_sqp_subproblem(
    g_t, B_t, h_t, J_t, d_lo_t, d_hi_t, M=1.0e4
)
d_k, mu_k = solve_kkt(B_t, g_t, J_t, h_t)

print(f"elastic QP : d = {d_e}, xi = {xi_e:.3e}, lambda = {lam_e}")
print(f"plain KKT  : d = {d_k}, lambda = {mu_k}")

assert abs(xi_e) < 1e-8, "xi should be 0 when the linearization is consistent"
assert np.allclose(d_e, d_k, atol=1e-8), "elastic step should match the KKT step"
assert np.allclose(lam_e, mu_k, atol=1e-8), "multipliers should match"
print("\nSubproblem assembly passes.")

def damped_bfgs_update(B, s, y):
    """
    Damped BFGS update, Nocedal & Wright Procedure 18.2 and
    Eqs. (18.15)-(18.17), p. 537.

    Two things differ from the unconstrained BFGS of Algorithms 3:

      1. The secant pair is built from the *Lagrangian*, not the
         objective (that happens in the caller):
             s = x^{k+1} - x^k
             y = grad_x L(x^{k+1}, lam^{k+1}) - grad_x L(x^k, lam^{k+1})
      2. s'y <= 0 genuinely occurs, arbitrarily close to the solution,
         because grad^2_xx L need only be positive definite on the null
         space of the active constraint Jacobian. Instead of skipping
         the update we *damp* it.

    Set

        r_k     = theta_k y_k + (1 - theta_k) B^k s_k

        theta_k = 1                              if s'y >= 0.2 s'Bs
                = 0.8 s'Bs / (s'Bs - s'y)        otherwise

    and apply ordinary BFGS with y replaced by r. theta = 1 is plain
    BFGS; theta = 0 gives r = B s, for which the update leaves B
    unchanged.

    Inputs:
        B : current Hessian approximation, (n, n), positive definite
        s : primal step, (n,)
        y : change in the gradient of the Lagrangian, (n,)

    Outputs:
        B_new : updated approximation, (n, n)
        theta : the damping weight actually used
    """
    # Compute Bs, s'Bs, s'y; then theta, then r; then the update.
    # Add your solution here

B0 = np.eye(3)

# Case 1: good curvature. Expect theta = 1 and the secant condition B s = y.
s_a = np.array([1.0, 0.0, 0.0])
y_a = np.array([2.0, 0.5, 0.0])
B_a, theta_a = damped_bfgs_update(B0, s_a, y_a)
print(f"case 1: theta = {theta_a:.6f} (expected 1.0)")
print(f"        B s - y = {B_a @ s_a - y_a}")
assert abs(theta_a - 1.0) < 1e-12
assert np.allclose(B_a @ s_a, y_a, atol=1e-10), "secant condition violated"

# Case 2: negative curvature, s'y < 0. Expect 0 < theta < 1, B s = r, B pd.
s_b = np.array([1.0, 0.0, 0.0])
y_b = np.array([-1.0, 0.0, 0.0])
B_b, theta_b = damped_bfgs_update(B0, s_b, y_b)
r_b = theta_b * y_b + (1 - theta_b) * (B0 @ s_b)
print(f"\ncase 2: theta = {theta_b:.6f} (expected 0.8 / (1 - (-1)) = 0.4)")
print(f"        s'y = {s_b @ y_b}, s'r = {s_b @ r_b} (must be > 0)")
print(f"        eigenvalues of B = {np.round(linalg.eigvalsh(B_b), 8)}")
assert 0.0 < theta_b < 1.0
assert np.allclose(B_b @ s_b, r_b, atol=1e-10), "damped secant condition violated"
assert np.all(linalg.eigvalsh(B_b) > 0), "damped BFGS must stay positive definite"
print("\nDamped BFGS passes.")

def convexify(W, delta0=1e-4, max_tries=40):
    """
    Return W + delta I with the smallest delta in {0, 1e-4, 1e-3, ...}
    that makes the matrix positive definite (tested by Cholesky).

    The QP subproblem needs positive curvature on the null space of the
    active constraints. This is a crude substitute: a production code
    corrects the *inertia* of the KKT matrix instead, which is weaker and
    therefore cheaper. See Biegler section 6.2.1, p. 137.
    """
    delta = 0.0
    n = W.shape[0]
    for _ in range(max_tries):
        try:
            linalg.cholesky(W + delta * np.eye(n))
            return W + delta * np.eye(n)
        except linalg.LinAlgError:
            delta = delta0 if delta == 0.0 else 10.0 * delta
    return np.eye(n)

def sqp(
    calc_f,
    calc_grad,
    calc_h,
    calc_jac,
    x0,
    x_lo=None,
    x_hi=None,
    calc_hess_lag=None,
    hessian="bfgs",  # "bfgs" (damped) or "exact"
    max_iter=100,
    eps_d=1e-8,  # step norm tolerance, epsilon_1
    eps_h=1e-8,  # constraint tolerance, epsilon_2
    eta=1e-4,  # Armijo parameter
    tau=0.5,  # backtracking factor
    delta_rho=0.1,  # penalty safety margin
    M=1.0e4,  # elastic penalty weight
    soc=False,  # apply a second order correction?
    verbose=False,
):
    """
    Line search SQP -- Algorithm 6.1 of Biegler (2010), p. 141 -- for

        min_x  f(x)   s.t.   h(x) = 0,   x_L <= x <= x_U

    using the l1 merit function

        phi_1(x; rho) = f(x) + rho ||h(x)||_1

    whose directional derivative along the QP step is, Biegler (6.19),

        D phi_1 = grad f(x^k)' d - rho (1 - xi) ||h(x^k)||_1.

    Inputs:
        calc_f, calc_grad   : objective and its gradient
        calc_h, calc_jac    : equality constraints h(x) and Jacobian J_h
                              (m x n -- our row convention, not Biegler's)
        x0                  : starting point
        x_lo, x_hi          : bounds (None means unbounded)
        calc_hess_lag       : grad^2_xx L(x, lam), needed for hessian="exact"
        hessian             : "bfgs" (damped, the default) or "exact"

    Outputs:
        a dictionary with the solution, multipliers, status and history
    """
    n = len(x0)
    x = np.array(x0, dtype=float)
    x_lo = -BIG * np.ones(n) if x_lo is None else np.array(x_lo, dtype=float)
    x_hi = BIG * np.ones(n) if x_hi is None else np.array(x_hi, dtype=float)

    # Every iterate stays inside the box, so the merit function never
    # needs a bound term.
    x = np.clip(x, x_lo, x_hi)

    lam = np.zeros(len(np.atleast_1d(calc_h(x))))
    B = np.eye(n)
    rho = 0.0
    status = "maximum iterations"
    hist = {
        "x": [x.copy()],
        "f": [],
        "h": [],
        "alpha": [],
        "xi": [],
        "rho": [],
        "soc": [],
        "theta": [],
    }

    for k in range(max_iter):
        # ---- Step 1: evaluate everything at the current iterate ------
        f = calc_f(x)
        g = calc_grad(x)
        h = np.atleast_1d(calc_h(x))
        J = np.atleast_2d(calc_jac(x))
        hist["f"].append(f)
        hist["h"].append(linalg.norm(h, np.inf))

        if hessian == "exact":
            B = convexify(calc_hess_lag(x, lam))

        # ---- Step 2: solve the QP subproblem -------------------------
        d_lo = np.maximum(x_lo - x, -BIG)
        d_hi = np.minimum(x_hi - x, BIG)
        d, xi, lam_bar, zeta, _ = solve_sqp_subproblem(g, B, h, J, d_lo, d_hi, M)
        hist["xi"].append(xi)

        if xi > 1.0 - 1e-8:
            status = "inconsistent linearization (xi = 1)"
            break

        # Stopping test. Theorem 6.1 (Biegler, p. 136): d = 0 is not
        # merely "no progress", it certifies a KKT point.
        if linalg.norm(d, np.inf) <= eps_d and linalg.norm(h, np.inf) <= eps_h:
            lam = lam_bar
            status = "converged"
            break

        # ---- Step 4: update the penalty parameter --------------------
        # Biegler's rule is rho^k = max(rho^{k-1}, ||lam_bar||_q + delta)
        # with 1/p + 1/q = 1, so q = infinity for the l1 merit function.
        #
        # We add one safeguard he does not: when the subproblem had to be
        # relaxed (xi large), lam_bar is a shadow price on the *elastic*
        # constraint and is O(M), not an estimate of the NLP multiplier.
        # Feeding it into a rule that never decreases rho poisons the
        # merit function for the rest of the solve. So in that case we
        # ask only for enough penalty to make D phi_1 negative.
        h_norm1 = linalg.norm(h, 1)
        if h_norm1 > 0.0:
            rho_descent = (g @ d) / ((1.0 - xi) * h_norm1) + delta_rho
        else:
            rho_descent = 0.0
        if xi < 0.5:
            rho = max(rho, linalg.norm(lam_bar, np.inf) + delta_rho, rho_descent)
        else:
            rho = max(rho, rho_descent)
        hist["rho"].append(rho)

        # ---- Steps 3, 5, 6: backtracking line search -----------------
        phi0 = f + rho * h_norm1
        Dphi = g @ d - rho * (1.0 - xi) * h_norm1

        alpha = 1.0
        used_soc = False
        for _ in range(60):
            x_trial = np.clip(x + alpha * d, x_lo, x_hi)
            phi_trial = calc_f(x_trial) + rho * linalg.norm(
                np.atleast_1d(calc_h(x_trial)), 1
            )
            if phi_trial <= phi0 + eta * alpha * Dphi:
                break

            # Second order correction, tried once, on the full step only.
            if soc and alpha == 1.0:
                h_trial = np.atleast_1d(calc_h(x + d))
                d_bar = h_trial - J @ d
                d_soc, _ = solve_kkt(B, g, J, d_bar)
                x_soc = np.clip(x + d_soc, x_lo, x_hi)
                phi_soc = calc_f(x_soc) + rho * linalg.norm(
                    np.atleast_1d(calc_h(x_soc)), 1
                )
                if phi_soc <= phi0 + eta * Dphi:
                    x_trial, d, used_soc = x_soc, d_soc, True
                    break

            alpha *= tau

        hist["alpha"].append(alpha)
        hist["soc"].append(used_soc)

        # ---- Quasi-Newton update ------------------------------------
        x_new = x_trial
        lam_new = lam_bar
        s = x_new - x
        # Both terms use lam^{k+1}: we are running quasi-Newton on
        # x -> L(x, lam) with lam held fixed.
        y = (calc_grad(x_new) + np.atleast_2d(calc_jac(x_new)).T @ lam_new) - (
            g + J.T @ lam_new
        )

        if hessian == "bfgs":
            B, theta = damped_bfgs_update(B, s, y)
            hist["theta"].append(theta)

        x, lam = x_new, lam_new
        hist["x"].append(x.copy())

        if verbose:
            print(
                f"{k:3d}  f = {calc_f(x): .8f}  |h|_inf = "
                f"{linalg.norm(np.atleast_1d(calc_h(x)), np.inf):.3e}  "
                f"alpha = {alpha:.4f}  xi = {xi:.4f}  rho = {rho:.4f}"
                + ("  [SOC]" if used_soc else "")
            )

    return {
        "x": x,
        "lam": lam,
        "f": calc_f(x),
        "h": np.atleast_1d(calc_h(x)),
        "status": status,
        "n_iter": len(hist["f"]),
        "hist": hist,
    }

def f1(x):
    return x[0] + x[1]


def grad1(x):
    return np.array([1.0, 1.0])


def h1(x):
    return np.array([x[0] ** 2 + x[1] ** 2 - 2.0])


def jac1(x):
    return np.array([[2 * x[0], 2 * x[1]]])


def hess_lag1(x, lam):
    # grad^2 f = 0; grad^2 h = 2 I
    return 2 * lam[0] * np.eye(2)


prob1 = dict(
    calc_f=f1, calc_grad=grad1, calc_h=h1, calc_jac=jac1, calc_hess_lag=hess_lag1
)

x_star_1 = np.array([-1.0, -1.0])

for hess in ["exact", "bfgs"]:
    res = sqp(**prob1, x0=[-2.0, -1.0], hessian=hess, eps_d=1e-12, eps_h=1e-12)
    print(
        f"{hess:>6s}: x = {res['x']}, f = {res['f']:.12f}, "
        f"lambda = {res['lam']}, {res['n_iter']} iterations, {res['status']}"
    )
    print(f"        ||x - x*|| = {linalg.norm(res['x'] - x_star_1):.3e}")

ref = minimize(
    f1,
    [-2.0, -1.0],
    jac=grad1,
    method="SLSQP",
    constraints=[{"type": "eq", "fun": h1, "jac": jac1}],
    options={"ftol": 1e-12},
)
assert ref.success, f"SLSQP failed: {ref.message}"
print(f" SLSQP: x = {ref.x}, f = {ref.fun:.12f}")
print(f"exact : x* = {x_star_1}, f* = -2, lambda* = 0.5")

for hess in ["exact", "bfgs"]:
    res = sqp(**prob1, x0=[-2.0, -1.0], hessian=hess, eps_d=1e-12, eps_h=1e-12)
    err = [linalg.norm(xk - x_star_1) for xk in res["hist"]["x"]]
    print(f"{hess:>6s}: " + "  ".join(f"{e:.2e}" for e in err))

def f2(x):
    return (x[0] - 2) ** 2 + (x[1] - 2) ** 2


def grad2(x):
    return np.array([2 * (x[0] - 2), 2 * (x[1] - 2)])


def h2(x):
    return np.array([x[0] ** 2 + x[1] ** 2 - 2.0])


def jac2(x):
    return np.array([[2 * x[0], 2 * x[1]]])


def hess_lag2(x, lam):
    return 2 * np.eye(2) + 2 * lam[0] * np.eye(2)


prob2 = dict(
    calc_f=f2, calc_grad=grad2, calc_h=h2, calc_jac=jac2, calc_hess_lag=hess_lag2
)

x2_lo = np.array([-BIG, 0.0])
x2_hi = np.array([BIG, 0.8])

# Analytic solution: the bound x2 <= 0.8 is active, so
# x2 = 0.8 and x1 = sqrt(2 - 0.64).
x_star_2 = np.array([np.sqrt(2 - 0.8**2), 0.8])
lam_star_2 = -grad2(x_star_2)[0] / jac2(x_star_2)[0, 0]
zeta_star_2 = -(grad2(x_star_2)[1] + lam_star_2 * jac2(x_star_2)[0, 1])
print(f"x*      = {x_star_2}")
print(f"f*      = {f2(x_star_2):.12f}")
print(f"lambda* = {lam_star_2:.12f}   (equality multiplier)")
print(f"zeta*   = {zeta_star_2:.12f}   (multiplier on x2 <= 0.8, must be >= 0)")

for hess in ["exact", "bfgs"]:
    res = sqp(
        **prob2,
        x0=[1.4, 0.2],
        x_lo=x2_lo,
        x_hi=x2_hi,
        hessian=hess,
        eps_d=1e-12,
        eps_h=1e-12,
    )
    print(
        f"{hess:>6s}: x = {res['x']}, f = {res['f']:.12f}, "
        f"lambda = {res['lam']}, {res['n_iter']} iterations, {res['status']}"
    )
    print(f"        ||x - x*|| = {linalg.norm(res['x'] - x_star_2):.3e}")

ref = minimize(
    f2,
    [1.4, 0.2],
    jac=grad2,
    method="SLSQP",
    bounds=[(None, None), (0.0, 0.8)],
    constraints=[{"type": "eq", "fun": h2, "jac": jac2}],
    # SLSQP reports "positive directional derivative for linesearch" on this
    # problem for ftol below about 1e-11 -- it reaches the same point, but its
    # own convergence test cannot certify it. Our step-norm test can.
    options={"ftol": 1e-10, "maxiter": 300},
)
assert ref.success, f"SLSQP failed: {ref.message}"
print(f" SLSQP: x = {ref.x}, f = {ref.fun:.12f}")

def f3(x):
    return x[0] + x[1]


def grad3(x):
    return np.array([1.0, 1.0, 0.0, 0.0])


def h3(x):
    return np.array(
        [
            1 + x[0] - x[1] ** 2 + x[2],
            1 - x[0] - x[1] ** 2 + x[3],
        ]
    )


def jac3(x):
    return np.array(
        [
            [1.0, -2 * x[1], 1.0, 0.0],
            [-1.0, -2 * x[1], 0.0, 1.0],
        ]
    )


def hess_lag3(x, lam):
    return np.diag([0.0, -2 * lam[0] - 2 * lam[1], 0.0, 0.0])


prob3 = dict(
    calc_f=f3, calc_grad=grad3, calc_h=h3, calc_jac=jac3, calc_hess_lag=hess_lag3
)

x3_lo = np.array([-BIG, 0.0, 0.0, 0.0])
x3_hi = np.array([BIG, 2.0, BIG, BIG])

x_star_3 = np.array([-3.0, 2.0, 6.0, 0.0])
print(f"stated global solution x* = {x_star_3}")
print(f"  h(x*) = {h3(x_star_3)}  (both zero, so x* is feasible)")
print(f"  f(x*) = {f3(x_star_3)}")

x0 = np.array([0.0, 0.1, 0.0, 0.0])

print(f"h(x0) = {h3(x0)}")
print("The two linearized rows at x0 are")
print("    0.99 + d1 - 0.2 d2 + d3 = 0")
print("    0.99 - d1 - 0.2 d2 + d4 = 0")
print("with -0.1 <= d2 <= 1.9 and d3, d4 >= 0. Adding them,")
print("    0.4 d2 = 1.98 + d3 + d4 >= 1.98,  i.e.  d2 >= 4.95 > 1.9.")
print("Inconsistent: the un-relaxed subproblem (6.11) has no feasible point.\n")

# Relaxing by (1 - xi) needs 1.98 (1 - xi) <= 0.4 (1.9) = 0.76,
# i.e. xi >= 1 - 0.76/1.98.
xi_hand = 1.0 - 0.76 / 1.98
print(f"Hand calculation: a feasible region first appears at xi = {xi_hand:.6f}")

d, xi, lam_bar, zeta, nit = solve_sqp_subproblem(
    grad3(x0), np.eye(4), h3(x0), jac3(x0), x3_lo - x0, x3_hi - x0, M=1.0e4
)
print(f"Elastic QP (6.15)-(6.16), B = I, M = 1e4:")
print(f"    xi = {xi:.6f}")
print(f"    d  = {np.round(d, 6)}")
print(f"    xi - xi_hand = {xi - xi_hand:.2e}")

for x0 in [[0.0, 0.1, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]:
    res = sqp(**prob3, x0=x0, x_lo=x3_lo, x_hi=x3_hi, hessian="bfgs", max_iter=100)
    print(f"x0 = {x0}")
    print(f"    x = {np.round(res['x'], 8)}, f = {res['f']:.10f}")
    print(f"    {res['n_iter']} iterations, status: {res['status']}")
    print(f"    xi history: {np.round(res['hist']['xi'][:4], 6)}")

    ref = minimize(
        f3,
        x0,
        jac=grad3,
        method="SLSQP",
        bounds=[(None, None), (0, 2), (0, None), (0, None)],
        constraints=[{"type": "eq", "fun": h3, "jac": jac3}],
        options={"ftol": 1e-12, "maxiter": 300},
    )
    print(
        f"    SLSQP: x = {np.round(ref.x, 8)}, f = {ref.fun:.10f}, "
        f"success = {ref.success}"
    )
    print()

def f_m(x):
    return 2 * (x[0] ** 2 + x[1] ** 2 - 1) - x[0]


def grad_m(x):
    return np.array([4 * x[0] - 1.0, 4 * x[1]])


def h_m(x):
    return np.array([x[0] ** 2 + x[1] ** 2 - 1.0])


def jac_m(x):
    return np.array([[2 * x[0], 2 * x[1]]])


def hess_lag_m(x, lam):
    return (4 + 2 * lam[0]) * np.eye(2)


prob_m = dict(
    calc_f=f_m, calc_grad=grad_m, calc_h=h_m, calc_jac=jac_m, calc_hess_lag=hess_lag_m
)

x_star_m = np.array([1.0, 0.0])
print(f"grad f(x*)  = {grad_m(x_star_m)}")
print(f"J_h(x*)     = {jac_m(x_star_m)}")
print(f"lambda*     = {-grad_m(x_star_m)[0] / jac_m(x_star_m)[0, 0]}  (should be -3/2)")
print(f"grad^2_xx L = \n{hess_lag_m(x_star_m, [-1.5])}  (should be the identity)")

vartheta = 0.4
xk = np.array([np.cos(vartheta), np.sin(vartheta)])

# Closed form, N&W (15.35), p. 441.
d_closed = np.array([np.sin(vartheta) ** 2, -np.sin(vartheta) * np.cos(vartheta)])

# The same step from the KKT system, with grad^2_xx L = I.
d_qp, mu_qp = solve_kkt(np.eye(2), grad_m(xk), jac_m(xk), h_m(xk))

print(f"closed form d = {d_closed}")
print(f"KKT solve   d = {d_qp}")
print(f"difference    = {linalg.norm(d_closed - d_qp):.3e}\n")

x_trial = xk + d_closed
print(f"f(x^k)        = {f_m(xk): .10f}")
print(
    f"f(x^k + d)    = {f_m(x_trial): .10f}   (sin^2 - cos = "
    f"{np.sin(vartheta) ** 2 - np.cos(vartheta): .10f})"
)
print(f"h(x^k)        = {h_m(xk)[0]: .10f}")
print(
    f"h(x^k + d)    = {h_m(x_trial)[0]: .10f}   (sin^2 = "
    f"{np.sin(vartheta) ** 2: .10f})"
)
print("\nBoth the objective AND the constraint violation increase, so every")
print("merit function of the form f + rho * eta(h) rejects the step.\n")

num = linalg.norm(x_trial - x_star_m)
den = linalg.norm(xk - x_star_m) ** 2
print(f"||x^k + d - x*|| / ||x^k - x*||^2 = {num / den:.12f}   (exactly 1/2)")

print(f"{'vartheta':>10s} {'ratio':>12s} {'f increase':>14s} {'|h| increase':>14s}")
for vartheta in [0.8, 0.4, 0.2, 0.1, 0.05, 0.01]:
    xk = np.array([np.cos(vartheta), np.sin(vartheta)])
    d = np.array([np.sin(vartheta) ** 2, -np.sin(vartheta) * np.cos(vartheta)])
    xt = xk + d
    ratio = linalg.norm(xt - x_star_m) / linalg.norm(xk - x_star_m) ** 2
    print(
        f"{vartheta:10.3f} {ratio:12.8f} {f_m(xt) - f_m(xk):14.3e} "
        f"{abs(h_m(xt)[0]) - abs(h_m(xk)[0]):14.3e}"
    )

vartheta = 0.4
xk = np.array([np.cos(vartheta), np.sin(vartheta)])
d = np.array([np.sin(vartheta) ** 2, -np.sin(vartheta) * np.cos(vartheta)])

# d_bar = h(x^k + d) - J_h^k d.
d_bar = h_m(xk + d) - jac_m(xk) @ d
d_soc, _ = solve_kkt(np.eye(2), grad_m(xk), jac_m(xk), d_bar)

print(f"d_bar   = {d_bar}")
print(f"d_soc   = {d_soc}\n")
print(f"{'trial point':>28s} {'x':>34s} {'f':>14s} {'|h|':>12s}")
for name, xt in [
    ("x^k (current iterate)", xk),
    ("x^k + d (SQP step)", xk + d),
    ("x^k + d_soc (corrected)", xk + d_soc),
]:
    print(
        f"{name:>28s} {str(np.round(xt, 8)):>34s} {f_m(xt):14.8f} "
        f"{abs(h_m(xt)[0]):12.3e}"
    )

print(f"{'vartheta':>10s} {'SOC':>6s} {'iters':>7s}  step lengths (first five)")
for vartheta in [0.1, 0.2, 0.4]:
    x0 = np.array([np.cos(vartheta), np.sin(vartheta)])
    for use_soc in [False, True]:
        res = sqp(**prob_m, x0=x0, hessian="exact", soc=use_soc, max_iter=60)
        assert res["status"] == "converged", res["status"]
        assert linalg.norm(res["x"] - x_star_m) < 1e-6
        print(
            f"{vartheta:10.2f} {str(use_soc):>6s} {res['n_iter']:7d}  "
            f"{np.round(res['hist']['alpha'][:5], 4)}"
        )

ref = minimize(
    f_m,
    [np.cos(0.4), np.sin(0.4)],
    jac=grad_m,
    method="SLSQP",
    constraints=[{"type": "eq", "fun": h_m, "jac": jac_m}],
    options={"ftol": 1e-14},
)
assert ref.success, f"SLSQP failed: {ref.message}"
print(f"\nSLSQP reference: x = {ref.x}, f = {ref.fun:.12f}")
print(f"analytic       : x* = {x_star_m}, f* = -1.0, lambda* = -1.5")

fig, ax = plt.subplots(figsize=(5.5, 5.5))

# The feasible set: the unit circle.
t = np.linspace(0, 2 * np.pi, 400)
ax.plot(np.cos(t), np.sin(t), "k-", lw=1, label="$h(x) = 0$")

for use_soc, style, marker in [(False, "--", "o"), (True, "-", "s")]:
    res = sqp(
        **prob_m,
        x0=np.array([np.cos(0.4), np.sin(0.4)]),
        hessian="exact",
        soc=use_soc,
        max_iter=60,
    )
    path = np.array(res["hist"]["x"])
    ax.plot(
        path[:, 0],
        path[:, 1],
        style,
        marker=marker,
        ms=5,
        label=f"SQP, SOC = {use_soc} ({res['n_iter']} iters)",
    )

ax.plot(1.0, 0.0, "k*", ms=14, label="$x^*$")
ax.set_xlabel("$x_1$")
ax.set_ylabel("$x_2$")
ax.set_xlim(0.6, 1.25)
ax.set_ylim(-0.25, 0.55)
ax.set_aspect("equal")
ax.grid(True)
ax.legend(loc="upper right", fontsize=9)
plt.tight_layout()
plt.show()