Blame for tests/test_lecture_1.py

"""Chapter 1's hand-computed numbers, asserted against the six Lecture 1 scripts.

Every row below is a value Chapter 1 works out by hand, paired with the label under which one of
the Lecture 1 scripts prints its own computation of it. The tolerance on each row is the rounding the
chapter itself prints: we check a value given to three decimals to 5e-4, one given to five decimals
to 5e-6, and a value the chapter states exactly to machine precision. A value the chapter predicts
and a script measures by sampling has a test of its own, with a band wide enough for the Monte
Carlo error at the sample size the test uses.

    cd cis6270 && PYTHONPATH=. python -m pytest tests/test_lecture_1.py -q
"""

import math
import subprocess
import sys
from pathlib import Path

import pytest

from cis6270.runner import set_seed
from lecture_1 import calculus, fields, information, linear_algebra, optimization, probability

SEED = 0  # The seed every fixture sets before running a script's sections.
SAMPLES = 4000  # Monte Carlo draws and random competitors, fixed so the printed labels are fixed.
GRID = 1000  # Points in the numerical checks of calculus.py and fields.py.
STEPS = 600  # Gradient-descent steps, chain steps, Euler steps and estimator repeats.
EXACT = 1e-12  # The tolerance for a value the chapter states exactly.
ROOT = Path(__file__).resolve().parents[1]  # The repository root, for the command-line check.


# %% 01. Probability, Sections 1.3 to 1.11
@pytest.fixture(scope="module")
def probability_worked() -> dict:
    """Run every worked example of probability.py and return the values it computed."""
    set_seed(SEED)
    probability.WORKED.clear()
    probability.two_token_distribution(quiet=True)
    probability.continuous_densities(quiet=True)
    probability.expected_adenine(quiet=True)
    probability.monte_carlo(SAMPLES, STEPS, quiet=True)
    probability.worked_variance(quiet=True)
    probability.peptide_covariance(quiet=True)
    probability.noising_covariance(quiet=True)
    probability.joint_marginal_conditional(quiet=True)
    probability.bayes(quiet=True)
    probability.independence(quiet=True)
    return dict(probability.WORKED)


