Delete tests

tests/conftest.py

@@ -1,23 +0,0 @@
-"""Shared setup for the test suite.
-
-Several lecture files print five decimal places of an arithmetic result, which float32 cannot
-separate, so those files raise the default precision when they are imported. A test may import a
-file that does not, so the suite raises the precision here for every test alike.
-
-This runs at module scope, not in a fixture, because pytest imports a conftest before it
-imports or collects anything else, and a module-scoped fixture in a test file would otherwise
-build its values first, in the lower precision.
-"""
-
-import pytest
-import torch
-
-torch.set_default_dtype(torch.float64)  # Match what every script does in main().
-torch.set_num_threads(1)  # These networks are small; one thread is the fastest here.
-
-
-@pytest.fixture(scope="session", autouse=True)
-def _restore_defaults():
-    """Put the process back the way it was once the session ends."""
-    yield
-    torch.set_default_dtype(torch.float32)

tests/test_lecture_1.py

@@ -1,554 +0,0 @@
-"""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: a value given to three decimals is checked to 5e-4, one given to five
-decimals to 5e-6, and a value the chapter states exactly is checked 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, which is 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

tests/test_lecture_2.py

@@ -1,235 +0,0 @@
-"""The worked numbers of Chapter 2, asserted against the Lecture 2 scripts.
-
-Every figure here is quoted in the chapter: the Euler landings at one, two and four steps, the flux
-balance on [-1, 1], both sides of the continuity equation at x = 2, the single endpoint pair, the
-average at the crossing, the nine-unit gap between the two objectives, the table of three
-interpolants at t = 0.25, and the parameter count of the MNIST U-Net. Nothing in this file trains.
-"""
-
-import math
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # Run from a fresh clone.
-
-from lecture_2 import flow_matching_2d as fm2d  # noqa: E402
-from lecture_2 import interpolants as interp  # noqa: E402
-from lecture_2 import velocity_fields as vf  # noqa: E402
-
-LOG2 = math.log(2.0)
-
-
-# %% 01. Velocity fields, trajectories and the flow map (Sections 2.2 to 2.4)
-
-
-def test_trajectories_match_the_chapter():
-    report = vf.trajectories()
-    assert report["constant_field_trajectory"] == [1.0, 1.5, 2.0, 2.5, 3.0]
-    assert report["expansion_x_half_from_two"] == pytest.approx(2.828, abs=1e-3)
-    assert report["expansion_endpoints"] == [-2.0, 0.0, 2.0]
-    assert report["expansion_outer_spacing"] == pytest.approx(4.0)
-    assert report["contraction_endpoints"] == [-0.5, 0.0, 0.5]
-    assert report["contraction_outer_spacing"] == pytest.approx(1.0)
-
-
-def test_rotation_preserves_the_radius():
-    report = vf.trajectories()
-    assert report["rotation_quarter_turn"] == [0.0, 1.0]
-    assert report["rotation_half_turn"] == [-1.0, 0.0]
-    assert report["rotation_radius"] == pytest.approx(1.0, abs=1e-6)
-    assert report["rotation_squared_norm_rate"] == pytest.approx(0.0, abs=1e-6)
-
-
-def test_flow_map_is_the_identity_at_zero_and_doubles_at_one():
-    report = vf.flow_map_report()
-    assert report["flow_map_identity_error"] == pytest.approx(0.0, abs=1e-6)
-    assert report["flow_map_jacobian_det"] == pytest.approx(2.0, abs=1e-6)
-    assert report["trajectory_gap_minimum"] == pytest.approx(2.0, abs=1e-6)
-
-
-# %% 02. Euler's method (Section 2.5)
-
-
-def test_four_euler_steps_land_at_the_chapter_value():
-    landing = vf.euler(lambda x, t: x, torch.tensor(1.0), 4)  # dx/dt = x from x_0 = 1.
-    assert float(landing) == pytest.approx(2.44140625, abs=1e-6)
-
-
-def test_euler_ladder_errors_and_ratios():
-    report = vf.euler_ladder(4)
-    assert report["euler_step_counts"] == [1, 2, 4]
-    assert report["euler_landings"] == [2.0, 2.25, pytest.approx(2.4414, abs=1e-4)]
-    assert report["euler_errors"] == [pytest.approx(0.7183, abs=1e-4),
-                                      pytest.approx(0.4683, abs=1e-4),
-                                      pytest.approx(0.2769, abs=1e-4)]
-    assert report["euler_error_halving_ratios"] == [pytest.approx(1.53, abs=5e-3),
-                                                    pytest.approx(1.69, abs=5e-3)]
-    assert report["euler_relative_error_four_steps"] == pytest.approx(0.1019, abs=1e-4)
-
-
-def test_the_euler_error_closes_like_one_over_the_step_count():
-    report = vf.euler_ladder(256)  # The error times the step count settles at e/2 = 1.359.
-    scaled = report["euler_error_times_steps"]
-    assert scaled[-1] == pytest.approx(math.e / 2.0, abs=0.02)
-    assert report["euler_error_halving_ratios"][-1] == pytest.approx(2.0, abs=0.02)
-
-
-# %% 03. The pushforward, the flux and the continuity equation (Sections 2.6 and 2.7)
-
-
-def test_the_expansion_field_doubles_the_standard_deviation():
-    report = vf.pushforward_moments(4096, seed=0)
-    assert report["pushforward_mean"] == pytest.approx(0.0, abs=0.05)
-    assert report["pushforward_variance"] == pytest.approx(4.0, abs=0.2)
-
-
-def test_flux_balance_on_the_unit_interval():
-    report = vf.flux_balance()
-    assert report["density_at_one"] == pytest.approx(0.2420, abs=1e-4)
-    assert report["flux_at_minus_one"] == pytest.approx(-0.1677, abs=1e-4)
-    assert report["flux_at_one"] == pytest.approx(0.1677, abs=1e-4)
-    assert report["flux_balance"] == pytest.approx(-0.3354, abs=1e-4)
-    assert report["enclosed_mass"] == pytest.approx(0.6827, abs=1e-4)
-    # The balance and the derivative of the enclosed mass agree to four figures.
-    assert report["enclosed_mass_rate"] == pytest.approx(report["flux_balance"], abs=1e-4)
-
-
-def test_both_sides_of_the_continuity_equation_at_x_equals_two():
-    report = vf.continuity_check()
-    assert report["density_at_two"] == pytest.approx(0.05399, abs=1e-5)
-    assert report["density_gradient_at_two"] == pytest.approx(-0.10798, abs=1e-5)
-    assert report["velocity_divergence"] == pytest.approx(LOG2, abs=1e-6)
-    assert report["continuity_time_derivative"] == pytest.approx(0.11227, abs=1e-5)
-    assert report["continuity_minus_flux_divergence"] == pytest.approx(0.11227, abs=1e-5)
-
-
-def test_log_density_along_a_trajectory_falls_by_log_two():
-    report = vf.log_density_along_trajectory()
-    assert report["log_density_rate"] == pytest.approx(-LOG2)
-    assert report["log_density_change"] == pytest.approx(-LOG2, abs=1e-6)
-    assert report["log_density_accumulated"] == pytest.approx(-LOG2)
-    assert report["prior_density_at_origin"] == pytest.approx(0.3989, abs=1e-4)
-    assert report["final_density_at_origin"] == pytest.approx(0.1995, abs=1e-4)
-    assert report["density_ratio_at_origin"] == pytest.approx(2.0, abs=1e-6)
-
-
-# %% 04. The conditional path and the crossing (Sections 2.12 and 2.13)
-
-
-def test_the_single_endpoint_pair():
-    report = fm2d.worked_pair()
-    assert report["worked_pair_path"] == [-2.0, -0.75, 0.5, 1.75, 3.0]
-    assert report["worked_pair_velocity"] == pytest.approx(5.0)
-
-
-def test_the_average_at_the_crossing():
-    report = fm2d.crossing_average()
-    assert report["crossing_state"] == pytest.approx(2.0)
-    assert report["crossing_velocities"] == [4.0, -2.0]
-    assert report["crossing_average_velocity"] == pytest.approx(1.0)
-
-
-def test_the_two_losses_differ_by_nine_at_the_crossing():
-    report = fm2d.crossing_gap()
-    assert report["crossing_conditional_loss"] == [10.0, 9.0, 10.0]
-    assert report["crossing_marginal_loss"] == [1.0, 0.0, 1.0]
-    assert report["crossing_constant_gap"] == [9.0, 9.0, 9.0]
-    assert report["crossing_target_variance"] == pytest.approx(9.0)
-
-
-# %% 05. The marginal velocity and the equivalence (Sections 2.13 and 2.15)
-
-
-def test_the_marginal_velocity_at_the_two_ends():
-    x = torch.randn(64, 2)
-    # v_0(x) = E[X_1] - x = -x, and v_1(x) = x - E[X_0] = x, since the coupling is independent.
-    assert torch.allclose(fm2d.marginal_velocity(x, torch.zeros(64)), -x, atol=1e-5)
-    assert torch.allclose(fm2d.marginal_velocity(x, torch.ones(64)), x, atol=1e-5)
-
-
-def test_the_marginal_velocity_transports_the_prior_onto_the_data():
-    from cis6270.metrics import mean_log_density, mode_report
-
-    torch.manual_seed(0)
-    x_t = torch.randn(512, 2)  # X_0 ~ p_0.
-    for step in range(50):  # Euler on the exact field, with no network anywhere.
-        t = torch.full((512,), step / 50.0)
-        x_t = x_t + fm2d.marginal_velocity(x_t, t) / 50.0
-    assert mode_report(x_t)["modes_found"] == 4
-    assert mean_log_density(x_t) > -3.0  # The data themselves measure -2.14.
-
-
-def test_the_losses_decompose_exactly_on_any_batch():
-    from cis6270.nets import VelocityMLP
-
-    torch.manual_seed(0)
-    model = VelocityMLP(dim=2, width=32)
-    x_0, x_1 = torch.randn(256, 2), torch.randn(256, 2) + 2.0
-    t = torch.rand(256)
-    loss_cfm, loss_fm, gap, cross = fm2d.objectives(model, x_0, x_1, t)
-    # Expanding the square gives L_CFM = L_FM + cross + gap term by term, on every batch.
-    assert float(loss_cfm) == pytest.approx(float(loss_fm + cross + gap), rel=1e-5)
-
-
-def test_the_two_objectives_share_a_gradient():
-    from cis6270.data import gaussian_mixture_2d
-    from cis6270.nets import VelocityMLP
-
-    torch.manual_seed(0)
-    model = VelocityMLP(dim=2, width=64)
-    data = gaussian_mixture_2d(4096, seed=0)
-    report = fm2d.equivalence_check(model, data, 8192)
-    assert report["loss_cfm"] == pytest.approx(report["loss_fm_plus_gap"], abs=0.3)
-    assert abs(report["loss_cross_term"]) < 0.3  # Zero in expectation, small on a large batch.
-    assert report["gradient_cosine"] > 0.99  # The residual is the batch's own sampling error.
-
-
-# %% 06. The three interpolants (Section 2.18)
-
-
-def test_every_schedule_reaches_the_right_endpoints():
-    for name in interp.SCHEDULES:
-        assert interp.endpoint_and_table(name)[f"{name}_endpoint_error"] == pytest.approx(0.0,
-                                                                                          abs=1e-6)
-
-
-def test_the_table_of_three_interpolants_at_one_quarter():
-    expected = {  # Table 2.1, at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
-        "linear": ([1.50, 0.50], [-2.0, 2.0]),
-        "quadratic": ([1.875, 0.125], [-1.0, 1.0]),
-        "trigonometric": ([1.848, 0.765], [-1.202, 2.902]),
-    }
-    for name, (position, velocity) in expected.items():
-        report = interp.endpoint_and_table(name)
-        assert report[f"{name}_position_at_quarter"] == pytest.approx(position, abs=1e-3)
-        assert report[f"{name}_velocity_at_quarter"] == pytest.approx(velocity, abs=1e-3)
-
-
-def test_only_the_trigonometric_schedule_keeps_identity_covariance():
-    trig = interp.covariance_check("trigonometric", 8192, seed=0)
-    assert trig["trigonometric_coefficient_squares"] == [1.0, 1.0, 1.0, 1.0, 1.0]
-    assert trig["trigonometric_covariance_deviation"] < 0.05
-    linear = interp.covariance_check("linear", 8192, seed=0)
-    assert linear["linear_coefficient_squares"] == [1.0, 0.625, 0.5, 0.625, 1.0]
-    assert linear["linear_covariance_deviation"] > 0.4  # Cov(X_{1/2}) is half the identity.
-
-
-def test_the_quadratic_schedule_has_covered_a_sixteenth_at_one_quarter():
-    # beta_{1/4} = (1/4)^2 = 0.0625, so the state has covered 6.25 percent of the segment.
-    _, beta, _, _ = interp.quadratic(torch.tensor([0.25]))
-    assert float(beta) == pytest.approx(0.0625)
-    _, beta_half, _, _ = interp.quadratic(torch.tensor([0.5]))
-    assert float(beta_half) == pytest.approx(0.25)  # A quarter of the distance at the midpoint.
-
-
-# %% 07. The MNIST U-Net (Section 2.20)
-
-
-def test_the_unet_has_the_parameter_count_the_chapter_quotes():
-    from lecture_2 import flow_matching_mnist as mnist
-
-    model = mnist.FlowUNet()  # The default base width of 32, so the widths are 32, 64 and 128.
-    assert sum(p.numel() for p in model.parameters()) == 471_265

tests/test_lecture_3.py

@@ -1,445 +0,0 @@
-"""Lecture 3: the numbers Chapter 3 works by hand, and the invariants the code must keep.
-
-Every assertion here is either a quantity the chapter computes in the text or a property the
-derivation guarantees. Nothing trains, so the whole file runs in a few seconds.
-"""
-
-import math
-
-import pytest
-import torch
-
-from cis6270.metrics import mode_report
-from lecture_3 import (
-    ddpm, forward_diffusion, guidance, latent_diffusion, samplers, score_matching, sde,
-)
-
-
-# %% The forward process, Sections 3.2 through 3.10
-
-
-def test_euler_maruyama_step_matches_the_worked_numbers():
-    landed, drift, noise = sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, 1.5)
-    assert landed == pytest.approx(2.11)
-    assert (drift, noise) == pytest.approx((-0.04, 0.15))
-
-
-def test_three_draws_average_above_the_drift_prediction():
-    endpoints = [sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, e)[0] for e in (1.5, -0.5, 0.2)]
-    assert endpoints == pytest.approx([2.11, 1.91, 1.98])
-    assert sum(endpoints) / 3.0 == pytest.approx(2.00, abs=1e-6)  # The drift alone reaches 1.96.
-
-
-def test_increment_scales_as_the_square_root_of_the_step():
-    assert sde.increment_scale(0.04)["std"] == pytest.approx(0.2)
-    assert sde.increment_scale(0.02)["std"] == pytest.approx(0.1414, abs=1e-4)
-    assert sde.increment_scale(0.04, -1.0, 0.5)["noise_over_drift"] == pytest.approx(2.5)
-
-
-def test_ornstein_uhlenbeck_law_at_log_two():
-    law = sde.ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), math.log(2.0))
-    assert law["mean"] == pytest.approx(1.0)
-    assert law["variance"] == pytest.approx(0.75)
-    assert law["std"] == pytest.approx(0.8660, abs=1e-4)
-    assert law["stationary_variance"] == pytest.approx(1.0)
-
-
-def test_ito_recovers_the_variance_and_the_chain_rule_alone_does_not():
-    generator = torch.Generator().manual_seed(0)
-    check = sde.check_ito_on_square(8192, 200, 0.25, generator)
-    assert check["ito_mean"] == pytest.approx(0.25, abs=0.02)
-    assert check["realized_mean_square"] == pytest.approx(0.25, abs=0.02)
-    assert abs(check["no_curvature_mean"]) < 0.02  # The curvature term carries all the growth.
-
-
-def test_fokker_planck_holds_on_the_spreading_gaussian():
-    grid = torch.linspace(-3.0, 3.0, 41)
-    for t in (0.0, 0.5, 2.0):
-        residual = sde.fokker_planck_residual(grid, torch.full_like(grid, t))
-        assert float(residual.abs().max()) < 1e-6
-
-
-def test_peak_of_the_gaussian_falls_at_the_rate_the_equation_gives():
-    peak = sde.peak_extrapolation(0.1)
-    assert peak["p_0"] == pytest.approx(0.3989, abs=1e-4)
-    assert peak["rate"] == pytest.approx(-0.1995, abs=1e-4)
-    assert peak["linear"] == pytest.approx(0.3790, abs=1e-4)
-    assert peak["exact"] == pytest.approx(0.3804, abs=1e-4)
-    assert peak["gap"] == pytest.approx(0.0014, abs=1e-4)
-
-
-# %% Score matching, Sections 3.14 through 3.17
-
-
-def test_posterior_average_at_the_running_observation():
-    report = score_matching.posterior_report(7.4, (10.0, 9.0), (0.3, 0.7), epsilon_theta=-0.8)
-    assert report["weights"] == pytest.approx([0.2156, 0.7844], abs=1e-4)
-    assert report["clean_mean"] == pytest.approx(9.2156, abs=1e-4)
-    assert report["conditional_scores"] == pytest.approx([5 / 3, -5 / 9], abs=1e-4)
-    assert report["marginal"] == pytest.approx([-0.0765, -0.0765], abs=1e-4)
-    assert report["residual"] == pytest.approx(0.835, abs=1e-3)
-    assert report["score_hat"] == pytest.approx(4 / 3, abs=1e-4)
-    assert report["x0"] == pytest.approx([9.85, 9.85, 9.2156], abs=1e-4)
-
-
-def test_one_parameter_fit_is_minimized_at_the_inverse_variance():
-    assert score_matching.one_parameter_loss(0.5, 1.0) == pytest.approx(-0.375)
-    assert score_matching.one_parameter_loss(1.0, 1.0) == pytest.approx(-0.5)
-    assert score_matching.one_parameter_loss(0.25, 4.0) == pytest.approx(-0.125)
-    grid = [0.05 * i for i in range(1, 41)]
-    best = min(grid, key=lambda a: score_matching.one_parameter_loss(a, 4.0))
-    assert best == pytest.approx(0.25)  # a = 1 / v with v = 4.
-
-
-def test_the_three_objectives_share_one_minimizer():
-    generator = torch.Generator().manual_seed(1)
-    clean = score_matching.sample_clean(32768, generator)
-    epsilon = torch.randn(clean.shape, generator=generator)
-    found = score_matching.minimizers(clean, epsilon)
-    assert found["explicit"] == pytest.approx(1.0)
-    assert found["implicit"] == pytest.approx(1.0, abs=0.03)
-    assert found["denoising"] == pytest.approx(1.0, abs=0.03)
-
-
-def test_integration_by_parts_holds_on_the_exact_field():
-    generator = torch.Generator().manual_seed(2)
-    clean = score_matching.sample_clean(32768, generator)
-    points = score_matching.A_T * clean + score_matching.SIGMA_T * torch.randn(
-        clean.shape, generator=generator,
-    )
-    identity = score_matching.integration_by_parts_gap(score_matching.exact_score, points)
-    assert identity["gap"] < 0.05 * abs(identity["left"])
-
-
-def test_hutchinson_estimates_the_trace_the_exact_loop_computes():
-    generator = torch.Generator().manual_seed(3)
-    means = torch.randn(2, score_matching.TRACE_DIM, generator=generator) * 1.5
-    points = torch.randn(8, score_matching.TRACE_DIM, generator=generator)
-
-    def field(z):
-        return score_matching.high_dimensional_score(z, means, 1.0)
-
-    exact, passes = score_matching.exact_trace(field, points)
-    assert passes == score_matching.TRACE_DIM  # One backward pass per coordinate.
-    estimate, draws = score_matching.hutchinson_trace(field, points, 256, generator)
-    error = float(draws.std(dim=0).mean()) / math.sqrt(256)
-    assert float((estimate - exact).mean().abs()) < 4.0 * error
-
-
-# %% The schedule, the corruption and DDPM, Sections 3.11, 3.12, 3.19 and 3.24 through 3.28
-
-
-def test_schedule_matches_the_continuous_curve():
-    assert float(forward_diffusion.alphas_cumprod[207]) == pytest.approx(0.6402, abs=1e-4)
-    assert float(forward_diffusion.a[207]) == pytest.approx(0.8001, abs=1e-4)
-    assert float(forward_diffusion.sigma[207]) == pytest.approx(0.5999, abs=1e-4)
-    assert forward_diffusion.continuous_signal(0.207) == pytest.approx(0.7997, abs=1e-4)
-    assert float(forward_diffusion.alphas_cumprod[1000]) == pytest.approx(4.036e-5, rel=1e-3)
-    assert float(forward_diffusion.a[1000]) == pytest.approx(0.0064, abs=1e-4)
-    assert forward_diffusion.continuous_signal(1.0) == pytest.approx(0.0066, abs=1e-4)
-    assert float(forward_diffusion.a[500]) == pytest.approx(0.2803, abs=1e-4)
-
-
-def test_every_copy_of_the_schedule_agrees():
-    """No lecture file imports another, so each rebuilds the schedule. The copies must match."""
-    for name in ("betas", "alphas", "alphas_cumprod"):
-        for other in (ddpm, latent_diffusion):
-            assert torch.equal(getattr(forward_diffusion, name), getattr(other, name)), name
-    assert torch.equal(ddpm.posterior_variance[1:], latent_diffusion.posterior_variance[1:])
-    # The samplers read the same path off the continuous curve, which agrees at a fine grid.
-    for grid in (samplers.make_schedule(1000), guidance.schedule_tensors(1000)):
-        assert float(grid["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3)
-        assert torch.allclose(grid["alphas_cumprod"], ddpm.alphas_cumprod, atol=1e-3)
-
-
-def test_noise_levels_of_the_continuous_schedule():
-    for signal, time, rate in ((0.98, 0.0589, 1.2721), (0.8, 0.2068, 4.2157),
-                               (0.35, 0.4544, 9.1420), (0.04, 0.7994, 16.0073)):
-        located = forward_diffusion.level_time(signal)
-        assert located["t"] == pytest.approx(time, abs=1e-4)
-        assert located["beta_t"] == pytest.approx(rate, abs=1e-3)
-        assert located["sigma_t"] == pytest.approx(math.sqrt(1.0 - signal ** 2), abs=1e-9)
-    assert forward_diffusion.level_time(0.8)["sigma_t"] == pytest.approx(0.6, abs=1e-4)
-
-
-def test_dropped_elbo_weight_at_step_five_hundred():
-    weight = ddpm.elbo_weight(
-        ddpm.betas[500], ddpm.posterior_variance[500],
-        ddpm.alphas[500], ddpm.alphas_cumprod[500],
-    )
-    assert float(ddpm.posterior_variance[500]) == pytest.approx(0.01003, abs=1e-5)
-    assert float(weight) == pytest.approx(0.005508, abs=1e-6)
-    assert ddpm.elbo_weight(0.2, 1 / 9, 0.8, 0.64) == pytest.approx(0.625, abs=1e-6)
-
-
-def test_one_pixel_through_the_forward_code():
-    a_k = float(forward_diffusion.a[207])
-    sigma_k = float(forward_diffusion.sigma[207])
-    assert a_k * 0.5 + sigma_k * (-1.0) == pytest.approx(-0.1998, abs=1e-4)
-
-
-def test_q_sample_uses_the_supplied_noise_exactly():
-    x_0 = torch.randn(4, 1, 28, 28)
-    noise = torch.randn_like(x_0)
-    k = torch.tensor([1, 207, 500, 1000])
-    for module in (forward_diffusion, ddpm):
-        x_k, returned = module.q_sample(x_0, k, noise)
-        abar = module.alphas_cumprod[k].view(-1, 1, 1, 1)
-        assert torch.allclose(x_k, abar.sqrt() * x_0 + (1 - abar).sqrt() * noise, atol=1e-6)
-        assert returned is noise
-
-
-def test_corruption_law_matches_the_gaussian_it_is_drawn_from():
-    generator = torch.Generator().manual_seed(0)
-    x_0 = torch.randn(1, 1, 28, 28, generator=generator)
-    levels = [1, 207, 500, 1000]
-    law = forward_diffusion.measure_law(x_0, levels, 1024, generator)
-    assert law["mean_gap"] < 0.05  # Centered on a_k x_0, to one over the square root of 1024.
-    assert law["variance_gap"] < 0.05  # Spread sigma_k^2.
-    # Four times the draws halves the gap, which a systematic error in a coefficient would not.
-    finer = forward_diffusion.measure_law(x_0, levels, 4096, generator)
-    assert finer["mean_gap"] < 0.6 * law["mean_gap"]
-
-
-def test_worked_training_batch_at_abar_point_eight():
-    batch = forward_diffusion.worked_batch(
-        0.8, (2.0, -1.0, 0.5), (0.5, -0.2, 1.0), (0.4, -0.1, 0.8),
-    )
-    assert batch["a_k"] == pytest.approx(0.8944, abs=1e-4)
-    assert batch["sigma_k"] == pytest.approx(0.4472, abs=1e-4)
-    assert batch["x_k"] == pytest.approx([2.0125, -0.9839, 0.8944], abs=1e-3)
-    assert batch["loss"] == pytest.approx(0.02, abs=1e-6)
-    assert batch["score_targets"] == pytest.approx([-1.118, 0.447, -2.236], abs=1e-3)
-    assert batch["x_0_hat"] == pytest.approx([2.05, -1.05, 0.60], abs=1e-2)
-
-
-def test_two_step_reverse_transition_agrees_in_both_forms():
-    step = ddpm.two_step_transition(0.2, 7.4, -1.0, 0.3)
-    assert step["x_0_hat"] == pytest.approx(10.0, abs=1e-6)
-    assert step["mean_from_noise"] == pytest.approx(8.646, abs=1e-3)
-    assert step["mean_from_clean"] == pytest.approx(step["mean_from_noise"], abs=1e-6)
-    assert step["weights"] == pytest.approx([0.4969, 0.4969], abs=1e-4)
-    assert step["variance"] == pytest.approx(1 / 9)
-    assert step["x_{k-1}"] == pytest.approx(8.746, abs=1e-3)
-    assert step["elbo_weight"] == pytest.approx(0.625, abs=1e-6)
-
-
-def test_unet_keeps_the_image_shape_and_the_last_reverse_step_adds_no_noise():
-    torch.manual_seed(0)
-    model = ddpm.UNet(in_channels=1, base_channels=8, time_embed_dim=32)
-    x = torch.randn(2, 1, 28, 28)
-    k = torch.tensor([1, 1])
-    assert model(x, k).shape == x.shape
-    first, second = ddpm.p_sample(model, x, k), ddpm.p_sample(model, x, k)
-    assert torch.allclose(first, second)  # keep = 0 at k = 1, so the step is the mean.
-
-
-# %% The three samplers, Sections 3.22, 3.26, 3.32, 3.34 and 3.35
-
-
-def test_one_backward_step_under_both_dynamics():
-    step = samplers.one_step_comparison(1.2, 2.0, 0.5, 0.8, 0.1)
-    assert step["score"] == pytest.approx(-1.6)
-    assert step["velocity"] == pytest.approx(0.4)
-    assert step["ode_next"] == pytest.approx(1.16)
-    assert step["sde_drift"] == pytest.approx(2.0)
-    assert step["sde_mean"] == pytest.approx(1.0)
-    assert step["sde_noise_std"] == pytest.approx(0.4472, abs=1e-4)
-
-
-def test_langevin_step_from_the_worked_state():
-    assert samplers.langevin_step(2.0, -2.0, 0.01, 0.5) == pytest.approx(2.0507, abs=1e-4)
-
-
-def test_schedule_stays_a_valid_discretization_at_any_grid_size():
-    for num_steps in (20, 200, 1000):
-        schedule = samplers.make_schedule(num_steps)
-        assert float(schedule["betas"][1:].max()) < 1.0
-        assert float(schedule["alphas_cumprod"][0]) == pytest.approx(1.0)
-        assert float(schedule["alphas_cumprod"][num_steps]) < 1e-4  # Terminal law is Gaussian.
-    thousand = samplers.make_schedule(1000)
-    assert float(thousand["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3)
-
-
-def test_every_sampler_finds_all_four_modes():
-    budget = 60
-    for name, points in (
-        ("ancestral", samplers.ancestral_sample(
-            samplers.ExactScore(), samplers.make_schedule(budget), 512,
-            torch.Generator().manual_seed(0))),
-        ("flow", samplers.probability_flow_sample(
-            samplers.ExactScore(), 512, budget, torch.Generator().manual_seed(0))),
-        ("langevin", samplers.annealed_langevin_sample(
-            samplers.ExactScore(), 512, [0.7 - 0.69 * i / 19 for i in range(20)], 3, 0.15,
-            torch.Generator().manual_seed(0))),
-    ):
-        measured = mode_report(points)
-        assert measured["modes_found"] == 4, name
-        assert measured["mode_entropy"] > 1.30, name  # log 4 = 1.386 when the modes are even.
-
-
-def test_each_sampler_reports_the_evaluations_it_spent():
-    score = samplers.ExactScore()
-    samplers.probability_flow_sample(score, 64, 37, torch.Generator().manual_seed(0))
-    assert score.calls == 37
-    score = samplers.ExactScore()
-    samplers.annealed_langevin_sample(score, 64, [0.5, 0.2, 0.05], 4, 0.1,
-                                      torch.Generator().manual_seed(0))
-    assert score.calls == 12  # Three levels, four iterations each.
-
-
-# %% Guidance, Sections 3.37 through 3.45
-
-
-def test_worked_guided_score():
-    score, gradient = torch.tensor([[-1.0, -0.1]]), torch.tensor([[1.3, 1.1]])
-    at_one = guidance.guided_score(score, gradient, 1.0)[0]
-    at_two = guidance.guided_score(score, gradient, 2.0)[0]
-    assert at_one.tolist() == pytest.approx([0.3, 1.0], abs=1e-6)
-    assert float(at_one.norm()) == pytest.approx(1.0440, abs=1e-4)
-    assert at_two.tolist() == pytest.approx([1.6, 2.1], abs=1e-6)
-    assert float(at_two.norm()) == pytest.approx(2.6401, abs=1e-4)
-
-
-def test_classifier_free_step_and_the_score_it_implies():
-    eps = guidance.classifier_free_field(torch.tensor(0.50), torch.tensor(0.10), 2.0)
-    assert float(eps) == pytest.approx(-0.30, abs=1e-6)
-    assert float(-eps / 0.5) == pytest.approx(0.6, abs=1e-6)
-
-
-def test_guided_backward_step_agrees_in_both_parameterizations():
-    guided_eps = float(guidance.guided_noise(torch.tensor(0.8), torch.tensor(0.6), 2.0, 0.5))
-    assert guided_eps == pytest.approx(0.2, abs=1e-6)
-    velocity = -0.5 * 2.0 * (1.2 + (-guided_eps / 0.5))  # The collected form of the velocity.
-    assert velocity == pytest.approx(-0.8, abs=1e-6)
-    assert 1.2 - 0.1 * velocity == pytest.approx(1.28, abs=1e-6)
-
-
-def test_classifier_gradient_equals_the_difference_of_the_two_fields():
-    points = torch.randn(64, 2, generator=torch.Generator().manual_seed(0))
-    gradient = guidance.classifier_gradient(points, 0.8, 0)
-    difference = (guidance.conditional_score(points, 0.8, 0)
-                  - guidance.unconditional_score(points, 0.8))
-    assert torch.allclose(gradient, difference, atol=1e-6)
-
-
-def test_weights_are_projected_onto_the_simplex():
-    weights = guidance.normalize_weights((-1.0, 3.0, 1.0), torch.device("cpu"))
-    assert float(weights.sum()) == pytest.approx(1.0)
-    assert float(weights[0]) == 0.0  # A negative weight is clamped away.
-
-
-def test_scalarized_rewards_reverse_the_ranking_with_the_weights():
-    first, second = (1.2, 0.2), (0.5, 1.1)
-    assert guidance.scalarize(first, (0.7, 0.3)) == pytest.approx(0.90, abs=1e-4)
-    assert guidance.scalarize(second, (0.7, 0.3)) == pytest.approx(0.68, abs=1e-4)
-    assert guidance.scalarize(first, (0.3, 0.7)) == pytest.approx(0.50, abs=1e-4)
-    assert guidance.scalarize(second, (0.3, 0.7)) == pytest.approx(0.92, abs=1e-4)
-
-
-def test_competing_gradients_cancel_at_the_stated_weights():
-    combined = 0.6 * torch.tensor([1.0, 0.8]) + 0.4 * torch.tensor([0.5, -1.2])
-    assert combined.tolist() == pytest.approx([0.80, 0.0], abs=1e-6)
-
-
-def test_selection_weights_and_the_guidance_schedule():
-    weights = (2.0 * torch.tensor([1.2, 0.8, 0.3])).softmax(0)
-    assert weights.tolist() == pytest.approx([0.6193, 0.2783, 0.1024], abs=1e-4)
-    assert guidance.schedule_value(0.6) == pytest.approx(0.96)
-    assert guidance.schedule_value(0.0) == 0.0 and guidance.schedule_value(1.0) == 0.0
-
-
-def test_guidance_moves_mass_onto_the_target_mode():
-    schedule = guidance.schedule_tensors(60)
-    fractions = []
-    for w in (0.0, 0.25, 1.0):
-        points = guidance.sample_diffusion(
-            schedule, 512, guidance.TARGET_LABEL, w=w,
-            generator=torch.Generator().manual_seed(0),
-        )
-        fractions.append(mode_report(points)["mode_fractions"][guidance.TARGET_LABEL])
-    assert fractions[0] < 0.35  # Unconditional, so about a quarter.
-    assert fractions[0] < fractions[1] < fractions[2]
-    assert fractions[2] > 0.9
-
-
-def test_reward_guidance_prefers_the_right_hand_modes():
-    schedule = guidance.schedule_tensors(60)
-    points = guidance.sample_diffusion(
-        schedule, 512, guidance.TARGET_LABEL, eta=2.0, lambdas=(1.0, 0.0),
-        generator=torch.Generator().manual_seed(0),
-    )
-    fractions = mode_report(points)["mode_fractions"]
-    assert fractions[0] + fractions[1] > 0.9  # The two modes at x = +2.
-
-
-# %% Latent diffusion, Sections 3.21, 3.29 and 3.46
-
-
-def test_latent_coordinate_corrupted_and_carried_back():
-    abar = float(latent_diffusion.alphas_cumprod[500])
-    a_500, sigma_500 = math.sqrt(abar), math.sqrt(1.0 - abar)
-    assert (a_500, sigma_500) == pytest.approx((0.2803, 0.9599), abs=1e-4)
-    z_500 = a_500 * 1.2 + sigma_500 * 0.4
-    assert z_500 == pytest.approx(0.7204, abs=1e-4)
-    eps_hat = (z_500 - a_500 * 1.05) / sigma_500
-    assert eps_hat == pytest.approx(0.4438, abs=1e-4)
-    beta_k, alpha_k = float(latent_diffusion.betas[500]), float(latent_diffusion.alphas[500])
-    mean = (z_500 - beta_k / sigma_500 * eps_hat) / math.sqrt(alpha_k)
-    assert mean == pytest.approx(0.7193, abs=1e-4)
-    variance = float(latent_diffusion.posterior_variance[500])
-    assert math.sqrt(variance) == pytest.approx(0.1002, abs=1e-4)
-
-
-def test_masked_corruption_leaves_padding_at_zero():
-    z0 = torch.randn(3, 8, 16)
-    mask = torch.ones(3, 8)
-    mask[:, 5:] = 0.0  # The last three positions are padding.
-    z_k, noise = latent_diffusion.q_sample_latent(z0, torch.tensor([1, 500, 1000]), mask)
-    assert torch.count_nonzero(z_k[:, 5:, :]) == 0
-    assert noise.shape == z0.shape
-
-
-def test_latent_loss_counts_valid_residues_only():
-    z0 = torch.zeros(2, 4, 3)
-    z0_hat = torch.ones(2, 4, 3)  # Unit error at every coordinate.
-    mask = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]])
-    assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0)
-    z0_hat[:, 2:, :] = 100.0  # A large error, but only on padding.
-    assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0)
-
-
-def test_denoiser_returns_a_latent_of_the_same_shape():
-    torch.manual_seed(0)
-    model = latent_diffusion.LatentDenoiser(embed_dim=16, num_layers=1, num_heads=2, max_len=8)
-    z_k = torch.randn(2, 8, 16)
-    mask = torch.ones(2, 8)
-    assert model(z_k, torch.tensor([7, 900]), mask).shape == z_k.shape
-
-
-def test_composition_proxies_on_a_known_peptide():
-    charge, polar = latent_diffusion.composition_proxies("KRDE")
-    assert charge == pytest.approx(0.0)  # Two positive and two negative over four residues.
-    assert polar == pytest.approx(1.0)  # All four are in the polar set.
-    assert latent_diffusion.composition_proxies("AAAA") == pytest.approx((0.0, 0.0))
-
-
-def test_enforce_polar_reaches_the_minimum_at_the_cheapest_positions():
-    torch.manual_seed(0)
-    logits = torch.randn(2, 12, 20)
-    logits[:, :, latent_diffusion.POLAR_COLUMNS] -= 5.0  # Make the argmax non-polar everywhere.
-    choices = logits.argmax(dim=-1)
-    sequences = latent_diffusion.enforce_polar(logits, choices, 5)
-    for sequence in sequences:
-        assert len(sequence) == 12
-        assert sum(a in latent_diffusion.POLAR_RESIDUES for a in sequence) >= 5
-    unchanged = latent_diffusion.enforce_polar(logits, choices, 0)  # Nothing is missing.
-    assert unchanged == ["".join(latent_diffusion.AMINO_ACIDS[int(i)] for i in row)
-                         for row in choices]
-
-
-def test_the_cheapest_two_substitutions_cost_what_the_notes_says():
-    costs = torch.tensor([0.9, 1.4, 2.6, 3.1, 3.3, 4.0, 4.2, 5.5, 6.1])
-    cheapest, _ = costs.topk(2, largest=False)
-    assert cheapest.tolist() == pytest.approx([0.9, 1.4])
-    assert float(cheapest.sum()) == pytest.approx(2.3, abs=1e-5)

tests/test_lecture_4.py

@@ -1,535 +0,0 @@
-"""Lecture 4: the worked numbers of Chapter 4, asserted against the files that print them.
-
-Every assertion here is a number the chapter computes by hand. The tests import the lecture files
-directly, so they check the same functions the scripts run, and none of them trains anything.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-LECTURE = Path(__file__).resolve().parents[1] / "lecture_4"
-
-
-def load(name: str):
-    """Import one lecture file by path, since lecture_4 is a folder of scripts and not a package."""
-    spec = importlib.util.spec_from_file_location(f"lecture_4_{name}", LECTURE / f"{name}.py")
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-ctmc = load("ctmc")
-mdlm = load("mdlm")
-uniform = load("uniform")
-d3pm = load("d3pm")
-sedd = load("sedd")
-block = load("block")
-guidance = load("guidance")
-planner = load("planner")
-search = load("search")
-simplex = load("simplex")
-
-
-# %% 01. ctmc: rates, the master equation, the exponential and the reversal
-
-
-def test_two_rules_and_the_dna_row():
-    """A jumps to C at rate 2 and to G and T at rate 1, so its diagonal entry is -4."""
-    assert ctmc.obeys_the_two_rules(ctmc.DNA_RATES)
-    assert float(ctmc.DNA_RATES[0, 0]) == pytest.approx(-4.0)
-    assert float(ctmc.exit_rates(ctmc.DNA_RATES)[0]) == pytest.approx(4.0)
-
-
-def test_short_interval_row_and_the_step_bound():
-    """P_{t,t+0.1}(A, .) = (0.6, 0.2, 0.1, 0.1), and the step must satisfy h <= 0.25."""
-    row = ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.1)
-    assert [round(float(v), 4) for v in row] == [0.6, 0.2, 0.1, 0.1]
-    assert math.exp(-0.4) == pytest.approx(0.6703, abs=5e-5)
-    assert ctmc.largest_valid_step(ctmc.DNA_RATES) == pytest.approx(0.25)
-    assert float(ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.3)[0]) == pytest.approx(-0.2)
-
-
-def test_waiting_time_from_the_chapters_uniform():
-    """Inverting S(s) = exp(-4s) at u = 0.42 gives a wait of 0.2169."""
-    assert -math.log(0.42) / 4.0 == pytest.approx(0.2169, abs=5e-5)
-    assert math.log(2.0) / 4.0 == pytest.approx(0.1733, abs=5e-5)
-
-
-def test_master_equation_on_four_dna_probabilities():
-    """p_dot = (0.6, 0.2, -0.2, -0.6) and one Euler step gives (0.16, 0.22, 0.28, 0.34)."""
-    derivative = ctmc.master_equation(ctmc.START, ctmc.UNIT_RATES)
-    assert [round(float(v), 4) for v in derivative] == [0.6, 0.2, -0.2, -0.6]
-    arriving, leaving = ctmc.arrivals_and_departures(ctmc.START, ctmc.UNIT_RATES)
-    assert float(arriving[0]) == pytest.approx(0.9)
-    assert float(leaving[0]) == pytest.approx(0.3)
-    stepped = ctmc.START + 0.1 * derivative
-    assert [round(float(v), 4) for v in stepped] == [0.16, 0.22, 0.28, 0.34]
-    assert float(derivative.sum()) == pytest.approx(0.0, abs=1e-12)
-
-
-def test_exact_evolution_matches_the_matrix_exponential():
-    """1/4 + (p_0 - 1/4) exp(-4t) agrees with exp(tR) at t = 0.1 and at t = 1."""
-    for t, expected in ((0.1, [0.1495, 0.2165, 0.2835, 0.3505]),
-                        (1.0, [0.2473, 0.2491, 0.2509, 0.2527])):
-        exact = ctmc.exact_evolution(ctmc.START, t)
-        assert [round(float(v), 4) for v in exact] == expected
-        through_matrix = ctmc.START @ ctmc.transition_matrix(ctmc.UNIT_RATES, t)
-        assert float((exact - through_matrix).abs().max()) == pytest.approx(0.0, abs=1e-12)
-
-
-def test_the_matrix_exponential_worked_for_dna():
-    """At tau = log 2 the uniform kernel is 0.625 on the diagonal and 0.125 elsewhere."""
-    pi = torch.full((4,), 0.25, dtype=torch.float64)
-    tau = math.log(2.0)
-    kernel = ctmc.transition_matrix(ctmc.r_pi(pi), tau)
-    assert float(kernel[0, 0]) == pytest.approx(0.625, abs=1e-12)
-    assert float(kernel[0, 1]) == pytest.approx(0.125, abs=1e-12)
-    assert float(kernel.sum(dim=1).min()) == pytest.approx(1.0)
-    closed = ctmc.pi_kernel(pi, math.exp(-tau))
-    assert float((kernel - closed).abs().max()) == pytest.approx(0.0, abs=1e-12)
-    euler = {n: float(ctmc.euler_product(ctmc.r_pi(pi), tau, n)[0, 0]) for n in (1, 2, 10)}
-    assert [round(euler[n], 3) for n in (1, 2, 10)] == [0.480, 0.570, 0.616]
-
-
-def test_masking_kernel_is_absorbing():
-    """At alpha = 1/2 the masking kernel keeps each base with 0.5 and sends 0.5 to the mask."""
-    mask_pi = torch.zeros(5, dtype=torch.float64)
-    mask_pi[4] = 1.0
-    kernel = ctmc.pi_kernel(mask_pi, 0.5)
-    assert float(kernel[0, 0]) == pytest.approx(0.5)
-    assert float(kernel[0, 4]) == pytest.approx(0.5)
-    assert float(kernel[4, 4]) == pytest.approx(1.0)
-
-
-def test_composition_eigenvalues_and_detailed_balance():
-    """Intervals compose, R_pi has eigenvalues 0 and -1, and the cycle breaks detailed balance."""
-    whole = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.7)
-    halves = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.3) @ ctmc.transition_matrix(
-        ctmc.UNIT_RATES, 0.4)
-    assert float((whole - halves).abs().max()) == pytest.approx(0.0, abs=1e-12)
-    pi = torch.full((4,), 0.25, dtype=torch.float64)
-    eigenvalues = sorted(round(float(v.real), 6) for v in torch.linalg.eigvals(ctmc.r_pi(pi)))
-    assert eigenvalues == [-1.0, -1.0, -1.0, 0.0]
-    assert float((pi @ ctmc.r_pi(pi)).abs().max()) == pytest.approx(0.0, abs=1e-12)
-    assert ctmc.detailed_balance_gap(ctmc.r_pi(pi), pi) == pytest.approx(0.0, abs=1e-12)
-    cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64)
-    assert float((cycle_pi @ ctmc.CYCLE).abs().max()) == pytest.approx(0.0, abs=1e-12)
-    assert ctmc.detailed_balance_gap(ctmc.CYCLE, cycle_pi) == pytest.approx(1.0 / 3.0)
-
-
-def test_time_reversal_retraces_the_marginals():
-    """With p_t = (0.5, 0.25, 0.125, 0.125) the reverse rate out of C toward A is 2."""
-    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64)
-    backward = ctmc.reverse_rates(ctmc.UNIT_RATES, marginal)
-    assert float(backward[1, 0]) == pytest.approx(2.0)
-    assert float(backward[1, 2]) == pytest.approx(0.5)
-    assert float(ctmc.exit_rates(backward)[1]) == pytest.approx(3.0)
-    net = marginal @ backward + marginal @ ctmc.UNIT_RATES
-    assert float(net.abs().max()) == pytest.approx(0.0, abs=1e-12)
-
-
-# %% 02. mdlm: the corruption, the loss and the sampler
-
-
-def test_reverse_step_worked():
-    """At alpha_s = 0.7 and alpha_t = 0.4 the reveal probability is 0.5."""
-    assert mdlm.reveal_probability(0.7, 0.4) == pytest.approx(0.5)
-    law = mdlm.reverse_step_distribution(0.7, 0.4, [0.6, 0.4, 0.0, 0.0])
-    assert [round(v, 4) for v in law] == [0.3, 0.2, 0.0, 0.0, 0.5]
-    assert sum(law) == pytest.approx(1.0)
-
-
-def test_one_term_of_the_loss_worked():
-    """w(0.5) = 2 and the masked cross-entropy of 1.04982 give a weighted term of 2.09964."""
-    assert mdlm.reveal_rate(0.5) == pytest.approx(2.0)
-    masked = -math.log(0.7) - math.log(0.5)
-    assert masked == pytest.approx(1.04982, abs=5e-6)
-    assert 2.0 * masked == pytest.approx(2.09964, abs=5e-6)
-
-
-def test_the_rate_error_collapses_to_the_cross_entropy():
-    """At a masked position the rate errors sum to w(t) times the negative log probability."""
-    rates = [2.0 * d for d in (0.2, 0.6, 0.1, 0.1)]
-    assert [round(v, 4) for v in rates] == [0.4, 1.2, 0.2, 0.2]
-    total = mdlm.rate_error(2.0, rates[1]) + sum(
-        mdlm.rate_error(0.0, r) for i, r in enumerate(rates) if i != 1)
-    assert total == pytest.approx(2.0 * -math.log(0.6), abs=1e-9)
-    assert total == pytest.approx(1.0217, abs=5e-5)
-    assert 2.0 * -math.log(0.4) == pytest.approx(1.8326, abs=5e-5)
-
-
-def test_corruption_and_sampling_shapes():
-    """The corruption keeps or masks, and the sampler returns clean DNA with no mask left."""
-    torch.manual_seed(0)
-    x0 = torch.randint(0, 4, (8, 6))
-    zt = mdlm.corrupt(x0, torch.full((8, 1), 0.5))
-    assert bool(((zt == x0) | (zt == mdlm.MASK)).all())
-    kept = mdlm.corrupt(x0, torch.zeros(8, 1)).eq(x0).to(torch.get_default_dtype())
-    assert float(kept.mean()) == 1.0
-    assert int(mdlm.corrupt(x0, torch.ones(8, 1)).eq(mdlm.MASK).all())
-    model = _tiny_denoiser(5, 6)
-    drawn = mdlm.sample(model, 4, 6, steps=4)
-    assert drawn.shape == (4, 6) and int(drawn.max()) < 4
-
-
-def _tiny_denoiser(vocab: int, length: int, num_classes: int = 0):
-    """A very small SequenceDenoiser, so the sampler tests cost nothing."""
-    from cis6270.nets import SequenceDenoiser
-
-    torch.manual_seed(0)
-    return SequenceDenoiser(vocab, length, width=16, depth=1, heads=2, num_classes=num_classes)
-
-
-# %% 03. uniform: the posterior average of 2, and the rate error that keeps its staying term
-
-
-def test_the_uniform_kernel_and_the_clean_factor():
-    """At alpha_t = 1/2 the kernel is 0.625 at the clean base and 0.125 elsewhere."""
-    assert float(uniform.kernel(0.5)[0, 0]) == pytest.approx(0.625)
-    assert float(uniform.kernel(0.5)[0, 1]) == pytest.approx(0.125)
-    assert uniform.clean_factor(0.5) == pytest.approx(5.0)
-
-
-def test_the_worked_posterior_average_is_two():
-    """The chapter's posterior average of the conditional reverse rates on one base comes out 2."""
-    prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64)
-    posterior = uniform.posterior_over_clean(0.5, 1, prior)
-    assert [round(float(v), 4) for v in posterior] == [0.375, 0.625, 0.0, 0.0]
-    assert float(uniform.conditional_rates(0.5, current=1, clean=0)[0]) == pytest.approx(5.0)
-    assert float(uniform.conditional_rates(0.5, current=1, clean=1)[0]) == pytest.approx(0.2)
-    assert uniform.marginal_reverse_rate(0.5, 1, 0, prior) == pytest.approx(2.0)
-    assert uniform.marginal_reverse_rate(0.5, 1, 2, prior) == pytest.approx(0.5)
-    marginal = prior @ uniform.kernel(0.5)
-    assert [round(float(v), 4) for v in marginal] == [0.5, 0.25, 0.125, 0.125]
-    assert float(marginal[0] / marginal[1]) == pytest.approx(2.0)
-
-
-def test_the_rate_error_on_one_base():
-    """A denoiser of (0.5, 0.3, 0.1, 0.1) gives rates 2.76 and 1.16 and a rate error of 0.754."""
-    denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64)
-    current = torch.tensor([[1]])
-    alpha = torch.tensor([[0.5]], dtype=torch.float64)
-    rates = uniform.model_rates(denoiser, uniform.ratio_table(current, alpha))[0, 0]
-    rates = rates / (uniform.BETA / 4)  # Quote the rates in the chapter's unit forward rate.
-    assert [round(float(v), 4) for v in rates] == [2.76, 1.0, 1.16, 1.16]
-    target = uniform.conditional_rates(0.5, current=1, clean=0)
-    error = float(uniform.rate_error(target[[0, 2, 3]], rates[[0, 2, 3]]).sum())
-    assert error == pytest.approx(0.754, abs=5e-4)
-    assert float(uniform.rate_error(torch.tensor(5.0), torch.tensor(2.76))) == pytest.approx(
-        0.731, abs=5e-4)
-
-
-def test_the_rate_error_is_minimized_at_the_posterior_average():
-    """0.375 l(5, B) + 0.625 l(0.2, B) is smallest at B = 2, where it equals 1.430."""
-    def mixture(b: float) -> float:
-        return float(0.375 * uniform.rate_error(torch.tensor(5.0), torch.tensor(b))
-                     + 0.625 * uniform.rate_error(torch.tensor(0.2), torch.tensor(b)))
-
-    assert mixture(2.0) == pytest.approx(1.430, abs=5e-4)
-    assert mixture(5.0) == pytest.approx(2.598, abs=5e-4)
-    assert mixture(0.2) == pytest.approx(4.235, abs=5e-4)
-    assert mixture(2.0) < min(mixture(1.9), mixture(2.1))
-
-
-# %% 04. d3pm: the finite step
-
-
-def test_cumulative_transition_matrix():
-    """With beta_1 = 0.2 and beta_2 = 0.3 the product is 0.67 on the diagonal and 0.11 off it."""
-    product = d3pm.one_step(0.2) @ d3pm.one_step(0.3)
-    assert float(product[0, 0]) == pytest.approx(0.67, abs=5e-3)
-    assert float(product[0, 1]) == pytest.approx(0.11, abs=5e-3)
-    assert float((product - d3pm.cumulative(0.56)).abs().max()) == pytest.approx(0.0, abs=1e-6)
-    assert (1.0 - 1.0 / 1000) ** 1000 == pytest.approx(0.3677, abs=5e-5)
-
-
-def test_reversing_one_finite_step():
-    """Observing C under a clean A gives (0.5795, 0.3523, 0.0341, 0.0341) at the earlier step."""
-    first, second = d3pm.one_step(0.2), d3pm.one_step(0.3)
-    clean_a = torch.tensor([[[1.0, 0.0, 0.0, 0.0]]])
-    posterior = d3pm.exact_posterior(torch.tensor([[1]]), clean_a, second[None], first[None])[0, 0]
-    assert [round(float(v), 4) for v in posterior] == [0.5795, 0.3523, 0.0341, 0.0341]
-
-
-def test_rate_matrices_chosen_for_the_data():
-    """The Gaussian row is 0.4088 to a neighbour, and exp(0.3 R) has first row (0.6747, ...)."""
-    row = d3pm.gaussian_row()
-    assert float(row[1]) == pytest.approx(0.4088, abs=5e-5)
-    assert float(row[0]) == pytest.approx(0.0912, abs=5e-5)
-    adjacency = torch.tensor([[0.0, 1.0, 0.5, 0.0], [1.0, 0.0, 1.0, 0.0],
-                              [0.5, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]])
-    R = d3pm.graph_rates(adjacency)
-    assert [round(float(v), 4) for v in R[0]] == [-1.5, 1.0, 0.5, 0.0]
-    kernel = torch.matrix_exp(0.3 * R)
-    assert [round(float(v), 4) for v in kernel[0]] == [0.6747, 0.1964, 0.1122, 0.0167]
-
-
-# %% 05. sedd: score entropy
-
-
-def test_score_entropy_at_three_predictions():
-    """With the target ratio 2, a prediction of 1.8 costs 0.0107 nats and 1.0 costs 0.3863."""
-    target = torch.tensor(2.0)
-    assert float(sedd.score_entropy(target, torch.tensor(1.8))) == pytest.approx(0.0107, abs=5e-5)
-    assert float(sedd.score_entropy(target, torch.tensor(1.0))) == pytest.approx(0.3863, abs=5e-5)
-    assert float(sedd.score_entropy(target, torch.tensor(3.0))) == pytest.approx(0.1891, abs=5e-5)
-    assert float(sedd.score_entropy(target, target)) == pytest.approx(0.0, abs=1e-6)
-
-
-def test_the_target_ratio_from_the_denoiser_route():
-    """0.375(5) + 0.625(0.2) is 2, and the kernel ratio reproduces it at a clean A."""
-    assert 0.375 * 5.0 + 0.625 * 0.2 == pytest.approx(2.0)
-    x0 = torch.tensor([[0]])  # A clean A.
-    z = torch.tensor([[1]])  # Observed as a C.
-    ratio = sedd.kernel_ratio(x0, z, torch.tensor([[0.5]]))
-    assert float(ratio[0, 0, 0]) == pytest.approx(5.0)
-    assert float(ratio[0, 0, 2]) == pytest.approx(1.0)
-
-
-# %% 06. block: the dial between diffusion and autoregression
-
-
-def test_the_worked_block_bound():
-    """ACGTAC in blocks of two at t = 0.5 gives 1.8326, 0.7134 and 0.2107, summing to 2.7567."""
-    weight = 2.0
-    contributions = [weight * (-math.log(0.5) - math.log(0.8)),
-                     weight * -math.log(0.7), weight * -math.log(0.9)]
-    assert [round(v, 4) for v in contributions] == [1.8326, 0.7133, 0.2107]
-    assert sum(contributions) == pytest.approx(2.7567, abs=5e-5)
-
-
-def test_the_one_position_identity():
-    """At blocks of one position the time integral of w(t)(1 - alpha_t) is exactly one."""
-    assert block.one_position_integral() == pytest.approx(1.0, abs=1e-9)
-    generator = torch.Generator().manual_seed(0)
-    estimate = block.one_position_estimate(0.7, 40000, generator)
-    assert estimate == pytest.approx(-math.log(0.7), abs=0.05)
-
-
-def test_block_input_keeps_the_prefix_and_hides_the_future():
-    """The prefix stays clean, the current block is corrupted, and the suffix is all mask."""
-    x0 = torch.arange(8).remainder(4)[None, :].repeat(2, 1)  # [2, 8].
-    z = block.block_input(x0, torch.tensor([1, 1]), 2, torch.full((2, 1), 1.0))
-    assert bool((z[:, :2] == x0[:, :2]).all())  # The prefix is clean.
-    assert bool((z[:, 2:4] == block.MASK).all())  # The block is fully masked at t = 1.
-    assert bool((z[:, 4:] == block.MASK).all())  # The future is hidden.
-
-
-# %% 07. guidance: the h-transform and the worked guided jump
-
-
-def test_h_transform_keeps_a_rate_matrix():
-    """Reweighting by h(y)/h(x) leaves the row sums at zero, and a constant h changes nothing."""
-    R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0))
-    transformed = guidance.h_transform(R, torch.tensor([2.0, 1.0, 0.5, 0.5]))
-    assert float(transformed.sum(dim=1).abs().max()) == pytest.approx(0.0, abs=1e-6)
-    constant = guidance.h_transform(R, torch.full((4,), 3.0))
-    assert float((constant - R).abs().max()) == pytest.approx(0.0, abs=1e-6)
-
-
-def test_the_reverse_process_as_an_h_transform():
-    """h_t = p_t/pi is (2, 1, 0.5, 0.5), and the rates out of C come back as 2, 0.5 and 0.5."""
-    pi = torch.full((4,), 0.25)
-    values = guidance.value_from_marginal(torch.tensor([0.5, 0.25, 0.125, 0.125]), pi)
-    assert [round(float(v), 4) for v in values] == [2.0, 1.0, 0.5, 0.5]
-    R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0))
-    reverse = guidance.h_transform(R, values)
-    assert float(reverse[1, 0]) == pytest.approx(2.0, abs=1e-5)
-    assert float(reverse[1, 2]) == pytest.approx(0.5, abs=1e-5)
-    clean_a = guidance.value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi)
-    assert float(clean_a[0] / clean_a[1]) == pytest.approx(5.0)
-    assert float(clean_a[1] / clean_a[0]) == pytest.approx(0.2)
-
-
-def test_the_worked_guided_jump():
-    """The classifier turns the rates 2, 0.5, 0.5 into 1, 1.5, 0.5, and gamma = 2 sharpens them."""
-    unguided = torch.tensor([2.0, 0.5, 0.5])
-    ratio = torch.tensor([0.1, 0.6, 0.2]) / 0.2
-    guided = unguided * ratio
-    assert [round(float(v), 4) for v in guided] == [1.0, 1.5, 0.5]
-    assert float(guided.sum()) == pytest.approx(3.0)
-    shares = guided / guided.sum()
-    assert [float(v) for v in shares] == pytest.approx([1 / 3, 1 / 2, 1 / 6], abs=1e-6)
-    sharper = unguided * ratio ** 2
-    assert [round(float(v), 4) for v in sharper] == [0.5, 4.5, 0.5]
-    assert float(sharper.sum()) == pytest.approx(5.5)
-    assert float((sharper / sharper.sum())[1]) == pytest.approx(0.8182, abs=5e-5)
-    assert 1.0 / float(sharper.sum()) == pytest.approx(0.1818, abs=5e-5)
-
-
-def test_exact_and_gradient_ratios_agree_at_first_order():
-    """Both estimates of p(c | z^{i <- a})/p(c | z) have the right shape and a positive value."""
-    torch.manual_seed(0)
-    classifier = guidance.NoisyClassifier(6, width=16)
-    z = torch.tensor([[0, guidance.MASK, 2, guidance.MASK, 1, 3]])
-    t = torch.tensor([0.5])
-    with torch.no_grad():
-        exact = guidance.exact_ratio(classifier, z, t)
-    approximate = guidance.gradient_ratio(classifier, z, t)
-    assert exact.shape == approximate.shape == (1, 6, 4)
-    assert bool((exact > 0).all()) and bool((approximate > 0).all())
-    # The ratio at the token already in place is one under both estimates.
-    assert float(exact[0, 0, 0]) == pytest.approx(1.0, abs=1e-5)
-    assert float(approximate[0, 0, 0]) == pytest.approx(1.0, abs=1e-5)
-
-
-# %% 08. planner: the order of the reveals and the weights it puts in the loss
-
-
-def test_the_worked_planner_selection():
-    """Scores (0.42, 0.91, 0.55, 0.78) at the masked positions reveal 3 and 6, and remask 5."""
-    scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]])
-    masked = torch.tensor([[True, False, True, True, False, True]])
-    revealed = sorted(int(i) + 1 for i in planner.topk_where(scores, masked, 2).nonzero()[:, 1])
-    returned = sorted(int(i) + 1 for i in planner.topk_where(-scores, ~masked, 1).nonzero()[:, 1])
-    assert revealed == [3, 6]
-    assert returned == [5]
-    assert int(planner.topk_where(scores, masked, 0).sum()) == 0
-
-
-def test_the_corrector_fluxes_match():
-    """alpha_t beta(t) and (1 - alpha_t) w(t) both equal -alpha_dot, which is one here."""
-    for t in (0.5, 0.25, 0.75):
-        forward, reverse = planner.corrector_fluxes(t)
-        assert forward == pytest.approx(1.0)
-        assert reverse == pytest.approx(1.0)
-
-
-def test_the_worked_papl_weights():
-    """A planner distribution of (0.10, 0.50, 0.15, 0.25) turns 2.9 nats into 2.12."""
-    rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]])
-    cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]])
-    lam = 4.0 * rho
-    assert [round(float(v), 4) for v in lam[0]] == [0.4, 2.0, 0.6, 1.0]
-    assert float(lam.sum()) == pytest.approx(4.0)
-    assert float(cross_entropies.sum()) == pytest.approx(2.9)
-    assert float((lam * cross_entropies).sum()) == pytest.approx(2.12)
-
-
-def test_planner_weights_average_to_one_and_reduce_to_mdlm():
-    """lambda_i averages to one over the masked positions, and strength zero gives every weight one."""
-    score = torch.tensor([[0.9, 0.1, 0.5, 0.2]])
-    masked = torch.tensor([[True, True, False, True]])
-    weights = planner.planner_weights(score, masked, 1.0)
-    assert float((weights * masked).sum()) == pytest.approx(float(masked.sum()), abs=1e-5)
-    plain = planner.planner_weights(score, masked, 0.0)
-    assert float((plain - 1.0).abs().max()) == pytest.approx(0.0, abs=1e-6)
-
-
-# %% 09. simplex: the Dirichlet belief and the thinning representation
-
-
-def test_the_worked_dirichlet_parameters():
-    """At alpha_t = 1/2 the parameters are (5, 1, 1, 1) and at alpha_s = 0.8 they are (17, 1, 1, 1)."""
-    assert simplex.concentration(0.5) == pytest.approx(8.0)
-    assert simplex.concentration(0.8) == pytest.approx(20.0)
-    mean_t = simplex.belief_mean(0.5, 0)
-    assert [round(float(v), 4) for v in mean_t] == [0.625, 0.125, 0.125, 0.125]
-    assert [round(float(v), 4) for v in simplex.belief_mean(0.8, 0)] == [0.85, 0.05, 0.05, 0.05]
-    at_t = simplex.dirichlet_parameters(0.5, 0)
-    at_s = simplex.dirichlet_parameters(0.8, 0)
-    assert [round(float(v), 4) for v in at_t] == [5.0, 1.0, 1.0, 1.0]
-    assert [round(float(v), 4) for v in at_s] == [17.0, 1.0, 1.0, 1.0]
-    total = at_t.sum()
-    deviation = (at_t * (total - at_t) / (total ** 2 * (total + 1))).sqrt()
-    assert float(deviation[0]) == pytest.approx(0.161, abs=5e-4)
-
-
-def test_the_increment_sits_on_the_clean_coordinate():
-    """The increment is (12, 0, 0, 0), so the innovation is the one-hot vector of the clean base."""
-    step = simplex.increment(0.5, 0.8, 0)
-    assert [round(float(v), 4) for v in step] == [12.0, 0.0, 0.0, 0.0]
-    first, second = simplex.thinning_law(0.5, 0.8)
-    assert (round(first, 4), round(second, 4)) == (8.0, 12.0)
-    weight = first / (first + second)
-    assert weight == pytest.approx(0.4)
-    mixed = weight * simplex.belief_mean(0.5, 0) + (1.0 - weight) * torch.eye(4)[0]
-    assert [round(float(v), 4) for v in mixed] == [0.85, 0.05, 0.05, 0.05]
-
-
-def test_the_thinning_representation_reproduces_the_dirichlet():
-    """W V + (1 - W) onehot(A) has the mean and the spread of Dir(17, 1, 1, 1)."""
-    torch.manual_seed(0)
-    at_t, at_s = simplex.dirichlet_parameters(0.5, 0), simplex.dirichlet_parameters(0.8, 0)
-    first, second = simplex.thinning_law(0.5, 0.8)
-    draws = 200_000
-    thinned = torch.distributions.Dirichlet(at_t).sample((draws,))
-    weight = torch.distributions.Beta(torch.tensor(first), torch.tensor(second)).sample((draws,))
-    drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(4)[0]
-    total = at_s.sum()
-    exact = (at_s * (total - at_s) / (total ** 2 * (total + 1))).sqrt()
-    assert [round(float(v), 4) for v in exact] == [0.0779, 0.0476, 0.0476, 0.0476]
-    assert float(drawn.mean(0)[0]) == pytest.approx(0.85, abs=2e-3)
-    assert float(drawn.std(0)[0]) == pytest.approx(float(exact[0]), abs=2e-3)
-
-
-def test_the_parameter_field_and_the_reverse_step():
-    """The batched parameters match the per-position ones, and one step stays on the simplex."""
-    x0 = torch.tensor([[0, 1], [2, 3]])
-    field = simplex.parameter_field(x0, torch.full((2, 1), 0.5))
-    assert [round(float(v), 4) for v in field[0, 0]] == [5.0, 1.0, 1.0, 1.0]
-    assert [round(float(v), 4) for v in field[1, 1]] == [1.0, 1.0, 1.0, 5.0]
-    torch.manual_seed(0)
-    belief = simplex.corrupt(x0, torch.full((2, 1), 0.5))
-    assert belief.shape == (2, 2, 4)
-    assert float(belief.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6)
-    moved = simplex.reverse_step(belief, x0, 0.5, 0.8)
-    assert float(moved.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6)
-    assert bool((moved >= 0).all())
-
-
-def test_a_one_hot_belief_recovers_the_token_denoiser():
-    """The only architectural change is at the input, so one-hot beliefs give identical logits."""
-    torch.manual_seed(0)
-    model = simplex.BeliefDenoiser(6, width=16)
-    tokens = torch.tensor([[0, 1, 2, 3, 0, 1]])
-    t = torch.tensor([0.4])
-    one_hot = torch.nn.functional.one_hot(tokens, 4).to(torch.get_default_dtype())
-    with torch.no_grad():
-        through_belief = model(one_hot, t)
-        through_tokens = model.inner(tokens, t)
-    assert float((through_belief - through_tokens).abs().max()) == pytest.approx(0.0, abs=1e-6)
-    drawn = simplex.sample(model, 3, 6, steps=4)
-    assert drawn.shape == (3, 6) and int(drawn.max()) < 4
-
-
-# %% 10. search: rollouts, the archive and the two objectives
-
-
-def test_the_worked_value_estimate():
-    """Rollouts averaging 0.35 and 0.75 reweight a rate of 1.2 to 2.571."""
-    parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean())
-    child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean())
-    assert parent == pytest.approx(0.35)
-    assert child == pytest.approx(0.75)
-    assert child / parent == pytest.approx(2.143, abs=5e-4)
-    assert 1.2 * child / parent == pytest.approx(2.571, abs=5e-4)
-    assert 0.2 / parent == pytest.approx(0.571, abs=5e-4)
-
-
-def test_the_worked_pareto_archive():
-    """P3 at (0.4, 0.2) is dominated by P1 at (0.8, 0.3), and P1 and P2 both stay."""
-    archive = search.ParetoArchive()
-    for name, scores in (((1,), (0.8, 0.3)), ((2,), (0.5, 0.7)), ((3,), (0.4, 0.2))):
-        archive.add(name, scores)
-    assert sorted(q for q, _ in archive.entries) == [(1,), (2,)]
-    archive.add((4,), (0.9, 0.9))  # A candidate that dominates both of them.
-    assert [q for q, _ in archive.entries] == [(4,)]
-
-
-def test_the_two_objectives():
-    """A GC-free sequence scores zero on the first axis, and the motif scores one on the second."""
-    tokens = torch.tensor([[0, 1, 2, 3, 0, 0], [0, 0, 3, 3, 0, 3]])
-    scores = search.objectives(tokens)
-    assert float(scores[0, 0]) == pytest.approx(2 / 6)
-    assert float(scores[0, 1]) == pytest.approx(1.0)  # ACGT sits at the start.
-    assert float(scores[1, 0]) == pytest.approx(0.0)
-    assert float(scores[1, 1]) < 1.0

