Blame for tests/test_lecture_4.py

"""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():
    """We invert S(s) = exp(-4s) at u = 0.42 and get 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():
    """We observe C under a clean A and get (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():
    """When we reweight by h(y)/h(x) the row sums stay 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 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, or 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 architecture changes only 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