PROBABILITY_ROWS = [
    # Section 1.3.1, the two-token distribution and the event that the first token is A.
    ("sum_x p_X(x)", 1.0, EXACT),
    ("P(x_1 = A) = p(AA) + p(AC)", 0.5, EXACT),
    # Section 1.4, densities whose height exceeds one.
    ("uniform[0, 0.1] density height", 10.0, EXACT),
    ("P(0 <= X <= 0.02)", 0.2, EXACT),
    ("peak of N(0, 0.1^2)", 3.989, 5e-4),
    # Sections 1.5.1 and 1.5.3, the expected adenine count and the Monte Carlo error scale.
    ("E[h(X)], expected A count", 1.1, EXACT),
    ("expected A count in 1000 draws", 1100.0, EXACT),
    ("E[h(X)^2]", 1.9, EXACT),
    ("Var(h(X))", 0.69, EXACT),
    ("sd of a 10-sample average", 0.263, 5e-4),
    ("sd of a 1000-sample average", 0.0263, 5e-5),
    # Section 1.6.2, the variance of the values 0, 1, 2 by both routes.
    ("E[X]", 1.1, EXACT),
    ("E[X^2]", 1.7, EXACT),
    ("Var(X) from the two moments", 0.49, EXACT),
    ("Var(X) from the definition", 0.49, EXACT),
    ("sigma_X", 0.7, EXACT),
    # Sections 1.7.1 and 1.7.2, the three peptides and their covariance matrix.
    ("mu_X", 2.0, EXACT),
    ("mu_Y", 4.0, EXACT),
    ("Cov(X, Y)", 5.0 / 3.0, EXACT),
    ("Var(X) = Sigma_11", 2.0 / 3.0, EXACT),
    ("Var(Y) = Sigma_22", 14.0 / 3.0, EXACT),
    ("sigma_X across the peptides", 0.8165, 5e-5),
    ("sigma_Y across the peptides", 2.160, 5e-4),
    ("rho_XY", 0.9449, 5e-5),
    ("Sigma symmetry, Sigma_21 - Sigma_12", 0.0, EXACT),
    # Section 1.8, the noising equation at alpha_t = 0.6 and sigma_t = 0.8.
    ("alpha_t^2 + sigma_t^2", 1.0, EXACT),
    ("Cov(X_t)_11 with Cov(X_0) = I", 1.0, EXACT),
    ("Cov(X_t)_11 with diag(4, 1)", 2.08, EXACT),
    ("Cov(X_t)_22 with diag(4, 1)", 1.00, EXACT),
    # Sections 1.9.1 to 1.9.3, the motif-and-binding table with its marginals and conditionals.
    ("sum over the four cells", 1.0, EXACT),
    ("p(X = 1), motif present", 0.60, EXACT),
    ("p(X = 0), motif absent", 0.40, EXACT),
    ("p(Y = 1), high binding", 0.44, EXACT),
    ("p(Y = 0), low binding", 0.56, EXACT),
    ("p(Y = 1 | X = 1)", 0.6, EXACT),
    ("p(Y = 0 | X = 1)", 0.4, EXACT),
    ("p(Y = 1 | X = 0)", 0.2, EXACT),
    ("p(Y = 0 | X = 0)", 0.8, EXACT),
    # Section 1.10, Bayes' rule and the law of total probability.
    ("evidence p(Y = 1)", 0.44, EXACT),
    ("posterior p(X = 1 | Y = 1)", 0.818, 5e-4),
    ("likelihood ratio 0.6 / 0.2", 3.0, EXACT),
    # Section 1.11, the failure of independence on the same table.
    ("p(X = 1) p(Y = 1)", 0.264, EXACT),
    ("p(X = 1, Y = 1)", 0.36, EXACT),
    ("association missed by factorizing", 0.096, EXACT),
]


@pytest.mark.parametrize("label, chapter, tolerance", PROBABILITY_ROWS)
def test_probability(probability_worked, label, chapter, tolerance):
    assert probability_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_monte_carlo_lands_inside_its_own_error_bar():
    """Section 1.5.3: the sample average lands within four standard deviations of E[h(X)] = 1.1."""
    set_seed(SEED)
    estimates = probability.monte_carlo(SAMPLES, STEPS, quiet=True)
    size, value, error = estimates["ladder"][-1]  # The last rung, at the largest N the ladder reached.
    assert 10 * size > SAMPLES >= size
    assert abs(value - estimates["exact"]) < 4.0 * error


def test_the_measured_spread_of_ten_sample_averages_matches_0_263(probability_worked):
    """Section 1.5.3: repeating the ten-sample average measures the 0.263 the chapter predicts.

    The spread of a sample standard deviation over 600 rounds is a few percent, so the band is
    ten percent of the predicted value, looser than the rounding of the printed digits.
    """
    predicted = math.sqrt(0.69 / 10.0)
    assert probability_worked[f"sd of {STEPS} ten-sample averages"] == pytest.approx(
        predicted, rel=0.10)


def test_the_same_seed_gives_the_same_monte_carlo_ladder():
    """Rule 8 of the conventions: a run is reproducible from its seed alone."""
    set_seed(SEED)
    first = probability.monte_carlo(1000, 20, quiet=True)["ladder"]
    set_seed(SEED)
    second = probability.monte_carlo(1000, 20, quiet=True)["ladder"]
    assert first == second