tests/test_lecture_5.py

@@ -1,550 +0,0 @@
-"""The worked numbers of Chapter 5, asserted against the files of ``lecture_5/``.
-
-Every number checked here is printed in the notes. Where the chapter runs a simulation, the test
-asserts the analytic value the simulation estimates, with a tolerance wide enough for the sample
-size the file uses by default.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-from cis6270.nets import SequenceDenoiser
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-def load(name):
-    """Import one lecture file by path, so that ``lecture_5`` needs no package marker."""
-    spec = importlib.util.spec_from_file_location(f"lecture5_{name}", ROOT / "lecture_5" / f"{name}.py")
-    module = importlib.util.module_from_spec(spec)
-    sys.path.insert(0, str(ROOT))
-    spec.loader.exec_module(module)
-    return module
-
-
-dfm = load("dfm")
-velocity = load("velocity")
-corrector = load("corrector")
-kinetic = load("kinetic_optimal")
-edits = load("edit_flows")
-dirichlet = load("dirichlet")
-fisher = load("fisher")
-gumbel = load("gumbel")
-guidance = load("guidance")
-objectives = load("multi_objective")
-rectify = load("rectify")
-
-
-def close(values, expected, tolerance=5e-4):
-    """Compare a list of numbers against the chapter's, entry by entry."""
-    assert len(values) == len(expected)
-    for got, want in zip(values, expected):
-        assert abs(float(got) - want) <= tolerance, (values, expected)
-
-
-# %% Section 5.2 and 5.3: a velocity, its two rules, and the master equation
-
-
-def test_worked_master_equation():
-    u, p, p_dot, arriving, leaving = velocity.worked_master_equation()
-    assert velocity.obeys_the_two_rules(u)
-    assert torch.allclose(u, velocity.set_diagonal(u))
-    close(p_dot.tolist(), [-0.1, -0.2, 0.7, -0.4])
-    assert abs(float(p_dot.sum())) < 1e-9
-    assert abs(arriving - 0.7) < 1e-9 and abs(leaving) < 1e-9
-    close((p + 0.1 * p_dot).tolist(), [0.09, 0.18, 0.37, 0.36])
-
-
-def test_hazard_and_survival():
-    assert abs(velocity.hazard(0.5) - 2.0) < 1e-9
-    assert abs(velocity.hazard(0.9) - 10.0) < 1e-9
-    assert abs(velocity.survival(4000) - 0.1) < 2e-3  # S(t) = 1 - kappa_t.
-
-
-def test_continuous_and_discrete_side_by_side():
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
-    moved, drawn, lam = velocity.side_by_side(2, p1, t=0.5, h=0.1)
-    close(moved.tolist(), [0.02, 0.12, 0.84, 0.02])
-    close(drawn.tolist(), [0.02, 0.12, 0.84, 0.02])
-    assert abs(lam - 0.2) < 1e-9
-
-
-# %% Sections 5.5 to 5.11: the mixture path, the velocity, the loss and the sampler
-
-
-def test_dfm_worked_numbers():
-    report = dfm.worked_numbers()
-    assert report["mixture_draw"] == "ATGA"
-    close(report["velocity_at_G"], [0.2, 1.2, -1.6, 0.2])
-    close(report["euler_step_at_G"], [0.02, 0.12, 0.84, 0.02])
-    assert abs(report["redraw_probability"] - 0.2) < 1e-9
-    assert abs(report["example_loss_nats"] - 4.780) < 1e-3
-    close(report["mask_posterior_C_G"], [0.625, 0.375])
-    assert report["mask_posterior_unchanged"]  # Proposition 6 of Discrete Flow Matching.
-    assert abs(report["posterior_entropy_nats"] - 0.6616) < 5e-4
-    assert abs(report["kl_gap_to_uniform"] - 0.0316) < 5e-4
-
-
-def test_conditional_velocity_generates_the_conditional_path():
-    """Section 5.6.2: the two terms of the master equation on the mixture path."""
-    kappa, x0, x1 = 0.4, 0, 1
-    path = torch.zeros(4, dtype=torch.double)
-    path[x0], path[x1] = 1.0 - kappa, kappa
-    total = sum(path[z] * dfm.conditional_velocity(z, x1, kappa) for z in (x0, x1))
-    derivative = torch.zeros(4, dtype=torch.double)
-    derivative[x1], derivative[x0] = 1.0, -1.0  # kappa_dot_t times the two indicators.
-    assert torch.allclose(total, derivative, atol=1e-9)
-
-
-def test_corrupt_reaches_both_endpoints():
-    torch.manual_seed(0)
-    process = dfm.DiscreteFlowMatching(4)
-    x1 = torch.randint(0, 4, (64, 16))
-    assert torch.equal(process.corrupt(x1, torch.ones(64)), x1)  # kappa_t = 1 gives the data.
-    source = process.corrupt(x1, torch.zeros(64))  # kappa_t = 0 gives the source.
-    assert float((source == x1).double().mean()) < 0.4  # A uniform draw agrees about a quarter of the time.
-    halfway = process.corrupt(x1, torch.full((64,), 0.5))
-    assert 0.55 < float((halfway == x1).double().mean()) < 0.75  # Heads with probability one half.
-
-
-def test_two_sources():
-    assert torch.equal(dfm.DiscreteFlowMatching(4, mask=4).source((3, 5)),
-                       torch.full((3, 5), 4, dtype=torch.long))  # The all-mask source.
-    uniform = dfm.DiscreteFlowMatching(4).source((200, 20))
-    assert int(uniform.min()) == 0 and int(uniform.max()) == 3  # Every base appears.
-
-
-# %% Sections 5.13 and 5.14: zero-flux velocities and the corrector
-
-
-def test_zero_flux_velocity_on_three_states():
-    p, u, v, with_u, with_both = corrector.three_state_example()
-    assert torch.allclose(corrector.net_flux(v, p), torch.zeros(3), atol=1e-12)
-    close(with_u.tolist(), [-0.2, -0.3, 0.5])
-    close(with_both.tolist(), [-0.2, -0.3, 0.5])
-    assert abs(float(-u[0, 0]) - 1.0) < 1e-12 and abs(float(-(u + v)[0, 0]) - 5.0) < 1e-12
-
-
-def test_corrector_step_table():
-    """Section 5.14.3 at t = 0.4 with h = 0.01 and the denoiser (0.1, 0.6, 0.2, 0.1)."""
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
-    expected = {
-        0.0: ([0.00167, 0.01000, 0.00333, 0.00167], 0.98333),
-        2.0: ([0.00300, 0.01800, 0.00600, 0.00300], 0.97000),
-        10.0: ([0.00833, 0.05000, 0.01667, 0.00833], 0.91667),
-    }
-    for eta, (probabilities, stay) in expected.items():
-        bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, eta)
-        close((0.01 * bar_u[:4]).tolist(), probabilities, tolerance=1e-5)
-        assert abs(float(1.0 + 0.01 * bar_u[corrector.MASK]) - stay) < 1e-5
-        assert abs(corrector.masked_fraction_derivative(0.4, eta) + 1.0) < 1e-12
-
-
-def test_corrector_rates_and_admissible_steps():
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])
-    bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, 2.0)
-    assert abs(float(-bar_u[corrector.MASK]) - 3.0) < 1e-12  # (1 + 2 x 0.4)/(1 - 0.4).
-    visible = torch.zeros(4)
-    visible[1] = 1.0
-    assert abs(float(corrector.corrector_velocity(visible, 1, 0.4, 2.0)[corrector.MASK]) - 2.0) < 1e-12
-    close([corrector.admissible_step(0.9, eta) for eta in (0.0, 2.0, 10.0)],
-          [0.1, 0.0357, 0.01], tolerance=1e-4)
-
-
-def test_jump_counts_are_one_plus_eta():
-    """Section 5.14.3: the masked fraction stays at 1 - t and the jumps come to 1 + eta."""
-    p1 = torch.tensor([0.25, 0.25, 0.25, 0.25])
-    for eta in (0.0, 2.0, 10.0):
-        half, mean = corrector.simulate_jumps(p1, eta, positions=10_000, steps=300)
-        assert abs(half - 0.5) < 0.03
-        assert abs(mean - (1.0 + eta)) < 0.25
-
-
-def test_parallel_step_error():
-    assert abs(corrector.joint_reveal_rate(64) - 1.0 / 64) < 1e-12
-    assert abs(corrector.pair_survival(2.0) - 0.1353) < 5e-4
-    assert abs(corrector.pair_survival(10.0) - 4.54e-5) < 1e-7
-    measured = corrector.simulate_pair(0.0, samples=40_000, steps=64)
-    assert abs(measured - 0.5 / 64) < 2e-3  # Half of the joint reveals land outside the data.
-
-
-# %% Section 5.16: kinetic-optimal paths
-
-
-def test_kinetic_optimal_recovers_the_mixture_velocity():
-    q = torch.tensor([0.1, 0.6, 0.2, 0.1])
-    p, p_dot = kinetic.mask_path(q, kappa_t=0.5)
-    u = kinetic.kinetic_optimal(p, p_dot)
-    assert abs(float(p_dot[kinetic.MASK] / p[kinetic.MASK]) + 2.0) < 1e-12
-    assert abs(float(p_dot[1] / p[1]) - 2.0) < 1e-12
-    assert abs(float(u[1, kinetic.MASK]) - 1.2) < 1e-12  # 0.30 times 4.
-    assert abs(float(u[1, kinetic.MASK]) - kinetic.mixture_rate(0.5) * 0.6) < 1e-12
-    assert abs(float(u[1, 2])) < 1e-12  # No probability moves between two visible tokens.
-    assert torch.allclose(u @ p, p_dot, atol=1e-12)
-
-
-def test_worked_metric_path_for_dna():
-    beta_t, beta_dot = kinetic.beta_schedule(0.5)
-    assert abs(beta_t - 1.0) < 1e-12 and abs(beta_dot - 4.0) < 1e-12
-    p = kinetic.metric_path(2, beta_t)
-    close(p.tolist(), [0.2245, 0.0826, 0.6103, 0.0826])
-    u = kinetic.metric_velocity(2, beta_t, beta_dot)
-    assert abs(float(u[0, 1]) - 0.898) < 5e-4  # C to A.
-    assert abs(float(u[2, 1]) - 4.882) < 5e-4  # C to G.
-    assert abs(float(u[2, 0]) - 2.441) < 5e-4  # A to G.
-    assert abs(float(-u[1, 1]) - 5.780) < 1e-3  # The total rate out of C.
-    close((torch.eye(4)[1] + 0.05 * u[:, 1]).tolist(), [0.0449, 0.7110, 0.2441, 0.0])
-    close((u @ p).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774])
-    close(kinetic.path_derivative(2, 0.5).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774], 1e-3)
-
-
-def test_energy_of_two_schedules():
-    assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-6) - 27.63) < 5e-3
-    assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-9) - 41.45) < 5e-3
-    assert abs(kinetic.schedule_energy(*kinetic.GEODESIC, 600, 1e-6) - math.pi ** 2) < 1e-3
-    assert abs(kinetic.geodesic_schedule(0.25) - 0.1464) < 5e-4
-    assert abs(kinetic.geodesic_schedule(0.75) - 0.8536) < 5e-4
-    assert abs(kinetic.geodesic_schedule(0.5, 0.25) - 0.6667) < 5e-4
-
-
-# %% Section 5.17: Edit Flows
-
-
-def test_alignment_and_column_wise_path():
-    z0 = [0, 1, 2, 3, edits.BLANK]  # (A, C, G, T, eps).
-    z1 = [0, edits.BLANK, 2, 0, 1]  # (A, eps, G, A, C).
-    assert edits.column_kinds(z0, z1) == ["none", "del", "none", "sub", "ins"]
-    drawn = edits.draw_state(z0, z1, [True, False, True, True, True])
-    assert drawn == [0, 1, 2, 0, 1]  # ACGAC, the draw of probability 0.5 ** 3 = 0.125.
-    assert abs(edits.expected_length([0, 1, 2, 3], [0, 2, 0, 1], 0.5) - 4.0) < 1e-12
-    recovered = edits.align([0, 1, 2, 3], [0, 2, 0, 1])
-    assert len(recovered[0]) == len(recovered[1])  # Any alignment pads to a common width.
-
-
-def test_alignment_posterior_and_averaged_rates():
-    first = ([0, 1, edits.BLANK], [0, 0, 1])  # Substitute C with A, then insert C.
-    second = ([0, edits.BLANK, 1], [0, 0, 1])  # Insert A, and leave the last column alone.
-    for kappa_t, expected in ((0.25, 0.4286), (0.5, 1 / 3), (0.75, 0.2)):
-        weights = edits.alignment_posterior([first, second], [0, 1], kappa_t)
-        assert abs(weights[0] - expected) < 5e-4
-    posterior, rates = edits.averaged_rates([first, second], [0, 1], 0.5)
-    assert abs(rates["sub(2,A)"] - 0.667) < 1e-3
-    assert abs(rates["ins(1,A)"] - 1.333) < 1e-3
-    assert abs(rates["ins(2,C)"] - 0.667) < 1e-3
-    assert abs(sum(rates.values()) - 2.667) < 1e-3
-
-
-def test_edit_loss_and_sampler_step():
-    assert abs(edits.edit_loss(3.0, [0.8], kappa_t=0.5) - 3.446) < 1e-3
-    grid = [0.01 * i for i in range(1, 801)]
-    assert abs(edits.loss_minimizer(2.0, grid) - 2.0) < 1e-6
-    intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]}
-    insert, change = edits.step_probabilities(intensities, 0.05)
-    close(insert, [0.010, 0.015, 0.020])
-    close(change, [0.025, 0.060, 0.035])
-    exact, first_order = edits.untouched_probability(intensities, 0.05)
-    assert abs(exact - 0.8452) < 5e-4 and abs(first_order - 0.835) < 5e-4
-    assert abs(exact - first_order - 0.0102) < 5e-4
-    assert abs(edits.expected_length_change(intensities, 0.05) + 0.015) < 1e-9
-    assert edits.edit_rates(intensities)[0] == 11  # Three intensities and two token distributions.
-
-
-# %% Section 5.20: Dirichlet Flow Matching
-
-
-def test_dirichlet_path_and_speeds():
-    assert abs(dirichlet.dirichlet_density(torch.full((4,), 0.25), 0, 0.0) - 6.0) < 1e-9
-    assert dirichlet.straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3) == 3
-    assert dirichlet.straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6) == 1
-    closed = dirichlet.speed_two_letters(0.3, 1.0)
-    assert abs(closed - 0.2580) < 5e-5
-    assert abs(float(dirichlet.speed(0.3, 1.0, K=2)) - closed) < 1e-6
-    x = torch.tensor([0.40, 0.30, 0.20, 0.10])
-    close(dirichlet.speed_vector(x, 1.0).tolist(), [0.232, 0.199, 0.159, 0.104], 1e-3)
-
-
-def test_dirichlet_velocity_and_posterior_average():
-    x = torch.tensor([0.40, 0.30, 0.20, 0.10])
-    toward_a = dirichlet.conditional_velocity(x, 0, 1.0)
-    close(toward_a.tolist(), [0.139, -0.070, -0.046, -0.023], 1e-3)
-    assert abs(float(toward_a.sum())) < 1e-12
-    p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05])
-    v = dirichlet.marginal_velocity(x, p_hat, 1.0)
-    close(v.tolist(), [0.0553, -0.0133, -0.0262, -0.0158], 2e-4)
-    close((x + 0.1 * v).tolist(), [0.4055, 0.2987, 0.1974, 0.0984])
-    assert abs(float((x + 0.1 * v).sum()) - 1.0) < 1e-12
-
-
-def test_dirichlet_continuity():
-    time_term, space_term = dirichlet.continuity_terms(torch.tensor([0.5, 0.3, 0.2]), 0, 0.8)
-    assert abs(time_term - 0.6355) < 5e-4
-    assert abs(space_term + 0.6355) < 5e-4
-    assert abs(time_term + space_term) < 1e-5
-
-
-# %% Section 5.21: Fisher-Flow
-
-
-def test_two_moves_under_two_metrics():
-    near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6])
-    rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999])
-    assert abs(fisher.euclidean_distance(*near) - 0.1414) < 5e-4
-    assert abs(fisher.euclidean_distance(*rare) - 0.0127) < 5e-4
-    assert abs(fisher.fisher_rao_distance(*near) - 0.201) < 1e-3
-    assert abs(fisher.fisher_rao_distance(*rare) - 0.137) < 1e-3
-    orthogonal = fisher.fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0]))
-    assert abs(orthogonal - math.pi) < 1e-6  # Any two letters are the farthest apart.
-
-
-def test_sphere_and_geodesic():
-    p0 = torch.full((4,), 0.25)
-    p1 = fisher.smooth(torch.eye(4)[0])
-    close(p1.tolist(), [0.91, 0.03, 0.03, 0.03])
-    x0, x1 = fisher.to_sphere(p0), fisher.to_sphere(p1)
-    close(x1.tolist(), [0.954, 0.173, 0.173, 0.173], 1e-3)
-    assert abs(float((x0 * x1).sum()) - 0.737) < 1e-3
-    assert abs(math.acos(float((x0 * x1).sum())) - 0.743) < 1e-3
-    assert abs(fisher.fisher_rao_distance(p0, p1) - 1.485) < 1e-3
-    mid = fisher.geodesic(x0, x1, torch.tensor(0.5))
-    close(mid.tolist(), [0.780, 0.361, 0.361, 0.361], 1e-3)
-    close(fisher.to_simplex(mid).tolist(), [0.609, 0.130, 0.130, 0.130], 1e-3)
-    remaining = fisher.log_map(mid, x1)
-    close(remaining.tolist(), [0.2323, -0.1672, -0.1672, -0.1672], 1e-3)
-    assert abs(float(remaining.norm()) - 0.3713) < 5e-4
-    assert abs(float(fisher.conditional_velocity(mid, x1, 0.5).norm()) - 0.743) < 1e-3
-    stepped = fisher.exp_map(mid, 0.1 * fisher.conditional_velocity(mid, x1, 0.5))
-    close(stepped.tolist(), [0.824, 0.327, 0.327, 0.327], 1e-3)
-    close(fisher.to_simplex(stepped).tolist(), [0.680, 0.107, 0.107, 0.107], 1e-3)
-    assert abs(float(stepped.norm()) - 1.0) < 1e-10
-
-
-def test_tangent_projection():
-    x0, p1 = fisher.to_sphere(torch.full((4,), 0.25)), fisher.smooth(torch.eye(4)[0])
-    mid = fisher.geodesic(x0, fisher.to_sphere(p1), torch.tensor(0.5))
-    raw = torch.tensor([0.3, -0.2, 0.5, 0.1])
-    assert abs(float((mid * raw).sum()) - 0.378) < 1e-3
-    tangent = fisher.project(mid, raw)
-    close(tangent.tolist(), [0.005, -0.337, 0.363, -0.037], 1e-3)
-    assert abs(float((mid * tangent).sum())) < 1e-12
-
-
-# %% Section 5.22: Gumbel-Softmax Flow Matching
-
-
-def test_gumbel_max_and_temperature():
-    pi = torch.tensor([0.5, 0.2, 0.2, 0.1])
-    torch.manual_seed(0)
-    frequencies = gumbel.gumbel_max(pi, 200_000)
-    close(frequencies.tolist(), pi.tolist(), tolerance=5e-3)
-    close([gumbel.temperature(t) for t in (0.0, 0.5, 0.9, 1.0)], [10.0, 2.2313, 0.6721, 0.4979], 5e-4)
-
-
-def test_annealed_path_and_replicator():
-    logits = torch.tensor([0.0, -0.45, 1.40, -0.05])
-    expected = {
-        0.0: [0.2438, 0.2331, 0.2805, 0.2426],
-        0.5: [0.2142, 0.1751, 0.4012, 0.2095],
-        0.9: [0.0955, 0.0489, 0.7669, 0.0887],
-        1.0: [0.0528, 0.0214, 0.8781, 0.0477],
-    }
-    for t, values in expected.items():
-        close(gumbel.annealed_path(logits, t).tolist(), values)
-    x = gumbel.annealed_path(logits, 0.5)
-    assert abs(float((x * logits).sum()) - 0.4724) < 5e-4
-    noisy = gumbel.replicator_velocity(x, logits, 0.5)
-    close(noisy.tolist(), [-0.136, -0.217, 0.500, -0.147], 1e-3)
-    assert abs(float(noisy.sum())) < 1e-12
-    close(gumbel.inference_velocity(x, 2, 0.5).tolist(), [-0.116, -0.094, 0.323, -0.113], 1e-3)
-
-
-def test_gumbel_marginal_velocity():
-    late = torch.tensor([0.138, 0.084, 0.647, 0.131])
-    p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15])
-    coefficients = (gumbel.LAMBDA / gumbel.temperature(0.8)) * late * p_hat
-    close(coefficients.tolist(), [0.0228, 0.0278, 1.4976, 0.0650], 5e-4)
-    v = gumbel.marginal_velocity(late, p_hat, 0.8)
-    close(v.tolist(), [-0.200, -0.108, 0.454, -0.146], 1e-3)
-    close((late + 0.01 * v).tolist(), [0.1360, 0.0829, 0.6515, 0.1295])
-
-
-# %% Sections 5.24 to 5.29 and 5.33: guidance
-
-
-def test_guided_base():
-    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4])
-    rates = p1 / (1.0 - 0.6)
-    scores = torch.tensor([0.10, 0.35, 0.55, 0.25])
-    close(rates.tolist(), [0.25, 0.50, 0.75, 1.00])
-    assert abs(float(rates.sum()) - 2.50) < 1e-12
-    close((scores / 0.20).tolist(), [0.50, 1.75, 2.75, 1.25])
-    guided = guidance.tilt(rates, scores, 0.20)
-    close(guided.tolist(), [0.125, 0.875, 2.0625, 1.250])
-    assert abs(float(guided.sum()) - 4.3125) < 1e-12
-    close((guided / guided.sum()).tolist(), [0.029, 0.203, 0.478, 0.290], 1e-3)
-    close((p1 * scores).tolist(), [0.010, 0.070, 0.165, 0.100])
-    h = guidance.h_value(p1, scores)
-    assert abs(h - 0.345) < 1e-12
-    assert abs(float(guidance.tilt(rates, scores, h).sum()) - 2.50) < 1e-12
-    assert guidance.reachable(100, 21) == 2001
-
-
-def test_guidance_strength_and_classifier_free():
-    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4])
-    rates, scores = p1 / 0.4, torch.tensor([0.10, 0.35, 0.55, 0.25])
-    harder = guidance.tilt(rates, scores, 0.20, gamma=2.0)
-    close(harder.tolist(), [0.0625, 1.5312, 5.6719, 1.5625], 1e-3)
-    assert abs(float(harder.sum()) - 8.828) < 1e-3
-    close((harder / harder.sum()).tolist(), [0.007, 0.173, 0.642, 0.177], 1e-3)
-    mixed = guidance.classifier_free(guidance.tilt(rates, scores, 0.20), rates, gamma=2.0)
-    assert torch.allclose(mixed, harder, atol=1e-12)
-
-
-def test_taylor_expansion():
-    gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61])
-    approximate = guidance.taylor_log_ratio(gradient, guidance.MASK)[:4]
-    close(approximate.tolist(), [-0.49, 0.46, 0.86, 0.31])
-    exact = (torch.tensor([0.10, 0.35, 0.55, 0.25]) / 0.20).log()
-    close(exact.tolist(), [-0.693, 0.560, 1.012, 0.223], 1e-3)
-    rates = torch.tensor([0.25, 0.50, 0.75, 1.00]) * approximate.exp()
-    close((rates / rates.sum()).tolist(), [0.038, 0.194, 0.434, 0.334], 1e-3)
-
-
-def test_guided_step_never_remasks_a_visible_position():
-    """The denoiser in the sampler is a distribution over the data tokens, which exclude the mask."""
-    torch.manual_seed(0)
-    model = SequenceDenoiser(guidance.V, 8, width=32, depth=1)
-    prop = guidance.PropertyModel(8, width=32)
-    xt = torch.randint(0, 4, (16, 8))  # Every position already holds a base.
-    stepped = guidance.guided_step(model, prop, xt, torch.full((16,), 0.9), 0.5, gamma=1.0)
-    assert int((stepped == guidance.MASK).sum()) == 0  # A revealed token never returns to the mask.
-
-
-# %% Sections 5.30 to 5.32: several objectives at once
-
-
-def test_pareto_front():
-    scores = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]])
-    assert objectives.pareto_front(scores) == [0, 1]  # The third and fourth are dominated.
-    close((scores @ torch.tensor([0.6, 0.4]))[:2].tolist(), [0.60, 0.64])
-    close((scores @ torch.tensor([0.8, 0.2]))[:2].tolist(), [0.70, 0.62])  # The order reverses.
-
-
-def test_mogdfm_step():
-    changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]])
-    delta_s, ranks, direction = objectives.combined_score(
-        changes, torch.tensor([1.2, 0.8]), torch.tensor([0.6, 0.4]))
-    close(ranks.tolist(), [0.500, 0.700, 0.850, 0.450])
-    close(direction.tolist(), [0.0, 0.100, 0.090, -0.028])
-    close(delta_s.tolist(), [-1.511, 1.541, 2.297, -2.327], 1e-3)
-    close(delta_s.exp().tolist(), [0.221, 4.667, 9.948, 0.098], 1e-2)
-    base = torch.tensor([0.3, 0.3, 0.1]) / 0.3  # The pretrained rates out of A toward C, G and T.
-    tilted = base * delta_s[1:].exp()
-    close(tilted.tolist(), [4.667, 9.948, 0.033], 1e-2)
-    angles, cosines = objectives.cone_angles(changes[1:], torch.tensor([0.6, 0.4]))
-    close(cosines.tolist(), [0.6727, 0.7894, -0.3032], 5e-4)
-    close(angles.tolist(), [47.7, 37.9, 107.7], 0.05)
-    survivors = tilted[angles <= 60.0]
-    close((survivors / survivors.sum()).tolist(), [0.319, 0.681], 1e-3)
-    assert abs(objectives.update_cone(60.0, 1 / 3) - 64.1) < 0.05
-
-
-def test_pcomole_rollout():
-    omega, reference = torch.tensor([0.5, 0.5]), torch.zeros(2)
-    first = [torch.tensor(s) for s in ((0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5))]
-    value, rewards = objectives.rollout_value(first, [1, 1, 0, 1], omega, reference, 2.0)
-    close(rewards, [1.682, 2.054, 0.0, 1.822], 1e-3)
-    assert abs(value - 1.390) < 1e-3
-    second = [torch.tensor(s) for s in ((1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1))]
-    other, _ = objectives.rollout_value(second, [1, 1, 1, 1], omega, reference, 2.0)
-    assert abs(other - 1.181) < 1e-3
-    products = torch.tensor([0.6 * value, 1.2 * other, 0.0])
-    close(products.tolist(), [0.834, 1.417, 0.0], 1e-3)
-    close((products / products.sum()).tolist(), [0.371, 0.629, 0.0], 1e-3)
-
-
-# %% Sections 5.36 to 5.41: rectification and AReUReDi
-
-
-def test_total_correlation_on_two_bases():
-    source = torch.full((rectify.STATES,), 1.0 / rectify.STATES)
-    q, support = rectify.complementary_pairs()
-    coupling = q[None, :].repeat(rectify.STATES, 1)
-    assert abs(rectify.total_correlation(coupling, source) - math.log(4)) < 1e-9
-    assert abs(rectify.hamming(coupling, source) - 1.5) < 1e-9
-    assert abs(rectify.one_step_success(coupling, source, support) - 0.25) < 1e-9
-
-
-def test_redi_rounds():
-    source = torch.full((rectify.STATES,), 1.0 / rectify.STATES)
-    q, support = rectify.complementary_pairs()
-    coupling = q[None, :].repeat(rectify.STATES, 1)
-    coupling = rectify.recouple(coupling, source, steps=300)
-    assert abs(rectify.total_correlation(coupling, source) - 1.101) < 2e-3
-    assert abs(rectify.hamming(coupling, source) - 1.022) < 2e-3
-    assert abs(rectify.one_step_success(coupling, source, support) - 0.403) < 2e-3
-    columns = [a * 4 + b for a, b in support]
-    close([float(coupling[0, c]) for c in columns], [0.400, 0.400, 0.100, 0.100], 2e-3)
-    close([float(coupling[3, c]) for c in columns], [0.755, 0.082, 0.082, 0.082], 2e-3)
-    coupling = rectify.recouple(coupling, source, steps=300)
-    assert abs(rectify.total_correlation(coupling, source) - 0.843) < 2e-3
-    assert abs(rectify.one_step_success(coupling, source, support) - 0.525) < 2e-3
-
-
-def test_areuredi_proposal_and_acceptance():
-    scores = torch.tensor([[0.60, 0.80], [0.70, 0.76], [0.50, 0.90], [0.64, 0.62]])
-    weights = torch.tensor([0.5, 0.5])
-    assert abs(rectify.anneal(11, 20) - 12.0) < 1e-12
-    tchebycheff = rectify.tchebycheff(scores, weights)
-    close(tchebycheff.tolist(), [0.300, 0.350, 0.250, 0.310])
-    tilts = (12.0 * tchebycheff).exp()
-    close(tilts.tolist(), [36.60, 66.69, 20.09, 41.26], 5e-2)
-    model = torch.tensor([0.40, 0.25, 0.20, 0.15])
-    q_current, z_current, ratios = rectify.proposal(model, tilts, 0)
-    close(ratios.tolist(), [1.000, 1.822, 0.549, 1.127], 1e-3)
-    close(rectify.barker(ratios).tolist(), [0.500, 0.646, 0.354, 0.530], 1e-3)
-    assert abs(z_current - 0.5118) < 5e-4
-    close(q_current.tolist(), [0.391, 0.315, 0.138, 0.155], 1e-3)
-    _, z_proposed, _ = rectify.proposal(model, tilts, 1)
-    assert abs(z_proposed - 0.3704) < 5e-4
-    assert abs(rectify.acceptance(z_current, z_proposed) - 1.0) < 1e-12
-    assert abs(rectify.acceptance(z_proposed, z_current) - 0.724) < 1e-3
-
-
-def test_areuredi_concentration():
-    base = torch.tensor([0.50, 0.25, 0.15, 0.10])
-    objective = torch.tensor([0.20, 0.30, 0.45, 0.45])
-    tilted = rectify.tilted_distribution(base, objective, 20.0)
-    close(tilted.tolist(), [0.0127, 0.0468, 0.5643, 0.3762], 1e-3)
-    assert abs(float(tilted[2:].sum()) - 0.9405) < 1e-3
-    bound, gap = rectify.concentration_bound(base, objective, 20.0)
-    assert abs(gap - 0.15) < 1e-9
-    assert abs(bound * float(base[1]) - 0.0498) < 5e-4
-    assert bound * float(base[1]) >= float(tilted[1])
-    sharper = rectify.tilted_distribution(base, objective, 60.0)
-    assert abs(float(sharper[2:].sum()) - 0.9999) < 1e-3
-    assert abs(float(sharper[1]) - 1.234e-4) < 1e-6
-
-
-# %% The command line of every file
-
-
-@pytest.mark.parametrize("script", [
-    "velocity", "kinetic_optimal", "rectify", "edit_flows", "multi_objective",
-])
-def test_script_runs(script):
-    """The five files that need no training run end to end at the smallest grid."""
-    result = subprocess.run(
-        [sys.executable, f"lecture_5/{script}.py", "--steps", "20", "--quiet"],
-        cwd=ROOT, capture_output=True, text=True, env={**os.environ, "PYTHONPATH": str(ROOT)},
-    )
-    assert result.returncode == 0, result.stderr[-2000:]
-    assert "nan" not in result.stdout.lower()

