Upload 87 files

0ba9d09 verified about 1 month ago

16.6 kB

"""Continuous-time Markov chains on one DNA base, the mathematics with no learning in it.
A rate matrix fixes what a base can do in the next instant, and the master equation, the matrix
exponential, the stationary distribution and the time reversal all follow from it. Every number
this file prints is a number the chapter works by hand.
Chapter 4, "The two rules a rate matrix obeys" through "Running the jumps backward". Time runs from
t = 0 clean to t = 1 all masks, the survival probability is alpha_t, and the rates R_t(x, y) act on
probability rows, p_dot = p R. Prints every worked quantity of those sections.
    python lecture_4/ctmc.py                   # the worked numbers, 10 Euler steps
    python lecture_4/ctmc.py --steps 20        # a finer Euler grid in Section 09
"""

# %% 01. Imports and configuration
import math  # Logarithms and exponentials of the worked numbers.

import torch  # Matrices, eigenvalues and the matrix exponential.

from cis6270.data import DNA, DNA_SIZE  # The alphabet ACGT and K = 4.
from cis6270.runner import common_parser, print_report, save_run, set_seed  # Flags and the report.

MASK_STATES = DNA_SIZE + 1  # A, C, G, T and m, the five states of the masked alphabet.


# %% 02. The two rules a rate matrix obeys
def rate_matrix(off_diagonal: torch.Tensor) -> torch.Tensor:
    """Fill the diagonal of a nonnegative array so that R_t(x, x) = -sum_{y != x} R_t(x, y)."""
    R = off_diagonal.clone().double()  # [K, K] the chosen jump rates.
    R.fill_diagonal_(0.0)  # The diagonal adds no information of its own.
    return R - torch.diag(R.sum(dim=1))  # [K, K] every row now sums to zero.


def obeys_the_two_rules(R: torch.Tensor, tol: float = 1e-9) -> bool:
    """Nonnegative off the diagonal, and zero row sums. That is the definition."""
    off = R - torch.diag(torch.diag(R))  # [K, K] the diagonal removed.
    return bool((off >= -tol).all()) and bool(R.sum(dim=1).abs().max() < tol)


def exit_rates(R: torch.Tensor) -> torch.Tensor:
    """lambda_t(x) = -R_t(x, x), the total rate of leaving each state."""
    return -torch.diag(R)  # [K] one exit rate per state.


# %% 03. One row of a rate matrix for DNA
#
# A jumps to C at rate 2 and to G and to T at rate 1 each. The chapter has that one row; the
# other three repeat the pattern, with rate 2 on the next base in ACGT order.
DNA_RATES = rate_matrix(torch.tensor([
    [0.0, 2.0, 1.0, 1.0],  # A: rate 2 toward C, rate 1 toward G and T.
    [1.0, 0.0, 2.0, 1.0],  # C: rate 2 toward G.
    [1.0, 1.0, 0.0, 2.0],  # G: rate 2 toward T.
    [2.0, 1.0, 1.0, 0.0],  # T: rate 2 toward A.
]))


def short_interval_row(R: torch.Tensor, state: int, h: float) -> torch.Tensor:
    """One row of P_{t,t+h} = I + h R_t + o(h), the first-order transition probabilities."""
    identity_row = torch.zeros(R.shape[0], dtype=R.dtype)  # [K] the row of I at this state.
    identity_row[state] = 1.0  # Staying put is the only certainty at h = 0.
    return identity_row + h * R[state]  # [K] probabilities, valid while the diagonal stays positive.


def largest_valid_step(R: torch.Tensor) -> float:
    """The bound h * sum_{y != x} R_t(x, y) <= 1, read at the state with the largest exit rate."""
    return float(1.0 / exit_rates(R).max())  # One bad state is enough to give a negative entry.


# %% 04. Waiting times and destinations
#
# A state x holds for an exponential time of rate lambda(x), because S(s + h) = S(s)(1 - lambda h)
# gives dS/ds = -lambda S. The destination is then the row of rates normalized to sum to one.
def sample_path(R: torch.Tensor, start: int, horizon: float, generator) -> list[tuple[float, int]]:
    """Gillespie simulation: an exponential wait, then a destination drawn from the normalized row."""
    time_now, state, path = 0.0, start, [(0.0, start)]  # The path begins at the starting base.
    while True:
        rate = float(exit_rates(R)[state])  # lambda(x) for the base we are sitting on.
        u = float(torch.rand((), generator=generator))  # One uniform variate for the wait.
        time_now += -math.log(u) / rate  # Inverting S(s) = exp(-lambda s) at u.
        if time_now > horizon:
            return path  # The last jump falls outside the interval we are watching.
        row = R[state].clamp(min=0.0)  # [K] the nonnegative rates out of this base.
        state = int(torch.multinomial(row / row.sum(), 1, generator=generator))  # The destination.
        path.append((time_now, state))  # Record when the jump happened and where it went.