# %% 02. Linear algebra, Sections 1.12 to 1.18
@pytest.fixture(scope="module")
def linear_algebra_worked() -> dict:
    """Run every worked example of linear_algebra.py and return the values it computed."""
    set_seed(SEED)
    linear_algebra.WORKED.clear()
    linear_algebra.simplex(quiet=True)
    linear_algebra.softmax(quiet=True)
    linear_algebra.one_transition(quiet=True)
    spectrum = linear_algebra.eigen(quiet=True)
    linear_algebra.powers(spectrum["values"], spectrum["vectors"], quiet=True)
    linear_algebra.diagonalization(spectrum["values"], spectrum["vectors"], quiet=True)
    linear_algebra.stationary(10, quiet=True)
    linear_algebra.matrix_exponential(spectrum["values"], spectrum["vectors"], quiet=True)
    return dict(linear_algebra.WORKED)


LINEAR_ALGEBRA_ROWS = [
    # Sections 1.12.2 and 1.12.3, expectation as a dot product and a mixture of two points.
    ("E[h(X)] = p^T h", 1.3, EXACT),
    ("mixture r_1 at lambda = 0.5", 0.30, EXACT),
    ("mixture r_2 at lambda = 0.5", 0.20, EXACT),
    ("mixture r_3 at lambda = 0.5", 0.50, EXACT),
    ("sum_i r_i, still a distribution", 1.0, EXACT),
    # Section 1.13.1, the softmax of z = (2, 1, 0).
    ("e^{z_1}", 7.389, 5e-4),
    ("e^{z_2}", 2.718, 5e-4),
    ("e^{z_3}", 1.000, EXACT),
    ("partition function sum_j e^{z_j}", 11.107, 5e-4),
    ("softmax(z)_1", 0.665, 5e-4),
    ("softmax(z)_2", 0.245, 5e-4),
    ("softmax(z)_3", 0.090, 5e-4),
    ("sum_i softmax(z)_i", 1.0, EXACT),
    ("shift invariance, largest change", 0.0, 1e-12),
    # Section 1.15.1, one step of the two-state chain from an even split.
    ("p_{t+1}(1)", 0.55, EXACT),
    ("p_{t+1}(2)", 0.45, EXACT),
    # Sections 1.16 and 1.16.1, the eigenvalues and eigenvectors of A = [[2, 1], [1, 2]].
    ("trace A, lambda_1 + lambda_2", 4.0, EXACT),
    ("det A, lambda_1 lambda_2", 3.0, EXACT),
    ("lambda_1", 3.0, 1e-12),
    ("lambda_2", 1.0, 1e-12),
    ("v_1 second entry over first", 1.0, 1e-12),
    ("v_2 second entry over first", -1.0, 1e-12),
    ("largest |A v - lambda v|", 0.0, 1e-12),
    # Section 1.16.2, powers along an eigendirection from x_0 = (3, 1).
    ("c_1", 2.0, 1e-12),
    ("c_2", 1.0, 1e-12),
    ("lambda_1^10", 59049.0, 1e-6),
    ("x_2 first entry, eigen route", 19.0, 1e-10),
    ("x_2 second entry, eigen route", 17.0, 1e-10),
    ("x_10 first entry", 118099.0, 1e-6),
    ("x_10 second entry", 118097.0, 1e-6),
    # Sections 1.16.3 and 1.16.4, diagonalization and the dominant eigenvalue.
    ("V^{-1}_11", 0.5, 1e-12),
    ("largest |V Lambda V^{-1} - A|", 0.0, 1e-12),
    ("A^3_11 through the eigenvalues", 14.0, 1e-10),
    ("A^3_12 through the eigenvalues", 13.0, 1e-10),
    ("dominant coefficient 2^10", 1024.0, EXACT),
    ("decaying coefficient 0.5^10", 0.000977, 5e-7),
    ("ratio of the two, in millions", 1.05, 5e-3),
    # Section 1.17, the stationary distribution and the rate of approach.
    ("p*_1", 0.6, 1e-10),
    ("p*_2", 0.4, 1e-10),
    ("largest |P^T p* - p*|", 0.0, 1e-10),
    ("trace P", 1.5, EXACT),
    ("det P", 0.5, EXACT),
    ("second eigenvalue, the rate", 0.5, 1e-12),
    ("p_1(1)", 0.55, EXACT),
    ("p_2(1)", 0.575, EXACT),
    ("p_3(1)", 0.5875, EXACT),
    ("p_10(1)", 0.59990, 5e-6),
    # Section 1.18.3, the four-term series against the eigenvalue formula at t = 0.1.
    ("four-term series, entry 11", 1.22733, 5e-6),
    ("four-term series, entry 12", 0.12217, 5e-6),
    ("eigenvalue route, entry 11", 1.22751, 5e-6),
    ("eigenvalue route, entry 12", 0.12234, 5e-6),
    ("largest gap between the routes", 1.8e-4, 5e-6),
    ("first discarded term, entry 11", 1.71e-4, 5e-7),
    ("first discarded term, entry 12", 1.67e-4, 5e-7),
    ("x_1 first entry, 2e^3 + e", 42.889, 5e-4),
    ("x_1 second entry, 2e^3 - e", 37.453, 5e-4),
    # Sections 1.18.4 and 1.18.5, stability and the covariance read through its eigenvalues.
    ("decay e^{-2t} at t = 1", 0.1353, 5e-5),
    ("decay e^{-2t} at t = 3", 0.002479, 5e-7),
    ("sqrt(lambda_1), Gaussian long axis", 1.732, 5e-4),
]


