lecture_1/optimization.py

"""Unconstrained descent, Lagrange multipliers, maximum entropy and the KKT conditions.

A differentiable minimum has a zero gradient, and a descent step moves against that gradient. An
equality constraint comes with a multiplier, and at the constrained optimum the two gradients are
parallel. We can apply that condition to entropy under an expectation constraint, and an
exponential family follows. An inequality comes with a sign requirement and complementary
slackness. Under the softmax the simplex constraints hold by construction.

Chapter 1, Sections 1.36, 1.37 and 1.38 of the notes. Runs on the quadratic loss of
Section 1.36, the constrained distance of Section 1.37.1, the GC-content target of
Section 1.37.3 and the bounded quadratic of Section 1.38.

    python lecture_1/optimization.py              # every worked example
    python lecture_1/optimization.py --steps 20   # a short descent, far from the minimizer
    python lecture_1/optimization.py --quiet      # the closing report only
"""

# %% 01. Imports and configuration
from __future__ import annotations

import argparse
import math

import torch

from cis6270.runner import print_report, save_run, set_seed

torch.set_default_dtype(torch.float64)  # Double precision, because the chapter prints five decimal places.


BASE_COUNTS = torch.tensor([4.0, 3.0, 2.0, 1.0])  # n(A), n(C), n(G), n(T) over ten DNA bases.
GC_INDICATOR = torch.tensor([0.0, 1.0, 1.0, 0.0])  # h(x) = 1 on the strong bases C and G.
GC_TARGET = 0.6  # The required expected GC content, the constraint value c of Section 1.37.3.


WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.


def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
    if not quiet:  # Under --quiet only the closing report is printed.
        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
              f"   gap {abs(by_hand - computed):.1e}")


def entropy(p: torch.Tensor) -> float:
    """H(p) = -sum_i p_i log p_i, the quantity maximized in both constrained problems below."""
    return float(-(p * p.clamp_min(1e-300).log()).sum())  # Clamp keeps 0 log 0 at zero.


# %% 02. Optimization without constraints
def unconstrained(steps: int, quiet: bool = False) -> dict:
    """Section 1.36: on L(theta) = (theta - 3)^2 the error contracts by 0.8 at every step."""
    learning_rate, theta = 0.1, 0.0  # The step size eta and the starting parameter theta_0.
    history = []  # The iterates, so we can compare the first three with the chapter.
    for _ in range(max(steps, 1)):
        theta = theta - learning_rate * 2.0 * (theta - 3.0)  # theta - eta L'(theta) = 0.8 theta + 0.6.
        history.append(theta)
    closed_form = 3.0 - 3.0 * 0.8 ** len(history)  # theta_k = 3 - 3(0.8)^k, from the contraction.
    for index, value in ((0, 0.600), (1, 1.080), (2, 1.464)):
        if index < len(history):  # The chapter prints theta_1, theta_2 and theta_3.
            show(f"theta_{index + 1}", value, history[index], quiet)
    show(f"theta_{len(history)} from the closed form", closed_form, history[-1], quiet)
    show("gradient at the minimizer theta = 3", 0.0, 2.0 * (3.0 - 3.0), quiet)
    return {"history": history, "closed_form": closed_form}


# %% 03. An equality constraint and its multiplier
def equality_constraint(quiet: bool = False) -> dict:
    """Section 1.37.1: minimize x_1^2 + x_2^2 on the line x_1 + x_2 = 1, through the Lagrangian."""
    system = torch.tensor([[2.0, 0.0, 1.0],  # 2 x_1 + lambda = 0, the first stationarity condition.
                           [0.0, 2.0, 1.0],  # 2 x_2 + lambda = 0, the second.
                           [1.0, 1.0, 0.0]])  # x_1 + x_2 = 1, the constraint itself.
    solution = torch.linalg.solve(system, torch.tensor([0.0, 0.0, 1.0]))  # (x_1, x_2, lambda); [3].
    optimum = float(solution[:2].square().sum())  # f at the constrained minimizer.
    endpoint = float(torch.tensor([1.0, 0.0]).square().sum())  # f at either end of the line.
    grad_f = 2.0 * solution[:2]  # grad f = (2 x_1, 2 x_2) at the solution; [2].
    parallel = float((grad_f + solution[2] * torch.ones(2)).abs().max())  # grad f + lambda grad g.
    show("x_1 at the constrained minimum", 0.5, float(solution[0]), quiet)
    show("x_2 at the constrained minimum", 0.5, float(solution[1]), quiet)
    show("lambda, the equality multiplier", -1.0, float(solution[2]), quiet)
    show("f at the minimum", 0.5, optimum, quiet)
    show("f at the endpoint (1, 0)", 1.0, endpoint, quiet)
    show("largest |grad f + lambda grad g|", 0.0, parallel, quiet)
    return {"solution": solution, "optimum": optimum}


