Blame for lecture_1/linear_algebra.py

"""The simplex, the softmax, transition matrices and the matrix exponential.

A distribution over K states is a point of the simplex, a softmax maps unrestricted logits onto
that simplex, and a row-stochastic matrix moves one point of it to the next. Along an
eigendirection a matrix power is a power of a number. That fact gives the stationary distribution,
the rate at which a chain reaches it, and the matrix exponential of continuous time.

Chapter 1, Sections 1.12, 1.13, 1.15, 1.16, 1.17 and 1.18 of the notes. Runs on the three-base
simplex of Section 1.12.2, the two-state chain of Section 1.15.1 and the symmetric matrix
A = [[2, 1], [1, 2]] of Section 1.16. Prints the chapter's number beside the computed one.

    python lecture_1/linear_algebra.py             # every worked example
    python lecture_1/linear_algebra.py --steps 40  # run the chain further toward p*
    python lecture_1/linear_algebra.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.


P_BASES = torch.tensor([0.5, 0.3, 0.2])  # The three-state distribution over A, C, G; Section 1.12.2.
H_BASES = torch.tensor([1.0, 0.0, 4.0])  # Hydrogen-bond contacts per base, the property h.
Q_BASES = torch.tensor([0.1, 0.1, 0.8])  # The second point of the simplex, mixed with P_BASES.
LOGITS = torch.tensor([2.0, 1.0, 0.0])  # The logit vector z of Section 1.13.1.
P_CHAIN = torch.tensor([[0.8, 0.2],  # Row 1: stay in state one with 0.8, move to state two with 0.2.
                        [0.3, 0.7]])  # Row 2: move to state one with 0.3, stay with 0.7.
A_MATRIX = torch.tensor([[2.0, 1.0], [1.0, 2.0]])  # The running matrix of Sections 1.16 and 1.18.
X_START = torch.tensor([3.0, 1.0])  # The initial vector x_0 = 2 v_1 + v_2 of Section 1.16.2.


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}")


# %% 02. Probability vectors and the simplex
def simplex(quiet: bool = False) -> dict:
    """Sections 1.12.2 and 1.12.3: expectation as a dot product, and a mixture of two points."""
    expectation = float(P_BASES @ H_BASES)  # E[h(X)] = p^T h = 0.5 + 0 + 0.8 = 1.3.
    mixture = 0.5 * P_BASES + 0.5 * Q_BASES  # r = lambda p + (1 - lambda) q at lambda = 0.5; [3].
    corners = 0.3 * torch.tensor([1.0, 0.0]) + 0.7 * torch.tensor([0.0, 1.0])  # Two-state mixture.
    show("sum_i p_i", 1.0, float(P_BASES.sum()), quiet)
    show("E[h(X)] = p^T h", 1.3, expectation, quiet)
    show("mixture r_1 at lambda = 0.5", 0.30, float(mixture[0]), quiet)
    show("mixture r_2 at lambda = 0.5", 0.20, float(mixture[1]), quiet)
    show("mixture r_3 at lambda = 0.5", 0.50, float(mixture[2]), quiet)
    show("sum_i r_i, still a distribution", 1.0, float(mixture.sum()), quiet)
    show("0.3 e_1 + 0.7 e_2, first entry", 0.3, float(corners[0]), quiet)
    return {"expectation": expectation, "mixture": mixture}


# %% 03. Logits and the softmax
def softmax(quiet: bool = False) -> dict:
    """Section 1.13.1: softmax(z)_i = e^{z_i} / sum_j e^{z_j} on the logits z = (2, 1, 0)."""
    exponentials = LOGITS.exp()  # e^2, e^1, e^0; [3].
    partition = float(exponentials.sum())  # The normalizing sum, 7.389 + 2.718 + 1.000.
    probabilities = exponentials / partition  # The point of the simplex the logits name; [3].
    shifted = torch.softmax(LOGITS + 5.0, dim=0)  # Adding a constant c cancels in the ratio.
    show("e^{z_1}", 7.389, float(exponentials[0]), quiet)
    show("e^{z_2}", 2.718, float(exponentials[1]), quiet)
    show("e^{z_3}", 1.000, float(exponentials[2]), quiet)
    show("partition function sum_j e^{z_j}", 11.107, partition, quiet)
    show("softmax(z)_1", 0.665, float(probabilities[0]), quiet)
    show("softmax(z)_2", 0.245, float(probabilities[1]), quiet)
    show("softmax(z)_3", 0.090, float(probabilities[2]), quiet)
    show("sum_i softmax(z)_i", 1.0, float(probabilities.sum()), quiet)
    show("shift invariance, largest change", 0.0, float((shifted - probabilities).abs().max()), quiet)
    return {"probabilities": probabilities, "partition": partition}


# %% 04. A probability vector under a transition matrix
def one_transition(quiet: bool = False) -> dict:
    """Section 1.15.1: p_{t+1} = P^T p_t moves an even split to (0.55, 0.45) in one step."""
    p_t = torch.tensor([0.5, 0.5])  # The current distribution over the two states; [2].
    p_next = P_CHAIN.T @ p_t  # Each new entry collects the mass arriving from every state; [2].
    show("row sums of P, row 1", 1.0, float(P_CHAIN[0].sum()), quiet)
    show("row sums of P, row 2", 1.0, float(P_CHAIN[1].sum()), quiet)
    show("p_{t+1}(1)", 0.55, float(p_next[0]), quiet)
    show("p_{t+1}(2)", 0.45, float(p_next[1]), quiet)
    show("sum_j p_{t+1}(j)", 1.0, float(p_next.sum()), quiet)
    return {"p_next": p_next}


# %% 05. Eigenvalues and eigenvectors
def eigen(quiet: bool = False) -> dict:
    """Section 1.16: the characteristic equation of A, its two eigenvalues and its eigenvectors."""
    trace = float(A_MATRIX.trace())  # The sum of the eigenvalues, 4.
    determinant = float(torch.linalg.det(A_MATRIX))  # The product of the eigenvalues, 3.
    values, vectors = torch.linalg.eigh(A_MATRIX)  # A is symmetric, so the eigenvalues are real.
    order = values.argsort(descending=True)  # The chapter indexes lambda_1 = 3 before lambda_2 = 1.
    values, vectors = values[order], vectors[:, order]  # [2] and [2, 2].
    vectors = vectors / vectors[0]  # Rescale each one to first entry one, here (1, ±1).
    residual = (A_MATRIX @ vectors - vectors * values).abs().max()  # A v - lambda v at both pairs.
    show("trace A, lambda_1 + lambda_2", 4.0, trace, quiet)
    show("det A, lambda_1 lambda_2", 3.0, determinant, quiet)
    show("lambda_1", 3.0, float(values[0]), quiet)
    show("lambda_2", 1.0, float(values[1]), quiet)
    show("v_1 second entry over first", 1.0, float(vectors[1, 0]), quiet)
    show("v_2 second entry over first", -1.0, float(vectors[1, 1]), quiet)
    show("largest |A v - lambda v|", 0.0, float(residual), quiet)
    return {"values": values, "vectors": vectors}


# %% 06. Powers along an eigendirection
def powers(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
    """Section 1.16.2: x_n = c_1 lambda_1^n v_1 + c_2 lambda_2^n v_2, against A^n x_0 taken directly."""
    coefficients = torch.linalg.solve(vectors, X_START)  # c_1 = 2 and c_2 = 1 for x_0 = (3, 1).
    eigen_route = vectors @ (coefficients * values ** 2)  # The n = 2 state along the two directions.
    direct = torch.linalg.matrix_power(A_MATRIX, 2) @ X_START  # A^2 x_0, squaring the matrix first.
    far = vectors @ (coefficients * values ** 10)  # The same decomposition carried to n = 10.
    show("c_1", 2.0, float(coefficients[0]), quiet)
    show("c_2", 1.0, float(coefficients[1]), quiet)
    show("lambda_1^10", 59049.0, float(values[0]) ** 10, quiet)
    show("x_2 first entry, eigen route", 19.0, float(eigen_route[0]), quiet)
    show("x_2 second entry, eigen route", 17.0, float(eigen_route[1]), quiet)
    show("x_2 first entry, A^2 x_0", 19.0, float(direct[0]), quiet)
    show("x_10 first entry", 118099.0, float(far[0]), quiet)
    show("x_10 second entry", 118097.0, float(far[1]), quiet)
    return {"x_two": eigen_route, "x_ten": far}


# %% 07. Diagonalization and the dominant eigenvalue
def diagonalization(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
    """Sections 1.16.3 and 1.16.4: A = V Lambda V^{-1}, so A^n = V Lambda^n V^{-1} telescopes."""
    inverse = torch.linalg.inv(vectors)  # V^{-1} = 0.5 [[1, 1], [1, -1]]; [2, 2].
    rebuilt = vectors @ torch.diag(values) @ inverse  # V Lambda V^{-1} returns A itself; [2, 2].
    cube = vectors @ torch.diag(values ** 3) @ inverse  # V diag(27, 1) V^{-1}, the n = 3 power.
    direct_cube = torch.linalg.matrix_power(A_MATRIX, 3)  # A^3 computed by repeated multiplication.
    dominant = torch.tensor([2.0, 0.5]) ** 10  # The coefficients 2^10 and 0.5^10 of Section 1.16.4.
    show("V^{-1}_11", 0.5, float(inverse[0, 0]), quiet)
    show("largest |V Lambda V^{-1} - A|", 0.0, float((rebuilt - A_MATRIX).abs().max()), quiet)
    show("A^3_11 through the eigenvalues", 14.0, float(cube[0, 0]), quiet)
    show("A^3_12 through the eigenvalues", 13.0, float(cube[0, 1]), quiet)
    show("largest |V Lambda^3 V^{-1} - A^3|", 0.0, float((cube - direct_cube).abs().max()), quiet)
    show("dominant coefficient 2^10", 1024.0, float(dominant[0]), quiet)
    show("decaying coefficient 0.5^10", 0.000977, float(dominant[1]), quiet)
    show("ratio of the two, in millions", 1.05, float(dominant[0] / dominant[1]) / 1e6, quiet)
    return {"inverse": inverse, "cube": cube}


# %% 08. Stationary distributions
def stationary(steps: int, quiet: bool = False) -> dict:
    """Section 1.17: p* solves (P^T - I) p* = 0 with sum one, and the chain halves its deviation."""
    drift = (P_CHAIN.T - torch.eye(2))[0]  # The first row of P^T - I, or 0.3 p_2 = 0.2 p_1.
    system = torch.stack([drift, torch.ones(2)])  # That one equation, with normalization; [2, 2].
    p_star = torch.linalg.solve(system, torch.tensor([0.0, 1.0]))  # The lambda = 1 eigenvector; [2].
    trace, determinant = float(P_CHAIN.trace()), float(torch.linalg.det(P_CHAIN))  # 1.5 and 0.5.
    second = (trace - math.sqrt(trace ** 2 - 4.0 * determinant)) / 2.0  # The smaller root, 0.5.
    p_t, history = torch.tensor([0.5, 0.5]), []  # Start from the even split of Section 1.15.1.
    for _ in range(max(steps, 1)):  # Iterate p_{t+1} = P^T p_t for as many steps as requested.
        p_t = P_CHAIN.T @ p_t  # One more application of the transition matrix; [2].
        history.append(p_t.clone())
    show("p*_1", 0.6, float(p_star[0]), quiet)
    show("p*_2", 0.4, float(p_star[1]), quiet)
    show("largest |P^T p* - p*|", 0.0, float((P_CHAIN.T @ p_star - p_star).abs().max()), quiet)
    show("trace P", 1.5, trace, quiet)
    show("det P", 0.5, determinant, quiet)
    show("second eigenvalue, the rate", 0.5, second, quiet)
    for index, value in ((0, 0.55), (1, 0.575), (2, 0.5875), (9, 0.59990)):
        if index < len(history):  # The chapter prints p_1, p_2, p_3 and p_10 of this sequence.
            show(f"p_{index + 1}(1)", value, float(history[index][0]), quiet)
    deviation = float((history[-1] - p_star).abs().max())  # |d_n| = 0.5^n |d_0| after n steps.
    show(f"|d_n| after {len(history)} steps", 0.1 * second ** len(history), deviation, quiet)
    return {"p_star": p_star, "second": second, "history": history, "deviation": deviation}


# %% 09. From discrete steps to continuous time
def matrix_exponential(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
    """Section 1.18.3: the four-term series and the eigenvalue formula for e^{tA} at t = 0.1."""
    t = 0.1  # The small time at which the chapter truncates the series.
    series, power, factorial = torch.eye(2), torch.eye(2), 1.0  # Accumulate I + tA + t^2A^2/2 + ...
    for k in range(1, 4):  # Terms k = 1, 2, 3, and with the identity they make four.
        power, factorial = power @ A_MATRIX, factorial * k  # A^k and k factorial.
        series = series + t ** k * power / factorial  # Add t^k A^k / k!; [2, 2].
    discarded = t ** 4 * (power @ A_MATRIX) / 24.0  # The first omitted term, t^4 A^4 / 4!.
    scaled = torch.diag((t * values).exp())  # e^{t Lambda}, the diagonal of entry exponentials; [2, 2].
    eigen_route = vectors @ scaled @ torch.linalg.inv(vectors)  # V e^{t Lambda} V^{-1}; [2, 2].
    gap = float((series - eigen_route).abs().max())  # How far four terms fall short at t = 0.1.
    coefficients = torch.linalg.solve(vectors, X_START)  # c_1 = 2, c_2 = 1, as in Section 1.16.2.
    x_one = vectors @ (coefficients * values.exp())  # x_t = sum_i c_i e^{t lambda_i} v_i at t = 1.
    show("four-term series, entry 11", 1.22733, float(series[0, 0]), quiet)
    show("four-term series, entry 12", 0.12217, float(series[0, 1]), quiet)
    show("eigenvalue route, entry 11", 1.22751, float(eigen_route[0, 0]), quiet)
    show("eigenvalue route, entry 12", 0.12234, float(eigen_route[0, 1]), quiet)
    show("largest gap between the routes", 1.8e-4, gap, quiet)
    show("first discarded term, entry 11", 1.71e-4, float(discarded[0, 0]), quiet)
    show("first discarded term, entry 12", 1.67e-4, float(discarded[0, 1]), quiet)
    show("x_1 first entry, 2e^3 + e", 42.889, float(x_one[0]), quiet)
    show("x_1 second entry, 2e^3 - e", 37.453, float(x_one[1]), quiet)
    show("decay e^{-2t} at t = 1", 0.1353, math.exp(-2.0), quiet)
    show("decay e^{-2t} at t = 3", 0.002479, math.exp(-6.0), quiet)
    show("sqrt(lambda_1), Gaussian long axis", 1.732, math.sqrt(float(values[0])), quiet)
    return {"series": series, "eigen_route": eigen_route, "gap": gap, "x_one": x_one}


# %% 10. The report
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--steps", type=int, default=10, help="iterations of p_{t+1} = P^T p_t")
    parser.add_argument("--seed", type=int, default=0, help="random seed, for a repeatable run")
    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 run repeats under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.12 to 1.18, worked example by worked example\n")
    triangle = simplex(args.quiet)
    logits = softmax(args.quiet)
    step = one_transition(args.quiet)
    spectrum = eigen(args.quiet)
    grown = powers(spectrum["values"], spectrum["vectors"], args.quiet)
    diagonalization(spectrum["values"], spectrum["vectors"], args.quiet)
    chain = stationary(args.steps, args.quiet)
    flow = matrix_exponential(spectrum["values"], spectrum["vectors"], args.quiet)

    report = {
        "E[h(X)] on the three bases": triangle["expectation"],
        "softmax(z)_1": float(logits["probabilities"][0]),
        "p_{t+1}(1) after one step": float(step["p_next"][0]),
        "lambda_1 of A": float(spectrum["values"][0]),
        "lambda_2 of A": float(spectrum["values"][1]),
        "x_10 first entry": float(grown["x_ten"][0]),
        "p*_1, stationary": float(chain["p_star"][0]),
        "second eigenvalue of P": chain["second"],
        "chain steps run": len(chain["history"]),
        "deviation from p* at the end": chain["deviation"],
        "e^{tA}_11 at t = 0.1": float(flow["eigen_route"][0, 0]),
        "four-term truncation error": flow["gap"],
    }
    print_report("Linear algebra, Sections 1.12 to 1.18", report)
    save_run(args.out, config=vars(args), report=report)


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