@pytest.mark.parametrize("label, chapter, tolerance", LINEAR_ALGEBRA_ROWS)
def test_linear_algebra(linear_algebra_worked, label, chapter, tolerance):
    assert linear_algebra_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_the_chain_halves_its_deviation_at_every_step():
    """Section 1.17: the deviation from p* lies along the second eigenvector, so d_n = 0.5^n d_0."""
    chain = linear_algebra.stationary(12, quiet=True)
    for step, point in enumerate(chain["history"], start=1):
        deviation = float((point - chain["p_star"]).abs().max())
        assert deviation == pytest.approx(0.1 * 0.5 ** step, abs=1e-12)


# %% 03. Calculus, Sections 1.22 to 1.30
@pytest.fixture(scope="module")
def calculus_worked() -> dict:
    """Run every worked example of calculus.py and return the values it computed."""
    set_seed(SEED)
    calculus.WORKED.clear()
    calculus.jacobian(quiet=True)
    calculus.determinant(quiet=True)
    calculus.change_of_variables(GRID, quiet=True)
    calculus.composition(quiet=True)
    calculus.branches(GRID, quiet=True)
    calculus.gradient(quiet=True)
    calculus.divergences(quiet=True)
    calculus.continuity(GRID, STEPS, quiet=True)
    return dict(calculus.WORKED)