# %% 04. Normalization as a constraint
def normalization(quiet: bool = False) -> dict:
    """Section 1.37.2: the uniform distribution maximizes entropy under normalization alone, and
    the observed frequencies minimize the empirical negative log-likelihood."""
    states = 4  # K, the size of the DNA alphabet.
    uniform = torch.full((states,), 1.0 / states)  # p_i = e^{lambda - 1}, the same at every state.
    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # p(x) = n(x) / N, the solution at lambda = 1.
    loss = float(-(frequencies * frequencies.log()).sum())  # The minimum empirical loss, its entropy.
    nearby = frequencies + torch.tensor([0.05, -0.05, 0.0, 0.0])  # A different categorical model.
    worse = float(-(frequencies * nearby.log()).sum())  # Its loss on the same ten bases.
    show("uniform p_i at K = 4", 0.25, float(uniform[0]), quiet)
    show("max entropy under normalization", 1.386, entropy(uniform), quiet)
    show("p_hat(A) = n(A) / N", 0.4, float(frequencies[0]), quiet)
    show("p_hat(T) = n(T) / N", 0.1, float(frequencies[3]), quiet)
    show("minimum empirical loss", 1.280, loss, quiet)
    show("loss drop from moving off p_hat", 0.0, max(0.0, loss - worse), quiet)
    if not quiet:  # The displaced model has no chapter number, so we print it on its own.
        print(f"  {'loss of a model moved by 0.05':<36s} loss    {worse:>11.5f}"
              f"   at p_hat {loss:>11.5f}")
    return {"uniform": uniform, "frequencies": frequencies, "loss": loss, "worse": worse}


# %% 05. Maximum entropy under a property constraint
def maximum_entropy(samples: int, quiet: bool = False) -> dict:
    """Section 1.37.3: with one multiplier per property the entropy maximizer is p* proportional to
    exp(lambda h(x)), and the constraint fixes lambda."""
    strong, weak = float(GC_INDICATOR.sum()), float((1.0 - GC_INDICATOR).sum())  # Two bases each.
    exp_lambda = (GC_TARGET / (1.0 - GC_TARGET)) * (weak / strong)  # From 2e^L/(2e^L + 2) = 0.6.
    multiplier = math.log(exp_lambda)  # lambda = log 1.5, the tilt toward the strong bases.
    unnormalized = (multiplier * GC_INDICATOR).exp()  # exp(lambda h(x)) at the four bases; [4].
    partition = float(unnormalized.sum())  # Z(lambda) = 2 e^lambda + 2 = 5.
    p_star = unnormalized / partition  # The maximum-entropy distribution (0.2, 0.3, 0.3, 0.2); [4].
    uniform = torch.full((4,), 0.25)  # The distribution at lambda = 0, with no property requirement.
    weights = torch.rand(max(samples, 1), 2)  # One split per base class, drawn independently; [N, 2].
    competitors = torch.stack([  # Distributions that also meet the GC target, in A, C, G, T order.
        (1.0 - GC_TARGET) * weights[:, 0], GC_TARGET * weights[:, 1],
        GC_TARGET * (1.0 - weights[:, 1]), (1.0 - GC_TARGET) * (1.0 - weights[:, 0])], dim=1)
    drawn = -(competitors * competitors.clamp_min(1e-300).log()).sum(1)  # H of every competitor; [N].
    best = float(drawn.max())  # The largest entropy reached by any distribution meeting the target.
    show("e^lambda", 1.5, exp_lambda, quiet)
    show("lambda = log 1.5", 0.4055, multiplier, quiet)
    show("partition function Z", 5.0, partition, quiet)
    show("p*(A)", 0.2, float(p_star[0]), quiet)
    show("p*(C)", 0.3, float(p_star[1]), quiet)
    show("GC content of p*", 0.6, float(p_star @ GC_INDICATOR), quiet)
    show("H(p*)", 1.366, entropy(p_star), quiet)
    show("H of the unconstrained uniform", 1.386, entropy(uniform), quiet)
    show("entropy cost of the constraint", 0.0201, entropy(uniform) - entropy(p_star), quiet)
    show(f"H above H(p*) over {len(drawn)} draws", 0.0, max(0.0, best - entropy(p_star)), quiet)
    return {"multiplier": multiplier, "p_star": p_star, "partition": partition, "best_drawn": best}


