CIS6270 / lecture_5 /kinetic_optimal.py

Upload 87 files

0ba9d09 verified about 1 month ago

12.3 kB

"""Kinetic-optimal paths: choosing the velocity, and then the path, by minimizing an energy.
Several velocities generate the same path and several paths connect the same endpoints, so we need a
number to choose among them. In the flux both constraints are linear and the energy is convex, so
the minimizer is the probability of the destination times the positive part of a difference of
logarithmic growth rates. The mixture velocity of this chapter turns out to be that minimizer for
the path it generates. If we pick a distance on the tokens and cool a softmax of negative distances,
we get a path that moves probability only toward tokens closer to the data token.
Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
Chapter 5, Sections 5.16.1 to 5.16.5 and 5.16.7, with the geodesic schedule of 5.16.6. The mask-path
check takes its data marginal from the letter frequency of the synthetic DNA of ``cis6270.data``.
Prints the kinetic-optimal rates, the worked metric path for DNA, and the energy of two schedules.
    python lecture_5/kinetic_optimal.py                     # 600 quadrature points per integral
    python lecture_5/kinetic_optimal.py --steps 20 --quiet  # a coarse rule, visibly less accurate
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import DNA, DNA_SIZE, dna_splits
from cis6270.runner import common_parser, print_report, save_run, set_seed


torch.set_default_dtype(torch.double)  # Full precision, because the report prints rates to six decimals.
MASK = DNA_SIZE  # Index 4 is the mask symbol m, visited by the mask path.
TRANSITION = torch.tensor([[0, 2, 1, 2], [2, 0, 2, 1], [1, 2, 0, 2], [2, 1, 2, 0]], dtype=torch.double)
# d = 1 for a transition, A for G or C for T, and d = 2 for a transversion across the two classes.


# %% 02. The flux and its energy
#
# The flux along an edge is the velocity times the probability of the state it leaves, and the master
# equation is linear in it. The cost is a weight over the probability of the current state, times the
# squared flux. In the velocity this is an expected squared rate under p_t.


def flux(u, p):
    """j_t(y, x) = u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K, K]."""
    return u * p[None, :]  # Each column scaled by the probability of its current state.


def energy_density(u, p, weight=None):
    """sum_{y != x} w_t(y, x) u_t(y, x)^2 p_t(x), the cost at one time written in the velocity."""
    off = u - torch.diag(torch.diagonal(u))  # The diagonal holds no flux of its own.
    weight = torch.ones_like(u) if weight is None else weight  # w = 1 gives the expected squared rate.
    return float((weight * off ** 2 * p[None, :]).sum())


# %% 03. The velocity of least energy
#
# We minimize the energy edge by edge under the master equation, with the weight 1/p_t(y) that makes
# the cost symmetric in the two ends of an edge. The multiplier then comes out as the logarithmic
# growth rate, and the flux becomes p_t(y) p_t(x) times the positive part of its difference. If we
# divide by p_t(x) we get back the velocity.


def kinetic_optimal(p, p_dot):
    """u*_t(y, x) = p_t(y) [d/dt log p_t(y) - d/dt log p_t(x)]_+. p: [K], p_dot: [K] -> [K, K]."""
    growth = torch.where(p > 0, p_dot / p.clamp(min=1e-300), torch.zeros_like(p))  # d/dt log p_t.
    difference = growth[:, None] - growth[None, :]  # [K, K] destination minus current state.
    u = p[:, None] * difference.clamp(min=0.0)  # Probability moves only toward the faster grower.
    u = u * (p[None, :] > 0)  # A state with no probability sends none.
    u.fill_diagonal_(0.0)
    return u - torch.diag(u.sum(dim=0))  # The two rules fix the diagonal.


# %% 04. Kinetic optimality of the mixture path
#
# On the conditional mixture path the two growth rates are kappa_dot/kappa at the data token and
# minus kappa_dot/(1 - kappa) at the source token, and their difference collapses to the conditional
# velocity built earlier in this chapter. The marginal mask path behaves the same way.


def mask_path(q, kappa_t):
    """p_t and p_dot_t over (A, C, G, T, m) for the mask source with data marginal q. Returns two [5]."""
    p = torch.cat([kappa_t * q, torch.tensor([1.0 - kappa_t])])  # Tokens share kappa_t, m holds 1-kappa.
    p_dot = torch.cat([q.clone(), torch.tensor([-1.0])])  # kappa_dot = 1 on the linear schedule.
    return p, p_dot


def mixture_rate(kappa_t, kappa_dot=1.0):
    """kappa_dot_t/(1 - kappa_t), the conditional velocity out of the source token."""
    return kappa_dot / (1.0 - kappa_t)  # 2 at t = 0.5 on the linear schedule.


# %% 05. A distance on the tokens
#
# The formula for u* holds for any path, so we can design one. If we cool a softmax of negative
# distances, the conditional path concentrates on the data token. The normalizing constant cancels
# between the two logarithms, and the velocity reduces to the probability of the destination times
# beta_dot times the distance the move gains.


def metric_path(x1, beta_t, distance=TRANSITION):
    """p_t(. | x_1) = softmax(-beta_t d(., x_1)). Returns [K]."""
    return torch.softmax(-beta_t * distance[:, x1], dim=0)  # Uniform at beta = 0, a point mass later.


def metric_velocity(x1, beta_t, beta_dot, distance=TRANSITION):
    """u*_t(y, x | x_1) = p_t(y | x_1) beta_dot_t [d(x, x_1) - d(y, x_1)]_+. Returns [K, K]."""
    p = metric_path(x1, beta_t, distance)  # [K] the path at this time.
    d = distance[:, x1]  # [K] distance from each token to the data token.
    gain = (d[None, :] - d[:, None]).clamp(min=0.0)  # [K, K] what the move from x to y gains.
    u = p[:, None] * beta_dot * gain  # Faster when the step in distance is larger.
    u.fill_diagonal_(0.0)
    return u - torch.diag(u.sum(dim=0))  # Columns sum to zero.


# %% 06. A worked metric path for DNA
#
# A transition swaps a purine for a purine or a pyrimidine for a pyrimidine, and a transversion
# crosses between the two classes, so d = 1 and d = 2 record which mistakes are small. With
# beta_t = t/(1 - t) and the data token G, the path at t = 0.5 puts more mass on the purine A than
# on the two pyrimidines, and the rates out of C run only toward A and G.


def beta_schedule(t):
    """beta_t = t/(1 - t) with beta_dot_t = 1/(1 - t)^2, and beta_0 = 0."""
    return t / (1.0 - t), 1.0 / (1.0 - t) ** 2  # 1 and 4 at t = 0.5.


def path_derivative(x1, t, h=1e-6, distance=TRANSITION):
    """d p_t/dt by a central difference, for checking the velocity against the path it generates."""
    forward, backward = metric_path(x1, beta_schedule(t + h)[0], distance), metric_path(
        x1, beta_schedule(t - h)[0], distance)
    return (forward - backward) / (2.0 * h)  # [K], compared against u* p_t.


# %% 07. The energy of two schedules
#
# On a two-state conditional mixture path the energy is the time integral of kappa_dot squared over
# kappa times one minus kappa. The linear schedule has a non-integrable singularity at each endpoint,
# so its energy grows with the truncation and never settles. The geodesic schedule has the constant
# integrand pi squared, the squared arc length of a quarter great circle.


def schedule_energy(kappa, kappa_dot, steps, cut=1e-6):
    """The integral of kappa_dot^2/(kappa (1 - kappa)) over [cut, 1 - cut], by a substituted midpoint.
    The grid is uniform in s with t = 1/(1 + exp(-s)), and its Jacobian t(1 - t) cancels the
    endpoint singularity exactly for the linear schedule and tames it for any other.
    """
    span = math.log((1.0 - cut) / cut)  # The s range that the truncation [cut, 1 - cut] maps to.
    h, total = 2.0 * span / steps, 0.0
    for k in range(steps):
        s = -span + (k + 0.5) * h  # Midpoint of the k-th cell.
        t = 1.0 / (1.0 + math.exp(-s))  # Back to the time variable.
        k_t = min(max(kappa(t), 1e-12), 1.0 - 1e-12)  # Hold the schedule off both endpoints.
        total += h * t * (1.0 - t) * kappa_dot(t) ** 2 / (k_t * (1.0 - k_t))
    return total


LINEAR = (lambda t: t, lambda t: 1.0)  # kappa_t = t.
GEODESIC = (lambda t: math.sin(math.pi * t / 2) ** 2,  # The arc of Section 5.16.6 read as a schedule.
            lambda t: (math.pi / 2) * math.sin(math.pi * t))


def geodesic_schedule(t, source_probability=0.0):
    """kappa_t(x_1) = 1 - sin^2((1 - t) Omega)/sin^2 Omega with Omega = arccos(sqrt(p(x_1)))."""
    omega = math.acos(math.sqrt(source_probability))  # pi/2 under the mask source.
    return 1.0 - math.sin((1.0 - t) * omega) ** 2 / math.sin(omega) ** 2


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    args = parser.parse_args()
    set_seed(args.seed)

    # Section 5.16.3: the mask path at t = 0.5 with the chapter's data marginal.
    q = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
    p, p_dot = mask_path(q, kappa_t=0.5)
    u_star = kinetic_optimal(p, p_dot)
    report = {
        "mask_path_ACGTm": [round(float(v), 4) for v in p],  # (0.05, 0.30, 0.10, 0.05, 0.50).
        "growth_at_mask": round(float(p_dot[MASK] / p[MASK]), 4),  # -2.
        "growth_at_token": round(float(p_dot[1] / p[1]), 4),  # +2.
        "rate_mask_to_C": round(float(u_star[1, MASK]), 4),  # 0.30 times 4, or 1.2.
        "mixture_rate_to_C": round(mixture_rate(0.5) * float(q[1]), 4),  # 2 times 0.6, the same.
        "rate_between_tokens": round(float(u_star[1, 2]), 6),  # Zero, the no-revision mask path.
        "generates_the_path": bool(torch.allclose(u_star @ p, p_dot, atol=1e-9)),
    }

    # The same check with the letter frequency of the synthetic DNA as the data marginal.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
    p_dna, p_dot_dna = mask_path(frequency, kappa_t=0.5)
    u_dna = kinetic_optimal(p_dna, p_dot_dna)
    report["dna_marginal"] = [round(float(v), 4) for v in frequency]
    report["dna_rate_mask_to_G"] = round(float(u_dna[2, MASK]), 4)  # 2 times the frequency of G.
    report["dna_energy_density"] = round(energy_density(u_dna, p_dna), 4)  # Expected squared rate.
    report["dna_flux_mask_to_G"] = round(float(flux(u_dna, p_dna)[2, MASK]), 4)  # Rate times p_t(m).

    # Section 5.16.5: the metric path toward G at t = 0.5.
    g, c = DNA.index("G"), DNA.index("C")
    beta_t, beta_dot = beta_schedule(0.5)  # 1 and 4.
    p_metric = metric_path(g, beta_t)
    u_metric = metric_velocity(g, beta_t, beta_dot)
    report["metric_path_at_0.5"] = [round(float(v), 4) for v in p_metric]
    report["rate_C_to_A"] = round(float(u_metric[0, c]), 3)  # 0.898.
    report["rate_C_to_G"] = round(float(u_metric[g, c]), 3)  # 4.882.
    report["rate_A_to_G"] = round(float(u_metric[g, 0]), 3)  # 2.441.
    report["rate_out_of_C"] = round(float(-u_metric[c, c]), 3)  # 5.780.
    step = torch.eye(DNA_SIZE)[c] + 0.05 * u_metric[:, c]  # One Euler step of size h = 0.05 from C.
    report["euler_step_from_C"] = [round(float(v), 4) for v in step]
    report["p_dot_from_velocity"] = [round(float(v), 4) for v in u_metric @ p_metric]
    report["p_dot_from_path"] = [round(float(v), 4) for v in path_derivative(g, 0.5)]

    # Section 5.16.7: the energy of the two schedules, and the geodesic schedule itself.
    report["energy_linear_1e-6"] = round(schedule_energy(*LINEAR, args.steps, 1e-6), 2)  # 27.63.
    report["energy_linear_1e-9"] = round(schedule_energy(*LINEAR, args.steps, 1e-9), 2)  # 41.45.
    report["energy_geodesic"] = round(schedule_energy(*GEODESIC, args.steps, 1e-6), 4)  # 9.8696.
    report["pi_squared"] = round(math.pi ** 2, 4)  # The squared arc length of a quarter circle.
    report["geodesic_kappa_0.25"] = round(geodesic_schedule(0.25), 4)  # 0.1464.
    report["geodesic_kappa_0.75"] = round(geodesic_schedule(0.75), 4)  # 0.8536.
    report["geodesic_rate_0.9"] = round(math.pi * math.tan(math.pi * 0.9 / 2), 4)  # 19.84.
    report["uniform_source_kappa_0.5"] = round(geodesic_schedule(0.5, 0.25), 4)  # 0.6667.
    print_report("kinetic-optimal paths", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":
    main()

	"""Kinetic-optimal paths: choosing the velocity, and then the path, by minimizing an energy.

	Several velocities generate the same path and several paths connect the same endpoints, so we need a
	number to choose among them. In the flux both constraints are linear and the energy is convex, so
	the minimizer is the probability of the destination times the positive part of a difference of
	logarithmic growth rates. The mixture velocity of this chapter turns out to be that minimizer for
	the path it generates. If we pick a distance on the tokens and cool a softmax of negative distances,
	we get a path that moves probability only toward tokens closer to the data token.

	Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
	u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).

	Chapter 5, Sections 5.16.1 to 5.16.5 and 5.16.7, with the geodesic schedule of 5.16.6. The mask-path
	check takes its data marginal from the letter frequency of the synthetic DNA of ``cis6270.data``.
	Prints the kinetic-optimal rates, the worked metric path for DNA, and the energy of two schedules.

	python lecture_5/kinetic_optimal.py # 600 quadrature points per integral
	python lecture_5/kinetic_optimal.py --steps 20 --quiet # a coarse rule, visibly less accurate
	"""

	# %% 01. Imports and configuration
	import math

	import torch

	from cis6270.data import DNA, DNA_SIZE, dna_splits
	from cis6270.runner import common_parser, print_report, save_run, set_seed


	torch.set_default_dtype(torch.double) # Full precision, because the report prints rates to six decimals.
	MASK = DNA_SIZE # Index 4 is the mask symbol m, visited by the mask path.
	TRANSITION = torch.tensor([[0, 2, 1, 2], [2, 0, 2, 1], [1, 2, 0, 2], [2, 1, 2, 0]], dtype=torch.double)
	# d = 1 for a transition, A for G or C for T, and d = 2 for a transversion across the two classes.


	# %% 02. The flux and its energy
	#
	# The flux along an edge is the velocity times the probability of the state it leaves, and the master
	# equation is linear in it. The cost is a weight over the probability of the current state, times the
	# squared flux. In the velocity this is an expected squared rate under p_t.


	def flux(u, p):
	"""j_t(y, x) = u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K, K]."""
	return u * p[None, :] # Each column scaled by the probability of its current state.


	def energy_density(u, p, weight=None):
	"""sum_{y != x} w_t(y, x) u_t(y, x)^2 p_t(x), the cost at one time written in the velocity."""
	off = u - torch.diag(torch.diagonal(u)) # The diagonal holds no flux of its own.
	weight = torch.ones_like(u) if weight is None else weight # w = 1 gives the expected squared rate.
	return float((weight * off ** 2 * p[None, :]).sum())


	# %% 03. The velocity of least energy
	#
	# We minimize the energy edge by edge under the master equation, with the weight 1/p_t(y) that makes
	# the cost symmetric in the two ends of an edge. The multiplier then comes out as the logarithmic
	# growth rate, and the flux becomes p_t(y) p_t(x) times the positive part of its difference. If we
	# divide by p_t(x) we get back the velocity.


	def kinetic_optimal(p, p_dot):
	"""u*_t(y, x) = p_t(y) [d/dt log p_t(y) - d/dt log p_t(x)]_+. p: [K], p_dot: [K] -> [K, K]."""
	growth = torch.where(p > 0, p_dot / p.clamp(min=1e-300), torch.zeros_like(p)) # d/dt log p_t.
	difference = growth[:, None] - growth[None, :] # [K, K] destination minus current state.
	u = p[:, None] * difference.clamp(min=0.0) # Probability moves only toward the faster grower.
	u = u * (p[None, :] > 0) # A state with no probability sends none.
	u.fill_diagonal_(0.0)
	return u - torch.diag(u.sum(dim=0)) # The two rules fix the diagonal.


	# %% 04. Kinetic optimality of the mixture path
	#
	# On the conditional mixture path the two growth rates are kappa_dot/kappa at the data token and
	# minus kappa_dot/(1 - kappa) at the source token, and their difference collapses to the conditional
	# velocity built earlier in this chapter. The marginal mask path behaves the same way.


	def mask_path(q, kappa_t):
	"""p_t and p_dot_t over (A, C, G, T, m) for the mask source with data marginal q. Returns two [5]."""
	p = torch.cat([kappa_t * q, torch.tensor([1.0 - kappa_t])]) # Tokens share kappa_t, m holds 1-kappa.
	p_dot = torch.cat([q.clone(), torch.tensor([-1.0])]) # kappa_dot = 1 on the linear schedule.
	return p, p_dot


	def mixture_rate(kappa_t, kappa_dot=1.0):
	"""kappa_dot_t/(1 - kappa_t), the conditional velocity out of the source token."""
	return kappa_dot / (1.0 - kappa_t) # 2 at t = 0.5 on the linear schedule.


	# %% 05. A distance on the tokens
	#
	# The formula for u* holds for any path, so we can design one. If we cool a softmax of negative
	# distances, the conditional path concentrates on the data token. The normalizing constant cancels
	# between the two logarithms, and the velocity reduces to the probability of the destination times
	# beta_dot times the distance the move gains.


	def metric_path(x1, beta_t, distance=TRANSITION):
	"""p_t(. \| x_1) = softmax(-beta_t d(., x_1)). Returns [K]."""
	return torch.softmax(-beta_t * distance[:, x1], dim=0) # Uniform at beta = 0, a point mass later.


	def metric_velocity(x1, beta_t, beta_dot, distance=TRANSITION):
	"""u*_t(y, x \| x_1) = p_t(y \| x_1) beta_dot_t [d(x, x_1) - d(y, x_1)]_+. Returns [K, K]."""
	p = metric_path(x1, beta_t, distance) # [K] the path at this time.
	d = distance[:, x1] # [K] distance from each token to the data token.
	gain = (d[None, :] - d[:, None]).clamp(min=0.0) # [K, K] what the move from x to y gains.
	u = p[:, None] * beta_dot * gain # Faster when the step in distance is larger.
	u.fill_diagonal_(0.0)
	return u - torch.diag(u.sum(dim=0)) # Columns sum to zero.


	# %% 06. A worked metric path for DNA
	#
	# A transition swaps a purine for a purine or a pyrimidine for a pyrimidine, and a transversion
	# crosses between the two classes, so d = 1 and d = 2 record which mistakes are small. With
	# beta_t = t/(1 - t) and the data token G, the path at t = 0.5 puts more mass on the purine A than
	# on the two pyrimidines, and the rates out of C run only toward A and G.


	def beta_schedule(t):
	"""beta_t = t/(1 - t) with beta_dot_t = 1/(1 - t)^2, and beta_0 = 0."""
	return t / (1.0 - t), 1.0 / (1.0 - t) ** 2 # 1 and 4 at t = 0.5.


	def path_derivative(x1, t, h=1e-6, distance=TRANSITION):
	"""d p_t/dt by a central difference, for checking the velocity against the path it generates."""
	forward, backward = metric_path(x1, beta_schedule(t + h)[0], distance), metric_path(
	x1, beta_schedule(t - h)[0], distance)
	return (forward - backward) / (2.0 * h) # [K], compared against u* p_t.


	# %% 07. The energy of two schedules
	#
	# On a two-state conditional mixture path the energy is the time integral of kappa_dot squared over
	# kappa times one minus kappa. The linear schedule has a non-integrable singularity at each endpoint,
	# so its energy grows with the truncation and never settles. The geodesic schedule has the constant
	# integrand pi squared, the squared arc length of a quarter great circle.


	def schedule_energy(kappa, kappa_dot, steps, cut=1e-6):
	"""The integral of kappa_dot^2/(kappa (1 - kappa)) over [cut, 1 - cut], by a substituted midpoint.

	The grid is uniform in s with t = 1/(1 + exp(-s)), and its Jacobian t(1 - t) cancels the
	endpoint singularity exactly for the linear schedule and tames it for any other.
	"""
	span = math.log((1.0 - cut) / cut) # The s range that the truncation [cut, 1 - cut] maps to.
	h, total = 2.0 * span / steps, 0.0
	for k in range(steps):
	s = -span + (k + 0.5) * h # Midpoint of the k-th cell.
	t = 1.0 / (1.0 + math.exp(-s)) # Back to the time variable.
	k_t = min(max(kappa(t), 1e-12), 1.0 - 1e-12) # Hold the schedule off both endpoints.
	total += h * t * (1.0 - t) * kappa_dot(t) ** 2 / (k_t * (1.0 - k_t))
	return total


	LINEAR = (lambda t: t, lambda t: 1.0) # kappa_t = t.
	GEODESIC = (lambda t: math.sin(math.pi * t / 2) ** 2, # The arc of Section 5.16.6 read as a schedule.
	lambda t: (math.pi / 2) * math.sin(math.pi * t))


	def geodesic_schedule(t, source_probability=0.0):
	"""kappa_t(x_1) = 1 - sin^2((1 - t) Omega)/sin^2 Omega with Omega = arccos(sqrt(p(x_1)))."""
	omega = math.acos(math.sqrt(source_probability)) # pi/2 under the mask source.
	return 1.0 - math.sin((1.0 - t) * omega) 2 / math.sin(omega) 2


	# %% 08. The report
	def main():
	parser = common_parser(__doc__.split("\n")[0])
	args = parser.parse_args()
	set_seed(args.seed)

	# Section 5.16.3: the mask path at t = 0.5 with the chapter's data marginal.
	q = torch.tensor([0.1, 0.6, 0.2, 0.1]) # [4] over A, C, G, T.
	p, p_dot = mask_path(q, kappa_t=0.5)
	u_star = kinetic_optimal(p, p_dot)
	report = {
	"mask_path_ACGTm": [round(float(v), 4) for v in p], # (0.05, 0.30, 0.10, 0.05, 0.50).
	"growth_at_mask": round(float(p_dot[MASK] / p[MASK]), 4), # -2.
	"growth_at_token": round(float(p_dot[1] / p[1]), 4), # +2.
	"rate_mask_to_C": round(float(u_star[1, MASK]), 4), # 0.30 times 4, or 1.2.
	"mixture_rate_to_C": round(mixture_rate(0.5) * float(q[1]), 4), # 2 times 0.6, the same.
	"rate_between_tokens": round(float(u_star[1, 2]), 6), # Zero, the no-revision mask path.
	"generates_the_path": bool(torch.allclose(u_star @ p, p_dot, atol=1e-9)),
	}

	# The same check with the letter frequency of the synthetic DNA as the data marginal.
	splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
	counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
	frequency = counts / counts.sum() # [4] the GC-rich background the lecture shares.
	p_dna, p_dot_dna = mask_path(frequency, kappa_t=0.5)
	u_dna = kinetic_optimal(p_dna, p_dot_dna)
	report["dna_marginal"] = [round(float(v), 4) for v in frequency]
	report["dna_rate_mask_to_G"] = round(float(u_dna[2, MASK]), 4) # 2 times the frequency of G.
	report["dna_energy_density"] = round(energy_density(u_dna, p_dna), 4) # Expected squared rate.
	report["dna_flux_mask_to_G"] = round(float(flux(u_dna, p_dna)[2, MASK]), 4) # Rate times p_t(m).

	# Section 5.16.5: the metric path toward G at t = 0.5.
	g, c = DNA.index("G"), DNA.index("C")
	beta_t, beta_dot = beta_schedule(0.5) # 1 and 4.
	p_metric = metric_path(g, beta_t)
	u_metric = metric_velocity(g, beta_t, beta_dot)
	report["metric_path_at_0.5"] = [round(float(v), 4) for v in p_metric]
	report["rate_C_to_A"] = round(float(u_metric[0, c]), 3) # 0.898.
	report["rate_C_to_G"] = round(float(u_metric[g, c]), 3) # 4.882.
	report["rate_A_to_G"] = round(float(u_metric[g, 0]), 3) # 2.441.
	report["rate_out_of_C"] = round(float(-u_metric[c, c]), 3) # 5.780.
	step = torch.eye(DNA_SIZE)[c] + 0.05 * u_metric[:, c] # One Euler step of size h = 0.05 from C.
	report["euler_step_from_C"] = [round(float(v), 4) for v in step]
	report["p_dot_from_velocity"] = [round(float(v), 4) for v in u_metric @ p_metric]
	report["p_dot_from_path"] = [round(float(v), 4) for v in path_derivative(g, 0.5)]

	# Section 5.16.7: the energy of the two schedules, and the geodesic schedule itself.
	report["energy_linear_1e-6"] = round(schedule_energy(*LINEAR, args.steps, 1e-6), 2) # 27.63.
	report["energy_linear_1e-9"] = round(schedule_energy(*LINEAR, args.steps, 1e-9), 2) # 41.45.
	report["energy_geodesic"] = round(schedule_energy(*GEODESIC, args.steps, 1e-6), 4) # 9.8696.
	report["pi_squared"] = round(math.pi ** 2, 4) # The squared arc length of a quarter circle.
	report["geodesic_kappa_0.25"] = round(geodesic_schedule(0.25), 4) # 0.1464.
	report["geodesic_kappa_0.75"] = round(geodesic_schedule(0.75), 4) # 0.8536.
	report["geodesic_rate_0.9"] = round(math.pi * math.tan(math.pi * 0.9 / 2), 4) # 19.84.
	report["uniform_source_kappa_0.5"] = round(geodesic_schedule(0.5, 0.25), 4) # 0.6667.
	print_report("kinetic-optimal paths", report)
	save_run(args.out, config=vars(args), report=report)


	if __name__ == "__main__":
	main()