CALCULUS_ROWS = [
    # Sections 1.22.2 and 1.22.3, the Jacobian at (1, 0) and the displacement it predicts.
    ("J_phi(1, 0)_11", 2.0, EXACT),
    ("J_phi(1, 0)_21 = cos(1)", 0.5403, 5e-5),
    ("largest |autograd - closed form|", 0.0, EXACT),
    ("predicted change, first output", 0.0, 1e-12),
    ("predicted change, second output", -0.0546, 5e-5),
    ("exact change, first output", 0.0001, 5e-6),
    ("exact change, second output", -0.05464, 5e-6),
    # Section 1.23, the determinant as the local volume factor.
    ("det A for diag(2, 3)", 6.0, EXACT),
    ("area of the image rectangle", 6.0, EXACT),
    ("det of the collapsing map", 0.0, EXACT),
    ("det J_phi at x_1 = 2", 4.0, 1e-12),
    ("det J_phi at x_1 = 0.25", 0.5, 1e-12),
    # Sections 1.24.1 and 1.24.2, change of variables in one and in two dimensions.
    ("preimage of [1.4, 2.0], lower end", 0.2, EXACT),
    ("preimage of [1.4, 2.0], upper end", 0.5, EXACT),
    ("P(X in the preimage)", 0.3, EXACT),
    ("p_Y on [1, 3] after Y = 2X + 1", 0.5, EXACT),
    ("integral of p_Y over [1, 3]", 1.0, EXACT),
    ("p_Y inside the image of A", 1.0 / 6.0, EXACT),
    ("density times area for A", 1.0, EXACT),
    ("p_Y(y_1 = 1) for the x_1^2 map", 0.5, EXACT),
    ("largest |p_Y |det J| - p_X|", 0.0, 1e-12),
    ("total mass after the x_1^2 map", 1.0, 1e-12),
    # Section 1.25, the composition and its log-determinant.
    ("det J_psi", 6.0, 1e-12),
    ("chained Jacobian, entry 11", 8.0, 1e-12),
    ("det of the composition", 24.0, 1e-10),
    ("log 4", 1.3863, 5e-5),
    ("log 6", 1.7918, 5e-5),
    ("log 4 + log 6 = log 24", 3.1781, 5e-5),
    ("log p_Y from a unit base density", -3.1781, 5e-5),
    # Section 1.26, the two inverse branches of the squaring map.
    ("p_Y(0.25), two branches at ±0.5", 1.0, 1e-12),
    ("p_Y(0.04), two branches at ±0.2", 2.5, 1e-12),
    ("integral of p_Y over [0.01, 1]", 0.9, 5e-5),
    # Sections 1.27.1 and 1.27.2, the gradient and the directional derivative.
    ("grad f(1, 2), first entry", 2.0, EXACT),
    ("grad f(1, 2), second entry", 8.0, EXACT),
    ("|grad f(1, 2)| = sqrt(68)", math.sqrt(68.0), 1e-12),
    ("steepest direction, first entry", 0.243, 5e-4),
    ("steepest direction, second entry", 0.970, 5e-4),
    ("D_u f along (1, 0)", 2.0, EXACT),
    ("D_u f along (0, 1)", 8.0, EXACT),
    # Section 1.29, the divergence of the three fields.
    ("div of (-x_1, -2 x_2)", -3.0, EXACT),
    ("div of (2 x_1, -x_2)", 1.0, EXACT),
    ("div of (-x_2, x_1)", 0.0, EXACT),
    ("area factor e^{div t} at t = 0.5", 1.649, 5e-4),
    # Section 1.30.2, the worked flow v(x) = x on a uniform source.
    ("largest continuity residual", 0.0, 1e-6),
    ("log p along the trajectory at t = 1", -1.0, 1e-12),
    ("p_t(x_t) at t = 1", 0.3679, 5e-5),
    ("density at t = log 2", 0.5, EXACT),
    ("mass at t = log 2", 1.0, EXACT),
    ("mass at t = 1", 1.000, 5e-4),
]


@pytest.mark.parametrize("label, chapter, tolerance", CALCULUS_ROWS)
def test_calculus(calculus_worked, label, chapter, tolerance):
    assert calculus_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_the_euler_walk_approaches_the_exact_trajectory(calculus_worked):
    """Section 1.30.2: x_t = x_0 e^t, so a first-order walk closes on 0.5e as the step shrinks."""
    coarse = calculus.continuity(GRID, 50, quiet=True)["x_t"]
    fine = calculus.continuity(GRID, 4 * 50, quiet=True)["x_t"]
    assert abs(fine - 0.5 * math.e) < abs(coarse - 0.5 * math.e) / 3.0


# %% 04. Information, Sections 1.33 to 1.35
@pytest.fixture(scope="module")
def information_worked() -> dict:
    """Run every worked example of information.py and return the values it computed."""
    set_seed(SEED)
    information.WORKED.clear()
    information.information(quiet=True)
    information.entropies(quiet=True)
    information.cross_entropies(quiet=True)
    information.kl(SAMPLES, quiet=True)
    information.monte_carlo(SAMPLES, STEPS, quiet=True)
    information.maximum_likelihood(quiet=True)
    return dict(information.WORKED)