# %% 05. The master equation
def master_equation(p: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
    """p_dot = p R, the Kolmogorov forward equation written on probability rows."""
    return p @ R  # [K] the y-th entry collects the y-th column of R, every rate pointing into y.


def arrivals_and_departures(p: torch.Tensor, R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """The same derivative split into the flux j_t(x -> y) arriving at y and the flux leaving y."""
    flux = p[:, None] * R  # [K, K] j_t(x -> y) = p_t(x) R_t(x, y), off the diagonal.
    arriving = flux.sum(dim=0) - torch.diag(flux)  # [K] sum over x != y of p_t(x) R_t(x, y).
    leaving = -torch.diag(flux)  # [K] p_t(y) times the exit rate of y.
    return arriving, leaving  # Their difference is the master equation.


# %% 06. A worked evolution of four DNA probabilities
#
# Every base jumps to every other at rate 1, so each diagonal entry is -3 and each coordinate obeys
# p_dot(y) = 1 - 4 p(y), with solution 1/4 + (p_0(y) - 1/4) exp(-4t).
UNIT_RATES = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # [4, 4] all pairwise rates equal to 1.
START = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float64)  # p_t over (A, C, G, T).


def exact_evolution(p0: torch.Tensor, t: float) -> torch.Tensor:
    """The closed-form solution of p_dot(y) = 1 - 4 p(y) for the all-ones rate matrix."""
    return 0.25 + (p0 - 0.25) * math.exp(-4.0 * t)  # [K] each coordinate decays toward 1/4.


# %% 07. The matrix exponential and its composition over two intervals
def transition_matrix(R: torch.Tensor, elapsed: float) -> torch.Tensor:
    """P_{0,t} = exp(t R), the solution of the master equation for a fixed rate matrix."""
    return torch.matrix_exp(elapsed * R)  # [K, K] rows sum to one and no entry is negative.


def euler_product(R: torch.Tensor, elapsed: float, n: int) -> torch.Tensor:
    """(I + (t/n) R)^n, the product of n short-interval steps. It approaches exp(t R)."""
    step = torch.eye(R.shape[0], dtype=R.dtype) + (elapsed / n) * R  # [K, K] one first-order step.
    return torch.linalg.matrix_power(step, n)  # [K, K] the steps compose by matrix multiplication.


# %% 08. The R_pi family and its closed-form kernel
#
# At rate one the chain discards the current token and draws a fresh one from pi, so every row of
# the rate matrix is pi with the current state subtracted. Because M = 1 pi^T is idempotent, the
# exponential collapses to a mixture of staying put and drawing afresh.
def r_pi(pi: torch.Tensor) -> torch.Tensor:
    """R_pi = 1 pi^T - I, the one-free-distribution family."""
    return torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :] - torch.eye(len(pi), dtype=pi.dtype)


def pi_kernel(pi: torch.Tensor, alpha: float) -> torch.Tensor:
    """Q(alpha) = alpha I + (1 - alpha) 1 pi^T, the kernel at survival probability alpha."""
    M = torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :]  # [K, K] every row equals pi.
    return alpha * torch.eye(len(pi), dtype=pi.dtype) + (1.0 - alpha) * M  # [K, K] a convex mixture.


# %% 09. Eigenvalues, stationarity and detailed balance
def detailed_balance_gap(R: torch.Tensor, pi: torch.Tensor) -> float:
    """The largest mismatch in pi(x) R(x, y) = pi(y) R(y, x), zero when the pair is reversible."""
    flux = pi[:, None] * R  # [K, K] the stationary flux along each edge.
    return float((flux - flux.T).abs().max())  # Zero exactly when every edge balances.


CYCLE = rate_matrix(torch.tensor([  # A -> C -> G -> A at rate 1, with no jumps the other way.
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 1.0],
    [1.0, 0.0, 0.0],
]))