tests/test_lecture_6.py

@@ -1,398 +0,0 @@
-"""The worked numbers of Chapter 6, asserted against the Lecture 6 files.
-
-Every number checked here appears in the chapter text, and every one is computed by the file under
-test and not stored in it. The training loops are not run; the tests call the closed-form and
-numerical functions the lectures build their reports from.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-ROOT = Path(__file__).resolve().parents[1]
-sys.path.insert(0, str(ROOT))
-
-
-def load(name: str):
-    """Import one lecture file by path, so that ``lecture_6`` needs no package marker."""
-    spec = importlib.util.spec_from_file_location(f"lecture_6_{name}",
-                                                  ROOT / "lecture_6" / f"{name}.py")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module
-
-
-def close(values, expected, tolerance=5e-4):
-    """Compare a list of numbers against the chapter's, coordinate by coordinate."""
-    assert len(values) == len(expected)
-    for got, want in zip(values, expected):
-        assert got == pytest.approx(want, abs=tolerance)
-
-
-# %% 01. flow_map.py, Sections 6.3 to 6.12
-def test_flow_map_listing():
-    flow_map = load("flow_map")
-    x = torch.tensor(0.7, dtype=torch.float64)
-    assert float(flow_map.flow(x, 0.0, 1.0)) == pytest.approx(4.1, abs=1e-9)
-    assert float(flow_map.flow(flow_map.flow(x, 0.0, 0.5), 0.5, 1.0)) == pytest.approx(4.1, abs=1e-9)
-    assert float(flow_map.flow(flow_map.flow(x, 0.0, 1.0), 1.0, 0.0)) == pytest.approx(0.7, abs=1e-9)
-    from torch.func import jacrev
-
-    assert float(jacrev(lambda z: flow_map.flow(z, 0.0, 1.0))(x)) == pytest.approx(3.0, abs=1e-9)
-
-
-def test_flow_map_axioms_and_liouville():
-    flow_map = load("flow_map")
-    x = torch.tensor(0.7, dtype=torch.float64)
-    identity, cocycle, tangent = flow_map.check_axioms(x, 200)
-    assert identity < 1e-12 and cocycle < 1e-10 and tangent < 1e-4
-    assert flow_map.liouville_log_det(0.0, 1.0, 4000) == pytest.approx(math.log(3.0), abs=1e-6)
-    assert float(flow_map.jacobian_by_ode(x, 0.0, 1.0, 4000)) == pytest.approx(3.0, abs=1e-3)
-
-
-def test_flow_map_transport_and_likelihood():
-    flow_map = load("flow_map")
-    source = flow_map.normal_density(0.7, 0.0, 1.0)
-    data = flow_map.normal_density(4.1, 2.0, 3.0)
-    assert source == pytest.approx(0.3123, abs=1e-4)
-    assert data == pytest.approx(0.1041, abs=1e-4)
-    assert source / data == pytest.approx(3.0, abs=1e-6)
-    assert math.log(source) - math.log(3.0) == pytest.approx(-2.2625, abs=1e-4)
-    assert math.log(data) == pytest.approx(-2.2625, abs=1e-4)
-
-
-# %% 02. residuals.py, Sections 6.16 and 6.17
-def test_stay_still_residuals():
-    residuals = load("residuals")
-    report = residuals.stay_still_residuals()
-    assert report["stay-still Lagrangian"] == pytest.approx(1.0, abs=1e-9)
-    assert report["stay-still Eulerian"] == pytest.approx(1.0, abs=1e-9)
-    assert report["stay-still composition"] == pytest.approx(0.0, abs=1e-12)
-    assert report["stay-still diagonal gap"] == pytest.approx(0.0, abs=1e-12)
-
-
-def test_identity_axiom_is_free():
-    residuals = load("residuals")
-    torch.manual_seed(0)
-    model = residuals.Map(dim=2, width=16)
-    x, t = torch.randn(8, 2), torch.rand(8)
-    assert torch.allclose(model.X(x, t, t), x)
-
-
-# %% 03. meanflow.py, Sections 6.13 and 6.18
-def test_meanflow_scalar_identity():
-    meanflow = load("meanflow")
-    report = meanflow.scalar_check()
-    assert report["v_0,1(1)"] == pytest.approx(math.e - 1.0, abs=1e-5)
-    assert report["b(1, 0)"] == pytest.approx(1.0, abs=1e-6)
-    assert report["partial_s v"] == pytest.approx(-1.0, abs=1e-5)
-    assert report["grad_x v_0,1"] == pytest.approx(math.e - 1.0, abs=1e-5)
-    assert report["d/ds v_0,1"] == pytest.approx(math.e - 2.0, abs=1e-5)
-    assert report["identity gap"] == pytest.approx(0.0, abs=1e-5)
-    assert report["one map evaluation"] == pytest.approx(math.e, abs=1e-5)
-    assert report["one Euler step"] == pytest.approx(2.0, abs=1e-6)
-
-
-def test_meanflow_correction_vanishes_on_the_gaussian_flow():
-    meanflow = load("meanflow")
-    report = meanflow.gaussian_check()
-    assert report["v_0,t(0.7) spread"] == pytest.approx(0.0, abs=1e-6)
-    assert report["v_0,1(0.7)"] == pytest.approx(3.4, abs=1e-5)
-    assert report["b(0.7, 0)"] == pytest.approx(3.4, abs=1e-5)
-
-
-# %% 04. shortcut.py, Section 6.20
-def test_shortcut_doubling():
-    shortcut = load("shortcut")
-    report = shortcut.scalar_doubling()
-    assert report["s(1, 0, d)"] == pytest.approx(1.136102, abs=1e-6)
-    assert report["state after the first step"] == pytest.approx(1.284025, abs=1e-6)
-    assert report["s(x', d, d)"] == pytest.approx(1.458783, abs=1e-6)
-    assert report["average of the two"] == pytest.approx(1.297443, abs=1e-6)
-    assert report["s(1, 0, 2d) direct"] == pytest.approx(1.297443, abs=1e-6)
-    assert report["doubling gap"] == pytest.approx(0.0, abs=1e-12)
-    assert report["state after the doubled step"] == pytest.approx(math.exp(0.5), abs=1e-9)
-    assert shortcut.ladder()[0] == 1.0 / 128.0 and shortcut.ladder()[-1] == 1.0
-
-
-# %% 05. consistency.py, Section 6.20
-def test_consistency_boundary_parameterization():
-    consistency = load("consistency")
-    report = consistency.scalar_consistency()
-    assert report["f(0.5)"] == pytest.approx(math.exp(0.5), abs=1e-6)
-    assert report["f(0.9)"] == pytest.approx(math.exp(0.1), abs=1e-6)
-    assert report["f(1.0)"] == pytest.approx(1.0, abs=1e-9)
-    assert report["c_out(1.0)"] == pytest.approx(0.0, abs=1e-12)
-    assert report["v_0.5,1(1)"] == pytest.approx(1.297443, abs=1e-6)
-
-
-def test_consistency_teacher_floor():
-    consistency = load("consistency")
-    report = consistency.teacher_floor()
-    assert report["f(1, 0.5)"] == pytest.approx(1.648721, abs=1e-6)
-    assert report["f(1.01, 0.51)"] == pytest.approx(1.648639, abs=1e-6)
-    assert report["distillation residual"] == pytest.approx(8.189e-5, rel=1e-3)
-    assert report["local Euler error"] == pytest.approx(5.0167e-5, rel=1e-3)
-    assert report["carried to the data end"] == pytest.approx(
-        report["distillation residual"], rel=1e-6)
-
-
-def test_consistency_trajectory_head():
-    consistency = load("consistency")
-    report = consistency.scalar_ctm()
-    assert report["a_0.0,0.5"] == pytest.approx(0.5, abs=1e-12)
-    assert report["g_0.0,0.5"] == pytest.approx(2.297443, abs=1e-6)
-    assert report["a_0.0,1.0"] == pytest.approx(0.0, abs=1e-12)
-    assert report["g_0.0,1.0"] == pytest.approx(math.e, abs=1e-6)
-    assert report["two routes apart"] == pytest.approx(0.0, abs=1e-12)
-
-
-# %% 06. fmlm.py, Sections 6.26 and 6.27
-def test_fmlm_denoiser_cocycle():
-    fmlm = load("fmlm")
-    report = fmlm.scalar_cocycle()
-    assert report["gamma"] == pytest.approx(0.2, abs=1e-12)
-    assert report["delta_s,t"] == pytest.approx(2.096158, abs=1e-6)
-    assert report["delta_s,u"] == pytest.approx(1.932957, abs=1e-6)
-    assert report["delta_u,t"] == pytest.approx(2.136959, abs=1e-6)
-    assert report["two legs combined"] == pytest.approx(2.096158, abs=1e-6)
-    assert report["convex weights"] == (0.25, 0.75)
-    assert report["convex form"] == pytest.approx(math.exp(0.6), abs=1e-6)
-    assert report["X_s,t(x)"] == pytest.approx(math.exp(0.6), abs=1e-9)
-
-
-def test_fmlm_worked_kl():
-    fmlm = load("fmlm")
-    report = fmlm.worked_kl()
-    assert report["worked gamma"] == pytest.approx(0.2, abs=1e-12)
-    close(report["composed target"], [0.06, 0.12, 0.76, 0.06], 1e-6)
-    assert report["target mass"] == pytest.approx(1.0, abs=1e-6)
-    close(report["the four terms"], [-0.030650, -0.026777, 0.118823, -0.030650], 1e-6)
-    assert report["worked KL"] == pytest.approx(0.030747, abs=1e-6)
-
-
-def test_fmlm_gradient_comparison():
-    """Section 6.27.4, the two gradients in the logits at the chapter's target and student."""
-    fmlm = load("fmlm")
-    report = fmlm.gradient_comparison()
-    close(report["KL gradient"], [0.04, 0.03, -0.11, 0.04], 1e-6)
-    close(report["squared-error gradient"], [0.0099, 0.01335, -0.03315, 0.0099], 1e-6)
-    close(report["squared over KL"], [0.2475, 0.445, 0.3014, 0.2475], 1e-4)  # A quarter to a half.
-    close(report["saturated student"], [1e-6, 0.1666665, 0.7222215, 0.111111], 1e-6)
-    assert report["saturated KL gradient"] == pytest.approx(-0.0600, abs=1e-4)
-    # The chapter quotes -6.10e-08 here. Differentiating its own expression
-    #   d/dz_i (1/2) sum_j (q_j - p_j)^2 = q_i (q_i - p_i) - q_i sum_j q_j (q_j - p_j)
-    # at q_1 = 1e-06 gives -4.6171e-08. The first term on its own is -6.00e-08, so the quoted
-    # number is the first term with the second dropped, and this file prints the full expression.
-    assert report["saturated squared-error gradient, in 1e-08"] == pytest.approx(-4.6171, abs=1e-3)
-    assert report["orders of magnitude between them"] == pytest.approx(6.11, abs=0.01)
-
-
-def test_fmlm_decoding_clock():
-    fmlm = load("fmlm")
-    tau, start = fmlm.clock(fmlm.BOOK_VOCAB)
-    assert start == pytest.approx(0.999967, abs=1e-5)
-    assert tau(0.60) == pytest.approx(0.006, abs=5e-4)
-    assert tau(0.75) == pytest.approx(0.141, abs=5e-4)
-    assert tau(0.80) == pytest.approx(0.456, abs=1e-3)
-    assert tau(0.85) == pytest.approx(0.931, abs=1e-3)
-    grid = fmlm.even_in_tau(fmlm.BOOK_VOCAB)
-    close(grid, [0.0, 0.773490, 0.804481, 0.828095, 1.0], 5e-5)
-    close([tau(t) for t in (0.25, 0.5, 0.75)], [0.00009, 0.00124, 0.14109], 1e-5)
-
-
-def test_fmlm_clock_weights():
-    fmlm = load("fmlm")
-    rounded = [0.0, 0.773, 0.804, 0.828, 1.0]  # The chapter quotes the grid to three decimals.
-    pairs = [((1 - b) / (1 - a), (b - a) / (1 - a)) for a, b in zip(rounded[:-1], rounded[1:])]
-    close([p[0] for p in pairs], [0.2270, 0.8634, 0.8776, 0.0], 1e-4)
-    close([p[1] for p in pairs], [0.7730, 0.1366, 0.1224, 1.0], 1e-4)
-    times = [torch.tensor([v]) for v in rounded[1:4]]
-    assert float(fmlm.composition_weight(*times)) == pytest.approx(0.494620, abs=1e-5)
-    early = [torch.tensor([v]) for v in (0.0, 0.773, 0.804)]
-    assert float(fmlm.composition_weight(*early)) == pytest.approx(0.830144, abs=1e-5)
-
-
-def test_fmlm_endpoint_slice_is_trained():
-    """Section 6.26 asks for t = 1 on purpose, since a continuous draw never lands there."""
-    fmlm = load("fmlm")
-    assert 0.0 < fmlm.ENDPOINT_FRACTION < 1.0
-    tokens, _ = fmlm.synthetic_dna(num_sequences=64, length=16, seed=0)
-    model = fmlm.TwoTimeMLP(dim=16 * 4, width=16)
-    torch.manual_seed(0)
-    loss_fn = fmlm.make_loss(model, tokens, 16, 4, 256, torch.device("cpu"))
-    assert torch.isfinite(loss_fn(0))  # The rows at t = 1 go through the loss like any other.
-    # The map at t = 1 is the prediction itself, so the one-step sampler reads that slice alone.
-    x = torch.randn(2, 16, 4)
-    s, t = torch.zeros(2), torch.ones(2)
-    prediction = fmlm.predict(model, x, s, t, 16, 4)
-    assert torch.allclose(fmlm.step_map(x, prediction, s, t), prediction)
-
-
-# %% 07. categorical_map.py, Section 6.28
-def test_categorical_worked_cancellation():
-    categorical = load("categorical_map")
-    report = categorical.worked_cancellation()
-    close(report["state x"], [-0.2591, -0.6675, -0.0467], 1e-4)
-    close(report["X_s,t(x)"], [0.1611, -0.2420, 0.0943], 1e-4)
-    close(report["pi_s,t(x)"], [0.5813, 0.1835, 0.2352], 1e-4)
-    close(report["pi_t,t(X_s,t(x))"], [0.6927, 0.0917, 0.2157], 1e-4)
-    assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12)
-    close(report["d pi / dt"], [0.5570, -0.4592, -0.0979], 1e-4)
-    close(report["central difference"], [0.5570, -0.4592, -0.0979], 1e-4)
-    close(report["gap a"], [-0.111, 0.092, 0.020], 1e-3)
-    close(report["drift b"], [0.111, -0.092, -0.020], 1e-3)
-
-
-def test_categorical_bound_is_positive_where_the_residual_is_zero():
-    categorical = load("categorical_map")
-    report = categorical.worked_cancellation()
-    assert report["residual |a + b|^2"] < 1e-20
-    assert report["4 KL"] == pytest.approx(0.156, abs=1e-3)
-    assert report["2 |b|^2"] == pytest.approx(0.042, abs=1e-3)
-    assert report["the bound"] == pytest.approx(0.199, abs=2e-3)
-    assert report["teacher entropy"] == pytest.approx(0.8043, abs=1e-3)
-    assert report["cross-entropy"] == pytest.approx(0.8434, abs=1e-3)
-    assert report["cross-entropy"] == pytest.approx(
-        0.25 * report["4 KL"] + report["teacher entropy"], abs=1e-6)
-
-
-def test_categorical_averaged_bound():
-    categorical = load("categorical_map")
-    report = categorical.averaged_bound(1000, seed=0)
-    assert report["averaged residual"] < 1e-18
-    assert report["averaged 4 KL"] == pytest.approx(0.39, abs=0.06)
-    assert report["averaged 2 |b|^2"] == pytest.approx(0.095, abs=0.03)
-
-
-# %% 08. discrete_map.py, Sections 6.26 and 6.29
-def test_discrete_derivative_target_leaves_the_simplex():
-    discrete = load("discrete_map")
-    report = discrete.worked_derivative()
-    assert report["derivative prefactor"] == pytest.approx(6.6667, abs=1e-3)
-    close(report["derivative target"], [-0.3333, -0.6667, 1.3333, -0.3333], 1e-4)
-    assert report["target sum"] == pytest.approx(0.0, abs=1e-6)
-    assert report["negative entries"] == 3
-
-
-def test_discrete_logit_repair():
-    discrete = load("discrete_map")
-    report = discrete.worked_repair()
-    assert report["weighted average of dz/dt"] == pytest.approx(0.0, abs=1e-6)
-    close(report["d psi/dt"], [1.92, -1.14, -0.78], 1e-4)
-    close(report["additive target"], [-0.284, 0.828, 0.456], 1e-4)
-    assert report["additive target sum"] == pytest.approx(1.0, abs=1e-6)
-    assert report["additive negative entries"] == 1
-    close(report["multiplicative factor"], [2.920, 0.544, 0.480], 1e-4)
-    close(report["quotient"], [0.0342, 1.1029, 0.6250], 1e-4)
-    assert report["quotient mass"] == pytest.approx(1.7622, abs=1e-4)
-    close(report["repaired target"], [0.0194, 0.6259, 0.3547], 1e-4)
-    close(report["log of the factor"], [1.0716, -0.6088, -0.7340], 1e-4)
-    assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12)
-    assert report["gamma^E_s,t"] == pytest.approx(0.8, abs=1e-12)
-
-
-# %% 09. stochastic_map.py, Section 6.31
-def test_stochastic_weak_floor():
-    stochastic = load("stochastic_map")
-    report = stochastic.weak_floor(draws=200_000, seed=0)
-    assert report["K_1"] == pytest.approx(3.194528, abs=1e-5)
-    assert report["sqrt K_1"] == pytest.approx(1.7873, abs=1e-4)
-    assert report["sqrt K_0.5"] == pytest.approx(0.9269, abs=1e-4)
-    assert report["one full step"] == pytest.approx(2.2453, abs=1e-3)
-    assert report["two half steps"] == pytest.approx(0.4673, abs=1e-3)
-    assert report["gap"] == pytest.approx(1.778, abs=1e-3)
-    assert report["squared gap"] == pytest.approx(3.161, abs=1e-3)
-    assert report["floor 2 K_1"] == pytest.approx(math.e ** 2 - 1.0, abs=1e-9)
-    assert report["carried variance equals K_1"] == pytest.approx(report["K_1"], abs=1e-9)
-    assert report["the draw that would close it"] == pytest.approx(-0.195, abs=1e-3)
-    assert report["Monte Carlo floor"] == pytest.approx(6.389, abs=0.05)
-
-
-def test_stochastic_chen_relations():
-    stochastic = load("stochastic_map")
-    report = stochastic.worked_chen()
-    assert report["I0 over [0,1]"] == pytest.approx(0.50, abs=1e-9)
-    assert report["I1 over [0,1]"] == pytest.approx(-0.20, abs=1e-9)
-    assert report["fused variance at h=1"] == pytest.approx(1.0 / 3.0, abs=1e-9)
-
-
-def test_stochastic_truncation():
-    stochastic = load("stochastic_map")
-    report = stochastic.truncation_error(paths=512, grid=2048, seed=0)
-    for n, expected in enumerate((1.0, 1.0 / 3.0, 1.0 / 5.0)):
-        assert report[f"coefficient variance I{n}"] == pytest.approx(expected, rel=0.15)
-    errors = [report[f"pathwise error with {k}"] for k in (1, 2, 3)]
-    assert errors[0] > 5.0 * errors[1] > 25.0 * errors[2]
-
-
-# %% 10. expanding.py, Section 6.32
-def test_expanding_local_clocks():
-    expanding = load("expanding")
-    report = expanding.worked_clocks()
-    close(report["local clocks at t=0.8"], [0.800, 0.800, 0.692, 0.333], 1e-3)
-    close(report["clock speeds"], [1.0, 1.0, 1.538, 3.333], 1e-3)
-    close(report["source share"], [0.200, 0.200, 0.308, 0.667], 1e-3)
-    assert report["last clock from its speed"] == pytest.approx(0.667, abs=1e-3)
-    assert report["live coordinates at 0.8"] == 4 and report["live coordinates at 0.6"] == 3
-
-
-def test_expanding_composition_weight():
-    expanding = load("expanding")
-    report = expanding.worked_omega()
-    assert report["omega at (0.3, 0.6, 0.9)"] == pytest.approx(0.125, abs=1e-6)
-    assert report["omega for a birth at 0.0"] == pytest.approx(0.125, abs=1e-6)
-    assert report["omega for a birth at 0.7"] == pytest.approx(0.0, abs=1e-9)
-    assert 0.0 < report["omega for a birth at 0.35"] < report["omega for a birth at 0.0"]
-
-
-# %% 11. posterior_map.py, Sections 6.33 and 6.34
-def test_posterior_tilt_and_jensen_gap():
-    posterior = load("posterior_map")
-    report = posterior.two_point()
-    close(report["tilted target"], [0.1192, 0.8808], 1e-4)
-    close(report["posterior at x=0.7"], [0.9820, 0.0180], 1e-4)
-    close(report["tilted posterior"], [0.8808, 0.1192], 1e-4)
-    assert report["V_t(0.7)"] == pytest.approx(0.1088, abs=1e-4)
-    assert report["V over the data"] == pytest.approx(1.4338, abs=1e-4)
-    assert report["single-endpoint estimate"] == pytest.approx(1.0, abs=1e-9)
-    assert report["Jensen gap"] == pytest.approx(0.434, abs=1e-3)
-
-
-def test_posterior_value_bias_and_gradient():
-    posterior = load("posterior_map")
-    report = posterior.value_bias()
-    close([report[f"E[V_{k}]"] for k in (1, 2, 4, 8)], [1.000, 1.217, 1.340, 1.393], 1e-3)
-    close([report[f"bias at K={k}"] for k in (1, 2, 4, 8)], [0.434, 0.217, 0.093, 0.041], 1e-3)
-    weights, weighted, plain = posterior.value_gradient(
-        torch.tensor([0.2, 1.5, 0.9]), torch.tensor([0.4, -1.2, 0.5]))
-    close(weights.tolist(), [0.1496, 0.5491, 0.3013], 1e-4)
-    assert weighted == pytest.approx(-0.4483, abs=1e-4)
-    assert plain == pytest.approx(-0.1000, abs=1e-6)
-
-
-def test_posterior_renoising_and_meta_fixed_point():
-    posterior = load("posterior_map")
-    report = posterior.renoise_example()
-    assert report["renoise scale"] == pytest.approx(0.75, abs=1e-9)
-    assert report["renoise variance"] == pytest.approx(0.1375, abs=1e-9)
-    assert report["renoise standard deviation"] == pytest.approx(0.3708, abs=1e-4)
-    assert report["proposal weight on the far mode"] == pytest.approx(0.420, abs=1e-3)
-    assert report["penalty at z=2"] == pytest.approx(10.125, abs=1e-6)
-    assert report["penalty at z=0"] == pytest.approx(6.125, abs=1e-6)
-    assert report["corrected weight on the far mode"] == pytest.approx(0.0131, abs=1e-4)
-    meta = posterior.stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]),
-                                               torch.tensor([0.4, 1.0, 3.6]))
-    assert meta["E[w]"] == pytest.approx(8.0 / 3.0, abs=1e-6)
-    assert meta["E[G]"] == pytest.approx(5.0 / 3.0, abs=1e-6)
-    assert meta["stop-gradient fixed point"] == pytest.approx(0.625, abs=1e-6)
-    assert meta["per-sample ratio average"] == pytest.approx(0.540, abs=1e-6)
-    assert meta["gap"] == pytest.approx(-0.085, abs=1e-3)
-    assert meta["covariance of G/w and w"] == pytest.approx(0.2267, abs=1e-4)
-    assert meta["gap from the covariance"] == pytest.approx(meta["gap"], abs=1e-6)