INFORMATION_ROWS = [
    # Section 1.33.1, the information content of a common and of a rare outcome.
    ("I(x) at p(x) = 0.5", 0.693, 5e-4),
    ("I(x) at p(x) = 0.01", 4.605, 5e-4),
    ("I(x) at p(x) = 1", 0.0, EXACT),
    ("ratio of the two surprises", 6.6, 5e-2),
    # Section 1.33.2, entropy on a skewed, a balanced and a uniform distribution.
    ("H((0.9, 0.1))", 0.325, 5e-4),
    ("H((0.5, 0.5)) = log 2", 0.693, 5e-4),
    ("H of a point mass", 0.0, EXACT),
    ("H of the uniform on 4 states", 1.386, 5e-4),
    # Section 1.33.3, the cross-entropy of p = (0.8, 0.2) scored by q = (0.6, 0.4).
    ("H(p, q)", 0.592, 5e-4),
    ("H(p)", 0.500, 5e-4),
    ("excess cost H(p, q) - H(p)", 0.092, 5e-4),
    ("H(p, p) - H(p)", 0.0, EXACT),
    # Section 1.34, the divergence, its split, its nonnegativity and its asymmetry.
    ("KL(p || q)", 0.0915, 5e-5),
    ("KL(q || p)", 0.1046, 5e-5),
    ("H(p) + KL(p || q) = H(p, q)", 0.592, 5e-4),
    ("KL(p || p)", 0.0, EXACT),
    ("reverse over forward, percent", 14.0, 0.5),
    ("largest log t - (t - 1)", 0.0, 1e-12),
    (f"largest KL breach over {SAMPLES} pairs", 0.0, EXACT),
    # Section 1.35, the empirical loss on ten DNA bases and its gap to the uniform model.
    ("term for A, 0.4 log(1/0.4)", 0.36652, 5e-6),
    ("term for C, 0.3 log(1/0.3)", 0.36119, 5e-6),
    ("term for G, 0.2 log(1/0.2)", 0.32189, 5e-6),
    ("term for T, 0.1 log(1/0.1)", 0.23026, 5e-6),
    ("loss at the frequencies", 1.280, 5e-4),
    ("loss under the uniform model", 1.386, 5e-4),
    ("gap between the two losses", 0.1064, 5e-5),
    ("KL(p_hat || uniform)", 0.1064, 5e-5),
]


@pytest.mark.parametrize("label, chapter, tolerance", INFORMATION_ROWS)
def test_information(information_worked, label, chapter, tolerance):
    assert information_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_cross_entropy_by_sampling_lands_inside_its_error_bar():
    """Section 1.33.3 estimated by the Monte Carlo of Section 1.5.3, to within four error bars."""
    set_seed(SEED)
    estimate = information.monte_carlo(SAMPLES, STEPS, quiet=True)
    assert abs(estimate["estimate"] - estimate["exact"]) < 4.0 * estimate["error"]
    assert estimate["measured"] == pytest.approx(estimate["predicted"], rel=0.15)


# %% 05. Optimization, Sections 1.36 to 1.38
@pytest.fixture(scope="module")
def optimization_worked() -> dict:
    """Run every worked example of optimization.py and return the values it computed."""
    set_seed(SEED)
    optimization.WORKED.clear()
    optimization.unconstrained(STEPS, quiet=True)
    optimization.equality_constraint(quiet=True)
    optimization.normalization(quiet=True)
    optimization.maximum_entropy(SAMPLES, quiet=True)
    optimization.kkt(quiet=True)
    optimization.over_the_simplex(STEPS, quiet=True)
    return dict(optimization.WORKED)