# %% 06. Inequality constraints and the KKT conditions
def kkt(quiet: bool = False) -> dict:
    """Section 1.38: on min (x - 3)^2 with x <= 1 the bound is active, so lambda is positive."""
    x_star = 1.0  # The unconstrained minimizer 3 breaks the bound, so the solution sits at the bound.
    multiplier = -2.0 * (x_star - 3.0)  # Stationarity 2(x - 3) + lambda = 0 solved at x = 1.
    slack = x_star - 1.0  # g(x*) = x* - 1, and it must be at or below zero.
    show("x* at the bound", 1.0, x_star, quiet)
    show("lambda, the inequality multiplier", 4.0, multiplier, quiet)
    show("g(x*) = x* - 1", 0.0, slack, quiet)
    show("complementary slackness lambda g", 0.0, multiplier * slack, quiet)
    show("stationarity 2(x* - 3) + lambda", 0.0, 2.0 * (x_star - 3.0) + multiplier, quiet)
    show("f at the unconstrained minimum", 0.0, (3.0 - 3.0) ** 2, quiet)
    show("f under the bound", 4.0, (x_star - 3.0) ** 2, quiet)
    show("value drop for epsilon = 0.001", 0.004, -((1.001 - 3.0) ** 2 - 4.0), quiet)
    return {"multiplier": multiplier, "value": (x_star - 3.0) ** 2}


# %% 07. Optimization over the simplex through the softmax
def over_the_simplex(steps: int, quiet: bool = False) -> dict:
    """Section 1.38.1: under logits the simplex constraints hold at every step, so the constrained
    likelihood problem of Section 1.37.2 becomes an unconstrained descent that reaches p_hat."""
    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # The minimizer the multiplier route produced.
    logits = torch.zeros(4, requires_grad=True)  # z in R^4, starting at the uniform distribution.
    learning_rate, losses = 1.0, []  # Plain gradient descent on z, with the loss after each step.
    for _ in range(max(steps, 1)):
        model = torch.softmax(logits, dim=0)  # p_i = e^{z_i} / sum_j e^{z_j}, always on the simplex.
        loss = -(frequencies * model.log()).sum()  # The empirical negative log-likelihood.
        gradient, = torch.autograd.grad(loss, logits)  # d L / d z, or softmax(z) - p_hat.
        with torch.no_grad():  # The update itself runs without a graph.
            logits -= learning_rate * gradient  # One descent step of Section 1.36, now on logits.
        losses.append(float(loss.detach()))
    final = torch.softmax(logits.detach(), dim=0)  # The distribution the descent arrived at; [4].
    show("sum_i p_i at the end", 1.0, float(final.sum()), quiet)
    show("entries at or below zero", 0.0, float((final <= 0.0).sum()), quiet)
    show("p(A) after the descent", 0.4, float(final[0]), quiet)
    show("p(T) after the descent", 0.1, float(final[3]), quiet)
    show("loss after the descent", 1.280, losses[-1], quiet)
    show("largest |p - p_hat|", 0.0, float((final - frequencies).abs().max()), quiet)
    return {"final": final, "losses": losses, "distance": float((final - frequencies).abs().max())}


# %% 08. The report
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--steps", type=int, default=600, help="gradient descent steps")
    parser.add_argument("--samples", type=int, default=4000, help="competing distributions drawn")
    parser.add_argument("--seed", type=int, default=0, help="random seed for the drawn competitors")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    set_seed(args.seed)  # Seed first, so the competitors repeat under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.36 to 1.38, worked example by worked example\n")
    descent = unconstrained(args.steps, args.quiet)
    constrained = equality_constraint(args.quiet)
    normalized = normalization(args.quiet)
    tilted = maximum_entropy(args.samples, args.quiet)
    bounded = kkt(args.quiet)
    fitted = over_the_simplex(args.steps, args.quiet)

    report = {
        "descent steps run": len(descent["history"]),
        "theta after the descent": descent["history"][-1],
        "x_1 under x_1 + x_2 = 1": float(constrained["solution"][0]),
        "lambda, equality constraint": float(constrained["solution"][2]),
        "minimum empirical loss, nats": normalized["loss"],
        "lambda for GC content 0.6": tilted["multiplier"],
        "p*(C) under that constraint": float(tilted["p_star"][1]),
        "H(p*), nats": entropy(tilted["p_star"]),
        "best drawn H at the same target": tilted["best_drawn"],
        "lambda, active inequality": bounded["multiplier"],
        "f under the bound": bounded["value"],
        "loss after the softmax descent": fitted["losses"][-1],
        "largest |p - p_hat|": fitted["distance"],
    }
    print_report("Optimization, Sections 1.36 to 1.38", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work every example of Sections 1.36 to 1.38 and print the report.