tests/test_lecture_7.py

@@ -1,709 +0,0 @@
-"""Lecture 7: every number the chapter works by hand, asserted against the file that prints it.
-
-Chapter 7 computes almost all of its results on paper, so these tests are the chapter itself read
-back. The tolerances are the number of digits the chapter quotes.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import math
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-LECTURE = Path(__file__).resolve().parent.parent / "lecture_7"
-
-
-def load(name: str):
-    """Import one lecture file by path, since ``lecture_7`` is a folder of scripts, not a package."""
-    spec = importlib.util.spec_from_file_location(f"lecture_7_{name}", LECTURE / f"{name}.py")
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-def tensor(*values):
-    return torch.tensor(values, dtype=torch.float64)
-
-
-# %% Sections 7.1 to 7.6, optimal_transport.py
-
-
-def test_four_points_and_two_pairings():
-    ot = load("optimal_transport")
-    pairings = ot.two_pairings()
-    assert pairings["arrows_ordered"] == pytest.approx(2.0)
-    assert pairings["arrows_crossed"] == pytest.approx(10.0)
-    assert pairings["mass_ordered"] == pytest.approx(1.0)
-    assert pairings["mass_crossed"] == pytest.approx(5.0)
-    assert pairings["mass_independent"] == pytest.approx(3.0)
-    assert pairings["C"].tolist() == [[1.0, 9.0], [1.0, 1.0]]
-
-
-def test_two_by_two_plan_and_prices():
-    ot = load("optimal_transport")
-    assert ot.feasible_interval() == (0.0, 0.3)
-    C = ot.cost_matrix(ot.SOURCE_X, ot.TARGET_Y)
-    assert ot.cost_along_the_segment(0.0, C) == pytest.approx(5.8)
-    assert ot.cost_along_the_segment(0.3, C) == pytest.approx(3.4)
-    plan = ot.feasible_plan(0.3)
-    assert plan.tolist() == [[0.3, 0.3], [0.0, 0.4]]
-    assert ot.is_coupling(plan, ot.SOURCE_A, ot.TARGET_B)
-    f, g = ot.prices_from_loaded_routes(C, plan)
-    assert f.tolist() == [0.0, -8.0] and g.tolist() == [1.0, 9.0]
-    assert ot.slack(C, f, g).tolist() == [[0.0, 0.0], [8.0, 0.0]]
-    assert ot.dual_value(ot.SOURCE_A, f, ot.TARGET_B, g) == pytest.approx(3.4)
-    assert ot.plan_cost(C, plan) ** 0.5 == pytest.approx(1.844, abs=5e-4)
-
-
-def test_no_crossing_and_three_points():
-    ot = load("optimal_transport")
-    expanded, collected = ot.crossing_difference(0.0, 2.0, 1.0, 3.0)
-    assert expanded == pytest.approx(collected) and collected == pytest.approx(-8.0)
-    x = tensor(4.0, 1.0, 6.0)
-    y = tensor(9.0, 2.0, 3.0)
-    costs = [c for c, _ in ot.monge_permutations(x.sort().values, y.sort().values)]
-    assert costs == pytest.approx([11 / 3, 17 / 3, 35 / 3, 45 / 3, 77 / 3, 81 / 3])
-    assert ot.wasserstein_on_the_line(x, y, 2.0, 3072) == pytest.approx(11 / 3)
-
-
-def test_gaussian_closed_form():
-    ot = load("optimal_transport")
-    assert ot.gaussian_w2_squared(0.0, 1.0, 3.0, 2.0) == pytest.approx(10.0)
-    assert math.sqrt(10.0) == pytest.approx(3.162, abs=5e-4)
-    points = tensor(-1.0, 0.0, 1.0)
-    assert ot.brenier_map(points, 0.0, 1.0, 3.0, 2.0).tolist() == [1.0, 3.0, 5.0]
-    source = torch.special.ndtri(ot.quantile_levels(3072))
-    mapped = ot.brenier_map(source, 0.0, 1.0, 3.0, 2.0)
-    assert float(((mapped - source) ** 2).mean()) == pytest.approx(10.0, abs=1e-3)
-
-
-# %% Sections 7.9 and 7.10, sinkhorn.py
-
-
-def test_entropies_of_the_three_couplings():
-    sk = load("sinkhorn")
-    independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B)
-    assert independent.flatten().tolist() == pytest.approx([0.18, 0.42, 0.12, 0.28])
-    assert sk.entropy(sk.EXACT_PLAN) == pytest.approx(1.089, abs=5e-4)
-    assert sk.entropy(independent) == pytest.approx(1.284, abs=5e-4)
-    assert sk.entropy(sk.SOURCE_A) == pytest.approx(0.673, abs=5e-4)
-    assert sk.entropy(sk.TARGET_B) == pytest.approx(0.611, abs=5e-4)
-    assert sk.relative_entropy(sk.EXACT_PLAN, independent) == pytest.approx(0.195, abs=5e-4)
-
-
-def test_sinkhorn_first_sweep_by_hand():
-    sk = load("sinkhorn")
-    K = sk.gibbs_kernel(sk.COST, 2.0)
-    assert K.flatten().tolist() == pytest.approx([0.606531, 0.011109, 0.606531, 0.606531],
-                                                 abs=5e-7)
-    v = torch.ones(2, dtype=torch.float64)
-    assert (K @ v).tolist() == pytest.approx([0.617640, 1.213061], abs=5e-7)
-    u = sk.SOURCE_A / (K @ v)
-    assert u.tolist() == pytest.approx([0.971440, 0.329744], abs=5e-7)
-    assert (K.T @ u).tolist() == pytest.approx([0.789208, 0.210792], abs=5e-7)
-    v = sk.TARGET_B / (K.T @ u)
-    assert v.tolist() == pytest.approx([0.38013, 3.32081], abs=5e-6)
-    plan = sk.scaling_form(u, K, v)
-    assert plan.flatten().tolist() == pytest.approx([0.2240, 0.0358, 0.0760, 0.6642], abs=5e-5)
-    assert plan.sum(dim=0).tolist() == pytest.approx([0.3, 0.7])
-    assert plan.sum(dim=1).tolist() == pytest.approx([0.2598, 0.7402], abs=5e-5)
-    after = sk.project_rows(plan, sk.SOURCE_A)
-    assert after.sum(dim=0).tolist() == pytest.approx([0.558, 0.442], abs=5e-4)
-
-
-def test_sinkhorn_convergence_and_epsilon():
-    sk = load("sinkhorn")
-    K = sk.gibbs_kernel(sk.COST, 2.0)
-    u, v, errors = sk.sinkhorn_sweeps(K, sk.SOURCE_A, sk.TARGET_B, 40)
-    assert errors[:5] == pytest.approx([0.680, 0.382, 0.155, 0.047, 0.012], abs=5e-4)
-    assert sk.sweeps_to_tolerance(K, sk.SOURCE_A, sk.TARGET_B) == 17
-    plan = sk.scaling_form(u, K, v)
-    assert plan.flatten().tolist() == pytest.approx(
-        [0.293122, 0.306878, 0.006878, 0.393122], abs=5e-7)
-    assert float((sk.COST * plan).sum()) == pytest.approx(3.45502, abs=5e-6)
-    assert sk.entropy(plan) == pytest.approx(1.124, abs=5e-4)
-    independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B)
-    assert sk.entropic_objective(plan, sk.COST, 2.0) == pytest.approx(1.208, abs=5e-4)
-    assert sk.entropic_objective(sk.EXACT_PLAN, sk.COST, 2.0) == pytest.approx(1.222, abs=5e-4)
-    assert sk.entropic_objective(independent, sk.COST, 2.0) == pytest.approx(1.792, abs=5e-4)
-    assert sk.relative_entropy(plan, independent) == pytest.approx(0.160, abs=5e-4)
-    f = sk.multipliers_from_scalings(u, 2.0)
-    g = sk.multipliers_from_scalings(v, 2.0)
-    assert f.tolist() == pytest.approx([3.766, -3.739], abs=5e-4)
-    assert g.tolist() == pytest.approx([-3.220, 4.872], abs=5e-4)
-    expected = {10.0: (4.0540165, 7), 2.0: (3.4550204, 17), 0.5: (3.4000004, 28),
-                0.1: (3.4000000, 80)}
-    for epsilon, (cost, sweeps) in expected.items():
-        kernel = sk.gibbs_kernel(sk.COST, epsilon)
-        su, sv, _ = sk.sinkhorn_sweeps(kernel, sk.SOURCE_A, sk.TARGET_B, 400)
-        assert float((sk.COST * sk.scaling_form(su, kernel, sv)).sum()) == pytest.approx(
-            cost, abs=5e-8)
-        assert sk.sweeps_to_tolerance(kernel, sk.SOURCE_A, sk.TARGET_B) == sweeps
-
-
-# %% Sections 7.7 to 7.7.5, displacement.py
-
-
-def test_one_particle_and_the_averaged_bound():
-    dp = load("displacement")
-    times = torch.linspace(0.0, 1.0, 1025, dtype=torch.float64)
-    assert dp.action_of_a_path(dp.steady_particle(times), times) == pytest.approx(4.0)
-    assert dp.action_of_a_path(dp.dawdling_particle(times), times) == pytest.approx(8.0)
-    x, y = tensor(0.0, 2.0), tensor(1.0, 3.0)
-    ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=torch.float64)
-    crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=torch.float64)
-    independent = torch.full((2, 2), 0.25, dtype=torch.float64)
-    assert dp.average_squared_displacement(x, y, ordered) == pytest.approx(1.0)
-    assert dp.average_squared_displacement(x, y, independent) == pytest.approx(3.0)
-    assert dp.average_squared_displacement(x, y, crossed) == pytest.approx(5.0)
-
-
-def test_displacement_interpolation_and_its_velocity():
-    dp = load("displacement")
-    mean, std = dp.displacement_law(0.5)
-    assert (mean, std) == pytest.approx((1.5, 1.5))
-    assert dp.gaussian_w2(0.0, 1.0, mean, std) == pytest.approx(math.sqrt(2.5))
-    assert math.sqrt(2.5) == pytest.approx(1.581, abs=5e-4)
-    assert dp.gaussian_w2(0.0, 1.0, 3.0, 2.0) == pytest.approx(3.162, abs=5e-4)
-    grid = torch.linspace(-6.0, 9.0, 15001, dtype=torch.float64)
-    peak = float(grid[int(dp.mixture_density(grid).argmax())])
-    assert peak == pytest.approx(0.13, abs=5e-3)
-    half = torch.tensor(0.5, dtype=torch.float64)
-    assert float(dp.velocity_field(torch.tensor(1.5, dtype=torch.float64), half)) == pytest.approx(3.0)
-    assert float(dp.velocity_field(torch.tensor(3.0, dtype=torch.float64), half)) == pytest.approx(4.0)
-    for t in (0.0, 0.25, 0.5, 0.75, 1.0):
-        assert dp.action_density(t, 4096) == pytest.approx(10.0, abs=1e-3)
-    assert abs(dp.continuity_residual(1.5, 0.5)) < 1e-12
-
-
-# %% Sections 7.8 and 7.8.1, ot_flow_matching.py
-
-
-def test_batch_coupling_costs():
-    fm = load("ot_flow_matching")
-    x = torch.tensor([[0.0], [2.0]])
-    y = torch.tensor([[1.0], [3.0]])
-    C = fm.batch_cost_matrix(x, y)
-    assert C.flatten().tolist() == pytest.approx([1.0, 9.0, 1.0, 1.0])
-    assert float(C.mean()) == pytest.approx(3.0)
-    pairing = fm.optimal_pairing(x, y)
-    assert pairing.tolist() == [0, 1]
-    assert float(C[torch.arange(2), pairing].mean()) == pytest.approx(1.0)
-    assert fm.vertex_count(256, 256) == {"unknowns": 65536, "constraints": 512,
-                                         "loaded_routes": 511}
-    assert fm.conditional_velocity_energy(x, fm.couple(x, y, "optimal")) == pytest.approx(1.0)
-
-
-def test_flow_matching_pipeline_runs():
-    fm = load("ot_flow_matching")
-    from cis6270.nets import VelocityMLP
-
-    torch.manual_seed(0)
-    model = VelocityMLP(dim=2, width=16)
-    x0 = torch.randn(8, 2)
-    x1 = torch.randn(8, 2)
-    assert float(fm.flow_matching_loss(model, x0, fm.couple(x0, x1, "optimal"))) > 0.0
-    samples, straightness = fm.integrate(model, x0, 5)
-    assert samples.shape == (8, 2) and straightness >= 0.0
-
-
-# %% Sections 7.12 to 7.15.3, static_bridge.py
-
-
-def test_eight_paths_and_the_chain_rule():
-    sb = load("static_bridge")
-    law = sb.path_law(sb.PI_0, [sb.A, sb.A])
-    assert law[(0, 0, 1)] == pytest.approx(0.105)
-    assert law[(0, 0, 0)] == pytest.approx(0.245)
-    assert law[(1, 1, 1)] == pytest.approx(0.180)
-    assert sum(law.values()) == pytest.approx(1.0)
-    table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=torch.float64)
-    uniform = torch.full((2, 2), 0.25, dtype=torch.float64)
-    total, marginal, conditional = sb.chain_rule_terms(table, uniform)
-    assert total == pytest.approx(0.2180, abs=5e-5)
-    assert marginal == pytest.approx(0.0201, abs=5e-5)
-    assert conditional == pytest.approx(0.1979, abs=5e-5)
-    assert marginal + conditional == pytest.approx(total)
-
-
-def test_proportional_fitting_on_the_chain():
-    sb = load("static_bridge")
-    R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2)
-    assert R01.flatten().tolist() == pytest.approx([0.305, 0.195, 0.26, 0.24])
-    assert float((R01.sum(dim=0) - sb.PI_1).abs().sum()) == pytest.approx(0.73)
-    f, g, errors = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 4)
-    assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3)
-    assert f.tolist() == pytest.approx([1.072002, 0.937019], abs=5e-7)
-    assert g.tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7)
-    plan = f[:, None] * R01 * g[None, :]
-    assert plan.flatten().tolist() == pytest.approx(
-        [0.114605, 0.385395, 0.085395, 0.414605], abs=5e-7)
-    assert sb.kl_vectors(plan.flatten(), R01.flatten()) == pytest.approx(0.2820, abs=5e-5)
-    independent = sb.PI_0[:, None] * sb.PI_1[None, :]
-    assert sb.kl_vectors(independent.flatten(), R01.flatten()) == pytest.approx(0.2847, abs=5e-5)
-
-
-def test_h_transform_and_the_bridge_path_law():
-    sb = load("static_bridge")
-    R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2)
-    f, g, _ = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 40)
-    potentials = sb.backward_potential(sb.A, g, 2)
-    assert potentials[2].tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7)
-    assert potentials[1].tolist() == pytest.approx([0.798453, 1.246389], abs=5e-7)
-    assert potentials[0].tolist() == pytest.approx([0.932834, 1.067215], abs=5e-7)
-    tilted = [sb.tilted_matrix(sb.A, potentials[k], potentials[k + 1]) for k in range(2)]
-    assert tilted[0].flatten().tolist() == pytest.approx(
-        [0.599160, 0.400840, 0.299266, 0.700734], abs=5e-7)
-    assert tilted[1].flatten().tolist() == pytest.approx(
-        [0.307297, 0.692703, 0.112490, 0.887510], abs=5e-7)
-    assert tilted[0].sum(dim=1).tolist() == pytest.approx([1.0, 1.0])
-    middle = sb.PI_0 @ tilted[0]
-    assert middle.tolist() == pytest.approx([0.449213, 0.550787], abs=5e-7)
-    assert (middle @ tilted[1]).tolist() == pytest.approx([0.2, 0.8])
-    assert (1.0 / potentials[0]).tolist() == pytest.approx(f.tolist(), abs=1e-9)
-    plan = f[:, None] * R01 * g[None, :]
-    bridge = sb.fill_with_reference_bridges(plan, sb.A)
-    reference = sb.path_law(sb.PI_0, [sb.A, sb.A])
-    assert bridge[(0, 0, 0)] == pytest.approx(0.092060, abs=5e-7)
-    assert bridge[(0, 0, 1)] == pytest.approx(0.207520, abs=5e-7)
-    assert bridge[(1, 1, 1)] == pytest.approx(0.310954, abs=5e-7)
-    ratios = {round(bridge[p] / reference[p], 6) for p in reference}
-    assert ratios == {0.375755, 1.976383, 0.328441, 1.727522}
-    assert sb.relative_entropy(bridge, reference) == pytest.approx(0.281961, abs=5e-7)
-    assert sum(v for k, v in bridge.items() if k[1] == 0) == pytest.approx(0.449213, abs=5e-7)
-
-
-def test_brownian_reference_recovers_the_entropic_plan():
-    sb = load("static_bridge")
-    plan = sb.brownian_reference_plan(1.0)
-    assert plan.flatten().tolist() == pytest.approx([0.2931, 0.3069, 0.0069, 0.3931], abs=5e-5)
-
-
-# %% Sections 7.16 to 7.18, dynamic_bridge.py
-
-
-def test_constant_control_two_ways():
-    db = load("dynamic_bridge")
-    assert db.control_energy_entropy(2.0, 1.0) == pytest.approx(2.0)
-    assert db.gaussian_entropy_gap(2.0, 1.0) == pytest.approx(2.0)
-    generator = torch.Generator().manual_seed(0)
-    _, log_ratio = db.simulate_controlled(2.0, 200, 8192, 1.0, generator)
-    assert float(log_ratio.mean()) == pytest.approx(2.0, abs=0.1)
-
-
-def test_the_bridge_potential_solves_both_equations():
-    db = load("dynamic_bridge")
-    value = db.hopf_cole(db.bridge_potential)
-    for x, t in ((-0.5, 0.2), (0.5, 0.5), (2.5, 0.8)):
-        assert abs(db.backward_residual(db.bridge_potential, x, t)) < 1e-12
-        assert abs(db.hjb_residual(value, x, t)) < 1e-12
-        drift = float(db.bridge_drift(torch.tensor(x, dtype=torch.float64),
-                                      torch.tensor(t, dtype=torch.float64)))
-        assert db.optimal_control(value, x, t) == pytest.approx(drift, abs=1e-10)
-
-
-def test_brownian_bridge_numbers():
-    db = load("dynamic_bridge")
-    mean, variance = db.bridge_law(0.25)
-    assert mean == pytest.approx(0.5)
-    assert variance == pytest.approx(0.1875)
-    assert math.sqrt(variance) == pytest.approx(0.4330, abs=5e-5)
-    drift = float(db.bridge_drift(torch.tensor(0.6732, dtype=torch.float64),
-                                  torch.tensor(0.25, dtype=torch.float64)))
-    assert drift == pytest.approx(1.7691, abs=5e-5)
-    assert db.bridge_law(0.5)[1] == pytest.approx(0.25)
-    for epsilon, spread in ((0.25, 0.25), (1.0, 0.5), (4.0, 1.0)):
-        assert math.sqrt(db.bridge_law(0.5, epsilon=epsilon)[1]) == pytest.approx(spread)
-
-
-# %% Section 7.19.1, dsb.py
-
-
-def test_dsb_half_steps_on_the_chain():
-    dsb = load("dsb")
-    reference = dsb.reference_path_law(dsb.PI_0, dsb.A, 2)
-    assert dsb.marginal(reference, 1).tolist() == pytest.approx([0.55, 0.45])
-    assert dsb.marginal(reference, 2).tolist() == pytest.approx([0.565, 0.435])
-    B2 = dsb.reverse_step(reference, 1)
-    assert B2.flatten().tolist() == pytest.approx(
-        [0.681416, 0.318584, 0.379310, 0.620690], abs=5e-7)
-    assert float(dsb.reverse_step(reference, 0)[0, 0]) == pytest.approx(0.636364, abs=5e-7)
-    half = dsb.fix_terminal(reference, dsb.PI_1)
-    assert dsb.marginal(half, 2).tolist() == pytest.approx([0.2, 0.8])
-    assert dsb.marginal(half, 1).tolist() == pytest.approx([0.439731, 0.560269], abs=5e-7)
-    assert dsb.marginal(half, 0).tolist() == pytest.approx([0.466585, 0.533415], abs=5e-7)
-    assert float((dsb.marginal(half, 0) - dsb.PI_0).abs().sum()) == pytest.approx(0.0668, abs=5e-5)
-    full = dsb.fix_initial(half, dsb.PI_0)
-    assert dsb.one_step_conditional(full, 0).flatten().tolist() == pytest.approx(
-        [0.599738, 0.400262, 0.299771, 0.700229], abs=5e-7)
-    assert dsb.one_step_conditional(full, 1).flatten().tolist() == pytest.approx(
-        [0.309924, 0.690076, 0.113725, 0.886275], abs=5e-7)
-    assert dsb.marginal(full, 2).tolist() == pytest.approx([0.201967, 0.798033], abs=5e-7)
-    gap = max(float((dsb.one_step_conditional(full, 0) - dsb.TILTED_0).abs().max()),
-              float((dsb.one_step_conditional(full, 1) - dsb.TILTED_1).abs().max()))
-    assert gap < 0.003
-    _, errors, _ = dsb.proportional_fitting(reference, dsb.PI_0, dsb.PI_1, 4)
-    assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3)
-    assert errors[0] / errors[1] == pytest.approx(187.5, abs=0.1)
-
-
-# %% Section 7.19.2, dsbm.py
-
-
-def test_markov_drift_at_one_point():
-    dm = load("dsbm")
-    assert dm.bridge_drift(3.0, 1.0, 0.5) == pytest.approx(4.0)
-    assert dm.bridge_drift(0.0, 1.0, 0.5) == pytest.approx(-2.0)
-    endpoints = tensor(3.0, 0.0)
-    weights = tensor(0.5, 0.5)
-    assert dm.markov_drift(endpoints, weights, 1.0, 0.5) == pytest.approx(1.0)
-    assert dm.least_squares_derivative(1.0, endpoints, weights, 1.0, 0.5) == pytest.approx(0.0)
-    assert dm.least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5) == pytest.approx(-2.0)
-
-
-def test_one_imf_round_on_the_chain():
-    dm = load("dsbm")
-    bridges = dm.reference_bridge_table(dm.A)
-    assert float(bridges[0, 0, 0]) == pytest.approx(0.803279, abs=5e-7)
-    assert float(bridges[0, 0, 1]) == pytest.approx(0.538462, abs=5e-7)
-    target = dm.bridge_coupling(dm.A, dm.PI_0, dm.PI_1)
-    independent = dm.PI_0[:, None] * dm.PI_1[None, :]
-    assert independent.flatten().tolist() == pytest.approx([0.1, 0.4, 0.1, 0.4])
-    assert float((independent - target).abs().sum()) == pytest.approx(0.0584, abs=5e-5)
-    coupling, gaps, first = dm.imf_rounds(independent, dm.A, 9, target)
-    assert first[0].flatten().tolist() == pytest.approx(
-        [0.591425, 0.408575, 0.307692, 0.692308], abs=5e-7)
-    assert first[1].flatten().tolist() == pytest.approx(
-        [0.298457, 0.701543, 0.119588, 0.880412], abs=5e-7)
-    projected = torch.diag(dm.PI_0) @ first[0] @ first[1]
-    assert projected.flatten().tolist() == pytest.approx(
-        [0.112688, 0.387312, 0.087312, 0.412688], abs=5e-7)
-    assert projected.sum(dim=0).tolist() == pytest.approx([0.2, 0.8])
-    assert projected.sum(dim=1).tolist() == pytest.approx([0.5, 0.5])
-    assert gaps[0] == pytest.approx(0.0077, abs=5e-5)
-    assert gaps[3] / gaps[4] == pytest.approx(7.37, abs=5e-3)
-    assert gaps[8] < 1e-9
-
-
-# %% Section 7.19.3, sf2m.py
-
-
-def test_the_two_regression_targets():
-    sm = load("sf2m")
-    x0, x1 = tensor(0.0)[0], tensor(2.0)[0]
-    z = tensor(0.4)[0]
-    for t_value, mean, sample, drift, flow, coefficient in (
-        (0.25, 0.5, 0.6732, 1.7691, 2.2309, 4 / 3),
-        (0.75, 1.5, 1.6732, 1.3072, 1.7691, -4 / 3),
-    ):
-        t = torch.tensor(t_value, dtype=torch.float64)
-        assert float(sm.bridge_mean(x0, x1, t)) == pytest.approx(mean)
-        assert float(sm.bridge_std(t)) == pytest.approx(0.4330, abs=5e-5)
-        x = sm.bridge_sample(x0, x1, t, z)
-        assert float(x) == pytest.approx(sample, abs=5e-5)
-        assert float(sm.conditional_drift(x1, x, t)) == pytest.approx(drift, abs=5e-5)
-        assert float(sm.conditional_score(x0, x1, x, t)) == pytest.approx(-0.9238, abs=5e-5)
-        assert float(sm.flow_velocity(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5)
-        assert float(sm.flow_velocity_collected(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5)
-        assert float(sm.flow_coefficient(t)) == pytest.approx(coefficient)
-        assert float(sm.scaled_score(x0, x1, x, t)) == pytest.approx(-0.4, abs=5e-7)
-
-
-def test_the_batch_plan_of_sf2m():
-    sm = load("sf2m")
-    x, y = tensor(0.0, 2.0), tensor(1.0, 3.0)
-    uniform = tensor(0.5, 0.5)
-    plan = sm.entropic_plan(x, y, uniform, uniform, 1.0)
-    assert plan.sum(dim=0).tolist() == pytest.approx([0.5, 0.5])
-    assert float(plan[0, 0] + plan[1, 1]) > float(plan[0, 1] + plan[1, 0])
-
-
-# %% Sections 7.21 to 7.23.3, discrete_bridge.py
-
-
-def test_rate_error_is_not_symmetric():
-    db = load("discrete_bridge")
-    assert db.scaling_cost(0.5) == pytest.approx(0.1534, abs=5e-5)
-    assert db.scaling_cost(2.0) == pytest.approx(0.3863, abs=5e-5)
-    assert float(db.rate_error(tensor(1.0)[0], tensor(1.0)[0])) == pytest.approx(0.0)
-
-
-def test_two_state_continuous_time_bridge():
-    db = load("discrete_bridge")
-    assert torch.matrix_exp(db.Q).flatten().tolist() == pytest.approx(
-        [0.6328, 0.3672, 0.5507, 0.4493], abs=5e-5)
-    R01 = db.endpoint_law(db.Q, db.PI_0)
-    assert R01.flatten().tolist() == pytest.approx([0.3164, 0.1836, 0.2754, 0.2246], abs=5e-5)
-    assert R01.sum(dim=0).tolist() == pytest.approx([0.5918, 0.4082], abs=5e-5)
-    f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1)
-    assert f.tolist() == pytest.approx([1.0716, 0.9373], abs=5e-5)
-    assert g.tolist() == pytest.approx([0.3349, 1.9643], abs=5e-5)
-    assert float(g[1] / g[0]) == pytest.approx(5.865, abs=5e-4)
-    assert plan.flatten().tolist() == pytest.approx([0.1136, 0.3864, 0.0864, 0.4136], abs=5e-5)
-    assert torch.matrix_exp(db.Q / 2).flatten().tolist() == pytest.approx(
-        [0.7146, 0.2854, 0.4281, 0.5719], abs=5e-5)
-    phi = db.backward_potential(db.Q, g, 0.5)
-    assert phi.tolist() == pytest.approx([0.7999, 1.2667], abs=5e-5)
-    tilted = db.tilted_generator(db.Q, phi)
-    assert tilted.flatten().tolist() == pytest.approx(
-        [-1.5836, 1.5836, 0.9472, -0.9472], abs=5e-5)
-    assert db.backward_potential(db.Q, g, 0.0).tolist() == pytest.approx([0.9331, 1.0669], abs=5e-5)
-    at_zero = db.tilted_generator(db.Q, db.backward_potential(db.Q, g, 0.0))
-    assert float(at_zero[0, 1]) == pytest.approx(1.1433, abs=5e-5)
-    assert float(at_zero[1, 0]) == pytest.approx(1.3120, abs=5e-5)
-    at_one = db.tilted_generator(db.Q, g)
-    assert float(at_one[0, 1]) == pytest.approx(5.8652, abs=5e-5)
-    assert float(at_one[1, 0]) == pytest.approx(0.2557, abs=5e-5)
-    marginal = db.bridge_marginal(db.Q, db.PI_0, f, g, 0.5)
-    assert marginal.tolist() == pytest.approx([0.4668, 0.5332], abs=5e-5)
-    assert db.bridge_marginal(db.Q, db.PI_0, f, g, 1.0).tolist() == pytest.approx([0.2, 0.8])
-    reverse = db.reverse_rates(tilted, marginal)
-    assert float(reverse[1, 0]) == pytest.approx(1.3862, abs=5e-5)
-    assert float(reverse[0, 1]) == pytest.approx(1.0821, abs=5e-5)
-    assert float(marginal[0] * tilted[0, 1]) == pytest.approx(
-        float(marginal[1] * reverse[1, 0]), abs=5e-5)
-
-
-def test_euler_sampler_and_the_cost_computed_twice():
-    db = load("discrete_bridge")
-    R01 = db.endpoint_law(db.Q, db.PI_0)
-    f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1)
-    errors = {}
-    for steps, expected in ((10, 0.0223), (20, 0.0106), (40, 0.0051), (100, 0.0021)):
-        terminal, leaving = db.euler_marginal(db.Q, db.PI_0, f, g, steps)
-        errors[steps] = float((terminal - db.PI_1).abs().sum() / 2)
-        assert errors[steps] == pytest.approx(expected, abs=1e-4)
-        if steps == 10:
-            assert leaving[:3] == pytest.approx([0.1143, 0.1187, 0.1245], abs=5e-5)
-            assert leaving[-1] == pytest.approx(0.3649, abs=5e-5)
-    assert db.relative_entropy(plan, R01) == pytest.approx(0.3235, abs=5e-5)
-    integral, at_half = db.path_cost(db.Q, db.PI_0, f, g)
-    assert at_half == pytest.approx(0.1300, abs=5e-5)
-    assert integral == pytest.approx(0.3235, abs=5e-5)
-    expected = {0.5: 0.3148, 1.0: 0.3235, 2.0: 0.3337, 4.0: 0.3348}
-    for lam, value in expected.items():
-        law = db.endpoint_law(lam * db.Q, db.PI_0)
-        _, _, scaled = db.sinkhorn(law, db.PI_0, db.PI_1)
-        assert db.relative_entropy(scaled, law) == pytest.approx(value, abs=5e-5)
-    assert db.relative_entropy(db.PI_1, db.STATIONARY) == pytest.approx(0.3348, abs=5e-5)
-
-
-# %% Section 7.24, ddsbm.py
-
-
-def test_ddsbm_single_bond():
-    dd = load("ddsbm")
-    kernel = dd.replacement_kernel(0.4, dd.PRIOR)
-    assert kernel.flatten().tolist() == pytest.approx(
-        [0.6, 0.2, 0.2, 0.2, 0.6, 0.2, 0.2, 0.2, 0.6])
-    reference = torch.diag(dd.BOND_0) @ kernel
-    assert reference.sum(dim=0).tolist() == pytest.approx([0.48, 0.28, 0.24])
-    assert (dd.BOND_1 / reference.sum(dim=0)).tolist() == pytest.approx(
-        [0.2083, 0.7143, 2.9167], abs=5e-5)
-    f, g, plan, errors = dd.sinkhorn(reference, dd.BOND_0, dd.BOND_1, 200)
-    assert f.tolist() == pytest.approx([1.1779, 0.9506, 0.5062], abs=5e-5)
-    assert g.tolist() == pytest.approx([0.1842, 0.6918, 3.0004], abs=5e-5)
-    assert errors[-1] < 1e-12
-    assert float(plan[0, 2]) == pytest.approx(0.495, abs=5e-4)
-    assert dd.held_mass(plan) == pytest.approx(0.261, abs=5e-4)
-    assert float((dd.BOND_0 * dd.BOND_1).sum()) == pytest.approx(0.18)
-    assert float(torch.minimum(dd.BOND_0, dd.BOND_1).sum()) == pytest.approx(0.4)
-    expected = {0.2: 0.2250, 0.4: 0.2612, 0.6: 0.2926, 0.8: 0.3248, 0.99: 0.3812}
-    for r, value in expected.items():
-        held = dd.sinkhorn(torch.diag(dd.BOND_0) @ dd.replacement_kernel(r, dd.PRIOR),
-                           dd.BOND_0, dd.BOND_1, 200)[2]
-        assert dd.held_mass(held) == pytest.approx(value, abs=5e-5)
-
-
-def test_ddsbm_pinned_rates_and_the_posterior_average():
-    dd = load("ddsbm")
-    beta = dd.resampling_rate(0.4)
-    assert beta == pytest.approx(0.9163, abs=5e-5)
-    rates = dd.replacement_generator(beta, dd.PRIOR)
-    assert float(rates[0, 1]) == pytest.approx(0.3054, abs=5e-5)
-    assert rates.sum(dim=1).abs().max() < 1e-12
-    assert dd.half_survival_time(0.4) == pytest.approx(math.log(0.8) / math.log(0.4))
-    assert dd.half_survival_time(0.4) == pytest.approx(0.2435, abs=5e-5)
-    pinned = dd.replacement_kernel(0.5, dd.PRIOR)
-    assert float(pinned[0, 0]) == pytest.approx(2 / 3, abs=5e-5)
-    assert float(pinned[0, 1]) == pytest.approx(1 / 6, abs=5e-5)
-    toward = dd.pinned_rate(rates, pinned, 0, 2, 2)
-    away = dd.pinned_rate(rates, pinned, 2, 0, 2)
-    between = dd.pinned_rate(rates, pinned, 0, 1, 2)
-    assert toward == pytest.approx(1.2217, abs=5e-5)
-    assert away == pytest.approx(0.0764, abs=5e-5)
-    assert between == pytest.approx(0.3054, abs=5e-5)
-    assert toward / away == pytest.approx(16.0, abs=5e-4)
-    assert dd.rate_error(toward, 1.0) == pytest.approx(0.0229, abs=5e-5)
-    posterior = 0.6 * toward + 0.4 * between
-    assert posterior == pytest.approx(0.8552, abs=5e-5)
-    rates_two = [toward, between]
-    weights = [0.6, 0.4]
-    assert dd.posterior_loss(rates_two, weights, posterior) == pytest.approx(0.1357, abs=5e-5)
-    assert dd.posterior_loss(rates_two, weights, 1.0) == pytest.approx(0.1467, abs=5e-5)
-    assert dd.posterior_loss(rates_two, weights, 0.5) == pytest.approx(0.2395, abs=5e-5)
-
-
-def test_ddsbm_factorization_counts():
-    dd = load("ddsbm")
-    assert dd.feature_count(9) == (9, 36, 45)
-    assert dd.state_space_size(9, 5, 4) == pytest.approx(9.2e27, rel=5e-3)
-    assert dd.total_jump_rate(36, dd.resampling_rate(0.4), 4) == pytest.approx(24.74, abs=5e-3)
-
-
-# %% Section 7.25, csbm.py
-
-
-def test_csbm_references():
-    cs = load("csbm")
-    A = cs.uniform_reference(0.35, 2)
-    assert A.flatten().tolist() == pytest.approx([0.65, 0.35, 0.35, 0.65])
-    assert torch.linalg.matrix_power(A, 3).flatten().tolist() == pytest.approx(
-        [0.5135, 0.4865, 0.4865, 0.5135], abs=5e-5)
-    assert cs.uniform_reference(0.35, 5).sum(dim=1).tolist() == pytest.approx([1.0] * 5)
-    ordered = cs.ordered_reference(1.0, 5)
-    assert ordered[2].tolist() == pytest.approx(
-        [0.0545, 0.2442, 0.4026, 0.2442, 0.0545], abs=5e-5)
-    assert cs.transition_loss(tensor(0.3160, 0.6840), tensor(0.25, 0.75)) == pytest.approx(
-        0.0110, abs=5e-5)
-
-
-def test_csbm_dimf_on_sixteen_paths():
-    cs = load("csbm")
-    A = cs.uniform_reference(cs.ALPHA, 2)
-    plan = cs.sinkhorn(torch.diag(cs.PI_0) @ torch.linalg.matrix_power(A, 3), cs.PI_0, cs.PI_1)
-    assert plan.flatten().tolist() == pytest.approx(
-        [0.10432, 0.39568, 0.09568, 0.40432], abs=5e-6)
-    assert float(cs.bridge_table(A, 3)[0, 0, 0, 0]) == pytest.approx(0.5348, abs=5e-5)
-    gap, start, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1)
-    assert float(target[0, 0, 0, 0]) == pytest.approx(0.0558, abs=5e-5)
-    assert float(target[1, 1, 1, 1]) == pytest.approx(0.2162, abs=5e-5)
-    assert float(start[0, 0, 0, 0]) == pytest.approx(0.0535, abs=5e-5)
-    assert float(start[1, 1, 1, 1]) == pytest.approx(0.2139, abs=5e-5)
-    assert gap == pytest.approx(0.0173, abs=5e-5)
-    for axis, value in ((1, [0.4754, 0.5246]), (2, [0.4107, 0.5893])):
-        marginal = start.sum(dim=[d for d in range(4) if d != axis])
-        assert marginal.tolist() == pytest.approx(value, abs=5e-5)
-    gaps, conditionals = cs.d_imf(cs.PI_0[:, None] * cs.PI_1[None, :], A, 3, 5, target)
-    assert conditionals[0].flatten().tolist() == pytest.approx(
-        [0.6243, 0.3757, 0.3265, 0.6735], abs=5e-5)
-    assert conditionals[1].flatten().tolist() == pytest.approx(
-        [0.5626, 0.4374, 0.2729, 0.7271], abs=5e-5)
-    assert conditionals[2].flatten().tolist() == pytest.approx(
-        [0.3160, 0.6840, 0.1192, 0.8808], abs=5e-5)
-    assert gaps == pytest.approx([2.9e-4, 5.0e-6, 8.4e-8, 1.4e-9, 2.4e-11], rel=0.05)
-    assert gaps[0] / gaps[1] == pytest.approx(59.0, abs=1.0)
-
-
-def test_csbm_mixing_and_length():
-    cs = load("csbm")
-    independent = cs.PI_0[:, None] * cs.PI_1[None, :]
-    expected = {0.2: (0.13589, 5.0), 0.4: (0.00512, 284.3)}
-    for alpha, (start, ratio) in expected.items():
-        A = cs.uniform_reference(alpha, 2)
-        gap, _, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1)
-        assert gap == pytest.approx(start, abs=5e-5)
-        later, _ = cs.d_imf(independent, A, 3, 2, target)
-        assert later[0] / later[1] == pytest.approx(ratio, rel=0.01)
-    A = cs.uniform_reference(cs.ALPHA, 2)
-    for steps, (start, ratio) in {2: (0.0574, 7.3), 5: (0.00156, 4674.0)}.items():
-        gap, _, target = cs.starting_gap(A, steps, cs.PI_0, cs.PI_1)
-        assert gap == pytest.approx(start, abs=5e-5)
-        later, _ = cs.d_imf(independent, A, steps, 2, target)
-        assert later[0] / later[1] == pytest.approx(ratio, rel=0.01)
-
-
-# %% Sections 7.26 and 7.27, branched.py
-
-
-def test_two_branches_worked_out():
-    br = load("branched")
-    linear = lambda t: 0.3 * t
-    late = lambda t: 0.3 * t ** 2
-    assert float(br.branch_weights(linear, torch.tensor(0.5))) == pytest.approx(0.15)
-    assert float(br.branch_weights(linear, torch.tensor(1.0))) == pytest.approx(0.3)
-    near_zero = br.weighted_energy(2.0, lambda t: 1.0 - linear(t))
-    near_one = br.weighted_energy(-2.0, linear)
-    assert near_zero == pytest.approx(1.7, abs=5e-5)
-    assert near_one == pytest.approx(0.3, abs=5e-5)
-    assert near_zero + near_one == pytest.approx(2.0, abs=5e-5)
-    far_linear = near_zero + br.weighted_energy(-4.0, linear)
-    far_late = br.weighted_energy(2.0, lambda t: 1.0 - late(t)) + br.weighted_energy(-4.0, late)
-    assert far_linear == pytest.approx(2.9, abs=5e-5)
-    assert far_late == pytest.approx(2.6, abs=5e-5)
-
-
-def test_the_interpolant():
-    br = load("branched")
-    grid = torch.tensor([0.0, 0.5, 1.0])
-    assert br.interpolant(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([0.0, 1.25, 2.0])
-    assert br.interpolant_velocity(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([3.0, 2.0, 1.0])
-    assert br.interpolant_energy(0.0, 2.0, 1.0) == pytest.approx(2.1667, abs=5e-5)
-    assert br.interpolant_energy(0.0, 2.0, 0.0) == pytest.approx(2.0, abs=5e-5)
-
-
-def test_entangled_tilt_and_parameterization():
-    br = load("branched")
-    normalizer = 0.1 * 1.0 + 0.9 * 0.01
-    assert normalizer == pytest.approx(0.109)
-    tilted = 0.1 / normalizer
-    assert tilted == pytest.approx(0.917, abs=5e-4)
-    assert tilted / 0.1 == pytest.approx(9.17, abs=5e-3)
-    reward = (1.0 - tilted) * math.log(0.01)
-    assert reward == pytest.approx(-0.380, abs=5e-4)
-    assert reward - math.log(normalizer) == pytest.approx(1.836, abs=5e-4)
-    assert -math.log(normalizer) == pytest.approx(2.216, abs=5e-4)
-    assert br.NOISE == pytest.approx(1.4142, abs=5e-5)
-    step = (-0.3 + 0.5) * 0.01 - 1.0 * 0.8 * 0.01 + br.NOISE * math.sqrt(0.01) * 0.4
-    assert step == pytest.approx(0.0506, abs=5e-5)
-    assert 0.5 / br.NOISE == pytest.approx(0.3536, abs=5e-5)
-    assert 0.5 * (0.5 / br.NOISE) ** 2 * 0.01 == pytest.approx(0.000625, abs=5e-7)
-    force = br.compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]),
-                            torch.tensor([[[1.0, 0.0]]]))[0, 0]
-    assert force.tolist() == pytest.approx([0.5, -0.4])
-    assert float(force.norm()) == pytest.approx(0.640, abs=5e-4)
-    weights = torch.softmax(torch.tensor([-1.0, 0.0, 1.0]), dim=0)
-    assert weights.tolist() == pytest.approx([0.090, 0.245, 0.665], abs=5e-4)
-
-
-def test_entangled_training_lowers_the_loss():
-    br = load("branched")
-    torch.manual_seed(0)
-    generator = torch.Generator().manual_seed(0)
-    force = br.BiasForce(width=16)
-    rollout = br.roll_out(force, 16, 8, torch.device("cpu"), generator)
-    weights = br.importance_weights(rollout)
-    assert float(weights.sum()) == pytest.approx(1.0)
-    assert br.force_check(rollout["positions"][-1]) < 1e-5
-    loss = br.cross_entropy_loss(force, rollout, weights)
-    assert torch.isfinite(loss)
-
-
-# %% Every file runs end to end
-
-
-@pytest.mark.parametrize("name,flags", [
-    ("optimal_transport", ["--steps", "300"]),
-    ("sinkhorn", ["--steps", "20"]),
-    ("displacement", ["--steps", "64"]),
-    ("ot_flow_matching", ["--steps", "3", "--samples", "64", "--sample-steps", "5"]),
-    ("static_bridge", ["--steps", "3"]),
-    ("dynamic_bridge", ["--steps", "20"]),
-    ("dsb", ["--steps", "2"]),
-    ("dsbm", ["--steps", "3"]),
-    ("sf2m", ["--steps", "21"]),
-    ("discrete_bridge", ["--steps", "4"]),
-    ("ddsbm", ["--steps", "20"]),
-    ("csbm", ["--steps", "3"]),
-    ("branched", ["--steps", "2", "--samples", "16"]),
-])
-def test_the_script_runs(name, flags, monkeypatch, capsys):
-    module = load(name)
-    monkeypatch.setattr(sys, "argv", [name] + flags + ["--quiet"])
-    module.main()
-    assert capsys.readouterr().out.strip()