OPTIMIZATION_ROWS = [
    # Section 1.36, three steps of gradient descent on (theta - 3)^2 at eta = 0.1.
    ("theta_1", 0.600, 5e-4),
    ("theta_2", 1.080, 5e-4),
    ("theta_3", 1.464, 5e-4),
    (f"theta_{STEPS} from the closed form", 3.0, 1e-12),
    ("gradient at the minimizer theta = 3", 0.0, EXACT),
    # Section 1.37.1, the squared distance to the origin along x_1 + x_2 = 1.
    ("x_1 at the constrained minimum", 0.5, 1e-12),
    ("x_2 at the constrained minimum", 0.5, 1e-12),
    ("lambda, the equality multiplier", -1.0, 1e-12),
    ("f at the minimum", 0.5, 1e-12),
    ("f at the endpoint (1, 0)", 1.0, EXACT),
    ("largest |grad f + lambda grad g|", 0.0, 1e-12),
    # Section 1.37.2, normalization alone, and the likelihood minimizer it produces.
    ("uniform p_i at K = 4", 0.25, EXACT),
    ("max entropy under normalization", 1.386, 5e-4),
    ("p_hat(A) = n(A) / N", 0.4, EXACT),
    ("p_hat(T) = n(T) / N", 0.1, EXACT),
    ("minimum empirical loss", 1.280, 5e-4),
    ("loss drop from moving off p_hat", 0.0, EXACT),
    # Section 1.37.3, maximum entropy under a GC-content target of 0.6.
    ("e^lambda", 1.5, 1e-12),
    ("lambda = log 1.5", 0.4055, 5e-5),
    ("partition function Z", 5.0, 1e-12),
    ("p*(A)", 0.2, 1e-12),
    ("p*(C)", 0.3, 1e-12),
    ("GC content of p*", 0.6, 1e-12),
    ("H(p*)", 1.366, 5e-4),
    ("H of the unconstrained uniform", 1.386, 5e-4),
    ("entropy cost of the constraint", 0.0201, 5e-5),
    (f"H above H(p*) over {SAMPLES} draws", 0.0, EXACT),
    # Section 1.38, the bounded quadratic and its four KKT conditions.
    ("x* at the bound", 1.0, EXACT),
    ("lambda, the inequality multiplier", 4.0, EXACT),
    ("g(x*) = x* - 1", 0.0, EXACT),
    ("complementary slackness lambda g", 0.0, EXACT),
    ("stationarity 2(x* - 3) + lambda", 0.0, EXACT),
    ("f at the unconstrained minimum", 0.0, EXACT),
    ("f under the bound", 4.0, EXACT),
    ("value drop for epsilon = 0.001", 0.004, 5e-6),
    # Section 1.38.1, the same likelihood problem solved on logits through the softmax.
    ("sum_i p_i at the end", 1.0, 1e-12),
    ("entries at or below zero", 0.0, EXACT),
    ("p(A) after the descent", 0.4, 5e-4),
    ("p(T) after the descent", 0.1, 5e-4),
    ("loss after the descent", 1.280, 5e-4),
    ("largest |p - p_hat|", 0.0, 5e-4),
]


@pytest.mark.parametrize("label, chapter, tolerance", OPTIMIZATION_ROWS)
def test_optimization(optimization_worked, label, chapter, tolerance):
    assert optimization_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_the_softmax_descent_reaches_the_multiplier_solution():
    """Section 1.38.1: the unconstrained route on logits reaches what Section 1.37.2 derived."""
    set_seed(SEED)
    fitted = optimization.over_the_simplex(2000, quiet=True)
    frequencies = optimization.BASE_COUNTS / optimization.BASE_COUNTS.sum()
    assert fitted["distance"] < 1e-5
    assert fitted["losses"][-1] == pytest.approx(
        float(-(frequencies * frequencies.log()).sum()), abs=1e-9)
    assert fitted["losses"][0] > fitted["losses"][-1]


# %% 06. Fields, Sections 1.31 and 1.32
@pytest.fixture(scope="module")
def fields_worked() -> dict:
    """Run every worked example of fields.py and return the values it computed."""
    set_seed(SEED)
    fields.WORKED.clear()
    fields.scores(quiet=True)
    fields.curvatures(quiet=True)
    fields.heat(GRID, quiet=True)
    fields.fokker_planck(GRID, quiet=True)
    fields.euler_maruyama(SAMPLES, STEPS, quiet=True)
    return dict(fields.WORKED)


