Download lecture_1/linear_algebra.py from ChatterjeeLab/CIS6270: direct link, hf CLI and curl.
- Browser
- Download file 15.7 kB
-
https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_1/linear_algebra.py
- Command line
-
hf download hf://ChatterjeeLab/CIS6270/lecture_1/linear_algebra.py
-
curl -L -o linear_algebra.py https://huggingface.co/ChatterjeeLab/CIS6270/resolve/main/lecture_1/linear_algebra.py
15.7 kB
| """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. | |