# %% 10. Time reversal
def reverse_rates(R: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
    """R_hat_t(y, x) = R_t(x, y) p_t(x) / p_t(y), the rates of the process watched backward."""
    reversed_rates = (R * p[:, None] / p[None, :]).T  # [K, K] transpose moves (x, y) to (y, x).
    reversed_rates.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule, not by the ratio.
    return reversed_rates - torch.diag(reversed_rates.sum(dim=1))  # [K, K] a rate matrix again.


# %% 11. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    # --steps is the Euler grid of Section 09 and --samples the number of Gillespie paths of
    # Section 04. The flags that describe a network go unread.
    parser.set_defaults(steps=10, samples=512)
    args = parser.parse_args()
    set_seed(args.seed)  # The same seed gives the same sampled path and the same histogram.
    generator = torch.Generator().manual_seed(args.seed)  # Local stream for the simulations.
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Sections 02 and 03, the row for A and its first-order transition probabilities.
    row_probabilities = short_interval_row(DNA_RATES, 0, 0.1)  # [4] at h = 0.1.
    exact_stay = math.exp(-4.0 * 0.1)  # The probability of never leaving A over that interval.
    talk(f"rate row for A      {[round(float(v), 2) for v in DNA_RATES[0]]}  rules hold "
         f"{obeys_the_two_rules(DNA_RATES)}  exit rate {float(exit_rates(DNA_RATES)[0]):.0f}")
    talk(f"P(A, .) at h = 0.1  {[round(float(v), 2) for v in row_probabilities]}  "
         f"exact stay {exact_stay:.4f}  valid while h <= {largest_valid_step(DNA_RATES):.2f}, "
         f"since at h = 0.3 the staying entry is "
         f"{float(short_interval_row(DNA_RATES, 0, 0.3)[0]):.1f}")

    # Section 04, the waiting time and the destination, first from the two uniforms of the chapter.
    wait = -math.log(0.42) / 4.0  # u = 0.42 inverted through S(s) = exp(-4s).
    destination = DNA[int(torch.cumsum(DNA_RATES[0].clamp(min=0) / 4.0, 0).lt(0.31).sum())]
    path = sample_path(DNA_RATES, 0, 1.0, generator)  # One sampled trajectory over the unit interval.
    ends = torch.tensor([sample_path(DNA_RATES, 0, 1.0, generator)[-1][1]
                         for _ in range(args.samples)])  # [N] the base each path finishes on.
    histogram = torch.bincount(ends, minlength=DNA_SIZE).double() / args.samples  # [4] proportions.
    exact_end = transition_matrix(DNA_RATES, 1.0)[0]  # [4] the exact law of the base at t = 1.
    talk(f"wait from u = 0.42  {wait:.4f}  destination from v = 0.31  {destination}  "
         f"mean {1 / 4:.2f}  median {math.log(2) / 4:.4f}")
    talk(f"sampled path        {' -> '.join(f'{DNA[s]}@{time:.3f}' for time, s in path)}")
    talk(f"{args.samples} paths at t = 1   {[round(float(v), 3) for v in histogram]}  against the "
         f"exact {[round(float(v), 3) for v in exact_end]}")

    # Sections 05 and 06, the master equation and the worked evolution.
    derivative = master_equation(START, UNIT_RATES)  # [4] p_dot = (0.6, 0.2, -0.2, -0.6).
    arriving, leaving = arrivals_and_departures(START, UNIT_RATES)  # The two halves of that number.
    euler_step = START + 0.1 * derivative  # One step of size h = 0.1, inside the bound h <= 1/3.
    talk(f"p_dot at A          {float(arriving[0]):.1f} in minus {float(leaving[0]):.1f} out "
         f"= {float(derivative[0]):.1f}   p_dot {[round(float(v), 1) for v in derivative]}")
    talk(f"Euler step h = 0.1  {[round(float(v), 2) for v in euler_step]}  exact "
         f"{[round(float(v), 4) for v in exact_evolution(START, 0.1)]}  "
         f"total change {float(derivative.sum()):.0e}")
    talk(f"exact at t = 1      {[round(float(v), 4) for v in exact_evolution(START, 1.0)]}")

    # Section 07, the matrix exponential and the Chapman-Kolmogorov composition.
    whole = transition_matrix(UNIT_RATES, 0.7)  # P_{s,u} over an interval of length 0.7.
    first, second = transition_matrix(UNIT_RATES, 0.3), transition_matrix(UNIT_RATES, 0.4)
    halves = first @ second  # P_{s,t} P_{t,u}, the two intervals composed.
    talk(f"exp(tR) rows        sum to {float(whole.sum(1).min()):.6f}, smallest entry "
         f"{float(whole.min()):.6f}, composition gap {float((whole - halves).abs().max()):.2e}")

    # Section 08, the R_pi family, its closed form and the multiplication of survival probabilities.
    uniform = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE, dtype=torch.float64)  # pi uniform over ACGT.
    tau = math.log(2.0)  # The clock the linear schedule reaches at t = 0.5, where alpha = 1/2.
    closed_form = pi_kernel(uniform, math.exp(-tau))  # alpha I + (1 - alpha) 1 pi^T.
    series = transition_matrix(r_pi(uniform), tau)  # The same matrix from the series.
    composed = pi_kernel(uniform, 0.8) @ pi_kernel(uniform, 0.5 / 0.8)  # Q(0.8) Q(0.625).
    talk(f"closed form vs exp  gap {float((closed_form - series).abs().max()):.2e}   "
         f"Q(0.8) Q(0.625) at (A, A) {float(composed[0, 0]):.5f} = Q(0.5) entry "
         f"{float(pi_kernel(uniform, 0.5)[0, 0]):.5f}")

    # Section 09, the matrix exponential worked for DNA, and the masking kernel beside it.
    euler_entries = {n: float(euler_product(r_pi(uniform), tau, n)[0, 0]) for n in (1, 2, args.steps)}
    mask_pi = torch.zeros(MASK_STATES, dtype=torch.float64)  # pi = delta_m on the five states.
    mask_pi[DNA_SIZE] = 1.0  # All of the target probability sits on the mask.
    mask_kernel = pi_kernel(mask_pi, 0.5)  # [5, 5] the masking kernel at alpha = 1/2.
    talk(f"exp(tau R_pi)       diagonal {float(closed_form[0, 0]):.3f}, off-diagonal "
         f"{float(closed_form[0, 1]):.3f}, Euler " +
         ", ".join(f"n={n}: {v:.3f}" for n, v in euler_entries.items()) +
         f";  masking A row {[round(float(v), 2) for v in mask_kernel[0]]}, m row "
         f"{[round(float(v), 2) for v in mask_kernel[DNA_SIZE]]}")

    # Section 10, the eigenvalues and the stationary distribution, then detailed balance.
    eigenvalues = sorted(float(v.real) for v in torch.linalg.eigvals(r_pi(uniform)))  # 0 and -1.
    cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64)  # Uniform over the three bases.
    talk(f"eigenvalues of R_pi {[round(v, 4) for v in eigenvalues]}, pi R_pi "
         f"{float((uniform @ r_pi(uniform)).abs().max()):.2e}, uniform flux "
         f"{float((uniform[:, None] * r_pi(uniform))[0, 1]):.4f} each way with gap "
         f"{detailed_balance_gap(r_pi(uniform), uniform):.2e}")
    talk(f"cycle of three      stationary gap {float((cycle_pi @ CYCLE).abs().max()):.2e}, "
         f"edge gap {detailed_balance_gap(CYCLE, cycle_pi):.4f}")

    # Section 11, the reverse rates and the marginals they retrace.
    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64)  # p_t at alpha_t = 1/2.
    unit_forward = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # Forward rate 1 between any two bases.
    backward = reverse_rates(unit_forward, marginal)  # [4, 4] R_hat_t(y, x).
    talk(f"reverse row for C   {[round(float(v), 3) for v in backward[1]]}  exit rate "
         f"{float(exit_rates(backward)[1]):.1f}  mean holding time "
         f"{1 / float(exit_rates(backward)[1]):.4f}")

    report = {
        "exit rate out of A": float(exit_rates(DNA_RATES)[0]),
        "P(A, A) at h = 0.1": float(row_probabilities[0]),
        "largest valid step": largest_valid_step(DNA_RATES),
        "wait from u = 0.42": wait,
        "p_dot at A": float(derivative[0]),
        "Euler step at A": float(euler_step[0]),
        "composition gap": float((whole - halves).abs().max()),
        "exp(tau R_pi) at (A, A)": float(closed_form[0, 0]),
        "exp(tau R_pi) off diagonal": float(closed_form[0, 1]),
        f"Euler product, n = {args.steps}": euler_entries[args.steps],
        "detailed balance gap, cycle": detailed_balance_gap(CYCLE, cycle_pi),
        "reverse rate C to A": float(backward[1, 0]),
        "reverse rate C to G": float(backward[1, 2]),
        "p R_hat plus p R": float((marginal @ backward + marginal @ unit_forward).abs().max()),
    }
    print_report("ctmc: rates, marginals and their reversal", report)
    save_run(args.out, config=vars(args), report=report)