FIELDS_ROWS = [
    # Section 1.31, the score of the standard Gaussian and of N(1, 4).
    ("s(x) at x = 3", -3.0, EXACT),
    ("s(x) at x = -2", 2.0, EXACT),
    ("s(x) at the mode x = 0", 0.0, EXACT),
    ("largest |s(x) + x|", 0.0, EXACT),
    ("largest change from dropping Z", 0.0, EXACT),
    ("s(5) under N(1, 4)", -1.0, EXACT),
    # Section 1.32, the two curvatures of the bowl and the cancelling pair of the saddle.
    ("d^2 f / d x_1^2", 2.0, EXACT),
    ("d^2 f / d x_2^2", 6.0, EXACT),
    ("Laplacian of x_1^2 + 3 x_2^2", 8.0, EXACT),
    ("Laplacian from the four neighbors", 8.0, 5e-6),
    ("Laplacian of the saddle", 0.0, EXACT),
    # Section 1.32.1, the heat equation on the widening Gaussian p_t = N(0, t) at D = 1/2.
    ("p_1(0) = 1 / sqrt(2 pi)", 0.39894, 5e-6),
    ("d p_t(0) / d t at t = 1", -0.19947, 5e-6),
    ("Laplacian of p_1 at 0", -0.39894, 5e-6),
    ("-p_1(0) / sigma^2 at sigma^2 = 1", -0.39894, 5e-6),
    ("D times the Laplacian of p_1 at 0", -0.19947, 5e-6),
    ("largest heat-equation residual", 0.0, 1e-6),
    # Section 1.32.2, the Fokker-Planck equation and the two equations it reduces to.
    ("center of p_t at t = 1", 0.5, EXACT),
    ("largest residual, drift and noise", 0.0, 1e-6),
    ("largest residual at sigma_t = 0", 0.0, 1e-6),
]


@pytest.mark.parametrize("label, chapter, tolerance", FIELDS_ROWS)
def test_fields(fields_worked, label, chapter, tolerance):
    assert fields_worked[label] == pytest.approx(chapter, abs=tolerance)


def test_the_heat_equation_matches_the_time_derivative_at_the_center(fields_worked):
    """Section 1.32.1: D times the Laplacian at the origin equals d p_t(0)/dt to every digit."""
    assert fields_worked["D times the Laplacian of p_1 at 0"] == pytest.approx(
        fields_worked["d p_t(0) / d t at t = 1"], abs=EXACT)


def test_the_simulated_walk_tracks_the_density_it_solves():
    """Section 1.32.2: dX_t = b dt + sigma dW_t has density N(bt, t), so the variance tracks t."""
    paths = 20000  # Enough paths that four standard errors is a tight band on both moments.
    set_seed(SEED)
    walk = fields.euler_maruyama(paths, 200, quiet=True)
    assert abs(walk["mean"] - 0.5) < 4.0 * math.sqrt(1.0 / paths)
    assert abs(walk["variance"] - 1.0) < 4.0 * math.sqrt(2.0 / paths)
    assert walk["variance_gap"] < 0.05


# %% 07. Every script runs from the command line
@pytest.mark.parametrize("script", ["probability", "linear_algebra", "calculus", "information",
                                    "optimization", "fields"])
def test_the_script_runs_clean(script):
    """Rule 7 of the conventions: --steps 20 --quiet finishes in a few seconds on a CPU."""
    finished = subprocess.run(
        [sys.executable, f"lecture_1/{script}.py", "--steps", "20", "--quiet"],
        cwd=ROOT, env={"PYTHONPATH": str(ROOT), "PATH": "/usr/bin:/bin"},
        capture_output=True, text=True, timeout=120,
    )
    assert finished.returncode == 0, finished.stderr
    assert finished.stderr == ""
    assert "Sections" in finished.stdout