# %% 12. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work every number of the sections above.

	"""Continuous-time Markov chains on one DNA base, the mathematics with no learning in it.

	A rate matrix fixes what a base can do in the next instant, and the master equation, the matrix
	exponential, the stationary distribution and the time reversal all follow from it. Every number
	this file prints is a number the chapter works by hand.

	Chapter 4, "The two rules a rate matrix obeys" through "Running the jumps backward". Time runs from
	t = 0 clean to t = 1 all masks, the survival probability is alpha_t, and the rates R_t(x, y) act on
	probability rows, p_dot = p R. Prints every worked quantity of those sections.

	python lecture_4/ctmc.py # the worked numbers, 10 Euler steps
	python lecture_4/ctmc.py --steps 20 # a finer Euler grid in Section 09
	"""

	# %% 01. Imports and configuration
	import math # Logarithms and exponentials of the worked numbers.

	import torch # Matrices, eigenvalues and the matrix exponential.

	from cis6270.data import DNA, DNA_SIZE # The alphabet ACGT and K = 4.
	from cis6270.runner import common_parser, print_report, save_run, set_seed # Flags and the report.

	MASK_STATES = DNA_SIZE + 1 # A, C, G, T and m, the five states of the masked alphabet.


	# %% 02. The two rules a rate matrix obeys
	def rate_matrix(off_diagonal: torch.Tensor) -> torch.Tensor:
	"""Fill the diagonal of a nonnegative array so that R_t(x, x) = -sum_{y != x} R_t(x, y)."""
	R = off_diagonal.clone().double() # [K, K] the chosen jump rates.
	R.fill_diagonal_(0.0) # The diagonal adds no information of its own.
	return R - torch.diag(R.sum(dim=1)) # [K, K] every row now sums to zero.


	def obeys_the_two_rules(R: torch.Tensor, tol: float = 1e-9) -> bool:
	"""Nonnegative off the diagonal, and zero row sums. That is the definition."""
	off = R - torch.diag(torch.diag(R)) # [K, K] the diagonal removed.
	return bool((off >= -tol).all()) and bool(R.sum(dim=1).abs().max() < tol)


	def exit_rates(R: torch.Tensor) -> torch.Tensor:
	"""lambda_t(x) = -R_t(x, x), the total rate of leaving each state."""
	return -torch.diag(R) # [K] one exit rate per state.


	# %% 03. One row of a rate matrix for DNA
	#
	# A jumps to C at rate 2 and to G and to T at rate 1 each. The chapter has that one row; the
	# other three repeat the pattern, with rate 2 on the next base in ACGT order.
	DNA_RATES = rate_matrix(torch.tensor([
	[0.0, 2.0, 1.0, 1.0], # A: rate 2 toward C, rate 1 toward G and T.
	[1.0, 0.0, 2.0, 1.0], # C: rate 2 toward G.
	[1.0, 1.0, 0.0, 2.0], # G: rate 2 toward T.
	[2.0, 1.0, 1.0, 0.0], # T: rate 2 toward A.
	]))


	def short_interval_row(R: torch.Tensor, state: int, h: float) -> torch.Tensor:
	"""One row of P_{t,t+h} = I + h R_t + o(h), the first-order transition probabilities."""
	identity_row = torch.zeros(R.shape[0], dtype=R.dtype) # [K] the row of I at this state.
	identity_row[state] = 1.0 # Staying put is the only certainty at h = 0.
	return identity_row + h * R[state] # [K] probabilities, valid while the diagonal stays positive.


	def largest_valid_step(R: torch.Tensor) -> float:
	"""The bound h * sum_{y != x} R_t(x, y) <= 1, read at the state with the largest exit rate."""
	return float(1.0 / exit_rates(R).max()) # One bad state is enough to give a negative entry.


	# %% 04. Waiting times and destinations
	#
	# A state x holds for an exponential time of rate lambda(x), because S(s + h) = S(s)(1 - lambda h)
	# gives dS/ds = -lambda S. The destination is then the row of rates normalized to sum to one.
	def sample_path(R: torch.Tensor, start: int, horizon: float, generator) -> list[tuple[float, int]]:
	"""Gillespie simulation: an exponential wait, then a destination drawn from the normalized row."""
	time_now, state, path = 0.0, start, [(0.0, start)] # The path begins at the starting base.
	while True:
	rate = float(exit_rates(R)[state]) # lambda(x) for the base we are sitting on.
	u = float(torch.rand((), generator=generator)) # One uniform variate for the wait.
	time_now += -math.log(u) / rate # Inverting S(s) = exp(-lambda s) at u.
	if time_now > horizon:
	return path # The last jump falls outside the interval we are watching.
	row = R[state].clamp(min=0.0) # [K] the nonnegative rates out of this base.
	state = int(torch.multinomial(row / row.sum(), 1, generator=generator)) # The destination.
	path.append((time_now, state)) # Record when the jump happened and where it went.


	# %% 05. The master equation
	def master_equation(p: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
	"""p_dot = p R, the Kolmogorov forward equation written on probability rows."""
	return p @ R # [K] the y-th entry collects the y-th column of R, every rate pointing into y.


	def arrivals_and_departures(p: torch.Tensor, R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
	"""The same derivative split into the flux j_t(x -> y) arriving at y and the flux leaving y."""
	flux = p[:, None] * R # [K, K] j_t(x -> y) = p_t(x) R_t(x, y), off the diagonal.
	arriving = flux.sum(dim=0) - torch.diag(flux) # [K] sum over x != y of p_t(x) R_t(x, y).
	leaving = -torch.diag(flux) # [K] p_t(y) times the exit rate of y.
	return arriving, leaving # Their difference is the master equation.


	# %% 06. A worked evolution of four DNA probabilities
	#
	# Every base jumps to every other at rate 1, so each diagonal entry is -3 and each coordinate obeys
	# p_dot(y) = 1 - 4 p(y), with solution 1/4 + (p_0(y) - 1/4) exp(-4t).
	UNIT_RATES = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE)) # [4, 4] all pairwise rates equal to 1.
	START = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float64) # p_t over (A, C, G, T).


	def exact_evolution(p0: torch.Tensor, t: float) -> torch.Tensor:
	"""The closed-form solution of p_dot(y) = 1 - 4 p(y) for the all-ones rate matrix."""
	return 0.25 + (p0 - 0.25) * math.exp(-4.0 * t) # [K] each coordinate decays toward 1/4.


	# %% 07. The matrix exponential and its composition over two intervals
	def transition_matrix(R: torch.Tensor, elapsed: float) -> torch.Tensor:
	"""P_{0,t} = exp(t R), the solution of the master equation for a fixed rate matrix."""
	return torch.matrix_exp(elapsed * R) # [K, K] rows sum to one and no entry is negative.


	def euler_product(R: torch.Tensor, elapsed: float, n: int) -> torch.Tensor:
	"""(I + (t/n) R)^n, the product of n short-interval steps. It approaches exp(t R)."""
	step = torch.eye(R.shape[0], dtype=R.dtype) + (elapsed / n) * R # [K, K] one first-order step.
	return torch.linalg.matrix_power(step, n) # [K, K] the steps compose by matrix multiplication.


	# %% 08. The R_pi family and its closed-form kernel
	#
	# At rate one the chain discards the current token and draws a fresh one from pi, so every row of
	# the rate matrix is pi with the current state subtracted. Because M = 1 pi^T is idempotent, the
	# exponential collapses to a mixture of staying put and drawing afresh.
	def r_pi(pi: torch.Tensor) -> torch.Tensor:
	"""R_pi = 1 pi^T - I, the one-free-distribution family."""
	return torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :] - torch.eye(len(pi), dtype=pi.dtype)


	def pi_kernel(pi: torch.Tensor, alpha: float) -> torch.Tensor:
	"""Q(alpha) = alpha I + (1 - alpha) 1 pi^T, the kernel at survival probability alpha."""
	M = torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :] # [K, K] every row equals pi.
	return alpha * torch.eye(len(pi), dtype=pi.dtype) + (1.0 - alpha) * M # [K, K] a convex mixture.


	# %% 09. Eigenvalues, stationarity and detailed balance
	def detailed_balance_gap(R: torch.Tensor, pi: torch.Tensor) -> float:
	"""The largest mismatch in pi(x) R(x, y) = pi(y) R(y, x), zero when the pair is reversible."""
	flux = pi[:, None] * R # [K, K] the stationary flux along each edge.
	return float((flux - flux.T).abs().max()) # Zero exactly when every edge balances.


	CYCLE = rate_matrix(torch.tensor([ # A -> C -> G -> A at rate 1, with no jumps the other way.
	[0.0, 1.0, 0.0],
	[0.0, 0.0, 1.0],
	[1.0, 0.0, 0.0],
	]))


	# %% 10. Time reversal
	def reverse_rates(R: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
	"""R_hat_t(y, x) = R_t(x, y) p_t(x) / p_t(y), the rates of the process watched backward."""
	reversed_rates = (R * p[:, None] / p[None, :]).T # [K, K] transpose moves (x, y) to (y, x).
	reversed_rates.fill_diagonal_(0.0) # The diagonal is fixed by the row-sum rule, not by the ratio.
	return reversed_rates - torch.diag(reversed_rates.sum(dim=1)) # [K, K] a rate matrix again.


	# %% 11. The report
	def main() -> None:
	parser = common_parser(__doc__.split("\n")[0])
	# --steps is the Euler grid of Section 09 and --samples the number of Gillespie paths of
	# Section 04. The flags that describe a network go unread.
	parser.set_defaults(steps=10, samples=512)
	args = parser.parse_args()
	set_seed(args.seed) # The same seed gives the same sampled path and the same histogram.
	generator = torch.Generator().manual_seed(args.seed) # Local stream for the simulations.
	talk = (lambda *a: None) if args.quiet else print # --quiet keeps only the report block.

	# Sections 02 and 03, the row for A and its first-order transition probabilities.
	row_probabilities = short_interval_row(DNA_RATES, 0, 0.1) # [4] at h = 0.1.
	exact_stay = math.exp(-4.0 * 0.1) # The probability of never leaving A over that interval.
	talk(f"rate row for A {[round(float(v), 2) for v in DNA_RATES[0]]} rules hold "
	f"{obeys_the_two_rules(DNA_RATES)} exit rate {float(exit_rates(DNA_RATES)[0]):.0f}")
	talk(f"P(A, .) at h = 0.1 {[round(float(v), 2) for v in row_probabilities]} "
	f"exact stay {exact_stay:.4f} valid while h <= {largest_valid_step(DNA_RATES):.2f}, "
	f"since at h = 0.3 the staying entry is "
	f"{float(short_interval_row(DNA_RATES, 0, 0.3)[0]):.1f}")

	# Section 04, the waiting time and the destination, first from the two uniforms of the chapter.
	wait = -math.log(0.42) / 4.0 # u = 0.42 inverted through S(s) = exp(-4s).
	destination = DNA[int(torch.cumsum(DNA_RATES[0].clamp(min=0) / 4.0, 0).lt(0.31).sum())]
	path = sample_path(DNA_RATES, 0, 1.0, generator) # One sampled trajectory over the unit interval.
	ends = torch.tensor([sample_path(DNA_RATES, 0, 1.0, generator)[-1][1]
	for _ in range(args.samples)]) # [N] the base each path finishes on.
	histogram = torch.bincount(ends, minlength=DNA_SIZE).double() / args.samples # [4] proportions.
	exact_end = transition_matrix(DNA_RATES, 1.0)[0] # [4] the exact law of the base at t = 1.
	talk(f"wait from u = 0.42 {wait:.4f} destination from v = 0.31 {destination} "
	f"mean {1 / 4:.2f} median {math.log(2) / 4:.4f}")
	talk(f"sampled path {' -> '.join(f'{DNA[s]}@{time:.3f}' for time, s in path)}")
	talk(f"{args.samples} paths at t = 1 {[round(float(v), 3) for v in histogram]} against the "
	f"exact {[round(float(v), 3) for v in exact_end]}")

	# Sections 05 and 06, the master equation and the worked evolution.
	derivative = master_equation(START, UNIT_RATES) # [4] p_dot = (0.6, 0.2, -0.2, -0.6).
	arriving, leaving = arrivals_and_departures(START, UNIT_RATES) # The two halves of that number.
	euler_step = START + 0.1 * derivative # One step of size h = 0.1, inside the bound h <= 1/3.
	talk(f"p_dot at A {float(arriving[0]):.1f} in minus {float(leaving[0]):.1f} out "
	f"= {float(derivative[0]):.1f} p_dot {[round(float(v), 1) for v in derivative]}")
	talk(f"Euler step h = 0.1 {[round(float(v), 2) for v in euler_step]} exact "
	f"{[round(float(v), 4) for v in exact_evolution(START, 0.1)]} "
	f"total change {float(derivative.sum()):.0e}")
	talk(f"exact at t = 1 {[round(float(v), 4) for v in exact_evolution(START, 1.0)]}")

	# Section 07, the matrix exponential and the Chapman-Kolmogorov composition.
	whole = transition_matrix(UNIT_RATES, 0.7) # P_{s,u} over an interval of length 0.7.
	first, second = transition_matrix(UNIT_RATES, 0.3), transition_matrix(UNIT_RATES, 0.4)
	halves = first @ second # P_{s,t} P_{t,u}, the two intervals composed.
	talk(f"exp(tR) rows sum to {float(whole.sum(1).min()):.6f}, smallest entry "
	f"{float(whole.min()):.6f}, composition gap {float((whole - halves).abs().max()):.2e}")

	# Section 08, the R_pi family, its closed form and the multiplication of survival probabilities.
	uniform = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE, dtype=torch.float64) # pi uniform over ACGT.
	tau = math.log(2.0) # The clock the linear schedule reaches at t = 0.5, where alpha = 1/2.
	closed_form = pi_kernel(uniform, math.exp(-tau)) # alpha I + (1 - alpha) 1 pi^T.
	series = transition_matrix(r_pi(uniform), tau) # The same matrix from the series.
	composed = pi_kernel(uniform, 0.8) @ pi_kernel(uniform, 0.5 / 0.8) # Q(0.8) Q(0.625).
	talk(f"closed form vs exp gap {float((closed_form - series).abs().max()):.2e} "
	f"Q(0.8) Q(0.625) at (A, A) {float(composed[0, 0]):.5f} = Q(0.5) entry "
	f"{float(pi_kernel(uniform, 0.5)[0, 0]):.5f}")

	# Section 09, the matrix exponential worked for DNA, and the masking kernel beside it.
	euler_entries = {n: float(euler_product(r_pi(uniform), tau, n)[0, 0]) for n in (1, 2, args.steps)}
	mask_pi = torch.zeros(MASK_STATES, dtype=torch.float64) # pi = delta_m on the five states.
	mask_pi[DNA_SIZE] = 1.0 # All of the target probability sits on the mask.
	mask_kernel = pi_kernel(mask_pi, 0.5) # [5, 5] the masking kernel at alpha = 1/2.
	talk(f"exp(tau R_pi) diagonal {float(closed_form[0, 0]):.3f}, off-diagonal "
	f"{float(closed_form[0, 1]):.3f}, Euler " +
	", ".join(f"n={n}: {v:.3f}" for n, v in euler_entries.items()) +
	f"; masking A row {[round(float(v), 2) for v in mask_kernel[0]]}, m row "
	f"{[round(float(v), 2) for v in mask_kernel[DNA_SIZE]]}")

	# Section 10, the eigenvalues and the stationary distribution, then detailed balance.
	eigenvalues = sorted(float(v.real) for v in torch.linalg.eigvals(r_pi(uniform))) # 0 and -1.
	cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64) # Uniform over the three bases.
	talk(f"eigenvalues of R_pi {[round(v, 4) for v in eigenvalues]}, pi R_pi "
	f"{float((uniform @ r_pi(uniform)).abs().max()):.2e}, uniform flux "
	f"{float((uniform[:, None] * r_pi(uniform))[0, 1]):.4f} each way with gap "
	f"{detailed_balance_gap(r_pi(uniform), uniform):.2e}")
	talk(f"cycle of three stationary gap {float((cycle_pi @ CYCLE).abs().max()):.2e}, "
	f"edge gap {detailed_balance_gap(CYCLE, cycle_pi):.4f}")

	# Section 11, the reverse rates and the marginals they retrace.
	marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64) # p_t at alpha_t = 1/2.
	unit_forward = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE)) # Forward rate 1 between any two bases.
	backward = reverse_rates(unit_forward, marginal) # [4, 4] R_hat_t(y, x).
	talk(f"reverse row for C {[round(float(v), 3) for v in backward[1]]} exit rate "
	f"{float(exit_rates(backward)[1]):.1f} mean holding time "
	f"{1 / float(exit_rates(backward)[1]):.4f}")

	report = {
	"exit rate out of A": float(exit_rates(DNA_RATES)[0]),
	"P(A, A) at h = 0.1": float(row_probabilities[0]),
	"largest valid step": largest_valid_step(DNA_RATES),
	"wait from u = 0.42": wait,
	"p_dot at A": float(derivative[0]),
	"Euler step at A": float(euler_step[0]),
	"composition gap": float((whole - halves).abs().max()),
	"exp(tau R_pi) at (A, A)": float(closed_form[0, 0]),
	"exp(tau R_pi) off diagonal": float(closed_form[0, 1]),
	f"Euler product, n = {args.steps}": euler_entries[args.steps],
	"detailed balance gap, cycle": detailed_balance_gap(CYCLE, cycle_pi),
	"reverse rate C to A": float(backward[1, 0]),
	"reverse rate C to G": float(backward[1, 2]),
	"p R_hat plus p R": float((marginal @ backward + marginal @ unit_forward).abs().max()),
	}
	print_report("ctmc: rates, marginals and their reversal", report)
	save_run(args.out, config=vars(args), report=report)


	# %% 12. Run the script
	if __name__ == "__main__": # Execute when launched directly.
	main() # Work every number of the sections above.